export simulate!
export propagate_init!

#---------------constructor first time step-------------------------------------
function propagate_init!(s::Simulation) #check seems okay #everything good.
    compute_eigen_H!(s)
    H_derivatives!(s)
    generate_therm_surf!(s)

    #set initial wavefunction
    @views s.ψ .= convert(Array{ComplexF64, 2}, s.Γ[:, s.surfp])

    # KE = sum(0.5 ./m_spread .* s.v.^2)
    # PE = s.hp[1] + sum(s.Γ[s.surfp])

    #calculate initial forces
    update_forces!(s)
    #fix position of the last layer of atoms
    s.F[399:530, :] .= 0.0
    s.v[399:530, :] .= 0.0

    #convert to Forces
    s.F .= -s.F
    if logopt == 1
        storage_init!(s)
    end
end

#stores data from initialization into first columns of storage arrays
@inline function storage_init!(s::Simulation)
    get_temperature!(s, 1)
    s.storage_K[1] = 0.5 * sum(s.v.^2 .* m_spread)
    s.storage_P[1] = s.hp[1] + sum(s.λ[s.surfp])
    rtemp = norm(s.Δ_no)
    @views s.storage_e[1, 1] = 0.5 * (mass_arr[1] + mass_arr[2]) *
     sum(((mass_arr[1].*s.v[1, :] + mass_arr[2].*s.v[2, :])/(mass_arr[1] + mass_arr[2])).^2)
    @views T_vib = 0.50 * μ * sum(((s.v[1, :] .- s.v[2, :]).*s.Δ_no/rtemp).^2)
    U_vib = F_n *(1.0 - exp(-δ_n * (rtemp -  r_0_NO)))^2
    E_vib = U_vib + T_vib
    s.storage_e[2, 1] = E_vib
    @views T_tot = 0.5 * mass_arr[1] * sum(s.v[1, :].^2) + 0.5*mass_arr[2]*sum(s.v[2, :].^2)
    @views s.storage_e[3, 1] = T_tot - T_vib - s.storage_e[1, 1]
    s.storage_e[4, 1] = sum(s.λ[s.surfp]) - sum(s.λ[s.surfpinit])       #das hier macht keinen Sinn
    @views s.storage_xno[1:3, 1] = s.x[1, :]
    @views s.storage_xno[4:6, 1] = s.x[2, :]
    @views s.storage_vno[1:3, 1] = s.v[1, :]
    @views s.storage_vno[4:6, 1] = s.v[2, :]#check seems okay
end

@inline function compute_eigen_H!(s::Simulation)
    s.H .= Symmetric(h0 .+ vm .* s.hp[3] ./ sqrt_de)
    s.H[1, 1] = s.H[1, 1] + s.hp[2] - s.hp[1]
    s.λ .= eigvals(s.H)
    s.Γ .= eigvecs(s.H)#check seems okay

end

@inbounds @inline function H_derivatives!(s::Simulation) #seems ok, fixed error
    for i in 1:Ms+1
        s.dhdea[:, i] .= view(s.Γ, 1, :) .* s.Γ[1, i]
    end
    temp = vm * s.Γ  #needs preallocation
    mul!(s.dhdv, transpose(s.Γ), temp)
    s.dhdv .= s.dhdv ./ sqrt_de
end

@inline @inbounds function generate_therm_surf!(s::Simulation) #check seems okay
    #generate thermalized populations for electronic states
    for j in 1:Ne
        s.surfp[j] = j
    end

    for j in 1:Ne
        s.occnum[s.surfp[j]] = 1
    end

    k = 1
    for j in 1:Ms+1
        if s.occnum[j] == 0
            s.surfh[k] = j
            k = k + 1
        end
    end

    for t in 1:Ne
        for k in 1:Ne*(Ms + 1 - Ne)
            hoprand = rand(Float64, 3)
            #hoprand = collect([0.2, 0.4, 0.6])
            jp = Int(floor(hoprand[1]*Ne)) + 1
            jh = Int(floor(hoprand[2] * (Ms + 1 - Ne))) + 1
            rtemp = exp(-(s.λ[s.surfh[jh]] - s.dhp_neutral[s.surfp[jp]])) / (kb * tsurf)
            if hoprand[3] < rtemp
                m = s.surfh[jh]
                s.surfh[jh] = s.surfp[jp]
                s.surfp[jp] = m
            end
        end
    end

    s.surfpinit .= deepcopy(s.surfp)
end

@inline @inbounds function update_forces!(s::Simulation) #check seems okay
    for j in 1:Ne
        s.F .= s.F .+ (s.dhdea[s.surfp[j], s.surfp[j]] .* (s.dhp_ion .- s.dhp_neutral) .+
        s.dhdv[s.surfp[j], s.surfp[j]] .* s.dhp_coup)

        # @views s.F .= s.F .+ s.dhdv[s.surfp[j], s.surfp[j]] .* s.dhp_coup
    end
    s.F .= s.F .+ s.dhp_neutral
end

#---------------------simulation main time loop for n=2:tsteps-----------------
function simulate!(s::Simulation)
    for n in 2:tsteps
        do_we_stop = v_z_condition_check(s)
        if do_we_stop == -1
            break
        end
        # Record minimial distance of z approach and orientational angle
        @views min_z_no = minimum(s.x[1:2, 3])
        if min_z_no <= s.trajzmin[1]
            s.trajzmin .= min_z_no
            s.trajtheta .= acos(s.Δ_no[3]/norm(s.Δ_no))
        end
        t_prog = round(n/tsteps * 100.0; digits=3)
        str_t_prog = string(t_prog)*"%"
        traj_prog = round(min_z_no/Å; digits=3)
        str_trajzmin_prog = string(traj_prog)
        println(str_t_prog*", "*str_trajzmin_prog)
        # get state corresponding to current surface
        @inbounds for j in 1:Ne
            s.ϕ[:, j] = view(s.Γ, :, s.surfp[j])
        end

        #Calculate nonadiabatic coupling matrix DM between adiabatic orbitals
        get_dm!(s::Simulation)
        # Get occupied (particle), unoccupied (hole) states corresponding to current surface specified by surfp
        # surfp is Ne x 1 array of particle states
        # surfh is (Ms+1-Ne) x 1 array of hole states
        get_occupied_unoccpied_states!(s)
        # Get b_{kl} where k is current many-electron adiabatic state and l is possible new
        # 	many-electron adiabatic state
        # k corresponds to occupation of states specified by surfp
        # l corresponds to occupation of states specified by a one electron-hole-pair excitation out of surfp
        get_blk2_pbmaxest!(s)
        # Calculate upper bound of maximum hopping probability, Pbmaxest

        #beginn hopping
        #Generate random number for surface hopping
        hoprand = rand()
        #! Only attempt hop if hoprand < Pbmaxest to avoid expensive calculation
        #! of real hopping elements blk
        if hoprand < s.Pbmaxest[1]
            hopping!(s, hoprand, n)
        end
        # Propagate electronic Hamiltonian
        uu = MVector{Ms+1, ComplexF64}(zeros(ComplexF64, Ms+1))
        uuu = zeros(ComplexF64, Ms+1, Ne)
        propagate_hamiltonian!(s, uu, uuu)
        #storage
        if logopt == 1
            storage_sim1!(s, n)
        end
        #propagate nuclear motion
        propagate_nuclear!(s)
        #get adiabatic eigenvalues, eigenvectors
        get_eigen_sim!(s)
        # Get dH / d Ep matrices where Ep is parameter in Newns-Anderson Hamiltonian
        # (i.e. Ep = V or Ep = E_a = (EI-EN) )
        get_dhdea_dhdv_loop!(s)

        # Calculate forces
        s.F .= 0
        update_forces!(s)
        s.F[399:530, :] .= 0
        s.v[399:530, :] .= 0
        s.F .= -s.F

        s.v .= s.vtemp .+ 0.5 ./m_spread .* s.F * dt
        if logopt == 1
            @views s.storage_xno[1:3, 1] = s.x[1, :]
            @views s.storage_xno[4:6, 1] = s.x[2, :]
            @views s.storage_vno[1:3, 1] = s.v[1, :]
            @views s.storage_vno[4:6, 1] = s.v[2, :]#check seems okay
        end
        s.nf .= s.nf .+ one(Int64)
    end
    s.storage_e[5, :] = sum(s.storage_e[1:4, :], dims=1) .- sum(s.storage_e[1:4, 1])

end
#-------------------------------------------------------------------------------

#check if the projectile has reached simulation boundaries
function v_z_condition_check(s::Simulation) #should be ok
    @views v_cond = (s.v[1, 3]*m_N + s.v[2, 3]*m_O)/(m_N + m_O)
    @views z_cond = (s.x[1, 3] + s.x[2, 3])/2.0
    if z_cond > z_end && v_cond > 0.0
        println("Boundary reached:End Simulation")
        return -1
    end
    return 1
end

@inbounds @inline function get_dm!(s::Simulation) #should be ok
    # Calculate nonadiabatic coupling elements DM between adiabatic orbitals s.dm
    # phipsi = <phi_k|psi> is overlap between "current" adiabatic state |phi_k> and electronic state |psi>
    temp_1 = transpose(s.ϕ) * s.ψ   #preallocation
    s.phipsi[1] = det(temp_1)
    # Calculate akk as defined in Tully JCP 1990
    s.akk[1] = abs(s.phipsi[1])^2 #corrected error

    # Calculate nonadiabatic coupling elements DM between adiabatic orbitals

    temp_2 = s.dhp_ion .- s.dhp_neutral #preallocation
    temp_2 .= temp_2 .* s.v
    temp_2_sum = sum(temp_2)
    s.dm .= temp_2_sum .* s.dhdea

    temp_3 = s.dhp_coup .* s.v
    temp_3_sum = sum(temp_3)
    s.dm .= s.dm .+ temp_3_sum .* s.dhdv

    temp_rep1 = repeat(s.λ, 1, Ms+1)
    temp_rep2 = temp_rep1' .- temp_rep1
    s.dm .= s.dm ./ temp_rep2
    for j in 1:Ms+1
        s.dm[j,j] = zero(eltype(s.dm))
    end #should be ok
end

@inline @inbounds function get_occupied_unoccpied_states!(s::Simulation) #should be ok
    s.occnum .= 0
    for j in 1:Ne
        s.occnum[s.surfp[j]] = 1
    end
    k = 1
    for j in 1:Ms+1
        if s.occnum[j] == 0
            s.surfh[k] = j
            k = k+1
        end
    end #should be ok
end

@inline @inbounds function get_blk2_pbmaxest!(s::Simulation) #should be ok
    rtemp = abs(s.phipsi[1] * s.phipsi[1])
    if rtemp > one(eltype(rtemp))
        rtemp = 1.0
    end
    for jp in 1:Ne
        for jh in 1:Ms+1-Ne
            s.blk2[jp, jh] = 2.0 * abs(s.phipsi[1] * s.dm[s.surfp[jp], s.surfh[jh]])
        end
    end

    s.Pbmaxest[1] = sqrt(sum(s.blk2.^2)) * sqrt(1.0 - rtemp)/s.akk[1] * dt * Float64(thop) #should be ok
end

#---------------------start hopping! hoprand < pbmaxest -----------------------
function hopping!(s::Simulation, hoprand::Float64, n::Int64)
    hopping!_pbmaxest!(s, hoprand, n)
end

function hopping!_pbmaxest!(s::Simulation, hoprand::Float64, n::Int64)
    get_blk_akl!(s)
    get_pb!(s)

    s.storage_phop[n] = s.Pb[Ne*(Ms+1-Ne)]
    #attempt hop conditioned on random number
    if hoprand < s.Pb[Ne*(Ms+1-Ne)]
        print("Hopping! ")
        hopping!_pbmaxest!_attempt!(s, hoprand, n)
    end
end

@inline function get_blk_akl!(s::Simulation) #seems ok
    ctemp1 = Matrix{ComplexF64}(undef, Ne, Ne)
    for jp in 1:Ne
        for jh in 1:Ms+1-Ne
            s.surfpnew .= s.surfp
            s.surfpnew[jp] = s.surfh[jh]
            for j in 1:Ne
                @views s.ϕ[:, j] = s.Γ[:, s.surfpnew[j]]
            end
            #Ctemp = <phi_l|psi> is overlap between "new" adiabatic state
            #|phi_l> and electronic state |psi>
            mul!(ctemp1, transpose(s.ϕ) , s.ψ)
            ctemp = det(ctemp1)
            s.akl[1] = s.phipsi[1] * conj(ctemp)
            @views s.blk[jp, jh] = 2.0 * real(s.akl[1] * s.dm[s.surfp[jp], s.surfh[jh]])
        end
    end
end

@inline function get_pb!(s::Simulation) #seems ok
    rtemp = 0.0
    for jh in 1:Ms+1-Ne
        for jp in 1:Ne
            s.Pb[(jh-1) * Ne + jp] = rtemp
            @views temp = s.blk[jp,jh]
            if temp <= 0.0
                temp = 0.0
            end
            s.Pb[(jh-1) * Ne + jp] = s.Pb[(jh-1)*Ne + jp] + temp
            rtemp = s.Pb[(jh - 1) * Ne + jp]
        end
    end
    s.Pb .= s.Pb ./ s.akk[1] * dt * Float64(thop)
end
#-----------------start hopping! hoprand < pbmaxest && hoprand < Pbmax ---------
function hopping!_pbmaxest!_attempt!(s::Simulation, hoprand::Float64, n::Int64)
    hopstate = which_state(s, hoprand) #ok
    jh= which_orbital_jh(hopstate) #ok
    jp= which_orbital_jp(hopstate) #ok
    rhop = which_direction(s, jp, jh) #needs durther testing, not sure about rescaling
    e_el_old, e_el_new, k_rhop, k_temp = check_energy1(s, jp, jh, rhop)
    bool_e = check_energy2(e_el_old, e_el_new, k_rhop, k_temp)
    s.attnum .= s.attnum .+ one(Int64)

    if bool_e == true
        update_vel_populations!(s, jp, jh, rhop, k_temp, k_rhop, e_el_old, e_el_new, n) #needs more storage
    end
end

@inline function which_state(s::Simulation, hoprand::Float64)::Int64 #ok
    hopstate = one(Int64)
    while s.Pb[hopstate] <= hoprand
        hopstate = hopstate + one(Int64)
    end
    return hopstate
end

@inline function which_orbital_jh(hopstate::Int64)::Int64 #ok
    jh = Int(floor((hopstate - 1) / Ne)) + one(Int64)
    return jh
end
@inline function which_orbital_jp(hopstate::Int64)::Int64 #ok
    jp = ((hopstate - one(Int64)) % Ne) + one(Int64)
    return jp
end

@inline function which_direction(s::Simulation, jp::Int64, jh::Int64)::float_array #needs checking
    @views temp = s.dhdea[s.surfp[jp], s.surfh[jh]] .* (s.dhp_ion .- s.dhp_neutral)
    @views temp .= temp .+ s.dhdv[s.surfp[jp], s.surfh[jh]] .* s.dhp_coup
    temp2 = sqrt(sum(temp.^2)/sum(s.v .^2))
    temp .= temp ./ temp2
    temp3 = sum(temp .* s.v)
    temp .= temp .* copysign(1.0, temp3)
    return temp
end

@inline function check_energy1(s::Simulation, jp::Int64, jh::Int64, rhop::float_array) # seems ok
    s.surfpnew .= s.surfp
    s.surfpnew[jp] = s.surfh[jh]
    e_el_old = sum(s.λ[s.surfp])
    e_el_new = sum(s.λ[s.surfpnew])
    k_rhop = 0.5 * sum(1.0 ./ m_spread .* rhop .^2)
    k_temp = sum(s.v .* rhop)

    return e_el_old, e_el_new, k_rhop, k_temp
end

@inline function check_energy2(e_el_old::Float64, e_el_new::Float64, k_rhop::Float64, k_temp::Float64)::Bool
    bool_e = k_temp^2 - 4.0*k_rhop * (e_el_new - e_el_old) > 0.0
    return bool_e
end

@inline function update_vel_populations!(s::Simulation, jp::Int64, jh::Int64, rhop::float_array,
    k_temp::Float64, k_rhop::Float64, e_el_old::Float64, e_el_new::Float64, n::Int64)
    rtemp = s.vscale[1]
    s.vscale[1] = (k_temp - sqrt(k_temp^2 - 4.0 * k_rhop * (e_el_new - e_el_old)))/2.0/k_rhop
    @views k_no = 0.50 * sum(view(m_spread, 1:2, :).* view(s.v, 1:2, :).^2)
    @views k_au = 0.50 * sum(view(m_spread, 3:N+2, :) .* view(s.v, 3:N+2, :).^2)
    s.v .= s.v .- s.vscale[1] .* (1.0 ./m_spread .* rhop)
    s.storage_deltaKNO[s.exnum[1]+1] = k_no - 0.5*sum(view(m_spread, 1:2, :) .* view(s.v, 1:2, :) .^2)
    s.storage_deltaKAu[s.exnum[1]+1] = k_au - 0.5*sum(view(m_spread, 3:N+2, :) .* view(s.v, 3:N+2, :) .^2)
    s.storage_state[2*s.exnum[1] + 1] = s.surfp[jp]
    s.storage_state[2*s.exnum[1] + 2] = s.surfh[jh]
    s.surfp .= s.surfpnew
    s.storage_hoptimes[s.exnum[1]+1] = n
    s.exnum .= s.exnum .+ one(Int64)
    #endif
end
#-----------------END hopping! hoprand < pbmaxest && hoprand < Pbmax ---------

@inline function propagate_hamiltonian!(s::Simulation, uu::MVector{Ms+1, ComplexF64},
     uuu::Matrix{ComplexF64}) # seems ok
    s.ψ .= transpose(s.Γ) * s.ψ

    uu .= exp.(-1im * s.λ * dt/hbar)
    uuu .= repeat(uu, 1, Ne)
    s.ψ .= uuu .* s.ψ
    s.ψ .= s.Γ * s.ψ
end

#storage_op_e
@inline function storage_sim1!(s::Simulation, n::Int64) #seems ok
    get_temperature!(s, n)
    s.storage_K[n] = 0.5 * sum(s.v .^2 .* m_spread)
    s.storage_P[n] = s.hp[1] + sum(s.λ[s.surfp])
    @views s.storage_aop[n] = sum(abs.(s.ψ[1, :]).^2) #rhoa
    store_temp = sum(abs.(transpose(s.Γ) * s.ψ).^2, dims=2)
    s.storage_op[:, n] .= store_temp[:, 1] #psip

    rtemp = norm(s.Δ_no)
    s.storage_e[1, n] = 0.5 * (mass_arr[1] + mass_arr[2]) *
     sum(((mass_arr[1].*s.v[1, :] + mass_arr[2].*s.v[2, :])/(mass_arr[1] + mass_arr[2])).^2)
    @views T_vib = 0.50 * μ * sum(((s.v[1, :] .- s.v[2, :]).*s.Δ_no/rtemp).^2)
    U_vib = F_n *(1.0 - exp(-δ_n * (rtemp -  r_0_NO)))^2
    E_vib = U_vib + T_vib
    s.storage_e[2, n] = E_vib
    @views T_tot = 0.5 * mass_arr[1] * sum(s.v[1, :].^2) + 0.5*mass_arr[2]*sum(s.v[2, :].^2)
    s.storage_e[3, n] = T_tot - T_vib - s.storage_e[1, n]
    s.storage_e[4, n] = sum(s.λ[s.surfp]) - sum(s.λ[s.surfpinit])

end
@inline function propagate_nuclear!(s::Simulation)
    s.vdot .= 1.0./m_spread .* s.F
    s.x .= s.x .+ s.v * dt .+ 0.50*s.vdot * dt^2
    s.Δ_no .=  s.x[1, :] - s.x[2, :]
    s.vtemp .= s.v .+ 0.50*s.vdot*dt #ARRAY ALLOCATION !!!! solve later
end

@inline function get_eigen_sim!(s::Simulation)
    s.hp .= get_E_all(s)    #update energy surface
    get_F_all(s)            #update partial derivatives dhp_i
    s.H .= h0 .+ vm .* s.hp[3] ./sqrt_de
    s.H[1, 1] = s.H[1, 1] + s.hp[2] - s.hp[1]
    Γ_hold = deepcopy(s.Γ) #ARRAY ALLOCATION !!!! solve later
    s.λ .= eigvals(s.H)
    s.Γ .= eigvecs(s.H)
    temp = zeros(Float64, Ms+1)
    @inbounds for j in 1:Ms+1
        @views temp .= abs.(transpose(Γ_hold) * s.Γ[:, j])
        Γmaxloc = argmax(temp)
        temp2 = dot(Γ_hold[:, Γmaxloc], s.Γ[:, j])
        s.Γ[:, j] = s.Γ[:, j]/copysign(1.0, temp2)
    end
end

function get_dhdea_dhdv_loop!(s::Simulation)
    s.blk .= 0.0
    @inbounds for j in 1:Ms+1
        s.dhdea[:, j] = s.Γ[1, :] .* s.Γ[1, j]
    end
    temp = (vm * s.Γ) #ARRAY ALLOCATION !!!! solve later
    mul!(s.dhdv, transpose(s.Γ), temp)
    s.dhdv .= s.dhdv ./sqrt_de
end

function get_temperature!(s::Simulation, n::Int)
    speed = [norm(s.v[i, :]) for i in 3:398]
    s.storage_temp[n] = 2*m_au*sum(speed.^2)/(3*(N*0.75-3)*kb)
end


module FastAI

using Base: NamedTuple
using Colors: colormaps_sequential
using Reexport
@reexport using DLPipelines
@reexport using FluxTraining
@reexport using DataLoaders
@reexport using Flux

using Animations
using Makie
using Colors
using DataAugmentation
using DataAugmentation: getbounds, Bounds
import DLPipelines: methoddataset, methodmodel, methodlossfn, methoddataloaders,
    mockmodel, mocksample, predict, predictbatch, mockmodel, encode, encodeinput,
    encodetarget, decodeŷ, decodey
using IndirectArrays: IndirectArray
using LearnBase: getobs, nobs
using FilePathsBase
using FixedPointNumbers
using Flux
using Flux.Optimise
import Flux.Optimise: apply!, Optimiser, WeightDecay
using FluxTraining: Learner, handle
using FluxTraining.Events
using JLD2: jldsave, jldopen
using Markdown
using MLDataPattern
using Parameters
using PrettyTables
using StaticArrays
using Setfield
using ShowCases
using Tables
using Statistics
using Test: @testset, @test, @test_nowarn

include("plotting.jl")
include("learner.jl")

# Data block API
include("datablock/block.jl")
include("datablock/encoding.jl")
include("datablock/method.jl")
include("datablock/describe.jl")
include("datablock/checks.jl")
include("datablock/wrappers.jl")

# submodules
include("datasets/Datasets.jl")
@reexport using .Datasets

include("models/Models.jl")
using .Models

# Encodings
include("encodings/tabularpreprocessing.jl")
include("encodings/onehot.jl")
include("encodings/imagepreprocessing.jl")
include("encodings/projective.jl")
include("encodings/keypointpreprocessing.jl")

# Training interface
include("datablock/models.jl")
include("datablock/loss.jl")
include("datablock/plot.jl")

# training
include("training/paramgroups.jl")
include("training/discriminativelrs.jl")
include("training/utils.jl")
include("training/onecycle.jl")
include("training/finetune.jl")
include("training/lrfind.jl")
include("training/metrics.jl")

include("serialization.jl")

include("fasterai/methodregistry.jl")
include("fasterai/learningmethods.jl")
include("fasterai/defaults.jl")

export
    # submodules
    Datasets,
    Models,
    datasetpath,
    loadtaskdata,
    mapobs,
    groupobs,
    filterobs,
    shuffleobs,
    datasubset,

    # method API
    methodmodel,
    methoddataset,
    methoddataloaders,
    methodlossfn,
    getobs,
    nobs,
    predict,
    predictbatch,

    # plotting API
    plotbatch,
    plotsample,
    plotsamples,
    plotpredictions,
    makebatch,

    # blocks
    Image,
    Mask,
    Label,
    LabelMulti,
    Keypoints,
    Many,
    TableRow,
    Continuous,

    # encodings
    encode,
    decode,
    setup,
    ProjectiveTransforms,
    ImagePreprocessing,
    OneHot,
    KeypointPreprocessing,
    Only,
    Named,
    augs_projection, augs_lighting,
    TabularPreprocessing,

    BlockMethod,
    describemethod,
    checkblock,

    # learning methods
    findlearningmethods,
    ImageClassificationSingle,
    ImageClassificationMulti,
    ImageSegmentation,
    ImageKeypointRegression,
    TabularClassificationSingle,
    TabularRegression,

    # training
    methodlearner,
    Learner,
    fit!,
    fitonecycle!,
    finetune!,
    lrfind,
    savemethodmodel,
    loadmethodmodel,
    accuracy_thresh,

    gpu,
    plot

end  # module


# %% imports

using JLD
using EvolutionaryStrategies

include("cartpole_fitness.jl")
#include("swingup_fitness.jl")
#include("mountcar_fitness.jl")
include("cambrianrun.jl")

# %% create snestopgen
init = 0
save("Julia/Actors/topgenxNES.jld","init", init)
# %% xNES optimisation

for opti in 1:50
    cfg = get_config("Julia/Actors/cartpole_esconfig.yml")
    es = xNES(cfg,fitness)
    a = run!(es, opti)
    println("optim $(opti)")
end

# %% plot the distribution of gen convergence

f = load("Julia/Actors/topgenxNES.jld")
gens = Int64[]
for i in 1:length(f)-1
    push!(gens, f["$(i)_run"])
end

histogram(f, bins = 0:10:300,
            label = "First R>195 individual (mean = 82.68)",
            color = :green,
            xlabel = "Generation",
            ylabel = "Number of optimisations")


"""
    Image{N}() <: Block

[`Block`](#) for an N-dimensional image. `obs` is valid for `Image{N}()`
if it is an N-dimensional array with color or number element type.

## Examples

Creating a block:

```julia
Image{2}()  # 2D-image
Image{3}()  # 3D-image
```

Example valid images:

```julia
@test checkblock(Image{2}(), rand(RGB, 10, 10))         # Color image
@test checkblock(Image{2}(), rand(10, 10))              # Numbers treated as grayscale
@test checkblock(Image{3}(), rand(Gray{N0f8}, 10, 10, 10))  # Grayscale 3D-image
```

The color channels (if any) are not counted as a dimension and represented
through color types like `RGB{N0f8}`:

```julia
@test !checkblock(Image{2}, rand(10, 10, 3))  # Not a 2D image
```

You can create a random observation using [`mockblock`](#):

{cell=main}
```julia
FastAI.mockblock(Image{2}())
```

To visualize a 2D-image observation, use [`showblock`](#). This is supported for
both the `ShowText` and the `ShowMakie` backend.

```julia
showblock(Image{2}(), rand(RGB{N0f8}, 10, 10))
```

"""
struct Image{N} <: Block end

checkblock(::Image{N}, ::AbstractArray{T,N}) where {T<:Union{Colorant,Number},N} = true
mockblock(::Image{N}) where {N} = rand(RGB{N0f8}, ntuple(_ -> 16, N))

setup(::Type{Image}, data) = Image{ndims(getobs(data, 1))}()

# Visualization

function showblock!(io, ::ShowText, block::Image{2}, obs)
    ImageInTerminal.imshow(io, obs)
end


function FieldDelaunayArguments(arg0::FieldAbsoluteDate, arg1::RealFieldElement, arg10::RealFieldElement, arg11::RealFieldElement, arg12::RealFieldElement, arg13::RealFieldElement, arg2::RealFieldElement, arg3::RealFieldElement, arg4::RealFieldElement, arg5::RealFieldElement, arg6::RealFieldElement, arg7::RealFieldElement, arg8::RealFieldElement, arg9::RealFieldElement)
    return FieldDelaunayArguments((FieldAbsoluteDate, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg10, arg11, arg12, arg13, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end


import Base.+;

const MarketDataSet = Dict{String, AbstractMonteCarloProcess};
export MarketDataSet;

#Link an underlying to a single name model
function →(x::String, y::AbstractMonteCarloProcess)
    out = MarketDataSet(x => y)
    return out
end

#Link an multiname underlying to a multi variate model
function →(x::String, y::NDimensionalMonteCarloProcess)
    sep = findfirst("_", x)
    if (isnothing(sep))
        error("Nvariate process must follow the format: INDEX1_INDEX2")
    end
    len_ = length(y.models)
    sep = split(x, "_")
    @assert len_ == length(sep)
    out = MarketDataSet(x => y)

    return out
end

#Joins two different MarketDataSet s
function +(x::MarketDataSet, y::MarketDataSet)
    out = copy(x)
    y_keys = keys(y)
    for y_key in y_keys
        if haskey(out, y_key) || any(out_el -> any(out_el_sub -> out_el_sub == y_key, split(out_el, "_")), keys(out)) || any(y_key_el -> haskey(out, y_key_el), split(y_key, "_"))
            error("check keys")
        else
            out[y_key] = y[y_key]
        end
    end
    return out
end


#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*
#*----------------------------------------------------------------------
#* File containg models of columns: distillation, stripping, absorbers
#* rectifier, ....
#*
#* The default nomenclature is:
#*		Type_Column_reboilertype_condensertyper
#*
#* where:
#*	Type = refluxed or reboiled or section
#*	Column = Stripping, Absorption, Rectifier, Distillation
#*	Reboiler type (if exists) = kettle or thermosyphon
#*	Condenser type (if exists) = with subccoling or without subcooling
#*
#*-----------------------------------------------------------------------
#* Author: Paula B. Staudt
#* $Id: column.mso 511 2008-05-12 17:25:33Z paula $
#*---------------------------------------------------------------------
type InitializeSection
	InitializeSection()=begin
		NComp=outers.NComp
		new(
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of components"
			)),
			temperature(Dict{Symbol,Any}(
				:Brief=>"Tray Temperature at Column Top",
				:Default=>300
			)),
			temperature(Dict{Symbol,Any}(
				:Brief=>"Tray Temperature at Column Bottom",
				:Default=>330
			)),
			fraction(Dict{Symbol,Any}(
				:Brief=>"Tray Level Fraction",
				:Default=>0.5
			)),
			fill(fraction(Dict{Symbol,Any}(
				:Brief=>"Component Molar Fraction at Column Top",
				:Default=>0.30
			)),(NComp)),
			fill(fraction(Dict{Symbol,Any}(
				:Brief=>"Component Molar Fraction at Column Bottom",
				:Default=>0.30
			)),(NComp)),
			[:NComp,:TopTemperature,:BottomTemperature,:LevelFraction,:TopComposition,:BottomComposition,],
		)
	end
	NComp::DanaInteger
	TopTemperature::temperature
	BottomTemperature::temperature
	LevelFraction::fraction
	TopComposition::Array{fraction}
	BottomComposition::Array{fraction}
	parameters::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export InitializeSection
function atributes(in::InitializeSection,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Brief]="Initial conditions for column section."
	drive!(fields,_)
	return fields
end
InitializeSection(_::Dict{Symbol,Any})=begin
	newModel=InitializeSection()
	newModel.attributes=atributes(newModel,_)
	newModel
end


module SingleHouseholdCalculations

#
# This module drives all the calculations for single household and a single set of parameters.
#

using ScottishTaxBenefitModel

using .Definitions

using .RunSettings:
    Settings

using .Results:
    IndividualResult,
    BenefitUnitResult,
    HouseholdResult,
    init_household_result,
    aggregate!

using .STBParameters:
    TaxBenefitSystem

using .ModelHousehold:
    BenefitUnit,
    BenefitUnits,
    BUAllocation,
    Household,
    People_Dict,
    PeopleArray,
    Person,
    default_bu_allocation,
    get_benefit_units,
    get_head,
    get_spouse,
    num_people,
    printpids

# using .LegacyMeansTestedBenefits

using .NonMeansTestedBenefits:
    calc_pre_tax_non_means_tested!,
    calc_post_tax_non_means_tested!

using .IncomeTaxCalculations:
    calc_income_tax!

using .NationalInsuranceCalculations:
    calculate_national_insurance!

using .STBIncomes

using .Intermediate:
    MTIntermediate,
    make_intermediate

using .Results:
    BenefitUnitResult,
    HouseholdResult,
    IndividualResult,
    ITResult,
    NIResult

using .LocalLevelCalculations:
    calc_council_tax

using .LegacyMeansTestedBenefits:
    calc_legacy_means_tested_benefits!

using .UniversalCredit:
    calc_universal_credit!

using .BenefitCap:
    apply_benefit_cap!

using .UCTransition: route_to_uc_or_legacy!

using .ScottishBenefits:
    calc_scottish_child_payment!,
    calc_bedroom_tax_mitigation!

using .UBI: calc_UBI!, make_ubi_post_adjustments!

export do_one_calc

"""
One complete calculation for a single household and tb system.
"""
function do_one_calc(
    hh :: Household{T},
    sys :: TaxBenefitSystem{T},
    settings :: Settings = Settings() ) :: HouseholdResult{T} where T
    # hh = deepcopy( mhh ) # for minwage and so on

    bus = get_benefit_units( hh )
    hres :: HouseholdResult{T} = init_household_result(hh)
    intermed = make_intermediate(
        hh,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits )

    hd :: BigInt = get_head( hh ).pid
    calc_pre_tax_non_means_tested!(
        hres,
        hh,
        sys.nmt_bens,
        sys.hours_limits,
        sys.age_limits )
    calc_UBI!(
        hres,
        hh,
        sys.ubi )
    buno = 1
    for bu in bus

        # income tax, with some nonsense for
        # what remains of joint taxation..
        head = get_head( bu )
        spouse = get_spouse( bu )
        calc_income_tax!(
            hres.bus[buno],
            head,
            spouse,
            sys.it )
        for chno in bu.children
            child = bu.people[chno]
            calc_income_tax!(
                hres.bus[buno].pers[child.pid],
                child,
                sys.it )
        end
        for (pid,pers) in bu.people
            if pers.age >= 16 # must be 16+
                calculate_national_insurance!(
                    hres.bus[buno].pers[pers.pid],
                    pers,
                    sys.ni )
                end
        end
        buno += 1
    end # bus loop
    calc_post_tax_non_means_tested!(
        hres,
        hh,
        sys.nmt_bens,
        sys.age_limits )

    hres.bus[1].pers[hd].income[LOCAL_TAXES] =
        calc_council_tax( hh, intermed.hhint, sys.loctax.ct )

    calc_legacy_means_tested_benefits!(
        hres,
        hh,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr )

    calc_universal_credit!(
        hres,
        hh,
        intermed,
        sys.uc,
        sys.age_limits,
        sys.hours_limits,
        sys.child_limits,
        sys.hr,
        sys.minwage
    )

    route_to_uc_or_legacy!(
        hres,
        settings,
        hh,
        intermed )

    for buno in eachindex( bus )
        calc_scottish_child_payment!(
            hres.bus[buno],
            bus[buno],
            intermed.buint[buno],
            sys.scottish_child_payment )

        apply_benefit_cap!(
            hres.bus[buno],
            hh.region,
            bus[buno],
            intermed.buint[buno],
            sys.bencap,
            hres.bus[buno].route )
    end
    # do this after the benefit cap
    # since the DISCRESIONARY_HOUSING_PAYMENT must be <= hb/uc housing costs
    calc_bedroom_tax_mitigation!( hres, hh )
    if ! sys.ubi.abolished
        make_ubi_post_adjustments!( hres, sys.ubi )
    end
    aggregate!( hh, hres )
    return hres
end

end # module


@sumProductRule(:node_type     => Transition,
                :outbound_type => Message{Categorical},
                :inbound_types => (Nothing, Message{PointMass}, Message{PointMass}),
                :name          => SPTransitionOutNPP)

@sumProductRule(:node_type     => Transition,
                :outbound_type => Message{Categorical},
                :inbound_types => (Message{PointMass}, Nothing, Message{PointMass}),
                :name          => SPTransitionIn1PNP)

@sumProductRule(:node_type     => Transition,
                :outbound_type => Message{Categorical},
                :inbound_types => (Nothing, Message{Categorical}, Message{PointMass}),
                :name          => SPTransitionOutNCP)

@sumProductRule(:node_type     => Transition,
                :outbound_type => Message{Categorical},
                :inbound_types => (Message{Categorical}, Nothing, Message{PointMass}),
                :name          => SPTransitionIn1CNP)

@naiveVariationalRule(:node_type     => Transition,
                      :outbound_type => Message{Categorical},
                      :inbound_types => (Nothing, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VBTransitionOut)

@naiveVariationalRule(:node_type     => Transition,
                      :outbound_type => Message{Categorical},
                      :inbound_types => (ProbabilityDistribution, Nothing, ProbabilityDistribution),
                      :name          => VBTransitionIn1)

@naiveVariationalRule(:node_type     => Transition,
                      :outbound_type => Message{Dirichlet},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, Nothing),
                      :name          => VBTransitionA)

@structuredVariationalRule(:node_type     => Transition,
                           :outbound_type => Message{Categorical},
                           :inbound_types => (Nothing, Message{Categorical}, ProbabilityDistribution),
                           :name          => SVBTransitionOutVCD)

@structuredVariationalRule(:node_type     => Transition,
                           :outbound_type => Message{Categorical},
                           :inbound_types => (Message{Categorical}, Nothing, ProbabilityDistribution),
                           :name          => SVBTransitionIn1CVD)

@structuredVariationalRule(:node_type     => Transition,
                           :outbound_type => Message{Dirichlet},
                           :inbound_types => (ProbabilityDistribution, Nothing),
                           :name          => SVBTransitionADV)

@marginalRule(:node_type     => Transition,
              :inbound_types => (Message{Categorical}, Message{Categorical}, ProbabilityDistribution),
              :name          => MTransitionCCD)

@marginalRule(:node_type     => Transition,
              :inbound_types => (Message{Categorical}, Message{Categorical}, Nothing), # "Nothing" indicates marginalization
              :name          => MTransitionCCN)


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

function create_root_model!(optimizer::MOI.ModelLike, jp::JuniperProblem; fix_start=false)
    ps = jp.options.log_levels

    jp.model = MOI.instantiate(jp.nl_solver, with_bridge_type=Float64)
    index_map = MOI.copy_to(jp.model, IntegerRelaxation(optimizer); copy_names=false)
    # all continuous we solve relaxation first
    jp.x = [index_map[vi] for vi in MOI.get(optimizer, MOI.ListOfVariableIndices())]
end

function fix_primal_start!(jp::JuniperProblem)
    lb = jp.l_var
    ub = jp.u_var
    x = jp.x
    for i=1:jp.num_var
        if jp.var_type[i] != :Cont && lb[i] <= jp.primal_start[i] <= ub[i]
            set_bounds(jp.model, jp.x[i], jp.primal_start[i], jp.primal_start[i])
        else
            MOI.set(jp.model, MOI.VariablePrimalStart(), jp.x[i], jp.primal_start[i])
        end
    end
end

function unfix_primal_start!(jp::JuniperProblem)
    lb = jp.l_var
    ub = jp.u_var
    x = jp.x
    for i=1:jp.num_var
        set_bounds(jp.model, jp.x[i], lb[i], ub[i])
        MOI.set(jp.model, MOI.VariablePrimalStart(), jp.x[i], jp.primal_start[i])
    end
end

function solve_root_incumbent_model(jp::JuniperProblem)
    MOI.optimize!(jp.model)
    status = MOI.get(jp.model, MOI.TerminationStatus())
    incumbent = nothing
    ps = jp.options.log_levels
    if state_is_optimal(status; allow_almost=jp.options.allow_almost_solved)
        # set incumbent
        objval = MOI.get(jp.model, MOI.ObjectiveValue())
        solution = MOI.get(jp.model, MOI.VariablePrimal(), jp.x)
        if are_type_correct(solution, jp.var_type, jp.disc2var_idx, jp.options.atol)
            if only_almost_solved(status) && jp.options.allow_almost_solved_integral
                @warn "Start value incumbent only almost locally solved. Disable with `allow_almost_solved_integral=false`"
            end

            incumbent = Incumbent(objval, solution, only_almost_solved(status))
            (:All in ps || :Info in ps) && println("Incumbent using start values: $objval")
        end
    else
        (:All in ps || :Info in ps) && println("Start values are not feasible.")
    end
    return incumbent
end

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

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