using Test
using StatsModels

@testset "Grouping" begin
    g = Grouping()
    @test isnothing(g.levels)
end

@testset "Grouping pseudo-contrasts" begin
    d = (y = rand(2_000_000), grp=string.([1:1_000_000; 1:1_000_000]))
    ## OOM seems to result in the process being killed on Mac so this messes up CI
    # @test_throws OutOfMemoryError schema(d)
    sch = schema(d, Dict(:grp => Grouping()))
    t = sch[term(:grp)]
    @test t isa CategoricalTerm{Grouping}
    @test size(t.contrasts.matrix) == (0,0)
    @test length(t.contrasts.levels) == 1_000_000
    @test_throws ErrorException StatsModels.modelcols(t, (a = 1.,))

    levs = sort(string.(1:1_000_000))

    @test all(t.contrasts.invindex[lev] == i for (i,lev) in enumerate(levs))
    @test all(t.contrasts.levels[i] == lev for (i,lev) in enumerate(levs))

    # @test_throws OutOfMemoryError fit(MixedModel, @formula(y ~ 1 + (1 | grp)), d)
    @test fit(MixedModel, @formula(y ~ 1 + (1 | grp)), d, contrasts=Dict(:grp => Grouping()), progress=false) isa LinearMixedModel
end


using Clingo
x,y = Clingo.testsolve("#const n=4. c(1..n). 1 {color(X,I) : c(I)} 1 :- v(X). :- color(X,I), color(Y,I), e(X,Y), c(I). v(1..100). e(1,55).")


############################################################################
#
#  Software.jl
#
#  Purpose:    Software file to control a quadrotor
#
#  Author:     Edward Daughtery
#
#  Date:       05 August 2016
#
#
############################################################################
include("Guidance.jl")

#######################################################################
# Intialization
function Initialize_SW(time::Float64, DT::Float64, nav_state::NavStates)
    zz = 0.0;
    zi = UInt(0);
    zvec = [0.0;0.0;0.0];
    dt = DT;
    count = UInt(0);
    mode = UInt(0);

    ## Controls Initialization
    throttle = 0.0;
    alpha = 0.0;
    alphadot = 0.0;
    phi = 0.0;
    phidot = 0.0;
    controls = Controls(throttle,alpha,alphadot,phi,phidot);

    ## Guidance Initialization
#    num_waypoints = 3;
#    GuidInit(wyptlist, num_waypoints)

    ## Software Message
    tgo = 999999.0;
    missvec = [9999999.0;9999999.0;9999999.0];
    miss = norm(missvec);
    swMsg = SoftwareMsg( time, dt, mode, controls, nav_state, zi, zz, zz, zz, zz, zz, tgo, miss, missvec, zz );

    return swMsg;
end

#######################################################################
# Flight Control
function FlightControl( softwareMsg::SoftwareMsg, apcmd::Float64, phicmd::Float64 )

    ## Software Telemetry Variable
    swOut = softwareMsg
    time = swOut.time;
    dt =  swOut.dt;

    ## Control Variables
    alpha = swOut.controls.alpha;
    alphadot = swOut.controls.alphadot;
    phi = swOut.controls.phi;
    phidot = swOut.controls.phidot;
    throttle = swOut.controls.throttle;

    ## IMU Measurements
    axmeas = swOut.navstates.accel_meas[1];
    aymeas = swOut.navstates.accel_meas[2];
    azmeas = swOut.navstates.accel_meas[3];

    # Alpha Gains
    Kpa = 0.20;
    Kpad = 0.3;

    # Roll Gain
    Kr = 0.30;

    # Z Acceleration Error
    alphacmd = -Kpa*(apcmd  - azmeas);

    # Alpha Response
    alphadot = Kpad*(alphacmd - alpha);
    tmpalpha = alpha + dt*alphadot;

    # Limits on the Angle of Attack
    if (tmpalpha*57.3 < -6.0 )
        alpha = -6.0/57.3;
    elseif (tmpalpha*57.3 > 9.0 )
        alpha = 9.0/57.3;
    else
        alpha = tmpalpha;
    end

    # Phi Response
    phidot = Kr*(phicmd - phi);
    tmpphi = phi + dt*phidot;

    # Limits on the Roll Angle
    if ( tmpphi*57.3 > 80.0)
        phi = 80.0/57.3;
    elseif ( tmpphi*57.3 < -80.0)
        phi = -80.0/57.3;
    else
        phi = tmpphi;
    end

    swOut.controls.alpha = alpha;
    swOut.controls.alphadot = alphadot;
    swOut.controls.alphadot = phi;
    swOut.controls.alphadot = phidot;

    return swOut
end

#######################################################################
function Steering(nav_state::NavStates, acc_cmd_ned::Array{Float64,1})

    # Computing NED to Body DCM
    Tr_BN = nav_state.euler.Tr_bn;
    roll = nav_state.euler.roll;
    pitch = nav_state.euler.pitch;
    heading = nav_state.euler.heading;

    # Computing Non-Rolled NED To Body DCM
    Tr_BN_NR = TR_BN(0.0, pitch, heading);

    # Acceleration Commands Body
    acc_cmd_body = Tr_BN*acc_cmd_ned;
    acc_cmd_bodyNR = Tr_BN_NR*acc_cmd_ned;

    # Acceleration Commands Non Rolled
    ay_cmd_ned_nr = acc_cmd_bodyNR[2];
    az_cmd_ned_nr = acc_cmd_bodyNR[3];

    # Zero Denominator Protection
    if ( abs(az_cmd_ned_nr) < 0.00000001 )
        phicmd = 0.5*eta_turn_cmd;
    else
        phicmd = atan(ay_cmd_ned_nr/abs(az_cmd_ned_nr + 0.001) );
    end

    # Phi Limiting
    philim = 60.0*d2r;
    phicmd = Bound(phicmd, -philim, philim);

    # Pitch Acceleration Command
    apcmd = acc_cmd_body[3];

    return  [apcmd; phicmd];
end

#######################################################################
function SpeedControl(nav_state::NavStates, vcmd_coop::Float64, softwareMsg::SoftwareMsg )

    ## Parameter Initialization
    vm_ned = nav_state.velocity.vm_ned;
    gamma = nav_state.velocity.gamma;
    alt = nav_state.wgs_pos.altitude;
    Vmag = nav_state.velocity.vmag;
    vel_err_int = softwareMsg.vel_err_int;
    dt50 = softwareMsg.dt;
    Throttle = softwareMsg.controls.throttle;

    ## Atmosphere Model
    atmdata = Atmosphere(alt, Vmag);
    Press = atmdata[1];
    Temp = atmdata[2];
    rho = atmdata[3];
    aa = atmdata[4];
    Mach = atmdata[5];
    qbar = atmdata[6];

    ## Mach Command converted to speed command
    machcmd = 0.7;
    vcmd = machcmd*aa + vcmd_coop;

    ## Velocity Error and Integrator
    verr = (vcmd - Vmag);
    tmp2 = vel_err_int + dt50*verr;

    ## PI Control to compute Throttle Command
    Kpv = 1.5;
    Kiv = 0.66;
    tmp = Kpv*verr + Kiv*tmp2;

    if ( tmp > 100.0 )
        tmp = 100.0;
    elseif ( tmp < 0.0)
        tmp = 0.0;
    else
        Throttle = tmp;
        vel_err_int = tmp2;
    end

    ## Bounding Throttle to be between 0 and 100
    Throttle = Bound(Throttle, 0.0, 100.0);

    return Throttle;
end

#######################################################################
# Software Update
function Update_SW(swMsg::SoftwareMsg)

    swOut = swMsg;
    alpha  = swOut.controls.alpha;
    phi  = swOut.controls.phi;
    gamma  = swOut.navstates.velocity.gamma;
    azimuth  = swOut.navstates.velocity.azimuth;

    ## Update
    swOut.time = swOut.time + swOut.dt;

    ## Acceleration Commands
    acc_cmd_ned = zeros(3);

    ## Guidance Section
    alt_cmd = 4000.0;
    eta_turn_cmd = 0.0;
    eta_dive_cmd = AltitudeControl(swOut.navstates, alt_cmd);
    println(eta_dive_cmd)

    ## Velocity Commands in Steering Frame (Velocity Frame)
    acc_cmd_vel = zeros(3);
    acc_cmd_vel[1] = 0.0;
    acc_cmd_vel[2] = eta_turn_cmd;
    acc_cmd_vel[3] = eta_dive_cmd;

    Tr_vn = TR_VN(gamma, azimuth);
    Tr_nv  = Tr_vn';
    acc_cmd_ned = Tr_nv*acc_cmd_vel;
    acc_cmd_ned[3] = acc_cmd_ned[3]  - 9.80665;

    ## Steering Function
    steerout = Steering(swOut.navstates, acc_cmd_ned);
    apcmd = steerout[1];
    phicmd = steerout[2];

    ## Flight Controls
    throttle = SpeedControl(swOut.navstates, 0.0, swOut);
    FlightControl( swOut, apcmd, phicmd );

    swOut.controls.throttle = throttle;
    return swOut;
end


push!(LOAD_PATH, pwd())
using NQCDynamics
using NQCDynamics.InitialConditions
using LinearAlgebra
using Unitful
using UnitfulAtomic
using Revise
using Random, Distributions
using Plots
using DelimitedFiles
#using Profile
# using DataFrames
#using BenchmarkTools

# IESH should proceed by the following steps:
#   1) initialize classical positions and momenta, initialize
#      electronic states, define surface on which to propagate
#   2) Do Verlet algorithm for the atomic positions, velocities electronic
#   3) Integrate the electronic wave function w/ 1e- Schroedinger equationa
#   4) calculate switching probability, calculate if hop should occur
#   5) If no hop, GOTO 2; if hop, see if enough energy to do it, GOTO 2

ps = 98.22694788464062
Random.seed!(13)
#n_states = 202
n_states = 42
#n_states = 502
# # number of electronic n_states
vinit = 0.005 # initial velocity
mass = 1.5 # atomic mass#
B_na = 1.0 # parameter that defines nonadiabatic coupling strength
n_DOF = 1
ntrajes = 20
tspan = (0.0, 5.0u"ps")

# 1a) Define adiabatic model, fom Model/scattering_anderson_holsteins.jl
# model = Models.ScatteringAndersonHolstein(N=n_states, a0=1, Bg=B_na,
# Cg=1, alpha=-5.0, beta=.5)
#model = Models.MiaoSubotnik(n_states=n_states, W=0.064, Gamma=0.0064)
model = Models.MiaoSubotnik(n_states=n_states, W=0.064, Gamma=0.0064)
#model = Models.TullyModelOne()
#println(model.n_states)

# 1b) Initialize atomic parameters, i.e. moving atoms, of which there is one
atoms = Atoms{Float64}([:H])

##############################################################################
#
# Try dynamics
#
###############################################################################

# create Boltzmann distribution of momenta
# kT = 9.5*10^-4 a.u. = 300 K
temperature = 5*298u"K"

# create distribution of atomic positions from Monte Carlo sampling of the PESs
# for the number of trajectories and atoms
# atoms and step size
Δ = Dict([(:H, 0.3)])
tsim = Simulation{Classical}(atoms, model; temperature=temperature)
R0 = zeros(3, length(tsim.atoms))
Δ = Dict([(:H, 0.3)])
R0 = zeros(1, length(tsim.atoms))
passes = 1000
outbolz = MetropolisHastings.run_monte_carlo_sampling(tsim, R0, Δ, passes)

# If sampled, gives an array of momenta and positions from the above initialized sampling
#bolz_pos_momenta = DynamicalDistribution(boltzmann, outbolz.R, (1, 1))

# # Initialize the simulation problem; Simulation is defined in src/simulation_constructors.jl
sim = Simulation{wave_IESH}(atoms, model)

# # Do dynamics
for i=1:ntrajes
    #i = 1
     nname=lpad(i,8,"0")
#     r, v = rand(bolz_pos_momenta)
#     v = -v
#     println(i," ", r, " ", v)

    #r = fill(-5.0, sim.DoFs, length(sim.atoms))
    r = rand(outbolz.R)
    v = fill(sqrt(rand(outbolz.energy)*2/sim.atoms.masses[1]), size(sim))
    #v = fill(5/sim.atoms.masses[1], sim.DoFs, length(sim.atoms))
    println(r, v)

    # intial state (how to initialize a vector that shows which states are occupied)
    # For now, just initialize in ground state. Makes sense, anyhow, when the molecule
    # is hurling towards the surface, inititially
    # This is defined in src/Dynamics/SurfaceHopping/surface_hopping_variables.jl
    # as: SurfaceHoppingVariables(v::AbstractArray, r::AbstractArray, n_states::Integer, state::Integer)
    k = round.(Int64, (zeros(n_states)))
    l = Int(n_states / 2)
    k[1:l] .= 1
    println(k)
    z = SurfaceHoppingVariablesIESH(v, r, n_states, k)

    #run_trajectory defined in src/Dynamics.jl
    # callbacks defined in: src/Dynamics/callbacks.jl and src/Models/Models.jl
    # Save impurity does not seem to be defined at the moment?
    #@time solution = Dynamics.run_trajectory(z, (0.0, 10000000.0), sim; output=(:save_impurity))
    @time solution = Dynamics.run_trajectory(z, (0.0, 100000.0), sim, dt=10, adaptive=false;
                                             output=(:position, :save_impurity))
    println("Finished")
    println(length(solution.t))

    aka = Int(ceil(length(solution)/10))-1
    j = 0
    #aka = Int(length(solution)
    println(aka)
    outarray = zeros(aka, Int(5 + model.n_states))
    open("trajectory_$nname.txt", "w") do fi
        write(fi, "step, r (a.u.), epot (a.u.), state, impurity population\n")
        for k = 1:length(solution)
            if (mod(k,10)==0)
                j = j + 1
                outarray[j,1] = solution.t[k]
                outarray[j,2] = solution.position[k][1]
                outarray[j,3] = solution.save_impurity[k][2]
                outarray[j,4] = solution.save_impurity[k][3]
                outarray[j,5] = solution.save_impurity[k][4]
                for imp in 1:model.n_states
                    outarray[j,5+imp] = solution.save_impurity[k][4+imp]
                end

            # outarray[i,3] = solution.save_impurity[i][2]
            # outarray[i,4] = solution.save_impurity[i][3]
            # outarray    [i,5] = solution.save_impurity[i][4]
            end
        end
        writedlm(fi, outarray,'\t')
    end

# b = plot(outarray[:,1], outarray[:,2], label="energy", marker=2)
# xlabel!("t")
# ylabel!("x")
a = plot(outarray[:,2], outarray[:,3], label="energy", marker=2)
xlabel!("x")
ylabel!("Energy (a.u.)")
b = plot(outarray[:,1], outarray[:,Int(5+42/2)], label="energy", marker=2)
xlabel!("time")
ylabel!("Impurity Population")

display(plot(b))
display(plot(a))

end
# # # savefig(a,"traj_subA_G4Em4_W10G_test5.png")
# # # savefig(b,"traj_imppop_G4Em4_W10G_test2.png")

#display(plot(b))


export apply_back!, apply_back

const Rotor{D,T} = Union{RotationGate{D,T},PutBlock{D,<:Any,<:RotationGate{<:Any,T}}}

as_scalar(arr::AbstractArray{T,0}) where {T} = arr[]
as_scalar(arr) = arr

"""
    generator(rot::Rotor) -> AbstractBlock

Return the generator of rotation block.
"""
generator(rot::RotationGate) = rot.block
generator(rot::PutBlock) =
    PutBlock(rot.n, generator(rot |> content), rot |> occupied_locs)

function apply_back!(st, block::AbstractBlock, collector) #,AbstractContainer{<:PrimitiveBlock}
    out, outδ = st
    if nparameters(block) == 0
        adjblock = block'
        in = apply!(out, adjblock)
        inδ = apply!(outδ, adjblock)
        return (in, inδ)
    else
        throw(MethodError(apply_back, (st, block, collector)))
    end
end

function apply_back!(st, block::Subroutine{D}, collector) where {D}
    out, outδ = st
    focus!(out, block.locs)
    focus!(outδ, block.locs)
    apply_back!((out, outδ), content(block), collector)
    relax!(out, block.locs; to_nactive = nqudits(block))
    relax!(outδ, block.locs; to_nactive = nqudits(block))
    return (out, outδ)
end

function apply_back!(st, block::Rotor{D}, collector) where {D}
    out, outδ = st
    adjblock = block'
    backward_params!((out, outδ), block, collector)
    in = apply!(out, adjblock)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::TimeEvolution, collector)
    out, outδ = st
    adjblock = block'

    out, outδ = st
    input = apply!(out, adjblock)
    for i=1:YaoArrayRegister._asint(nbatch(outδ))
        o = viewbatch(outδ, i)
        !all(x -> x ≈ 0.0im, o.state) && apply!(o, adjblock)
    end
    pushfirst!(collector, -sum(imag(input' * apply!(copy(outδ), block.H))))
    return (input, outδ)
end

function apply_back!(st, block::PutBlock, collector)
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::KronBlock, collector)
    apply_back!(st, chain(nqudits(block), [put(loc => block[loc]) for loc in block.locs]), collector)
end

function apply_back!(st, block::ControlBlock, collector)
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::Daggered, collector)
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(in, outδ)
    mat_back!(datatype(in), content(block), adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::Scale, collector)
    out, outδ = st
    # recoverting inputs
    if isunitary(block)
        in = apply!(out, block')
    elseif isunitary(content(block))
        in = apply!(regscale!(out, inv(factor(block))), content(block)')
    else
        @warn "rescaling a non-unitary block, call into inefficient back-propgation routine."
        in = apply!(out, matblock(inv(Matrix(block))))
    end
    # adjoint of matrix
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)

    inδ = apply!(outδ, block')

    # differentiate scaling
    #if niparams(block) > 0
    #    g = dropdims(sum(conj.(state(in)) .* state(outδ), dims = (1, 2)), dims = (1, 2)) |> as_scalar
    #    pushfirst!(collector, g')
    #end
    return (in, inδ)
end

function apply_back!(st, circuit::ChainBlock, collector)
    for blk in Base.Iterators.reverse(subblocks(circuit))
        st = apply_back!(st, blk, collector)
    end
    return st
end

function apply_back!(st, circuit::Add, collector; in)
    out, outδ = st
    adjmat = outerprod(outδ, in)
    for blk in Base.Iterators.reverse(subblocks(circuit))
        mat_back!(datatype(in), blk, adjmat, collector)
    end
    inδ = apply!(outδ, circuit')
    (in, inδ)
end

function apply_back!(st, block::RepeatedBlock{D,C}, collector) where {D,C}
    if nparameters(content(block)) == 0
        return apply!.(st, Ref(block'))
    end
    res = Any[]
    st = apply_back!(st, chain(nqudits(block), [put(loc => content(block)) for loc in block.locs]), res)
    res = dropdims(sum(reshape(res, :, C), dims = 2), dims = 2) |> as_scalar
    prepend!(collector, res)
    return st
end

# TODO: concentrator, repeat, kron
apply_back!(st, block::Measure, collector) =
    throw(MethodError(apply_back!, (st, block, collector)))

function backward_params!(st, block::Rotor, collector)
    in, outδ = st
    Σ = generator(block)
    g =
        dropdims(sum(conj.(state(in |> Σ)) .* state(outδ), dims = (1, 2)), dims = (1, 2)) |>
        as_scalar
    pushfirst!(collector, -imag(g) / 2)
    in |> Σ
    nothing
end

"""
    apply_back(st::Tuple{<:AbstractArrayReg, <:AbstractArrayReg}, block::AbstractBlock; kwargs...) -> (out, outδ), paramsδ

The backward function of `apply!`. Returns a tuple of ((input register, gradient of input register), parameter gradients)
"""
function apply_back(st::Tuple{<:AbstractArrayReg,<:AbstractArrayReg}, block::AbstractBlock; kwargs...)
    col = []
    in, inδ = apply_back!(st, block, col; kwargs...)
    (in, inδ), col
end


export iauAtciqz
"""
Quick ICRS to CIRS transformation, given precomputed star-
independent astrometry parameters, and assuming zero parallax and
proper motion.

Use of this function is appropriate when efficiency is important and
where many star positions are to be transformed for one date.  The
star-independent parameters can be obtained by calling one of the
functions iauApci[13], iauApcg[13], iauApco[13] or iauApcs[13].

The corresponding function for the case of non-zero parallax and
proper motion is iauAtciq.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  support function.

Given:
   rc,dc  double     ICRS astrometric RA,Dec (radians)
   astrom iauASTROM* star-independent astrometry parameters:
      pmt    double       PM time interval (SSB, Julian years)
      eb     double[3]    SSB to observer (vector, au)
      eh     double[3]    Sun to observer (unit vector)
      em     double       distance from Sun to observer (au)
      v      double[3]    barycentric observer velocity (vector, c)
      bm1    double       sqrt(1-|v|^2): reciprocal of Lorenz factor
      bpn    double[3][3] bias-precession-nutation matrix
      along  double       longitude + s' (radians)
      xpl    double       polar motion xp wrt local meridian (radians)
      ypl    double       polar motion yp wrt local meridian (radians)
      sphi   double       sine of geodetic latitude
      cphi   double       cosine of geodetic latitude
      diurab double       magnitude of diurnal aberration vector
      eral   double       "local" Earth rotation angle (radians)
      refa   double       refraction constant A (radians)
      refb   double       refraction constant B (radians)

Returned:
   ri,di  double     CIRS RA,Dec (radians)

Note:

   All the vectors are with respect to BCRS axes.

References:

   Urban, S. & Seidelmann, P. K. (eds), Explanatory Supplement to
   the Astronomical Almanac, 3rd ed., University Science Books
   (2013).

   Klioner, Sergei A., "A practical relativistic model for micro-
   arcsecond astrometry in space", Astr. J. 125, 1580-1597 (2003).

Called:
   iauS2c       spherical coordinates to unit vector
   iauLdsun     light deflection due to Sun
   iauAb        stellar aberration
   iauRxp       product of r-matrix and p-vector
   iauC2s       p-vector to spherical
   iauAnp       normalize angle into range +/- pi

This revision:   2013 October 9

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""
function iauAtciqz(rc::Real, dc::Real,  astrom::iauASTROM)

   # Allocate return value
   ref_astrom = Ref{iauASTROM}(astrom)
   ref_ri     = Ref{Float64}(0.0)
   ref_di     = Ref{Float64}(0.0)

   ccall((:iauAtciqz, libsofa_c), Cvoid,
            (Cdouble, Cdouble, Ref{iauASTROM},
            Ref{Cdouble}, Ref{Cdouble}),
            convert(Float64, rc), convert(Float64, dc),
            ref_astrom, ref_ri, ref_di)

   return ref_ri[], ref_di[]
end


using CLIMA
using CLIMA.Mesh.Topologies: StackedCubedSphereTopology, cubedshellwarp, grid1d
using CLIMA.Mesh.Grids
using CLIMA.Mesh.Filters
using CLIMA.DGmethods: DGModel, init_ode_state
using CLIMA.DGmethods.NumericalFluxes: Rusanov, CentralGradPenalty,
                                       CentralNumericalFluxDiffusive
using CLIMA.ODESolvers: solve!, gettime
using CLIMA.LowStorageRungeKuttaMethod: LSRK144NiegemannDiehlBusch
using CLIMA.VTK: writevtk, writepvtu
using CLIMA.GenericCallbacks: EveryXWallTimeSeconds, EveryXSimulationSteps
using CLIMA.MPIStateArrays: euclidean_distance
using CLIMA.PlanetParameters: R_d, grav, MSLP, planet_radius, cp_d, cv_d, day
using CLIMA.MoistThermodynamics: air_density, total_energy, soundspeed_air, internal_energy, air_temperature
using CLIMA.Atmos: AtmosModel, SphericalOrientation, NoReferenceState,
                   DryModel, NoRadiation, NoSubsidence, NoFluxBC,
                   ConstantViscosityWithDivergence,
                   vars_state, vars_aux,
                   Gravity, Coriolis,
                   HydrostaticState, IsothermalProfile
using CLIMA.VariableTemplates: flattenednames

using MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test

const output_vtk = true
const ArrayType = CLIMA.array_type()

function main()
  CLIMA.init()
  mpicomm = MPI.COMM_WORLD

  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = Dict("DEBUG" => Logging.Debug,
                  "WARN"  => Logging.Warn,
                  "ERROR" => Logging.Error,
                  "INFO"  => Logging.Info)[ll]

  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))

  polynomialorder = 5
  numelem_horz = 6
  numelem_vert = 8
  timeend = 60 # 400day
  outputtime = 2day

  for FT in (Float64,)

    run(mpicomm, polynomialorder, numelem_horz, numelem_vert,
        timeend, outputtime, ArrayType, FT)
  end
end

function run(mpicomm, polynomialorder, numelem_horz, numelem_vert,
             timeend, outputtime, ArrayType, FT)

  setup = HeldSuarezSetup{FT}()

  vert_range = grid1d(FT(planet_radius), FT(planet_radius + setup.domain_height), nelem = numelem_vert)
  topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)

  grid = DiscontinuousSpectralElementGrid(topology,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = polynomialorder,
                                          meshwarp = cubedshellwarp)

  model = AtmosModel(SphericalOrientation(),
                     HydrostaticState(IsothermalProfile(setup.T_initial), FT(0)),
                     ConstantViscosityWithDivergence(FT(0)),
                     DryModel(),
                     NoRadiation(),
                     NoSubsidence{FT}(),
                     (Gravity(), Coriolis(), held_suarez_forcing!),
                     NoFluxBC(),
                     setup)

  dg = DGModel(model, grid, Rusanov(),
               CentralNumericalFluxDiffusive(), CentralGradPenalty())

  # determine the time step
  element_size = (setup.domain_height / numelem_vert)
  acoustic_speed = soundspeed_air(FT(315))
  lucas_magic_factor = 5
  dt = lucas_magic_factor * min_node_distance(grid) / acoustic_speed

  Q = init_ode_state(dg, FT(0))
  lsrk = LSRK144NiegemannDiehlBusch(dg, Q; dt = dt, t0 = 0)

  filterorder = 14
  filter = ExponentialFilter(grid, 0, filterorder)
  cbfilter = EveryXSimulationSteps(1) do
    Filters.apply!(Q, 1:size(Q, 2), grid, filter)
    nothing
  end

  eng0 = norm(Q)
  @info @sprintf """Starting
                    ArrayType       = %s
                    FT              = %s
                    polynomialorder = %d
                    numelem_horz    = %d
                    numelem_vert    = %d
                    filterorder     = %d
                    dt              = %.16e
                    norm(Q₀)        = %.16e
                    """ "$ArrayType" "$FT" polynomialorder numelem_horz numelem_vert filterorder dt eng0

  # Set up the information callback
  starttime = Ref(now())
  cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s=false)
    if s
      starttime[] = now()
    else
      energy = norm(Q)
      runtime = Dates.format(convert(DateTime, now() - starttime[]), dateformat"HH:MM:SS")
      @info @sprintf """Update
                        simtime = %.16e
                        runtime = %s
                        norm(Q) = %.16e
                        """ gettime(lsrk) runtime energy
    end
  end
  callbacks = (cbinfo, cbfilter)

  if output_vtk
    # create vtk dir
    vtkdir = "vtk_heldsuarez" *
      "_poly$(polynomialorder)_horz$(numelem_horz)_vert$(numelem_vert)" *
      "_filter$(filterorder)_$(ArrayType)_$(FT)"
    mkpath(vtkdir)

    vtkstep = 0
    # output initial step
    do_output(mpicomm, vtkdir, vtkstep, dg, Q, model)

    # setup the output callback
    cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do
      vtkstep += 1
      Qe = init_ode_state(dg, gettime(lsrk))
      do_output(mpicomm, vtkdir, vtkstep, dg, Q, model)
    end
    callbacks = (callbacks..., cbvtk)
  end

  solve!(Q, lsrk; timeend=timeend, callbacks=callbacks)

  # final statistics
  engf = norm(Q)
  @info @sprintf """Finished
  norm(Q)                 = %.16e
  norm(Q) / norm(Q₀)      = %.16e
  norm(Q) - norm(Q₀)      = %.16e
  """ engf engf/eng0 engf-eng0
end

Base.@kwdef struct HeldSuarezSetup{FT}
  p_ground::FT = MSLP
  T_initial::FT = 255
  domain_height::FT = 30e3
end

function (setup::HeldSuarezSetup)(state, aux, coords, t)
  # callable to set initial conditions
  FT = eltype(state)

  r = norm(coords, 2)
  h = r - FT(planet_radius)

  scale_height = R_d * setup.T_initial / grav
  p = setup.p_ground * exp(-h / scale_height)

  state.ρ = air_density(setup.T_initial, p)
  state.ρu = SVector{3, FT}(0, 0, 0)
  state.ρe = state.ρ * (internal_energy(setup.T_initial) + aux.orientation.Φ)
  nothing
end

function held_suarez_forcing!(source, state, aux, t::Real)
  FT = eltype(state)

  ρ = state.ρ
  ρu = state.ρu
  ρe = state.ρe
  coord = aux.coord
  Φ = aux.orientation.Φ
  e = ρe / ρ
  u = ρu / ρ
  e_int = e - u' * u / 2 - Φ
  T = air_temperature(e_int)
  # Held-Suarez constants
  k_a = FT(1 / (40 * day))
  k_f = FT(1 / day)
  k_s = FT(1 / (4 * day))
  ΔT_y = FT(60)
  Δθ_z = FT(10)
  T_equator = FT(315)
  T_min = FT(200)
  σ_b = FT(7 / 10)
  r = norm(coord, 2)
  @inbounds λ = atan(coord[2], coord[1])
  @inbounds φ = asin(coord[3] / r)
  h = r - FT(planet_radius)
  scale_height = FT(7000) #from Smolarkiewicz JAS 2001 paper
  σ = exp(-h / scale_height)
  # TODO: use
  #  p = air_pressure(T, ρ)
  #  σ = p/p0
  exner_p = σ ^ (R_d / cp_d)
  Δσ = (σ - σ_b) / (1 - σ_b)
  height_factor = max(0, Δσ)
  T_equil = (T_equator - ΔT_y * sin(φ) ^ 2 - Δθ_z * log(σ) * cos(φ) ^ 2 ) * exner_p
  T_equil = max(T_min, T_equil)
  k_T = k_a + (k_s - k_a) * height_factor * cos(φ) ^ 4
  k_v = k_f * height_factor
  source.ρu += -k_v * ρu
  source.ρe += -k_T * ρ * cv_d * (T - T_equil) - k_v * ρu' * ρu / ρ
end

function do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, testname = "heldsuarez")
  ## name of the file that this MPI rank will write
  filename = @sprintf("%s/%s_mpirank%04d_step%04d",
                      vtkdir, testname, MPI.Comm_rank(mpicomm), vtkstep)

  statenames = flattenednames(vars_state(model, eltype(Q)))
  auxnames = flattenednames(vars_aux(model, eltype(Q)))
  writevtk(filename, Q, dg, statenames, dg.auxstate, auxnames)

  ## Generate the pvtu file for these vtk files
  if MPI.Comm_rank(mpicomm) == 0
    ## name of the pvtu file
    pvtuprefix = @sprintf("%s/%s_step%04d", vtkdir, testname, vtkstep)

    ## name of each of the ranks vtk files
    prefixes = ntuple(MPI.Comm_size(mpicomm)) do i
      @sprintf("%s_mpirank%04d_step%04d", testname, i - 1, vtkstep)
    end

    writepvtu(pvtuprefix, prefixes, (statenames..., auxnames...))

    @info "Done writing VTK: $pvtuprefix"
  end
end

main()


export BCQDLearner

"""
    BCQDLearner(;kwargs)

See paper: [Benchmarking Batch Deep Reinforcement Learning Algorithms](https://arxiv.org/abs/1910.01708).

# Keyword arguments

- `approximator`::[`ActorCritic`](@ref): used to get Q-values (Critic) and logits (Actor) of a state.
- `target_approximator`::[`ActorCritic`](@ref): similar to `approximator`, but used to estimate the target.
- `γ::Float32 = 0.99f0`, reward discount rate.
- `τ::Float32 = 0.005f0`, the speed at which the target network is updated.
- `θ::Float32 = 0.99f0`, regularization coefficient.
- `threshold::Float32 = 0.3f0`, determine whether the action can be used to calculate the Q value.
- `batch_size::Int=32`
- `update_freq::Int`: the frequency of updating the `approximator`.
- `update_step::Int = 0`
- `rng = Random.GLOBAL_RNG`
"""
mutable struct BCQDLearner{
    Aq<:ActorCritic,
    At<:ActorCritic,
    R<:AbstractRNG,
} <: AbstractLearner
    approximator::Aq
    target_approximator::At
    γ::Float32
    τ::Float32
    θ::Float32
    threshold::Float32
    batch_size::Int
    update_freq::Int
    update_step::Int
    rng::R
    # for logging
    actor_loss::Float32
    critic_loss::Float32
end

function BCQDLearner(;
    approximator::Aq,
    target_approximator::At,
    γ::Float32 = 0.99f0,
    τ::Float32 = 0.005f0,
    θ::Float32 = 1f-2,
    threshold::Float32 = 0.3f0,
    batch_size::Int = 32,
    update_freq::Int = 10,
    update_step::Int = 0,
    rng = Random.GLOBAL_RNG,
) where {Aq<:ActorCritic, At<:ActorCritic}
    copyto!(approximator, target_approximator)
    BCQDLearner(
        approximator,
        target_approximator,
        γ,
        τ,
        θ,
        threshold,
        batch_size,
        update_freq,
        update_step,
        rng,
        0.0f0,
        0.0f0,
    )
end

Flux.functor(x::BCQDLearner) = (Q = x.approximator, Qₜ = x.target_approximator),
y -> begin
    x = @set x.approximator = y.Q
    x = @set x.target_approximator = y.Qₜ
    x
end

function (learner::BCQDLearner)(env)
    s = state(env)
    s = Flux.unsqueeze(s, ndims(s) + 1)
    s = send_to_device(device(learner), s)
    q = learner.approximator.critic(s)
    prob = softmax(learner.approximator.actor(s), dims=1)
    mask = Float32.((prob ./ maximum(prob, dims=1)) .> learner.threshold)
    new_q = q .* mask .+ (1.0f0 .- mask) .* -1f8
    new_q |> vec |> send_to_host
end

function RLBase.update!(learner::BCQDLearner, batch::NamedTuple)
    AC = learner.approximator
    target_AC = learner.target_approximator
    γ, τ, θ = learner.γ, learner.τ, learner.θ
    batch_size = learner.batch_size
    D = device(AC)

    s, a, r, t, s′ = (send_to_device(D, batch[x]) for x in SARTS)
    a = CartesianIndex.(a, 1:batch_size)

    prob = softmax(AC.actor(s′))
    mask = Float32.((prob ./ maximum(prob, dims=1)) .> learner.threshold)
    q′ = AC.critic(s′)
    a′ = argmax(q′ .* mask .+ (1.0f0 .- mask) .* -1f8, dims=1)
    target_q = target_AC.critic(s′)

    target = r .+ γ .* (1 .- t) .* vec(target_q[a′])

    ps = Flux.params(AC)
    gs = gradient(ps) do
        # Critic loss
        q_t = AC.critic(s)
        critic_loss = Flux.Losses.huber_loss(q_t[a], target)

        # Actor loss
        logit = AC.actor(s)
        log_prob = -log.(softmax(logit, dims=1))
        actor_loss = mean(log_prob[a])

        ignore() do
            learner.actor_loss = actor_loss
            learner.critic_loss = critic_loss
        end

        actor_loss + critic_loss + θ * mean(logit .^ 2)
    end

    update!(AC, gs)

    # polyak averaging
    for (dest, src) in zip(
        Flux.params([learner.target_approximator]),
        Flux.params([learner.approximator]),
    )
        dest .= (1 - τ) .* dest .+ τ .* src
    end
end


# Perform actomyosin network simulations
# Alex Tam, 12/10/2020

# Load packages
using Revise # Modify code without restarting Julia
using Parameters # Tools for storing parameters in data structures
using LinearAlgebra # Matrix and vector operations
using StatsBase # Sampling without replacement
using Plots # Plotting library
using Plots.PlotMeasures # Enable re-sizing plot margins
using LaTeXStrings # Display LaTeX output
using ForwardDiff # Automatic differentiation tools
using Optim # Optimisation routines
using LineSearches # Line searches for optimisation
using Printf # C-style printing macros
using Distributions # Probability distributions
using DelimitedFiles # Read and write delimited data

# Include code from external files
include("model_parameters.jl")
include("actomyosin_network.jl")
include("State.jl")
include("initial_condition.jl")
include("Actin_Filament.jl")
include("periodic.jl")
include("Cross_Link.jl")
include("Myosin_Motor.jl")
include("intersection_search.jl")
include("network_statistics.jl")
include("network_force.jl")
include("draw_network.jl")
include("turnover.jl")
include("optimise_network.jl")
include("energy.jl")

function run_simulations()
    par_ref = 0.04;
    nTrials = 1;
    for i = 1:length(par_ref)
        par = par_ref[i]; # Select parameter
        bulk_stress = Vector{Float64}(); # Pre-allocate time-averaged bulk stress
        curvature = Vector{Float64}(); # Pre-allocate curvature
        index = Vector{Float64}(); # Pre-allocate integrated tension
        filament_speed = Vector{Float64}(); # Pre-allocate actin filament node speed
        motor_speed = Vector{Float64}(); # Pre-allocate motor head speed
        angle_roc = Vector{Float64}(); # Pre-allocate rate of change of angles
        motor_pos = Vector{Float64}(); # Pre-allocate time-averaged motor position
        motor_angle = Vector{Float64}(); # Pre-allocate time-averaged motor angle
        for j = 1:nTrials
            # Specify parameters
            parN = Numerical_Parameters(); # Initialise struct of numerical parameters
            parA = Actin_Properties(k_off = par); # Initialise struct of actin filament properties
            parM = Myosin_Properties(); # Initialise struct of myosin motor properties

            # Run simulations
            @time state, af, mm, xl, Bulk_Stress, Curvature, Index, Filament_Speed, Motor_Speed, Angle_ROC, Motor_Pos, Motor_Angle, Bins, Hist = actomyosin_network(parN, parA, parM, i, j);
            push!(bulk_stress, Bulk_Stress); # Store time-averaged net stress
            push!(curvature, Curvature); # Store time-averaged integrated curvature
            push!(index, Index); # Store time-averaged two-filament index
            push!(filament_speed, Filament_Speed); # Store time-averaged filament node speed
            push!(motor_speed, Motor_Speed); # Store time-averaged motor head speed
            push!(angle_roc, Angle_ROC); # Store time-averaged rate of change of angle
            push!(motor_pos, Motor_Pos); # Store time-averaged motor position
            push!(motor_angle, Motor_Angle); # Store time-averaged motor angle
            writedlm("bulk_stress-$par.txt", bulk_stress); # Write time-averaged bulk stress per trial to file
            writedlm("curvature-$par.txt", curvature); # Write time-averaged curvature per trial to file
            writedlm("index-$par.txt", index); # Write time-averaged two-filament index per trial to file
            writedlm("filament_speed-$par.txt", filament_speed); # Write time-averaged actin filament node speed per trial to file
            writedlm("motor_speed-$par.txt", motor_speed); # Write time-averaged motor head speed per trial to file
            writedlm("angle_roc-$par.txt", angle_roc); # Write time-averaged angle rate of change to file
            writedlm("motor_pos-$par.txt", motor_pos); # Write time-averaged motor position to file
            writedlm("motor_angle-$par.txt", motor_angle); # Write time-averaged motor angle to file
            writedlm("pd_end_bins-$par-trial-$j.txt", Bins); # Write paried-distances bins to file
            writedlm("pd_end_counts-$par-trial-$j.txt", Hist); # Write paired-distances counts to file
        end
    end
end

run_simulations();


# struct Bounds{L,U} end
# Bounds() = Bounds{-Inf,Inf}()
# Base.min(::Bounds{L}) where {L} = L
# Base.max(::Bounds{L,U}) where {L,U} = U

struct RealScalar{L,U} <: AbstractParameter
    RealScalar{L,U}() where {L,U} = new{Float64(L),Float64(U)}()
end
struct RealArray{S<:Tuple,L,U,D} <: AbstractParameter
    dynamicsizes::NTuple{D,Int}
    RealArray{S,L,U,D}(d::NTuple{D,Int}) where {S,L,U,D} = ((@assert U > L); new{S,Float64(L),Float64(U),D}(d))
end
const RealVector{N,L,U,D} = RealArray{Tuple{N},L,U,D}
const RealMatrix{M,N,L,U,D} = RealArray{Tuple{M,N},L,U,D}
RealScalar() = RealScalar{-Inf,Inf}()
RealScalar{L}() where {L} = RealScalar{Float64(L),Inf}()
RealArray{S}() where {S} = RealArray{S,-Inf,Inf,0}(tuple())
RealArray{S,L}() where {S,L} = RealArray{S,Float64(L),Inf,0}(tuple())
RealArray{S,L,U}() where {S,L,U} = RealArray{S,Float64(L),Float64(U),0}(tuple())
RealArray{S,L,U,0}() where {S,L,U} = RealArray{S,Float64(L),Float64(U),0}(tuple())
RealArray{S}(ds::NTuple{D}) where {S<:Tuple,D} = RealArray{S,-Inf,Inf,D}(ds)
RealArray{S,L}(ds::NTuple{D}) where {S<:Tuple,L,D} = RealArray{S,Float64(L),Inf,D}(ds)
RealArray{S,L,U}(ds::NTuple{D}) where {S<:Tuple,L,U,D} = RealArray{S,Float64(L),Float64(U),D}(ds)
RealArray(ds::NTuple{1,Int}) = RealArray{Tuple{-1},-Inf,Inf,1}(ds)
RealArray(ds::NTuple{2,Int}) = RealArray{Tuple{-1,-1},-Inf,Inf,2}(ds)
RealArray{L}(ds::NTuple{1,Int}) where {L} = RealArray{Tuple{-1},Float64(L),Inf,1}(ds)
RealArray{L}(ds::NTuple{2,Int}) where {L} = RealArray{Tuple{-1,-1},Float64(L),Inf,2}(ds)
RealArray{L,U}(ds::NTuple{1,Int}) where {L,U} = RealArray{Tuple{-1},Float64(L),Float64(U),1}(ds)
RealArray{L,U}(ds::NTuple{2,Int}) where {L,U} = RealArray{Tuple{-1,-1},Float64(L),Float64(U),2}(ds)

using PaddedMatrices
import PaddedMatrices: memory_length
memory_length(::RealScalar) = 1
memory_length(::RealVector{M}) where {M} = M
memory_length(::RealMatrix{M,N}) where {M,N} = M*N
memory_length(::Type{<:RealScalar}) = 1
memory_length(::Type{<:RealVector{M}}) where {M} = M
memory_length(::Type{<:RealMatrix{M,N}}) where {M,N} = M*N
@generated PaddedMatrices.memory_length(::RealArray{S}) where {S} = PaddedMatrices.simple_vec_prod(S.parameters)
@generated PaddedMatrices.memory_length(::Type{<:RealArray{S}}) where {S} = PaddedMatrices.simple_vec_prod(S.parameters)

abstract type AbstractMissingDataArray{M,L,U,T,N,A <: AbstractArray{T,N}} <: AbstractParameter end
memory_length(::AbstractMissingDataArray{M}) where {M} = M
"""
Parameters are:
M: number Missing
B: Bounds
T: DataType
N: Dimensionality of Array (how many axis?)
"""
struct MissingDataArray{M,L,U,T,N,A} <: AbstractMissingDataArray{M,L,U,T,N,A}
    data::A
    inds::Vector{CartesianIndex{N}}
end
struct ThreadedMissingDataArray{M,L,U,T,N,A} <: AbstractMissingDataArray{M,L,U,T,N,A}
    data::Vector{A}
    inds::Vector{CartesianIndex{N}}
end

@enum ParameterType::Int8 begin
    RealScalarParameter
    RealArrayParameter
    LKJCorrCholeskyParameter
    SimplexParameter
end

"""
The basic plan is to use Val{Bounds}() for dispatch.
"""
struct Bounds
    lb::Float64
    ub::Float64
end

Bounds(::RealScalar{L,U}) where {L,U} = Bounds(L,U)
Bounds(::RealArray{S,L,U}) where {S,L,U} = Bounds(L,U)

Base.min(b::Bounds) = b.lb
Base.max(b::Bounds) = b.ub
ismin(x::T) where {T} = ((x == typemin(T)) | (x == -floatmax(T))) # floatmin
ismax(x::T) where {T} = ((x == typemax(T)) | (x ==  floatmax(T)))
function isunbounded(b::Bounds)
    ismin(b.lb) & ismax(b.ub)
end
function islowerbounded(b::Bounds)
    (!ismin(b.lb)) & ismax(b.ub)
end
function isupperbounded(b::Bounds)
    ismin(b.lb) & (!ismax(b.ub))
end
function isbounded(b::Bounds)
    (!ismin(b.lb)) & (!ismax(b.ub))
end

# struct ParamDescriptionCore
#     b::Bounds
#     ind::Int32
#     t::ParameterType
# end

# struct SizedParamDescription{N}
#     s::NTuple{N,Int}
#     d::ParamDescriptionCore
# end
struct LengthParamDescription
    l::Int
    ind::Int32
    r::Int8
    LengthParamDescription(l::Int, d) = new(l, d % Int32, (l & ((VectorizationBase.REGISTER_SIZE>>>3)-1) ) % Int8)
end
# LengthDescription(::Type{A}, ind) where {A} = LengthDescription(memory_length(A)::Int, ind % Int32)
LengthParamDescription(A, ind) = LengthParamDescription(memory_length(A)::Int, ind % Int32)

# LengthParamDescription(spd::SizedParamDescription{0}) = LengthParamDescription(1, spd.d)#, one(Int8))
# function LengthParamDescription(spd::SizedParamDescription)
#     s = spd.s
#     l = minimum(s) < 0 ? -1 : prod(s)
#     LengthParamDescription(l, spd.d)
# end
# function update_rem!(descript::AbstractVector{LengthParamDescription}, R)
#     Rm1 = R - 1
#     for j ∈ 1:length(descript)
#         d = descript[j]
#         d = LengthParamDescription(d.l, d.ind, (d.l & Rm1) % Int8)
#         descript[j] = d
#     end
# end

function Base.isless(lpd1::LengthParamDescription, lpd2::LengthParamDescription)
    l1 = lpd1.l; l2 = lpd2.l
    l1dR = iszero(lpd1.r)
    l2dR = iszero(lpd2.r)
    if l1dR ⊻ l2dR
        l1dR
    else#if l1dR & l2dR
        l1 > l2
    # else
        # isless(lpd1, lpd2, R >>> 1)
    end
end
function shift_entries_right!(x, by = 1, start = firstindex(x), stop = lastindex(x))
    @assert stop ≤ lastindex(x)
    @assert start ≥ firstindex(x)
    @inbounds for i ∈ stop-by:-1:start
        x[i+by] = x[i]
    end
end
# Search for parameter pairs that preserve alignemnt, and move them to the front of unaligned params.
# E.g., given [5,5,5,3,3,1] and 8-Float64 alignment, move the (5,3)s in front, reordering to [5,3,5,3,5,1]
function combine_param_pairs!(descript::AbstractVector{LengthParamDescription}, R, start = firstindex(descript), jstart = start)
    N = length(descript)
    start > N - 2 && return start
    for j ∈ jstart:N
        descriptⱼ = descript[j]
        rⱼ = descriptⱼ.r
        Rdiff = R - rⱼ
        for k ∈ j+1:N
            descriptₖ = descript[k]
            rₖ = descriptₖ.r
            if rₖ == Rdiff
                shift_entries_right!(descript, 1, j + 1, k)
                shift_entries_right!(descript, 2, start, j + 1)
                descript[start] = descriptⱼ
                descript[start+1] = descriptₖ
                return combine_param_pairs!(descript, R, start + 2, j + 2)
            end
        end
    end
    start
end
# Search for parameter tripples that preserve alignment, and move them in front of unaligned params.
# E.g., given [5,5,5,2,2,1] and 8-Float64 alignment, move (5,2,1) in front so that it is [5,2,1,5,5,2]
function combine_param_tripples!(descript::AbstractVector{LengthParamDescription}, R, start = firstindex(descript), istart = start)
    N = length(descript)
    start > N - 3 && return start
    for i ∈ start:N
        descriptᵢ = descript[i]
        rᵢ = descriptᵢ.r
        Rdiffᵢ = R - rᵢ
        for j ∈ i+1:N
            descriptⱼ = descript[j]
            rⱼ = descriptⱼ.r
            Rdiffⱼ = Rdiffᵢ - rⱼ
            for k ∈ j+1:N
                descriptₖ = descript[k]
                rₖ = descriptₖ.r
                if rₖ == Rdiffⱼ
                    shift_entries_right!(descript, 1, j + 1, k)
                    shift_entries_right!(descript, 2, i + 1, j + 1)
                    shift_entries_right!(descript, 3, start, i + 1)
                    descript[start  ] = descriptᵢ
                    descript[start+1] = descriptⱼ
                    descript[start+2] = descriptₖ
                    return combine_param_tripples!(descript, R, start + 3, i + 3)
                end
            end
        end
    end
    start
end
function sort_by_rem!(descript::Vector{LengthParamDescription}, R = VectorizationBase.REGISTER_SIZE >>> 3)
    i = 1
    N = length(descript)
    # Original version looped while reducing R
    # We only do this once, because that risks misaligning larger objects.
    # E.g., align to 8 doubles, if we have [(l = 10007, id=1, r = 8), (l = 4, id=2, r = 4)]
    # The algorithm would now choose (id=1,id=2) order, but with the loop it would favor (id=2,id=1)
    #
    # descriptv = @view(descript[1:end])
    # while R > 1
    # descriptv = @view(descriptv[i:end])
    descriptv = descript
    # update_rem!(descriptv, R)
    sort!(descriptv)
    i = findfirst(d -> !iszero(d.r), descriptv)
    i === nothing && return
    i = combine_param_pairs!(descriptv, R, i)
    i = combine_param_tripples!(descriptv, R, i)
    nothing
        # N < i || return
        # R >>>= 1
    # end
end

function parameter_offsets(descript::Vector{LengthParamDescription})
    sort_by_rem!(descript) # should hopefully align loads and stores reasonably well
    offsets = similar(descript, Int)
    offset = 0
    for d ∈ descript
        # @show d.d.ind, d.l, offset
        offsets[d.ind] = offset
        offset += d.l
    end
    offsets
end
@generated function parameter_offsets_noinline(::Val{descript}) where {descript}
    lpd = LengthParamDescription[LengthParamDescription(d)::LengthParamDescription for d ∈ descript]
    offsetvec = parameter_offsets(lpd)
    tup = Expr(:tuple)
    append!(tup.args, offsetvec)
    Expr(:block, Expr(:meta, :noinline), tup)
end
@generated parameter_offsets(::Val{descript}) where {descript} = parameter_offsets_noinline(Val{descript}())
@generated parameter_offset(::Val{descript}, ::Val{I}) where {descript,I} = parameter_offsets_noinline(Val{descript}())[I]


# Correia, Sérgio P.C. and Dickinson, Anthony and Clayton, Nicola S., Western
# Scrub-Jays Anticipate Future Needs Independently of Their Current Motivational
# State, Current Biology 17 (10), 856-861, http://dx.doi.org/10.1016/j.cub.2007.03.063
#
function _summarize(::Experiment{:Correia07_exp1}, results)
	combine(groupby(results, [:prefed, :action, :foodtype]),
				 df -> DataFrame(μ = mean(df.count), sem = sem(df.count)))
end
function statistical_tests(exp::Experiment{:Correia07_exp1}, data)
    groups = groupby(data, [:prefed])
    tests = Test[]
    id = "interaction"
    push!(tests, Test(id,
                      splitplotanova(data,
                                     exp.tests[id].locals,
                                     betweenfactors = [],
                                     withinfactors = [:prefed, :foodtype],
                                     id = :id, data = :count),
                      exp.tests))
    for g in groups
        if g.prefed[1] in ("both", "none")
            push!(tests, Test("prefed $(g.prefed[1])",
                              isnot_significantly_different(g,
                                                            conditionkey = :foodtype,
                                                            condition1 = "kibble",
                                                            condition2 = "pinenut",
                                                            valuekey = :count),
                              exp.tests))
        else
            conditions = ["kibble", "pinenut"]
            if g.prefed[1] == "kibble"
                reverse!(conditions)
            end
            push!(tests, Test("prefed $(g.prefed[1])",
                              is_significantly_bigger(g, conditionkey = :foodtype,
                                                      condbigger = conditions[1],
                                                      condsmaller = conditions[2],
                                                      valuekey = :count),
                              exp.tests))
        end
    end
    TestCollection(tests)
end
function run!(exp::Experiment{:Correia07_exp1}, models; N = 11)
    results = DataFrame(prefed = String[], action = String[],
                        foodtype = String[], id = Int[],
						count = Int[])
	for id in 1:N
		m = models[id]
		for prefed in ("pinenut", "kibble", "both", "none")
			if prefed in ("pinenut", "both"); add!(m, Pinenut, 100); end
			if prefed in ("kibble", "both"); add!(m, Kibble, 100); end
			wait!(m, 3u"hr")
			remove!(m, Any)
			add!(m, Pinenut, 30)
			add!(m, Kibble, 30)
			wait!(m, 10u"minute")
			push!(results, [prefed, "eat", "pinenut", id, 30 - countfooditems(m, Pinenut)])
			push!(results, [prefed, "eat", "kibble", id, 30 - countfooditems(m, Kibble)])
			remove!(m, Any)
			add!(m, MaintenanceDiet)
			wait!(m, 50u"minute" + 20u"hr") # next day; not specified in paper
			remove!(m, MaintenanceDiet)
		end
	end
	results
end

function _summarize(::Experiment{:Correia07_exp2}, results)
	combine(groupby(results, [:trial, :group, :foodtype, :action]),
				 df -> DataFrame(μ = mean(df.count), sem = sem(df.count)))
end
function statistical_tests(exp::Experiment{:Correia07_exp2}, data)
    groups = groupby(data, :action)
    tests = Test[]
    for g in groups
        if g.action[1] == "eat"
            id = "eat"
            test = Test(id,
                        splitplotanova(g,
                                       exp.tests[id].locals,
                                       betweenfactors = :group,
                                       withinfactors = [:foodtype, :trial],
                                       id = :id, data = :count),
                        exp.tests)
            push!(tests, test)
        else
            id = "cache"
            test = Test(id,
                        splitplotanova(g,
                                       exp.tests[id].locals,
                                       betweenfactors = [:group, :prefed],
                                       withinfactors = [:foodtype, :trial],
                                       id = :id, data = :count),
                        exp.tests)
            push!(tests, test)
            ddata = combine(groupby(g, [:trial, :group, :action, :id]),
                       df ->  begin
#                            df.foodtype[1] != df.foodtype[2] &&
#                            length(df.foodtype) == 2 || error(df)
                           sum(df.count) > 0 ? df.count[df.foodtype[1] == "prefed" ? 1 : 2]/sum(df.count) : missing end)
            dropmissing!(ddata, disallowmissing = true)
            ddata.x1 = float.(ddata.x1)
            for i in 1:3 # trials
                x = @where ddata :trial .== i :group .== "same"
                y = @where ddata :trial .== i :group .== "different"
                if length(x.x1) > 0 && length(y.x1) > 0
                    test = Test("trial $i: proportion cache",
                                utest(x.x1, y.x1),
                                exp.tests)
                    push!(tests, test)
                elseif length(y.x1) > 0
                    push!(tests, Test("trial $i: proportion cache",
                                      NoTest("no caching in same condition"),
                                      exp.tests,
                                      [1.]))
                else
                    push!(tests, Test("trial $i: proportion cache",
                                      NoTest("no caching at all"),
                                      exp.tests,
                                      [1.]))
                end
            end
        end
    end
    TestCollection(tests)
end
function run!(exp::Experiment{:Correia07_exp2}, models; N = 11)
	results = DataFrame(trial = Int[], group = String[],
                        prefed = String[], foodtype = String[],
						action = String[], id = Int[], count = Int[])
	for id in 1:N
		m = models[id]
		for trial in 1:3
		    tray = Tray()
            if id % 2 == 0
                prefed = Pinenut
            else
                prefed = Kibble
            end
            add!(m, prefed, 100)
            wait!(m, 3u"hr")
            remove!(m, Any)
            add!(m, Pinenut, 30)
            add!(m, Kibble, 30)
            add!(m, tray)
            wait!(m, 10u"minute")
            group = id > N/2 ? "different" : "same"
            for foodtype in (Pinenut, Kibble)
                food = foodtype == prefed ? "prefed" : "nonprefed"
                push!(results, [trial, group, string(prefed), food,
                                "eat", id, 30 - countfooditems(m, foodtype)])
                push!(results, [trial, group, string(prefed), food,
                                "cache", id, countcache(tray, foodtype)])
            end
            remove!(m, Any)
            wait!(m, 30u"minute")
            if group == "same"
                add!(m, prefed, 100)
            else
                add!(m, setdiff([Pinenut, Kibble], [prefed])[1], 100)
            end
            wait!(m, 3u"hr")
            remove!(m, Any)
            add!(m, tray)
            wait!(m, 10u"minute")
            remove!(m, Any)
            add!(m, MaintenanceDiet)
            wait!(m, 10u"minute" + 16u"hr") # next day; not specified in paper
            remove!(m, MaintenanceDiet)
		end
	end
	results
end


"""
$(TYPEDEF)

Base for types defining system bath interactions in open quantum system models.
"""
abstract type AbstractInteraction end

"""
$(TYPEDEF)

An object to hold coupling operator and the corresponding bath object.

$(FIELDS)
"""
struct Interaction <: AbstractInteraction
    "system operator"
    coupling::AbstractCouplings
    "bath coupling to the system operator"
    bath::AbstractBath
end

"""
$(TYPEDEF)

A Lindblad operator, define by a rate ``γ`` and corresponding operator ``L```.

$(FIELDS)
"""
struct Lindblad <: AbstractInteraction
    "Lindblad rate"
    γ::Any
    "Lindblad operator"
    L::Any
    "size"
    size::Tuple
end

Lindblad(γ::Number, L::Matrix) = Lindblad((s) -> γ, (s) -> L, size(L))
Lindblad(γ::Number, L) = Lindblad((s) -> γ, L, size(L(0)))
Lindblad(γ, L::Matrix) = Lindblad(γ, (s) -> L, size(L))

function Lindblad(γ, L)
    if !(typeof(γ(0)) <: Number)
        throw(ArgumentError("γ should return a number."))
    end
    if !(typeof(L(0)) <: Matrix)
        throw(ArgumentError("L should return a matrix."))
    end
    Lindblad(γ, L, size(L(0)))
end

Base.size(lind::Lindblad) = lind.size

"""
$(TYPEDEF)

An container for different system-bath interactions.

$(FIELDS)
"""
struct InteractionSet{T <: Tuple}
    "A tuple of Interaction"
    interactions::T
end

InteractionSet(inters::AbstractInteraction...) = InteractionSet(inters)
Base.length(inters::InteractionSet) = Base.length(inters.interactions)
Base.getindex(inters::InteractionSet, key...) =
    Base.getindex(inters.interactions, key...)
Base.iterate(iters::InteractionSet, state=1) =
    Base.iterate(iters.interactions, state)

# The following functions are used to build different Liouvillians from
# the `InteractionSet`. They are not publicly available in the current
# release.
function redfield_from_interactions(iset::InteractionSet, U, Ta, atol, rtol)
    kernels = [build_redfield_kernel(i) for i in iset if !(typeof(i.bath) <: StochasticBath)]
    [RedfieldLiouvillian(kernels, U, Ta, atol, rtol)]
end

function cg_from_interactions(iset::InteractionSet, U, tf, Ta, atol, rtol)
    if Ta === nothing || ndims(Ta) == 0
        kernels = [build_cg_kernel(i, tf, Ta) for i in iset if !(typeof(i.bath) <: StochasticBath)]
    else
        if length(Ta) != length(iset)
            throw(ArgumentError("Ta should have the same length as the interaction sets."))
        end
        kernels = [build_cg_kernel(i, tf, t) for (i, t) in zip(iset, Ta) if !(typeof(i.bath) <: StochasticBath)]
    end
    [CGLiouvillian(kernels, U, atol, rtol)]
end

function ule_from_interactions(iset::InteractionSet, U, Ta, atol, rtol)
    kernels = [build_ule_kernel(i) for i in iset if !(typeof(i.bath) <: StochasticBath)]
    [ULELiouvillian(kernels, U, Ta, atol, rtol)]
end

function davies_from_interactions(iset::InteractionSet, ω_range, lambshift::Bool, lambshift_S)
    davies_list = []
    for i in iset
        coupling = i.coupling
        bath = i.bath
        if !(typeof(bath) <: StochasticBath)
            γfun = build_spectrum(bath)
            Sfun = build_lambshift(ω_range, lambshift, bath, lambshift_S)
            if typeof(bath) <: CorrelatedBath
                push!(davies_list, CorrelatedDaviesGenerator(coupling, γfun, Sfun, build_inds(bath)))
            else
                push!(davies_list, DaviesGenerator(coupling, γfun, Sfun))
            end
        end
    end
    davies_list
end

function onesided_ame_from_interactions(iset::InteractionSet, ω_range, lambshift::Bool, lambshift_S)
    l_list = []
    for i in iset
        coupling = i.coupling
        bath = i.bath
        if !(typeof(bath) <: StochasticBath)
            γfun = build_spectrum(bath)
            Sfun = build_lambshift(ω_range, lambshift, bath, lambshift_S)
            if typeof(i.bath) <: CorrelatedBath
                inds = build_inds(bath)
            else
                inds = ((i, i) for i in 1:length(coupling))
                γfun = SingleFunctionMatrix(γfun)
                Sfun = SingleFunctionMatrix(Sfun)
            end
            push!(l_list, OneSidedAMELiouvillian(coupling, γfun, Sfun, inds))
        end
    end
    l_list
end

function fluctuator_from_interactions(iset::InteractionSet)
    f_list = []
    for i in iset
        coupling = i.coupling
        bath = i.bath
        if typeof(bath) <: EnsembleFluctuator
            num = length(coupling)
            dist = construct_distribution(bath)
            b0 = [x.b for x in bath.f] .* rand([-1, 1], length(dist), num)
            next_τ, next_idx = findmin(rand(dist, num))
            push!(f_list, FluctuatorLiouvillian(coupling, dist, b0, next_idx, next_τ, sum(b0, dims=1)[:]))
        end
    end
    f_list
end

lindblad_from_interactions(iset::InteractionSet) =
    [LindbladLiouvillian([i for i in iset if typeof(i) <: Lindblad])]

function build_redfield_kernel(i::Interaction)
    coupling = i.coupling
    bath = i.bath
    cfun = build_correlation(bath)
    rinds = typeof(cfun) == SingleFunctionMatrix ?
        ((i, i) for i = 1:length(coupling)) : build_inds(bath)
    # the kernels is current set as a tuple
    (rinds, coupling, cfun)
end

function build_cg_kernel(i::Interaction, tf, Ta)
    coupling = i.coupling
    cfun = build_correlation(i.bath)
    Ta = Ta === nothing ? coarse_grain_timescale(i.bath, tf)[1] : Ta
    rinds = typeof(cfun) == SingleFunctionMatrix ?
        ((i, i) for i = 1:length(coupling)) : build_inds(i.bath)
    # the kernels is current set as a tuple
    (rinds, coupling, cfun, Ta)
end

function build_ule_kernel(i::Interaction)
    coupling = i.coupling
    cfun = build_jump_correlation(i.bath)
    rinds = typeof(cfun) == SingleFunctionMatrix ?
        ((i, i) for i = 1:length(coupling)) : build_inds(i.bath)
    # the kernels is current set as a tuple
    (rinds, coupling, cfun)
end


using Flux, Test, CUDA
using Zygote: pullback

@info "Testing GPU Support"
CUDA.allowscalar(false)

function gpu_gradtest(f, args...)
  args_gpu = gpu.(args)

  l_cpu, back_cpu = pullback((x...) -> f(x...), args...)
  g_cpu = back_cpu(1f0)[1]

  l_gpu, back_gpu = pullback((x...) -> f(x...), args_gpu...)
  g_gpu = back_gpu(1f0)[1]

  @test l_cpu ≈ l_gpu   rtol=1e-4 atol=1e-4
  @test g_gpu isa CuArray
  @test g_cpu ≈ collect(g_gpu)   rtol=1e-4 atol=1e-4
end

include("cuda.jl")
include("losses.jl")
include("layers.jl")

if CUDA.has_cudnn()
  @info "Testing Flux/CUDNN"
  include("cudnn.jl")
  include("curnn.jl")
else
  @warn "CUDNN unavailable, not testing GPU DNN support"
end


export BasicDQNLearner

using Random
using Flux

struct BasicDQNLearner{Q,F,R} <: AbstractLearner
    approximator::Q
    loss_func::F
    γ::Float32
    batch_size::Int
    min_replay_history::Int
    rng::R
end

(learner::BasicDQNLearner)(obs) =
    obs |> get_state |>
    x ->
        send_to_device(device(learner.approximator), x) |> learner.approximator |>
        send_to_host

function BasicDQNLearner(;
    approximator::Q,
    loss_func::F,
    γ = 0.99f0,
    batch_size = 32,
    min_replay_history = 32,
    seed = nothing,
) where {Q,F}
    rng = MersenneTwister(seed)
    BasicDQNLearner{Q,F,typeof(rng)}(
        approximator,
        loss_func,
        γ,
        batch_size,
        min_replay_history,
        rng,
    )
end

function RLBase.update!(learner::BasicDQNLearner, batch)
    Q, γ, loss_func, batch_size =
        learner.approximator, learner.γ, learner.loss_func, learner.batch_size
    s, r, t, s′ = map(
        x -> send_to_device(device(Q), x),
        (batch.state, batch.reward, batch.terminal, batch.next_state),
    )
    a = CartesianIndex.(batch.action, 1:batch_size)

    gs = gradient(params(Q)) do
        q = batch_estimate(Q, s)[a]
        q′ = vec(maximum(batch_estimate(Q, s′); dims = 1))
        G = r .+ γ .* (1 .- t) .* q′
        loss_func(G, q)
    end

    update!(Q, gs)
end

function RLBase.extract_experience(
    t::CircularCompactSARTSATrajectory,
    learner::BasicDQNLearner,
)
    if length(t) > learner.min_replay_history
        inds = rand(learner.rng, 1:length(t), learner.batch_size)
        map(get_trace(t, :state, :action, :reward, :terminal, :next_state)) do x
            consecutive_view(x, inds)
        end
    else
        nothing
    end
end


using FreeTypeAbstraction: height_insensitive_boundingbox

"""
Calculates the exact boundingbox of a Scene/Plot, without considering any transformation
"""
raw_boundingbox(x::Atomic) = data_limits(x)

rootparent(x) = rootparent(parent(x))
rootparent(x::Scene) = x

# function raw_boundingbox(x::Annotations)
#     bb = raw_boundingbox(x.plots)
#     inv(modelmatrix(rootparent(x))) * bb
# end

raw_boundingbox(x::Combined) = raw_boundingbox(x.plots)
boundingbox(x) = raw_boundingbox(x)

function combined_modelmatrix(x)
    m = Mat4f0(I)
    while true
        m = modelmatrix(x) * m
        if parent(x) !== nothing && parent(x) isa Combined
            x = parent(x)
        else
            break
        end
    end
    return m
end

function modelmatrix(x)
    t = transformation(x)
    transformationmatrix(t.translation[], t.scale[], t.rotation[])
end

function boundingbox(x::Atomic)
    bb = raw_boundingbox(x)
    return combined_modelmatrix(x) * bb
end

boundingbox(scene::Scene) = raw_boundingbox(scene)
function raw_boundingbox(scene::Scene)
    if scene[OldAxis] !== nothing
        return raw_boundingbox(scene[OldAxis])
    elseif scene.limits[] !== automatic
        return scene_limits(scene)
    elseif cameracontrols(scene) == EmptyCamera()
        # Empty camera means this is a parent scene that itself doesn't display anything
        return raw_boundingbox(scene.children)
    else
        plots = plots_from_camera(scene)
        children = filter(scene.children) do child
            child.camera == scene.camera
        end
        return raw_boundingbox([plots; children])
    end
end

function raw_boundingbox(plots::Vector)
    isempty(plots) && return FRect3D()
    plot_idx = iterate(plots)
    bb = FRect3D()
    while plot_idx !== nothing
        plot, idx = plot_idx
        plot_idx = iterate(plots, idx)
        # isvisible(plot) || continue
        bb2 = boundingbox(plot)
        isfinite_rect(bb) || (bb = bb2)
        isfinite_rect(bb2) || continue
        bb = union(bb, bb2)
    end
    return bb
end

function project_widths(matrix, vec)
    pr = project(matrix, vec)
    zero = project(matrix, zeros(typeof(vec)))
    return pr - zero
end

function rotate_bbox(bb::FRect3D, rot)
    points = decompose(Point3f0, bb)
    FRect3D(Ref(rot) .* points)
end

function gl_bboxes(gl::GlyphCollection)
    scales = gl.scales.sv isa Vec2f0 ? (gl.scales.sv for _ in gl.extents) : gl.scales.sv
    map(gl.extents, gl.fonts, scales) do ext, font, scale
        unscaled_hi_bb = height_insensitive_boundingbox(ext, font)
        hi_bb = FRect2D(
            Makie.origin(unscaled_hi_bb) * scale,
            widths(unscaled_hi_bb) * scale
        )
    end
end

function boundingbox(glyphcollection::GlyphCollection, position::Point3f0, rotation::Quaternion)

    if isempty(glyphcollection.glyphs)
        return FRect3D(position, Vec3f0(0, 0, 0))
    end

    chars = glyphcollection.glyphs
    glyphorigins = glyphcollection.origins
    glyphbbs = gl_bboxes(glyphcollection)

    bb = FRect3D()
    for (char, charo, glyphbb) in zip(chars, glyphorigins, glyphbbs)
        # ignore line breaks
        # char in ('\r', '\n') && continue

        charbb = rotate_bbox(FRect3D(glyphbb), rotation) + charo + position
        if !isfinite_rect(bb)
            bb = charbb
        else
            bb = union(bb, charbb)
        end
    end
    !isfinite_rect(bb) && error("Invalid text boundingbox")
    bb
end

function boundingbox(layouts::AbstractArray{<:GlyphCollection}, positions, rotations)

    if isempty(layouts)
        FRect3D((0, 0, 0), (0, 0, 0))
    else
        bb = FRect3D()
        broadcast_foreach(layouts, positions, rotations) do layout, pos, rot
            if !isfinite_rect(bb)
                bb = boundingbox(layout, pos, rot)
            else
                bb = union(bb, boundingbox(layout, pos, rot))
            end
        end
        !isfinite_rect(bb) && error("Invalid text boundingbox")
        bb
    end
end

function boundingbox(x::Text{<:Tuple{<:GlyphCollection}})
    boundingbox(
        to_value(x[1]),
        to_ndim(Point3f0, to_value(x[:position]), 0),
        convert_attribute(to_value(x[:rotation]), key"rotation"())
    )
end

function boundingbox(x::Text{<:Tuple{<:AbstractArray{<:GlyphCollection}}})
    boundingbox(
        to_value(x[1]),
        to_ndim.(Point3f0, to_value(x[:position]), 0),
        convert_attribute(to_value(x[:rotation]), key"rotation"())
    )
end

#     if x.space[] == :data
#         if x[1][] isa AbstractArray
#             bb = FRect3D()
#             for
#         else
#             boundingbox(x, to_value(x[1]))
#         end
#     elseif x.space[] == :screen
#         data_limits(x)
#     end
# end

# boundingbox(x::Text) = boundingbox(x, to_value(x[1]))

# function boundingbox(
#         text::String, position, textsize;
#         font = "default", align = (:left, :bottom), rotation = 0.0
#     )
#     return boundingbox(
#         text, position, textsize,
#         to_font(font), to_align(align), to_rotation(rotation)
#     )
# end

"""
Calculate an approximation of a tight rectangle around a 2D rectangle rotated by `angle` radians.
This is not perfect but works well enough. Check an A vs X to see the difference.
"""
function rotatedrect(rect::Rect{2}, angle)
    ox, oy = rect.origin
    wx, wy = rect.widths
    points = @SMatrix([
        ox oy;
        ox oy+wy;
        ox+wx oy;
        ox+wx oy+wy;
    ])
    mrot = @SMatrix([
        cos(angle) -sin(angle);
        sin(angle) cos(angle);
    ])
    rotated = mrot * points'

    rmins = minimum(rotated, dims = 2)
    rmaxs = maximum(rotated, dims = 2)

    return Rect2D(rmins..., (rmaxs .- rmins)...)
end

function quaternion_to_2d_angle(quat)
    # this assumes that the quaternion was calculated from a simple 2d rotation as well
    return 2acos(quat[4]) * (signbit(quat[1]) ? -1 : 1)
end

function boundingbox(
        text::String, position, textsize, font,
        align, rotation, model, justification, lineheight;
        # use the font's ascenders and descenders for the bounding box
        # this means that a string's boundingbox doesn't change in the vertical
        # dimension when characters change (for example numbers during an animation)
        # this is not wanted in most cases because of the jitter it creates when
        # the boundingbox slightly changes size in each frame (in MakieLayout mostly)
        use_vertical_dimensions_from_font = true
    )
    atlas = get_texture_atlas()
    N = length(text)
    ctext_state1 = iterate(text)
    ctext_state1 === nothing && return FRect3D()

    # call the layouting algorithm to find out where all the glyphs end up
    # this is kind of a doubling, maybe it could be avoided if at creation all
    # positions would be populated in the text object, but that seems convoluted
    if position isa VecTypes
        glyph_offsets = layout_text(text, textsize, font, align,
            rotation, model, justification, lineheight)
        position = [go .+ to_ndim(Point3f0, position, 0) for go in glyph_offsets]
    end

    bbox = Ref(FRect3D())            # use ref to avoid julia#15276
    ctext_state = Ref{Union{Nothing,typeof(ctext_state1)}}(ctext_state1)

    let position = position
        broadcast_foreach(1:N, rotation, font, textsize) do i, rotation, font, scale
            c, text_state = ctext_state[]::typeof(ctext_state1)
            ctext_state[] = iterate(text, text_state)

            if !(c in ('\r', '\n'))
                bb_unitspace = if use_vertical_dimensions_from_font
                    height_insensitive_boundingbox(
                        FreeTypeAbstraction.get_extent(font, c), font)
                else
                    inkboundingbox(FreeTypeAbstraction.get_extent(font, c))
                end

                scaled_bb = bb_unitspace * scale

                # TODO this only works in 2d
                rot_2d_radians = quaternion_to_2d_angle(rotation)
                rotated_bb = rotatedrect(scaled_bb, rot_2d_radians)

                # bb = rectdiv(bb, 1.5)
                shifted_bb = FRect3D(rotated_bb) + position[i]
                bboxc = bbox[]
                if !isfinite_rect(bboxc)
                    bbox[] = shifted_bb
                else
                    bbox[] = union(bboxc, shifted_bb)
                end
            end
        end
    end
    return bbox[]
end


using Flux
using Flux: onehot
using Flux.Data.Sentiment
using Flux.Batches: Tree, leaves

traintrees = Sentiment.train()

# Get the raw labels and phrases as separate trees.
labels  = map.(x -> x[1], traintrees)
phrases = map.(x -> x[2], traintrees)

# All tokens in the training set.
tokens = vcat(map(leaves, phrases)...)

# Count how many times each token appears.
freqs = Dict{String,Int}()
for t in tokens
  freqs[t] = get(freqs, t, 0) + 1
end

# Replace singleton tokens with an "unknown" marker.
# This roughly cuts our "alphabet" of tokens in half.
phrases = map.(t -> get(freqs, t, 0) == 1 ? "UNK" : t, phrases)

# Our alphabet of tokens.
alphabet = unique(vcat(map(leaves, phrases)...))

# One-hot-encode our training data with respect to the alphabet.
phrases_e = map.(t -> t == nothing ? t : onehot(t, alphabet), phrases)
labels_e  = map.(t -> onehot(t, 0:4), labels)

train = map.(tuple, phrases_e, labels_e)


module ExtendedVariantModeToggleDashDirectionTrigger

using ..Ahorn, Maple

@mapdef Trigger "ExtendedVariantMode/ToggleDashDirectionTrigger" ToggleDashDirectionTrigger(x::Integer, y::Integer, width::Integer=16, height::Integer=16,
	dashDirection::Integer=0b1000000000, enable::Bool=true, revertOnLeave::Bool=false, revertOnDeath::Bool=true)

const placements = Ahorn.PlacementDict(
	"Toggle Dash Direction (Extended Variant Mode)" => Ahorn.EntityPlacement(
		ToggleDashDirectionTrigger,
		"rectangle"
	)
)

Ahorn.editingOptions(trigger::ToggleDashDirectionTrigger) = Dict{String, Any}(
	"dashDirection" => Dict{String, Any}(
		"Top" =>          0b1000000000,
		"Top-Right" =>    0b0100000000,
		"Right" =>        0b0010000000,
		"Bottom-Right" => 0b0001000000,
		"Bottom" =>       0b0000100000,
		"Bottom-Left" =>  0b0000010000,
		"Left" =>         0b0000001000,
		"Top-Left" =>     0b0000000100
	)
)

end


@testset "Image Partitioning" begin
    img = testimage("mandrill")
    # Window size divides image size without remainder.
    subdivision = partition_image(AllowContraction(), img, 32)
    size(subdivision) == (16,16)
    for i in subdivision
        @test length.(i) == (32,32)
    end

    # Window size divides image size with remainder.
    subdivision = partition_image(AllowContraction(), img, 31)
    for r = 1:18
        for c = 1:18
            t = subdivision[r,c]
            if c == 1 || c == 18
                @test length(t[2]) == 8
            end
            if  1 < c < 18
                @test length(t[2]) == 31
            end
            if r == 1 || r == 18
                @test length(t[1]) == 8
            end
            if  1 < r < 18
                @test length(t[1]) == 31
            end
        end
    end
end


export BERL_NEAT, process

function BERL_NEAT(cfg::Dict, fitness::Function)::Cambrian.Evolution
    evaluate!::Function = e::Evolution -> NEAT_evaluate!(e, fitness)
    selection::Function = i::Array{NEATIndiv} -> i[1]

    if cfg["selection_type"] == "tournament"
        selection =
            i::Array{NEATIndiv} -> NEAT_tournament(i, cfg["tournament_size"])
    elseif cfg["selection_type"] == "random_top"
        selection =
            i::Array{NEATIndiv} -> NEAT_random_top(i, cfg["survival_threshold"])
    else
        throw(ArgumentError("Wrong selection type: " + cfg["selection_type"]))
    end

    populate!::Function = e::Evolution -> NEAT_populate!(e, selection)
    Evolution(NEATIndividual, cfg; id=cfg["id"], evaluate = evaluate!, populate = populate!)
end

function process(i::NEATIndiv, X)
    NeuroEvolution.process(i, X)
end

algorithms["NEAT"]=BERL_NEAT


# This is to make sure the swap operation does not make a mess
function swaptest()
    @testset "swap" begin
        r=Repo{String}(100,3)
        for i in 1:500           # We shuffle the array
            f=mod1(i*12345,100)  # Make some quick random numbers
            t=mod1(i*54321,100)
            Repos._swap(r,f,t)
        end
        for i=1:100          # Now test if the state is still consistent
            @test r.pos[r.ids[i]]==i
        end
    end
end

# Test for setclass and getclass functionality
function classtest()
    @testset "class" begin
        r=Repo{String}(4,4)
        setclass!(r,4,1)
        setclass!(r,3,2)
        setclass!(r,2,3)
        setclass!(r,1,4)
        @test getclass(r,1)==4
        @test getclass(r,2)==3
        @test getclass(r,3)==2
        @test getclass(r,4)==1
    end
end


using DataFrames
# Test evolve_function.jl in Evolve.jl for the specific purpose of reducing the number of gates
#    in a 2-component goal where circuits has already been evolved to map to each single-component goal.

# Goal is the 2-component goal
function test_evolve_function( goal::Goal, p::Parameters, maxsteps::Int64 )
  @assert p.numouputs==2
  p1 = Parameters( div(p.numinputs,2), 1, p.numlevelsback, p.numinteriors, p.numlevelsback )
  maxtries = 10
  tries = 1
  done = :false
  while !done && tries < maxtries
    ec = random_chromosome(p1)
    (c1,steps,maxsteps) = neutral_evolution( ec, [goal[1]], maxsteps )
    if steps < maxsteps
      done = :true
    end
  end
  if !done
    error("unable to evolve circuit to map to goal ",[goal[1]])
  end
  tries = 1
  done = :false
  while !done && tries < maxtries
    ec = random_chromosome(p1)
    (c2,steps,maxsteps) = neutral_evolution( ec, [goal[2]], maxsteps )
    if steps < maxsteps
      done = :true
    end
  end
  if !done
    error("unable to evolve circuit to map to goal ",[goal[1]])
  end
end

# Combines two 1-output circuits c1 and c2 with the same paramterers and number of gates into a 2-output circuit
# The output values of the result is the 2-output goal whose components are the output values of c1 and of c2.
function combine_circuits( c1::Chromosome, c2::Chromosome )
  # Two functions used only within this function
  inp1(i,numinputs) = i <= numinputs ? i : 2*(i-numinputs)+numinputs-1
  inp2(i,numinputs) = i <= numinputs ? i : 2*(i-numinputs)+numinputs
  @assert c1.params.numinputs == c2.params.numinputs
  @assert c1.params.numoutputs == 1
  @assert c2.params.numoutputs == 1
  if length(c1.interiors) < length(c2.interiors)
    println("  c1ints: ",length(c1.interiors),"  c2ints: ",length(c2.interiors))
    ctemp = c1
    c1 =c2
    c2 = ctemp
  end
  #println("c1ints: ",length(c1.interiors),"  c2ints: ",length(c2.interiors))
  p = Parameters( c1.params.numinputs, 2, c1.params.numinteriors+c2.params.numinteriors,
    c1.params.numlevelsback+c2.params.numlevelsback )
  inputs = c1.inputs
  interiors = InteriorNode[]
  new_numinteriors = 2*length(c2.interiors)
  #println("new_numints: ",new_numinteriors)
  #print_parameters(p)
  for i = 1:length(c2.interiors)
    int_node1 = deepcopy(c1.interiors[i])
    input1 = inp1(c1.interiors[i].inputs[1],p.numinputs)
    input2 = inp1(c1.interiors[i].inputs[2],p.numinputs)
    int_node1.inputs = [input1,input2]
    #println("i: ",i,"  int_node1: ",int_node1)
    push!(interiors,int_node1)
    int_node2 = deepcopy(c2.interiors[i])
    input1 = inp2(c2.interiors[i].inputs[1],p.numinputs)
    input2 = inp2(c2.interiors[i].inputs[2],p.numinputs)
    int_node2.inputs = [input1,input2]
    #println("i: ",i,"  int_node2: ",int_node2)
    push!(interiors,int_node2)
  end
  #println("interiors: ",interiors)
  outputs = [OutputNode(p.numinputs+new_numinteriors-1),OutputNode(p.numinputs+new_numinteriors)]
  Chromosome( p, c1.inputs, interiors, outputs, 0.0, 0.0 )
end

function test_combine_circuits( p::Parameters )
  c1 = random_chromosome(p); print(output_values(c1),"  "); print_circuit(c1)
  c2 = random_chromosome(p); print(output_values(c2),"  "); print_circuit(c2)
  c = combine_circuits(c1,c2); print(output_values(c),"  "); print_circuit(c)
end

# Evaluate the performance of function evolve_reduce_numactive() which is in evolve_funct.jl.
# For each pair of chromosomes
#   Creates two chromosomes using parameter settings p.
#   Combines them into one chromosome with one more input
#   For each trial run evolve_reduce_numactive() and record average and minimum numactive.
function test_reduce_numactive( p::Parameters, numpairs::Int64, numtrials::Int64, maxsteps::Int64;
      num_mutations::Int64=1, csvfile::String="" )
  df = DataFrame()
  df.combined_goal = Goal[]
  df.numinputs = Int64[]
  df.numgates = Int64[]
  df.numlevelsback = Int64[]
  df.numpairs = Int64[]
  df.numtrials = Int64[]
  df.maxsteps = Int64[]
  df.nactive0 = Float64[]
  df.nactive1_mean = Float64[]
  df.nactive1_std = Float64[]
  if num_mutations >= 2
    df.nactive2_mean = Float64[]
    df.nactive2_std = Float64[]
  end
  if num_mutations == 3
    df.nactive3_mean = Float64[]
    df.nactive3_std = Float64[]
  end
  for npair = 1:numpairs
    c1 = random_chromosome(p); print(output_values(c1),"  "); print_circuit(c1)
    c2 = random_chromosome(p); print(output_values(c2),"  "); print_circuit(c2)
    c = combine_circuits(c1,c2)
    rowlist = [ output_values(c), p.numinputs, p.numinteriors, p.numlevelsback, numpairs, numtrials, maxsteps ]
    #@assert [output_values(c1)[1],output_values(c1)[1]] == output_values(c)
    println("output_values(c): ",output_values(c))
    println("goal: ",[output_values(c1)[1],output_values(c2)[1]])
    #nactive_lists = map(i->reduce_nactive_helper(c,maxsteps,num_mutations=num_mutations), collect(1:numtrials ) )
    nactive_lists = pmap(i->reduce_nactive_helper(c,maxsteps,num_mutations=num_mutations), collect(1:numtrials ) )
    nactive0 = vcat( [ na[1] for na in nactive_lists]... )
    nactive1 = vcat( [ na[2] for na in nactive_lists]... )
    if num_mutations == 1
      row_results = [mean(nactive0),mean(nactive1),std(nactive1)]
    end
    if num_mutations >= 2
      nactive2 = vcat( [ na[3] for na in nactive_lists]... )
      row_results = [mean(nactive0),mean(nactive1),std(nactive1),mean(nactive2),std(nactive2)]
    end
    if num_mutations >= 3
      nactive3 = vcat( [ na[4] for na in nactive_lists]... )
      row_results = [mean(nactive0),mean(nactive1),std(nactive1),mean(nactive2),std(nactive2),mean(nactive3),std(nactive3)]
    end
    row = deepcopy(rowlist)
    println("length(row): ",length(row),"  length(row_results): ",length(row_results),"  size(df): ",size(df))
    #println("npair: ",npair,"  mean(nactive1): ",mean(nactive1),"  mean(nactive2): ",mean(nactive2),"  mean(nactive3): ",mean(nactive3))
    push!(df,vcat(row,row_results))
  end
  hostname = chomp(open("/etc/hostname") do f read(f,String) end)
  if length(csvfile) > 0
    open( csvfile, "w" ) do f
      println(f,"# date and time: ",Dates.now())
      println(f,"# host: ",hostname," with ",nprocs()-1,"  processes: " )
      println(f,"# funcs: ", Main.CGP.default_funcs(p.numinputs))
      print_parameters(f,p,comment=true)
      println(f,"# numpairs: ", numpairs)
      println(f,"# numtrials: ", numtrials)
      println(f,"# maxsteps: ", maxsteps)
      CSV.write( f, df, append=true, writeheader=true )
    end
  end
  df
end

# nactive0 is number_active_gates() after evolution to find circuit that maps to goal
#    but before running evolve_reduce_numactive()
function reduce_nactive_helper( c::Chromosome, maxsteps::Int64; num_mutations::Int64=1 )
  #println("num_mutations: ",num_mutations)
  goal = output_values( c )
  (ec,steps) = neutral_evolution( random_chromosome(c.params), goal, maxsteps )
  nactive0 = number_active_gates(ec)
  println("nactive0: ",nactive0)
  nactive1 = Int64[]
  nc = deepcopy(c)
  (ec,cur_nactive) = evolve_reduce_numactive( nc, maxsteps, num_mutations=1 )
  @assert cur_nactive == number_active_gates(ec)
  push!(nactive1,number_active_gates(ec))
  if num_mutations == 1
    return(nactive0, nactive1)
  end
  nactive2 = Int64[]
  nc = deepcopy(c)
  (ec,cur_nactive) = evolve_reduce_numactive( nc, maxsteps, num_mutations=2 )
  push!(nactive2,number_active_gates(ec))
  if num_mutations == 2
    return(nactive0, nactive1, nactive2)
  end
  nactive3 = Int64[]
  nc = deepcopy(c)
  (ec,cur_nactive) = evolve_reduce_numactive( nc, maxsteps, num_mutations=3 )
  push!(nactive3,number_active_gates(ec))
  #println("trial: ",trial,"  mean(nactive1): ",mean(nactive1),"  mean(nactive2): ",mean(nactive2),"  mean(nactive2): ",mean(nactive2))
  (nactive0, nactive1, nactive2, nactive3)
end

# folder_path is relative to data folder
function compare_dataframe( folder_path::String, suffix::String )
  rdf = read_dataframe("$(folder_path)/reduce_numactive_$(folder_path)_$(suffix).csv")
  kdf = read_dataframe("$(folder_path)/k_complexity$(folder_path)$(suffix).csv")
  cdf = DataFrame()
  cdf.goal = kdf.goal
  cdf.kng = kdf.num_gates
  cdf.rng = rdf.nactive1_mean
  cdf.k_minus_r_diff = cdf.kng .- cdf.rng
  cdf
end


__precompile__()

module SnpArrays

using CodecZlib, CodecXz, CodecBzip2, CodecZstd,  TranscodingStreams
using Glob, LinearAlgebra, Missings, Mmap, SparseArrays, Statistics, StatsBase, LoopVectorization
using Requires, Adapt
import Base: IndexStyle, convert, copyto!, eltype, getindex, setindex!, length, size, wait
import DataFrames: DataFrame, rename!, eachrow
import DelimitedFiles: readdlm, writedlm
import CSV: categorical!
import CSV # for CSV.read, to avoid clash with Base.read
import LinearAlgebra: mul!
import Statistics: mean, std, var
import StatsBase: counts
import SpecialFunctions: gamma_inc
import VectorizationBase: gesp
import GeneticVariation.VCF
export AbstractSnpArray, SnpArray, SnpBitMatrix, SnpLinAlg, SnpData
export compress_plink, decompress_plink, split_plink, merge_plink, write_plink
export counts, grm, maf, mean, minorallele, missingpos, missingrate, std, var
export vcf2plink, kinship_pruning
export ADDITIVE_MODEL, DOMINANT_MODEL, RECESSIVE_MODEL
export CuSnpArray

const ADDITIVE_MODEL = Val(1)
const DOMINANT_MODEL = Val(2)
const RECESSIVE_MODEL = Val(3)

include("codec.jl")
include("snparray.jl")
include("filter.jl")
include("cat.jl")
include("snpdata.jl")
include("grm.jl")
include("kinship_pruning.jl")
include("linalg_direct.jl")
include("linalg_bitmatrix.jl")
include("reorder.jl")
include("vcf2plink.jl")

datadir(parts...) = joinpath(@__DIR__, "..", "data", parts...)

function __init__()
    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
end

end # module


function search11(r::Robot)
    i = 2
    while (i > 0)
        if !ismarker(r)
            move!(r,Nord)
            move!(r,West)
            i += rrrruond(r,i)
            i += 2
        end
    end
end

function rrrruond(r::Robot,i::Int64)
    j=0
    for side in (Ost,Sud,West,Nord)
        a = i
        while (a > 0)
            move!(r,side)
            if ismarker(r)
                a = -1
                break
            end
            a -= 1
        end
        if (a == -1)
            j = -i-2
            break

        end

    end
    return(j)
end


include("../../__lib__/std/datastructs/src/datastructs_module.jl")

# dependencies:
# - datastructs

module autodiff
using ..datastructs
using LinearAlgebra

abstract type AbstractMode end

struct ForwardMode <: AbstractMode
end

struct ReverseMode <: AbstractMode
end

include("forward_mode.jl")
include("reverse_mode.jl")

end


using StringDistances, Random
using BenchmarkTools

N = if length(ARGS) > 0
    try
        parse(Int, ARGS[1])
    catch _
        100
    end
else
    100 # default value
end

Maxlength = if length(ARGS) > 1
    try
        parse(Int, ARGS[2])
    catch _
        100
    end
else
    100 # default value
end

# If there are strings already cached to disk we start with them and only
# add new ones if needed.
using Serialization
const CacheFile = joinpath(@__DIR__(), "perfteststrings_$(Maxlength).juliabin")
SaveCache = false

S = if isfile(CacheFile)
    try
        res = deserialize(CacheFile)
        println("Read $(length(res)) strings from cache file: $CacheFile")
        res
    catch err
        String[]
    end
else
    println("Creating $N random strings.")
    SaveCache = true
    String[randstring(rand(3:Maxlength)) for _ in 1:N]
    SaveCache = true
    String[randstring(rand(3:Maxlength)) for _ in 1:N]
end

if length(S) < N
    for i in (length(S)+1):N
        push!(S, randstring(rand(3:Maxlength)))
    end
    SaveCache = true
end

if SaveCache
    println("Saving cache file with $(length(S)) strings: $CacheFile")
    serialize(CacheFile, S)
end

println("For ", Threads.nthreads(), " threads and ", N, " strings of max length ", Maxlength, ":")

dist = Cosine(2)
t1 = @belapsed dm1 = pairwise(dist, S; preprocess = false)
t2 = @belapsed dm2 = pairwise(dist, S; preprocess = true)

println("  - time WITHOUT pre-calculation: ", round(t1, digits = 3))
println("  - time WITH    pre-calculation: ", round(t2, digits = 3))
println("  - speedup with pre-calculation: ", round(t1/t2, digits = 1))


using Test
using FileIO, Images, Colors
using CSV

# !!! commented out because of local repository of test images. !!!

# base_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/anhir/"
# dataset_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/anhir/dataset/"
# landmark_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/anhir/landmarks/"
# loc_table = CSV.read(base_path * "location_table.csv")
#
# train_rows = loc_table[loc_table[:status] .== "training", :]
#
# @testset "LAP_julia.jl" begin
#     TEST_SIZE = 2
#     @testset "Tests on random images $k" for k = 1:TEST_SIZE
#         random_rows = train_rows[rand(1:size(train_rows, 1), TEST_SIZE), :]
#         source_loc = (random_rows[k, :])[Symbol("Source image")]
#         target_loc = (random_rows[k, :])[Symbol("Target image")]
#         source_in = load(dataset_path * source_loc)
#         target_in = load(dataset_path * target_loc)
#         gray_source = Gray.(source_in)
#         gray_target = Gray.(target_in)
#
#         @testset "rescale!" begin
#             both((x,y)) = [x[:]; y[:]]
#             @test maximum(both(LAP_julia.helpers.rescale!(gray_source, gray_target))) == 1
#             @test minimum(both(LAP_julia.helpers.rescale!(gray_source, gray_target))) == 0
#         end
#     end
# end

# @testset "Registration Algorithms" begin
#     TimerOutputs.enable_debug_timings(LAP_julia.registration)
#     img, imgw, flow = gen_init(:lena, flow_args=[15, 140])
#
#     @testset "PFLAP" begin
#         timer = TimerOutput("Registration");
#         @timeit timer "polyfilter lap" begin
#             flow_est, source_reg = pflap(img, imgw, display=false, timer=timer)
#         end
#         print_timer(timer)
#
#         flow_est, source_reg = pflap(img, imgw, display=false)
#         showflow(flow)
#         showflow(flow_est)
#         @test angle_mae(flow, flow_est) < 10

@testset "helpers" begin
    @testset "pad_images $k, $l, $m, $n" for k = 1:3, l = 1:3, m = 1:3, n = 1:3
        im1 = ones(k, l)
        im2 = ones(m, n)
        @test size(LAP_julia.pad_images(im1, im2)[1]) <= (k+m-1, l+n-1)
        @test size(LAP_julia.pad_images(im1, im2)[1]) == ((k<m) ? m : k, (l<n) ? n : l)
    end
end

@testset "inpaint" begin
    u = ones(10,10) .+ 2im .* ones(10,10);
    nan_u = u
    nan_u[3:8, 3:8] .= NaN .+ NaN .* 1im;
    LAP_julia.inpaint.inpaint_nans!(nan_u)
    @test all(nan_u .== u)
    @test any(isnan.(nan_u)) == false
end


module CubeIO
using ..Cubes
using ..DAT
using ..Proc

import NetCDF.ncread, NetCDF.ncclose
import Base.Iterators
import ..Cubes: saveCube, check_overwrite, getsavefolder, cubechunks
import ..Cubes.ESDLZarr: ZArrayCube

function toPointAxis(aout,ain,loninds,latinds)
  iout = 1
  for (ilon,ilat) in zip(loninds,latinds)
    aout[iout]=ain[ilon,ilat]
    iout+=1
  end
end

"""
    extractLonLats(c::AbstractCubeData,pl::Matrix)

Extracts a list of longitude/latitude coordinates from a data cube. The coordinates
are specified through the matrix `pl` where `size(pl)==(N,2)` and N is the number
of extracted coordinates. Returns a data cube without `LonAxis` and `LatAxis` but with a
`SpatialPointAxis` containing the input locations.
"""
function extractLonLats(c::AbstractCubeData,pl::Matrix;kwargs...)
  size(pl,2)==2 || error("Coordinate list must have exactly 2 columns")
  axlist=caxes(c)
  ilon=findAxis(LonAxis,axlist)
  ilat=findAxis(LatAxis,axlist)
  ilon>0 || error("Input cube must contain a LonAxis")
  ilat>0 || error("input cube must contain a LatAxis")
  lonax=axlist[ilon]
  latax=axlist[ilat]
  pointax = SpatialPointAxis([(pl[i,1],pl[i,2]) for i in 1:size(pl,1)])
  loninds = map(ll->axVal2Index(lonax,ll[1]),pointax.values)
  latinds = map(ll->axVal2Index(latax,ll[2]),pointax.values)
  incubes=InDims("Lon","Lat")
  outcubes=OutDims(pointax)
  y=mapCube(toPointAxis,c,loninds,latinds;indims=incubes,outdims=outcubes,max_cache=1e8,kwargs...)
end
export extractLonLats

using NetCDF
import Base.Iterators: product
function writefun(xout,xin::AbstractArray{Union{Missing,T}},a,nd,cont_loop,filename;kwargs...) where T

  x = map(ix->ismissing(ix) ? convert(T,-9999.0) : ix,xin)

  count_vec = fill(-1,nd)
  start_vec = fill(1,nd)

  used_syms = Symbol[]
  for (k,v) in cont_loop
    count_vec[k] = 1
    start_vec[k] = a[Symbol(v)][1]
    push!(used_syms,Symbol(v))
  end

  splinds = Iterators.filter(i->!in(i,used_syms),keys(a))
  vn = join(string.([a[s][2] for s in splinds]),"_")
  isempty(vn) && (vn="layer")
  ncwrite(x,filename,vn,start=start_vec,count=count_vec)
end

"""
    exportcube(r::AbstractCubeData,filename::String)

Saves a cube object to a portable NetCDF file in `filename`.

When saving, every RangeAxis will be converted to an axis in the NetCDF cube,
while every categorical axis will be represented by a different variable
inside the resulting file. Dimensions will be ordered according to the
`priorities` keyword argument, which defaults to `Dict("LON"=>1,"LAT"=>2,"TIME"=>3)`,
which means that the file will be stored with longitudes varying fastest.
"""
function exportcube(r::AbstractCubeData,filename::String;priorities = Dict("LON"=>1,"LAT"=>2,"TIME"=>3), proj=epsg4326)

  ax = caxes(r)
  ax_cont = collect(filter(i->isa(i,RangeAxis),ax))
  ax_cat  = filter(i->!isa(i,RangeAxis),ax)
  prir = map(i->get(priorities,uppercase(axname(i)),10),ax_cont)
  ax_cont=ax_cont[sortperm(prir)]
  dims = map(NcDim,ax_cont)
  isempty(ax_cat) && (ax_cat=[VariableAxis(["layer"])])
  it = map(i->i.values,ax_cat)
  elt = Base.nonmissingtype(eltype(r))
  vars = NcVar[NcVar(join(collect(string.(a)),"_"),dims,t=elt,atts=Dict("missing_value"=>convert(elt,-9999.0), "grid_mapping" => proj["grid_mapping_name"])) for a in product(it...)]
  file = NetCDF.create(filename,vars)
  for d in dims
    ncwrite(d.vals,filename,d.name)
  end
  nccreate(filename, proj["grid_mapping_name"])
  ncputatt(filename, proj["grid_mapping_name"], proj)
  dl = map(i->i.dimlen,dims) |> cumprod
  isplit = findfirst(i->i>5e7,dl)
  isplit isa Nothing && (isplit=length(dl))
  incubes = InDims(ax_cont[1:(isplit-1)]...)
  cont_loop = Dict(ii=>axname(ax_cont[ii]) for ii in isplit:length(ax_cont))
  mapCube(writefun,r,length(ax_cont),cont_loop,filename,indims=incubes,include_loopvars=true,ispar=false,max_cache=5e8,
  nthreads=[1])
  NetCDF.close(file)
  nothing
end

import Zarr: NoCompressor

function saveCube(c::AbstractCubeData, name::AbstractString; overwrite = false, chunksize = cubechunks(c), compressor=NoCompressor(), max_cache=1e8)
  allax = caxes(c)
  firstaxes = findall(i->i>1,chunksize)
  lastaxes = setdiff(1:length(allax),firstaxes)
  allax = allax[[firstaxes;lastaxes]]
  dl = cumprod(length.(caxes(c)))
  isplit = findfirst(i->i>max_cache/sizeof(eltype(c)),dl)
  isplit isa Nothing && (isplit=length(dl)+1)
  forcesingle = (isplit+1)<length(firstaxes)
  axn = axname.(allax[1:isplit-1])
  indims = InDims(axn...)
  path = getsavefolder(name)
  check_overwrite(path,overwrite)
  outdims = OutDims(axn..., retcubetype=ZArrayCube,chunksize=chunksize[1:length(axn)], compressor=compressor, path = path)
  if forcesingle
    nprocs()>1 && println("Forcing single core processing because of bad chunk size")
    o = mapCube(copyto!,c,indims=indims, outdims=outdims,ispar=false,max_cache=max_cache)
  else
    o = mapCube(copyto!,c,indims=indims, outdims=outdims,max_cache=max_cache)
  end
end

export exportcube

global const projection = Dict(
"grid_mapping_name" => "transverse_mercator",
"longitude_of_central_meridian" => -9. ,
"false_easting" => 500000. ,
"false_northing" => 0. ,
"latitude_of_projection_origin" => 0. ,
"scale_factor_at_central_meridian" => 0.9996,
"long_name" => "CRS definition",
"longitude_of_prime_meridian" => 0.,
"semi_major_axis" => 6378137.,
"inverse_flattening" => 298.257223563 ,
"spatial_ref" => "PROJCS[\"WGS 84 / UTM zone 29N\",GEOGCS[\"WGS 84\",
DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,
AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],
PRIMEM[\"Greenwich\",0,AUTHORITY[\"EPSG\",\"8901\"]],
UNIT[\"degree\",0.0174532925199433,AUTHORITY[\"EPSG\",\"9122\"]],
AUTHORITY[\"EPSG\",\"4326\"]],
PROJECTION[\"Transverse_Mercator\"],
PARAMETER[\"latitude_of_origin\",0],
PARAMETER[\"central_meridian\",-9],
PARAMETER[\"scale_factor\",0.9996],
PARAMETER[\"false_easting\",500000],
PARAMETER[\"false_northing\",0],
UNIT[\"metre\",1,AUTHORITY[\"EPSG\",\"9001\"]],
AXIS[\"Easting\",EAST],
AXIS[\"Northing\",NORTH],
AUTHORITY[\"EPSG\",\"32629\"]]",
"GeoTransform" => "722576.2320803495 20 0 4115483.464603715 0 -20 ")

global const epsg4326=Dict(
  "grid_mapping_name" => "latitude_longitude",
  "long_name" => "CRS definition",
  "longitude_of_prime_meridian" => 0.,
  "semi_major_axis" => 6378137.,
  "inverse_flattening" => 298.257223563,
  "spatial_ref" => "GEOGCS[\"WGS 84\",DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],PRIMEM[\"Greenwich\",0],UNIT[\"degree\",0.0174532925199433],AUTHORITY[\"EPSG\",\"4326\"]]",
  "GeoTransform" => "-98.68712696894481 0.0001796630568239077 0 20.69179551612753 0 -0.0001796630568239077 "
)

end


function B_0_equation(cur_reactor::AbstractReactor, cur_value::Number)
  cur_G_T = cur_value

  cur_equation = Equation(0, 1, 0, cur_G_T)

  cur_equation
end


"""
Abstract type for all Components
"""
abstract type Component end

struct ComponentPoolIndex
    type::Type{<:Component}
    index::Int
end
ComponentPoolIndex(t::Component, i::Int) = ComponentPoolIndex(typeof(t), i)

struct ComponentSet{T<:Tuple} <: Component
    _raw::Vector{ComponentPoolIndex}
end

struct ComponentPool
    _raw::Dict{Type{<:Component}, Vector{Component}}
    _sets::Vector{Type{<:ComponentSet}}
end
ComponentPool()  = ComponentPool(Dict(), Vector())
Base.getindex(c::ComponentPool, t::Type{<:Component}, i::Int) = c._raw[t][i]
Base.getindex(c::ComponentPool, t::Type{T}) where T<:Component = if t in keys(c._raw) c._raw[t] else Vector() end
function Base.getindex(c::ComponentPool, i::ComponentPoolIndex)
    res = getindex(c, i.type, i.index)
    if typeof(res) <: ComponentSet
        res = Tuple(c[v] for v in res._raw)
    end
    res
end

ComponentSet(T::Vararg{Type{<:Component}}) = ComponentSet{T}
ComponentSet(v::Vararg{ComponentPoolIndex}) = ComponentSet(collect(v))
ComponentSet{T}(args...) where T<:Tuple = ComponentSet(args...)
function ComponentSet(c::ComponentPool, v::Vararg{Tuple{<:Component}})
    set = ComponentSet(v)
    if !(v in c._sets) push!(c._sets, set) end
end

function indexify(v::Vector{Component})
    i = 1
    a = []
    for x in v
        push!(a, ComponentPoolIndex(x, i))
        i+=1
    end
    a
end

tupify(v) = (v,)
veccify(v) = [v]
#+------------------------------------------------------+
#| This is the main idea                                |
#| Combinations are only generated if needed            |
#+------------------------------------------------------+
function create_combinations(scheme, c::ComponentPool)
    arrs = []
    # Get a ComponentPoolIndex version of the array in the ComponentPool
    # for the Type of each element in the scheme.
    # since we're iterating the scheme in order,
    # there is no need to worry about passing the
    # args to the system function in the correct order.
    for x in scheme
        # If one of these has no components defined for its type,
        # then the overall combination will just be empty,
        # so we might as well just return []
        if c[x] == [] return c[x] end
        push!(arrs, indexify(c[x]))
    end
    # if arrs is length 1 then a will be returned straight away,
    # so we need each of its elements to be a tuple for compat.
    a = veccify.(first(arrs))
    k = 2
    # Create pairs until you reach the end of the scheme.
    # This will be ordered since arrs is based on scheme.
    while k<=length(arrs)
        a = create_pair_combinations(a, arrs[k])
        k+=1
    end
    return a
end
create_pair_combinations(a, b) = reduce(vcat, [[ae..., be...] for ae in a for be in b])