using DelimitedFiles
using PyPlot
using RAFF

"""

    draw_problem(M; raff_output=nothing, model_str="logistic",
                 datafile="/tmp/output.txt")

    draw_problem(datafile::String="/tmp/output.txt"; kwargs...)

Draw the problem data given by a (`m`x`3`) `M` matrix. By default it
is assumed that the model is given by the logistic model. If no
arguments are given it is assumed that the data file is given in file
`/tmp/output.txt`.

If a [RAFFOutput](@ref) object is provided, then it plots the model
found and also the outliers (true and false positives).

Optional arguments:

  - `raff_output`: [RAFFOutput](@ref) object with the solution
    obtained
  - `model_str`: a string with the name of the model to be used to
    plot the solution. See [model_list](@ref) for details.

"""
function draw_problem(M; raff_output=nothing, model_str="logistic",
                      θsol=nothing)

    x = M[:, 1]
    y = M[:, 2]
    co = M[:, 3]

    true_outliers = findall(co .!= 0.0)

    PyPlot.scatter(x[co .== 0.0], y[co .== 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="o", s=50.0, linewidths=0.2)

    PyPlot.scatter(x[co .!= 0.0], y[co .!= 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="^", s=50.0, linewidths=0.2, label="Outliers")

    if raff_output != nothing
        
        n, model, modelstr = RAFF.model_list[model_str]

        modl1 = (x) -> model(x, raff_output.solution)

        t = minimum(x):0.01:maximum(x)
        PyPlot.plot(t, modl1.(t), color=PyPlot.cm."Set1"(2.0/9.0))

        # Draw outliers found by RAFF
        
        true_positives = intersect(true_outliers, raff_output.outliers)
        false_positives = setdiff(raff_output.outliers, true_positives)

        PyPlot.scatter(x[false_positives], y[false_positives],
                       color=PyPlot.cm."Pastel1"(0.0/9.0), marker="o",
                       linewidths=0.2, edgecolors="k", s=50.0, label="False positives")
        
        PyPlot.scatter(x[true_positives], y[true_positives],
                       color=PyPlot.cm."Pastel1"(0.0/9.0), marker="^",
                       s=50.0, linewidths=0.2, edgecolors="k", label="Identified outliers")

    end

    # Plot the true solution, if available
    if θsol != nothing
        
        modl2 = (x) -> model(x, θsol)

        PyPlot.plot(t, modl2.(t), color=PyPlot.cm."Set1"(1.0/9.0),
                    linestyle="--")

    end

    PyPlot.legend(loc="best")

    PyPlot.show()
    
    PyPlot.savefig("/tmp/figure.png", DPI=150)

end

function draw_problem(datafile::String="/tmp/output.txt"; kwargs...)

    fp = open(datafile, "r")

    N = parse(Int, readline(fp))

    M = readdlm(fp)

    close(fp)

    draw_problem(M; kwargs...)
    
end

function draw_comparison(M::Array{Float64, 2}, model_str, algorithms...)

    x = M[:, 1]
    y = M[:, 2]
    co = M[:, 3]

    true_outliers = findall(co .!= 0.0)

    PyPlot.scatter(x[co .== 0.0], y[co .== 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="o", s=50.0, linewidths=0.2)

    PyPlot.scatter(x[co .!= 0.0], y[co .!= 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="^", s=50.0, linewidths=0.2, label="Outliers")


    t = minimum(x):0.01:maximum(x)

    n, model, modelstr = RAFF.model_list[model_str]

    for (i, alg) in enumerate(algorithms)

        name, line, sol = alg
        
        modl1 = (x) -> model(x, sol)

        PyPlot.plot(t, modl1.(t), color=PyPlot.cm."Set1"(i/9.0),
                    linestyle=line, label=name)

        PyPlot.legend(loc="best")

        PyPlot.show()

        # PyPlot.savefig(outimage, DPI=600, bbox_inches="tight")

    end

end

function draw_comparison(model_str::String, algorithms...)

    fp = open("/tmp/output.txt", "r")

    N = parse(Int, readline(fp))

    M = readdlm(fp)

    close(fp)

    draw_comparison(M, model_str, algorithms...)
    
end

"""

Draw the points generated by the previous function.

"""
function draw_circle(data, outliers)

    np, = size(data)
    
    c = zeros(np)
    
    c[outliers] .= 1.0
    
    PyPlot.scatter(data[:, 1], data[:, 2], c=c, marker="o", s=50.0, linewidths=0.2,
                   cmap=PyPlot.cm["Paired"], alpha=0.9)
    
    PyPlot.axis("scaled")

    PyPlot.xticks([])

    PyPlot.yticks([])

    PyPlot.savefig("/tmp/circle.png", dpi=150, bbox_inches="tight")

end

"""

Draw the points and the solutions obtained. Save the picture in a file.

"""
function draw_circle_sol(M; model_str="circle", raff_output=nothing, other_sols...)

    x  = M[:, 1]
    y  = M[:, 2]
    co = M[:, 4]

    t = 0:0.1:2.1 * π
    
    ptx = (α, ρ, d) -> ρ * cos(α) + d[1]
    pty = (α, ρ, d) -> ρ * sin(α) + d[2]    

    # Plot data

    true_outliers = findall(co .!= 0.0)

    PyPlot.scatter(x[co .== 0.0], y[co .== 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="o", s=50.0, linewidths=0.2)

    PyPlot.scatter(x[co .!= 0.0], y[co .!= 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="^", s=25.0, linewidths=0.2, label="Outliers")

    if raff_output != nothing
        
        n, model, modelstr = RAFF.model_list[model_str]

        fSol = raff_output.solution
        
        modl1x = (α) -> ptx(α, fSol[3], fSol[1:2])
        modl1y = (α) -> pty(α, fSol[3], fSol[1:2])

        PyPlot.plot(modl1x.(t), modl1y.(t), color=PyPlot.cm."Set1"(2.0/9.0))

        # Draw outliers found by RAFF
        
        true_positives = intersect(true_outliers, raff_output.outliers)
        false_positives = setdiff(raff_output.outliers, true_positives)

        if length(false_positives) > 0
        
            PyPlot.scatter(x[false_positives], y[false_positives],
                           color=PyPlot.cm."Pastel1"(0.0/9.0), marker="o",
                           linewidths=0.2, edgecolors="k", s=50.0, label="False positives")

        end

        if length(true_positives) > 0
        
            PyPlot.scatter(x[true_positives], y[true_positives],
                           color=PyPlot.cm."Pastel1"(0.0/9.0), marker="^",
                           s=50.0, linewidths=0.2, edgecolors="k", label="Identified outliers")

        end

    end

    PyPlot.legend(loc="best")

    PyPlot.axis("scaled")

    PyPlot.xticks([])

    PyPlot.yticks([])

    PyPlot.savefig("/tmp/circle.png", dpi=150, bbox_inches="tight")
    
end

function draw_circle_sol(tSol, fSol, lsSol)

    datafile = "/tmp/output.txt"

    fp = open(datafile, "r")

    N = parse(Int, readline(fp))

    M = readdlm(fp)

    close(fp)

    x  = M[:, 1]
    y  = M[:, 2]
    ρ  = M[:, 3]
    co = M[:, 4]

    t = [0:0.1:2.1 * π;]
    
    ptx = (α, ρ, d) -> ρ * cos(α) + d[1]
    pty = (α, ρ, d) -> ρ * sin(α) + d[2]

    # True solution
    
    pptx = (α) -> ptx(α, tSol[3], tSol[1:2])
    ppty = (α) -> pty(α, tSol[3], tSol[1:2])

    PyPlot.plot(pptx.(t), ppty.(t), "b--", label="True solution")
    
    # RAFF solution
    
    pptx = (α) -> ptx(α, fSol[3], fSol[1:2])
    ppty = (α) -> pty(α, fSol[3], fSol[1:2])

    PyPlot.plot(pptx.(t), ppty.(t), "g-", label="RAFF")
    
    # # LS solution
    
    # pptx = (α) -> ptx(α, lsSol[3], lsSol[1:2])
    # ppty = (α) -> pty(α, lsSol[3], lsSol[1:2])

    # PyPlot.plot(pptx.(t), ppty.(t), "r-", label="Least squares")

    PyPlot.scatter(x[co .== 0.0], y[co .== 0.0], color=PyPlot.cm."Pastel1"(2.0/9.0),
                   marker="o", s=50.0, linewidths=0.2)

    PyPlot.scatter(x[co .!= 0.0], y[co .!= 0.0], color=PyPlot.cm."Pastel1"(1.0/9.0),
                   marker=".", s=25.0, linewidths=0.2, label="Outliers")

    PyPlot.legend(loc=4)
    
    PyPlot.axis("scaled")

    PyPlot.xticks([])

    PyPlot.yticks([])

    PyPlot.savefig("/tmp/circle.png", dpi=150, bbox_inches="tight")

end

"""

    draw_img_sol(imgfile, sol; thck::Int=2)

Draw the image and the solution found, using JuliaImage
packages. `imagefile` is the path to the image and `sol` is a
3-dimensional vector with the solution.

Extra parameter `thck` defines the thickness of the circle.

"""
function draw_img_sol(imgfile, sol; thck::Int=2)

    img = load(imgfile)

    h, w = size(img)
    
    t = [0:0.01:2.1 * π;]
    
    ptx = (α, ρ, d) -> ρ * cos(α) + d[1]
    pty = (α, ρ, d) -> ρ * sin(α) + d[2]

    # Solution
    
    pptx = (α) -> ptx(α, sol[3], sol[1:2])
    ppty = (α) -> pty(α, sol[3], sol[1:2])

    for α in t

        i = Int(round(pptx(α)))
        j = Int(round(ppty(α)))

        !((1 <= i <= h) && (1 <= j <= w)) && continue
        
        img[max(1, i - thck):min(i + thck, h), max(1, j - thck):min(j + thck, w)] .= RGB(1.0, 0.0, 0.0)

    end
    
    ImageView.imshow(img)
    
end
