using EntanglementDist
using SQLite
using Gadfly
using DataFrames

import Base.convert
convert(x::Type{Nullable{Float64}}, y::Int64) = convert(Nullable{Float64}, convert(Float64, y))

# The plots generated by this script are stored in:
# - `./pdfplots` in PDF format
# - `./plots` in PNG format

### settings, change these for different plots ###
# This can be any parameter, it depends on how it has been saved in the database
ps = collect(0.4:0.1:1)

# The collection of states that are to be plotted
states = AbstractString["pptRelax: Double Werner"]

# The database where the data is stored
db = SQLite.DB("test.sqlite")

# Title of the plot as a function of the parameter
plottitle(p) = "Double Werner (p = $p)"

# Due to a bug in SQLite it might be the case that
# an error is thrown because a certain column in the database
# includes both floats and integers. Untill this is resolved
# by the SQLite.jl team, one might resort to adding a very
# small number (e.g. 10^-8) to the relevant columns. Typically
# this happens to delta_min if `0`, `1` are present.
#
# Example query:
# table = SQLite.query(db, "UPDATE `RainsProb` SET `delta_max` = `delta_max` + 0.00000001")
for p in ps
  for state in states
    # calculate the performance of the protocols
    if state == "pptRelax: Double Werner" || state == "Triple Werner"
      stateFn = wernerState
    elseif state == "Double Ronald"
      stateFn = sState
    elseif state == "Double Ronald 2"
      stateFn = rState
    elseif state == "Ronald 2 averaged phase" || state == "ronald2StateCorrPhase"
      stateFn = rStateCorrPhase
    end

    # get data from DB
    q = string("SELECT * FROM `pptRelax` WHERE
      (`state` = '", state, "') AND
      `p` = '", p, "'")
      println(q)
    table = SQLite.query(db, q)
    # get rid of the nullables introduced by SQLite,
    # as they cannot be interpreted by the plotting
    # software
    for c in 1:ncol(table)
      table[c] = table[c].values
    end
    PPTdf = copy(table)

    q = string("SELECT * FROM `pptRelax` WHERE
      (`state` = '1ext ", state, "') AND
      `p` = '", p, "'")
      println(q)
    table = SQLite.query(db, q)
    # get rid of the nullables introduced by SQLite
    for c in 1:ncol(table)
      table[c] = table[c].values
    end
    kextdf = copy(table)

    q = string("SELECT * FROM `pptRelax` WHERE
      (`state` = '1ext ", state, "') AND
      `p` = '", p, "'")
      println(q)
    table = SQLite.query(db, q)
    # get rid of the nullables introduced by SQLite
    for c in 1:ncol(table)
      table[c] = table[c].values
    end
    onlySymdf = copy(table)


    PPTLayer = layer(PPTdf,
      x=:p_succ,
      y=:fidelity,
      Geom.line,
      # Geom.point,
      Theme(default_color=colorant"blue"))

    layers = typeof(PPTLayer)[]
    names = AbstractString[]
    colors = AbstractString[]
    if nrow(PPTdf) > 0 # only include this layer if it is nonempty
      push!(layers, PPTLayer)
      push!(names, "PPT")
      push!(colors, "blue")
    end

    if nrow(kextdf) > 0 && false # only include this layer if it is nonempty
      kextLayer = layer(kextdf,
        x=:p_succ,
        y=:fidelity,
        Geom.point,
        Theme(default_color=colorant"red"))
      push!(layers, kextLayer)
      push!(names, "1-ext")
      push!(colors, "red")
    end

    if nrow(onlySymdf) > 0 # only include this layer if it is nonempty
      onlySymLayer = layer(onlySymdf,
      x=:p_succ,
      y=:fidelity,
      Geom.point,
      Theme(default_color=colorant"green"))
      push!(layers, onlySymLayer)
      push!(names, "1-ext")
      push!(colors, "green")
    end
    
    # Add for some input states the EPL protocol
      F, p_succ = EPLParam(sortAB(copies(stateFn(p),2),2,2))
      eplLayer = layer(y = [F], x = [p_succ], Geom.point, Theme(default_color=colorant"purple"))
      unshift!(layers, eplLayer)
      unshift!(names, "EPL protocol")
      unshift!(colors, "purple")
    
    
    # Include schemes in plots
    F, p_succ = DEJMPSParam(stateFn(p))
    deutschLayer = layer(y = [F], x = [p_succ], Geom.point, Theme(default_color=colorant"cyan"))
    F, p_succ = BBPSSWParam(stateFn(p))
    bennettLayer = layer(y = [F], x = [p_succ], Geom.point, Theme(default_color=colorant"orange"))

    unshift!(layers, deutschLayer)
    unshift!(layers, bennettLayer)
    unshift!(names, "DEJMPS protocol")
    unshift!(names, "BBPSSW protocol")
    unshift!(colors, "cyan")
    unshift!(colors, "orange")

    pl = plot(
      Theme(panel_fill=colorant"white"),
      layers...,
      Stat.xticks(ticks=collect(0:0.1:1)),
      Stat.yticks(ticks=collect(0.5:0.1:1)),
      Guide.xlabel("Probability of success"),
      Guide.ylabel("Fidelity"),
      Guide.title(plottitle(p)),
      Guide.manual_color_key("Legend",
        names,
        colors)
        )

    filename = replace(replace(string(state,p), " ", ""), ".", "")

    draw(PNG(string("plots/", filename ,".png"), 600px, 300px), pl)
    draw(PDF(string("pdfplots/", filename ,".pdf"), 4inch, 3inch), pl)
  end
end
