module Pixels

using Images
using ArgParse

include("ui.jl")
include("bag.jl")
include("treap.jl")
include("sss.jl")

s = ArgParseSettings()
@add_arg_table s begin
    "--file", "-f"
        help = "filename for image to use as input"
        arg_type = String
        default = "images/Bachalpseeflowers.jpg"
    "--type", "-t"
        help = "function type to use: min, avg, or evo"
        arg_type = String
        default = "min"
end
parsed_args = parse_args(ARGS, s)
const target_file = parsed_args["file"]

const fntype = symbol(parsed_args["type"])
(fntype in [:min,:avg,:evo]) || error("Invalid function type provide. Got $fntype; expected min, avg, or evo.")

immutable NeighborStats
    n::Int
    avg::Rgb{Uint8}
    sumvar::Uint
    sqrt_sumvar::Uint8
    NeighborStats(n, avg, sumvar) = new(n, avg, sumvar, uint8(sqrt(sumvar)))
    NeighborStats() = new(0, Rgb{Uint8}(), 0, 0)
end

type Vertex
    x::Int
    y::Int
    color::GLPixel
    filled::Bool
    nominated::Bool
    bagindex::Int
    stats::NeighborStats
    neighbors::Vector{Vertex}
    t::Float64 # Used to preserve key uniqueness and break ties in <(Vertex, Vertex)
    function Vertex(x, y)
        neighbors = Vertex[]
        sizehint(neighbors, 8)
        new(x, y, GLPixel(), false, false, 0, NeighborStats(), neighbors, rand(Float64))
    end
end

Base.start(p::Rgb) = 1
Base.next{T}(p::Rgb{T}, state::Int) =
    state == 1 ? (p.r, 2) :
    state == 2 ? (p.g, 3) :
    state == 3 ? (p.b, 4) : throw("Iteration state error.")
Base.done(p::Rgb, state::Int) = state == 4

type Canvas
    grid::Array{Vertex}
    frontier
    function Canvas(width, height, frontier)
        grid = [Vertex(x, y) for y in 1:height, x in 1:width]

        for x in 1:width, y in 1:height
            for nx in max(x - 1, 1):min(x + 1, width), ny in max(y - 1, 1):min(y + 1, height)
                !(nx == x && ny == y) && push!(grid[y, x].neighbors, grid[ny, nx])
            end
        end
        new(grid, frontier)
    end
end
frontier(canvas) = canvas.frontier
Base.getindex(c::Canvas, a, b) = c.grid[a, b]
neighbors(vertex::Vertex) = vertex.neighbors
isnominated(vertex::Vertex) = vertex.bagindex != 0 || vertex.nominated # TODO: Fix.

function nominate!(vertex, frontier::Bag{Vertex})
    @assert !isnominated(vertex)

    # Bag add
    push!(frontier, vertex)
    vertex.bagindex = length(frontier)
end

function withdraw!(vertex, frontier::Bag{Vertex})
    @assert isnominated(vertex)

    # Bag remove
    index = vertex.bagindex
    deleteat!(frontier, index)
    length(frontier) >= index && (frontier[index].bagindex = index)
    vertex.bagindex = 0
end

function nominate!(vertex, frontier::Treap)
    @assert !vertex.nominated
    add!(frontier, vertex)
    vertex.nominated = true
end

function withdraw!(vertex, frontier::Treap)
    @assert vertex.nominated
    remove!(frontier, vertex)
    vertex.nominated = false
end

hasemptyneighbors(vertex) = vertex.stats.n < length(neighbors(vertex))
emptyneighbors(vertex) = filter((v) -> !v.filled, neighbors(vertex))
function randemptyneighbor(vertex)
    prob = 1.0
    local neighbor
    for v in neighbors(vertex)
        if !v.filled
            if rand() < prob
                neighbor = v
            end
            prob /= 2
        end
    end
    neighbor
end

function fill!(pixels, frontier, vertex, color, outer::Bool)
    # TODO: Don't pass pixels in.
    pixels[vertex.y, vertex.x] = color

    vertex.color = color
    vertex.filled = true
    if outer
        # Outer frontier
        for v in neighbors(vertex)
            # Withdraw in order to re-insert neighbors already on the frontier,
            # as their stats have changed in a way that may influence their
            # storage in the frontier
            isnominated(v) && withdraw!(v, frontier)
            stats!(v)
            !isnominated(v) && !v.filled && nominate!(v, frontier)
        end
        isnominated(vertex) && withdraw!(vertex, frontier)
    else
        # Inner frontier
        for v in neighbors(vertex)
            # NOTE: Does something like the re-insertion above need to happen here as well?
            stats!(v)
            isnominated(v) && !hasemptyneighbors(v) && withdraw!(v, frontier)
        end
        hasemptyneighbors(vertex) && nominate!(vertex, frontier)
    end
end

# TODO: To calculate average without overflow: (R - L)/2 + L
function stats!(vertex)
    n = zero(vertex.stats.n)
    avg = var = Rgb{Float64}() #zero(vertex.stats.avg)

    for v in neighbors(vertex)
        if v.filled
            n += 1
            avg += v.color
            # Promote types before squaring, since Rgb{Uint8} * Rgb{Uint8} => Rgb{Uint8}
            var += oftype(var, v.color) * v.color
        end
    end

    if n == 0
        vertex.stats = NeighborStats()
    else
        avg = avg / n
        var = var / n - avg * avg
        vertex.stats = NeighborStats(n, Rgb{Uint8}(uint8(avg.r), uint8(avg.g), uint8(avg.b)), uint(sum(var)))
    end
end

r() = rand() - 0.5
mutate(color::GLPixel, scale) = GLPixel(
    clamp(color.r + int(scale * r()), typemin(color.r), typemax(color.r)),
    clamp(color.g + int(scale * r()), typemin(color.g), typemax(color.g)),
    clamp(color.b + int(scale * r()), typemin(color.b), typemax(color.b))
)

function nearest(frontier::Bag{Vertex}, color)
    best = frontier[1]
    best_dist::Float64 = Inf
    for v in frontier
        d = sqdiff(v.stats.avg, color) + v.stats.sumvar
        if d < best_dist
            best = v
            best_dist = d
        end
    end
    best
end

function loadimage(imgpath::String)
    im = convert(Array, imread(imgpath))
    w, h = size_spatial(im)
    colors = Pixel[]
    for i in 1:w, j in 1:h
        push!(colors, Pixel(uint8(im[i, j].r * 255), uint8(im[i, j].g * 255), uint8(im[i, j].b * 255)))
    end
    colors
end

SPEED = 1000

function evolve(pixels)
    (height, width) = size(pixels)
    canvas = Canvas(width, height, Bag{Vertex}(height * width))

    fill!(pixels, frontier(canvas), canvas[div(height, 2), div(width, 2)], GLPixel(255, 255, 255), true)
    for i in 1:length(canvas.grid) - 1
        vertex = frontier(canvas)[rand(1:length(frontier(canvas)))]
        color = filter(n -> n.filled, neighbors(vertex))[1].color
        fill!(pixels, frontier(canvas), vertex, mutate(color, 15), true)
        i % SPEED == 0 && produce(pixels)
    end
    pixels
end

function gencolors(imgpath::String)
    colors = shuffle!(loadimage(imgpath))
    sort!(colors, alg=QuickSort, lt=(p::Pixel, q::Pixel) -> hue(p) < hue(q))
    colors
end

function gencolors(n::Int)
    colors = [rand(GLPixel) for i in 1:n]
    # sort!(colors, alg=QuickSort, lt=(p::Pixel, q::Pixel) -> hue(p) < hue(q))
    sort!(colors, rev=true, alg=QuickSort, lt=(p::Pixel, q::Pixel) -> p.r < q.r && p.g < q.g && p.b < q.b)
    # sort!(colors, rev=true, alg=QuickSort, by=(p::Pixel) -> p.b + p.g^2)
    colors
end

function placeavg(pixels)
    (height, width) = size(pixels)
    canvas = Canvas(width, height, Treap{Vertex}())
    # colors = gencolors(length(pixels))
    colors = gencolors(target_file)
    # colors = gencolors(1000)
    i = 1
    fill!(pixels, frontier(canvas), canvas[div(height, 2), div(width, 2)], colors[i += 1], true)
    while i < length(pixels)
        color = colors[mod1(i, length(colors))]
        vertex = nearest(frontier(canvas), color)
        fill!(pixels, frontier(canvas), vertex, color, true)
        i % SPEED == 0 && produce(pixels)
        i += 1
    end
    pixels
end

function placemin(pixels)
    (height, width) = size(pixels)
    canvas = Canvas(width, height, Treap{Vertex}())
    # colors = gencolors(length(pixels))
    # colors = gencolors("Bachalpseeflowers.jpg")
    colors = gencolors(target_file)
    i = 1

    fill!(pixels, frontier(canvas), canvas[div(height, 2), div(width, 2)], colors[i += 1], false)
    # fill!(pixels, frontier(canvas), canvas[200, 200], colors[i += 1], false)
    # fill!(pixels, frontier(canvas), canvas[1, 1], colors[i += 1], false)

    while i < length(pixels)
        color = colors[mod1(i, length(colors))]
        candidate = nearest(frontier(canvas), color)
        vertex = randemptyneighbor(candidate)
        fill!(pixels, frontier(canvas), vertex, color, false)
        i % SPEED == 0 && produce(pixels)
        i += 1
    end
    pixels
end

fn = fntype == :min ? placemin :
     fntype == :avg ? placeavg :
     fntype == :evo ? evolve : throw("fntype must be :min, :avg, or :evo.")

if fntype == :min || fntype == :evo
    # Minimum selection
    Base.size(p::Vertex, n) = n == 1 ? 3 : error("Vertices are one-dimensional.")
    Base.getindex(p::Vertex, n) = p.color[n]
    <(p::Vertex, q::Vertex) = p.color == q.color ? p.t < q.t : shuffless(p.color, q.color)
    <(p::Vertex, q::Rgb) = shuffless(p.color, q)
    <(p::Rgb, q::Vertex) = shuffless(p, q.color)
    dist_sq{T}(p::Rgb{T}, q::Rgb{T}) = sqdiff(p, q)
    dist_sq{T}(p::Vertex, q::Rgb{T}) = sqdiff(p.color, q)
    dist_sq{T}(p::Rgb{T}, q::Vertex) = sqdiff(p, q.color)
else
    # Average selection
    Base.size(p::Vertex, n) = n == 1 ? 4 : error("n must be 1.")
    Base.getindex(p::Vertex, n) = n < 4 ? uint8(p.stats.avg[n]) : n == 4 ? uint8(p.stats.sqrt_sumvar) : error("Index out of bounds: $n")
    # If the colors are equal, compare t values; we need this because the treap enforces uniqueness.
    <(p::Vertex, q::Vertex) = shuffless(p, q) ? true : shuffmore(p, q) ? false : p.t < q.t
    <(p::Vertex, q::Rgb) = shuffless(p, q)
    <(p::Rgb, q::Vertex) = shuffless(p, q)
    dist_sq{T}(p::Rgb{T}, q::Rgb{T}) = sqdiff(p, q)
    dist_sq{T}(p::Vertex, q::Rgb{T}) = sqdiff(p.stats.avg, q) + p.stats.sumvar
    dist_sq{T}(p::Rgb{T}, q::Vertex) = sqdiff(p, q.stats.avg) + q.stats.sumvar

    # Avg color: Hacked version with p[4] == 0
    Base.start(p::Rgb) = 1
    Base.next{T}(p::Rgb{T}, state::Int) =
        state == 1 ? (p.r, 2) :
        state == 2 ? (p.g, 3) :
        state == 3 ? (p.b, 4) :
        state == 4 ? (zero(T), 5) : throw("Iteration state error.")
    Base.done(p::Rgb, state::Int) = state == 5
    Base.getindex{T}(p::Rgb{T}, n::Int) = n == 1 ? p.r : n == 2 ? p.g : n == 3 ? p.b : n == 4 ? zero(T) : throw("Color indexing error.")
    Base.size(p::Rgb) = (4,)
    Base.size(p::Rgb, n) = n == 1 ? 4 : throw("Invalid dimension.")
end

# placemin(zeros(GLPixel, 100, 100))
display(1000, 750, fn; title="Pixels")

end


# see results in README.md

using BenchmarkTools
using PEG # this package
using StringParserPEG # another similar package
using JSON # hand-written JSON parser for comparison

# PEG grammar
commasPEG(head, tail) = [head, (t[2] for t ∈ tail)...]
bracketsPEG(left, middle, right) = (isempty(middle) ? [] : middle[1])
@rule object =
  r"{"p & (pair & (r","p & pair)[*] > commasPEG)[:?] & r"}"p > Dict∘bracketsPEG
@rule pair = string & r":"p & value > (k,c,v) -> (k => v)
@rule array =
  r"\["p & (value & (r","p & value)[*] > commasPEG)[:?] & r"]"p > bracketsPEG
@rule value =
  string, number, object, array,
  r"true|false"w |> Meta.parse, r"null"w |> x->nothing
@rule string = r"\"([^\"\\]|\\[\"\\/bfnrt]|\\u[0-9A-Fa-f]{4})*\""p |> s->Meta.parse(replace(s, r"\$"=>"\\\$"))
@rule number = r"-?(0|[1-9]\d*)(\.\d+)?([Ee][+-]?\d+)?"p |> Meta.parse
@rule startPEG = r"\s*" & value > (w,v)->v

# StringParserPEG grammar
commasSPP(r,v,f,l,c) = (length(c) == 1 ? [c[1]] : [c[1], (t.children[2] for t ∈ c[2].children)...])
bracketsSPP(r,v,f,l,c) = (isempty(c[2].children) ? [] : c[2].children[1])
braces = Dict∘bracketsSPP
parsevalue(r,v,f,l,c) = Meta.parse(v)
parsestr(r,v,f,l,c) = Meta.parse(replace(v, r"\$"=>"\\\$"))
grammarSPP = Grammar("""
  object => (r({\\s*)r & ?((pair & *(r(,\\s*)r & pair)) {Main.commasSPP}) & r(}\\s*)r) {Main.braces}
  pair => (string & r(:\\s*)r & value) {(r,v,f,l,c)->(c[1] => c[3])}
  array => (r(\\[\\s*)r & ?((value & *(r(,\\s*)r & value)) {Main.commasSPP}) & r(]\\s*)r) {Main.bracketsSPP}
  value => string | number | object | array | r((true|false)\\b\\s*)r {Main.parsevalue} | r(null\\b\\s*)r {(r,v,f,l,c)->:null}
  string => r(\\"([^\\"\\\\]|\\\\[\\"\\\\/bfnrt]|\\\\u[0-9A-Fa-f]{4})*\\"\\s*)r {Main.parsestr}
  number => r(-?(0|[1-9]\\d*)(\\.\\d+)?([Ee][+-]?\\d+)?\\s*)r {Main.parsevalue}
  start => (r(\\s*)r & value) {(r,v,f,l,c)->c[1]}
""")

# load JSON string from file from
# https://github.com/Chevrotain/chevrotain/blob/gh-pages/performance/samples/10K_json.js
json_str = open(x->read(x,String), "10K.json")

println("JSON:")
display(@benchmark JSON.parse(json_str) seconds=60)
println("PEG:")
display(@benchmark parse_whole(startPEG, json_str) seconds=60)
println("StringParserPEG:")
display(@benchmark parse(grammarSPP, json_str) seconds=60)


#these map model types to model functions
fluxmodels = Dict{Symbol, Function}(
  :levelmle => fluxlevelmlemodel,
  :levellsq => fluxlevellsqmodel,
)

fluxresultsname(λ⁺) = "results_$(PARAM[:crspdatalabel])" *
  "_$(PARAM[:fluxopt])" *
  "_$(PARAM[:fluxmodel])" *
  "_$(PARAM[:fluxpredictiontype])" *
  "_$(PARAM[:fluxresultssuffix])" *
  "_$λ⁺"

#primary function for estimating the measure using flux
function formfluxregression(panel::DataFrame, ms::MeasureSpec;
  fluxmodel::Function = fluxmodels[PARAM[:fluxmodel]])

  issorted(panel, [:permno, :date]) || error("panel must be sorted")

  zs::ZSpec = condition!(panel, ms)

  #testing
  if PARAM[:testmeasure]
    testvolumepartsindiceslsq()
    testvolumepartsindicesmle()
  end

  if PARAM[:refreshmeasure]
    λ⁺, results::DataFrame = fluxmodel(panel, zs)
    destandardize!(results, zs)

    #this will be the name of the results file
    resultname::String = fluxresultsname(λ⁺)

    results |> CSV.write("$(PARAM[:outputpath])\\$(resultname).csv")

  end

end

#pre-run on optimimization to find optimal learning rate
function findηmax(optgen, ::Type{T}=Float32;
  loss::Function = error("loss function is required"),
  Θ::AbstractVolumeParts = error("original VolumeParts model is required"),
  X::NamedTuple= error("X is required"),
  y = error("y is required"),
  ηsearchmax::T = T(PARAM[:fluxsearchmax]),
  ηstart::T = T(PARAM[:fluxsearchstart]),
  Δη::T = T(PARAM[:fluxsearchdelta]),
  searchiter::Int=PARAM[:fluxsearchiter],
  reportinterval::NInt = nothing) where T<:Real

  Θ₀ = copytrainable(Θ)  #the best model
  Π = Flux.params(Θ) #formatted parameters

  #record η here
  ηmax::T=ηstart
  λmax::T = -T(Inf)
  η::T=ηstart
  λmin::T = T(Inf)

  #set up an iterator so we can use the built-in training loop
  repdat = Iterators.repeated((X,y), searchiter)

  #callback for reporting
  reportcb() = println("η: $η\tηmax: $ηmax\t ηsupermax: $ηsearchmax λmin: $λmin")
  throttledreportcb = reportinterval === nothing ? reportcb : throttle(reportcb, reportinterval)

  @info "Beginning search for ηmax"
  while η ≤ ηsearchmax
    opt = optgen(η)
    Flux.train!(loss, Π, repdat, opt)
      #cb=throttle(()->println("Searching...loss=", loss(ws, RLws,V)), reportinterval))
    λ::T = loss(X,y)

    if λ < λmin
      λmin = λ
      ηmax = η
    end

    if λ > λmax
      λmax = λ
    end
    throttledreportcb()

    if λ === T(Inf)
      @warn "infinite loss. Assuming ηmax ($ηmax) < η($η) and breaking."
      break
    end

    #reset and repeat
    updatexfromy!(Θ, Θ₀)
    η *= exp(Δη)
  end

  updatexfromy!(Θ, Θ₀)
  ((λmax - λmin)/(abs(λmin) + abs(λmax)) ≤ 0.01) && error("
    λmin ($λmin) nearly equal to λmax ($λmax), consider expanding search bounds")
  (λmin === T(Inf)) && error("no ηmax found!!")

  @info "ηmax search complete."
  reportcb()
  return ηmax
end

function solvevol!(optgen, ::Type{T}=Float32;
    loss::Function = error("loss function is required"),
    Θ::AbstractVolumeParts = error("original vol parts model is required"),
    X::NamedTuple = error("X is required"),
    y = error("y is required"),
    debugcb::Function = ()->nothing,
    maxiter = PARAM[:fluxmaxiter],
    maxiternoimprovement = PARAM[:fluxmaxiternoimprovement],
    reportinterval = PARAM[:fluxreportinterval],
    ηsupermindenom = T(PARAM[:fluxlearningsupermindenom]),
    Δη⁺ = T(PARAM[:fluxlearningdeltaup]),
    Δη⁻ = T(PARAM[:fluxlearningdeltadown]),
    cyclemaxdenom = T(PARAM[:fluxcyclemaxdenom]),
    cyclemindenom = T(PARAM[:fluxcyclemindenom]),
    cycledeltadenom = T(PARAM[:fluxcycledeltadenom]),
    cycles = PARAM[:fluxcycles],
    ∇rmstol = PARAM[:fluxgradrmstol],
    stopnoimprovement = PARAM[:fluxstopnoimprovement],
    stopsmallgradient = PARAM[:fluxstopsmallgradient]
  ) where T<:Real

  #get the data
  Π = Flux.params(Θ) #formatted parameters
  Θ⁺ = copytrainable(Θ)  #the best model
  Θtemp = copytrainable(Θ)  #a temporary storage location

  PARAM[:fluxopt] == :adadelta && error("udpate not supported with :adadelta")

  #initialize the current state
  local ∞::T = T(Inf)
  local λ ::T= ∞
  local λ⁺::T = ∞ #the best loss

  local itersincebest::Int = 0
  local iter::Int = 0

  #cycle parameters
  local ηmax::T = findηmax(optgen, T, loss=loss, Θ=Θ, X=X, y=y, reportinterval=reportinterval)
  ηmax = ηmax / cyclemaxdenom
  local ηmin::T = ηmax / cyclemindenom

  local ηsupermin::T = ηmin / ηsupermindenom
  (ηmax ≤ ηmin) && error("ηmax $ηmax ≤ ηmin $ηmin")
  local Δcycle::T = ηmin / cycledeltadenom

  #learning parameters
  η::T = ηmin
  local η⁺::T = η #learnign rate at best
  local ηupmult::T = exp(Δη⁺) #multiply η by this when we decrease the learning rate
  local ηdownmult::T = exp(-Δη⁻) #multiply η by this when we reduce the learning rate
  local isbestλ::Bool

  #this will yield an updated optimizer given a learning rate
  #updateopt(η) = (η, optgen(η))
  @inline function updateopt!(η, opt)
    opt.eta = η
  end
  updateopt(η) = (η, optgen(η))

  #initialize the optimizer
  opt = optgen(η)
  opt⁺ = deepcopy(opt)

  #main gradient function
  get∇Θ() = gradient(Π) do
      λ = loss(X,y) #store the loss so that we don't have to calculate it again

      if λ < λ⁺ #check if this is the best we have
        λ⁺ = λ
        η⁺ = η
        itersincebest = 0
        isbestλ = true
      else
        itersincebest += 1
        isbestλ=false
      end

      return λ
    end

  #Report progress to the user
  function reportcb(∇Θ)

    #get the gradients for monitoring
    rmsnow = ∇rms(Θ, ∇Θ)
    updatexfromy!(Θtemp, Θ) #temp=Θ
    updatexfromy!(Θ, Θ⁺) #Θ=Θ⁺
    ∇Θ = get∇Θ()
    rmsbest = ∇rms(Θ, ∇Θ)
    updatexfromy!(Θ, Θtemp) #Θ=temp

    println("loss: $λ\tbestloss: $(λ⁺)\tη:$η\tcurrent∇rms: $(rmsnow)\tbest∇rms: $(rmsbest)\t",
      iter==0 ? "" : "iter: $iter/$maxiter\titersincebest: $itersincebest/$maxiternoimprovement")

  end
  throttledreportcb = reportinterval === nothing ? reportcb : throttle(reportcb, reportinterval)

  #three exit conditions: 1) maxiter 2) no improvement for some time 3) an error
  @info "beginning cycle conditioning. ηmax=$ηmax, ηmin=$ηmin, η=$η"
  @progress for cycle ∈ 0:cycles
    @info "beginning cycle $cycle of $(cycles). Est. $((ηmax-ηmin)/Δcycle*2.0) iter/cycle"
    (cycle==cycles) && @info "YAY! ITS THE FINAL CYCLE"
    while η ≤ ηmax
      ∇Θ = get∇Θ()
      (λ==∞) && debugcb() #infinity means something went wrong
      Flux.update!(opt, Π, ∇Θ)
      if isbestλ
        updatexfromy!(Θ⁺, Θ)
        opt₊ = deepcopy(opt)
      end
      η, opt = updateopt(η + Δcycle)
      #η = updateopt!(η + Δcycle, opt)
      throttledreportcb(∇Θ) #keep user appraised
    end
    while η ≥ ηmin
      ∇Θ = get∇Θ()
      (λ==∞) && debugcb() #infinity means something went wrong
      Flux.update!(opt, Π, ∇Θ)
      if isbestλ
        updatexfromy!(Θ⁺, Θ)
        opt₊ = deepcopy(opt)
      end
      η, opt = updateopt(η - Δcycle)
      #η = updateopt!(η - Δcycle, opt)
      throttledreportcb(∇Θ) #keep user appraised
    end
  end

  @info "Final cycle complete. Beginning final convex optimization from best point"
  updatexfromy!(Θ, Θ⁺)
  itersincebest = 0
  isbestλ = false

  opt = opt⁺
  @progress for i ∈ 1:maxiter
    iter = i #for reporting purposes
    ∇Θ = get∇Θ()

    isbestλ && updatexfromy!(Θ⁺, Θ)
    if λ==∞ #reset and lower learning rate
      debugcb()  #this is an error scenario where learning rate can't decrease
    elseif isbestλ && (η ≤ ηmax) #if its the best answer, note model and accelerate learning
      #η, opt = updateopt(η * ηupmult)
      η = updateopt!(η * ηupmult, opt)
    elseif (!isbestλ) && η ≥ ηsupermin #if its not, decellerate learning
      #η, opt = updateopt(η * ηdownmult)
      η = updateopt!(η * ηdownmult, opt)
    end

    #run through the stopping criteria
    if stopnoimprovement && (itersincebest ≥ maxiternoimprovement) #this is the main stopping criteria
      @info "No improvement in $maxiternoimprovement iterations."
      reportcb(∇Θ) #run the callback once more
      break
    elseif stopsmallgradient && ∇rms(Θ, ∇Θ) < ∇rmstol
      @info "GOOD NEWS: ∇rms ($(∇rms(Θ, ∇Θ))) < tol ($∇rmstol)."
      reportcb(∇Θ) #run the callback once more
      break
    elseif iter == maxiter
      @warn "max iter reached, returning best result"
      reportcb(∇Θ)
    else
      throttledreportcb(∇Θ) #keep user appraised
      Flux.update!(opt, Π, ∇Θ)
    end
  end

  ∇Θ = get∇Θ()
  printτ = :τ ∈ propertynames(Θ) #check if τ (risk or precision) is used in the optimization process
  open("$(PARAM[:outputpath])\\$(fluxresultsname(λ⁺))_info.txt", "w+") do f
    println(f, "******Loss: loss: $λ\tbestloss: $(λ⁺)\tη:$η\tgradrms: $(∇rms(Θ, ∇Θ))",
    "\n******Current A:\n", ∇Θ[Θ.A]',
    "\n******Current V₀:\n", ∇Θ[Θ.V₀],
    (printτ ? "\n******Current τ:\n$(display(∇Θ[Θ.τ]'))" : ""),
    "\n++++++Switch to best: ", isnothing(updatexfromy!(Θ,Θ⁺)),
    "\n++++++Best A:\n", ∇Θ[Θ.A]',
    "\n++++++Best V₀:\n", ∇Θ[Θ.V₀],
    (printτ ? "\n++++++Best τ:\n$(display(∇Θ[Θ.τ]'))" : ""),
    )
  end
  return λ⁺
end


using Random, JSON, Statistics

noFunc() = nothing

include("utils/GeneticProgramming/nodes.jl")
include("utils/GeneticProgramming/nodeSetsParser.jl")
include("utils/miscFunctions.jl")
include("tools/individual.jl")
include("tools/experimentSummary.jl")

"
Auxiliary structure that keeps some generic information for GP experiments.
It is the abstract type from which every other GP experiment info inherits.

See also: [`ExperimentInfo`](@ref), [`CGPInfo`](@ref), [`STGPInfo`](@ref),
[`GEPInfo`](@ref), [`GEInfo`](@ref)
"
abstract type GPExperimentInfo end
# abstract type

"""
    getVariables(gpExperimentInfo::GPExperimentInfo)

Gets the identifiers of the variables of GP problem.

# Arguments
- `gpExperimentInfo::GPExperimentInfo`: information about the GP experiment.

# Returns
Array with variables' identifiers as `strings`.
"""
getVariables(gpExperimentInfo::GPExperimentInfo) = gpExperimentInfo._variables
# function

"
Auxiliary structure that keeps some generic information for experiments.

# Fields
- `_individualType::DataType`: type of the individual that determines which
    methods are going to be used.
- `_GPExperimentInfo::GPExperimentInfo`: if the problem to solve is GP, this field
    will contain its information.
- `_randomSeed::UInt`: seed used for rand functions. By default, it will be 1234.
- `_rng::Random.AbstractRNG`: random number generator used by the system. By
    default, it will be set to Random.GLOBAL_RNG.
- `_algorithm::Function`: genetic algorithm used for the experiment. By default,
    it will be set to the classic genetic algorithm (see [`basicExperiment`](@ref)).
- `_algorithmArgs::Array{Any}`: arguments for the genetic algorithm, if needed.

See also: [`GeneticJulia`](@ref)
"
mutable struct ExperimentInfo
    _individualType::DataType
    _GPExperimentInfo::GPExperimentInfo
    _experimentSummary::ExperimentSummary
    _randomSeed::UInt
    _rng::Random.AbstractRNG
    _algorithm::Function
    _algorithmArgs::Array{Any}

    ExperimentInfo() = new()
end # struct

include("tools/stopCondition.jl")
include("tools/evaluator.jl")
include("tools/initialization.jl")
include("tools/selector.jl")
include("tools/crossover.jl")
include("tools/mutation.jl")
include("tools/replacement.jl")

"
Main structure of the system, gathers all the tools necessary to run a
population-based algorithm. By default, the main structure for `code users` is
called `GenJ`.

# Fields
- `_experimentInfo::ExperimentInfo`: auxiliary structure that keeps some generic
    information for experiments.
- `_stopCondition::StopCondition`: tool that defines the stop conditions and
    checks them throughout the experiment.
- `_population::Array{Individual}`: array that stores the population as Individual.
- `_evaluator::Evaluator`: defines the methods and structure for evaluate individuals.
- `_generator::Generator`: defines the methods and structure for generate individuals.
- `_selector::SelectionOperator`: defines the methods and structure for select
    \"parents\" among the population.
- `_crossoverOp::CrossoverOperator`: defines the methods and structure for cross
    \"parents\" and obtain the offspring.
- `_mutationOp::MutationOperator`: defines the methods and structure for mutate
    an offspring's individual.
- `_replacementOp::ReplacementOperator`: defines the methods and structure for
    replace the old population with the offspring.
"
mutable struct GeneticJuliaStructure
    _experimentInfo::ExperimentInfo
    _stopCondition::StopCondition
    _population::Array{Individual}
    _evaluator::Evaluator
    _generator::Generator
    _selector::SelectionOperator
    _crossoverOp::CrossoverOperator
    _mutationOp::MutationOperator
    _replacementOp::ReplacementOperator

    GeneticJuliaStructure() = new()
end # struct

# Global structure generation for code user transparency.
GenJ = GeneticJuliaStructure()

# Due to optional options, it is created here.
GenJ._experimentInfo = ExperimentInfo()
GenJ._experimentInfo._randomSeed = 1234
GenJ._experimentInfo._rng = Random.GLOBAL_RNG

include("utils/GeneticProgramming/includes.jl")
include("utils/GeneticAlgorithms/includes.jl")
include("utils/algorithms.jl")
include("utils/defaultSettings.jl")
include("utils/error.jl")
include("utils/parser.jl")
include("utils/verboseFunctions.jl")
include("interface/includes.jl")

"""
    getIndividualType(experimentInfo::ExperimentInfo)

Gets the individual type of the main structure.
"""
getIndividualType(experimentInfo::ExperimentInfo) = experimentInfo._individualType
 # function

"""
    getIndividualType(genj::GeneticJulia)

Gets the individual type of the main structure.
"""
getIndividualType(genj::GeneticJuliaStructure) = getIndividualType(genj._experimentInfo)
 # function

"""
    setIndividualType(type::DataType; genj::GeneticJulia = GenJ)

Sets the type of the individuals' representation for the experiment.

# Arguments
- `type::DataType`: data type of the individuals of the experiment. Should be
    one of the implemented ones (see [`IndividualRep`](@ref)). The user could
    provide their own type, as long as they also provide their own methods for
    population generation, crossover and mutation.

# Keyword Arguments
- `genj::GeneticJulia = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.

# Examples
```jldoctest
julia> setIndividualType(BinaryGenotype)
BinaryGenotype
```
"""
function setIndividualType(type::DataType; genj::GeneticJuliaStructure = GenJ)

    if findfirst(x->type <: x, [BinaryGenotype, IntegerGenotype, CGPGenotype,
                                          STGPGenotype, GEPGenotype, GEGenotype]) == nothing
        @warn "The individual type provided does not match any of the predefined
               ones. Personalized methods for generation, crossover and mutation
               shall be provided"
    end

    genj._experimentInfo._individualType = type
end # function

"""
    setAlgorithm(algorithm::Function, varArgs...; genj::GeneticJulia = GenJ)

Sets the genetic algorithm that is going to be run for the experiment.

# Arguments
- `algorithm::Function`: genetic algorithm that is going to run the experiment.
    By default, it is set to the classic one, which follows this steps:
    - Generation of the population.
    - Fitness evaluation.
    - Parent selection.
    - Crossover.
    - Mutation.
    - Fitness evaluation.
    - Replacement of the population.
- `varArgs`: arguments for the algorithm, if needed.

# Keyword Arguments
- `genj::GeneticJulia = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.

# Examples
```jldoctest
julia> setAlgorithm(basicExperiment)
```
"""
function setAlgorithm(algorithm::Function, varArgs...; genj::GeneticJuliaStructure = GenJ)
    genj._experimentInfo._algorithm = algorithm
    varArgs = [i for i in varArgs]
    genj._experimentInfo._algorithmArgs = varArgs
end # function

"""
    setRandomSeed(seed::UInt; genj::GeneticJulia = GenJ)

Sets the random number generator as a MersenneTwister generator initialized with
the specified seed.

# Arguments
- `seed::Integer`: seed to initialize the random number generator with.

# Keyword Arguments
- `genj::GeneticJulia = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.

# Examples
```jldoctest
julia> setRandomSeed(2)
```
"""
function setRandomSeed(seed::Integer; genj::GeneticJuliaStructure = GenJ)
    if seed < 0
        error("RandomSeed ($seed) must be a positive integer")
    end
    genj._experimentInfo._randomSeed = convert(UInt, seed)
    genj._experimentInfo._rng = Random.MersenneTwister(seed)
end # function

"""
    getRNG(experimentInfo::ExperimentInfo)

Gets the individual type of the main structure.
"""
getRNG(experimentInfo::ExperimentInfo) = experimentInfo._rng
 # function

"""
    getRNG(genj::GeneticJulia)

Gets the random number generator of the main structure.
"""
getRNG(genj::GeneticJuliaStructure) = getRNG(genj._experimentInfo)
# function

"""
    setRNG(rng::Random.AbstractRNG; genj::GeneticJulia = GenJ)

Sets the random number generator for the system to run the experiment.

# Arguments
- `seed::Integer`: seed to initialize the random number generator with.

# Keyword Arguments
- `genj::GeneticJulia = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.

# Examples
```jldoctest
julia> rng = Random.MersenneTwister()

julia> setRNG(rng)
```
"""
function setRNG(rng::Random.AbstractRNG; genj::GeneticJuliaStructure = GenJ)
    genj._experimentInfo._rng = rng
end # function

"""
    runGenJ(jsonFile::String; verbose::Bool = true, outputFile::String = "")

Runs the experiments specified in the configuration file.

# Arguments
- `configFile::String`: the configuration file (in JSON format) with all the information about the
    experiments that are going to be run.

# Keyword Arguments
- `verbose::Bool = true`: specifies wether information about the experiment is going
    to be displayed or not.
- `outputFile::String = ""`: specifies the output stream in which the information
    displayed by verbose is going to be dumped.

# Returns
The array of experiments run.

For guidance about how to create the information file about the experiments
visit our GitHub repository at:
[https://github.com/SergioGmezM/Evolutionary\\_Computation\\_Julia](https://github.com/SergioGmezM/Evolutionary_Computation_Julia)
"""
function runGenJ(configFile::String; verbose::Bool = true, outputFile::String = "")
    experiments = generateMainStructure(configFile)
    nExperiments = length(experiments)

    if !verbose
        verbose = outputFile != ""
    end

    for i=1:nExperiments
        # Sets the default settings according to the individual type provided in the experiment
        setDefaultSettings(genj=experiments[i], i=i)
        # Checks that all has been set
        checkParametersDefined(experiments[i], i)
        # Checks that all has been set correctly
        checkParametersCorrect(experiments[i], i)
    end

    for i=1:nExperiments
        println("Experiment $i \n==========================")
        # Prints the information about the experiment
        verbose && printExperimentInfo(outputFile, experiments[i])
        # Runs the experiment
        experiments[i]._experimentInfo._algorithm(experiments[i]._experimentInfo._algorithmArgs..., genj=experiments[i])
    end

    return experiments
end # function

"""
    runGenJ(genj::GeneticJulia = GenJ; verbose::Bool = true, outputFile::String = "")

Runs a single experiment for `code users`.

# Arguments
- `genj::GeneticJulia = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.

# Keyword Arguments
- `verbose::Bool = true`: specifies wether information about the experiment is going
    to be displayed or not.
- `outputFile::String = ""`: specifies the output stream in which the information
    displayed by verbose is going to be dumped.

For guidance about how to create the information file about the experiments
visit our GitHub repository at:
[https://github.com/SergioGmezM/Evolutionary\\_Computation\\_Julia](https://github.com/SergioGmezM/Evolutionary_Computation_Julia)
"""
function runGenJ(genj::GeneticJuliaStructure = GenJ; verbose::Bool = true, outputFile::String = "")

    if !verbose
        verbose = outputFile != ""
    end

    # Sets the default settings according to the individual type provided in the experiment
    setDefaultSettings(genj=genj)
    # Checks that all has been set
    checkParametersDefined(genj)
    # Checks that all has been set correctly
    checkParametersCorrect(genj)
    # Prints the information about the experiment
    verbose && printExperimentInfo(outputFile, genj)
    # Runs the experiment
    genj._experimentInfo._algorithm(genj._experimentInfo._algorithmArgs..., genj=GenJ)
end # function


using Test
using LavoisierCore

model = GERG2008(:C2,:C3)
x0 = [0.5,0.5]
T0 = 300
@testset "initial" begin
    @test length(model) == 2
end


using Random
Random.seed!(101)

@testset "Values tests" begin
  @testset "String values" begin
    @test rand(String) == "Tw2GszEb"
    @test rand(String) == "3U3vvin8"
    @test rand(String, 20) == "fhJmWB8TAixQ057e4eSF"
  end

  @testset "Pair values" begin
    @test rand(Pair{String,String}) == ("Tw2GszEb" => "3U3vvin8")
  end
end


# Copyright 2019, Patricia Levi, Elias Kueppper, and contributors

## dependencies
using StatsPlots, DataFrames, CSV

"""
    groupedbar(var::OptVariable{T,N};
                    title::String="",
                    export_name="",
                    kwargs...) where {T,N}
Convert the provided OptVariable into a groupedbar-plot and export it
"""
import StatsPlots.groupedbar
function groupedbar(var::OptVariable{T,N};
                    title::String="",
                    export_name="",
                    kwargs...) where {T,N}
    if N==2
        outplot = groupedbar(var.axes[1], var.data; labels=var.axes[2], title = title, kwargs...)
    end
    if export_name!=""
        savefig(outplot, export_name*".png")
    end
    return outplot
end

"""
        dataframe(var::OptVariable{T,N};
                    export_name="") where {T,N}
Convert the provided OptVariable into a `DataFrame` and export it
"""
function dataframe(var::OptVariable{T,N};
                    export_name="") where {T,N}
    if N==2
        table=DataFrame(var.data)
        names!(table,Symbol.(var.axes[2]))
        table[Symbol(var.axes_names[1])]=var.axes[1]
    end
    showall(table)
    if export_name!=""
        CSV.write(export_name*".csv",table)
    end
end

"""
        plotcapacity(results::Array{OptResult,1},
                    outputname::String;
                    kwargs...)
Plot the capacity of the `results`-Array and form a table
"""
function plotcapacity(results::Array{OptResult,1},
                    outputname::String;
                    kwargs...)
    desc=Array{String,1}()
    techs=Array{String,1}()
    for result in results
        push!(desc, result.config["descriptor"])
        unique!(union!(techs, result.sets["tech"]["all"]))
    end

    cap=OptVariable{Number}(undef, techs, desc; axes_names=["tech", "results"])
    cap.=0
    for result in results
        for tech in result.sets["tech"]["node"]
            cap[tech, result.config["descriptor"]]=sum(result.variables["CAP"][tech,:,:])
        end
        if haskey(result.sets["tech"], "line")
            lines=load_cep_data_provided(result.config["region"]).lines
            for tech in result.sets["tech"]["line"]
                c=0
                for line in axes(lines, "line")
                    c+=sum(result.variables["TRANS"][tech,:,line])*lines[tech,line].length
                end
                cap[tech, result.config["descriptor"]]=c./1e3
            end
        end
    end
    dataframe(cap; export_name=outputname*"capdata")
    return groupedbar(cap; title="Installed Capacities", export_name=outputname*"capdata", yaxis="Capacities [MW] and Transmission [GW*km]", kwargs...)
end

"""
        plotcost(results::Array{OptResult,1},
                    outputname::String;
                    kwargs...)
Plot the costs of the `results`-Array and form a table
"""
function plotcost(results::Array{OptResult,1},
                outputname::String;
                kwargs...)
    desc=Array{String,1}()
    for result in results
        push!(desc, result.config["descriptor"])
    end

    cost=OptVariable{Number}(undef, desc, ["var", "cap_fix"]; axes_names=["results", "cost_type"])
    for result in results
        for cost_type in axes(cost, "cost_type")
            cost[result.config["descriptor"], cost_type]=sum(result.variables["COST"][cost_type,first(result.sets["impact"]["mon"]),:])
        end
    end
    dataframe(cost; export_name=outputname*"costdata")
    return groupedbar(cost; title="Costs", export_name=outputname*"costdata", bar_position=:stack, yaxis="Cost [$(first(first(results).sets["impact"]["mon"]))]", kwargs...)
end

"""
        plotgen(results::Array{OptResult,1},
                    outputname::String;
                    kwargs...)
Plot the gen of the `results`-Array and form a table
"""
function plotgen(results::Array{OptResult,1},
                outputname::String;
                kwargs...)
    desc=Array{String,1}()
    techs=Array{String,1}()
    for result in results
        push!(desc, result.config["descriptor"])
        unique!(union!(techs, result.sets["tech"]["power"]))
    end

    gen=OptVariable{Number}(undef, techs, desc; axes_names=["results", "tech"])
    gen.=0
    for result in results
        for tech in result.sets["tech"]["power"]
            g=0
            for k in result.sets["time_K"]["all"]
                for t in result.sets["time_T_period"]["all"]
                    g+=sum(result.variables["GEN"][tech,"electricity",t,k,:])*result.config["time_series"]["delta_t"][t,k]*result.config["time_series"]["delta_t"][k]
                end
            end
            gen[tech, result.config["descriptor"]]=g
        end
    end
    dataframe(gen; export_name=outputname*"gendata")
    return groupedbar(gen; title="Total Generation, Demand, Loss", export_name=outputname*"gendata", yaxis="Energy [MWh]", kwargs...)
end


using Test
using Random
using StaticArrays: SVector
Random.seed!(1234)

using ClimateMachine.VariableTemplates: @vars, varsize, Vars
using ClimateMachine.BalanceLaws
const BL = BalanceLaws
import ClimateMachine.BalanceLaws: vars_state, eq_tends, prognostic_vars

struct TestBL <: BalanceLaw end
struct X <: PrognosticVariable end
struct Y <: PrognosticVariable end
struct F1{PV} <: TendencyDef{Flux{FirstOrder}, PV} end
struct F2{PV} <: TendencyDef{Flux{SecondOrder}, PV} end
struct S{PV} <: TendencyDef{Source, PV} end

prognostic_vars(::TestBL) = (X(), Y())
eq_tends(::X, ::TestBL, ::Flux{FirstOrder}) = (F1{X}(),)
eq_tends(::Y, ::TestBL, ::Flux{FirstOrder}) = (F1{Y}(),)
eq_tends(::X, ::TestBL, ::Flux{SecondOrder}) = (F2{X}(),)
eq_tends(::Y, ::TestBL, ::Flux{SecondOrder}) = (F2{Y}(),)
eq_tends(::X, ::TestBL, ::Source) = (S{X}(),)
eq_tends(::Y, ::TestBL, ::Source) = (S{Y}(),)

@testset "BalanceLaws" begin
    bl = TestBL()
    @test prognostic_vars(bl) == (X(), Y())
    @test fluxes(bl, FirstOrder()) == (F1{X}(), F1{Y}())
    @test fluxes(bl, SecondOrder()) == (F2{X}(), F2{Y}())
    @test sources(bl) == (S{X}(), S{Y}())
    show_tendencies(bl)
    show_tendencies(bl; include_params = true)
end

vars_state(bl::TestBL, st::Prognostic, FT) = @vars begin
    ρ::FT
    ρu::SVector{3, FT}
end

vars_state(bl::TestBL, st::Auxiliary, FT) = @vars()

@testset "Prognostic-Primitive conversion (identity)" begin
    FT = Float64
    bl = TestBL()
    vs_prog = vars_state(bl, Prognostic(), FT)
    vs_prim = vars_state(bl, Primitive(), FT)
    vs_aux = vars_state(bl, Auxiliary(), FT)
    prim_arr = zeros(varsize(vs_prim))
    prog_arr = zeros(varsize(vs_prog))
    aux_arr = zeros(varsize(vs_aux))

    # Test prognostic_to_primitive! identity
    prog_arr .= rand(varsize(vs_prog))
    prog_0 = deepcopy(prog_arr)
    prim_arr .= 0
    BL.prognostic_to_primitive!(bl, prim_arr, prog_arr, aux_arr)
    BL.primitive_to_prognostic!(bl, prog_arr, prim_arr, aux_arr)
    @test all(prog_arr .≈ prog_0)

    # Test primitive_to_prognostic! identity
    prim_arr .= rand(varsize(vs_prim))
    prim_0 = deepcopy(prim_arr)
    prog_arr .= 0
    BL.primitive_to_prognostic!(bl, prog_arr, prim_arr, aux_arr)
    BL.prognostic_to_primitive!(bl, prim_arr, prog_arr, aux_arr)
    @test all(prim_arr .≈ prim_0)
end


# scripts to analyse behavioral data in various ways

using SpecialFunctions, Distributions, Optim, Calculus, MAT, DataFrames, Gadfly, Colors

# constants and common functions
const sqrttwo = √(2)
const sqrttwoπ = √(2 * π)
Φ(x) = (1 + erf(x / sqrttwo)) / 2

###############################################################################
# loading experimental data
#
# Most of the dataset is self-explanatory.
# Stimulus ranges from 1 to 8 (Stimulus = 1..4 -> right, 5..8 -> left). OSD is
# odor sampling duration [s]. Outcome is 1 (rewarded), 0 (not rewarded due to
# wrong choice), or NaN (not rewarded due to too fast choice).
#
# For the interleaved :itlvd dataset, the format is slightly different, in
# that Stimulus ranges from 1 to 32, with the following ordering:
#  1.. 4 (right easy->hard)  8.. 5 (left hard->easy) concentration 0.1
#  9..12 (right easy->hard) 16..13 (left hard->easy) concentration 0.01
# 17..20                    24..21                   concentration 0.001
# 25..28                    32..29                   concentration 0.0001
# In the below, the Id's for left will be inverted to recover the 1-8, 9-16, ...
# ordering.
#
# On top of the stimulus id, the method adds the following stimulus identifier
# mixw 1  mixmag 4  for 100/  0 mix | concentr 4  for 0.1     concentration
# mixw 2  mixmag 3  for  80/ 20 mix | concentr 3  for 0.01    concentration
# mixw 3  mixmag 2  for  68/ 32 mix | concentr 2  for 0.001   concentration
# mixw 4  mixmag 1  for  56/ 44 mix | concentr 1  for 0.0001  concentration
# mixw 5  mixmag 1  for  44/ 56 mix
# mixw 6  mixmag 2  for  32/ 68 mix
# mixw 7  mixmag 3  for  20/ 80 mix
# mixw 8  mixmag 4  for   0/100 mix

# dataset is :categ, :ident, :itlvd
# a ratid of 0 causes the function to return all rats
function load_data(dataset::Symbol, ratid::Int = 0)
    @assert dataset == :categ || dataset == :ident || dataset == :itlvd
    @assert ratid ∈ [0, 1, 2, 3, 4]  # only 4 rats
    # load MAT file
    d = dataset == :categ ?
        matread("../data/DATA CAT.mat")["DATA_cat"] :
        (dataset == :ident ?
         matread("../data/DATA ID.mat")["DATA_id"] :
         matread("../data/DATA INT.mat")["DATA_int"])
    stimid = round.(Int, vec(d["Stimulus"]))
    resp = round.(Int, vec(d["ChoiceDir"]))
    trials = length(stimid)
    if dataset == :itlvd
        # :itlvd dataset
        @assert minimum(stimid) == 1 && maximum(stimid) == 32
        corrresp = zeros(Int, trials)
        stimmag = zeros(Int, trials)
        corr = zeros(Bool, trials)
        mixw = zeros(Int, trials)
        mixmag = zeros(Int, trials)
        concentr = zeros(Int, trials)
        # process stimulus properties separately per concentration (1-8, 9-16, ...)
        for i = 1:4
            # reverse stimulus ordering for 5 -> 8, 13 -> 16, etc.
            righttrials = (stimid .>= 8i-3) .& (stimid .<= 8i)
            lefttrials = (stimid .>= 8i-7) .& (stimid .<= 8i-4)
            alltrials = righttrials .| lefttrials
            stimid[righttrials] = (16i - 3) .- stimid[righttrials]
            corrresp[righttrials] .= 1 # left = 0, right = 1
            stimmag[righttrials] .= stimid[righttrials]
            stimmag[lefttrials] .= (16i - 7) .- stimid[lefttrials]
            corr[lefttrials] = resp[lefttrials] .== 0
            corr[righttrials] = resp[righttrials] .== 1
            mixw[alltrials] = stimid[alltrials] .- 8(i-1)       # to 1-8 range
            mixmag[alltrials] = stimmag[alltrials] .- (8i - 4)  # to 1-4 range
            concentr[alltrials] .= 5-i
        end
    else
        # :categ or :ident dataset - those are simpler
        @assert minimum(stimid) == 1 && maximum(stimid) == 8
        corrresp = Int[i < 5 ? 0 : 1 for i in stimid]
        stimmag = Int[i < 5 ? 9 - i : i for i in stimid]
        corr = (stimid .>= 5) .& (resp .== 1) .| (stimid .<= 4) .& (resp .== 0)
        if dataset == :categ
            mixw = copy(stimid)
            mixmag = stimmag .- 4
            concentr = 4 * ones(Int, trials)
        else
            mixw = 7corrresp .+ 1  # either 1 or 8
            mixmag = 4 * ones(Int, trials)
            concentr = stimmag .- 4
        end
    end
    rew = [!isnan(i) && round(Int, i) == 1 for i in vec(d["Outcome"])]
    # turn into DataFrame
    df = DataFrame(rat = round.(Int, vec(d["RAT"])),
                   session = round.(Int, vec(d["Session"])),
                   trial = round.(Int, vec(d["Trial"])),
                   stimid = stimid,
                   stimmag = stimmag,
                   mixw = mixw,
                   mixmag = mixmag,
                   concentr = concentr,
                   corrresp = corrresp,
                   previd = [stimid[end]; stimid[1:end-1]],    # ign. boundaries
                   prevmag = [stimmag[end]; stimmag[1:end-1]], # "
                   resp = resp,
                   t = vec(d["OSD"]),
                   corr = corr,
                   rew = rew,
                   prevcorr = [corr[end]; corr[1:end-1]],      # ign. boundaries
                   prevrew = [rew[end]; rew[1:end-1]])
    if ratid != 0
        df = df[df[:rat] .== ratid, :]
    end
    return df
end

###############################################################################
# fitting psychometric function, assuming the probabilistic model
#
# p(y = 1 | x) = (1-pl) Phi((x - μ) / σ) + pl pb
#
# with mean μ, standard deviation σ, lapse rate pl and lapse bias bp.
# THe log-likelihood is defined over J xj's, nj's and pj's, where nj is the
# number of observations for the corresponding xj, and pj = sum yj / nj is the
# fraction of y's that are 1 at this xj.
#
# The parameter vector is θ = [μ, σ, pl, pb]

# p(y = 1 | x) for given parameters
psychpy1(θ::AbstractVector, x) = (1 - θ[3]) .* Φ.((x .- θ[1]) ./ θ[2]) .+ θ[3] * θ[4]

# log-likelihood
function psychllh(θ::AbstractVector,
                  x::AbstractVector, n::AbstractVector, p::AbstractVector)
    @assert length(x) == length(n) == length(p)
    @assert length(θ) == 4

    # parameter bound checking
    μ, σ, pl, pb = θ
    σ > 0.0 && 0.0 <= pl <= 1.0 && 0.0 <= pb <= 1.0 || return -Inf
    # compute llh
    llh = 0.0
    for j = 1:length(x)
        res = (x[j] - μ) / σ
        llh += n[j] * (p[j] * log((1-pl) * Φ(res) + pl * pb) +
                       (1-p[j]) * log((1-pl) * Φ(-res) + pl * (1-pb)))
    end
    return llh
end

# gradient g
function psychgrad(θ::AbstractVector,
                   x::AbstractVector, n::AbstractVector, p::AbstractVector)
    @assert length(x) == length(n) == length(p)
    @assert length(θ) == 4
    g = Vector{Float64}(undef, 4)

    # parameter bound checking
    μ, σ, pl, pb = θ
    if !(σ > 0.0 && 0.0 <= pl <= 1.0 && 0.0 <= pb <= 1.0)
        g[1] = 0.0
        g[2] = σ > 0.0 ? 0.0 : Inf
        g[3] = 0.0 <= pl <= 1.0 ? 0.0 : (pl < 0.0 ? Inf : -Inf)
        g[4] = 0.0 <= pb <= 1.0 ? 0.0 : (pb < 0.0 ? Inf : -Inf)
        return g
    end

    # compute gradient
    g[:] .= 0.0
    for j = 1:length(x)
        res = (x[j] - μ) / σ
        Nres = exp(- 0.5 * abs2(res)) / (sqrttwoπ * σ)
        Φres, Φnres = Φ(res), Φ(-res)
        plΦres, plΦnres = (1-pl) * Φres + pl * pb, (1-pl) * Φnres + pl * (1-pb)
        plfrac = (1-p[j]) / plΦnres - p[j] / plΦres

        g[1] += n[j] * (1-pl) * Nres * plfrac
        g[2] += n[j] * (1-pl) * Nres * res * plfrac
        g[3] += n[j] * (p[j] * (pb - Φres) / plΦres +
                        (1 - p[j]) * (1 - Φnres - pb) / plΦnres)
        g[4] -= n[j] * pl * plfrac
    end
    return g
end
# returing llh, modifying gradient, to be used by NLopt.jl
function psychllhgrad!(θ::AbstractVector, g::AbstractVector,
                       x::AbstractVector, n::AbstractVector, p::AbstractVector)
    if !isempty(g)
        g[:] = psychgrad(θ, x, n, p)
    end
    return psychllh(θ, x, n, p)
end
# modifying negative gradient, to be used by Optim.jl
function negpsychgrad!(θ::AbstractVector, g::AbstractVector,
                       x::AbstractVector, n::AbstractVector, p::AbstractVector)
    g[:] = -psychgrad(θ, x, n, p)
    return g
end

# Hessian h, by finite differences
psychhess(θ, x, n, p) = hessian(θt -> psychllh(θt, x, n, p), θ)
# Hessian diagonal by taking (finite difference) derivative of gradient
function psychhessdiag(θ, x, n, p)
    @assert length(θ) == 4
    psychgradi(i, θi) = (θtmp = copy(θ); θtmp[i] = θi; psychgrad(θtmp, x, n, p)[i])
    return Float64[derivative(θi -> psychgradi(i, θi), θ[i]) for i in 1:4]
end

###############################################################################
# functions analyzing behavioral data
#
# For psycometric/chronometric functions, objective correctness of the choices
# is used. For sequential dependencies, the measure of 'correctness' is if
# reward has been given in the previous trial.

# statistics that also work for empty lists
_0mean(x) = isempty(x) ? zero(eltype(x)) : mean(x)
_0var(x) = isempty(x) ? Inf : var(x)
_0sem(x, n) = n == 0 ? Inf : √(_0var(x) / n)
_0sem(x) = _0sem(x, length(x))

# returns various RT / choice statistics over the whole set
_tcstats(df, datasets::Int=1) = DataFrame(
    n = size(df, 1) / datasets,
    nr = sum(df[:resp] .== 1) / datasets,
    ncorr = sum(df[:corr]) / datasets,
    pcorrμ = _0mean(df[:corr]),
    pcorrsem = √(datasets * _0mean(df[:corr]) .* (1 - _0mean(df[:corr])) ./ size(df, 1)),
    prμ = _0mean(df[:resp] .== 1),
    prsem = √(datasets * _0mean(df[:resp] .== 1) .* (1 - _0mean(df[:resp] .== 1)) ./ size(df, 1)),
    tμ = _0mean(df[:t]),
    tσ = √(_0var(df[:t])),
    tsem = _0sem(df[:t], size(df, 1) / datasets),
    tcorrμ = _0mean(df[df[:corr], :t]),
    tcorrσ = √(_0var(df[df[:corr], :t])),
    tcorrsem = _0sem(df[df[:corr], :t], sum(df[:corr]) / datasets),
    tincorrμ = _0mean(df[.!df[:corr], :t]),
    tincorrσ = √(_0var(df[.!df[:corr], :t])),
    tincorrsem = _0sem(df[.!df[:corr], :t], sum(.!df[:corr]) / datasets)
    )

# returns psychometric/chronometric statistics, once ordered by stimid, once by
# stimmag.
# If datasets > 1, the returned sem's are average sem's per dataset, assuming
# that the provided data is from the given number of datasets.
getpsychchronstats(df, datasets::Int=1) = (
    by(df, :stimid, df -> _tcstats(df, datasets)),
    by(df, :stimmag, df -> _tcstats(df, datasets)))

# returns sequential statistics, seq_df, fit_df, biasid_df, biasmag_df
# seq_df returns sequential choice statistics, conditional on previous reward,
# previous stimulus id, and current stimulus id.
# fit_df fits these sequential psychometric curve, with one fit per
# (pre. rew, prev. stim) combination. The function first performs a global
# fit (stimid/stimmag = 0) to determine the slope of the psychometric curve,
# and fixes this slope in all subsequent (conditional) fits.
# biasid_df and biasmag_df use the fits to compute the biases, either per
# id or magnitude of the previous stimulus.
function getseqstats(df, datasets::Int=1)
    seq_df = by(df, [:stimid, :previd, :prevrew], df -> _tcstats(df, datasets))

    # find psychometric curve fit parameters
    fit_df = DataFrame(previd = Int[], prevrew = Bool[],
                       μ = Float64[], σ = Float64[], pl = Float64[], pb = Float64[],
                       μsd = Float64[], σsd = Float64[], plsd = Float64[], pbsd = Float64[])
    # fit overall psychometric curve
    psych_df = by(df, :stimid,
                  df -> DataFrame(n = size(df, 1), pr = _0mean(df[:resp])))
    x, n, p = psych_df[:stimid], psych_df[:n] ./ datasets, psych_df[:pr]
    θini = [mean(x), 1.0, 0.1, 0.5]
    θmin = [-Inf, 0.0, 0.0, 0.0]
    θmax = [Inf, Inf, 1.0, 1.0]
    res = optimize(
        OnceDifferentiable(
            θ -> -psychllh(θ, x, n, p),
            (g, θ) -> negpsychgrad!(θ, g, x, n, p), θini),
        θmin, θmax, θini, Fminbox(LBFGS()))
    globalθ = Optim.minimizer(res)
    globalθsd = sqrt.(max.(0.0, -1 ./ psychhessdiag(globalθ, x, n, p)))
    push!(fit_df, [0, true,
                   globalθ[1], globalθ[2], globalθ[3], globalθ[4],
                   globalθsd[1], globalθsd[2], globalθsd[3], globalθsd[4]])

    # fit per-condition psychometric curve, fixing slope to above slope
    stimids = sort(unique(psych_df[:stimid]))
    for prevrew in [true, false], sid in stimids
        psych_df = by(df[(df[:prevrew] .== prevrew) .& (df[:previd] .== sid), :], :stimid,
                      df -> DataFrame(n = size(df, 1), pr = _0mean(df[:resp])))
        x, n, p = psych_df[:stimid], psych_df[:n] ./ datasets, psych_df[:pr]
        # optimize over [μ, pl, pb], while using σ = globalθ[2]
        function psychllhgradσ!(subθ, subg)
            θ = [subθ[1], globalθ[2], subθ[2], subθ[3]]
            if !isempty(subg)
                subg[:] = -psychgrad(θ, x, n, p)[[1, 3, 4]] # ignore σ gradient
            end
        end
        res = optimize(
            OnceDifferentiable(
                subθ -> -psychllh([subθ[1], globalθ[2], subθ[2], subθ[3]], x, n, p),
                (subg, subθ) -> psychllhgradσ!(subθ, subg), θini[[1, 3, 4]]),
            θmin[[1, 3, 4]], θmax[[1, 3, 4]], θini[[1, 3, 4]], Fminbox(LBFGS()))
        θ = Optim.minimizer(res)
        θ = Float64[θ[1], globalθ[2], θ[2], θ[3]]
        # approximation, as off-diagonal terms might influence inv. Hessian
        θsd = sqrt.(max.(0.0, -1 ./ psychhessdiag(θ, x, n, p)))
        # ensure σ^2 < μ (1 - μ)
        θsd[3] = min(θsd[3], 0.999999 * sqrt.(θ[3] * (1 - θ[3])))
        θsd[4] = min(θsd[4], 0.999999 * sqrt.(θ[4] * (1 - θ[4])))
        push!(fit_df, [sid, prevrew,
                       θ[1], θ[2], θ[3], θ[4], θsd[1], θsd[2], θsd[3], θsd[4]])
    end
    fit_df[:prevmag] = [id == 0 ? 0 : (id < 5 ? 9 - id : id) for id in fit_df[:previd]]

    # compute biases
    bootstrapn = 1000
    # draw n samples from Beta distribution with mean μ, standard deviation σ
    samplebeta(μ, σ, n) = σ <= 1e-30 ? fill(μ, n) :
        (x = (μ - abs2(μ)) / abs2(σ) - 1; rand(Beta(μ * x, (1 - μ) * x), n))
    biasid_df = DataFrame(previd = Int[], prevrew = Bool[],
                          biasμ = Float64[], biasσ = Float64[])
    for prevrew in [true, false], sid in stimids
        μ, σ, pl, pb, μsd, σsd, plsd, pbsd = Vector{Float64}(
            fit_df[(fit_df[:previd] .== sid) .& (fit_df[:prevrew] .== prevrew),
                   [:μ, :σ, :pl, :pb, :μsd, :σsd, :plsd, :pbsd]][1,:])
        # bias = intersection of psychometric curve with 4.5
        biases = Float64[psychpy1([μ + μsd * randn(),
                                   σ + σsd * randn(),
                                   samplebeta(pl, plsd, 1)[1],
                                   samplebeta(pb, pbsd, 1)[1]], 4.5) - 0.5
                         for n in 1:bootstrapn]
        push!(biasid_df, [sid, prevrew, mean(biases), √(var(biases))])
    end
    # biases over magnitudes -> flip for id < 5, and combine by variance
    stimmags = sort(unique(fit_df[fit_df[:prevmag] .!= 0, :prevmag]))
    biasmag_df = DataFrame(prevmag = Int[], prevrew = Bool[],
                           biasμ = Float64[], biasσ = Float64[])
    for prevrew in [true, false], smag in stimmags
        μs = [-biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                         (biasid_df[:previd] .== 9 - smag), :biasμ][1],
              biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                         (biasid_df[:previd] .== smag), :biasμ][1]]
        σs = [biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                        (biasid_df[:previd] .== 9 - smag), :biasσ][1],
              biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                         (biasid_df[:previd] .== smag), :biasσ][1]]
        ws = 1 ./ abs2.(σs)
        push!(biasmag_df, [smag, prevrew, sum(μs .* ws) / sum(ws), √(1 / sum(ws))])
    end

    return seq_df, fit_df, biasid_df, biasmag_df
end

###############################################################################
# functions plotting behavioral data
#
# These functions use the above analysis functions to perform the plotting.

# returns chron/psych plot over stimulus id
function plotchronpsychid(d_bystim, stimx, xlabel, ttitle, t=Theme())
    d_bystim[:stimx] = stimx[d_bystim[:stimid]]
    xmin, xmax = minimum(stimx), maximum(stimx)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bystim[:tmin] = d_bystim[:tμ] - d_bystim[:tsem]
    d_bystim[:tmax] = d_bystim[:tμ] + d_bystim[:tsem]
    pchronid = plot(d_bystim, x="stimx", y="tμ", ymin="tmin", ymax="tmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bystim[:prmin] = d_bystim[:prμ] - d_bystim[:prsem]
    d_bystim[:prmax] = d_bystim[:prμ] + d_bystim[:prsem]
    ppsychid = plot(d_bystim, x="stimx", y="prμ", ymin="prmin", ymax="prmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("left choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
        tp)
    return pchronid, ppsychid
end
function plotchronpsychidm(d_bystim, m_bystim, stimx, xlabel, ttitle, t=Theme())
    d_bystim[:stimx] = stimx[d_bystim[:stimid]]
    m_bystim[:stimx] = stimx[m_bystim[:stimid]]
    xmin, xmax = minimum(stimx), maximum(stimx)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bystim[:tmin] = d_bystim[:tμ] - d_bystim[:tsem]
    d_bystim[:tmax] = d_bystim[:tμ] + d_bystim[:tsem]
    pchronid = plot(
        layer(d_bystim, x="stimx", y="tμ", ymin="tmin", ymax="tmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bystim, x="stimx", y="tμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bystim[:prmin] = d_bystim[:prμ] - d_bystim[:prsem]
    d_bystim[:prmax] = d_bystim[:prμ] + d_bystim[:prsem]
    ppsychid = plot(
        layer(d_bystim, x="stimx", y="prμ", ymin="prmin", ymax="prmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bystim, x="stimx", y="prμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("left choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
        tp)
    return pchronid, ppsychid
end

# returns psych/chron over stimulus magnitude
function plotchronpsychmag(d_bymag, magctrst, xlabel, ttitle, t=Theme())
    d_bymag[:stimctrst] = magctrst[d_bymag[:stimmag]]
    xmin, xmax = minimum(magctrst), maximum(magctrst)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bymag[:tmin] = d_bymag[:tμ] - d_bymag[:tsem]
    d_bymag[:tmax] = d_bymag[:tμ] + d_bymag[:tsem]
    pchronmag = plot(d_bymag, x="stimctrst", y="tμ", ymin="tmin", ymax="tmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bymag[:pcorrmin] = d_bymag[:pcorrμ] - d_bymag[:pcorrsem]
    d_bymag[:pcorrmax] = d_bymag[:pcorrμ] + d_bymag[:pcorrsem]
    ppsychmag = plot(d_bymag, x="stimctrst", y="pcorrμ", ymin="pcorrmin", ymax="pcorrmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("correct choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.5, maxvalue=1.0),
        tp)
    return pchronmag, ppsychmag
end
function plotchronpsychmagm(d_bymag, m_bymag, magctrst, xlabel, ttitle, t=Theme())
    d_bymag[:stimctrst] = magctrst[d_bymag[:stimmag]]
    m_bymag[:stimctrst] = magctrst[m_bymag[:stimmag]]
    xmin, xmax = minimum(magctrst), maximum(magctrst)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bymag[:tmin] = d_bymag[:tμ] - d_bymag[:tsem]
    d_bymag[:tmax] = d_bymag[:tμ] + d_bymag[:tsem]
    pchronmag = plot(
        layer(d_bymag, x="stimctrst", y="tμ", ymin="tmin", ymax="tmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bymag, x="stimctrst", y="tμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bymag[:pcorrmin] = d_bymag[:pcorrμ] - d_bymag[:pcorrsem]
    d_bymag[:pcorrmax] = d_bymag[:pcorrμ] + d_bymag[:pcorrsem]
    ppsychmag = plot(
        layer(d_bymag, x="stimctrst", y="pcorrμ", ymin="pcorrmin", ymax="pcorrmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bymag, x="stimctrst", y="pcorrμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("correct choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.5, maxvalue=1.0),
        tp)
    return pchronmag, ppsychmag
end

# returns array of conditional psychometric curves/fits
function plotcondpsych(d_seq, d_seqfit, stimx, xlabel, ttitle, t=Theme())
    # ignore the actual x stimulus distribution, plot on [0..1] scale
    stimx = range(0, stop=1, length=8)
    d_seq[:stimx] = stimx[d_seq[:stimid]]
    stimidcol(stimid) = RGB(0.0, 1 - 0.125stimid, 0.125stimid)
    pcondpsych = Vector{Any}(undef, 4)
    xcont = range(0.0, stop=1.0, length=100)
    globalθ = Vector{Float64}(
        d_seqfit[(d_seqfit[:prevrew] .== true) .& (d_seqfit[:previd] .== 0),
                 [:μ, :σ, :pl, :pb]][1,:])
    for stimmag = 1:4
        id1, id2 = stimmag, 9-stimmag
        pr1 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr1sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        θ1 = Vector{Float64}(
            d_seqfit[(d_seqfit[:prevrew] .== true) .& (d_seqfit[:previd] .== id1),
                     [:μ, :σ, :pl, :pb]][1,:])
        pr2 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr2sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        θ2 = Vector{Float64}(
            d_seqfit[(d_seqfit[:prevrew] .== true) .& (d_seqfit[:previd] .== id2),
                     [:μ, :σ, :pl, :pb]][1,:])

        pcondpsych[stimmag] = plot(
            Scale.x_continuous(minvalue=0.0, maxvalue=1.0),
            Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
            layer(x = xcont, y = psychpy1(globalθ, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=colorant"black"), order=0),
            layer(x = xcont, y = psychpy1(θ1, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=stimidcol(id1)), order=1),
            layer(x = xcont, y = psychpy1(θ2, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=stimidcol(id2)), order=2),
            layer(x = stimx, y = pr1, ymin = max.(0.0, pr1-pr1sem), ymax = min.(1.0, pr1+pr1sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id1)), order=3),
            layer(x = stimx, y = pr2, ymin = max.(0.0, pr2-pr2sem), ymax = min.(1.0, pr2+pr2sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id2)), order=4),
            Guide.xlabel(xlabel, orientation=:horizontal),
            Guide.ylabel("left choices (fraction)", orientation=:vertical),
            t)
    end

    return pcondpsych
end
# returns array of conditional psychometric curves/fits
function plotcondpsychm(d_seq, m_seq, m_seqfit, stimx, xlabel, ttitle, t=Theme())
    # ignore the actual x stimulus distribution, plot on [0..1] scale
    stimx = range(0, stop=1, length=8)
    d_seq[:stimx] = stimx[d_seq[:stimid]]
    stimidcol(stimid) = RGB(0.0, 1 - 0.125stimid, 0.125stimid)
    fitidcol(stimid) = RGB(0.5, 1 - 0.0625stimid, 0.5 + 0.0625stimid)
    pcondpsych = Vector{Any}(undef, 4)
    xcont = range(0.0, stop=1.0, length=100)
    for stimmag = 1:4
        id1, id2 = stimmag, 9-stimmag
        pr1 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr1sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        mpr1 = [m_seq[(m_seq[:prevrew] .== true) .& (m_seq[:previd] .== id1) .&
                      (m_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        θ1 = Vector{Float64}(
            m_seqfit[(m_seqfit[:prevrew] .== true) .& (m_seqfit[:previd] .== id1),
                     [:μ, :σ, :pl, :pb]][1,:])
        pr2 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr2sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        mpr2 = [m_seq[(m_seq[:prevrew] .== true) .& (m_seq[:previd] .== id2) .&
                      (m_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        θ2 = Vector{Float64}(
            m_seqfit[(m_seqfit[:prevrew] .== true) .& (m_seqfit[:previd] .== id2),
                     [:μ, :σ, :pl, :pb]][1,:])

        pcondpsych[stimmag] = plot(
            Scale.x_continuous(minvalue=0.0, maxvalue=1.0),
            Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
            layer(x = xcont, y = psychpy1(θ1, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=fitidcol(id1)), order=0),
            layer(x = xcont, y = psychpy1(θ2, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=fitidcol(id2)), order=1),
            layer(x = stimx, y = mpr1, Geom.line,
                  Theme(default_color=stimidcol(id1)), order=2),
            layer(x = stimx, y = mpr2, Geom.line,
                  Theme(default_color=stimidcol(id2)), order=3),
            layer(x = stimx, y = pr1, ymin = max.(0.0, pr1-pr1sem), ymax = min.(1.0, pr1+pr1sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id1)), order=4),
            layer(x = stimx, y = pr2, ymin = max.(0.0, pr2-pr2sem), ymax = min.(1.0, pr2+pr2sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id2)), order=5),
            Guide.xlabel(xlabel, orientation=:horizontal),
            Guide.ylabel("left choices (fraction)", orientation=:vertical),
            t)
    end

    return pcondpsych
end

# returns bias plot over simulus id
function plotbiasid(d_biasid, stimx, xlabel, ttitle, t=Theme())
    d_biasid[:prevx] = stimx[d_biasid[:previd]]
    d_biasid[:biasmin] = d_biasid[:biasμ] - d_biasid[:biasσ]
    d_biasid[:biasmax] = d_biasid[:biasμ] + d_biasid[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasid = plot(d_biasid[d_biasid[:prevrew],:],
        x="prevx", y="biasμ", ymin="biasmin", ymax="biasmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(stimx), maxvalue=maximum(stimx)),
        Scale.y_continuous(minvalue=-0.5, maxvalue=0.5),
        tp)
    return pbiasid
end
function plotbiasidm(d_biasid, m_biasid, stimx, xlabel, ttitle, t=Theme())
    d_biasid[:prevx] = stimx[d_biasid[:previd]]
    m_biasid[:prevx] = stimx[m_biasid[:previd]]
    d_biasid[:biasmin] = d_biasid[:biasμ] - d_biasid[:biasσ]
    d_biasid[:biasmax] = d_biasid[:biasμ] + d_biasid[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasid = plot(
        layer(d_biasid[d_biasid[:prevrew],:],
            x="prevx", y="biasμ", ymin="biasmin", ymax="biasmax",
            Geom.point, Geom.errorbar),
        layer(m_biasid[m_biasid[:prevrew],:],
            x="prevx", y="biasμ", Geom.line),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(stimx), maxvalue=maximum(stimx)),
        Scale.y_continuous(minvalue=-0.5, maxvalue=0.5),
        tp)
    return pbiasid
end

# returns bias plot over stimulus magnitude
function plotbiasmag(d_biasmag, magctrst, xlabel, ttitle, t=Theme())
    d_biasmag[:prevctrst] = magctrst[d_biasmag[:prevmag]]
    d_biasmag[:biasmin] = d_biasmag[:biasμ] - d_biasmag[:biasσ]
    d_biasmag[:biasmax] = d_biasmag[:biasμ] + d_biasmag[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasmag = plot(d_biasmag[d_biasmag[:prevrew],:],
        x="prevctrst", y="biasμ", ymin="biasmin", ymax="biasmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(magctrst), maxvalue=maximum(magctrst)),
        Scale.y_continuous(minvalue=-0.05, maxvalue=0.25),
        tp)
    return pbiasmag
end
function plotbiasmagm(d_biasmag, m_biasmag, magctrst, xlabel, ttitle, t=Theme())
    d_biasmag[:prevctrst] = magctrst[d_biasmag[:prevmag]]
    m_biasmag[:prevctrst] = magctrst[m_biasmag[:prevmag]]
    d_biasmag[:biasmin] = d_biasmag[:biasμ] - d_biasmag[:biasσ]
    d_biasmag[:biasmax] = d_biasmag[:biasμ] + d_biasmag[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasmag = plot(
        layer(d_biasmag[d_biasmag[:prevrew],:],
            x="prevctrst", y="biasμ", ymin="biasmin", ymax="biasmax",
            Geom.point, Geom.errorbar),
        layer(m_biasmag[m_biasmag[:prevrew],:],
            x="prevctrst", y="biasμ", Geom.line),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(magctrst), maxvalue=maximum(magctrst)),
        Scale.y_continuous(minvalue=-0.05, maxvalue=0.25),
        tp)
    return pbiasmag
end


using GPUifyLoops: @launch, @loop
using Oceananigans.Architectures: device
using Oceananigans.Utils: @loop_xyz

"""
    set!(model; kwargs...)

Set velocity and tracer fields of `model`. The keyword arguments
`kwargs...` take the form `name=data`, where `name` refers to one of the
fields of `model.velocities` or `model.tracers`, and the `data` may be an array,
a function with arguments `(x, y, z)`, or any data type for which a
`set!(ϕ::AbstractField, data)` function exists.

Example
=======
```julia
model = Model(grid=RegularCartesianGrid(size=(32, 32, 32), length=(1, 1, 1))

# Set u to a parabolic function of z, v to random numbers damped
# at top and bottom, and T to some silly array of half zeros,
# half random numbers.

u₀(x, y, z) = z/model.grid.Lz * (1 + z/model.grid.Lz)
v₀(x, y, z) = 1e-3 * rand() * u₀(x, y, z)

T₀ = rand(size(model.grid)...)
T₀[T₀ .< 0.5] .= 0

set!(model, u=u₀, v=v₀, T=T₀)
```
"""
function set!(model; kwargs...)
    for (fldname, value) in kwargs
        if fldname ∈ propertynames(model.velocities)
            ϕ = getproperty(model.velocities, fldname)
        elseif fldname ∈ propertynames(model.tracers)
            ϕ = getproperty(model.tracers, fldname)
        else
            throw(ArgumentError("name $fldname not found in model.velocities or model.tracers."))
        end
        set!(ϕ, value)
    end
    return nothing
end

function set!(Φ::NamedTuple; kwargs...)
    for (fldname, value) in kwargs
        ϕ = getproperty(Φ, fldname)
        set!(ϕ, value)
    end
    return nothing
end

set!(u::Field, v::Number) = @. u.data = v

set!(u::Field{X, Y, Z, A}, v::Field{X, Y, Z, A}) where {X, Y, Z, A} =
    @. u.data.parent = v.data.parent

"Set the CPU field `u` to the array `v`."
function set!(u::AbstractCPUField, v::Array)
    for k in 1:u.grid.Nz, j in 1:u.grid.Ny, i in 1:u.grid.Nx
        u[i, j, k] = v[i, j, k]
    end
    return nothing
end

# Set the GPU field `u` to the array `v`.
@hascuda function set!(u::AbstractGPUField, v::Array)
    v_field = Field(location(u), CPU(), u.grid)
    set!(v_field, v)
    set!(u, v_field)
    return nothing
end

# Set the GPU field `u` to the CuArray `v`.
@hascuda function set!(u::AbstractGPUField, v::CuArray)
    @launch device(GPU()) config=launch_config(u.grid, :xyz) _set_gpu!(u.data, v, u.grid)
    return nothing
end

function _set_gpu!(u, v, grid)
	@loop_xyz i j k grid begin
        u[i, j, k] = v[i, j, k]
    end
    return nothing
end

# Set the GPU field `u` data to the CPU field data of `v`.
@hascuda set!(u::AbstractGPUField, v::AbstractCPUField) = copyto!(u.data.parent, v.data.parent)

# Set the CPU field `u` data to the GPU field data of `v`.
@hascuda set!(u::AbstractCPUField, v::AbstractGPUField) = u.data.parent .= Array(v.data.parent)

"Set the CPU field `u` data to the function `f(x, y, z)`."
set!(u::AbstractCPUField, f::Function) = interior(u) .= f.(nodes(u)...)

# Set the GPU field `u` data to the function `f(x, y, z)`.
@hascuda function set!(u::AbstractGPUField, f::Function)
    u_cpu = Field(location(u), CPU(), u.grid)
    set!(u_cpu, f)
    set!(u, u_cpu)
    return nothing
end


using Elefridge
using PyPlot
using PyCall
xr = pyimport("xarray")
using TranscodingStreams, CodecZstd
using JLD

ZstdCompressorL3 = ZstdCompressor(level=3)
TranscodingStreams.initialize(ZstdCompressorL3)

ZstdCompressorL10 = ZstdCompressor(level=10)
TranscodingStreams.initialize(ZstdCompressorL10)

ZstdCompressorL22 = ZstdCompressor(level=22)
TranscodingStreams.initialize(ZstdCompressorL22)

path = "/Users/milan/cams/gridded"
filelist = filter(x->endswith(x,".grib"),readdir(path))

n = length(filelist)
varnames = fill("",n)
methods = ["ZstdL3","ZstdL3+bittranspose",
            "ZstdL10","ZstdL10+bittranspose",
            "ZstdL22","ZstdL22+bittranspose"]
sbits = Array(1:23)
C = fill(0.0,n,length(methods),length(sbits))

for (i,file) in enumerate(filelist)
    varname = split(split(file,"0_")[end],".")[1]
    varnames[i] = varname

    println("---")
    println("Reading $varname")
    X = xr.open_dataarray(joinpath(path,file),engine="cfgrib").data

    # transpose array
    X = permutedims(X,[3,2,1])  # unravels as longitude, latitude, level
    Xs = X[1:225,1:225,:]       # only 1/8th of the globe

    for (is,s) in enumerate(sbits)
        print(s)
        Xsr = round(Xs,s)
        Xsrt = bittranspose(Xsr)

        Xsr8 = unsafe_wrap(Array, Ptr{UInt8}(pointer(Xsr)), sizeof(Xsr))
        Xsrt8 = unsafe_wrap(Array, Ptr{UInt8}(pointer(Xsrt)), sizeof(Xsrt))

        C[i,1,is] = sizeof(Xs)/sizeof(transcode(ZstdCompressorL3,Xsr8))
        C[i,2,is] = sizeof(Xs)/sizeof(transcode(ZstdCompressorL3,Xsrt8))
        println("ZstdL3 done.")
        C[i,3,is] = sizeof(Xs)/sizeof(transcode(ZstdCompressorL10,Xsr8))
        C[i,4,is] = sizeof(Xs)/sizeof(transcode(ZstdCompressorL10,Xsrt8))
        println("ZstdL10 done.")
        C[i,5,is] = sizeof(Xs)/sizeof(transcode(ZstdCompressorL22,Xsr8))
        C[i,6,is] = sizeof(Xs)/sizeof(transcode(ZstdCompressorL22,Xsrt8))
        println("ZstdL22 done.")
    end
    @save "/Users/milan/cams/entropy/compression_transpose_all_zstd_1_8.jld" C varnames methods sbits
end