text:
using Revise
using AutomotiveDrivingModels
using AutoViz
using AutomotiveSensors
using AutomotivePOMDPs
using GridInterpolations

using POMDPPolicies
using POMDPModelTools
using DiscreteValueIteration

using JLD2
using FileIO
using Reel
using Random
using POMDPs
using POMDPModelTools
using LinearAlgebra

using PedestrianAvoidancePOMDP
using EmergencyBrakingSystem

using CSV
using DataFrames

using ArgParse

s = ArgParseSettings()

@add_arg_table s begin

    "--algorithm"
        arg_type = String
        default = "EmergencyBrakingSystem"
        help = "choose among 'EmergencyBrakingSystem', 'PedestrianAvoidancePOMDP', 'PedestrianAvoidancePOMDP_EmergencyBrakingSystem'"
    "--policy_name"
        arg_type = String
        default = "no"
        help = "name for the policy"
    "--collisioncost"
        arg_type = Float64
        default = -1000.0
        help = "reward function - cost for collision: -500.0"
    "--actionloncost"
        arg_type = Float64
        default = -1.0
        help = "reward function - cost for longitudinal action: -1.0"
    "--actionlatcost"
        arg_type = Float64
        default = -1.0
        help = "reward function - cost for lateral action: -1.0"
    "--keepvelocityreward"
        arg_type = Float64
        default = 5.0
        help = "reward function - reward for keeping the velocity: 5.0"
    "--keeplanereward"
        arg_type = Float64
        default = 5.0
        help = "reward function - reward for keeping the lane center: 5.0"
    "--pedestrianbirthprobability"
        arg_type = Float64
        default = 0.6
        help = "probability of a new pedestrian behind an obstruction: 0.6"

end

parsed_args = parse_args(ARGS, s)

algorithm = parsed_args["algorithm"]

pomdp = SingleOCFPOMDP()
if ( occursin("lateral", parsed_args["policy_name"]) )
    pomdp = SingleOCFPOMDP(lateral_actions = [1.0, 0.0, -1.0])
else
    pomdp = SingleOCFPOMDP(lateral_actions = [0.0])
end

pomdp.COLLISION_COST = parsed_args["collisioncost"]
pomdp.ACTION_LON_COST = parsed_args["actionloncost"]
pomdp.ACTION_LAT_COST = parsed_args["actionlatcost"]
pomdp.KEEP_VELOCITY_REWARD = parsed_args["keepvelocityreward"]
pomdp.KEEP_LANE_REWARD = parsed_args["keeplanereward"]
pomdp.PROBABILITY_PEDESTRIAN_BIRTH = parsed_args["pedestrianbirthprobability"]
pomdp.γ = 0.95

println("algorithm: ", algorithm)
println("lateral actions: ", pomdp.lateral_actions)
parameters_pomdp = string(pomdp.COLLISION_COST, "_", pomdp.PROBABILITY_PEDESTRIAN_BIRTH, "_", pomdp.ACTION_LON_COST, "_", pomdp.ACTION_LAT_COST, "_", pomdp.KEEP_VELOCITY_REWARD, "_", pomdp.KEEP_LANE_REWARD, "_", pomdp.γ, "_")
policy_name = string(parsed_args["policy_name"], "_", parameters_pomdp, ".jld2")
log_filename = string("results_", algorithm, "_", policy_name, ".csv")

# train policy if required
if algorithm != "EmergencyBrakingSystem"

    println("Policy name: ", policy_name)

    println("pomdp.COLLISION_COST: ", pomdp.COLLISION_COST)
    println("pomdp.ACTION_LON_COST: ", pomdp.ACTION_LON_COST)
    println("pomdp.ACTION_LAT_COST: ", pomdp.ACTION_LAT_COST)
    println("pomdp.KEEP_VELOCITY_REWARD: ", pomdp.KEEP_VELOCITY_REWARD)
    println("pomdp.KEEP_LANE_REWARD: ", pomdp.KEEP_LANE_REWARD)
    println("pomdp.PROBABILITY_PEDESTRIAN_BIRTH: ", pomdp.PROBABILITY_PEDESTRIAN_BIRTH)

    solver = SparseValueIterationSolver(max_iterations=200, belres=1e-4, include_Q=true, verbose=true)
    mdp = UnderlyingMDP(pomdp);

    vi_policy = solve(solver, mdp)
    global policy
    policy = AlphaVectorPolicy(pomdp, vi_policy.qmat, vi_policy.action_map)

    # save policy!
    FileIO.save(string("../policy/", policy_name), "policy", policy)
end

# Evaluate EuroNCAP scenarios

scenarios = ["CPCN", "CPAN25", "CPAN75", "CPFA", "FP"]
vut_speeds = [10., 15., 20., 25., 30., 35., 40., 45., 50., 55., 60.]
vut_speeds = [50.]
vut_speeds = vut_speeds / 3.6

println("Algorithm to evaluate: ", algorithm, " Policy: ", policy_name, " (if specified)")

#rec, timestep, env, ego_vehicle, sensor, sensor_observations, risk, ttc, collision_rate, emergency_brake_request, prediction_obstacle, collision, ego_a
results = Vector[]
for scenario in scenarios
    if (scenario == "FP")
        hit_points = [-100., 200.]
    else
        if (scenario == "CPCN")
            hit_points = [0., 10., 20., 30., 40., 50., 800.] # additional scenario without crossing pedestrian
        else
            hit_points = [0., 10., 20., 30., 40., 50.]
        end
    end
    for hit_point in hit_points
        for vut_speed in vut_speeds
            ego_v = vut_speed
            (ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles, scenario_id) = PedestrianAvoidancePOMDP.generate_scenario(scenario, ego_v, hit_point)
            if algorithm == "EmergencyBrakingSystem"
                (rec, timestep, env, sensor, sensor_observations, ego_vehicle, ego_a, collision, collision_rate, ttc, risk, emergency_brake_request, prediction_obstacle) = EmergencyBrakingSystem.evaluate_scenario(ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles)
            elseif algorithm == "PedestrianAvoidancePOMDP"
                (rec, timestep, env, sensor, sensor_observations, ego_vehicle, ego_a, collision, belief, action_pomdp, collision_rate, ttc, risk, emergency_brake_request, prediction_obstacle) = PedestrianAvoidancePOMDP.evaluate_scenario(ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles, policy, algorithm, pomdp.PROBABILITY_PEDESTRIAN_BIRTH, pomdp.lateral_actions)
            elseif algorithm == "PedestrianAvoidancePOMDP_EmergencyBrakingSystem"
                (rec, timestep, env, sensor, sensor_observations, ego_vehicle, ego_a, collision, belief, action_pomdp, collision_rate, ttc, risk, emergency_brake_request, prediction_obstacle) = PedestrianAvoidancePOMDP.evaluate_scenario(ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles, policy, algorithm, pomdp.PROBABILITY_PEDESTRIAN_BIRTH, pomdp.lateral_actions)
            else
                println("No valid algorithm defined!")
                return false
            end
            println(scenario, " HP: ", hit_point, " ", ego_x, " ", ego_y, " ", ego_v, " ", ped_x," ",  ped_y," ",  ped_v)
            (collision, emergency_brake_intervention, dv_collision, v_mean, a_mean, a_jerk, a_min) = PedestrianAvoidancePOMDP.evaluateScenarioMetric(ego_vehicle, emergency_brake_request, ego_a, collision, ped_x)
            println("Collision: ", collision, " eb: ", emergency_brake_intervention, " ", dv_collision, " ", v_mean, " ", a_mean, " ", a_jerk, " ", a_min)
            result = [scenario_id, ego_v, hit_point, collision, emergency_brake_intervention, dv_collision, v_mean, a_mean, a_jerk, a_min  ]
            push!(results, result)
        end
    end
end

# store results from scneario evaluation in log file
df = DataFrame(results)
df = DataFrame(Matrix(df)')
rename!(df, :x1 => :scenario_id, :x2 => :ego_v, :x3 => :hit_point, :x4 => :collision, :x5 => :eb_intervention)
rename!(df, :x6 => :dv_collision, :x7 => :v_mean, :x8 => :a_mean, :x9 => :a_jerk, :x10 => :a_min)

try
    mkdir(string("../results/", algorithm, "_", parsed_args["policy_name"]) )
catch
end
CSV.write(string("../results/", algorithm, "_", parsed_args["policy_name"], "/", log_filename), df);

(sum_collisions, sum_eb, dv, v_mean, a_mean, a_jerk, a_min) = PedestrianAvoidancePOMDP.evaluateScenariosMetric(results)

println("#collisions: ", sum_collisions, " #emergency brake interventions: ", sum_eb, " dv_m: ", dv, " v_mean: ", v_mean, " a_mean: ", a_mean, " a_jerk: ", a_jerk, " a_min: ", a_min)


text:
###
### parsing
###

const OPTIONS = Parsers.Options(
    sentinel=missing,
    quoted=true,
    openquotechar='\'',
    closequotechar='\'',
    delim=',',
)
# Change delimiter as way to handle end-of-line comments.
const EOL_OPTIONS = Parsers.Options(
    sentinel=missing,
    quoted=true,
    openquotechar='\'',
    closequotechar='\'',
    delim='/',
)

getbytes(source::Vector{UInt8}) = source, 1, length(source)
getbytes(source::IOBuffer) = source.data, source.ptr, source.size
getbytes(source) = getbytes(read(source))

"""
    parse_network(source) -> Network

Read a PSS/E-format `.raw` Power Flow Data file and return a [`Network`](@ref) object.
"""
function parse_network(source)
    bytes, pos, len = getbytes(source)

    caseid, pos = parse_caseid(bytes, pos, len, OPTIONS)
    @debug "caseid" pos

    # Skip the 2 lines of comments
    # TODO: confirm it is always only and exactly 2 lines of comments
    pos = next_line(bytes, pos, len)
    pos = next_line(bytes, pos, len)
    @debug "comments" pos

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "buses" nrows pos
    buses, pos = parse_records!(Buses(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "loads" nrows pos
    loads, pos = parse_records!(Loads(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "gens" nrows pos
    gens, pos = parse_records!(Generators(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    nrows = count_nrow(bytes, pos, len, OPTIONS)
    @debug "branches" nrows pos
    branches, pos = parse_records!(Branches(nrows), bytes, pos, len, OPTIONS, EOL_OPTIONS)

    return Network(caseid, buses, loads, gens, branches)
end

function parse_caseid(bytes, pos, len, options)
    ic, pos, code = parse_value(Int, bytes, pos, len, options)
    @debug codes(code) pos newline=newline(code)

    # Support files that have first row like:
    # 0,   100.00          / PSS/E-30.3    WED, SEP 15 2021  21:04
    # and those with first row like:
    # 0,100.0,30 / PSS(tm)E-30 RAW created      Wed, Sep 15 2021 21:04
    # TODO: avoid needing to extract `sbase` value from String?
    sbase = Parsers.tryparse(Float64, bytes, options, pos, len)
    if sbase === nothing
        str, pos, code = parse_value(String, bytes, pos, len, options)
        @debug codes(code) pos newline=newline(code)
        sbase = parse(Float64, first(split(str, '/')))
    end

    # if delimiter after `sbase` value, then won't have reached end of line.
    if !newline(code)
        pos = next_line(bytes, pos, len)
    end
    @debug codes(code) pos newline=newline(code)

    return CaseID(ic, sbase), pos
end

function parse_records!(rec::R, bytes, pos, len, options, eol_options)::Tuple{R, Int} where {R <: Records}
    nrows = length(getfield(rec, 1))
    nrows == 0 && return rec, pos
    for row in 1:nrows
        pos, code = parse_row!(rec, row, bytes, pos, len, options, eol_options)

        # Because we're working around end-of-line comments,
        # rows with comments won't have hit the newline character yet
        if !newline(code)
            pos = next_line(bytes, pos, len)
        end
    end

    # Data input is terminated by specifying a bus number of zero.
    # @assert peekbyte(bytes, pos) == UInt8('0')
    if !(eof(bytes, pos, len) || peekbyte(bytes, pos) == UInt8('0'))
        @warn "Not at end of $(typeof(rec)) records"
    end
    pos = next_line(bytes, pos, len)
    return rec, pos
end

function count_nrow(buf, pos, len, options)
    nlines = 0
    if eof(buf, pos, len) || peekbyte(buf, pos) == UInt8('0')
        return nlines
    end
    while true
        res = xparse(String, buf, pos, len, options)
        pos += res.tlen
        if newline(res.code) || eof(res.code)
            nlines += 1
            if eof(buf, pos, len) || peekbyte(buf, pos) == UInt8('0')
                break
            end
        end
    end
    return nlines
end

# Taken from `Parsers.checkcmtemptylines`
# TODO: move to Parsers.jl?
function next_line(bytes, pos, len)
    eof(bytes, pos, len) && return pos
    b = peekbyte(bytes, pos)
    while b !== UInt8('\n') && b !== UInt8('\r')
        pos += 1
        eof(bytes, pos, len) && break
        b = peekbyte(bytes, pos)
    end
    # Move forward to be past the `\r` or `\n` byte.
    pos += 1
    # if line ends `\r\n`, then we're at `\n`and need to move forward again.
    if b === UInt8('\r') && !eof(bytes, pos, len) && peekbyte(bytes, pos) === UInt8('\n')
        pos += 1
    end
    return pos
end

function parse_row!(rec::Records, row::Int, bytes, pos, len, options, eol_options)
    ncols = nfields(rec)
    local code::Parsers.ReturnCode
    for col in 1:ncols
        eltyp = eltype(fieldtype(typeof(rec), col))
        opts = ifelse(col == ncols, eol_options, options)
        val, pos, code = parse_value(eltyp, bytes, pos, len, opts)
        @inbounds getfield(rec, col)[row] = val

        @debug codes(code) row col pos newline=newline(code)
    end
    return pos, code
end

function parse_value(T, bytes, pos, len, options)
    res = xparse(T, bytes, pos, len, options)

    invalid(res.code) && @warn codes(res.code) pos

    pos += res.tlen
    code = res.code

    val = if T === String
        Parsers.getstring(bytes, res.val, options.e)
    else
        res.val
    end
    return val, pos, code
end

ext: jl


text:
################################################################
"""
$(TYPEDSIGNATURES)

Plot the observed data as a scatter plot.
"""
function plotobsdata(dobs)

    sdscon = dobs.sdsprotcon[:,1]
    procon = dobs.sdsprotcon[:,2]
    protein = dobs.protein

    PyPlot.figure(figsize=(11,6))
    PyPlot.title("$protein observed data")
    #tricontour(sdscon,procon,dobs,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.scatter(sdscon,procon,c=dobs.enthalpy,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")

    PyPlot.tight_layout()
    return
end

####################################################################

"""
$(TYPEDSIGNATURES)

Plot the fit to the enthalpy data as a result of the initial guess, i.e., the starting model parameters of the Beta functions.
"""
function plotinitialguess(betpar,dobs,mstart)

    dcalcstart = forwmod2D(betpar,dobs.sdsprotcon,mstart)
    sdscon = dobs.sdsprotcon[:,1]
    procon = dobs.sdsprotcon[:,2]
    protein = dobs.protein

    #lenm = length(mstart)
    nummodparam = size(mstart,1) #betpar.nummodpar
    #ncomp = size(mstart,2) #div(lenm,nummodparam)
    ym = LinRange(betpar.ymin,betpar.ymax,40)

    PyPlot.figure(figsize=(12,5))
    #suptitle("Observed data and initial guess")

    PyPlot.subplot(121)
    PyPlot.title("$protein observed data")
    #tricontour(sdscon,procon,dobs,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.scatter(sdscon,procon,c=dobs.enthalpy,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")
    plotmodelines(betpar,mstart,"start. m.")

    PyPlot.subplot(122)
    PyPlot.title("Calculated data from starting model")
    #tricontour(sdscon,procon,dcalcstart,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.scatter(sdscon,procon,c=dcalcstart,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")
    plotmodelines(betpar,mstart,"start. m.")

    PyPlot.tight_layout()
    return
end

############################################
"""
$(TYPEDSIGNATURES)

Plot the results of inverting the ITC data to fit the enthalpy function in 2D.
"""
function plotresults(betamix,dobs,mstart,outdir)

    mpost = betamix.modkonamp
    betpar = betamix.betpar
    dcalcstart = forwmod2D(betpar,dobs.sdsprotcon,mstart)
    dcalccur = forwmod2D(betpar,dobs.sdsprotcon,mpost)
    sdscon = dobs.sdsprotcon[:,1]
    procon = dobs.sdsprotcon[:,2]
    protein = dobs.protein

    ##==============================
    ## show data
    lenm = length(mstart)
    nummodparam = betpar.nummodpar
    #ncomp = div(lenm,nummodparam)
    ym = LinRange(betpar.ymin,betpar.ymax,40)

    PyPlot.figure(figsize=(14,7))
    PyPlot.subplot(231)
    PyPlot.title("$protein observed data")
    #tricontour(sdscon,procon,dobs,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.scatter(sdscon,procon,c=dobs.enthalpy,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")
    plotmodelines(betpar,mstart,"start. m.")
    # xlim([-2,betpar.b])
    # ylim([0.9*betpar.ymin,1.05*betpar.ymax])
    #plotmodelines(betpar,mpost,"mpost")

    PyPlot.subplot(232)
    PyPlot.title("Calculated data from starting model")
    #tricontour(sdscon,procon,dcalcstart,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.scatter(sdscon,procon,c=dcalcstart,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")
    plotmodelines(betpar,mstart,"start. m.")

    PyPlot.subplot(233)
    PyPlot.title("Model parameters")
    PyPlot.plot(vec(mstart),".-",label="starting model")
    PyPlot.plot(vec(mpost),".-",label="posterior model")
    PyPlot.legend()

    PyPlot.subplot(234)
    PyPlot.title("Calculated data from posterior model")
    #tricontour(sdscon,procon,dcalccur,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.scatter(sdscon,procon,c=dcalccur,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")
    plotmodelines(betpar,mpost,"post. m.")

    PyPlot.subplot(235)
    PyPlot.title("Calculated minus observed data")
    #tricontour(sdscon,procon,dcalccur,cmap=PyPlot.get_cmap("rainbow"))
    dcmo = dcalccur-dobs.enthalpy
    vmax = maximum(abs.(dcmo))
    PyPlot.scatter(sdscon,procon,c=dcmo,vmin=-vmax,vmax=vmax,cmap=PyPlot.get_cmap("RdBu"))
    PyPlot.colorbar(label="ΔEnthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")
    # xlim([-2,betpar.b])
    # ylim([0.9*betpar.ymin,1.05*betpar.ymax])

    PyPlot.subplot(236)
    PyPlot.title("Calculated vs. observed data")
    PyPlot.plot(dcalcstart,"--",label="calc. data from start. mod.",linewidth=0.8)
    PyPlot.plot(dobs.enthalpy,"-k",label="observed data")
    PyPlot.plot(dcalccur,"-r",label="calc. data from post. mod.")
    PyPlot.legend()

    # subplot(235)
    # title("Misfit")
    # plot(misfall,"o-")
    # xlabel("Iteration number")
    # ylabel("misfit")

    PyPlot.tight_layout()

    try
        mkdir(outdir)
    catch
        nothing
    end
    PyPlot.savefig(joinpath(outdir,protein*"_results.pdf"))
    return nothing
end

############################################

"""
$(TYPEDSIGNATURES)

Plot the lines defined by the modes (default) or, optionally, other
  parameters (see `firstidpar`).
"""
function plotmodelines(betpar,mcur,modname; firstidpar=1)
    #Npoints = 40
    ncomp = size(mcur,2)
    ym = [betpar.ymin,betpar.ymax] #LinRange(betpar.ymin,betpar.ymax,Npoints)
    y1 = betpar.ymin
    y2 = betpar.ymax
    for i=1:ncomp
        i2 = firstidpar+1
        i1 = firstidpar
        xm = (mcur[i2,i].-mcur[i1,i])./(y2.-y1).*(ym.-y1) .+ mcur[i1,i]
        if i==ncomp
            PyPlot.plot(xm,ym,"-k",linewidth=0.6,label="mode $modname")
        else
            PyPlot.plot(xm,ym,"-k",linewidth=0.6)
        end
        PyPlot.plot(xm[[1,end]],ym[[1,end]],"ok",markersize=2.0)
    end
    PyPlot.legend(loc="lower right")
    # xlim([-2,betpar.b])
    # ylim([0.9*betpar.ymin,1.05*betpar.ymax])
    return
end

############################################

"""
$(TYPEDSIGNATURES)

Plot the lines defined by the model parameters as a function of protein concentration.
"""
function plotparamlines(betamix,protcons=nothing,areas=nothing)

    betpar = betamix.betpar
    mcur = betamix.modkonamp

    @assert betpar.modefuny=="linear"
    #@assert betpar.konfuny=="linear"
    if betpar.konfuny=="constant"

    elseif betpar.konfuny=="linear"

    end
    @assert betpar.ampfuny=="linear"

    Npoints = 100
    ncomp = size(mcur,2)
    xm = LinRange(0.0,betpar.ymax,Npoints)
    parname = ["Mode","Confidence par.","Amplitude"]

    PyPlot.figure(figsize=(12,9))
    ylab = ["SDS concentration [mM]",
            "Confidence parameter",
            "Enthalpy"]

    for ip=1:3

        PyPlot.subplot(2,2,ip)
        PyPlot.title("$(parname[ip])")
        PyPlot.xlabel("Protein concentration [mM]")
        PyPlot.axvline(betpar.ymin,color="black")
        PyPlot.axvline(betpar.ymax,color="black")

        if ip==1
            PyPlot.axhline(betpar.a,color="black")
            PyPlot.axhline(betpar.b,color="black")
        end

        for i=1:ncomp
            if ip==1
                firstidpar = betpar.idxmode
                i2 = firstidpar+1
                i1 = firstidpar
                x1 = betpar.ymin
                x2 = betpar.ymax
                ym = (mcur[i2,i].-mcur[i1,i])./(x2.-x1).*(xm.-x1) .+ mcur[i1,i]

            elseif ip==2
                if betpar.konfuny=="constant"
                    firstidpar = betpar.idxkon
                    ym = mcur[firstidpar,i].*ones(length(xm))

                elseif betpar.konfuny=="linear"
                    firstidpar = betpar.idxkon
                    i2 = firstidpar+1
                    i1 = firstidpar
                    x1 = betpar.ymin
                    x2 = betpar.ymax
                    ym = (mcur[i2,i].-mcur[i1,i])./(x2.-x1).*(xm.-x1) .+ mcur[i1,i]

                end

            elseif ip==3
                firstidpar = betpar.idxamp
                i2 = firstidpar+1
                i1 = firstidpar
                x1 = betpar.ymin
                x2 = betpar.ymax
                ym = (mcur[i2,i].-mcur[i1,i])./(x2.-x1).*(xm.-x1) .+ mcur[i1,i]

            end

            PyPlot.plot(xm,ym,"-",linewidth=1,label="comp. $i")
            PyPlot.plot(xm[[1,end]],ym[[1,end]],".k",markersize=5.0)
            PyPlot.ylabel(ylab[ip])

        end
        PyPlot.legend()

        # xlim([-2,betpar.b])
        # ylim([0.9*betpar.ymin,1.05*betpar.ymax])
    end

    if areas!=nothing
        PyPlot.subplot(224)
        PyPlot.title("Areas of single components")
        for i=1:ncomp
            PyPlot.plot(protcons,areas[:,i],"o-",label="Beta comp. #$i")
        end
        PyPlot.legend()
        PyPlot.xlabel("Protein concentration [mM]")
        PyPlot.ylabel("Area [kJ]")
    end

    PyPlot.tight_layout()

    return
end

#####################################################

# """
# $(TYPEDSIGNATURES)

# Plot the results of binding isotherm calculations using the individual Beta functions.
# """
# function plotbindisotherm_singlebetas(protein,betpar,xy,mpost,sdsfNb,outdir; Npts=100)

#     dcalc = forwmod2D(betpar,xy,mpost)
#     sdscon = xy[:,1]
#     procon = xy[:,2]

#     figure(figsize=(13,5))

#     subplot(121)
#     title("Protein $protein fitted lines on given amplitude ratios")
#     scatter(sdscon,procon,c=dcalc,cmap=PyPlot.get_cmap("rainbow"))
#     colorbar()
#     xlabel("SDS concentration [mM]")
#     ylabel("Protein concentration [mM]")
#     ##
#     ##  y = m*x + y0
#     ##
#     ## One line per row
#     ##N = 100
#     firstplot=true
#     for ico=1:length(sdsfNb)
#         for r=1:size(sdsfNb[ico],1)
#             x = collect(LinRange(betpar.ymin,betpar.ymax,Npts))
#             y = sdsfNb[ico][r,2].*x .+ sdsfNb[ico][r,1]
#             if firstplot
#                 plot(y,x,"--r",linewidth=0.7,label="selected fitted amplitudes")
#                 firstplot = false
#             else
#                 plot(y,x,"--r",linewidth=0.7)
#             end
#         end
#     end

#     plotmodelines(betpar,mpost,"")

#     subplot(122)
#     title("Protein $protein: binding isotherm from individual Beta functions")
#     for ico=1:length(sdsfNb)
#         plot(sdsfNb[ico][:,1],sdsfNb[ico][:,2],"o-")
#     end
#     xlabel("free SDS conc. [mM]")
#     ylabel("Nbound")

#     tight_layout()

#     savefig(joinpath(outdir,protein*"_bind-isoth_singlebetas.pdf"))
#     return nothing
# end

######################################################################

"""
$(TYPEDSIGNATURES)

Save the fitting surface in the VTK format for Paraview.
"""
function saveresultVTK(protein,betpar,mpost)

    ## VTK stuff
    Nx=500
    Ny=500
    xvtk = collect(LinRange(betpar.a,betpar.b,Nx))
    yvtk = collect(LinRange(betpar.ymin,betpar.ymax,Ny))
    xy = zeros(Nx*Ny,2)
    for j=1:Ny, i=1:Nx
        xy[i+Nx*(j-1),:] .= (xvtk[i],yvtk[j])
    end
    vtkfile = vtk_grid("../output/"*protein*"_gridITCdata", xvtk, 100.0 .* yvtk)
    betparvtk = ScaledBeta2DParams(nummodpar=nummodpar,a=betpar.a,b=betpar.b,
                                   ymin=betpar.ymin,ymax=betpar.ymax)
    dcalcvtk = forwmod2D(betparvtk,xy,mpost)
    dcvtk2 = reshape(dcalcvtk,Nx,Ny)
    vtkfile["energyout",VTKPointData()] = dcvtk2
    outfiles = vtk_save(vtkfile)

    return
end

####################################################
"""
$(TYPEDSIGNATURES)

Plot the results of binding isotherm calculations using the Beta mix, i.e., sum of all Beta fitting functions.

"""
function plotbindisotherm(betamix,protcon,dobs,statpts,inflpts,freeSDS,Nbound,outdir;
                          resstdev=nothing)

    protein = betamix.protein
    N = 1000
    x = collect(LinRange(betamix.betpar.a,betamix.betpar.b,N))

    ##-------------------------------------
    PyPlot.figure(figsize=(12,10))

    PyPlot.subplot(211)
    PyPlot.title("Protein $protein: binding isotherm from stationary points of Beta mix")
    PyPlot.scatter(dobs.sdsprotcon[:,1],dobs.sdsprotcon[:,2],c=dobs.enthalpy,cmap=PyPlot.get_cmap("rainbow"))
    PyPlot.colorbar(label="Enthalpy [kJ/mol]")
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Protein concentration [mM]")

    ny = length(protcon)
    for i in 1:ny
        for z in 1:length(statpts[i])
            PyPlot.plot(statpts[i][z],protcon[i],"or")
            PyPlot.text(statpts[i][z],protcon[i],string(z),color="red")
        end

        for z in 1:length(inflpts[i])
            PyPlot.plot(inflpts[i][z],protcon[i],"ob")
            PyPlot.text(inflpts[i][z],protcon[i],string(z),color="blue")
        end
    end

    nlines = length(freeSDS)
    for i=1:nlines
        ## totSDS = freeSDS + Nbound * protc
        # y = angcoe.*x .+ intercept
        # x = (-intercept .+ 0)./angcoe
        ypl = (1.0/Nbound[i]).*x .+ (-freeSDS[i]/Nbound[i])
        PyPlot.plot(x,ypl,"-k",linewidth=0.5)
    end

    # xlim([betamix.betpar.a,betamix.betpar.b])
    PyPlot.ylim([0.92*betamix.betpar.ymin,1.05betamix.betpar.ymax])

    ##
    ## totSDS = freeSDS + Nbound * protc
    ##
    PyPlot.subplot(212)
    PyPlot.title("Binding isotherm for protein $protein")
    xbi = [0.0, freeSDS...]
    ybi = [0.0, Nbound...]
    if resstdev!=nothing
        resstdevbi = [0.0,resstdev...]
        PyPlot.scatter(xbi,ybi,c=resstdevbi,cmap=PyPlot.cm_get_cmap("rainbow"))
        PyPlot.colorbar(label="residuals standard dev.")
    else
        PyPlot.scatter(xbi,ybi)
    end
    PyPlot.plot(xbi,ybi,"-")
    PyPlot.xlabel("free SDS concentration [mM]")
    PyPlot.ylabel("Nbound")

    PyPlot.tight_layout()
    try
        mkdir(outdir)
    catch
        nothing
    end
    PyPlot.savefig(joinpath(outdir,protein*"_bind-isoth_betamix.pdf"))
    return
end

#########################################################

"""
$(TYPEDSIGNATURES)

Plot the points/features defined on the enthalpy curves to find the binding isotherm.
"""
function plotfoundfeatures(betamix,protcon,statpts,inflpts,outdir)

    protein = betamix.protein
    N = 1000
    x = collect(LinRange(betamix.betpar.a,betamix.betpar.b,N))
    ny = length(protcon)

    function fwd1ptxy(xpt::Real,ypt::Real)
        xy = hcat(xpt,ypt)
        out = forwmod2D(betamix.betpar,xy,betamix.modkonamp)
        return out[1]
    end

    PyPlot.figure(figsize=(12,6))
    PyPlot.title("Protein $protein: stationary and inflection points")
    for i in 1:ny
        y = [protcon[i] for _ in 1:N]
        slbet = fwd1ptxy.(x,y)
        PyPlot.plot(x,slbet,label="betamix at prot. conc. $(round(y[1],digits=4)) ")

        PyPlot.plot(statpts[i],fwd1ptxy.(statpts[i],protcon[i]),"or")
        for z in 1:length(statpts[i])
            PyPlot.text(statpts[i][z],fwd1ptxy.(statpts[i][z],protcon[i]),
                        string(z),color="red")
        end

        PyPlot.plot(inflpts[i],fwd1ptxy.(inflpts[i],protcon[i]),"ob")
        for z in 1:length(inflpts[i])
            PyPlot.text(inflpts[i][z],fwd1ptxy.(inflpts[i][z],protcon[i]),
                        string(z),color="blue")
        end
    end
    PyPlot.legend()
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Enthalpy [kJ/mol]")
    PyPlot.tight_layout()

    try
        mkdir(outdir)
    catch
        nothing
    end
    PyPlot.savefig(joinpath(outdir,protein*"_foundfeatures_betamix.pdf"))
    return
end

#####################################################

"""
$(TYPEDSIGNATURES)

Plot each single experiment, i.e., enthalpy for an initial protein concentration and increasing SDS concentration.
If the third argument `betamix` is passed, shows a comparison of measured and calculate data (from results of inversion) including each single Beta component.
"""
function plotsingleexperiments(outdir,dobs,betamix=nothing) #; expernumber=nothing)

    nexper = length(dobs.idxdata)
    obsdata = dobs.enthalpy
    protein = dobs.protein

    for i in 1:nexper

        idx = dobs.idxdata[i]
        xy = dobs.sdsprotcon[idx,:]
        obsexper = dobs.enthalpy[idx]

        PyPlot.figure(figsize=(12,6.4))
        PyPlot.title("Protein: $protein, initial concentration: $(xy[1,2])")
        PyPlot.plot(xy[:,1],obsexper,".-k",linewidth=2,label="measured data")

        if betamix!=nothing
            dcalcall = forwmod2D(betamix.betpar,xy,betamix.modkonamp)
            PyPlot.plot(xy[:,1],dcalcall,".-r",linewidth=2,label="calculated data")
            ncomp = size(betamix.modkonamp,2)
            for c=1:ncomp
                dcalc1 = singlescaledbeta2D(betamix.betpar,xy,betamix.modkonamp[:,c])
                PyPlot.plot(xy[:,1],dcalc1,".--",linewidth=0.7,markersize=0.75,label="Beta comp. $c")
            end
        end
        PyPlot.legend()
        PyPlot.xlabel("SDS concentration [mM]")
        PyPlot.ylabel("Enthalpy [kJ/mol")
        PyPlot.tight_layout()

        try
            mkdir(outdir)
        catch
            nothing
        end
        if betamix!=nothing
            PyPlot.savefig(joinpath(outdir,protein*"_experiment$(i)_fit.pdf"))
        else
            PyPlot.savefig(joinpath(outdir,protein*"_experiment$(i).pdf"))
        end
    end

    return
end

#####################################################
"""
$(TYPEDSIGNATURES)

"""
function plotareavsprotcon(proteinname,protcons,areas,linfitres,resstdev,outdir )

    ncomp = size(areas,2)
    PyPlot.figure()

    PyPlot.println("### Protein $proteinname, area vs. protein conc. ###")
    ##------------------------
    # linear fits
    ncomp = size(linfitres,1)
    x1 = protcons[1]
    x2 = protcons[end]
    xm = [x1,x2]
    for i=1:ncomp
        acoe   = linfitres[i,1]
        interc = linfitres[i,2]
        println("Linear fit, Beta comp. #$i, ang.coef.=$acoe,\n                           intecept=$interc")
        ym = acoe.*xm .+ interc
        PyPlot.plot(xm,ym,"-k",linewidth=1.5)
    end

    ##------------------------------
    # areas
    for i=1:ncomp
        rsd = round(resstdev[i],sigdigits=3)
        PyPlot.plot(protcons,areas[:,i],"o",
                    markersize=8,linewidth=2.5,label="Beta comp. #$i, res. std $rsd")
    end
    PyPlot.legend()
    PyPlot.xlabel("Protein concentration [mM]")
    PyPlot.ylabel("Area [kJ]")
    PyPlot.tight_layout()

    try
        mkdir(outdir)
    catch
        nothing
    end
    PyPlot.savefig(joinpath(outdir,proteinname*"_areavsprotcon.pdf"))
    return
end

#####################################################

"""
$(TYPEDSIGNATURES)

Plot the components and the sum of Beta functions for a given protein concentration.
"""
function plotbetacomp1D(betamix,protcon) #; expernumber=nothing)

    protein = betamix.protein

    # function fwd1ptxy(xpt::Real,ypt::Real)
    #     xy = hcat(xpt,ypt)
    #     out = forwmod2D(betamix.betpar,xy,betamix.modkonamp)
    #     return out[1]
    # end

    PyPlot.figure(figsize=(12,6.4))
    PyPlot.title("Protein: $protein, concentration: $protcon")

    N = 100
    x = collect(LinRange(betamix.betpar.a,betamix.betpar.b,N))
    y = protcon.*ones(N)
    xy = [x y]

    ncomp = size(betamix.modkonamp,2)
    for c=1:ncomp
        #dcalc1 = singlescaledbeta2D(betamix.betpar,xy,betamix.modkonamp[:,c])
        xy = [x y]
        enth = forwmod2D(betamix.betpar,xy,betamix.modkonamp[:,c:c])
        bpars = round.(betamix.modkonamp[:,c:c],digits=2)
        PyPlot.plot(x,enth,"--",linewidth=1,label="Beta comp. $c, $(bpars)")
        #@show c,betamix.modkonamp[:,c:c]
    end

    enth = forwmod2D(betamix.betpar,xy,betamix.modkonamp)
    PyPlot.plot(x,enth,"-",linewidth=2,color="red",label="sum of Beta func.")

    PyPlot.legend()
    PyPlot.xlabel("SDS concentration [mM]")
    PyPlot.ylabel("Enthalpy [kJ/mol")
    PyPlot.tight_layout()

    return
end

#############################################################################3


text:
using Serialization
using Parameters:@with_kw
using Statistics:mean
using CUDA
using Flux: cpu, Parallel
import Flux, Functors
using Base.Threads

CUDA.allowscalar(false)
array_on_gpu(::Array) = false
array_on_gpu(::CuArray) = true
array_on_gpu(arr) = error("Usupported array type: ", typeof(arr))

using Flux: relu, softmax, flatten
using Flux.Losses: mse, logitcrossentropy, crossentropy
using ParameterSchedulers: Cos,Stateful, next!
using Flux: Chain, Dense, Conv, BatchNorm, SkipConnection, MeanPool, MaxPool, AdaptiveMeanPool
import Zygote
using BSON: @save

#####
##### Support functions
#####

# This should be included in Flux
function loss_grad(f, args...)
    loss, back = Zygote.pullback(f, args...)
    grad = back(Zygote.sensitivity(loss))
    return loss, grad
end

# Invert the scaling (defined in https://arxiv.org/abs/1805.11593) of predicted_values #TODO
invert_scaling(x) = convert(Float32, sign(x) * (((sqrt(1 + 4 * 0.001 * (abs.(x) + 1 + 0.001)) - 1) / (2 * 0.001))^2 - 1))

# Reduce the scale (defined in https://arxiv.org/abs/1805.11593) of value and reward before feeding to network
scaling(x) = convert(Float32, sign(x) * (sqrt(abs(x) + 1) - 1 + 0.001*x)) # dims=(features, batch)

# function convert_input_tuple(nn, input::Tuple)
#     return map(input) do arr
#         array_on_gpu(nn.first_head[end].b) ? Flux.gpu(arr) : arr
#     end
# end

# function convert_output_tuple(output::Tuple)
# 	return map(output) do arr
# 		Flux.cpu(arr)
# 	end
# end

to_singletons(x) = reshape(x, size(x)..., 1)
squeeze(x) = reshape(x, size(x)[1:end - 1])
unsqueeze(xs::AbstractArray, dim::Integer) = reshape(xs, (size(xs)[1:dim - 1]..., 1, size(xs)[dim:end]...))

######
###### Networks
######

# custom split layer
struct Split{T}
  paths::T
end

Split(paths...) = Split(paths)

Flux.@functor Split

(m::Split)(x::AbstractArray) = map(f -> f(x), m.paths)

function make_dense(indim::Int, outdim::Int, bnmom::Float32, hyper::FeedForwardHP)
    if hyper.use_batch_norm
      Chain(
        Dense(indim, outdim),
        BatchNorm(outdim, relu, momentum=bnmom))
    else
      	Dense(indim, outdim, relu)
    end
end

hlayers(depth::Int, hsize, bnmom, hyper) = [make_dense(hsize, hsize, bnmom, hyper) for _ in 1:depth]

#######
####### FeedForward Networks
#######

function init_representation(hyper::FeedForwardHP)
	indim = prod([conf.observation_shape[1],conf.observation_shape[2], (conf.observation_shape[3]*(conf.stacked_observations+1)+conf.stacked_observations)])
	outdim = hyper.hidden_state_size
	bnmom = hyper.batch_norm_momentum
	hsize = hyper.width_hidden
	layers = Chain(flatten,
	make_dense(indim, hsize, bnmom, hyper),
	hlayers(hyper.depth_representation, hsize, bnmom, hyper)...,
	Dense(hsize, outdim)
	)
	return layers
end

function init_prediction(hyper::FeedForwardHP)
	bnmom = hyper.batch_norm_momentum
	indim = hyper.hidden_state_size
	outdim = length(conf.action_space)
	hsize = hyper.width_hidden
	common = Chain(flatten,
		make_dense(indim, hsize, bnmom, hyper),
		hlayers(hyper.depth_prediction, hsize, bnmom, hyper)...)
	value_head = Chain(
		hlayers(hyper.depth_value, hsize, bnmom, hyper)...,
		Dense(hsize, 1, tanh))
	policy_head = Chain(
		hlayers(hyper.depth_policy, hsize, bnmom, hyper)...,
		Dense(hsize, outdim),
		softmax)
	return Chain(common, Split(value_head, policy_head))
end

function init_dynamics(hyper::FeedForwardHP)
	bnmom = hyper.batch_norm_momentum
	indim = prod([conf.observation_shape[1],conf.observation_shape[2], (conf.observation_shape[3]+1)])
	outdim = hyper.hidden_state_size
	hsize = hyper.width_hidden
	# state_input = Chain(
	# flatten,
	# make_dense(indim, hsize, bnmom, hyper))
	# action_input = make_dense(outdim, hsize, bnmom, hyper)
	common = Chain(
		flatten,
		make_dense(indim, hsize, bnmom, hyper),
		# Parallel(vcat, state_input, action_input),
		# make_dense(2 * hsize, hsize, bnmom, hyper),
		hlayers(hyper.depth_dynamics, hsize, bnmom, hyper)...
	)
	state_head = Chain(
		hlayers(hyper.depth_state_head, hsize, bnmom, hyper)...,
		Dense(hsize, outdim)
	)
	reward_head = Chain(
		hlayers(hyper.depth_reward, hsize, bnmom, hyper)...,
		Dense(hsize, 1, hyper.reward_activation))
	return Chain(common, Split(state_head, reward_head))
end

########
######## ResNets
########

function resnet_block(size::Tuple{Int,Int}, n::Int, bnmom::Float32)
  pad = size .÷ 2
  layers = Chain(
    Conv(size, n => n, pad=pad),
    BatchNorm(n, relu, momentum=bnmom),
    Conv(size, n => n, pad=pad),
    BatchNorm(n, momentum=bnmom))
  return Chain(
    SkipConnection(layers, +),
    x -> relu.(x))
end

function init_representation(hyper::ResNetHP)
	indim = conf.observation_shape
	ksize = hyper.conv_kernel_size
	@assert all(ksize .% 2 .== 1)
	pad = ksize .÷ 2
	nf = hyper.num_filters
	bnmom = hyper.batch_norm_momentum

	common = Chain(
		Conv(ksize, indim[3] => nf, pad=pad),
		BatchNorm(nf, relu, momentum=bnmom),
		[resnet_block(ksize, nf, bnmom) for i in 1:hyper.num_blocks]...)

	hyper.representation_output_size = Flux.outputsize(common, ((indim)..., 1))

	if downsampling
	downsample = Chain(
		Conv(size, indim[3] => indim[3], stride=2, pad=pad),
		[resnet_block(size, indim[3], bnmom) for i in 1:2]...,
		Conv(size, indim[3] => indim[3] * 2, stride=2, pad=pad),
		[resnet_block(size, indim[3] * 2, bnmom) for i in 1:3]...,
		MeanPool(3, stride=2, pad=1),
		[resnet_block(size, indim[3] * 2, bnmom) for i in 1:3]...,
		MeanPool(3, stride=2, pad=1),
		Conv(ksize, indim[3] * 2 => nf, pad=pad),
		BatchNorm(nf, relu, momentum=bnmom),
		[resnet_block(ksize, nf, bnmom) for i in 1:hyper.num_blocks]...)
		return downsample
	else
		return common
	end
end

function init_prediction(hyper::ResNetHP)
	indim = hyper.representation_output_size
	outdim = length(conf.action_space)
	ksize = (1, 1)
	@assert all(ksize .% 2 .== 1)
	pad = ksize .÷ 2
	nf = hyper.num_filters
	npf = hyper.num_second_head_filters
	nvf = hyper.num_first_head_filters
	bnmom = hyper.batch_norm_momentum
	hsize = hyper.width_hidden
	common = Chain(
		Conv(ksize, indim[3] => nf, pad=pad),
		BatchNorm(nf, relu, momentum=bnmom),
		[resnet_block(ksize, nf, bnmom) for i in 1:hyper.num_blocks]...)

	value_head = Chain(
        Conv(ksize, nf => nvf),
        BatchNorm(nvf, relu, momentum=bnmom),
        flatten,
        Dense(indim[1] * indim[2] * nvf, hsize, relu),
		hlayers(hyper.depth_value, hsize, bnmom, hyper)...,
        Dense(hsize, 1, tanh))

	policy_head = Chain(
			Conv(ksize, nf => npf),
			BatchNorm(npf, relu, momentum=bnmom),
			flatten,
			Dense(indim[1] * indim[2] * npf, hsize),
			hlayers(hyper.depth_value, hsize, bnmom, hyper)...,
			Dense(hsize, outdim),
			softmax)
	return Chain(common, Split(value_head, policy_head))
end

function init_dynamics(hyper::ResNetHP)
	indim = hyper.representation_output_size
	ksize = (1, 1)
	@assert all(ksize .% 2 .== 1)
	pad = ksize .÷ 2
	nvf = hyper.num_first_head_filters
	bnmom = hyper.batch_norm_momentum
	hsize = hyper.width_hidden
	common = Chain(
		Conv(ksize, indim[3] + hyper.stacked_actions => indim[3], pad=pad),
		BatchNorm(indim[3], relu, momentum=bnmom),
		[resnet_block(ksize, indim[3], bnmom) for i in 1:hyper.num_blocks]...)

	state_head = Chain(
		Conv(ksize, indim[3] => indim[3], pad=pad),
		BatchNorm(indim[3], relu, momentum=bnmom),
		[resnet_block(ksize, indim[3], bnmom) for i in 1:hyper.num_blocks]...)

	reward_head = Chain(
        Conv(ksize, indim[3] => nvf),
        BatchNorm(nvf, relu, momentum=bnmom),
        flatten,
        Dense(indim[1] * indim[2] * nvf, hsize, relu),
		hlayers(hyper.depth_value, hsize, bnmom, hyper)...,
        Dense(hsize, 1, tanh))

		return Chain(common, Split(state_head, reward_head))
end

##########
##########  Training
##########

function loss(params, predictions::Tuple, targets::Tuple, weight_batch::Any, gradient_scale_batch::Matrix{Float32})::Float32
	value, reward, policy_logits = predictions
	target_values, target_rewards, target_policies = targets
	# By default we Correct PER bias by using importance-sampling (IS) weights
	# But if conf.PER is false then weight_batch is 1, and we avoid the correction
	if !conf.PER
		weight_batch = 1.0f0
	end

	# L2 regularization
	sqnorm(x) = sum(abs2, x)

    value_loss = mse(value, target_values, agg=x->mean((sum(x,dims=1)./gradient_scale_batch).*weight_batch))
	# TODO Scale the value loss, paper recommends by 0.25 (See paper appendix Reanalyze)

	if conf.intermediate_rewards
    	reward_loss = mse(reward, target_rewards, agg=x->mean((sum(x,dims=1)./gradient_scale_batch).*weight_batch))
	else
		reward_loss = 0.0f0
	end
    policy_loss = logitcrossentropy(policy_logits, target_policies, agg=x->mean((sum(x,dims=2)./gradient_scale_batch).*weight_batch))

	loss = sum([value_loss, reward_loss, policy_loss])
	# TODO value loss and reward loss are much smaller that the policy loss

	#loss + L2
    return loss + sum(sqnorm, params)
end

"""
Makes a state_action stack along the channels dimension, as input for the dynamics net
"""
function make_dynamics_input(states::Array{Float32,4},actions::Vector{Float32},conf::Config)::Array{Float32,4}
	actions ./= length(conf.action_space)
	state_actions=Array{Float32,4}(undef, conf.observation_shape[1], conf.observation_shape[2], conf.observation_shape[3]+1 ,0)
	for i in 1:length(actions)
		action = actions[i] * ones(Float32,(conf.observation_shape[1], conf.observation_shape[2]))
		# Scale the gradient by half at the start of the dynamics function (See paper appendix Training)
		state = states[:,:,:,i] * 2.0f0
		state_action = cat(state, action, dims=3)
		state_actions = cat(state_actions, state_action, dims=4)
	end
	return state_actions
end

function learning!(num_played_games,
training_step,
remote_NNs,
remote_buffer::RemoteChannel{BufferChannel})::Bool

    while fetch(num_played_games) < 1 # TODO, GPU should always have data available to train on.
        # @info "Waiting for replay buffer to be filled"
		sleep(0.1)
    end

	# @info "Training Started"

	optimiser = Flux.ADAMW()
  	schedule = Stateful(Cos(λ0=1e-4, λ1=1e-1, period=10))

	NNs = fetch(remote_NNs)
	representation= deepcopy(NNs.representation)
	prediction= deepcopy(NNs.prediction)
	dynamics= deepcopy(NNs.dynamics)
	training_step_ = 0

	while training_step_ ≤ conf.training_steps

		buffer = fetch(remote_buffer)

        next_batch = get_batch(buffer)
		index_batch, batch = next_batch
    	observation_batch, action_batch, target_values, target_rewards, target_policies, weight_batch, gradient_scale_batch = batch
		gradient_scale_batch = permutedims(gradient_scale_batch)
		conf.PER ? weight_batch = permutedims(weight_batch) : nothing
		# @info "Batch Sizes are:" size(observation_batch) size(action_batch) size(target_values) size(target_rewards) size(target_policies) size(gradient_scale_batch)
		# observation_batch: Width, height, channels, batch_size
        # action_batch: num_unroll_steps + 1, batch_size
        # target_values: num_unroll_steps+1, batch_size
        # target_rewards: num_unroll_steps+1, batch_size
        # target_policies: len(action_space), num_unroll_steps+1, batch_size
        # weight_batch, gradient_scale_batch: 1, batch_size

    	priorities = zeros(eltype(target_values), size(target_values))

		## Generate predictions, first for the observation then for num_unroll_steps*hidden_states
		hidden_state = representation(observation_batch)
		if ndims(hidden_state)==2
			hidden_state=reshape(hidden_state, (conf.observation_shape...,conf.batch_size))
		end
		predicted_values, predicted_policies = prediction(hidden_state)
		predicted_rewards = zeros((1,conf.batch_size))
		predicted_policies=Flux.unsqueeze(predicted_policies, 2)

		for i = 1:conf.num_unroll_steps
        	value, policy_logits = prediction(hidden_state)
			if ndims(hidden_state)==2
				hidden_state=reshape(hidden_state, (conf.observation_shape...,conf.batch_size))
			end
			policy_logits=Flux.unsqueeze(policy_logits, 2)
			state_action=make_dynamics_input(hidden_state, action_batch[i,:],conf)
			hidden_state, reward = dynamics(state_action)
			if ndims(hidden_state)==2
				hidden_state=reshape(hidden_state, (conf.observation_shape...,conf.batch_size))
			end
			# @info "Size of Predictions" size(value) size(reward) size(policy_logits) size(predicted_values) size(predicted_rewards) size(predicted_policies)
			predicted_values= vcat(predicted_values, value)
			predicted_rewards= vcat(predicted_rewards, reward)
			predicted_policies= cat(predicted_policies, policy_logits, dims=2)
		end

        # predictions & targets: if_any_other_dim, num_unroll_steps + 1, batch
		targets = (target_values, target_rewards, target_policies)
		predictions = (predicted_values, predicted_rewards, predicted_policies)

		# @info "Predictions and Targets generated successfully" size(target_values) size(target_rewards) size(target_policies) size(predicted_values) size(predicted_rewards) size(predicted_policies)

		params_representation = Flux.params(representation)
		params_prediction = Flux.params(prediction)
		params_dynamics = Flux.params(dynamics)

		optimiser[1].eta = next!(schedule) # changes with every minibatch

		# calculates losses over all the samples in this batch at once
		l_representation, grads_representation = loss_grad(params_representation) do
			loss(params_representation, predictions, targets, weight_batch, gradient_scale_batch)
		end
		l_prediction, grads_prediction = loss_grad(params_prediction) do
			loss(params_prediction, predictions, targets, weight_batch, gradient_scale_batch)
		end
		l_dynamics, grads_dynamics = loss_grad(params_dynamics) do
			loss(params_dynamics, predictions, targets, weight_batch, gradient_scale_batch)
		end

		Flux.update!(optimiser, params_representation, grads_representation)
		Flux.update!(optimiser, params_prediction, grads_prediction)
		Flux.update!(optimiser, params_dynamics, grads_dynamics)

		if conf.PER
			priorities = (abs.(predicted_values - target_values)).^conf.PER_alpha
            # Save new priorities in the replay buffer (See https://arxiv.org/abs/1803.00933)
            update_priorities!(buffer, priorities, index_batch)
        end

		training_step_ += 1

		# update_remote_counter!(training_step, 1)
		# take!(training_step)
		println(training_step_)
		put!(training_step, training_step_)
		println(training_step_)

		# Save to the shared storage(disk) #TODO make them availbal on memory
        if training_step_ % conf.checkpoint_interval == 0 && training_step_ > 1

			# old_net=take!(remote_NNs)
			put!(remote_NNs, (representation=representation, prediction=prediction, dynamics=dynamics))

			@info "Training Progress" training_step_
			@info "Representation loss =" l_representation
			@info "Prediction loss =" l_prediction
			@info "Dynamics loss=" l_dynamics

			if training_step_ > round(Int, 0.9*conf.training_steps)
				#saving networks to disk
				representation= cpu(representation)
				prediction= cpu(prediction)
				dynamics= cpu(dynamics)
				serialize(joinpath(conf.networks_path,"$(training_step_)_representation.bin"), representation)
				serialize(joinpath(conf.networks_path,"$(training_step_)_prediction.bin"), prediction)
				serialize(joinpath(conf.networks_path,"$(training_step_)_dynamics.bin"), dynamics)
			end
        end
	end
	return true
end


text:
# outputexample.jl
module outputexample
using Termbox
using Compat
@static if VERSION >= v"0.7-"
    using Printf
end

const chars = "nnnnnnnnnbbbbbbbbbuuuuuuuuuBBBBBBBBB"
const all_attrs = [0, TB_BOLD, TB_UNDERLINE, TB_BOLD | TB_UNDERLINE,]
function next_char(current::Int)
  current += 1
  if current >= length(chars) || current < 1
    current = 1
  end
  return current
end

function draw_line(x, y, bg)
  current_char = 1
  for a in all_attrs
    for c in TB_COLORS
      fg = a | c
      tb_change_cell(x, y, chars[current_char], fg, bg)
      current_char = next_char(current_char)
      x += 1
    end
  end
end
function print_combinations_table(sx, sy, attr)
  for c in TB_COLORS
    bg = attr | c
    draw_line(sx, sy, bg)
    sy += 1
  end
end
function draw_all()
  tb_clear()
  tb_select_output_mode(TB_OUTPUT_NORMAL)

  print_combinations_table(1, 1, 0x0)
  print_combinations_table(2 + length(chars), 1, TB_BOLD)

  print_combinations_table(1, 2 + length(TB_COLORS), TB_REVERSE)

  tb_present()

  tb_select_output_mode(TB_OUTPUT_GRAYSCALE)
  y = 3 + 2 * length(TB_COLORS)
  for x in 0:23
    tb_change_cell(x, y, '@', x, 0);
    tb_change_cell(x + 25, y, ' ', 0, x);
  end
  tb_present()

  tb_select_output_mode(TB_OUTPUT_216)
  y += 1
  x = 1
  for c in 0:215
      tb_change_cell(x, y, '@', c, 0);
      tb_change_cell(x + 26, y, ' ', 0, c);
    x += 1
    if mod(x, 24) == 0
      x = 0
      y += 1
    end
  end
  tb_present()

  tb_select_output_mode(TB_OUTPUT_256)
  y += 1
  x = 1
  for c in 0:255
      (mod(y, 2) == 0) && (c = c | TB_UNDERLINE)
      tb_change_cell(x, y, '+', c, 0)
      tb_change_cell(x + 26, y, ' ', 0, c)
    x += 1
    if mod(x, 24) == 0
      x = 0
      y += 1
    end
  end
  tb_present()
end
function main(sleepTime = 1)
    with_term() do
        draw_all()
        sleep(sleepTime)
        # ev = tb_event()
        # while tb_poll_event(ev) != 0
        #   (ev.etype == TB_EVENT_KEY && ev.key == TB_KEY_ESC) && break
        #   if ev.etype == TB_EVENT_RESIZE
        #     draw_all()
        #   end #if
        # end #while
        tb_present()
        return 0
    end
end # main

function demo_takeover(sleepTime = 1)
  tb_init()
  if ret != 0
    @printf stderr "tb_init() failed with error code %d\n" ret
    return 1
  end
  println("Taking over terminal...")
  println("Height: ", tb_height())
  println("Width: ", tb_width())
  sleep(sleepTime)
  tb_shutdown()
end
end # module


text:
module HyperbolicDrawSimpleGraphs

using SimpleGraphs, HyperbolicPlane, SimpleDrawing

import Base: show
import HyperbolicPlane: polar

export HyperbolicGraphEmbedding, hdraw, hembed

mutable struct HyperbolicGraphEmbedding
    G::SimpleGraph
    locs::Dict{Any,HPoint}
end

function show(io::IO, X::HyperbolicGraphEmbedding)
    print(io, "Hyperbolic embedding of $(X.G)")
end

polar(x::Real,y::Real) = (sqrt(x*x+y*y), atan(x,y) )

# convert SimpleGraph to Hyperbolic embedding
function h_convert(GG::SimpleGraph)::HyperbolicGraphEmbedding
    loc = getxy(GG)
    d = Dict{Any,HPoint}()

    for v in GG.V
        xy = loc[v]
        (r,theta) = polar(xy[1],xy[2])
        P = HPoint(r,theta)
        d[v] = P
    end
    return HyperbolicGraphEmbedding(GG,d)
end

function h_circular(GG::SimpleGraph)::HyperbolicGraphEmbedding
    n = NV(GG)
    VV = vlist(GG)
    r = sqrt(n)
    angs = [2*pi*t/n for t=0:n-1]
    d = Dict{Any,HPoint}()
    for j=1:n
        v = VV[j]
        P = HPoint(r,angs[j])
        d[v] = P
    end
    return HyperbolicGraphEmbedding(GG,d)
end

function h_random(G::SimpleGraph)::HyperbolicGraphEmbedding
    d = Dict{Any,HPoint}()
    for v in G.V
        d[v] = RandomHPoint()
    end
    return HyperbolicGraphEmbedding(G,d)
end

function private_adj(G::SimpleGraph)
    n = NV(G)
    A = zeros(Int,n,n)
    vv = vlist(G)
    for i=1:n-1
        u = vv[i]
        for j=(i+1):n
            w = vv[j]
            if has(G,u,w)
                A[i,j] = 1
                A[j,i] = 1
            end
        end
    end
    return A,vv
end

include("myspring.jl")

"""
`spring!(X)` gives the graph held in `X` with a spring embedding
(based on code in the `GraphLayout` module). If runs a default number of
iterations (100) of that algorithm. To change the number of
iterations, use `spring!(X,nits)`.
"""
function h_spring(G::SimpleGraph, nits::Int=100)::HyperbolicGraphEmbedding
    GG = deepcopy(G)
    embed(GG,:spring)
    X = h_convert(GG)
    n = NV(GG)
    locs_r = zeros(n)
    angs = zeros(n)
    i = 1
    for v in G.V
        xy = X.locs[v]
        (r,theta) = polar(xy)
        locs_r[i] = r
        angs[i] = theta
        i = i + 1
    end
    n = NV(G)
    A,vv = private_adj(G)

    locs_r, angs = h_layout_spring_adj(locs_r,angs,A,MAXITER=nits)

    d = Dict{Any,HPoint}()
    for i = 1:n
        v = vv[i]
        P = HPoint(locs_r[i],angs[i])
        d[v] = P
    end

    return HyperbolicGraphEmbedding(G,d)
end

"""
`hembed(G,method)` embeds `G` in the hyperbolic plane by placing vertices
around a circle.

Options for `method`:
+ `:circular` places vertices around a circle
+ `:random` places vertices at random
+ `:convert` converts the Euclidean embedding of `G` into a hyperbolic embedding
"""
function hembed(G::SimpleGraph, method::Symbol = :circular)

    if method == :circular
        X = h_circular(G)
    elseif method == :random
        X = h_random(G)
    elseif method == :convert
        X = h_convert(G)
    elseif method == :spring
        X = h_spring(G)
    else
        error("Unknown method $method")
    end
    cache_save(G,:HyperbolicGraphEmbedding,X)
end

"""
`hdraw(G)` draws the graph `G` in its current hyperoblic embedding
(or creates a default, circular embedding if need be).
"""
function hdraw(G::SimpleGraph)
    if !cache_check(G,:HyperbolicGraphEmbedding)
        hembed(G)
    end
    X = cache_recall(G,:HyperbolicGraphEmbedding)
    newdraw()
    for e in G.E
        u,v = e
        A = X.locs[u]
        B = X.locs[v]
        draw(A+B)
    end
    for v in G.V
        A = X.locs[v]
        A = HPoint(A)
        set_radius(A,2)
        draw(A)
    end
    # draw(HPlane()) # let the user decide to show the line at infinity
    finish()
end

end  # end of module


text:
include("imports.jl")

include("chamfer.jl")
include("euclidean.jl")
include("squared_euclidean.jl")
include("utils.jl")

@testset "CUDA" begin
    if CUDA.functional()
        include("cuda/runtests.jl")
    else
        @warn "CUDA unavailable, not testing GPU support"
    end
end


text:
# Profiling script

using MkCell

cellfcc = [-0.5 0 0.5; 0 0.5 0.5; -0.5 0.5 0];

MkCell.cellopt(cellfcc, 7, 2)

using Profile

Profile.clear_malloc_data()

# small benchmark
# MkCell.cellopt(cellfcc, 7, 2)

# large benchmark
MkCell.cellopt(cellfcc, 108, 5)


text:
#
# Constraint interface and implementations
#

import Base: clamp,
    show
using LinearAlgebra:
    dot,
    norm,
    normalize,
    logabsdet,
    logdet
using StatsFuns:
    logit,
    logistic

"""
    VariableConstraint{NC,NF}

Abstract type for transformations on constrained variables represented with
`NC`-dimensional vectors that produce `NF`-dimensional free (unconstrained)
vectors.

The supported interface is the 4 functions `constrain`, `free`,
`constrain_with_pushlogpdf` and `constrain_with_pushlogpdf_grad`.

To implement a new constraint, simply create a new type of `VariableConstraint`
and implement `constrain` and `free`. Various internal functions are used to
provide efficient and accurate defaults; these may be overriden for increased
efficiency when analytical gradients/jacobian determinants are known. The most
common override is `constrain_with_logpdf_correction`.
"""
abstract type VariableConstraint{NC,NF} end

"""
    OneToOneConstraint{N}

Alias for dispatch on one-to-one constraints `f: ℝⁿ → ℝⁿ`.
"""
const OneToOneConstraint{N} = VariableConstraint{N,N}

"""
    UnivariateConstraint

Alias for dispatch on univariate constraints `f: ℝ → ℝ`. Univariate
constraints can take scalar inputs and produce scalar outputs.
"""
const UnivariateConstraint = OneToOneConstraint{1}

###
### Basic interface for constraints
###

"""
    free_dimension(c::VariableConstraint)

Get the number of dimensions (length) of the freed vector.
"""
free_dimension(::VariableConstraint{NC,NF}) where {NC,NF} = NF

"""
    constrain_dimension(c::VariableConstraint)

Get the number of dimensions (length) of the constrained vector.
"""
constrain_dimension(::VariableConstraint{NC}) where {NC} = NC

"""
    clamp(c::VariableConstraint, x)

Return `x`, ensuring that its value satisfies the constraint. This is useful
to avoid numerical instability when a value nears the boundary condition.
The clamp is invisible during differentiation.
"""
function clamp end

# Passthrough adjoints for clamp
Zygote.@adjoint clamp(c::VariableConstraint, x) = clamp(c, x), Δ -> (nothing, Δ)

"""
    free(c::VariableConstraint, x)

From constrained variable `x`, construct free variable.
"""
function free end

Base.@propagate_inbounds function free(c::UnivariateConstraint,
                                       x::AbstractVector)
    return [free(c, x[1])]
end

"""
    constrain(c::VariableConstraint, y)

From free variable `y`, construct constrained variable.
"""
function constrain end

Base.@propagate_inbounds function constrain(c::UnivariateConstraint,
                                            y::AbstractVector)
    return [constrain(c, y[1])]
end

"""
    constrain_with_logpdf_correction(c::VariableConstraint, y)

From free variable `y = f(x)` get the constrained variable `x` and the
addtive correction to the log density `log π(x)` to get `log π(x)`. See
[`free_logpdf_correction`](@free_logpdf_correction).
"""
function constrain_with_logpdf_correction(c, y)
    x = constrain(c, y)
    logdetJ = free_logpdf_correction(c, y)
    return x, logdetJ
end

Base.@propagate_inbounds function constrain_with_logpdf_correction(
        c::UnivariateConstraint,
        y::AbstractVector
    )
    x, logdetJ = constrain_with_logpdf_correction(c, y[1])
    return [x], logdetJ
end

"""
    constrain_with_pushlogpdf(c::VariableConstraint, y)

From free variable `y = f(x)` get the constrained variable `x` and a new
function that pushes forward the log density `log π(x)` with respect to `x` and
to the corresponding log density `log π(y)` with respect to `y`.

```julia
x, pushlogpdf = constrain_with_pushlogpdf(c, y)
logπx = ... # Compute density in constrained space
logπy = pushlogpdf(logπx)
```
"""
function constrain_with_pushlogpdf(c, y)
    x, logdetJ = constrain_with_logpdf_correction(c, y)
    return x, logπx -> logπx + logdetJ
end

Base.@propagate_inbounds function constrain_with_pushlogpdf(
        c::UnivariateConstraint,
        y::AbstractVector
    )
    x, pushlogpdf = constrain_with_pushlogpdf(c, y[1])
    return [x], pushlogpdf
end

"""
    constrain_with_pushlogpdf_grad(c::VariableConstraint, y)

From free variable `y = f(x)` get the constrained variable `x` and a new
function that pushes forward the log density `log π(x)` with respect to `x` and
its gradient to the corresponding log density `log π(y)` with respect to `y` and
its gradient.

```julia
x, pushlogpdf_grad = constrain_with_pushlogpdf_grad(c, y)
logπx, ∇x_logπx = ... # Compute density and its gradient in constrained space
logπy, ∇y_logπy = pushlogpdf_grad(logπx, ∇x_logπx)
```
"""
function constrain_with_pushlogpdf_grad(c, y)
    (x, logdetJ), back = Zygote.pullback(constrain_with_logpdf_correction, c, y)
    nf = free_dimension(c)

    return x, function (logπx, ∇x_logπx)
        s = Zygote.sensitivity(logdetJ)
        logπy = logπx + logdetJ
        T = eltype(logπy)
        ∇y_logπy = similar(∇x_logπx, T, nf)
        copyto!(∇y_logπy, back((∇x_logπx, s))[2])
        return logπy, ∇y_logπy
    end
end

function constrain_with_pushlogpdf_grad(c::UnivariateConstraint, y::Real)
    (x, logdetJ), back = Zygote.pullback(constrain_with_logpdf_correction, c, y)

    return x, function (logπx, ∇x_logπx::Real)
        s = Zygote.sensitivity(logdetJ)
        logπy = logπx + logdetJ
        ∇y_logπy = back((∇x_logπx, s))[2]
        return logπy, ∇y_logπy
    end
end

Base.@propagate_inbounds function constrain_with_pushlogpdf_grad(
        c::UnivariateConstraint,
        y::AbstractVector
    )
    x, pushlogpdf_grad = constrain_with_pushlogpdf_grad(c, y[1])
    return [x], function (logπx, ∇x_logπx::AbstractVector)
        logπy, ∇y_logπy = pushlogpdf_grad(logπx, ∇x_logπx[1])
        return logπy, [∇y_logπy]
    end
end

"""
    constrain_jacobian(c::VariableConstraint, y)

From free vector `y = f(x)`, compute the Jacobian matrix of the inverse
transformation `x = f⁻¹(y)` with entries `Jᵢⱼ = ∂xᵢ/∂yⱼ`.
"""
function constrain_jacobian(c, y)
    nf = free_dimension(c)
    # NOTE: work-around to make forward_jacobian type-inferrable
    # see https://github.com/FluxML/Zygote.jl/issues/299
    v = Val(min(nf, ForwardDiff.DEFAULT_CHUNK_THRESHOLD))
    # NOTE: Zygote's (reverse-mode) Jacobians are adjoints
    J′ = last(Zygote.forward_jacobian(y -> constrain(c, y), y, v))
    return adjoint(J′)
end

# NOTE: Workaround until Zygote supports nesting Jacobians
# see https://github.com/FluxML/Zygote.jl/issues/305
Zygote.@adjoint function constrain_jacobian(c,
                                            y::AbstractVector)
    nf = free_dimension(c)
    nc = constrain_dimension(c)

    jac(y) = ForwardDiff.jacobian(y->constrain(c, y), y)
    J = similar(y, (nc, nf))
    diffres = DiffResults.JacobianResult(J, y)
    diffres = ForwardDiff.jacobian!(diffres, jac, y)
    J = DiffResults.value(diffres)
    ∇y_J = reshape(DiffResults.jacobian(diffres), (nc, nf, nf))

    return J, function (J̄)
        @einsum ȳ[k] := J̄[i,j] * ∇y_J[i,j,k]
        return (nothing, ȳ)
    end
end

"""
    constrain_jacobian(c::UnivariateConstraint, y)

From free scalar `y = f(x)`, compute the derivative of the inverse
transformation `x = f⁻¹(y)`, `dx/dy`.
"""
function constrain_jacobian(c::UnivariateConstraint, y)
    dx_dy = Zygote.gradient(constrain, c, y)[2]
    return dx_dy
end

Base.@propagate_inbounds function constrain_jacobian(
        c::UnivariateConstraint,
        y::AbstractArray
    )
    return [constrain_jacobian(c, y[1])]
end

"""
    halflogdetmul(x::AbstractMatrix)

For a matrix `x`, compute `½log(det (x' x))`.
"""
halflogdetmul(x) = logdet(x' * x) / 2

# custom adjoint for slight speed-up
Zygote.@adjoint function halflogdetmul(x::AbstractArray)
    s = x' * x
    return logdet(s) / 2, Δ -> (Δ * (x * inv(s)),) # `Δ * x⁺ᵀ`
end

"""
    free_logpdf_correction(c::VariableConstraint, y)

From free vector `y`, compute correction to log pdf for transformation. Given
a transformation `f: x ↦ y`, its inverse `f⁻¹: y ↦ x`, and pdf `π(x)`, the log
pdf of the transformed density is

`log π(y) = log π(f⁻¹(y)) + ½log(det G)`,

where `det G` is the determinant of the matrix `G = Jᵀ J`, and `J` is the
Jacobian matrix of the inverse transformation with entries `Jᵢⱼ = ∂xᵢ/∂yⱼ`.
This result is known as the area formula.

This function returns `½log(det G)` for the general case of `f: ℝᵐ → ℝⁿ`.
"""
function free_logpdf_correction(c, y)
    @assert free_dimension(c) <= constrain_dimension(c)
    J = constrain_jacobian(c, y)
    return halflogdetmul(J)
end

_logabsdet(x) = first(logabsdet(x))
_logabsdet(x::Real) = log(abs(x))

"""
    free_logpdf_correction(c::OneToOneConstraint, y)

From free variable `y`, compute correction to log pdf for transformation. For
a one-to-one transformation with a square Jacobian, the correction simplifies
to `log |det J|`.
"""
function free_logpdf_correction(c::OneToOneConstraint, y::AbstractVector)
    J = constrain_jacobian(c, y)
    return first(logabsdet(J))
end

###
### Constraint implementations
###

"""
    IdentityConstraint{N} <: OneToOneConstraint{N}

Do-nothing constraint on `ℝⁿ`, corresponding to the identity function on
`n`-dimensional variables. Included for convenient bundling of constrained
with unconstrained variables.

# Constructor

    IdentityConstraint(n::Int)
"""
struct IdentityConstraint{N} <: OneToOneConstraint{N} end

IdentityConstraint(n) = IdentityConstraint{n}()

function Base.show(io::IO, mime::MIME"text/plain",
                    c::IdentityConstraint{N}) where {N}
    print(io, "IdentityConstraint($N)")
end

clamp(::IdentityConstraint, x) = x
clamp(::IdentityConstraint, x::ForwardDiff.Dual) = x

free(::IdentityConstraint{1}, x::Real) = x
free(::IdentityConstraint, x::AbstractVector) = x

constrain(::IdentityConstraint{1}, y::Real) = y
constrain(::IdentityConstraint, y::AbstractVector) = y

free_logpdf_correction(::IdentityConstraint{1}, y::Real) = zero(eltype(y))
free_logpdf_correction(::IdentityConstraint, y::AbstractVector) = zero(eltype(y))

constrain_with_pushlogpdf(::IdentityConstraint{1}, y::Real) = y, identity
constrain_with_pushlogpdf(::IdentityConstraint, y::AbstractVector) = y, identity

function constrain_with_pushlogpdf_grad(::IdentityConstraint{1}, y::Real)
    return y, (logπx, ∇x_logπx) -> (logπx, ∇x_logπx)
end

function constrain_with_pushlogpdf_grad(::IdentityConstraint, y::AbstractVector)
    return y, (logπx, ∇x_logπx) -> (logπx, ∇x_logπx)
end

"""
    LowerBoundedConstraint{T} <: UnivariateConstraint

Constraint on a scalar that is strictly greater than a lower bound.

# Constructor

    LowerBoundedConstraint(lb)
"""
struct LowerBoundedConstraint{T} <: UnivariateConstraint
    lb::T
end

function Base.show(io::IO, mime::MIME"text/plain", c::LowerBoundedConstraint)
    print(io, "LowerBoundedConstraint($(c.lb))")
end

clamp(c::LowerBoundedConstraint, x::Real) = max(x, c.lb + eps(x))
clamp(::LowerBoundedConstraint, x::ForwardDiff.Dual) = x

free(c::LowerBoundedConstraint, x::Real) = log(clamp(c, x) - c.lb)

constrain(c::LowerBoundedConstraint, y::Real) = clamp(c, exp(y) + c.lb)

free_logpdf_correction(c::LowerBoundedConstraint, y::Real) = y

"""
    UpperBoundedConstraint{T} <: UnivariateConstraint

Constraint on a scalar that is strictly less than an upper bound.

# Constructor

    UpperBoundedConstraint(ub)
"""
struct UpperBoundedConstraint{T} <: UnivariateConstraint
    ub::T
end

function Base.show(io::IO, mime::MIME"text/plain", c::UpperBoundedConstraint)
    print(io, "UpperBoundedConstraint($(c.ub))")
end

clamp(c::UpperBoundedConstraint, x::Real) = min(x, c.ub - eps(x))
clamp(::UpperBoundedConstraint, x::ForwardDiff.Dual) = x

free(c::UpperBoundedConstraint, x::Real) = log(c.ub - clamp(c, x))

constrain(c::UpperBoundedConstraint, y::Real) = clamp(c, c.ub - exp(y))

free_logpdf_correction(c::UpperBoundedConstraint, y::Real) = y

"""
    BoundedConstraint{TL,TU,TD} <: UnivariateConstraint

Constraint on a scalar that is has both an upper and lower bound.

# Constructor

    BoundedConstraint(lb, ub)
"""
struct BoundedConstraint{TL,TU,TD} <: UnivariateConstraint
    lb::TL
    ub::TU
    delta::TD
end

function Base.show(io::IO, mime::MIME"text/plain", c::BoundedConstraint)
    print(io, "BoundedConstraint($(c.lb), $(c.ub))")
end

function BoundedConstraint(lb::Real, ub::Real)
    @assert lb < ub
    return BoundedConstraint(lb, ub, ub - lb)
end

clamp(c::BoundedConstraint, x::Real) = clamp(x, c.lb + eps(x), c.ub - eps(x))
clamp(::BoundedConstraint, x::ForwardDiff.Dual) = x

free(c::BoundedConstraint, x::Real) = logit((clamp(c, x) - c.lb) / c.delta)

constrain(c::BoundedConstraint, y::Real) = clamp(c, c.delta * logistic(y) + c.lb)

function constrain_with_logpdf_correction(c::BoundedConstraint, y::Real)
    z = logistic(y)
    delz = c.delta * z
    x = delz + c.lb
    dx_dy = delz * (1 - z)
    return clamp(c, x), log(dx_dy)
end

"""
    TransformConstraint(lb::Real, ub::Real)

Convenient constructor for lower-, upper-, lower- and upper-, and un-bounded
univariate constraints. The correct type is chosen based on the arguments.
"""
function TransformConstraint(lb=-Inf, ub=Inf)
    @assert lb < ub
    has_lb, has_ub = isfinite(lb), isfinite(ub)
    if has_lb && has_ub
        return BoundedConstraint(lb, ub)
    elseif has_lb
        return LowerBoundedConstraint(lb)
    elseif has_ub
        return UpperBoundedConstraint(ub)
    else
        return IdentityConstraint(1)
    end
end

"""
    UnitVectorConstraint{N} <: OneToOneConstraint{N}

Transformation from an `n`-dimensional unit-vector to an unconstrained
`n`-dimensional vector. Note that in this case the inverse transformation
(ℓ²-normalization) is not unique, and therefore the pushforward density cannot
be obtained using the usual Jacobian technique.

However, using the fact that a standard multivariate normally distributed
vector when normalized is uniformly distributed on a sphere, we can push
forward the uniform measure on the sphere by applying a standard multivariate
normal prior to `y`. The corresponding log density correction is `-½ yᵀ y = -½
|y|²`.

The Jacobian of the inverse transformation is a normalized projection matrix
onto the tangent space to the sphere at `x`: `J = Πₓ / |y|`, where
`Πₓ = I - xᵀ x`.

# Constructor

    UnitVectorConstraint(n::Int)
"""
struct UnitVectorConstraint{N} <: OneToOneConstraint{N} end

UnitVectorConstraint(n) = UnitVectorConstraint{n}()

function Base.show(io::IO, mime::MIME"text/plain",
                   c::UnitVectorConstraint{N}) where {N}
    print(io, "UnitVectorConstraint($N)")
end

clamp(::UnitVectorConstraint, x) = normalize(x)
clamp(::UnitVectorConstraint, x::AbstractArray{<:ForwardDiff.Dual}) = x

free(c::UnitVectorConstraint, x) = clamp(c, x)

function constrain(::UnitVectorConstraint, y)
    x = y ./ norm(y)
    return x
end

function normalize_with_norm(y)
    ny = norm(y)
    return y ./ ny, ny
end

Zygote.@adjoint function normalize_with_norm(y)
    x, ny = normalize_with_norm(y)
    return (x, ny), function (Δ)
        x̄, n̄ȳ = Δ
        return (x̄ ./ ny .- x .* (dot(x, x̄) / ny - n̄ȳ),)
    end
end

function constrain_with_logpdf_correction(::UnitVectorConstraint, y)
    x, ny = normalize_with_norm(y)
    logdetJ = -ny^2 / 2
    return x, logdetJ
end

"""
    UnitVectorScaledConstraint{N} <: OneToOneConstraint{N}

Transformation from an `n`-dimensional unit-vector to an unconstrained
`n`-dimensional vector, scaled by a factor of ``r``.

``r`` is interpreted as a radius. This is intended to bring directional/quaternion
vectors to the same approximate initial scale as other variables.

# Constructor

    UnitVectorScaledConstraint(n::Int, r = 1)
"""
struct UnitVectorScaledConstraint{N,T<:Real} <: OneToOneConstraint{N}
    r::T

    function UnitVectorScaledConstraint{N,T}(r::T) where {N,T}
        isinteger(N) && N > 0 || throw(DomainError("N must be a positive integer."))
        r > 0 || throw(DomainError("r must be strictly positive."))
        return new{N,T}(r)
    end
end

UnitVectorScaledConstraint(n, r = 1.0) = UnitVectorScaledConstraint{n,typeof(r)}(r)

function Base.show(io::IO, mime::MIME"text/plain",
                   c::UnitVectorScaledConstraint{N}) where {N}
    print(io, "UnitVectorScaledConstraint($N, c.r)")
end

clamp(::UnitVectorScaledConstraint, x) = normalize(x)
clamp(::UnitVectorScaledConstraint, x::AbstractArray{<:ForwardDiff.Dual}) = x

free(c::UnitVectorScaledConstraint, x) = normalize(x) .* c.r

function constrain(::UnitVectorScaledConstraint, y)
    x = y ./ norm(y)
    return x
end

function constrain_with_logpdf_correction(c::UnitVectorScaledConstraint, y)
    x, ny = normalize_with_norm(y)
    logdetJ = -(ny / c.r)^2 / 2
    return x, logdetJ
end

"""
    UnitSimplexConstraint{N,M} <: VariableConstraint{N,M}

Transformation from an `n`-dimensional vector of positive reals with a unit
ℓ¹-norm to an unconstrained `m = n - 1`-dimensional vector.

This constraint uses the stick-breaking process to define the transformation.
See https://en.wikipedia.org/wiki/Dirichlet_process#The_stick-breaking_process
for more details.

# Constructor

    UnitSimplexConstraint(n::Int)
"""
struct UnitSimplexConstraint{N,M} <: VariableConstraint{N,M} end

UnitSimplexConstraint(n::Int) = UnitSimplexConstraint{n,n-1}()

function Base.show(io::IO, mime::MIME"text/plain",
                   c::UnitSimplexConstraint{N}) where {N}
    print(io, "UnitSimplexConstraint($N)")
end

"""
    stick_ratio(x, Σx)

Break a piece of length `x` off of a unit-length stick off of which pieces of
total length `Σx` have already been broken. Return the ratio of the length `x`
to the remaining length of the stick.
"""
stick_ratio(x, Σx) = x / (1 - Σx)

"""
    stick_length(r, Σx)

Inverse of `stick_ratio`.
"""
stick_length(r, Σx) = r * (1 - Σx)

"""
    constrain_stick_ratio(k, y, K)

Constrain the free variable `y` to the `k`th of `K` stick ratios.
"""
constrain_stick_ratio(k, y, K) = logistic(y - log(K - k))

"""
    free_stick_ratio(k, r, K)

Free the `k`th stick ratio `r` of `K` stick ratios to an unconstrained variable.
"""
free_stick_ratio(k, r, K) = logit(r) + log(K - k)

function clamp(::UnitSimplexConstraint, x)
    ϵ = eps(eltype(x))
    return normalize(clamp.(x, ϵ, 1 - ϵ), 1)
end

clamp(::UnitSimplexConstraint, x::AbstractArray{<:ForwardDiff.Dual}) = x

function free(c::UnitSimplexConstraint, x)
    K = constrain_dimension(c)
    @assert length(x) == K

    x = clamp(c, x)
    T = eltype(x)
    y = similar(x, K - 1)
    Σx = zero(T)
    @inbounds begin
        y[1] = free_stick_ratio(1, x[1], K)
        @simd for k = 2:(K - 1)
            Σx += x[k - 1]
            zₖ = stick_ratio(x[k], Σx)
            y[k] = free_stick_ratio(k, zₖ, K)
        end
    end
    return y
end

function constrain(c::UnitSimplexConstraint, y)
    K = constrain_dimension(c)
    @assert length(y) == free_dimension(c)

    x = Zygote.Buffer(y, K)
    Σx = zero(eltype(y))
    @inbounds begin
        for k = 1:(K - 1)
            zₖ = constrain_stick_ratio(k, y[k], K)
            xₖ = stick_length(zₖ, Σx)
            x[k] = xₖ
            Σx += xₖ
        end
        x[K] = 1 - Σx
    end
    return clamp(c, copy(x))
end

function constrain_with_logpdf_correction(c::UnitSimplexConstraint, y)
    K = constrain_dimension(c)
    @assert length(y) == free_dimension(c)
    T = eltype(y)

    x = Zygote.Buffer(y, K)
    Σx = zero(eltype(y))
    logdetJ = log(K) / 2
    @inbounds begin
        for k = 1:(K - 1)
            zₖ = constrain_stick_ratio(k, y[k], K)
            xₖ = stick_length(zₖ, Σx)
            x[k] = xₖ
            logdetJ += log(xₖ * (1 - zₖ))
            Σx += xₖ
        end
        x[K] = 1 - Σx
    end

    return clamp(c, copy(x)), logdetJ
end

function constrain_with_pushlogpdf_grad(
        c::UnitSimplexConstraint,
        y::SubArray
    )
    return constrain_with_pushlogpdf_grad(c, collect(y))
end

"""
    JointConstraint{TC,RC,RF,NC,NF} <: VariableConstraint{NC,NF}

Joint transformation on a series of constraints. This constraint type
conveniently binds together a series of transformations on non-overlapping
subarrays of a longer parameter array.

# Constructor

    JointConstraint(constraints::VariableConstraint...)
"""
struct JointConstraint{TC,CR,CF,NC,NF} <: VariableConstraint{NC,NF}
    constraints::TC
    cranges::CR
    franges::CF

    function JointConstraint(constraints...)
        cs = merge_constraints(constraints...)

        ncs = map(constrain_dimension, cs)
        nc = sum(ncs)
        cranges = _ranges_from_lengths(ncs)

        nfs = map(free_dimension, cs)
        nf = sum(nfs)
        franges = _ranges_from_lengths(nfs)

        return new{typeof(cs),typeof(cranges),typeof(franges),nc,nf}(cs, cranges, franges)
    end
end

JointConstraint(c) = c

"""
    merge_constraints(cs::VariableConstraint...)

Merge adjacent constraints that can be merged. For example,
`IdentityConstraint(3)` and `IdentityConstraint(2)` can be merged into a single
`IdentityConstraint(5)`.
"""
merge_constraints(c1, c2, cs...) = reduce(merge_constraints, (c1, c2, cs...))

@inline merge_constraints(c1) = (c1,)

@inline merge_constraints(c1, c2) = (c1, c2)

@inline function merge_constraints(c1, c2::Tuple)
    return (merge_constraints(c1, c2[1])..., Base.tail(c2)...)
end

@inline function merge_constraints(c1::Tuple, c2)
    return (Base.front(c1)..., merge_constraints(c1[end], c2)...)
end

@inline function merge_constraints(c1::Tuple, c2::Tuple)
    return (Base.front(c1)...,
            merge_constraints(c1[end], c2[1])...,
            Base.tail(c2)...)
end

@inline function merge_constraints(::IdentityConstraint{M},
                                   ::IdentityConstraint{N}) where {M,N}
    return (IdentityConstraint(M + N),)
end

function _ranges_from_lengths(lengths)
    ranges = UnitRange{Int}[]
    k = 1
    for len in lengths
        push!(ranges, k:(k + len - 1))
        k += len
    end
    return tuple(ranges...)
end

function Base.show(io::IO, mime::MIME"text/plain", jc::JointConstraint)
    print(io, "JointConstraint(")
    nconstraints = length(jc.constraints)
    if nconstraints > 0
        print(io, "$(repr(mime, jc.constraints[1]))")
        if nconstraints > 10
            print(io, ", $(repr(mime, jc.constraints[2]))")
            print(io, ", ...")
            print(io, ", $(repr(mime, jc.constraints[end-1]))")
            print(io, ", $(repr(mime, jc.constraints[end]))")
        else
            for i in 2:nconstraints
                print(io, ", $(repr(mime, jc.constraints[i]))")
            end
        end
    end
    print(io, ")")
end

function _parallel_free(cs, cranges, franges, nf, x)
    y = similar(x, nf)
    @simd for i in 1:length(cs)
        @inbounds begin
            xᵢ = view(x, cranges[i])
            setindex!(y, free(cs[i], xᵢ), franges[i])
        end
    end
    return y
end

function free(jc::JointConstraint, x)
    @assert length(x) == constrain_dimension(jc)
    return _parallel_free(
        jc.constraints,
        jc.cranges,
        jc.franges,
        free_dimension(jc),
        x
    )
end

function _parallel_constrain(cs, cranges, franges, nc, y)
    x = similar(y, nc)
    @simd for i in 1:length(cs)
        @inbounds begin
            yᵢ = view(y, franges[i])
            setindex!(x, constrain(cs[i], yᵢ), cranges[i])
        end
    end
    return x
end

function constrain(jc::JointConstraint, y)
    @assert length(y) == free_dimension(jc)
    return _parallel_constrain(
        jc.constraints,
        jc.cranges,
        jc.franges,
        constrain_dimension(jc),
        y
    )
end

function _parallel_constrain_with_logpdf_correction(cs, cranges, franges, nc, nf, y)
    x = similar(y, nc)

    TL = eltype(y)
    cinds = 1:length(cs)

    logdetJ = zero(TL)
    @simd for i in cinds
        @inbounds begin
            yᵢ = view(y, franges[i])
            xᵢ, logdetJᵢ = constrain_with_logpdf_correction(cs[i], yᵢ)
            setindex!(x, xᵢ, cranges[i])
        end
        logdetJ += logdetJᵢ
    end

    return x, logdetJ
end

Zygote.@adjoint function _parallel_constrain_with_logpdf_correction(cs, cranges, franges, nc, nf, y)
    x = similar(y, nc)

    TL = eltype(y)
    cinds = 1:length(cs)

    backs = []
    logdetJ = zero(TL)
    @simd for i in cinds
        @inbounds begin
            yᵢ = view(y, franges[i])
            (xᵢ, logdetJᵢ), backᵢ = (
                Zygote.pullback(constrain_with_logpdf_correction, cs[i], yᵢ)
            )
            setindex!(x, xᵢ, cranges[i])
        end
        push!(backs, backᵢ)
        logdetJ += logdetJᵢ
    end

    return (x, logdetJ), function (Δ)
        x̄, logdetJ̄ = Δ
        TA = Base.promote_eltypeof(x̄, logdetJ̄)
        ȳ = similar(x̄, TA, nf)

        @simd for i in cinds
            @inbounds begin
                x̄ᵢ = view(x̄, cranges[i])
                backᵢ = backs[i]
                ȳᵢ = backᵢ((x̄ᵢ, logdetJ̄))[2]
                setindex!(ȳ, ȳᵢ, franges[i])
            end
        end

        return (nothing, nothing, nothing, nothing, nothing, ȳ)
    end
end

function constrain_with_logpdf_correction(jc::JointConstraint, y)
    nf = free_dimension(jc)
    @assert length(y) == nf
    return _parallel_constrain_with_logpdf_correction(
        jc.constraints,
        jc.cranges,
        jc.franges,
        constrain_dimension(jc),
        nf,
        y
    )
end

function _parallel_constrain_with_pushlogpdf(cs, cranges, franges, nc, nf, y)
    x = similar(y, nc)
    pushes = []
    cinds = 1:length(cs)
    logdetJs = []
    @simd for i in cinds
        @inbounds begin
            yᵢ = view(y, franges[i])
            xᵢ, pushᵢ = constrain_with_pushlogpdf(cs[i], yᵢ)::Tuple
            setindex!(x, xᵢ, cranges[i])
        end
        push!(pushes, pushᵢ)
    end
    y1 = @inbounds y[1]

    function pushlogpdf(logπx)
        TL = Base.promote_eltypeof(y1, logπx)

        logdetJ = zero(TL)
        @simd for i in cinds
            @inbounds pushᵢ = pushes[i]
            logdetJᵢ = pushᵢ(zero(TL))
            logdetJ += logdetJᵢ
        end

        logπy::TL = logπx + logdetJ
        return logπy
    end

    return x, pushlogpdf
end

function constrain_with_pushlogpdf(jc::JointConstraint, y)
    @assert length(y) == free_dimension(jc)
    return _parallel_constrain_with_pushlogpdf(
        jc.constraints,
        jc.cranges,
        jc.franges,
        constrain_dimension(jc),
        free_dimension(jc),
        y
    )
end

function _parallel_constrain_with_pushlogpdf_grad(cs, cranges, franges, nc, nf, y)
    x = similar(y, nc)
    pushes = []
    cinds = 1:length(cs)
    @simd for i in cinds
        @inbounds begin
            yᵢ = view(y, franges[i])
            xᵢ, pushᵢ = constrain_with_pushlogpdf_grad(cs[i], yᵢ)::Tuple
            setindex!(x, xᵢ, cranges[i])
        end
        push!(pushes, pushᵢ)
    end
    y1 = @inbounds y[1]

    function pushlogpdf_grad(logπx, ∇x_logπx)
        TL = Base.promote_eltypeof(y1, logπx, ∇x_logπx)
        ∇y_logπy = similar(∇x_logπx, nf)

        logdetJ = zero(TL)
        @simd for i in cinds
            @inbounds begin
                ∇x_logπxᵢ = view(∇x_logπx, cranges[i])
                pushᵢ = pushes[i]
                logdetJᵢ, ∇y_logπyᵢ = pushᵢ(zero(logπx), ∇x_logπxᵢ)::Tuple
                setindex!(∇y_logπy, ∇y_logπyᵢ, franges[i])
            end
            logdetJ += logdetJᵢ
        end

        logπy::TL = logπx + logdetJ
        return logπy, ∇y_logπy
    end

    return x, pushlogpdf_grad
end

function constrain_with_pushlogpdf_grad(jc::JointConstraint, y)
    @assert length(y) == free_dimension(jc)
    return _parallel_constrain_with_pushlogpdf_grad(
        jc.constraints,
        jc.cranges,
        jc.franges,
        constrain_dimension(jc),
        free_dimension(jc),
        y
    )
end


text:
"""
`process_streaks`
"""

function process_streaks(df::DataFrames.AbstractDataFrame; photometry = false)
    dayly_vars_list = [:MouseID, :Gen, :Drug, :Day, :Daily_Session, :Box, :Stim_Day, :Condition, :ExpDay, :Area, :Session];
    booleans=[:Reward,:Stim,:Wall,:Correct,:Stim_Day]#columns to convert to Bool
    for x in booleans
        df[!,x] = eltype(df[!,x]) == Bool ? df[!,x] : occursin.("true",df[!,x],)
    end
    streak_table = by(df, :Streak) do dd
        dt = DataFrame(
        Num_pokes = size(dd,1),
        Num_Rewards = length(findall(dd[!,:Reward].==1)),
        Start_Reward = dd[1,:Reward],
        Last_Reward = findlast(dd[!,:Reward] .== 1).== nothing ? 0 : findlast(dd[!,:Reward] .== 1),
        Prev_Reward = findlast(dd[!,:Reward] .== 1).== nothing ? 0 : findprev(dd[!,:Reward] .==1, findlast(dd[!,:Reward] .==1)-1),
        Trial_duration = (dd[end,:PokeOut]-dd[1,:PokeIn]),
        Start = (dd[1,:PokeIn]),
        Stop = (dd[end,:PokeOut]),
        Pre_Interpoke = size(dd,1) > 1 ? maximum(skipmissing(dd[!,:Pre_Interpoke])) : missing,
        Post_Interpoke = size(dd,1) > 1 ? maximum(skipmissing(dd[!,:Post_Interpoke])) : missing,
        PokeSequence = [SVector{size(dd,1),Bool}(dd[!,:Reward])],
        Stim = dd[1,:Stim],
        StimFreq = dd[1,:StimFreq],
        Wall = dd[1,:Wall],
        Protocol = dd[1,:Protocol],
        Correct_start = dd[1,:Correct],
        Correct_leave = !dd[end,:Correct],
        Block = dd[1,:Block],
        Streak_within_Block = dd[1,:Streak_within_Block],
        Side = dd[1,:Side],
        ReverseStreak = dd[1,:ReverseStreak]
        )
        for s in dayly_vars_list
            if s in names(df)
                dt[!,s] .= df[1, s]
            end
        end
        return dt
    end
    streak_table[!,:Prev_Reward] = [x .== nothing ? 0 : x for x in streak_table[:,:Prev_Reward]]
    streak_table[!,:AfterLast] = streak_table[!,:Num_pokes] .- streak_table[!,:Last_Reward];
    streak_table[!,:BeforeLast] = streak_table[!,:Last_Reward] .- streak_table[!,:Prev_Reward].-1;
    prov = lead(streak_table[!,:Start],default = 0.0) .- streak_table[!,:Stop];
    streak_table[!,:Travel_to] = [x.< 0 ? 0 : x for x in prov]
    if photometry
        frames = by(df, :Streak) do df
            dd = DataFrame(
            In = df[1,:In],
            Out = df[end,:Out],
            LR_In = findlast(df[:Reward])==0 ? NaN : df[findlast(df[:Reward]),:In],
            LR_Out = findlast(df[:Reward])==0 ? NaN : df[findlast(df[:Reward]),:Out]
            )
        end
        streak_table[:In] = frames[:In]
        streak_table[:Out] = frames[:Out]
        streak_table[:LR_In] = frames[:LR_In]
        streak_table[:LR_Out] = frames[:LR_Out]
    end

    return streak_table
end


text:
export findCenterOfDfFov

"""This function calculates the center of the DfFov of the given MPIFile bSF"""
function findCenterOfDfFov(bSF::MPIFile)
	S1 = getSF(bSF,0,2,0,1) #xdir
  S2 = getSF(bSF,2,0,0,2) #ydir
  S3 = getSF(bSF,0,0,2,3) #zdir

  center = ones(3)
  if size(S1,1) > 1
    u = floor.(Int, centerOfMass(abs.(S1)))
    x = argmin(abs.(vec((S1[(u[1]-5):(u[1]+5),u[2],u[3]]))))+u[1]-6
    y_0 = real(S1[x-1,u[2],u[3]])
    y_1 = real(S1[x+1,u[2],u[3]])
    xdiff = -y_0 / (y_1-y_0) *2
    center[1] = x#+xdiff-1
  end
  if size(S2,2) > 1
    u = floor.(Int, centerOfMass(abs.(S2)))
    x = argmin(abs.(vec((S2[u[1],(u[2]-5):(u[2]+5),u[3]]))))+u[2]-6
    y_0 = real(S2[u[1],x-1,u[3]])
    y_1 = real(S2[u[1],x+1,u[3]])
    xdiff = -y_0 / (y_1-y_0) *2
    center[2] = x#+xdiff-1
  end
  if size(S3,3) > 1
    u = floor.(Int, centerOfMass(abs.(S3)))
    x = argmin(abs.(vec((S3[u[1],u[2],(u[3]-5):(u[3]+5)]))))+u[3]-6
    y_0 = real(S3[u[1],u[2],x-1])
    y_1 = real(S3[u[1],u[2],x+1])
    xdiff = -y_0 / (y_1-y_0) *2
    center[3] = x#+xdiff-1
  end
  return center
end

function centerOfMass(data::Array{T,3}) where T
  data = abs.(data)
  maxData = maximum(data)
  data[ data .< maxData*0.5] .= 0

  lx=0
  ly=0
  lz=0
  for i3=1:size(data,3)
    for i2=1:size(data,2)
      for i1=1:size(data,1)
           lx += data[i1,i2,i3]*i1
           ly += data[i1,i2,i3]*i2
           lz += data[i1,i2,i3]*i3
       end
     end
  end
  cm = [lx/sum(data), ly/sum(data),lz/sum(data)]
  for d=1:3
    if size(data,d) == 1
      cm[d] = 1.0
    end
  end
  return cm
end


text:
type Rule
    dict::Dict{Transition, Tuple{Transition,Index}} # rule dictionary (fin=>(fout,rule_index))
    states::Array{Tuple{Partition,Count}}   # states of partition
    N::Count                                # num of transition rules
    Nrot::Count                             # num of transition rules including rotational symmetry
    function Rule()
        dict = Dict{Transition,Tuple{Transition,Index}}()
        states = Tuple{Partition,Count}[]
        new(dict,states,0,0)
    end
end

function save(fname::AbstractString, rule::Rule)
    open(fname, "w") do f
        show(f, rule)
    end
end

function load_rule(fname::AbstractString)
    flines = open(fname) do f
        readlines(f)
    end
    lines = AbstractString[]
    for line in flines
        # eliminate ret code
        line = chomp(line)

        # skip comment/brank line
        if ismatch(r"^\s*(#)",line) || length(line)==0
            continue
        end

        push!(lines, line)
    end

    load_rule(lines)
end

function load_rule(lines::Array)
    rule = Rule()

    for line in lines
        if !ismatch(r"^[0-9a-zA-Z]{8,8}\s[0-9a-zA-Z]{8,8}$", line)
            warn(string("Parse error: ", line))
            continue
        end
        str = split(line[1:end], ' ')
        fin = toval64(str[1])
        fout = toval64(str[2])
        add_rule!(rule, fin, fout)
    end
    gen_partition_states_list!(rule)

    return rule
end

function get_partition_states_list(rule::Rule)
    return map(x->x[1], rule.states)
end

function add_rule!(rule::Rule, fin::Transition, fout::Transition)
   rot_rules = gen_rotational_rule(fin, fout)
   check_rule_validity(rot_rules)
   has_rotational_rule(rule.dict, rot_rules)
   for r in rot_rules
       # push!(rule.dict, r[1], (r[2], rule.N))
       setindex!(rule.dict, (r[2], rule.N), r[1])
       rule.Nrot+=1
   end
   rule.N+=1
end

function gen_rotational_rule(fin::Transition, fout::Transition)
    rot_rules = Tuple{Transition,Transition}[]
    push!(rot_rules, (fin,fout))
    for i=1:3
        fin = rotate(fin)
        fout = rotate(fout)
        push!(rot_rules, (fin,fout))
    end
    unique(rot_rules)
end

function check_rule_validity(rot_rules::Array{Tuple{Transition,Transition}})
    fins = unique(map(x->x[1],rot_rules))
    fouts = unique(map(x->x[2],rot_rules))
    if length(fins)<length(fouts)
        warn(string("Bad rule: ",tostr(fins[1]),"->", tostr(fouts[1])))
    end
end

function has_rotational_rule(dict::Dict, rot_rules::Array{Tuple{Transition,Transition}})
    for rule in rot_rules
        if haskey(dict,rule[1])
            warn(string("Repeated rule: ",tostr(rule[1]),"->",tostr(rule[2])))
        end
    end
end

function gen_partition_states_list!(rule::Rule)
    states = Dict{Partition,Count}()
    for fin in keys(rule.dict)
        fout,ridx = rule.dict[fin]
        v = (UInt128(fin) << NBITS_T)|fout
        for p in reinterpret(Partition, UInt128[v])
            if haskey(states, p)
                states[p] += 1
            else
                setindex!(states, 1, p)
            end
        end
    end
    keylist = keys(states)
    for key in keylist
        push!(rule.states, (key, states[key]))
    end
    sort!(rule.states, by=x->x[2])
end

# 0x11223344_55667788 -> 0x44112233_88556677
@inline function rotate(val::Transition)
    hv = Cell(val >> BSHIFT_CELL)
    lv = Cell(val & BMASK_C)
    hv = rotate(hv)
    lv = rotate(lv)
    (Transition(hv) << BSHIFT_CELL)|Transition(lv)
end

# 0x11223344 -> 0x44112233
@inline function rotate(val::Cell)
    (val >> NBITS_P)|((val & BMASK_P) << NBITS_P*3)
end

function show(io::IO, rule::Rule, rotflag::Bool=false)
    if rotflag
        for fin in keys(rule.dict)
            write(io, string(tostr(fin)," ",tostr(rule.dict[fin][1]),"\n"))
        end
    else
        ruleidx = Transition[]
        for fin in keys(rule.dict)
            if !(rule.dict[fin][2] in ruleidx)
                write(io, string(tostr(fin)," ",tostr(rule.dict[fin][1]),"\n"))
                push!(ruleidx, rule.dict[fin][2])
            end
        end
    end
end

function print(rule::Rule, rotflag::Bool=false)
    show(STDOUT, rule, fotflag)
end

function println(rule::Rule, rotflag::Bool=false)
    show(STDOUT, rule, rotflag)
    println()
end


text:
__precompile__()
module RegTools

export avPlot,vif,rsquared,adjrsquared,rstudent,jackknife,halfnorm,cooksdistance

include("diagnostics.jl")
include("misc.jl")
include("modsel.jl")

end # module

max_stars_repo_name: atantos/RegTools

max_issues_repo_name: atantos/RegTools

max_forks_repo_name: atantos/RegTools


text:
"""
    gevfitpwm(...)

Estimate the GEV parameters with the probability weighted moments.

# Implementation

Estimation with the probability weighted moments, as described by [Hosking *et al*. (1985)](https://www.tandfonline.com/doi/abs/10.1080/00401706.1985.10488049),
is only possible in the stationary case.

See also [`gevfitpwm`](@ref) for the other methods, [`gevfit`](@ref), [`gevfitbayes`](@ref) and [`BlockMaxima`](@ref).

# Reference

Hosking, J. R. M., Wallis, J. R. and Wood, E. F. (1985). Estimation of the generalized extreme-value
    distribution by the method of probability-weighted moments. *Technometrics*, 27:251-261.
"""
function gevfitpwm end

"""
    gevfitpwm(y::Vector{<:Real})

Estimate the GEV parameters with the probability weighted moments.

"""
function gevfitpwm(y::Vector{<:Real})::pwmEVA

    model = BlockMaxima(Variable("y", y))

    fittedmodel = gevfitpwm(model)

    return fittedmodel

end

"""
    gevfitpwm(df::DataFrame, datacol::Symbol)

Estimate the GEV parameters with the probability weighted moments.

Block maxima data are in the column `datacol` of the dataframe `df`.

"""
function gevfitpwm(df::DataFrame, datacol::Symbol)::pwmEVA

    model = BlockMaxima(Variable(string(datacol), df[:, datacol]))

    fittedmodel = gevfitpwm(model)

    return fittedmodel

end

"""
    gevfitpwm(model::BlockMaxima)

Estimate the GEV parameters with the probability weighted moments.

"""
function gevfitpwm(model::BlockMaxima)::pwmEVA

    model = validatestationarity(model)

    y = model.data.value

    # Computing the estimates of the probability weighted moments M_{1,q,0} for q ∈ {0,1,2}.
    b₀ = pwm(y,1,0,0)
    b₁ = pwm(y,1,1,0)
    b₂ = pwm(y,1,2,0)

    # GEV parameters estimations. Expressions retrieved from Hosking et al. (1985).
    c = (2b₁ - b₀)/(3b₂ - b₀) - log(2)/log(3)
    k = 7.859c + 2.9554c^2
    σ̂ = k *( 2b₁-b₀ ) /(1-2^(-k))/SpecialFunctions.gamma(1+k)
    μ̂ = b₀ - σ̂/k*( 1-SpecialFunctions.gamma(1+k) )

    ξ̂ = -k
    ϕ̂ = log(σ̂)

    θ̂ = [μ̂; ϕ̂; ξ̂]

    fm = pwmEVA{BlockMaxima, GeneralizedExtremeValue}(model, θ̂)

    return fm

end