using CSV
using StatsBase

include("common.jl")

_shifted_geomean(vals, shift) = StatsBase.geomean(vals .+ shift) - shift

df = CSV.read(RESULT_PATH)

# Remove configurations that might be in data set, but that we
# do not wish to include in our analysis.
df = filter(df) do row
    for (solver, method) in CONFIGS_TO_ANALYZE
        if row.solver == solver && row.formulation == method
            return true
        end
    end
    return false
end

df.fam_instance = df.family .* "-" .* df.instance

# Figure out the best primal objective value available in dataset
# for each instance.
df[!, :best_primal] .= 0.0
for i in 1:size(df, 1)
    row = df[i, :]
    best_primal_value = minimum(df[(df.family .== row.family) .& (df.instance .== row.instance), :true_primal])
    df.best_primal[i] = best_primal_value
end

# Add a column to record MIP gap, with respect to best primal solution in data set.
function _mip_gap(row)
    return 100 * abs(row.dual - row.best_primal) / abs(row.best_primal)
end
df[!, :gap] .= 0.0
for i in 1:size(df, 1)
    df.gap[i] = _mip_gap(df[i,:])
end

# Split off problem classes based on how easy they were. We will do separate analysis
# on each problem class.
PROBLEM_CLASS_DFS = Dict()
PROBLEM_CLASS_DFS["solved"] = filter(df) do row
    family, instance = row.family, row.instance
    slice = df[(df.family .== row.family) .& (df.instance .== row.instance), :]
    return all(slice.optimal .== "true")
end

PROBLEM_CLASS_DFS["contested"] = filter(df) do row
    family, instance = row.family, row.instance
    slice = df[(df.family .== row.family) .& (df.instance .== row.instance), :]
    return any(slice.optimal .== "true") && !all(slice.optimal .== "true")
end

PROBLEM_CLASS_DFS["unsolved"] = filter(df) do row
    family, instance = row.family, row.instance
    slice = df[(df.family .== row.family) .& (df.instance .== row.instance), :]
    return !any(slice.optimal .== "true")
end

for (class, subset_df) in PROBLEM_CLASS_DFS
    total_instance_count = length(unique(subset_df.fam_instance))
    total_instance_count == 0 && continue
    println("Class: $class ($total_instance_count instances)")
    println("="^(30))

    fails = Dict(pair => 0 for pair in CONFIGS_TO_ANALYZE)
    best_bound = Dict(pair => 0 for pair in CONFIGS_TO_ANALYZE)
    for fam_instance in unique(subset_df.fam_instance)
        slice = subset_df[subset_df.fam_instance .== fam_instance, :]
        if size(slice, 1) == 0 error("Unexpected case, aborting.") end
        best_solve_time = minimum(slice.solve_time)
        best_gap = minimum(slice.gap)
        for (solver, formulation) in CONFIGS_TO_ANALYZE
            sub_slice = slice[(slice.solver .== solver) .& (slice.formulation .== formulation), :]
            solve_times = sub_slice.solve_time
            @assert size(solve_times, 1) == 1
            gaps = sub_slice.gap
            solve_time = minimum(solve_times)
            gap = minimum(gaps)
            best_primal = minimum(sub_slice.best_primal)
            if solve_time >= 0.999 * TIME_LIMIT
                fails[(solver, formulation)] += 1
            end
            # Either attain best bound, or reach Gurobi's optimality cutoff,
            # which is 1e-4 * |primal_bound|
            if gap == best_gap || gap < 1e-4 * abs(best_primal)
                best_bound[(solver, formulation)] += 1
            end
        end
    end

    time_shift = minimum(subset_df.solve_time)
    gap_shift = max(1e-4, minimum(subset_df.gap))

    for (solver, formulation) in CONFIGS_TO_ANALYZE
        slice = subset_df[(subset_df.solver .== solver) .& (subset_df.formulation .== formulation), :]
        println("$solver -- $formulation")
        println("-"^(length(solver) + length(formulation) + 4))
        println("  * solve time: ", _shifted_geomean(slice.solve_time, time_shift), " sec")
        println("  * MIP gap:    ", _shifted_geomean(slice.gap, gap_shift))
        num_fails = fails[(solver, formulation)]
        num_best_bound = best_bound[(solver, formulation)]
        println("  * Fails:      ", num_fails, " / ", total_instance_count)
        println("  * Best bound: ", num_best_bound, " / ", total_instance_count)
    end
end


module UnitTest

export testIV!, testGmm!, testFixedPoint!, printEstim!, histIncome
using TypeDef, Helper, Estim, Gmm, GMMInterface, Optim, UserFileIO
using Printf, DelimitedFiles, Random, Distributions, Plots

function testIV!(
  model::Int64,  # ID for estimation model
  stage::Int64,  # Estimation stage (0 for continuous updating)
  mcid ::Int64,  # ID for Monte Carlo estimation
  iv   ::String, # IV approach
  ver  ::String, # Version of estimation data
  impF ::String, # Type of import files
  simEmp=false;  # Simulation draws from empirical distribution
  nPerson=200,   # Number of simulation draws
  stageGmm=2,    # Parameter for output file
  lowBound=true) # Lower bound for simulation draws for robust simulation

  dat,dim,fxp,lab,mat,set,indx,pred,diags =
    initialize(model,iv,nPerson,ver,impF,simEmp;lowBound)

  # Export for testing δ
  pathName = string("out/",impF,"/m",model,"s",stage,iv,nPerson)
  imp = readdlm(string(pathName,"/gmmStage",stageGmm,".csv"),',')
  θ = convert(Vector{Float64},imp[mcid,2:(end-2)])
  θ₂ = θ[1:dim.θ₂]
  setupMu!(fxp.eμ,pred.eμ₀,pred.adj, θ₂,mat.Xᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)
  fixPoint!(fxp.eδ,fxp.∂δ∂θ, fxp.eμ,dat,dim,mat,indx,pred,diags,θ₂)
  unpackStackedArray!(fxp.eδ1d, fxp.eδ, dim.market, indx.mkt)
  fid = open(string(pathName,"/expDelta.csv"),"w")
  for i=1:dim.obs
    @printf(fid,"%.15f\n",fxp.eδ1d[i])
  end
  close(fid)

  if iv=="opt"
    optimIV!(mat,θ₂,dat,dim,fxp,indx,pred,diags)
  end
  # Export for Sanderson-Windmeijer test
  ∂δ∂θ₂ = zeros(dim.obs,dim.θ₂)
  for t=1:dim.market
    ∂δ∂θ₂[indx.mkt[:,t],:] = fxp.∂δ∂θ[t]
  end
  fid = open(string(pathName,"/swtest_",iv,".csv"),"w")
  @printf(fid,"%8s, %8s, %8s, %24s, %24s",
          "network","region","quarter","delta","price")
  for j=1:dim.θ₂
    @printf(fid,", %24s",string("jacob",j))
  end
  for j=1:size(mat.Xexo,2)
    @printf(fid,", %24s",string("x",j))
  end
  for j=1:(size(mat.Z,2)-size(mat.Xexo,2))
    @printf(fid,", %24s",string("z",j))
  end
  @printf(fid,"\n")
  for i=1:dim.obs
    @printf(fid,"%8d, %8d, %8d, ",dat.brand[i],dat.region[i],dat.qtr[i])
    @printf(fid,"%24.15f, %24.15f",log(fxp.eδ1d[i]),dat.price[i])
    for j=1:dim.θ₂
      @printf(fid,", %24.15f",∂δ∂θ₂[i,j])
    end
    for j=1:size(mat.Z,2)
      @printf(fid,", %24.15f",mat.Z[i,j])
    end
    @printf(fid,"\n")
  end
  close(fid)

  println("Export complete.")
end

function printEstim!(
  model  ::Int64,
  stage  ::Int64,
  mcid   ::Int64,
  iv     ::String,
  ver    ::String,
  impF   ::String,
  update ::Bool=true,
  mcid1st::Int64=0,
  nPerson::Int64=200,
  simEmp ::Bool=false)

  dim,lab,indx = importData(model,iv,nPerson,ver,impF;output=0)
  dat,dim,fxp,lab,mat,set,indx,pred,diags =
    initialize(model,iv,nPerson,ver,impF,simEmp)
  pathName = string("out/",impF,"/m",model,"s",stage,iv,nPerson)
  imp = readdlm(string(pathName,"/gmmStage2.csv"),',')
  θ = convert(Vector{Float64},imp[mcid,2:(end-2)])
  θ₂ = θ[1:dim.θ₂]

  seθ,θ₁,covθ = stdError(θ₂,model,stage,iv,nPerson,impF,update,
                         dat,dim,fxp,mat,set,indx,pred,diags,mcid1st)
  # println(seθ)
  Jstat::Float64 = 0
  θ[1] *= -1;
  if model==0
    θ[2] *= 10; seθ[2] *= 10;
    θ[3] *= -1;
  elseif model==15
    θ[2] *= -1;
    θ[4] *= 10; seθ[4] *= 10;
    θ[5] *= 10; seθ[5] *= 10;
  elseif model==27
    θ[2] *= -1;
  end
  param = paramType(θ₂,θ₁,seθ,Jstat)
  nameParam = [lab.L; lab.θ₁]
  fid = open(string(pathName,"/estimLatex.csv"),"w")
  for i=1:dim.θ
    z = abs(θ[i])/param.seθ[i]
    @printf(fid,"%18s,  %.3f",nameParam[i],θ[i])
    if z > 2.575829
      @printf(fid,"***");
    elseif z > 1.959964
      @printf(fid,"**");
    elseif z > 1.644854
      @printf(fid,"*");
    end
    @printf(fid,"\n%18s, (%.3f)\n"," ",param.seθ[i])
  end
  close(fid)

end

function histIncome(
  fName  ::fileNameType,
  model  ::Int64,
  stage  ::Int64,
  iv     ::String,
  ver    ::String,
  impF   ::String,
  simEmp ::Bool,
  figName::String;
  update ::Bool=true,
  seed   ::Int64=20,
  nPerson::Int64=200,
  nIter  ::Int64=1000,
  mcid   ::Int64=0)

  # update=true;seed=20; nPerson=200; nIter=1000;mcid=0
  # stage=0;iv="opt";impF="base";fName=fileName
  θ̂₂,θ₂,θ₁,seθ,fval,exitflag,θ₂seed,dim,lab,indx,fJLD,fCSV,fINT,fPAR =
    setupGmmOptim(fName,model,stage,iv,ver,impF,update,seed,nPerson,nIter,mcid)

  dat,dim,fxp,lab,mat,set,indx,pred,diags =
    initialize(model,iv,nPerson,ver,impF,simEmp)

  y = mat.y[dat.region.==11,:]*100
  histogram(y[1,:],normed=true,xlabel="Monthly income (€)",
            ylabel="Probability",label="Income draws")
  savefig(string("post/",figName,".pdf"))
  histogram(1 ./ y[1,:],normed=true,xlabel="Inverse income (€)",
            ylabel="Probability",label="Income draws")
  savefig(string("post/",figName,"Inv.pdf"))

end

end


#####
##### Grabbing properties for thermodynamic variables
#####

function thermodynamic_field(model)
    model.thermodynamic_variable isa Energy  && return model.tracers.ρe
    model.thermodynamic_variable isa Entropy && return model.tracers.ρs
end

function intermediate_thermodynamic_field(model)
    model.thermodynamic_variable isa Energy  && return model.time_stepper.intermediate_fields.tracers.ρe
    model.thermodynamic_variable isa Entropy && return model.time_stepper.intermediate_fields.tracers.ρs
end

#####
##### Kernels for diagnosing thermodynamic variables
#####

@kernel function compute_velocities!(velocities, grid, momenta, total_density)
    i, j, k = @index(Global, NTuple)

    @inbounds velocities.u[i, j, k] = momenta.ρu[i, j, k] / ℑxᶠᵃᵃ(i, j, k, grid, total_density)
    @inbounds velocities.v[i, j, k] = momenta.ρv[i, j, k] / ℑyᵃᶠᵃ(i, j, k, grid, total_density)
    @inbounds velocities.w[i, j, k] = momenta.ρw[i, j, k] / ℑzᵃᵃᶠ(i, j, k, grid, total_density)
end

@kernel function compute_p_over_ρ!(p_over_ρ, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
    i, j, k = @index(Global, NTuple)

    @inbounds p_over_ρ[i, j, k] = diagnose_p_over_ρ(i, j, k, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
end

@kernel function compute_temperature!(temperature, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
    i, j, k = @index(Global, NTuple)

    @inbounds temperature[i, j, k] = diagnose_temperature(i, j, k, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
end

function compute_temperature!(model)
    temperature = intermediate_thermodynamic_field(model)

    temperature, total_density, momenta, tracers =
        datatuples(temperature, model.total_density, model.momenta, model.tracers)

    compute_temperature_event =
        launch!(model.architecture, model.grid, :xyz, compute_temperature!,
                temperature, model.grid, model.thermodynamic_variable, model.gases,
                model.gravity, total_density, momenta, tracers,
                dependencies=Event(device(model.architecture)))

    wait(device(model.architecture), compute_temperature_event)

    return temperature
end

#####
##### CFL
#####

function cfl(model, Δt)
    # We will store the velocities in the time stepper's intermediate fields.
    velocities = (
        u = model.time_stepper.intermediate_fields.ρu,
        v = model.time_stepper.intermediate_fields.ρv,
        w = model.time_stepper.intermediate_fields.ρw
    )

    velocities, momenta, total_density =
        datatuples(velocities, model.momenta, model.total_density)

    velocity_event =
        launch!(model.architecture, model.grid, :xyz, compute_velocities!,
                velocities, model.grid, momenta, total_density,
                dependencies=Event(device(model.architecture)))

    wait(device(model.architecture), velocity_event)

    u_max = maximum(velocities.u.parent)
    v_max = maximum(velocities.v.parent)
    w_max = maximum(velocities.w.parent)

    Δx, Δy, Δz = model.grid.Δx, model.grid.Δy, model.grid.Δz
    CFL = Δt / min(Δx/u_max, Δy/v_max, Δz/w_max)

    return CFL
end

#####
##### Acoustic CFL
#####

function acoustic_cfl(model, Δt)
    # We will store p/ρ in the time stepper's intermediate field for the thermodynamic variable.
    p_over_ρ = intermediate_thermodynamic_field(model)

    p_over_ρ, total_density, momenta, tracers =
        datatuples(p_over_ρ, model.total_density, model.momenta, model.tracers)

    compute_p_over_ρ_event =
        launch!(model.architecture, model.grid, :xyz, compute_p_over_ρ!,
                p_over_ρ, model.grid, model.thermodynamic_variable, model.gases,
                model.gravity, total_density, momenta, tracers,
                dependencies=Event(device(model.architecture)))

    wait(device(model.architecture), compute_p_over_ρ_event)

    p_over_ρ_max = maximum(p_over_ρ.parent)

    γ_max = maximum(gas.cₚ / gas.cᵥ for gas in model.gases)
    c_max = √(γ_max * p_over_ρ_max)

    Δx, Δy, Δz = model.grid.Δx, model.grid.Δy, model.grid.Δz
    acoustic_CFL = Δt / min(Δx/c_max, Δy/c_max, Δz/c_max)

    return acoustic_CFL
end


# Generate voicings exercises.

import JSON

# This ID is incremented for every problem generated.  Barring changes
# to the generation, this should generate the same ids on subsequent
# runs.
cur_id = 1

function emit(question, answer)
    global cur_id
    println("INSERT INTO \"probs\" VALUES($cur_id, '$question', '$answer');")
    cur_id += 1
end

# The ii-V-I progressions through the cycle.  The DeGregg puts some
# work into keeping these in range.  However, since we will be mostly
# doing these out of order, the exact order here doesn't matter all
# that much, and the simple Circle of V is fine.  F♯ and G♭ and both
# printed (making for 13 of each voicing), since both are common.
major_iivi = [
    ["Dm7", "G7", "CΔ"],
    ["Am7", "D7", "GΔ"],
    ["Em7", "A7", "DΔ"],
    ["Bm7", "E7", "AΔ"],
    ["F♯m7", "B7", "EΔ"],
    ["C♯m7", "F♯7", "BΔ"],
    ["G♯m7", "C♯7", "F♯Δ"],
    ["A♭m7", "D♭7", "G♭Δ"],
    ["E♭m7", "A♭7", "D♭Δ"],
    ["B♭m7", "E♭7", "A♭Δ"],
    ["Fm7", "B♭7", "E♭Δ"],
    ["Cm7", "F7", "B♭Δ"],
    ["Gm7", "C7", "FΔ"] ]

# Diminished tri-tone sub.  The idea of this pattern is to alternate
# the inversion of the chords so that the pattern only moves slightly
# on the piano.  See p39 in the deGreg book for an example.
min_tritone_sub = [
    ["CΔ", "C♯°7", "Dm7", "D♯°7", "Em7", "E♭7", "Dm7", "D♭7", "CΔ"],
    ["FΔ", "F♯°7", "Gm7", "G♯°7", "Am7", "A♭7", "Gm7", "G♭7", "FΔ"],
    ["B♭Δ", "B°7", "Cm7", "C♯°7", "Dm7", "D♭7", "Cm7", "B7", "B♭Δ"],
    ["E♭Δ", "E°7", "Fm7", "F♯°7", "Gm7", "G♭7", "Fm7", "E7", "E♭Δ"],

    ["A♭Δ", "A°7", "B♭m7", "B°7", "Cm7", "B7", "B♭m7", "A7", "A♭Δ"],
    ["D♭Δ", "D°7", "E♭m7", "E°7", "Fm7", "E7", "E♭m7", "D7", "D♭Δ"],
    ["G♭Δ", "G°7", "A♭m7", "A°7", "B♭m7", "A7", "A♭m7", "G7", "G♭Δ"],
    ["BΔ", "C°7", "C♯m7", "D°7", "D♯m7", "D7", "C♯m7", "C7", "BΔ"],

    ["EΔ", "F°7", "F♯m7", "G°7", "G♯m7", "G7", "F♯m7", "F7", "EΔ"],
    ["AΔ", "A♯°7", "Bm7", "C°7", "C♯m7", "C7", "Bm7", "B♭7", "AΔ"],
    ["DΔ", "D♯°7", "Em7", "F°7", "F♯m7", "F7", "Em7", "E♭7", "DΔ"],
    ["GΔ", "G♯°7", "Am7", "A♯°7", "Bm7", "B♭7", "Am7", "A♭7", "GΔ"] ]

# Extract the root of a chord, the characters before any modifier.
function root_name(chord)
    m = match(r"^[A-G][♭♯]?", chord)
    if m === nothing
        error("Invalid chord name")
    else
        m.match
    end
end

base_notes = Dict("C" => 36, "C♯" => 37,
    "D♭" => 37, "D" => 38, "D♯" => 39,
    "E♭" => 39, "E" => 40, "E♯" => 41,
    "F♭" => 40, "F" => 41, "F♯" => 42,
    "G♭" => 42, "G" => 43, "G♯" => 44,
    "A♭" => 44, "A" => 45, "A♯" => 46,
    "B♭" => 46, "B" => 47, "B♯" => 48)

# A Chord Voicing describes how to play a particular chord.
type ChordVoicing
    # Indicates movement of the root from the previous chord.
    root :: Int8
    # Intervals beyond the "base" to be played by the left hand.
    # e.g., starting with a '0' would indicate that the root should be
    # played.  All intervals are in half steps.
    left :: Vector{Int8}
    # Intervals to be played by the right hand, starting with the last
    # note in the left hand.
    right :: Vector{Int8}
end

function make_notes(base, v :: ChordVoicing, first=false)
    notes = Vector{Int8}()
    cur = base
    for lh in v.left
        cur += lh
        push!(notes, cur)
    end
    rh_start = cur + v.right[1]
    for rh in v.right
        cur += rh
        push!(notes, cur)
    end
    # TODO: Adjust octave.
    while first && rh_start < 54
        rh_start += 12
        notes .+= 12
        base += 12
    end
    (base, notes)
end

# A voicing is then a full pattern, for multiple chords.
type Voicing
    name :: String
    chords :: Vector{ChordVoicing}
end

function make_notes(chords, v :: Voicing)
    base = base_notes[root_name(chords)]
    chords = Vector{Vector{Int8}}()
    first = true
    for ch in v.chords
        base += ch.root
        (base, notes) = make_notes(base, ch, first)
        first = false
        push!(chords, notes)
    end
    Dict("type" => "voicing",
        "chords" => chords)
end

# Generate patterns for a particular voicing.  The name will be
# printed before, and the hint afterwards.  This should be enough for
# the reader to understand how to voice the chord.  The chord start
# can be used to skip part of the progression.
function gen_cycle(exercise, progression, voicing, hint; chord_start=1)
    name = voicing.name
    for prog in progression
        chord = join(prog[chord_start:end], "-")
        notes = make_notes(prog[1], voicing)
        notes["chords"] = notes["chords"][chord_start:end]
        emit("$name $chord $hint ($exercise)", JSON.json(notes))
    end
end

# The "shell" voicing, with a 1-7 played first in the left hand.
# However, we'll encode the 7 in the right hand, so we can
# automatically find the right octave.
shell17_voicing = let
    minor = ChordVoicing(0, [0], [10])
    dominant = ChordVoicing(5, [0], [4])
    major = ChordVoicing(-7, [0], [11])
    Voicing("shell", [minor, dominant, major])
end

shell13_voicing = let
    minor = ChordVoicing(0, [0], [3])
    dominant = ChordVoicing(-7, [0], [10])
    major = ChordVoicing(5, [0], [4])
    Voicing("shell", [minor, dominant, major])
end

guide173_voicing = let
    minor = ChordVoicing(0, [0], [10, 5])
    dominant = ChordVoicing(-7, [0], [16, 6])
    major = ChordVoicing(5, [0], [11, 5])
    Voicing("guide", [minor, dominant, major])
end

guide137_voicing = let
    minor = ChordVoicing(0, [0], [15, 7])
    dominant = ChordVoicing(5, [0], [10, 6])
    major = ChordVoicing(-7, [0], [16, 7])
    Voicing("guide", [minor, dominant, major])
end

guide73_voicing = let
    minor = ChordVoicing(0, [], [10, 5])
    dominant = ChordVoicing(-7, [], [16, 6])
    major = ChordVoicing(5, [], [11, 5])
    Voicing("guide", [minor, dominant, major])
end

guide37_voicing = let
    minor = ChordVoicing(0, [], [15, 7])
    dominant = ChordVoicing(5, [], [10, 6])
    major = ChordVoicing(-7, [], [16, 7])
    Voicing("guide", [minor, dominant, major])
end

function generate()
    println("DELETE FROM probs;")

    # Exercise 1-1A, ii-V-I, voiced with a 1-7 shell first.
    gen_cycle("1-1A", major_iivi, shell17_voicing, "1-7")

    # Exercise 1-1B, ii-V-I, voiced with a 1-3 shell first.
    gen_cycle("1-1B", major_iivi, shell13_voicing, "1-3")

    # Exercise 1-1A, ii-V-I, voiced with a 1-7 shell first.
    gen_cycle("1-1A", major_iivi, shell13_voicing, "1-7", chord_start=2)

    # Exercise 1-1B, ii-V-I, voiced with a 1-3 shell first.
    gen_cycle("1-1B", major_iivi, shell17_voicing, "1-3", chord_start=2)

    # Exercise 1-1A, ii-V-I, voiced with a 1-7 shell first.
    gen_cycle("1-1A", major_iivi, shell17_voicing, "1-7", chord_start=3)

    # Exercise 1-1B, ii-V-I, voiced with a 1-3 shell first.
    gen_cycle("1-1B", major_iivi, shell13_voicing, "1-3", chord_start=3)

    # Exercise 2-1A, ii-V-I, guide tones, 1-7-3 first.
    gen_cycle("2-1A", major_iivi, guide173_voicing, "1-7-3")

    # Exercise 2-1B, ii-V-I, guide tones, 1-3-7 first.
    gen_cycle("2-1B", major_iivi, guide137_voicing, "1-3-7")

    # Exercise 2-1A, ii-V-I, guide tones, 1-7-3 first.
    gen_cycle("2-1A", major_iivi, guide137_voicing, "1-7-3", chord_start=2)

    # Exercise 2-1B, ii-V-I, guide tones, 1-3-7 first.
    gen_cycle("2-1B", major_iivi, guide173_voicing, "1-3-7", chord_start=2)

    # Exercise 2-1A, ii-V-I, guide tones, 1-7-3 first.
    gen_cycle("2-1A", major_iivi, guide173_voicing, "1-7-3", chord_start=3)

    # Exercise 2-1B, ii-V-I, guide tones, 1-3-7 first.
    gen_cycle("2-1B", major_iivi, guide137_voicing, "1-3-7", chord_start=3)

    # Repeat of the above, but rootless.
    gen_cycle("2-1Ar", major_iivi, guide73_voicing, "7-3")
    gen_cycle("2-1Br", major_iivi, guide37_voicing, "3-7")
    gen_cycle("2-1Ar", major_iivi, guide37_voicing, "7-3", chord_start=2)
    gen_cycle("2-1Br", major_iivi, guide73_voicing, "3-7", chord_start=2)
    gen_cycle("2-1Ar", major_iivi, guide73_voicing, "7-3", chord_start=3)
    gen_cycle("2-1Br", major_iivi, guide37_voicing, "3-7", chord_start=3)

    if false
        # Exercise 2-4A, Diminshed / Tri-tone sub, guide tones, 1-7-3 first.
        gen_cycle("2-4A", min_tritone_sub, "guide", "1-7-3")

        # Exercise 2-4B, Diminshed / Tri-tone sub, guide tones, 1-3-7 first.
        gen_cycle("2-4B", min_tritone_sub, "guide", "1-3-7")
    end
end

generate()


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using AbaqusReader

"""
    abaqus_read_mesh(fn::String)

Read and parse ABAQUS `.inp` file.

`fn` (filename) is the name of the file to parse.
"""
function abaqus_read_mesh(fn::String)
    m = AbaqusReader.abaqus_read_mesh(fn)
    return Mesh(m)
end

"""
    create_surface_elements(mesh::Mesh, surface_name::Symbol)

Create a set of surface elements from solid elements.

Notation follow what is defined in ABAQUS. For example, if solid elements
are Tet10, surface elements will be Tri6 and they can be used to define
boundary conditions.
"""
function create_surface_elements(mesh::Mesh, surface_name::Symbol)
    elements = Element[]
    for (elid, elsi) in mesh.surface_sets[surface_name]
        elty = mesh.element_types[elid]
        elco = mesh.elements[elid]
        chel, chcon = AbaqusReader.create_surface_element(elty, elsi, elco)
        ch = Element(getfield(JuliaFEM, chel), chcon)
        push!(elements, ch)
    end
    update!(elements, "geometry", mesh.nodes)
    return elements
end

"""
    create_surface_elements(mesh::Mesh, surface_name::String)

Create a set of surface elements from solid elements.

Notation follow what is defined in ABAQUS. For example, if solid elements
are Tet10, surface elements will be Tri6 and they can be used to define
boundary conditions.
"""
function create_surface_elements(mesh::Mesh, surface_name::String)
    return create_surface_elements(mesh, Symbol(surface_name))
end

"""
    create_nodal_elements(mesh::Mesh, node_set_name::String)

Create a set of "nodal" elements from node set.

Nodal elements are of type Poi1. In principle they don't have volume and it's
not possible to integrate over them. It is however possible to add boundary
conditions to them.
"""
function create_nodal_elements(mesh::Mesh, node_set_name::String)
    node_ids = mesh.node_sets[Symbol(node_set_name)]
    elements = [Element(Poi1, [j]) for j in node_ids]
    update!(elements, "geometry", mesh.nodes)
    return elements
end


struct RandomModel{T} <: Persa.Model{T}
    preference::Persa.Preference{T}
    users::Int
    items::Int
end

RandomModel(dataset::Persa.Dataset) = RandomModel(dataset.preference, Persa.users(dataset), Persa.items(dataset))

@testset "Model Learn Tests" begin
    df1 = createdummydatasetone()
    dataset1 = Persa.Dataset(df1)

    df2 = createdummydatasettwo()
    dataset2 = Persa.Dataset(df2, 10, 10)

    Persa.predict(model::RandomModel, user::Int, item::Int) = rand(model.preference.possibles)

    model1 = RandomModel(dataset1)
    model2 = RandomModel(dataset2)

    @testset "Dummy Tests" begin
        @test model1[1,1] > 0

        values = model1[dataset1]
        @test length(values) == length(dataset1)

        values = model2[dataset2]
        @test length(values) == length(dataset2)

        @test ismissing(dataset1[3,6])
    end
end


# Substitution Matrices
# =====================
#
# Substitution matrix types and predefined matrices.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/BioAlignments.jl/blob/master/LICENSE.md

"""
Supertype of substitution matrix.

The required method:

* `Base.getindex(submat, x, y)`: substitution score/cost from `x` to `y`
"""
abstract type AbstractSubstitutionMatrix{S<:Real} end

"""
Substitution matrix.
"""
struct SubstitutionMatrix{T,S} <: AbstractSubstitutionMatrix{S}
    # square substitution matrix
    data::Matrix{S}
    # score is defined or not
    defined::BitMatrix

    function SubstitutionMatrix{T,S}(data::Matrix{S}, defined::BitMatrix) where {T,S}
        @assert size(data) == size(defined)
        @assert size(data, 1) == size(data, 2) == length(BioSymbols.alphabet(T)) - 1
        return new{T,S}(data, defined)
    end

end

function SubstitutionMatrix(::Type{T},
                            submat::AbstractMatrix{S},
                            defined::AbstractMatrix{Bool},
                            default_match::S,
                            default_mismatch::S) where {T,S}
    alpha = alphabet_without_gap(T)
    n = length(alpha)
    data = Matrix{S}(undef, (n, n))
    for x in alpha, y in alpha
        i = index(x)
        j = index(y)
        if defined[i,j]
            data[i,j] = submat[i,j]
        else
            data[i,j] = x == y ? default_match : default_mismatch
        end
    end
    return SubstitutionMatrix{T,S}(data, copy(defined))
end

function SubstitutionMatrix(::Type{T},
                            submat::AbstractMatrix{S},
                            # assume scores of all substitutions are defined
                            defined=trues(size(submat));
                            default_match=S(0), default_mismatch=S(0)) where {T,S}
    return SubstitutionMatrix(T, submat, defined, default_match, default_mismatch)
end

function SubstitutionMatrix(scores::AbstractDict{Tuple{T,T},S};
                            default_match=S(0), default_mismatch=S(0)) where {T,S}
    n = length(BioSymbols.alphabet(T)) - 1
    submat = Matrix{S}(undef, (n, n))
    defined = falses(n, n)
    for ((x, y), score) in scores
        i = index(x)
        j = index(y)
        submat[i,j] = score
        defined[i,j] = true
    end
    return SubstitutionMatrix(T, submat, defined, default_match, default_mismatch)
end

Base.convert(::Type{Matrix}, submat::SubstitutionMatrix) = copy(submat.data)

@inline function Base.getindex(submat::SubstitutionMatrix{T}, x, y) where T
    i = index(convert(T, x))
    j = index(convert(T, y))
    return submat.data[i,j]
end

function Base.setindex!(submat::SubstitutionMatrix{T}, val, x, y) where T
    i = index(convert(T, x))
    j = index(convert(T, y))
    submat.data[i,j] = val
    submat.defined[i,j] = true
    return submat
end

function Base.copy(submat::SubstitutionMatrix{T,S}) where {T,S}
    return SubstitutionMatrix{T,S}(copy(submat.data), copy(submat.defined))
end

Base.minimum(submat::SubstitutionMatrix) = minimum(submat.data)
Base.maximum(submat::SubstitutionMatrix) = maximum(submat.data)

function Base.show(io::IO, submat::SubstitutionMatrix{T,S}) where {T,S}
    alpha = alphabet_without_gap(T)
    n = length(alpha)
    mat = Matrix{String}(undef, (n, n))
    for (i, x) in enumerate(alpha), (j, y) in enumerate(alpha)
        i′ = index(x)
        j′ = index(y)
        mat[i, j] = string(submat.data[i′, j′])
        if !submat.defined[i′, j′]
            mat[i, j] = underline(mat[i, j])
        end
    end

    # add rows and columns
    rows = map(string, alpha)
    cols = vcat(" ", rows)
    mat = hcat(rows, mat)
    mat = vcat(reshape(cols, 1, length(cols)), mat)

    println(io, summary(submat), ':')
    width = maximum(map(x -> length(x), mat))
    for i in 1:n + 1
        for j in 1:n + 1
            print(io, lpad(mat[i, j], width + ifelse(last(mat[i, j]) == '\U0332', 2, 1)))
        end
        println(io)
    end
    print(io, "(underlined values are default ones)")
end

underline(s) = join([string(c, '\U0332') for c in s])

# Return a vector of all symbols of `T` except the gap symbol.
function alphabet_without_gap(::Type{T}) where T
    return filter!(x -> x != BioSymbols.gap(T), collect(BioSymbols.alphabet(T)))
end

# Return the row/column index of `nt`.
function index(nt::BioSymbols.NucleicAcid)
    return convert(Int, nt)
end

# Return the row/column index of `aa`.
function index(aa::BioSymbols.AminoAcid)
    return convert(Int, aa) + 1
end

"""
Dichotomous substitution matrix.
"""
struct DichotomousSubstitutionMatrix{S} <: AbstractSubstitutionMatrix{S}
    match::S
    mismatch::S
end

function DichotomousSubstitutionMatrix(match::Real, mismatch::Real)
    typ = promote_type(typeof(match), typeof(mismatch))
    return DichotomousSubstitutionMatrix{typ}(match, mismatch)
end

function Base.convert(::Type{SubstitutionMatrix{T,S}},
                      submat::DichotomousSubstitutionMatrix) where {T,S}
    n = length(BioSymbols.alphabet(T)) - 1
    data = Matrix{S}(undef, (n, n))
    fill!(data, submat.mismatch)
    data[diagind(data)] .= submat.match
    return SubstitutionMatrix{T,S}(data, trues(n, n))
end

@inline function Base.getindex(submat::DichotomousSubstitutionMatrix, x, y)
    return ifelse(x == y, submat.match, submat.mismatch)
end

function Base.show(io::IO, submat::DichotomousSubstitutionMatrix)
    match = string(submat.match)
    mismatch = string(submat.mismatch)
    w = max(length(match), length(mismatch))
    println(io, typeof(submat), ':')
    # right aligned
    print(io, "     match = ", lpad(match, w), '\n')
    print(io, "  mismatch = ", lpad(mismatch, w))
end

# Predefined Substitution Matrices
# --------------------------------

function load_submat(::Type{T}, name) where T
    return parse_ncbi_submat(T, joinpath(dirname(@__FILE__), "data", "submat", name))
end

function parse_ncbi_submat(::Type{T}, filepath) where T
    scores = Dict{Tuple{T,T},Int}()
    cols = T[]
    open(filepath) do io
        for line in eachline(io)
            line = chomp(line)
            if startswith(line, "#") || isempty(line)
                continue  # skip comments and empty lines
            end
            if isempty(cols)
                for y in eachmatch(r"[A-Z*]", line)
                    @assert length(y.match) == 1
                    push!(cols, convert(T, first(y.match)))
                end
            else
                x = convert(T, first(line))
                ss = collect(eachmatch(r"-?\d+", line))
                @assert length(ss) == length(cols)
                for (y, s) in zip(cols, ss)
                    scores[(x, y)] = parse(Int, s.match)
                end
            end
        end
    end
    return SubstitutionMatrix(scores, default_match=0, default_mismatch=0)
end

"EDNAFULL (or NUC4.4) substitution matrix"
const EDNAFULL = load_submat(BioSymbols.DNA, "NUC.4.4")

"PAM30 substitution matrix"
const PAM30    = load_submat(BioSymbols.AminoAcid, "PAM30")

"PAM70 substitution matrix"
const PAM70    = load_submat(BioSymbols.AminoAcid, "PAM70")

"PAM250 substitution matrix"
const PAM250   = load_submat(BioSymbols.AminoAcid, "PAM250")

"BLOSUM45 substitution matrix"
const BLOSUM45 = load_submat(BioSymbols.AminoAcid, "BLOSUM45")

"BLOSUM50 substitution matrix"
const BLOSUM50 = load_submat(BioSymbols.AminoAcid, "BLOSUM50")

"BLOSUM62 substitution matrix"
const BLOSUM62 = load_submat(BioSymbols.AminoAcid, "BLOSUM62")

"BLOSUM80 substitution matrix"
const BLOSUM80 = load_submat(BioSymbols.AminoAcid, "BLOSUM80")

"BLOSUM90 substitution matrix"
const BLOSUM90 = load_submat(BioSymbols.AminoAcid, "BLOSUM90")


using .JSort
import CSV
import DelimitedFiles: writedlm
using Polynomials

const E_MIN = 5  # E Threshold

function sortevent!(matrices, event::Event, parameters::Parameters)
    #=
    Algorithm:
    For each ΔE, look for an E in the associated E-detector. If no E is found,
    or multiple correlations are made, ignore this count.
    =#
    minievent = correlate(event)

    # increment!(matrices[:emul], length(minievents))
    # increment!(matrices[:Δemul], length(minievents))

    # If there were more than 1 correlated event
    # in the time window, it is impossible to tell which
    # LaBr corresponds to which SiRi detection
    isnothing(minievent) && return

    # Calibrate the events
    e = calibrate(minievent, :e, parameters)
    Δe = calibrate(minievent, :Δe, parameters)

    # Reject event if under threshold
    e ≤ E_MIN || Δe ≤ E_MIN && return

    # Look for events within the banana defined by Qkinz
    thick::Float64 = parameters.range(e+Δe) - parameters.range(e)
    thickdev::Float64 = parameters.thickrange[2] + e*parameters.thickrange[3]

    # Increment the matrices
    b, f = minievent.back, minievent.front
    increment!(matrices[:front], Δe, minievent.frontchannel)
    increment!(matrices[:back], e, minievent.backchannel)
    increment!(matrices[:eΔe], e, Δe)
    # increment!(matrices[Symbol("eΔeb",b, "f", f)], e, Δe)
    increment!(matrices[:thick], thick)

    # Check if the event comes within the banana
    within_banana = abs(thick-parameters.thickrange[1]) < thickdev
    !within_banana && return

    # Increment even more matrices
    increment!(matrices[:eΔethick], e, Δe)
    increment!(matrices[:heΔe], e+Δe)
    # increment!(matrices[Symbol("heΔerf", f)], e+Δe)
    # increment!(matrices[Symbol("heΔeb", b, "f", f)], e+Δe)

    ex = calibrate(e+Δe, minievent, :ex, parameters)
    increment!(matrices[:hex], ex)
    # increment!(matrices[Symbol("hexrf", f)], ex)

    # TODO Find a way to support time evolution
    # Time evolution

    # LaBr
    increment!(matrices[:labrmul], event.labrnum)
    for (channel, labr) in eachlabr(event)
        # Check that both adc and tdc have been set
        !islabrgood(labr) && return
        labre = calibrate(labr.adc, :labre, parameters)
        # TODO why this magic /8?
        labrtime = calibrate(labr.tdc/8, :labrtime, parameters)
        labrtime_c = tLaBr(labrtime, labre, e)

        increment!(matrices[:labre], labre, channel)
        increment!(matrcies[:labrtimec], channel)

    end
    # for (channel, labr) in event.labr
        # labr.adc ≤ 0 && continue
        # e = 0
    # end
    return
    # e = ecalibrator(channel_b, channel_f)(si_e[channel_b])
    # Do thick stuff here
end

function oldmakeeΔe(events::Array{TrueEvent, 1}, parameters::Parameters)
    meΔe = OMatrix{Array{Int64, 2}}((NUMBINS_E, NUMBINS_E), max=(MAX_E, MAX_ΔE))
    meΔebf = [OMatrix{Array{Int64, 2}}((NUMBINS_E, NUMBINS_E), max=(MAX_E, MAX_ΔE))
              for f in 1:8, b in 1:8]
    meΔebf_raw = [[] for f in 1:8, b in 1:8]

    for (i, event) in enumerate(events)
        e = calibrate(event.e, 0.0, 5.0)
        Δe = calibrate(event.Δe, 0.0, 2.5)
        push!(meΔebf_raw[event.b, event.f], (e, Δe))
        increment!(meΔe, e, Δe)
        increment!(meΔebf[event.b, event.f], e, Δe)
    end
    #save(meΔe, "mede", parameters.savepath)
    for b in 1:8, f in 1:8
        #save(meΔebf[b, f], "medeb$(b)f$f", parameters.savepath)
        path = joinpath(getpath(parameters, :savepath), "/medeb$(b)f$f.csv")
        CSV.write(path, meΔebf_raw[b, f])
    end
    meΔe, meΔebf
end

function makeeΔe(events::AbstractVector{TrueEvent}, parameters::Parameters; docalibrate=true)
    eΔe = [Tuple{Float32, Float32}[] for f in 1:8, b in 1:8]
    ecalibrators  = parameters.calibrator[:e]
    Δecalibrators = parameters.calibrator[:Δe]
    e = Δe = 0.0
    for (i, event) in enumerate(events)
        e  = ecalibrators[event.f, event.b](event.e)
        Δe = Δecalibrators[event.f, event.b](event.Δe)
        push!(eΔe[event.f, event.b], (e, Δe))
    end
    for f in 1:8
        total = Tuple{Float32, Float32}[]
        for b in 1:8
            path = joinpath(getpath(parameters, :savepath), "edeb$(b)f$(f).csv")
            # Stupid writedlm syntax
            data = []
            push!(data, ("e", "de"))
            append!(data, eΔe[f, b])
            open(path, "w") do io
                writedlm(io, data, ',')
            end
            append!(total, eΔe[f, b])
        end
        path = joinpath(getpath(parameters, :savepath), "edef$(f).csv")
        data = []
        push!(data, ("e", "de"))
        append!(data, total)
        open(path, "w") do io
            writedlm(io, data, ',')
        end
    end
end

function makeeΔebin(events::AbstractVector{TrueEvent}, parameters::Parameters; T=Float32)
    eΔe = [Tuple{T, T}[] for f in 1:8, b in 1:8]
    ecalibrators  = parameters.calibrator[:e]
    Δecalibrators = parameters.calibrator[:Δe]
    e = Δe = 0.0
    for (i, event) in enumerate(events)
        e  = ecalibrators[event.f, event.b](event.e)
        Δe = Δecalibrators[event.f, event.b](event.Δe)
        push!(eΔe[event.f, event.b], (e, Δe))
    end
    for f in 1:8
        path = joinpath(getpath(parameters, :savepath), "edef$(f).bin")
        iototal = open(path, "w")
        # Write the total length
        totallength = length.(eΔe[f, :]) |> sum |> T
        write(iototal, totallength)
        for b in 1:8
            path = joinpath(getpath(parameters, :savepath), "edeb$(b)f$(f).bin")
            data = eΔe[f, b]
            # Stupid writedlm syntax
            open(path, "w") do io
                write(io, length(data) |> T)
                write(io, data)
            end
            # Append to the summed data
            write(iototal, data)
        end
        close(iototal)
    end
end

function makelabr(events::AbstractVector{TrueEvent}, parameters::Parameters; docalibrate=true)
    # TODO: OMatrices are *almost* useless
    # Use long arrays instead
    #gamma_e = OMatrix{Array{Int64, 2}}((2000, 32), max=(14000, 32))
    #gamma_t = OMatrix{Array{Int64, 2}}((500, 32))
    channels = 1:33
    Γe = Dict(i => Float32[] for i in channels)
    Γt = Dict(i => Float32[] for i in channels)

    ecalibrators = parameters.calibrator[:γ]
    tcalibrators = parameters.calibrator[:t]
    @debug "Making LaBr gamma events"
    e = 0.0
    t = 0.0
    for event in events
        for γ in event.labr
            e = ecalibrators[γ.channel+1](γ.adc)
            t = tcalibrators[γ.channel+1](γ.tdc/8)
            #increment!(gamma_e, e, gamma.channel)
            push!(Γe[γ.channel+1], e)
            push!(Γt[γ.channel+1], t)
            #increment!(gamma_t, t, gamma.channel)
        end
    end
    #save(gamma_e, "gammae", getpath(parameters, :savepath))
    #save(gamma_t, "gammat", getpath(parameters, :savepath))
    return Γe, Γt
end


using DataStructures: OrderedDict
using Flux, CuArrays
export Game, showAll, validMoves, randomPlace, place, getGameState, int, PLAYERS

## Game setting
# board size
BOARD_SIZE = Int8(9)
ACTION_SIZE = Int8(9)
HISTORY = Int8(4)

# trainsformation of matrix mapping
ALL_MOVES = (
    original_dont_move  = Int8[1, 2, 3, 4, 5, 6, 7, 8, 9],
    original_dim1_flip  = Int8[7, 8, 9, 4, 5, 6, 1, 2, 3],
    right_rotate_1x90°  = Int8[7, 4, 1, 8, 5, 2, 9, 6, 3],
    right_rot1x90_flip  = Int8[9, 6, 3, 8, 5, 2, 7, 4, 1],
    right_rotate_2x90°  = Int8[9, 8, 7, 6, 5, 4, 3, 2, 1],
    right_rot2x90_flip  = Int8[3, 2, 1, 6, 5, 4, 9, 8, 7],
    right_rotate_3x90°  = Int8[3, 6, 9, 2, 5, 8, 1, 4, 7],
    right_rot3x90_flip  = Int8[1, 4, 7, 2, 5, 8, 3, 6, 9]
)

# transition from players
PLAYERS = Dict{Union{String, Int8}, Union{String, Int8}}("O"=>Int8(1), Int8(1)=>"O", "X"=>Int8(-1), Int8(-1)=>"X", zero(Int8)=>".")
# NEXT_TURN = Dict{Union{String, Int8}, Union{String, Int8}}("O"=>"X", Int8(1)=>Int8(-1), "X"=>"O", Int8(-1)=>Int8(1))

# The state we need to evaluate winning situtaions
STATE = Array{Array{Int8,1}}([
        [1, 2, 3], [4, 5, 6], [7, 8, 9],
        [1, 4, 7], [2, 5, 8], [3, 6, 9],
        [1, 5, 9], [3, 5, 7]
        ])

# give an action, get all possible winning situations
SLOT = Array{Array{Array{Int8,1},1},1}([
        [[1, 2, 3], [1, 4, 7], [1, 5, 9]],
        [[1, 2, 3], [2, 5, 8]],
        [[1, 2, 3], [3, 5, 7], [3, 6, 9]],
        [[1, 4, 7], [4, 5, 6]],
        [[4, 5, 6], [2, 5, 8], [1, 5, 9], [3, 5, 7]],
        [[4, 5, 6], [3, 6, 9]],
        [[1, 4, 7], [7, 8, 9], [3, 5, 7]],
        [[7, 8, 9], [2, 5, 8]],
        [[7, 8, 9], [3, 6, 9], [1, 5, 9]]
        ])

# offset move -> actual board index
OFFSET = Dict{Tuple{Int8,Int8},Tuple{Int8,Int8}}(
  (3, 5) => (2, 8), (3, 6) => (2, 9), (6, 9) => (6, 9), (4, 4) => (5, 1), (9, 4) => (8, 7), (9, 1) => (7, 7), (8, 9) => (9, 6), (3, 1) => (1, 7), (4, 5) => (5, 2),
  (9, 9) => (9, 9), (2, 4) => (2, 4), (6, 5) => (5, 8), (4, 9) => (6, 3), (8, 4) => (8, 4), (1, 2) => (1, 2), (9, 2) => (7, 8), (5, 1) => (4, 4), (9, 8) => (9, 8),
  (3, 4) => (2, 7), (9, 6) => (8, 9), (8, 2) => (7, 5), (2, 3) => (1, 6), (2, 6) => (2, 6), (5, 4) => (5, 4), (2, 1) => (1, 4), (2, 5) => (2, 5), (8, 5) => (8, 5),
  (5, 9) => (6, 6), (1, 4) => (2, 1), (3, 9) => (3, 9), (2, 8) => (3, 5), (5, 8) => (6, 5), (4, 8) => (6, 2), (7, 8) => (9, 2), (1, 1) => (1, 1), (7, 7) => (9, 1),
  (7, 9) => (9, 3), (5, 7) => (6, 4), (7, 2) => (7, 2), (2, 2) => (1, 5), (6, 3) => (4, 9), (1, 7) => (3, 1), (4, 7) => (6, 1), (6, 1) => (4, 7), (6, 8) => (6, 8),
  (8, 8) => (9, 5), (7, 1) => (7, 1), (1, 3) => (1, 3), (7, 5) => (8, 2), (2, 9) => (3, 6), (3, 2) => (1, 8), (7, 6) => (8, 3), (9, 3) => (7, 9), (3, 7) => (3, 7),
  (6, 6) => (5, 9), (5, 5) => (5, 5), (2, 7) => (3, 4), (4, 1) => (4, 1), (8, 3) => (7, 6), (4, 2) => (4, 2), (1, 6) => (2, 3), (4, 3) => (4, 3), (9, 7) => (9, 7),
  (5, 6) => (5, 6), (8, 6) => (8, 6), (4, 6) => (5, 3), (8, 1) => (7, 4), (3, 3) => (1, 9), (6, 2) => (4, 8), (7, 4) => (8, 1), (1, 9) => (3, 3), (8, 7) => (9, 4),
  (6, 4) => (5, 7), (5, 2) => (4, 5), (1, 8) => (3, 2), (6, 7) => (6, 7), (7, 3) => (7, 3), (5, 3) => (4, 6), (3, 8) => (3, 8), (1, 5) => (2, 2), (9, 5) => (8, 8)
)

int(x::Union{Char, String}, T=Int) = parse(T, x)
convertState(state) = [int(i, Int8) for i in state]

mutable struct Game
    boards::AbstractArray{Float32,3}
    vmoves::Dict{Int8, AbstractArray}
    nextPlayer::Int8
    steps::Int8
    winner::Int8
    state::String
    function Game()
        boards = zeros(Float32, (BOARD_SIZE, BOARD_SIZE, 2HISTORY+1))
        vmoves = Dict(Int8(1)=>trues(9), Int8(-1)=>trues(9), Int8(0)=>zeros(Int8, 9,9))
        new(boards, vmoves, Int8(-1),  Int8(1), Int8(0), "")
    end
end

function nozero(x::AbstractArray)
    for i in x
        i==0 && return false
    end
    return true
end

@inline function rmInvalid!(self::Game, x::Integer, y::Integer)
    if @inbounds self.vmoves[0][x,y] != 0
        self.winner = -self.nextPlayer
        return
    end
    @inbounds self.vmoves[0][x,y] = self.nextPlayer

    @inbounds for triple in SLOT[y]
        @inbounds code = sum(self.vmoves[0][x, triple])
        if  code == 2self.nextPlayer
            @inbounds self.vmoves[-self.nextPlayer][x] = false
        end
        if @inbounds code ==-self.nextPlayer && nozero(self.vmoves[0][x, triple])
            count = 0
            for t in STATE
                @inbounds count+= (sum(self.vmoves[0][x, t]) == -2self.nextPlayer)
                count>0 && break
            end
            count >0 && continue
            @inbounds self.vmoves[self.nextPlayer][x] = true
        end
        if @inbounds code==3self.nextPlayer
            self.winner = self.nextPlayer
            return
        end
    end
    return
end

function Game(state::String)
    self = Game()
    isempty(state) && return self
    moves = convertState(state)
    num_moves =  length(moves)
    subboard = moves[1]
    update_bound = max(2, num_moves - HISTORY+1)
    upp = max(1, HISTORY+2-num_moves)
    @sync for i=2:num_moves
        @inbounds rmInvalid!(self, subboard, moves[i])
        @inbounds x,y = OFFSET[(subboard, moves[i])]
        @inbounds self.boards[x, y, upp:HISTORY] .= self.nextPlayer
        @inbounds self.boards[x, y, HISTORY+1:end-(upp-1)] .= -self.nextPlayer
        upp += (i ≥ update_bound);
        self.nextPlayer *= -1;
        @inbounds subboard = moves[i]
    end
    self.steps = num_moves
    @inbounds self.boards[:,:,HISTORY+1] .= self.nextPlayer
    self.state = state
    return self
end

function Base.show(io::IO, self::Game)
    "print out the board"
    compact = get(io, :compact, false)
    if !compact
        board = self.boards[:,:,HISTORY]
        for i=1:9
            print(io, PLAYERS[board[i,1]]," ", PLAYERS[board[i,2]]," ",PLAYERS[board[i,3]]," | ")
            print(io, PLAYERS[board[i,4]]," ", PLAYERS[board[i,5]]," ",PLAYERS[board[i,6]]," | ")
            print(io, PLAYERS[board[i,7]]," ", PLAYERS[board[i,8]]," ",PLAYERS[board[i,9]],"\n")
            if i in (3, 6)
                print(io, "------+-------+------\n")
            end
        end
        showState = isempty(self.state) ? "0" : self.state
        print(io, "\nSubboard: ",showState[end],"; next Player: ", PLAYERS[self.nextPlayer], ";\n")
        print(io, "\nGame States: ",showState,"; Steps: ",self.steps-1,"; ")
        iszero(self.winner) ? print(io, "Game continue...\n") : print(io, PLAYERS[self.winner], " wins\n")
    end
end

function checkWin(self::Game, move::Integer)
    @inbounds x = int(self.state[end])
    @inbounds for triple in SLOT[move]
        @inbounds if sum(self.vmoves[0][x, triple]) == 2self.nextPlayer
            return true
        end
    end
    return false
end

function validMoves(self::Game)
    "Returns all the legal moves for current subboard"
    isempty(self.state) && return keys(OFFSET)
    @inbounds x = int(self.state[end])
    @inbounds available_moves = Array{Int8}(findall(iszero, self.vmoves[0][x,:]))
    if @inbounds !self.vmoves[-self.nextPlayer][x]
        return filter(available_moves) do move
            checkWin(self, move)
        end
    end
    validmoves = filter(available_moves) do move
        @inbounds self.vmoves[self.nextPlayer][move]
    end
    @inbounds if isempty(validmoves)
        return isempty(available_moves) ? [Int8(1)] : available_moves
    end
    return validmoves
end

function place(self::Game, move::Union{Tuple, Integer})
    !iszero(self.winner) && return false
    if isempty(self.state)
        @inbounds rmInvalid!(self, move[1], move[2])
        @inbounds self.state = "$(move[1])$(move[2])"
    else
        @inbounds rmInvalid!(self, int(self.state[end]), move)
        x, y = OFFSET[(int(self.state[end]), move)]
        if self.boards[x, y, HISTORY] != 0
            self.winner = -self.nextPlayer
            return false
        end
        self.state = "$(self.state)$move"
    end
    num_moves = length(self.state)
    iters = max(1, HISTORY-num_moves+2)
    player = self.nextPlayer
    for i=HISTORY:-1:iters
        @inbounds subboard = int(self.state[end+i-HISTORY-1])
        @inbounds x, y = OFFSET[(subboard, int(self.state[end+i-HISTORY]))]
        @inbounds self.boards[x, y, i] = player
        @inbounds self.boards[x, y, end-i+1] = -player
        player *= -1
    end
    self.steps += 1
    self.nextPlayer *= -1
    @inbounds self.boards[:,:,HISTORY+1] .= self.nextPlayer
    return true
end

function randomPlace(self::Game, times::Integer=1)
    for i = 1:times
        !iszero(self.winner) && break
        moves = validMoves(self)
        place(self, rand(moves))
    end
end

function conv(state::String, typ::Int)
    result = ""
    for (i,e) in enumerate(state)
        result *= string(ALL_MOVES[typ][int(e)])
    end
    return result
end

function getGameState(self::Game)
    return reshape(self.boards, (BOARD_SIZE, BOARD_SIZE, 2HISTORY+1,1))
end

function getAugmentation(self::Game, p::Array{Float32, 1})
    """ The order is
        {game.state => M1, flip1 => M2,
         rotr1x90   => M3, flip2 => M4
         rotr2x90   => M5, flip3 => M6
         rotr3x90   => M7, flip4 => M8}
    """
    # strings = String[];
    # push!(strings, conv(self.state, 2))
    data = Array{Array{Float32, 3}, Array{Float32, 1}, Integer}[]
    gameState1 = self.boards
    push!(data, [gameState1, p, game.nextPlayer])
    # flip along with dim1
    # states = cat(dims=4, reverse(gameState1, dims=1))
    @inbounds push!(data, [reverse(gameState1, dims=1), p[ALL_MOVES[2]], game.nextPlayer])
    for i=3:2:8
        # right rotate 90°
        # push!(strings, conv(self.state, i))
        gameState1 = mapslices(rotr90, gameState1, dims=[1,2])
        @inbounds push!(data, [gameState1, p[ALL_MOVES[i]], game.nextPlayer])
        # states = cat(dims=4, states, gameState1)
        # flip along with dim1
        # push!(strings, conv(self.state, i+1))
        @inbounds push!(data, [reverse(gameState1, dims=1), p[ALL_MOVES[i]], game.nextPlayer])
        # states = cat(dims=4, states, reverse(gameState1, dims=1))
    end
    return data
end


using JsonGrinder, JSON, Test, SparseArrays, Mill

@testset "ExtractMultipleRepresentation" begin
	ex = MultipleRepresentation((ExtractCategorical(["Olda", "Tonda", "Milda"]),
		JsonGrinder.ExtractString(String)))
	e = ex("Olda")

	@test length(e.data) == 2
	@test size(e.data[1].data) == (4, 1)
	@test e.data[1].data[:] ≈ [0, 1, 0, 0]
	@test size(e.data[2].data) == (2053, 1)
	@test findall(x->x > 0, SparseMatrixCSC(e.data[2].data)) .|> Tuple == [(80, 1), (98, 1), (206, 1), (738, 1),
		(1062, 1), (1856, 1)]
end

@testset "show" begin
	ex = MultipleRepresentation((ExtractCategorical(["Olda", "Tonda", "Milda"]),
		JsonGrinder.ExtractString(String)))
	e = ex("Olda")

	buf = IOBuffer()
	Base.show(buf, ex)
	str_repr = String(take!(buf))
	@test str_repr ==
"""
: MultiRepresentation
 : Categorical d = 4
 : String
"""

	buf = IOBuffer()
	Mill.dsprint(buf, e)
	str_repr = String(take!(buf))
	@test str_repr ==
"""
TreeNode
  ├── ArrayNode(4, 1)
  └── ArrayNode(2053, 1)"""
end


module RTESolver

using StaticArrays
using CUDA
using Adapt

using ..Device: array_type, array_device, CUDADevice, CPU
using ..AngularDiscretizations
using ..Vmrs
using ..AtmosphericStates
using ..Fluxes
using ..BCs
using ..Sources
using ..RTE
using ..Optics
using ..LookUpTables

export solve_lw!, solve_sw!

include("RTESolverKernels.jl")

"""
    solve_lw!(
        slv::Solver{FT,I},
        max_threads::I,
        lkp_args...,
    ) where {I<:Int,FT<:AbstractFloat}

Solver for the longwave radiation problem
"""
function solve_lw!(
    slv::Solver{FT,I},
    max_threads::I,
    lkp_args...,
) where {I<:Int,FT<:AbstractFloat}
    (; as, op, bcs_lw, src_lw, flux_lw, fluxb_lw) = slv
    (; nlay, ncol) = as
    DA = array_type()
    nargs = length(lkp_args)
    @assert nargs < 3
    if nargs > 0
        @assert lkp_args[1] isa LookUpLW
        if nargs == 2
            @assert lkp_args[2] isa LookUpCld
        end
        (; n_gpt, major_gpt2bnd) = lkp_args[1]
        flux = fluxb_lw
    else
        n_gpt = 1
        major_gpt2bnd = DA(UInt8[1])
        flux = flux_lw
    end

    if op isa OneScalar
        rte_lw_noscat_solve!(
            flux,
            flux_lw,
            src_lw,
            bcs_lw,
            op,
            nlay,
            ncol,
            major_gpt2bnd,
            max_threads,
            as,
            lkp_args...,
        ) # no-scattering solver
    else
        rte_lw_2stream_solve!(
            flux,
            flux_lw,
            src_lw,
            bcs_lw,
            op,
            nlay,
            ncol,
            major_gpt2bnd,
            max_threads,
            as,
            lkp_args...,
        ) # 2-stream solver
    end
    return nothing
end

"""
    rte_lw_noscat_solve!(
        flux::FluxLW{FT},
        flux_lw::FluxLW{FT},
        src_lw::SourceLWNoScat{FT},
        bcs_lw::LwBCs{FT},
        op::OneScalar{FT},
        nlay,
        ncol,
        major_gpt2bnd::AbstractArray{UInt8,1},
        max_threads,
        as::AbstractAtmosphericState{FT},
        lkp_args...,
    ) where {FT<:AbstractFloat}

No scattering solver for the longwave problem.
LW fluxes, no scattering, mu (cosine of integration angle) specified by column
Does radiation calculation at user-supplied angles; converts radiances to flux
using user-supplied weights. Currently, only n_gauss_angles = 1 supported.
"""
function rte_lw_noscat_solve!(
    flux::FluxLW{FT},
    flux_lw::FluxLW{FT},
    src_lw::SourceLWNoScat{FT},
    bcs_lw::LwBCs{FT},
    op::OneScalar{FT},
    nlay,
    ncol,
    major_gpt2bnd::AbstractArray{UInt8,1},
    max_threads,
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    nlev = nlay + 1
    n_gpt = length(major_gpt2bnd)
    device = array_device(op.τ)
    if device === CUDADevice() # launch CUDA kernel
        tx = min(ncol, max_threads)
        bx = cld(ncol, tx)
        args = (
            flux,
            flux_lw,
            src_lw,
            bcs_lw,
            op,
            nlay,
            ncol,
            major_gpt2bnd,
            as,
            lkp_args...,
        )
        @cuda threads = (tx) blocks = (bx) rte_lw_noscat_solve_CUDA!(args...)
    else # use Julia native multithreading
        for igpt = 1:n_gpt
            @inbounds Threads.@threads for gcol = 1:ncol
                igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
                compute_optical_props!(op, as, src_lw, gcol, igpt, lkp_args...)
                rte_lw_noscat_source!(src_lw, op, gcol, nlay)
                rte_lw_noscat_transport!(
                    src_lw,
                    bcs_lw,
                    op,
                    gcol,
                    flux,
                    igpt,
                    major_gpt2bnd,
                    nlay,
                    nlev,
                )
                n_gpt > 1 && add_to_flux!(flux_lw, flux, gcol)
            end
        end
    end
    #------------------------------------------------
    return nothing
end
# ---------------------------------------------------------------
function rte_lw_noscat_solve_CUDA!(
    flux::FluxLW{FT},
    flux_lw::FluxLW{FT},
    src_lw::SourceLWNoScat{FT},
    bcs_lw::LwBCs{FT},
    op::OneScalar{FT},
    nlay,
    ncol,
    major_gpt2bnd,
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
    nlev = nlay + 1
    n_gpt = length(major_gpt2bnd)
    if gcol ≤ ncol
        for igpt = 1:n_gpt
            igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
            compute_optical_props!(op, as, src_lw, gcol, igpt, lkp_args...)
            rte_lw_noscat_source!(src_lw, op, gcol, nlay)
            rte_lw_noscat_transport!(
                src_lw,
                bcs_lw,
                op,
                gcol,
                flux,
                igpt,
                major_gpt2bnd,
                nlay,
                nlev,
            )
            n_gpt > 1 && add_to_flux!(flux_lw, flux, gcol)
        end
    end
    return nothing
end

"""
    rte_lw_2stream_solve!(
        flux::FluxLW{FT},
        flux_lw::FluxLW{FT},
        src_lw::SourceLW2Str{FT},
        bcs_lw::LwBCs{FT},
        op::TwoStream{FT},
        nlay,
        ncol,
        major_gpt2bnd::AbstractArray{UInt8,1},
        max_threads,
        as::AbstractAtmosphericState{FT},
        lkp_args...,
    ) where {FT<:AbstractFloat}

Two stream solver for the longwave problem.

RRTMGP provides source functions at each level using the spectral mapping
of each adjacent layer. Combine these for two-stream calculations
lw combine sources (rte_lw_2stream_combine_sources!)
combine RRTMGP-specific sources at levels
compute layer reflectance, transmittance
compute total source function at levels using linear-in-tau (rte_lw_2stream_source!)
transport (adding_lw!)
"""
function rte_lw_2stream_solve!(
    flux::FluxLW{FT},
    flux_lw::FluxLW{FT},
    src_lw::SourceLW2Str{FT},
    bcs_lw::LwBCs{FT},
    op::TwoStream{FT},
    nlay,
    ncol,
    major_gpt2bnd::AbstractArray{UInt8,1},
    max_threads,
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    nlev = nlay + 1
    n_gpt = length(major_gpt2bnd)
    device = array_device(op.τ)
    if device === CUDADevice() # launch CUDA kernel
        tx = min(ncol, max_threads)
        bx = cld(ncol, tx)
        args = (
            flux,
            flux_lw,
            src_lw,
            bcs_lw,
            op,
            nlay,
            ncol,
            major_gpt2bnd,
            as,
            lkp_args...,
        )
        @cuda threads = (tx) blocks = (bx) rte_lw_2stream_solve_CUDA!(args...)
    else # use Julia native multithreading
        for igpt = 1:n_gpt
            Threads.@threads for gcol = 1:ncol
                igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
                compute_optical_props!(op, as, src_lw, gcol, igpt, lkp_args...)
                rte_lw_2stream_combine_sources!(src_lw, gcol, nlev, ncol)
                rte_lw_2stream_source!(op, src_lw, gcol, nlay, ncol)
                adding_lw!(
                    flux,
                    src_lw,
                    bcs_lw,
                    gcol,
                    igpt,
                    major_gpt2bnd,
                    nlev,
                    ncol,
                )
                n_gpt > 1 && add_to_flux!(flux_lw, flux, gcol)
            end
        end
    end
    return nothing
end
# -------------------------------------------------------------------------------------------------
function rte_lw_2stream_solve_CUDA!(
    flux::FluxLW{FT},
    flux_lw::FluxLW{FT},
    src_lw::SourceLW2Str{FT},
    bcs_lw::LwBCs{FT},
    op::TwoStream{FT},
    nlay,
    ncol,
    major_gpt2bnd::AbstractArray{UInt8,1},
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
    nlev = nlay + 1
    n_gpt = length(major_gpt2bnd)
    if gcol ≤ ncol
        for igpt = 1:n_gpt
            igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
            compute_optical_props!(op, as, src_lw, gcol, igpt, lkp_args...)
            rte_lw_2stream_combine_sources!(src_lw, gcol, nlev, ncol)
            rte_lw_2stream_source!(op, src_lw, gcol, nlay, ncol)
            adding_lw!(
                flux,
                src_lw,
                bcs_lw,
                gcol,
                igpt,
                major_gpt2bnd,
                nlev,
                ncol,
            )
            n_gpt > 1 && add_to_flux!(flux_lw, flux, gcol)
        end
    end
    return nothing
end

"""
    solve_sw!(
        slv::Solver{FT,I},
        max_threads::I,
        lkp_args...,
    ) where {I<:Int,FT<:AbstractFloat}

Solver for the shortwave radiation problem
"""
function solve_sw!(
    slv::Solver{FT,I},
    max_threads::I,
    lkp_args...,
) where {I<:Int,FT<:AbstractFloat}
    (; as, op, bcs_sw, src_sw, flux_sw, fluxb_sw) = slv
    (; nlay, ncol) = as
    DA = array_type()
    nargs = length(lkp_args)
    @assert nargs < 3
    if nargs > 0
        @assert lkp_args[1] isa LookUpSW
        if nargs == 2
            @assert lkp_args[2] isa LookUpCld
        end
        (; n_gpt, major_gpt2bnd, solar_src_scaled) = lkp_args[1]
        flux = fluxb_sw
    else
        n_gpt = 1
        major_gpt2bnd = DA(UInt8[1])
        solar_src_scaled = DA(FT[1])
        flux = flux_sw
    end
    # solving radiative transfer equation
    if slv.op isa OneScalar
        rte_sw_noscat_solve!(
            flux,
            flux_sw,
            op,
            bcs_sw,
            nlay,
            ncol,
            solar_src_scaled,
            max_threads,
            as,
            lkp_args...,
        ) # no-scattering solver
    else
        rte_sw_2stream_solve!(
            flux,
            flux_sw,
            op,
            bcs_sw,
            src_sw,
            nlay,
            ncol,
            major_gpt2bnd,
            solar_src_scaled,
            max_threads,
            as,
            lkp_args...,
        ) # 2-stream solver
    end
    return nothing
end

"""
    rte_sw_noscat_solve!(
        flux::FluxSW{FT},
        flux_sw::FluxSW{FT},
        op::OneScalar{FT},
        bcs_sw::SwBCs{FT},
        nlay,
        ncol,
        solar_src_scaled::AbstractArray{FT,1},
        max_threads,
        as::AbstractAtmosphericState{FT},
        lkp_args...,
    ) where {FT<:AbstractFloat}

No-scattering solver for the shortwave problem.
(Extinction-only i.e. solar direct beam)
"""
function rte_sw_noscat_solve!(
    flux::FluxSW{FT},
    flux_sw::FluxSW{FT},
    op::OneScalar{FT},
    bcs_sw::SwBCs{FT},
    nlay,
    ncol,
    solar_src_scaled::AbstractArray{FT,1},
    max_threads,
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    nlev = nlay + 1
    n_gpt = length(solar_src_scaled)
    device = array_device(op.τ)
    # setting references for flux_sw
    if device === CUDADevice() # launching CUDA kernel
        tx = min(ncol, max_threads)
        bx = cld(ncol, tx)
        args = (
            flux,
            flux_sw,
            op,
            bcs_sw,
            nlay,
            ncol,
            solar_src_scaled,
            as,
            lkp_args...,
        )
        @cuda threads = (tx) blocks = (bx) rte_sw_noscat_solve_CUDA!(args...)
    else # launch juila native multithreading
        for igpt = 1:n_gpt
            @inbounds Threads.@threads for gcol = 1:ncol
                igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
                compute_optical_props!(op, as, gcol, igpt, lkp_args...)
                rte_sw_noscat_solve_kernel!(
                    flux,
                    op,
                    bcs_sw,
                    igpt,
                    solar_src_scaled,
                    gcol,
                    nlev,
                )
                n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
            end
        end
    end
    #-----------------------------------------------------------------------------------
    return nothing
end
#--------------------------------------------------------------------------------------------------
function rte_sw_noscat_solve_CUDA!(
    flux::FluxSW{FT},
    flux_sw::FluxSW{FT},
    op::OneScalar{FT},
    bcs_sw::SwBCs{FT},
    nlay,
    ncol,
    solar_src_scaled::AbstractArray{FT,1},
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
    nlev = nlay + 1
    n_gpt = length(solar_src_scaled)
    # setting references for flux_sw
    if gcol ≤ ncol
        for igpt = 1:n_gpt
            igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
            compute_optical_props!(op, as, gcol, igpt, lkp_args...)
            rte_sw_noscat_solve_kernel!(
                flux,
                op,
                bcs_sw,
                igpt,
                solar_src_scaled,
                gcol,
                nlev,
            )
            n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
        end
    end
    return nothing
end
#--------------------------------------------------------------------------------------------------
"""
    rte_sw_2stream_solve!(
        flux::FluxSW{FT},
        flux_sw::FluxSW{FT},
        op::TwoStream{FT},
        bcs_sw::SwBCs{FT},
        src_sw::SourceSW2Str{FT},
        nlay,
        ncol,
        major_gpt2bnd::AbstractArray{UInt8,1},
        solar_src_scaled::AbstractArray{FT,1},
        max_threads,
        as::AbstractAtmosphericState{FT},
        lkp_args...,
    ) where {FT<:AbstractFloat}

Two stream solver for the shortwave problem.
"""
function rte_sw_2stream_solve!(
    flux::FluxSW{FT},
    flux_sw::FluxSW{FT},
    op::TwoStream{FT},
    bcs_sw::SwBCs{FT},
    src_sw::SourceSW2Str{FT},
    nlay,
    ncol,
    major_gpt2bnd::AbstractArray{UInt8,1},
    solar_src_scaled::AbstractArray{FT,1},
    max_threads,
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    nlev = nlay + 1
    n_gpt = length(major_gpt2bnd)
    device = array_device(op.τ)
    if device === CUDADevice()
        tx = min(ncol, max_threads)
        bx = cld(ncol, tx)
        args = (
            flux,
            flux_sw,
            op,
            bcs_sw,
            src_sw,
            nlay,
            ncol,
            major_gpt2bnd,
            solar_src_scaled,
            as,
            lkp_args...,
        )
        @cuda threads = (tx) blocks = (bx) rte_sw_2stream_solve_CUDA!(args...)
    else # launch julia native multithreading
        # setting references for flux_sw
        for igpt = 1:n_gpt
            @inbounds Threads.@threads for gcol = 1:ncol
                igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
                compute_optical_props!(op, as, gcol, igpt, lkp_args...)
                sw_two_stream!(op, src_sw, bcs_sw, gcol, nlay) # Cell properties: transmittance and reflectance for direct and diffuse radiation
                sw_source_2str!(
                    src_sw,
                    bcs_sw,
                    gcol,
                    flux,
                    igpt,
                    solar_src_scaled,
                    major_gpt2bnd,
                    nlay,
                ) # Direct-beam and source for diffuse radiation
                adding_sw!(
                    src_sw,
                    bcs_sw,
                    gcol,
                    flux,
                    igpt,
                    major_gpt2bnd,
                    nlev,
                ) # Transport
                n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
            end
        end
    end
    return nothing
end
#--------------------------------------------------------------------------------------------------
function rte_sw_2stream_solve_CUDA!(
    flux::FluxSW{FT},
    flux_sw::FluxSW{FT},
    op::TwoStream{FT},
    bcs_sw::SwBCs{FT},
    src_sw::SourceSW2Str{FT},
    nlay,
    ncol,
    major_gpt2bnd::AbstractArray{UInt8,1},
    solar_src_scaled::AbstractArray{FT,1},
    as::AbstractAtmosphericState{FT},
    lkp_args...,
) where {FT<:AbstractFloat}
    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
    nlev = nlay + 1
    n_gpt = length(major_gpt2bnd)
    if gcol ≤ ncol
        for igpt = 1:n_gpt
            igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
            compute_optical_props!(op, as, gcol, igpt, lkp_args...)
            sw_two_stream!(op, src_sw, bcs_sw, gcol, nlay) # Cell properties: transmittance and reflectance for direct and diffuse radiation
            sw_source_2str!(
                src_sw,
                bcs_sw,
                gcol,
                flux,
                igpt,
                solar_src_scaled,
                major_gpt2bnd,
                nlay,
            ) # Direct-beam and source for diffuse radiation
            adding_sw!(src_sw, bcs_sw, gcol, flux, igpt, major_gpt2bnd, nlev) # Transport
            n_gpt > 1 && add_to_flux!(flux_sw, flux, gcol)
        end
    end
    return nothing
end
#--------------------------------------------------------------------------------------------------
end


# get all faces, including
#   bboxes, landmarks, features
using Glob
using PyCall
using JLD2
using FileIO

const FR = pyimport("evolveface")
const PIL = pyimport("PIL")

# get all templates
template_files = glob("data/*.jpg")
template_features = hcat(map(template_files) do f
    img = PIL.Image.open(f)
    bounding_boxes, landmarks = FR.detect_faces(img)
    features = FR.extract_feature_IR50A(img, landmarks)
    @assert size(features)[1] == 1
    return vec(features)
end...)

@save "out/template.jld2" template_features

map(glob("data/vid/*.jpg")) do f
    println(f)
    img = PIL.Image.open(f)
    bounding_boxes, landmarks = FR.detect_faces(img)
    size(landmarks)[1] == 0 && return
    features = FR.extract_feature_IR50A(img, landmarks)
    res = (bounding_boxes, landmarks, features)
    @save "out/$(splitext(basename(f))[1]).jld2" res
end


function opt_add(self::Tensor{T, N}, other::Tensor{T, N}) where {T, N}
    outputs__ = Int[0]
    __cret = ccall((:atg_add, libtorch_capi),
                 Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}),
                   outputs__, self.pointer, other.pointer)
    return Tensor{T, N}(Ptr{Cvoid}(outputs__[1]), nothing)
end

function opt_index_select(self::Tensor{T, N}, dim::Int64, index::Int64) where {T, N}
    ptr = ccall((:tensor_method_index_select_int64, libtorch_capi),
                Ptr{Cvoid}, (Ptr{Cvoid}, Clonglong, Clonglong),
                self.pointer, dim, index)
    return Tensor{T, N}(ptr, nothing)
end

function opt_index_select(self::Tensor{T, N}, dim::Int64, i::Array{Int64}) where {T, N}
    return opt_index_select(self, dim, Tensor(i))
end

function opt_index_select(self::Tensor{T, N}, dim::Int64, index::Tensor) where {T, N}
    outputs__ = Int[0]
    __cret = ccall((:atg_index_select, libtorch_capi),
                 Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Clonglong, Ptr{Cvoid}),
                 outputs__, self.pointer, dim, index.pointer)
    return Tensor{T, N}(Ptr{Cvoid}(outputs__[1]), nothing)
end

function opt_reshape(self::Tensor{T}, shape_data::Array{Int64}) where T
    outputs__ = Int[0]
    shape_len = length(shape_data)
    __cret = ccall((:atg_reshape, libtorch_capi),
                 Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Cint),
                 outputs__, self.pointer, shape_data, shape_len)
    return Tensor{T, shape_len}(Ptr{Cvoid}(outputs__[1]), nothing)
end


using Metalhead, Test
using Flux

@testset "ViT" begin
    @test size(ViT()(rand(Float32, 256, 256, 3, 2))) == (1000, 2)
    @test_skip gradtest(ViT(), rand(Float32, 256, 256, 3, 2))
end


module Proc
using YAXArrays.Cubes: YAXArray, cubechunks, caxes
using YAXArrays.Cubes.Axes: getAxis, findAxis, CategoricalAxis, axVal2Index,
  RangeAxis, get_bb, axisfrombb, CubeAxis, axname
using YAXArrays.DAT: mapCube, InDims, OutDims, NValid, AnyMissing
using YAXArrays.Datasets: getsavefolder, Cube
using Dates: year, Date, DateTime
"""
    getNpY(cube)

Get the number of time steps per year
"""
function getNpY(cube)
    timax = getAxis("Time",cube)
    years = year.(timax.values)
    years[end] > years[1] + 1 || error("Must have at least 3 years to estimate number of time steps per year")
    return count(i -> i == years[1] + 1, years)
end

include("MSC.jl")
include("Stats.jl")
include("TSDecomposition.jl")
include("remap.jl")
include("Shapes.jl")
include("extractlonlats.jl")

end


module Compose

using Dates
using HDF5
using Statistics: mean

#

# write(*,*) i2,' temperature T                                      [MeV]        '
# write(*,*) i2,' baryon number density n_b                          [fm^-3]      '
# write(*,*) i2,' hadronic charge fraction Y_q                       []           '
# write(*,*) i2,' magnetic field strength B                          [G]          '

const thermo_qtys = Dict(1 => "pressure p                                         [MeV fm^-3]  ",
                         2 => "entropy per baryon S                               []           ",
                         3 => "shifted baryon chemical potential mu_b-m_n         [MeV]        ",
                         4 => "charge chemical potential mu_q                     [MeV]        ",
                         5 => "lepton chemical potential mu_l                     [MeV]        ",
                         6 => "scaled free energy per baryon F/m_n-1              []           ",
                         7 => "scaled internal energy per baryon E/m_n-1          []           ",
                         8 => "scaled enthalpy energy per baryon H/m_n-1          []           ",
                         9 => "scaled free enthalpy per baryon G/m_n-1            []           ",
                         10 => "derivative dp/dn_b|E                               [MeV]        ",
                         11 => "derivative p/dE|n_b                                [fm^-3]      ",
                         12 => "square of speed of sound (c_s)^2                   []           ",
                         13 => "specific heat capacity at constant volume c_V      []           ",
                         14 => "specific heat capacity at constant pressure c_p    []           ",
                         15 => "adiabatic index Gamma                              []           ",
                         16 => "expansion coefficient at constant pressure alpha_p [MeV^-1]     ",
                         17 => "tension coefficient at constant volume beta_V      [fm^-3]      ",
                         18 => "isothermal compressibility kappa_T                 [MeV^-1 fm^3]",
                         19 => "isentropic compressibility kappa_S                 [MeV^-1 fm^3]",
                         20 => "free energy per baryon F                           [MeV]        ",
                         21 => "internal energy per baryon E                       [MeV]        ",
                         22 => "enthalpy per baryon H                              [MeV]        ",
                         23 => "free enthalpy per baryon G                         [MeV]        ",
                         24 => "energy density epsilon                             [MeV/fm^3]   ")

# Additional thermodynamic quantities

# Derivative quantities

# Pairs

# Quadruples
# write(*,*) i2,' total number fraction Y of particle set with index     ',idx_q(i1)
# write(*,*) i2,' average mass number A_av of particle set with index    ',idx_q(i1)
# write(*,*) i2,' average proton number Z_av of particle set with index  ',idx_q(i1)
# write(*,*) i2,' average neutron number N_av of particle set with index ',idx_q(i1)

# Microscopic quantities

# Error quantities

#

function clip(range::UnitRange, mask::UnitRange)
    @assert all(==(1), step(range))
    imin = max(first(range), first(mask))
    imax = min(last(range), last(mask))
    return imin:imax
end

function linterp(x0::T, y0::T, x1::T, y1::T, x::T,
                 extrap::Bool=false) where {T<:AbstractFloat}
    if !extrap
        if x0 < x1
            @assert x0 <= x && x <= x1
        elseif x0 > x1
            @assert x1 <= x && x <= x0
        else
            # We expect the input to have a non-empty range
            @assert false
            @assert x == x0
        end
    end
    y = (x1 - x) / (x1 - x0) * y0 + (x - x0) / (x1 - x0) * y1
    if !extrap
        if y0 < y1
            @assert y0 <= y && y <= y1
        elseif y0 > y1
            @assert y1 <= y && y <= y0
        else
            @assert y == y0
        end
    end
    return y
end

function linterp(x0::Number, y0::Number, x1::Number, y1::Number, x::Number,
                 extrap::Bool=false)
    return linterp(promote(x0, y0, x1, y1, x)..., extrap)
end

function linterp0(arr::Array{T,3}, i::Int, j::Int,
                  k::Int)::T where {T<:AbstractFloat}
    return arr[i, j, k]
end
function linterp1(arr::Array{T,3}, i::Int, j::Int, k::Int,
                  x::T)::T where {T<:AbstractFloat}
    return (1 - x) * linterp0(arr, i, j, k) + x * linterp0(arr, i + 1, j, k)
end
function linterp2(arr::Array{T,3}, i::Int, j::Int, k::Int, x::T,
                  y::T)::T where {T<:AbstractFloat}
    return (1 - y) * linterp1(arr, i, j, k, x) +
           y * linterp1(arr, i, j + 1, k, x)
end
function linterp3(arr::Array{T,3}, i::Int, j::Int, k::Int, x::T, y::T,
                  z::T)::T where {T<:AbstractFloat}
    @assert 0 ≤ x ≤ 1
    @assert 0 ≤ y ≤ 1
    @assert 0 ≤ z ≤ 1
    return (1 - z) * linterp2(arr, i, j, k, x, y) +
           z * linterp2(arr, i, j, k + 1, x, y)
end

#

# de/dt > 0
function find_outliers_dedt(table_e::AbstractArray{<:Real,3})
    outliers = falses(size(table_e))
    nnb, nt, nyq = size(table_e)

    for iyq in 1:nyq, inb in 1:nnb

        # Find outliers
        prev_it = 1
        it = prev_it + 1
        while it ≤ nt
            if table_e[inb, it, iyq] ≤ table_e[inb, prev_it, iyq]
                # We found an inconsistency. Either this point it or
                # the previous point prev_it is an outlier that needs
                # to be removed.
                inbrange = clip((inb - 1):(inb + 1), 1:nnb)
                itrange = clip((prev_it - 1):(it + 1), 1:nt)
                iyqrange = clip((iyq - 1):(iyq + 1), 1:nyq)
                avg_e = mean(table_e[inbrange, itrange, iyqrange])
                # Assume that the point further from the average is
                # the outlier
                if abs(table_e[inb, prev_it, iyq] - avg_e) >
                   abs(table_e[inb, it, iyq] - avg_e)
                    # Disable the previous point
                    outliers[inb, prev_it, iyq] = true
                    # Next, compare the current point to the one
                    # before the previous point
                    prev_it -= 1
                    while 1 ≤ prev_it && outliers[inb, prev_it, iyq]
                        prev_it -= 1
                    end
                    if 1 ≤ prev_it
                        # keep it
                    else
                        # All points before this are disabled: move on to
                        # the next point
                        prev_it = it
                        it += 1
                    end
                else
                    # Disable the current point
                    outliers[inb, it, iyq] = true
                    # keep prev_it
                    it += 1
                end
            else
                # There was no inconsistency. Move on to the next
                # point pair
                prev_it = it
                it += 1
            end
        end

        # Check consistency for all non-outliers

        # Find the first non-disabled point
        prev_it = 1
        while prev_it ≤ nt && outliers[inb, prev_it, iyq]
            prev_it += 1
        end
        # Check consistency for all following points
        for it in (prev_it + 1):nt
            if !outliers[inb, it, iyq]
                @assert table_e[inb, it, iyq] > table_e[inb, prev_it, iyq]
                prev_it = it
            end
        end
    end

    return outliers
end

function correct_outliers_dedt(table_e::AbstractArray{<:Real,3},
                               outliers::AbstractArray{Bool,3})
    @assert size(table_e) == size(outliers)
    table_e_new = copy(table_e)
    nnb, nt, nyq = size(table_e)

    for iyq in 1:nyq, it in 1:nt, inb in 1:nnb
        if outliers[inb, it, iyq]
            # Find two valid neighbouring points
            extrap = false
            prev_it = it - 1
            while 1 ≤ prev_it && outliers[inb, prev_it, iyq]
                prev_it -= 1
            end
            next_it = it + 1
            while next_it ≤ nt && outliers[inb, next_it, iyq]
                next_it += 1
            end
            if prev_it < 1
                extrap = true
                prev_it = next_it
                next_it += 1
                while next_it ≤ nt && outliers[inb, next_it, iyq]
                    next_it += 1
                end
            end
            if next_it > nt
                extrap = true
                next_it = prev_it
                prev_it -= 1
                while 1 ≤ prev_it && outliers[inb, prev_it, iyq]
                    prev_it -= 1
                end
            end
            @assert 1 ≤ prev_it ≤ nt
            @assert 1 ≤ next_it ≤ nt
            # Linear interpolation
            table_e_new[inb, it, iyq] = linterp(prev_it,
                                                table_e[inb, prev_it, iyq],
                                                next_it,
                                                table_e[inb, next_it, iyq], it,
                                                extrap)
        end
    end

    return table_e_new
end

#

function main()
    filename = "$(ENV["HOME"])/data/eos31/eos/eoscompose.h5"
    filename2 = "$(ENV["HOME"])/data/eos31/eos/eoscompose_corrected.h5"

    file = h5open(filename, "r")

    # Metadata
    date = Date(read(attributes(file["metadata"])["date"]), dateformat"y/m/d")
    time = Time(read(attributes(file["metadata"])["time"]),
                dateformat"H\hM\mS\ss")

    # Parameters
    pointsnb = read(attributes(file["Parameters"])["pointsnb"])[]
    pointst = read(attributes(file["Parameters"])["pointst"])[]
    pointsyq = read(attributes(file["Parameters"])["pointsyq"])[]
    tabulation_scheme = read(attributes(file["Parameters"])["tabulation_scheme"])[] # 0 explicit, 1 loop

    @assert pointsnb >= 0
    @assert pointst >= 0
    @assert pointsyq >= 0

    nb = read(file["Parameters"]["nb"])
    t = read(file["Parameters"]["t"])
    yq = read(file["Parameters"]["yq"])

    @assert length(nb) == pointsnb
    @assert length(t) == pointst
    @assert length(yq) == pointsyq

    @assert all(nb[i] > nb[i - 1] for i in 2:pointsnb)
    @assert all(t[i] > t[i - 1] for i in 2:pointst)
    @assert all(yq[i] > yq[i - 1] for i in 2:pointsyq)

    # Thermodynamic quantities
    pointsqty = read(attributes(file["Thermo_qty"])["pointsqty"])[]
    index_thermo = read(file["Thermo_qty"]["index_thermo"])
    thermo = read(file["Thermo_qty"]["thermo"])

    @assert pointsqty >= 0
    @assert length(index_thermo) == pointsqty
    @assert length(Set(index_thermo)) == length(index_thermo)
    @assert all(idx ∈ keys(thermo_qtys) for idx in index_thermo)

    # This order is hard-coded into `hdf5compose.f90`. It is different
    # from the order given in `eos.parameters`.
    @assert size(thermo) == (pointsnb, pointst, pointsyq, pointsqty)
    @info "Read EOS table with ($pointsnb, $pointst, $pointsyq) entries for $pointsqty quantities"

    inverse_index_thermo = Dict(idx => n
                                for (n, idx) in enumerate(index_thermo))

    # Ensure dE/dT > 0

    qty_nb(inb, it, iyq) = nb[inb] # baryon number density n_b
    qty_t(inb, it, iyq) = t[it]    # temperature T
    qty_yq(inb, it, iyq) = yq[iyq] # hadronic charge fraction Y_q
    ne = 21                        # internal energy per baryon E
    ie = inverse_index_thermo[ne]
    qty_e(inb, it, iyq) = thermo[inb, it, iyq, ie]

    @info "Looking for outliers..."
    table_e = @view thermo[:, :, :, ie]
    outliers = find_outliers_dedt(table_e)
    npoints = length(outliers)
    noutliers = count(outliers)
    @info "dE/dT ≤ 0 at $noutliers of $npoints points ($(100 * noutliers / npoints)%)"

    @info "Correcting outliers..."
    table_e_new = correct_outliers_dedt(table_e, outliers)
    outliers_new = find_outliers_dedt(table_e_new)
    any(outliers_new) && error("Could not correct outliers")
    @info "Corrected all outliers"

    # Save result
    @info "Writing result..."

    # Copy HDF5 file
    cp(filename, filename2; force=true)
    # Update thermodynamics table
    file2 = h5open(filename2, "r+")
    file2["Thermo_qty"]["thermo"][:, :, :, ie] = table_e
    close(file2)

    @info "Done."
    return nothing
end

end


module BoundaryConditions

export
    BCType, Flux, Gradient, Value, NormalFlow,

    BoundaryCondition, bctype, getbc, setbc!,

    PeriodicBoundaryCondition, NormalFlowBoundaryCondition, NoFluxBoundaryCondition,
    FluxBoundaryCondition, ValueBoundaryCondition, GradientBoundaryCondition,

    CoordinateBoundaryConditions,

    FieldBoundaryConditions, UVelocityBoundaryConditions, VVelocityBoundaryConditions,
    WVelocityBoundaryConditions, TracerBoundaryConditions, PressureBoundaryConditions,

    DiffusivityBoundaryConditions,

    apply_x_bcs!, apply_y_bcs!, apply_z_bcs!,

    fill_halo_regions!

using CUDA
using KernelAbstractions

using Oceananigans.Architectures: device
using Oceananigans.Utils: work_layout, launch!
using Oceananigans.Grids

include("boundary_condition_types.jl")
include("boundary_condition.jl")
include("discrete_boundary_function.jl")
include("continuous_boundary_function.jl")
include("coordinate_boundary_conditions.jl")
include("field_boundary_conditions.jl")
include("show_boundary_conditions.jl")

include("fill_halo_regions.jl")
include("fill_halo_regions_value_gradient.jl")
include("fill_halo_regions_normal_flow.jl")
include("fill_halo_regions_periodic.jl")
include("fill_halo_regions_flux.jl")
include("fill_halo_regions_nothing.jl")

include("apply_flux_bcs.jl")

end


using BayesianSumProductNetworks
using StatsBase
using StatsFuns
using AxisArrays
using LinearAlgebra
using Random
using DelimitedFiles
using JLD2
using FileIO
using ArgParse
using Logging
using SparseArrays
using Dates
using HDF5
using Serialization

include("helper.jl")

s = ArgParseSettings()
@add_arg_table s begin
    "--alpha"
        help = "Concentration parameter of sum nodes"
        default = 1.0
        arg_type = Float64
    "--beta"
        help = "Concentration parameter of product nodes"
        default = 10.0
        arg_type = Float64
    "--gamma"
        help = "Concentration parameter of leaf distributions"
        default = 0.1
        arg_type = Float64
    "--Ksum"
        help = "Number of children under a sum node"
        default = 2
        arg_type = Int
    "--Kprod"
        help = "Number of children under a product node"
        default = 2
        arg_type = Int
    "--Kleaf"
        help = "Number of leaves in an atomic region"
        default = 2
        arg_type = Int
    "--J"
        help = "Number of partitions."
        default = 2
        arg_type = Int
    "--L"
        help = "Number of layers."
        default = 2
        arg_type = Int
    "--seed"
        help = "Random seed = seed * Int(first(dataset)) * chain."
        default = 42
        arg_type = Int
    "--chain"
        help = "Chain"
        default = 1
        arg_type = Int
    "--maxcategories"
        default = 20
        arg_type = Int
    "--numsamples"
        arg_type = Int
        default = 1000
        help = "Number of MCMC samples we want to store."
    "dataset"
        help = "Dataset name, assuming naming convention of 20 terrible datasets."
        required = true
        arg_type = String
    "datafolder"
        help = "Path to folder containing datasets."
        required = true
        arg_type = String
    "outputfolder"
        help = "Path to folder in which output files should be stored."
        required = true
        arg_type = String
    "burnin"
        help = "Number of burnin iterations."
        required = true
        arg_type = Int
    "iterations"
        help = "Number of iterations."
        required = true
        arg_type = Int
end

args = parse_args(s)

# Load the dataset
dataset = args["dataset"]
datafolder = args["datafolder"]

# results dir
resultsdir = args["outputfolder"]
if !isdir(resultsdir)
    mkpath(resultsdir)
end

@info "Processing $dataset in $datafolder"
@info "Storing results in $resultsdir"

donefile = joinpath(resultsdir, "done")
if isfile(donefile)
    @warn "Sampling already done."
    exit()
end

# check if results are locked and exit if this is the case
lockfile = joinpath(resultsdir, "filelock.lck")

# Output files.
modelfile = joinpath(resultsdir, "bspn_model.jld2")
burninfile = joinpath(resultsdir, "bspn_burnin.jld2")
llhfile = joinpath(resultsdir, "llhfile.jld2")
tracefile = joinpath(resultsdir, "tracefile.jld2")
iterfile = joinpath(resultsdir, "iteration")

# Check if we should continue the sampling...
if (isfile(modelfile) || isfile(burninfile)) && isfile(iterfile)
    d = now() - Dates.unix2datetime(stat(iterfile).mtime)
    if round(d, Minute) > Minute(60)
        # remove lock file
        if isfile(lockfile)
            rm(lockfile)
        end
    end
end

# Check if we should continue the sampling...
if isfile(lockfile)
    @info "MCMC sampling is locked, exiting."
    exit(-2)
else

    t = rand(1:15) + rand(1:15)
    sleep(t)

    if isfile(lockfile)
        @info "MCMC sampling is locked, exiting."
        exit(-2)
    end

    touch(lockfile)
end

# Write out the configuration of the run.
save(joinpath(resultsdir, "configuration.jld2"), Dict("configuration" => args))

# Start script.
open(joinpath(resultsdir, "log.txt"), "w+") do io
    logger = SimpleLogger(io)
    with_logger(logger) do
        if Threads.nthreads() == 1
            @info "Run export JULIA_NUM_THREADS=nproc before running julia to use multithreading."
        else
            @info "Using $(Threads.nthreads()) threads."

            # Test if all threads are available.
            thrsa = zeros(Threads.nthreads())
            Threads.@threads for t in 1:Threads.nthreads()
                thrsa[t] = Threads.threadid()
            end

            avThreads = length(unique(thrsa[thrsa .> 0]))
            @assert avThreads == Threads.nthreads() "Only $(avThreads) threads available."
        end

        @info "Loading dataset: $dataset stored under: $datafolder"

        x_train = map(Float64, h5read(joinpath(datafolder, dataset*"PP", "train.hdf5"), "train"))
        x_train = transpose(x_train)

        x_valid = map(Float64, h5read(joinpath(datafolder, dataset*"PP", "valid.hdf5"), "valid"))
        x_train = cat(x_train, transpose(x_valid), dims = 1)

        x_test = map(Float64, h5read(joinpath(datafolder, dataset*"PP", "test.hdf5"), "test"))
        x_test = transpose(x_test)

        # Set the number of iterations, burnin, etc.
        iterations = args["iterations"]
        burnin = args["burnin"]

        # ######### #
        # MCMC code #
        # ######### #
        ks = min(iterations, args["numsamples"])
        @assert isinteger(iterations / ks)
        thinning = Int(iterations / ks)

        (N_train, D) = size(x_train)
        (N_test, _) = size(x_test)

        @info "Processing dataset of size N=$N_train, D=$D."

        # Preprocessing
        dids = findall(map(d -> all(isinteger, filter(x -> !isnan(x), x_train[:,d])), 1:D))
        dids = filter(d -> length(filter(x -> !isnan(x), unique(x_train[:,d]))) < args["maxcategories"], 1:D)
        for d in dids
            ix = findall(!isnan, x_train[:,d])
            l = unique(x_train[ix,d])

            x_train[ix,d] = map(x -> findfirst(l .== x), x_train[ix,d])

            ix = findall(!isnan, x_test[:,d])
            x_test[ix,d] = map(x -> x ∈ l ? findfirst(l .== x) : length(l)+1, x_test[ix,d])

            if length(l) == 2
                # Make binary.
                x_train .-= 1
                x_test .-= 1
            end
        end

        priors = Vector{Vector{Distribution}}(undef, D)
        likelihoods = Vector{Vector{Distribution}}(undef, D)
        sstats = Vector{Vector{AbstractSufficientStats}}(undef, D)

        for d in 1:D
            priors[d] = Vector{Distribution}()
            likelihoods[d] = Vector{Distribution}()
            sstats[d] = Vector{Distribution}()
            ix = findall(!isnan, x_train[:,d])

            if d ∈ dids # is discrete ?
                if minimum(x_train[ix,d]) == 0
                    # Binary data
                    push!(priors[d], Beta(1/2, 1/2))
                    push!(likelihoods[d], Bernoulli())
                    push!(sstats[d], BetaSufficientStats())
                else
                    # Categorical data
                    l = unique(x_train[:,d])
                    push!(priors[d], Dirichlet(length(l), 0.1))
                    push!(likelihoods[d], Categorical(rand(last(priors[d]))))
                    push!(sstats[d], DirichletSufficientStats(length(l)))
                end

                # Poisson data
                push!(priors[d], Gamma())
                push!(likelihoods[d], Poisson())
                push!(sstats[d], GammaSufficienStats())
            else # continuous data
                if all((x) -> x >= 0, x_train[ix,d])
                    # Positive continuous data
                    push!(priors[d], Gamma())
                    push!(likelihoods[d], Exponential())
                    push!(sstats[d], GammaSufficienStats())
                end

                # Normal distributed data
                push!(priors[d], NormalInverseGamma(mean(x_train[ix,d]), 1.0/std(x_train[ix,d]), 2.0, 3.0))
                push!(likelihoods[d], Normal())
                push!(sstats[d], NormalInvGammaSufficientStats())
            end
        end

        # Set the random seed.
        rnd_seed = args["seed"]+args["chain"]
        rng = Random.seed!(rnd_seed)

        # Construct a second network which we use for sampling.
        @info "Constructing spn for sampling."
        spn = templateRegion(args["alpha"],
                             args["beta"],
                             args["gamma"],
                             priors,
                             likelihoods,
                             sstats,
                             N_train, D,
                             args["Ksum"], args["Kprod"],
                             args["J"], args["Kleaf"],
                             1, args["L"], root = true)

        @info "Starting MCMC for $iterations iterations."

        startIter = 1
        storePredictions = true
        storeTrace = false

        # Initialize SPN.
        if isfile(modelfile)
            # Resume.
            spn = open(modelfile, "r") do fin
                deserialize(fin)
            end
            r_ = read(iterfile, Int)
            startIter = 1 + r_
            @info "Continue MCMC."
        elseif isfile(burninfile)
            spn = open(burninfile, "r") do fin
                deserialize(fin)
            end
            @info "Continue burnin or start using previously converged chain, not storing trace."
            r_ = read(iterfile, Int)
            startIter = 1 + r_
        else
            @warn "Starting from random position, not storing predictions!"
            randomscopes!(spn, x_train)
        end

        for iteration in startIter:(burnin + iterations)
            dateprint = Dates.format(now(), "HH:MM:SS")

            storeTrace = iteration <= burnin
            storePredictions = iteration > burnin

            @info "[$(dateprint)] iteration = $(iteration) "

            ancestralsampling!(spn, x_train)
            gibbssamplescopes!(spn, x_train)

            if ((iteration-burnin) % thinning) == 0
                # Serialize model.
                fname = joinpath(resultsdir, "tmp_"*randstring()*".jld2")
                open(fname, "w") do fout
                    serialize(fout, spn)
                end
                if iteration > burnin
                    mv(fname, modelfile, force=true)
                else
                    mv(fname, burninfile, force=true)
                end

                if storeTrace
                    # Store parameters of root.
                    Nparam = length(spn.logweights[:])
                    θ = isfile(tracefile) ? load(tracefile, "weights") : zeros(0,Nparam)
                    fname = joinpath(resultsdir, "tmp_"*randstring()*".jld2")
                    save(fname, "weights", vcat(θ, spn.logweights[:]'))
                    mv(fname, tracefile, force=true)
                end

                if storePredictions
                    # Evaluate network.
                    lp_test = logpdf(spn, x_test)
                    lp_train = logpdf(spn, x_train)

                    b_ = isfile(llhfile)
                    llh_test = b_ ? load(llhfile, "test") : zeros(0,N_test)
                    llh_train = b_ ? load(llhfile, "train") : zeros(0,N_train)
                    fname = joinpath(resultsdir, "tmp_"*randstring()*".jld2")
                    save(fname, "train", vcat(llh_train, lp_train'), "test", vcat(llh_test, lp_test'))
                    mv(fname, llhfile, force=true)
                end

                write(iterfile, iteration)
            end
        end
        # Save sampling results to disk.
        @info "Finished MCMC"

        # Create done file and release lock.
        touch(donefile)
        rm(lockfile)

        tmp_files = filter(f -> isfile(joinpath(resultsdir,f)) && startswith(f, "tmp_"), readdir(resultsdir))
        for f in tmp_files
            rm(joinpath(resultsdir,f))
        end
    end
end


include("debug.jl")
include("fpump.jl")

function create_root_model!(optimizer::MOI.AbstractOptimizer, jp::JuniperProblem)
    ps = jp.options.log_levels

    jp.model = Model()
    lb = jp.l_var
    ub = jp.u_var
    # all continuous we solve relaxation first
    @variable(jp.model, lb[i] <= x[i=1:jp.num_var] <= ub[i])

    for i=1:jp.num_var
        JuMP.set_start_value(x[i], jp.primal_start[i])
    end

    if jp.options.registered_functions !== nothing
        for reg_f in jp.options.registered_functions
            if reg_f.gradf === nothing
                JuMP.register(jp.model, reg_f.s, reg_f.dimension, reg_f.f; autodiff=reg_f.autodiff)
            elseif reg_f.grad2f === nothing
                JuMP.register(jp.model, reg_f.s, reg_f.dimension, reg_f.f, reg_f.gradf)
            else
                JuMP.register(jp.model, reg_f.s, reg_f.dimension, reg_f.f, reg_f.gradf, reg_f.grad2f)
            end
        end
    end

    # TODO check whether it is supported
    if optimizer.nlp_data.has_objective
        obj_expr = MOI.objective_expr(optimizer.nlp_data.evaluator)
        expr_dereferencing!(obj_expr, jp.model)
        JuMP.set_NL_objective(jp.model, optimizer.sense, obj_expr)
    elseif optimizer.objective !== nothing
        MOI.set(jp.model, MOI.ObjectiveFunction{typeof(optimizer.objective)}(), optimizer.objective)
        MOI.set(jp.model, MOI.ObjectiveSense(), optimizer.sense)
    end

    backend = JuMP.backend(jp.model);
    llc = optimizer.linear_le_constraints
    lgc = optimizer.linear_ge_constraints
    lec = optimizer.linear_eq_constraints
    qlc = optimizer.quadratic_le_constraints
    qgc = optimizer.quadratic_ge_constraints
    qec = optimizer.quadratic_eq_constraints
    for constr_type in [llc, lgc, lec, qlc, qgc, qec]
        for constr in constr_type
            MOI.add_constraint(backend, constr[1], constr[2])
        end
    end
    for i in 1:jp.num_nl_constr
        constr_expr = MOI.constraint_expr(optimizer.nlp_data.evaluator, i)
        expr_dereferencing!(constr_expr, jp.model)
        JuMP.add_NL_constraint(jp.model, constr_expr)
    end

    (:All in ps || :Info in ps) && print_info(jp)

    jp.x = x
end

function solve_root_model!(jp::JuniperProblem)
    status, backend = optimize_get_status_backend(jp.model; solver=jp.nl_solver)
    jp.relaxation_status = status
    restarts = 0
    max_restarts = jp.options.num_resolve_root_relaxation
    jp.options.debug && debug_init(jp.debugDict)
    while !state_is_optimal(jp.relaxation_status; allow_almost=jp.options.allow_almost_solved) &&
        restarts < max_restarts && time()-jp.start_time < jp.options.time_limit

        # TODO freemodel for Knitro
        restart_values = generate_random_restart(jp)
        jp.options.debug && debug_restart_values(jp.debugDict,restart_values)
        for i=1:jp.num_var
            JuMP.set_start_value(jp.x[i], restart_values[i])
        end
        JuMP.optimize!(jp.model)
        jp.relaxation_status = MOI.get(backend, MOI.TerminationStatus())
        restarts += 1
    end
    if only_almost_solved(jp.relaxation_status)
        @warn "The relaxation is only almost solved."
    end
    return restarts
end