##
#
# Visualization and related functions
#


"""

    RootsImage(leftbottom, topright, width, height)

Essentially a two dimensional histogram covering a rectangle in the
complex plane specified by complex numbers `leftbottom` and `topright`
with `width` horizontal and `height` vertical bins.
"""
struct RootsImage{T <: Real}
  remin::T
  remax::T
  immin::T
  immax::T
  height::Int
  width::Int
  hbin::T
  wbin::T
  data::Array{T, 2}

  function RootsImage(leftbottom::Complex{T}, topright::Complex{T}, width::Int, height::Int) where {T <: Real}
    remin = real(leftbottom)
    immin = imag(leftbottom)
    remax = real(topright)
    immax = imag(topright)
    new{T}(remin, remax, immin, immax, height, width, (remax - remin) / width, (immax - immin) / height, zeros(T, width, height))
  end
end

"""

    add_root!(img, z)

Add a root `z` to the image `img`.
"""
function add_root!(img::RootsImage{T}, z::Complex{T}) where {T <: Real}
  r = real(z); i = imag(z)

  if r < img.remin || r >= img.remax || i < img.immin || i >= img.immax
    # out of range, skipping...
    return
  end

  x = min(max(1, 1 + floor(Int64,(r - img.remin) / img.hbin)), img.width)
  y = min(max(1, 1 + floor(Int64,(i - img.immin) / img.wbin)), img.height)

  img.data[x, y] += T(1)
end


"""

    plot(img[; mode])

Show a graphical representation of the image `img`.
The optional parameter `mode` controls how to scale number of roots in
histogram bins:

 * `:sharp`,
 * `:log`,
 * `:log_cutoff`,
 * `:linear`,
 * a function.

"""
function plot(img::RootsImage{T}; mode=:sharp) where {T <: Real}
  if isa(mode, Function)
    output = mode.(img.data)
  elseif mode == :sharp
    output = (x -> min(1, x)).(img.data)
  elseif mode == :log_cutoff
    m      = mean(img.data)
    output = (x -> log(1 + min(2m, x))).(img.data)
  elseif mode == :log
    output = (x -> log(1 + x)).(img.data)
  elseif mode == :linear
    output = img.data
  else
    error("Unknown mode!")
  end

  # normalization to [0, 1]
  output = output / maximum(output)

  return colorview(Gray, transpose(output)[end:-1:1,1:end])
end


"""

    save_image(img, filename[; mode])

Save the `img` to file `filename`.
Optional parameter `mode` has the same meaning as in the [`plot`](@ref) function.

"""
function save_image(img::RootsImage{T}, filename::AbstractString; mode=:log_cutoff, latex_filename=nothing) where {T <: Real}
  @info "Saving the image $filename..."

  if isa(mode, Function)
    output = mode.(img.data)
  elseif mode == :sharp
    output = (x -> min(1, x)).(img.data)
  elseif mode == :log_cutoff
    m      = mean(img.data)
    output = (x -> log(1 + min(2m, x))).(img.data)
  elseif mode == :log
    output = (x -> log(1 + x)).(img.data)
  elseif mode == :linear
    output = img.data
  else
    error("Unknown mode!")
  end

  output = output / maximum(output)

  Images.save(filename, colorview(Gray, transpose(output)[end:-1:1,1:end]))

  if latex_filename != nothing
    @info "Constructing the LaTeX file $latex_filename..."
    write_latex(img, filename, latex_filename)
  end
end


"""

    fill_images!(imgs, filename)

Read roots from file `filename` and draw them onto images contained
in the array `imgs`.  

"""
function fill_images!(imgs::Array{RootsImage{T}}, filename::AbstractString) where {T <: Real}
  roots_count = div(filesize(filename), sizeof(Complex{T})) + 1
  progress = Progress(roots_count, 10)

  @info "Sifting through data (there are $roots_count roots)..."

  open(filename) do io
    while !eof(io)
      z = read(io, Complex{T})
      for img in imgs
        add_root!(img, z)
      end
      next!(progress)
    end
  end
end

"""

  animate_growth(poly_iter, img, filename[; output_filename])

Take roots of polynomials generated by `poly_iter` and saved in file `filename`
and draw them incrementaly (each step corresponds to change of degree) onto
image `img`.

"""
function animate_growth(poly_iter::PolynomialIterator{S}, img::RootsImage{T}, filename; output_filename=j->"animation_$j.png") where {S <: Number, T <: Real}
  progress = Progress(div(filesize(filename), sizeof(Complex{T})) + 1, 10)
  degree = 0

  open(filename) do io
    for p in poly_iter
      # save the current snapshot
      if degree != length(p) - 1
        save_image(img, output_filename(degree))
        degree += 1
      end

      # read roots
      for k in 1:degree
        z = read(io, Complex{T})
        add_root!(img, z)
        next!(progress)
      end
    end
  end 
end


"""

    fill_image!(img, filename)

An analog of [`fill_images!`](@ref), only with one image.

"""
function fill_image!(img::RootsImage{T}, filename::AbstractString) where {T <: Real}
  fill_images!([img], filename)
end

"""

  write_latex(img, imagename, latex_filename)

Create a LaTeX PGFPlots version of the image with axes.

"""
function write_latex(img::RootsImage{T}, imagename::AbstractString, latex_filename::AbstractString) where {T <: Real}
  open(latex_filename, "w") do io
    write(io, "\\documentclass{standalone}\n")
    write(io, "\\usepackage[T1]{fontenc}\n")
    write(io, "\\usepackage[utf8]{inputenc}\n")
    write(io, "\\usepackage{pgfplots}\n")
    write(io, "\\usepackage{lmodern}\n")
    write(io, "\\begin{document}\n")
    write(io, "\\begin{tikzpicture}\n")
    write(io, "\\begin{axis}[enlargelimits=false, axis on top, xlabel={\$\\Re\$}, ylabel={\$\\Im\$}]\n")
    write(io, "\\addplot graphics [xmin=$(img.remin), xmax=$(img.remax), ymin=$(img.immin), ymax=$(img.immax)] {$imagename};\n")
    write(io, "\\end{axis}\n")
    write(io, "\\end{tikzpicture}\n")
    write(io, "\\end{document}\n")
  end
end
