# Based on https://github.com/cloud-oak/Tqdm.jl by @cloud-oak under Mozilla Public License 2.0
# Modified for Knet by Deniz Yuret

using Printf
import Base: length, size, iterate, eltype, IteratorSize, IteratorEltype, haslength, @propagate_inbounds

"""
    progress(itr; width, alpha, interval)

Return an iterator which acts exactly like `itr`, but prints a progressbar as new values are
requested:

    2.70e-01  21.83%┣███▉              ┫ 13101/60000 [00:12/00:53, 1137.05i/s]

Here `2.70e-01` is the exponential average of values generated by itr (only displayed for
iterators with numeric values). `21.83%` is the percentage, `13101` is the number of
iterations completed, `60000` is the total number of iterations. `00:12` is elapsed seconds,
`00:53` is the estimated total seconds, `1137.05i/s` is the average number of iterations
completed per second. If the speed is less than 1, the average number of seconds per
iteration (s/i) is reported instead.  The percent, total iterations, and completion time are
omitted for iterators whose size is unknown.

`progress!(...)` is equivalent to `(for x in progress(...) end)`, i.e.  iterates over the
object created by `progress(...)` and returns `nothing`.

An integer `itr` is treated as `1:itr`, i.e. `progress(n::Integer)` is equivalent to
`progress(1:n)`

Keyword arguments:
* `width=max(64,displaysize()[2])`: controls display width. The default width can be
  controlled using `ENV["COLUMNS"]`.
* `interval=1.0`: minimum time interval in seconds between progressbar updates.
* `alpha=1.0`: controls the exponential average displayed for numeric iterators: 
  `avg = alpha * val + (1-alpha) * avg`
"""
progress, progress!

mutable struct Progress{I}
    iter::I
    current::Int
    nprint::Int
    start_time::UInt
    print_time::UInt
    print_interval::UInt
    width::Int
    alpha::Float64
    avg::Float64
end

progress(iter::I; width=max(64,displaysize()[2]), alpha=1.0, interval=1.0) where {I} =
    Progress{I}(iter,0,0,time_ns(),0,Int(1e9*interval),width,alpha,Inf)

progress(i::Integer; o...)=progress(1:n; o...)
progress!(x...; o...)=(for x in progress(x...; o...) end)

length(p::Progress) = length(p.iter)
size(p::Progress) = size(p.iter)
eltype(p::Progress) = eltype(p.iter)
IteratorSize(::Type{Progress{I}}) where {I} = IteratorSize(I)
IteratorEltype(::Type{Progress{I}}) where {I} = IteratorEltype(I)

@propagate_inbounds function iterate(p::Progress, s...)
    next = iterate(p.iter, s...)
    if next !== nothing
        p.current += 1
        (x, s) = next
        if p.alpha > 0 && value(x) isa Number
            p.avg = (p.avg === Inf ? value(x) : p.alpha * value(x) + (1-p.alpha) * p.avg)
        end
    end
    display_progress(p, next === nothing)
    return next
end

function display_progress(p::Progress, last=false)
    curr_time = time_ns()
    if !last && (curr_time < p.print_time + p.print_interval)
        return
    end
    p.print_time = curr_time
    p.nprint += 1
    seconds   = (curr_time - p.start_time) * 1e-9
    speed     = p.current / seconds
    

    fval_string = isfinite(p.avg) ? @sprintf("%.2e  ", p.avg) : ""

    if haslength(p)
        ETA = (length(p) - p.current) / speed
        percentage_string = string(@sprintf("%.2f%%",p.current/length(p)*100))
        status_string = string(p.current, "/", length(p),
                               " [", format_time(seconds), "/", format_time(seconds+ETA),
                               ", ", format_speed(speed),"]")

    else
        ETA = Inf
        percentage_string = ""
        status_string = string(p.current, # "/", length(p),
                               " [", format_time(seconds), # "-", format_time(ETA),
                               ", ", format_speed(speed),"]")
    end

    print("\r")
    print(fval_string)
    print(percentage_string)
    print("┣")

    width = p.width - length(fval_string) - length(percentage_string) - length(status_string) - 2

    if (haslength(p))
        cellvalue = length(p) / width
        full_cells, remain = divrem(p.current, cellvalue)
        full_cells = round(Int, full_cells)
        print(repeat("█", full_cells))
        if (full_cells < width)
	    part = floor(Int, 8 * remain / cellvalue)
	    print(EIGHTS[part])
            print(repeat(" ", width - full_cells - 1))
        end
    else
        offset = p.nprint % 8
        print(repeat(" ", offset))
        print("/")
        segments, remain = divrem(width - offset - 1, 8)
        print(repeat("       /", Int(segments)))
        print(repeat(" ", Int(remain)))
    end
    print("┫ ")

    print(status_string)
    last && println()
end

function format_time(seconds)
    if seconds != Inf
        mins,s  = divrem(round(Int,seconds), 60)
        h, m    = divrem(mins, 60)
    else
        h=0;m=Inf;s=Inf
    end
    if h!=0
         return @sprintf("%02d:%02d:%02d",h,m,s)
    else
         return @sprintf("%02d:%02d",m,s)
    end
end

format_speed(s)=(s >= 1 ? @sprintf("%.2fi/s",s) : @sprintf("%.2fs/i",1/s))

EIGHTS = Dict(0 => ' ',
	      1 => '▏',
	      2 => '▎',
	      3 => '▍',
	      4 => '▌',
	      5 => '▋',
	      6 => '▊',
	      7 => '▉',
	      8 => '█')
