include("../src/BlockCiphers.jl")

importall BlockCiphers

print("benching AES ")

plain = zeros(Uint8, 16)
k1 = AES(zeros(Uint8, int32(128/8)))

buff = zeros(Uint8, 16 * 1000000)

function gigs(len :: Int64)
    for i=0:(len - 1)
        s = slice(buff, ((i*16)+1) : ((i+1) * 16))
        encrypt!(k1, s, s)
    end
    show(buff[1:16])
    println()
end

gigs(1000)

@time gigs(100000)
@time gigs(1000000)
@time gigs(1000000)

println()


using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CSV
using DataFrames
using Turing
using Dates
using StatsBase
using Distributions
using CovidSurvey
using Underscores
# using Gadfly
using StatsPlots
using LaTeXStrings
using NamedTupleTools
using StatsFuns
using Random
using TransformVariables, LogDensityProblems
using Optim
using ReverseDiff
using BSON
import TimeSeriesDecompositions as TSD
plotlyjs()

# using DataFramesMeta
# plotlyjs()
const dateform = "yyyymmdd"
include( projectdir("scripts/survey/rawsurvey_utils.jl") )

## ============================================================================
function plot_contacts_per_day(df; region=0)
    gp = groupby(df[:,[:date]], :date, sort=true)
    cb = combine(gp, nrow)
    p = plot(cb.date, cb.nrow, ylab="# contacts", xticks=:native)
    savefig(p, projectdir("scripts/survey/","contacts_per_day_region=$region.html"))
    return p
end

function main()
    plotlyjs()
    for region in 1:5
        @info "load data region $region"
        rawcontacts = readcontacts( "/home/and/data/covidsurvey/rawcontacts.csv", select_region=region )
        plot_contacts_per_day(rawcontacts; region)
        threshold = (family=50, colleagues=100, friends=100, strangers=1000)
        contacts = filteroutliers(rawcontacts, threshold) #NOTE National: p = 99.9

        @info "zero inflated regression region $region"
        surveyproblem = SurveyInferenceProblem(
            ZINegBinomialProblem,
            ctype,
            contacts
        )

        results = optimize(surveyproblem, LBFGS(),
            Optim.Options(show_trace=true, iterations=2000);
        )

        contactrates = DataFrame(results)
        save(projectdir("data","raw_contact_rates_region=$region.csv"), contactrates)

        let
            fname = projectdir("scripts/survey/","survey_inference_results_method=$(results.problemtype)_region=$region.bson")
            save(fname, θ = results)
        end

        plot_confidence_timeseries(results)

        # pyplot()
        # p = plot_survey_histograms(results, day=nothing)
        # savefig(p, projectdir("figures/histogram_survey_data_vs_fake_data_region=$region.png"))
        # plotlyjs()

        @info "decomposition region $region"
        slowtrends = TSD.decompose(contactrates, Optim.Options(show_trace=true, iterations=2000); fdir=projectdir("figures/"))
        save(projectdir("data","smoothed_contact_rates_region=$region.csv"), slowtrends)
    end
end

# main()

## =============================================================================
# apply simple average

function rollingmean(xs, window)
    n  = length(xs)-window+1
    ys = Vector{Float64}(undef, n)
    for i in 1:n
        ys[i] = mean(xs[i:i+window-1])
    end
    return ys
end

function smoothing(df, window)
    @assert isodd(window)
    Δ = window ÷ 2
    smoothed = DataFrame( :date => df.date[1+Δ:end-Δ] )
    cols = filter(x->nonmissingtype( eltype( df[!,x] ) ) <: Number, names(df))
    for col in cols
        smoothed[!,col] = rollingmean(df[:,col], window)
    end
    smoothed
end

ps = (
    threshold = (family=50, colleagues=100, friends=100, strangers=1000),
    window = 7,
    fname = normpath( homedir(), "data/covidsurvey/rawcontacts.csv"),
    ctypes = [:family, :colleagues, :friends, :strangers],
)

function main(ps)
    contacts = @_ ps.fname |>
        readcontacts |>
        filteroutliers(__, ps.threshold) |>
        groupby(__, :date, sort = true) |>
        combine(__, ps.ctypes .=> sum, nrow, renamecols = false) |>
        DataFrames.select(__, All(), AsTable(ps.ctypes) => ByRow(sum) => :total) |>
        transform(__, [ps.ctypes..., :total] .=> ( x -> x ./ __.nrow ), renamecols = false) |>
        disallowmissing |>
        smoothing(__, ps.window)

    save(projectdir("data/contacts/dk","smoothed_contact_rates_window=$(ps.window).csv"), contacts)
    return contacts
end

contacts = main(ps)

## =========================================================================================================
@info "load data"
rawcontacts = readcontacts( normpath( homedir(), "data/covidsurvey/rawcontacts.csv"), select_region=0 )
p = plot_contacts_per_day(rawcontacts)
threshold = 99.9 #NOTE National
contacts = filteroutliers(rawcontacts, threshold)
contacts[!,:total] .= 0
for c in ctype
    contacts[!,:total] += contacts[:,c]
end
contacts = contacts[:, [:total, :strangers, :date]]
## =========================================================================================================
@info "zero inflated regression"
surveyproblem = SurveyInferenceProblem(
    ZINegBinomialProblem,
    datacols(contacts),
    contacts
)

results = optimize(surveyproblem, LBFGS(),
    Optim.Options(show_trace=true, iterations=2000);
)

contactrates = DataFrame(results)
save(projectdir("data","raw_total_contact_rates.csv"), contactrates)

let
    fname = projectdir("scripts/survey/","survey_inference_results_method=$(results.problemtype)_TOTAL.bson")
    save(fname, θ = results)
end

plot_confidence_timeseries(results)

pyplot()
p = plot_survey_histograms(results, day=nothing)
# savefig(p, projectdir("figures/histogram_survey_data_vs_fake_data_region=$region.png"))
plotlyjs()

@info "decomposition region $region"
slowtrends = TSD.decompose(contactrates, Optim.Options(show_trace=true, iterations=2000); fdir=projectdir("figures/"))
save(projectdir("data","smoothed_total_contact_rates.csv"), slowtrends)


# SECTION 3: DECREASE INDI-ACCOUNT'S INCOME
# --------------------------------
# NOTE: we use the expr, macros in SECTION 2

# 3.1 the parameter combinations of this section's reforms
ReformYear = 2020 - env.START_YEAR + 1
idx_plot = (2010:2110) .- env.START_YEAR .+ 1
reform_indiacct_expr = Dict(
   "Benchmark" => :(
      nothing
   ),  # benchmark, no change
   "NoFirmTrans2Indi" => :(
      Pt[:𝕒][ReformYear:end] .= 0.0
   ),  # firm's contrib now only go to pool & retirees's indi acct
   "NoFirmTrans2Retiree" => :(
      # Pt[:𝕒][ReformYear:end] .= 0.0;
      Pt[:𝕓][ReformYear:end] .= 0.0;
   ),  # canceling transfer payment to retirees
   "Everything2Pool" => :(
      # it is equvalent to the parameters let new piM = old piM, where new phi = 0
      tmpNewZeta = (1.0 + Pt[:z][ReformYear-1] * Pt[:η][ReformYear-1]) * (Pt[:ζ][ReformYear-1] + Pt[:ϕ][ReformYear-1]) ;
      tmpNewZeta /= 1.0 + Pt[:z][ReformYear-1] * Pt[:η][ReformYear-1] - Pt[:ϕ][ReformYear-1] ;
      # let new phi = 0, bba = 0, bbb = 0
      Pt[:ϕ][ReformYear:end] .= 0.0;
      Pt[:𝕒][ReformYear:end] .= 0.0;
      Pt[:𝕓][ReformYear:end] .= 0.0;
      # set new zeta to equilize the old piM & new piM
      Pt[:ζ][ReformYear:end] .= tmpNewZeta;
   ),  # hold personal contributions, both personal & firm's contrib all go to pool, i.e. no individual account
   "JustKillIndiAcct" => :(
      Pt[:ϕ][ReformYear:end] .= 0.0;
      Pt[:𝕒][ReformYear:end] .= 0.0;
      Pt[:𝕓][ReformYear:end] .= 0.0;
   ) # just kill individual accounts but change nothing (no contribution, no transfer payments, no benefits)
)

# 3.2 spaces to store results
res_IndiInDesc = Dict(
   "reformNames" => [ "Benchmark", "NoFirmTrans2Indi",       "NoFirmTrans2Retiree",     "Everything2Pool",                      "JustKillIndiAcct" ],
   "plotlegends" => [ "Benchmark", "No transfer to workers", "No transfer to retirees", "Keep contribution but all to pooling", "Canceling individual accounts" ],
   "plotstyles" => [ "-", "-.b", "-r", ":m", "--c" ],
   "gap/exp" => [],
   "gap/in" => [],
   "gap/gdp" => [],
   "gap/taxrev" => [],
   "Dt" => [],
)

# 3.3 evaluate reforms
for tmpreform in res_IndiInDesc["reformNames"]
   # ------ data loading
   include("$(pwd())\\src\\proc_VarsDeclare.jl"); # use pwd() to avoid possible ambiguity of working directory
   include("$(pwd())\\src\\proc_InitPars.jl");
   # ------ reform
   eval( reform_indiacct_expr[tmpreform] )
   # ------ evaluate model & save
   eval(expr_RunModel)
   push!( res_IndiInDesc["gap/exp"], 100 .* (Dt[:LI] ./ Dt[:AggPoolExp])[idx_plot] )
   push!( res_IndiInDesc["gap/in"] , 100 .* (Dt[:LI] ./ Dt[:AggPoolIn])[idx_plot] )
   push!( res_IndiInDesc["gap/gdp"], 100 .* (Dt[:LI] ./ Dt[:Y])[idx_plot] )
   push!( res_IndiInDesc["gap/taxrev"], 100 .* (Dt[:LI] ./ (Dt[:TRc] .+ Dt[:TRw]))[idx_plot] )
   push!( res_IndiInDesc["Dt"], copy(Dt) )
end # end loops

#


__precompile__()
module SeisXcorrelation

using SeisIO, Noise, Printf, Dates, FFTW, JLD2, MPI

export seisxcorrelation

"""

    seisxcorrelation(maxtimelag::Real, finame::String, foname::String; IsAllComponent::Bool=false)

Compute cross-correlation function and save data in jld2 file with SeisData format.

# Arguments
- `finame::String,`    : Input file name e.g. network = "BPnetwork"
- `maxtimelag::Real`    : Maximum lag time e.g. maxtimelag = 100 [s]
- `IsAllComponent::Bool`   : If true, compute 3 X 3 components cross-correlation

# Output
- `foname.jld2`                 : contains SeisData structure with a hierarchical structure (CC function, metadata)

"""
function seisxcorrelation(maxtimelag::Real, finame::String, foname::String; IsAllComponents::Bool=false) #Add whatever necessary arguments

    MPI.Init()
    # establish the MPI communicator and obtain rank
    comm = MPI.COMM_WORLD
    size = MPI.Comm_size(comm)
    rank = MPI.Comm_rank(comm)

    #Read validation data from ["../EXAMPLE/codevalidation"]

    #Make Station pairs (after validating CC part, read from actual dataset and use statino list in SeisData structure)
    #Take into account IsAllComponents == true or not

    mpiitrcount = 0
    baton = Array{Int32, 1}([0]) # for Relay Dumping algorithm

    for i = 1:size(data["info/timestamplist"])

        #parallelize one processor for one time stamp
        processID = stid - (size * mpiitrcount)

        # if this mpiitrcount is final round or not
        length(stlist) - size >= size * mpiitrcount ? anchor_rank = size-1 : anchor_rank = mod(length(stlist), size)-1

        if rank == processID-1

            #roop each station pair

            for j = 1:length("""stationpairlist`""")

                # compute crosscorrelation using Noise.jl

                # store in SeisChannel structure
                # time, fs, misc["stationpair"], misc["cc1"]

                # save dataset

                # Save data one by one
                if size == 1
                    save_SeisData2JLD2(foname, varname, S)

                else
                    if rank == 0
                        save_SeisData2JLD2(foname, varname, S)

                        if anchor_rank != 0
                            MPI.Send(baton, rank+1, 11, comm)
                            MPI.Recv!(baton, anchor_rank, 12, comm)
                        end

                    elseif rank == anchor_rank
                        MPI.Recv!(baton, rank-1, 11, comm)
                        save_SeisData2JLD2(foname, varname, S)
                        MPI.Send(baton, 0, 12, comm)

                    else
                        MPI.Recv!(baton, rank-1, 11, comm)
                        save_SeisData2JLD2(foname, varname, S)
                        MPI.Send(baton, rank+1, 11, comm)
                    end
                end
            end

            mpiitrcount += 1

        end
    end

    if rank == 0 println("Crosscorrelation and Saving data is successfully done.\njob ended at "*string(now())) end

    MPI.Finalize()

    return 0
end

end


#===============================================================================

    CalibrationData.jl

    container for calibration data

    data has 3 dimensions:
    1. Well     :0 - :15
    2. Channel  1, 2
    3. Dye      water, dye1, dye2

    Author: Tom Price
    Date:   June 2019

===============================================================================#

import DataStructures.OrderedDict
import StaticArrays: SArray

struct CalibrationData{N <: NumberOfChannels, R <: Union{Int_T,Float_T}}
    array       ::SArray{S,R,3} where {S <: Tuple}
    ## inner constructors supply type parameters
    ## and enforce constraints on size of data array
    CalibrationData(
        ::Type{SingleChannel},
        data ::AbstractArray{R}) where{R <: Union{Int_T,Float_T}} =
            new{SingleChannel,R}(SArray{Tuple{size(data,1),1,2},R}(data))
    CalibrationData(
        ::Type{DualChannel},
        data ::AbstractArray{R}) where{R <: Union{Int_T,Float_T}} =
            new{DualChannel,R}(SArray{Tuple{size(data,1),2,3},R}(data))
end

## parser / outer constructor
function CalibrationData(calib ::Associative)
    local water, dye1, dye2 ## not enum
    num_wells = count_wells(calib)
    water     = calib[WATER_KEY][FLUORESCENCE_VALUE_KEY]
    R         = water |> first |> first |> typeof
	if (R == Int64 && Int_T != Int64)
		R = Int_T
	end

    dye1      = calib[CHANNEL_KEY * "_1"][FLUORESCENCE_VALUE_KEY]
    if length(dye1) > 1 && thing(dye1[2]) && haskey(calib, CHANNEL_KEY * "_2")
        dye2  = calib[CHANNEL_KEY * "_2"][FLUORESCENCE_VALUE_KEY]
        data  = try
            [water, dye1, dye2] |> gather(vcat) |> gather(hcat) |>
                morph(num_wells,2,3) |> mold(bless(R))
        catch
            warn(logger, "dual channel calibration data are in the wrong format")
            rethrow()
        end ## try
        return CalibrationData(DualChannel, data)
    end ## if
    ## else single channel
    data  = try
        [water, dye1] |> moose(first, hcat) |> morph(num_wells,1,2) |> mold(bless(R))
    catch
        warn(logger, "single channel calibration data are in the wrong format")
        rethrow()
    end ## try
    return CalibrationData(SingleChannel, data)
end ## CalibrationData()


using JuMP, GLPK
using DataFrames, CSV
using LinearAlgebra
using JSON
using LightGraphs

#cd("/home/xiaotong/Dropbox (MIT)/13_Amazon_LM_competition/submission_new")
cd(@__DIR__)

tsp_opt_seq = JSON.parsefile("../data/model_apply_outputs/model_apply_output/mean_dist/opt_complete_seq_tour_post_process.json")

# Output and store sequences
model_apply_output = Dict()
for i in keys(tsp_opt_seq)
    model_apply_output[i] = Dict()
    model_apply_output[i]["proposed"] = Dict()
    for j in 1:length(tsp_opt_seq[i])
        model_apply_output[i]["proposed"][tsp_opt_seq[i][j]] = j-1
    end
end

open("../data/model_apply_outputs/proposed_sequences.json","w") do f
    JSON.print(f, model_apply_output)
end

println("Successfully store proposed sequences")


"""
    linearized_air_pressure(ρ, ρe_tot, ρe_pot, ρq_tot=0, ρq_liq=0, ρq_ice=0)

The air pressure, linearized around a dry rest state, from the equation of state
(ideal gas law) where:

 - `ρ` (moist-)air density
 - `ρe_tot` total energy density
 - `ρe_pot` potential energy density
and, optionally,
 - `ρq_tot` total water density
 - `ρq_liq` liquid water density
 - `ρq_ice` ice density
"""
function linearized_air_pressure(
    ρ::FT,
    ρe_tot::FT,
    ρe_pot::FT,
    ρq_tot::FT = FT(0),
    ρq_liq::FT = FT(0),
    ρq_ice::FT = FT(0),
) where {FT <: Real}
    ρ * FT(R_d) * FT(T_0) +
    FT(R_d) / FT(cv_d) * (
        ρe_tot - ρe_pot - (ρq_tot - ρq_liq) * FT(e_int_v0) +
        ρq_ice * (FT(e_int_i0) + FT(e_int_v0))
    )
end

@inline function linearized_pressure(
    ::DryModel,
    orientation::Orientation,
    state::Vars,
    aux::Vars,
)
    ρe_pot = state.ρ * gravitational_potential(orientation, aux)
    linearized_air_pressure(state.ρ, state.ρe, ρe_pot)
end
@inline function linearized_pressure(
    ::EquilMoist,
    orientation::Orientation,
    state::Vars,
    aux::Vars,
)
    ρe_pot = state.ρ * gravitational_potential(orientation, aux)
    linearized_air_pressure(state.ρ, state.ρe, ρe_pot, state.moisture.ρq_tot)
end

abstract type AtmosLinearModel <: BalanceLaw end

vars_state(lm::AtmosLinearModel, FT) = vars_state(lm.atmos, FT)
vars_gradient(lm::AtmosLinearModel, FT) = @vars()
vars_diffusive(lm::AtmosLinearModel, FT) = @vars()
vars_aux(lm::AtmosLinearModel, FT) = vars_aux(lm.atmos, FT)
vars_integrals(lm::AtmosLinearModel, FT) = @vars()
vars_reverse_integrals(lm::AtmosLinearModel, FT) = @vars()

function update_aux!(
    dg::DGModel,
    lm::AtmosLinearModel,
    Q::MPIStateArray,
    t::Real,
)
    return false
end
function flux_diffusive!(
    lm::AtmosLinearModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    nothing
end
integral_load_aux!(lm::AtmosLinearModel, integ::Vars, state::Vars, aux::Vars) =
    nothing
integral_set_aux!(lm::AtmosLinearModel, aux::Vars, integ::Vars) = nothing
reverse_integral_load_aux!(
    lm::AtmosLinearModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
) = nothing
reverse_integral_set_aux!(lm::AtmosLinearModel, aux::Vars, integ::Vars) =
    nothing
flux_diffusive!(
    lm::AtmosLinearModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
) = nothing
function wavespeed(lm::AtmosLinearModel, nM, state::Vars, aux::Vars, t::Real)
    ref = aux.ref_state
    return soundspeed_air(ref.T, lm.atmos.param_set)
end

function boundary_state!(
    nf::NumericalFluxNonDiffusive,
    atmoslm::AtmosLinearModel,
    args...,
)
    atmos_boundary_state!(nf, AtmosBC(), atmoslm, args...)
end
function boundary_state!(
    nf::NumericalFluxDiffusive,
    atmoslm::AtmosLinearModel,
    args...,
)
    nothing
end
init_aux!(lm::AtmosLinearModel, aux::Vars, geom::LocalGeometry) = nothing
init_state!(lm::AtmosLinearModel, state::Vars, aux::Vars, coords, t) = nothing

struct AtmosAcousticLinearModel{M} <: AtmosLinearModel
    atmos::M
    function AtmosAcousticLinearModel(atmos::M) where {M}
        if atmos.ref_state === NoReferenceState()
            error("AtmosAcousticLinearModel needs a model with a reference state")
        end
        new{M}(atmos)
    end
end

function flux_nondiffusive!(
    lm::AtmosAcousticLinearModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    FT = eltype(state)
    ref = aux.ref_state
    e_pot = gravitational_potential(lm.atmos.orientation, aux)

    flux.ρ = state.ρu
    pL =
        linearized_pressure(lm.atmos.moisture, lm.atmos.orientation, state, aux)
    flux.ρu += pL * I
    flux.ρe = ((ref.ρe + ref.p) / ref.ρ - e_pot) * state.ρu
    nothing
end
function source!(
    lm::AtmosAcousticLinearModel,
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)
    nothing
end

struct AtmosAcousticGravityLinearModel{M} <: AtmosLinearModel
    atmos::M
    function AtmosAcousticGravityLinearModel(atmos::M) where {M}
        if atmos.ref_state === NoReferenceState()
            error("AtmosAcousticGravityLinearModel needs a model with a reference state")
        end
        new{M}(atmos)
    end
end
function flux_nondiffusive!(
    lm::AtmosAcousticGravityLinearModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    FT = eltype(state)
    ref = aux.ref_state
    e_pot = gravitational_potential(lm.atmos.orientation, aux)

    flux.ρ = state.ρu
    pL =
        linearized_pressure(lm.atmos.moisture, lm.atmos.orientation, state, aux)
    flux.ρu += pL * I
    flux.ρe = ((ref.ρe + ref.p) / ref.ρ) * state.ρu
    nothing
end
function source!(
    lm::AtmosAcousticGravityLinearModel,
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)
    ∇Φ = ∇gravitational_potential(lm.atmos.orientation, aux)
    source.ρu -= state.ρ * ∇Φ
    nothing
end


module SpSkipGramB
# share sparse representations

using CUDArt: HostArray, CudaArray, synchronize, null_stream
using JLD: jldopen
using CycUtils: invertable_unique, get, set!
using Veclib: gamma_kernel!, gamma_kernel_loss!, normalize_columns!

export SparseSkipGramB, train, init, close_device

include("config.jl")
include("array_handler.jl")
include("samples.jl")
include("utils.jl")

type SparseSkipGramB{Tensor} # Tensor could be HostArray{Float32, 2}
    α::Tensor
    β::Tensor
    B::Tensor
    C::Tensor

    # the following fields are training parameters.
    m::Int
    k::Int
    lr::Float32
    lr0::Float32
    λ::Float32

    # maps
    idx2word::Array{String, 1}
    word2idx::Dict{String, UInt32}

    # the following fields are named memory for training
    ∇αβ::Tensor
    ᾱ::Tensor
    β̄::Tensor
    a⃗::Tensor
    b⃗::Tensor
    γ::Tensor
    σ::Tensor
    ∇α::Tensor
    N::Tensor
    M::Tensor
    ∇β::Tensor
    Q::Tensor
    ∇B::Tensor
    ∇C::Tensor
    idxes::Tensor
    sign::Tensor

    function SparseSkipGramB(d::Integer, # vector dimension
        num_base::Integer,     # number of atom words
        num_vocab::Integer,    # number of vocabulary
        mini_batch::Integer,     # mini-batch size
        negative::Integer,        # number of negative samples per positive sample
        lr::AbstractFloat=0.01,       # learning rate
        λ::AbstractFloat=0.1,         # sparseness weight
        init_model::String=""   # dumped init_model.
    )
        @assert num_base ≤ num_vocab "The number of atom words must be less than the vocabulary size"

        α_ = randn(Float32, num_base, num_vocab) ./ sqrt(num_base)
        B_ = randn(Float32, d, num_base) ./ sqrt(d)
        C_ = randn(Float32, d, num_base) ./ sqrt(d)

        i2w = Array{String, 1}()
        w2i = Dict{String, UInt32}()

        if init_model != ""
            e, c, i2w = jldopen(init_model, "r") do f
                (read(f, "embeddings"), read(f, "contexts"),
                 read(f, "idx2word"))
            end
            @assert num_vocab ≤ length(i2w) "Initial model must contains as many vocabularies as the model to be trained."
            @assert d == size(e, 1) "Initial vector dim. and specified dim. must agree."

            i2w = i2w[1:num_vocab]
            for (i, w) in enumerate(i2w)
                w2i[w] = i
            end

            α_[1:num_base, 1:num_base] = eye(num_base)
            copy!(B_, @view e[:, 1:num_base])
            copy!(C_, @view c[:, 1:num_base])
        end

        α = Tensor(α_)
        β = α
        B = Tensor(B_)
        C = Tensor(C_)

        ∇αβ = Tensor(Float32, num_base*(negative+2)*mini_batch)
        ∇α, ∇β = splitarray(∇αβ, (num_base, mini_batch),
                                   (num_base, negative+1, mini_batch))

        ᾱ = Tensor(Float32, num_base, mini_batch)
        β̄ = Tensor(Float32, num_base, negative+1, mini_batch)
        a⃗ = Tensor(Float32, d, mini_batch)
        b⃗ = Tensor(Float32, d, negative+1, mini_batch)
        γ = Tensor(Float32, negative+1, mini_batch)
        σ = Tensor(Float32, negative+1, mini_batch)
        # ∇α = Tensor(Float32, num_base, mini_batch)
        N = Tensor(Float32, d, mini_batch)
        M = Tensor(Float32, num_base, mini_batch)
        # ∇β = Tensor(Float32, num_base, negative+1, mini_batch)
        Q = Tensor(Float32, num_base, mini_batch)
        ∇B = Tensor(Float32, d, num_base)
        ∇C = Tensor(Float32, d, num_base)
        idxes = Tensor(UInt32, (negative+2)*mini_batch)

        fill!(∇αβ, 0.0)
        fill!(ᾱ, 0.0)
        fill!(β̄, 0.0)
        fill!(a⃗, 0.0)
        fill!(b⃗, 0.0)
        fill!(γ, 0.0)
        fill!(σ, 0.0)
        # fill!(∇α, 0.0)
        fill!(N, 0.0)
        fill!(M, 0.0)
        # fill!(∇β, 0.0)
        fill!(Q, 0.0)
        fill!(∇B, 0.0)
        fill!(∇C, 0.0)
        fill!(idxes, 0)

        tmp = -ones(Float32, (negative+1, mini_batch))
        tmp[1, :] = 1.0f0
        s = Tensor(tmp)

        new{Tensor}(α, β, B, C,                               # model parameters
            mini_batch, negative, lr, lr, λ,              # training parameters
            i2w, w2i,                                 # maps
            ∇αβ, ᾱ, β̄, a⃗, b⃗, γ, σ, ∇α, N, M, ∇β, Q, ∇B, ∇C, idxes, s
        )
    end
end

function save(model::SparseSkipGramB{Array}, fname::String; save_maps=false)
    jldopen(fname, "w") do f
        write(f, "α", model.α)
        write(f, "B", model.B)
        write(f, "C", model.C)

        # the following fields are training parameters.
        write(f, "m", model.m)
        write(f, "k", model.k)
        write(f, "lr", model.lr0)
        write(f, "λ", model.λ)

        # maps
        if save_maps == true
            write(f, "idx2word", model.idx2word)
            write(f, "word2idx", model.word2idx)
        end
    end
end

num_vocab(model::SparseSkipGramB) = size(model.α, 2)
num_base(model::SparseSkipGramB) = size(model.α, 1)
num_vecdim(model::SparseSkipGramB) = size(model.B, 1)
num_minibatch(model::SparseSkipGramB) = size(model.a⃗, 2)
num_negative(model::SparseSkipGramB) = size(model.b⃗, 2) - 1

function to_device{M <: Union{SparseSkipGramB{Array},
                    SparseSkipGramB{HostArray}}}(model::M)
    d = num_vecdim(model)
    num_base_ = num_base(model)
    num_vocab_ = num_vocab(model)
    mini_batch = num_minibatch(model)
    negative = num_negative(model)

    d_model = SparseSkipGramB{CudaArray}(d, num_base_, num_vocab_, mini_batch,
                                        negative, model.lr, model.λ)
    for fn in fieldnames(typeof(d_model))
        if fn == :β
            continue
        end
        dest = getfield(d_model, fn)
        if isa(dest, CudaArray)
             src = getfield(model, fn)
             copy!(dest, src)
        end
    end
    return d_model
end

@generated function get_training_model(model::SparseSkipGramB)
    if BACKEND == "GPU"
        return :(to_device(model))
    else
        return :(model)
    end
end

function compute_grad_sparse(model, batch)
    const d = num_vecdim(model)
    const m = model.m
    const k = model.k
    const B = model.B
    const C = model.C
    const N = model.N
    const M = model.M
    w, c = batch
    const nb = num_base(model)
    const ᾱ = takelast!(model.ᾱ, model.α, w)
    const β̄ = takelast!(model.β̄, model.β, c)
    const a⃗ = model.a⃗
    const b⃗ = model.b⃗
    const γ = model.γ
    const ∇α = model.∇α
    const ∇β = model.∇β

    flat_b⃗ = reshape(b⃗, (d, (k+1)*m))
    gemm!('N', 'N', 1.0f0, B, ᾱ, 0.0f0, a⃗)
    gemm!('N', 'N', 1.0f0, C, reshape(β̄, (nb, (k+1)*m)), 0.0f0, flat_b⃗)

    a⃗_ = reshape(a⃗, (d, 1, m))
    γ_ = reshape(γ, (k+1, 1, m))
    gemm_strided_batched!('T', 'N', 1.0f0, b⃗, a⃗_, 0.0f0, γ_)
    gamma_kernel_loss!(γ, model.σ, model.sign)

    #∇α
    N_ = reshape(N, (d, 1, m))
    gemm_strided_batched!('N', 'N', 1.0f0, b⃗, γ_, 0.0f0, N_)
    gemm!('T', 'N', 1.0f0, B, N, 0.0f0, ∇α)

    # ∇β
    gemm!('T', 'N', 1.0f0, C, a⃗, 0.0f0, M)
    M_ = reshape(M, (num_base(model), 1, m))
    gemm_strided_batched!('N', 'T', 1.0f0, M_, γ_, 0.0f0, ∇β)
    nothing
end

function compute_grad_dict(model)
    const m = model.m
    const k = model.k

    gemm!('N', 'T', 1.0f0, model.N, model.ᾱ, 0.0f0, model.∇B)
    β̄ = model.β̄                                   # |B|×(k+1)×m
    γ = reshape(model.γ, (k+1, 1, m))             # (k+1) × 1 × m
    Q = reshape(model.Q, (num_base(model), 1, m)) # |B| × 1 × m
    gemm_strided_batched!('N', 'N', 1.0f0, β̄, γ, 0.0f0, Q)
    gemm!('N', 'T', 1.0f0, model.a⃗, model.Q, 0.0f0, model.∇C)
    nothing
end

function update_dict_u1f1!(model)
    const m = model.m
    const k = model.k
    const lr_ = model.lr

    gemm!('N', 'T', -lr_, model.N, model.ᾱ, 1.0f0, model.B)
    β̄ = model.β̄  # |B|×(k+1)×m
    γ = reshape(model.γ, (k+1, 1, m))  # (k+1) × 1 × m
    Q = reshape(model.Q, (num_base(model), 1, m)) # |B| × 1 × m
    gemm_strided_batched!('N', 'N', 1.0f0, β̄, γ, 0.0f0, Q)
    gemm!('N', 'T', -lr_, model.a⃗, model.Q, 1.0f0, model.C)
    nothing
end

function update_dict_u2f1!(model)
    compute_grad_dict(model)
    update_dict_u2!(model)
end

function accumulate_grad_dict(model, t)
    const m = model.m
    const k = model.k
    w1 = (Float32(t)/(t+1)) :: Float32
    w2 = (1.0f0/(t+1)) :: Float32

    gemm!('N', 'T', w2, model.N, model.ᾱ, w1, model.∇B)
    β̄ = model.β̄                                   # |B|×(k+1)×m
    γ = reshape(model.γ, (k+1, 1, m))             # (k+1) × 1 × m
    Q = reshape(model.Q, (num_base(model), 1, m)) # |B| × 1 × m
    gemm_strided_batched!('N', 'N', 1.0f0, β̄, γ, 0.0f0, Q)
    gemm!('N', 'T', w2, model.a⃗, model.Q, w1, model.∇C)
    nothing
end

function update_dict_u1!(model)
    const lr_::Float32 = model.lr
    axpy!(-lr_, vec(model.∇B), vec(model.B))
    axpy!(-lr_, vec(model.∇C), vec(model.C))
    nothing
end

const ∇B_cache = Ref{HostArray}()
const ∇C_cache = Ref{HostArray}()
function update_dict_u2!(model)
    global ∇B_cache
    global ∇C_cache
    ∇B = ∇B_cache[]
    ∇C = ∇C_cache[]
    lr_::Float32 = model.lr

    copy!(∇B, model.∇B)
    CUDArt.synchronize(CUDArt.null_stream)
    normalize_columns!(∇B)
    copy!(model.∇B, ∇B)
    copy!(∇C, model.∇C)
    CUDArt.synchronize(CUDArt.null_stream)
    axpy!(-lr_, vec(model.∇B), vec(model.B))
    normalize_columns!(∇C)
    copy!(model.∇C, ∇C)
    axpy!(-lr_, vec(model.∇C), vec(model.C))
    nothing
end

function fetch_grad_sparse!(grads, model)
    copy!(grads, model.∇αβ)
    return grads
end

# function cpu_sparse_update!(arr, grad, idxes, lr, lambda)
#     for (j, idx) in zip(1:size(grad,2), idxes)
#         @inbounds for i in 1:size(arr, 1)
#             x = arr[i, idx] -  lr * grad[i, j]
#             s = (x > 0.0f0 ? 1.0f0 : -1.0f0)
#             x *= s
#             x -= lambda*lr
#             x = x > 0.0f0 ? x : 0.0f0
#             x = x < 1.0f0 ? x : 1.0f0
#             arr[i, idx] = x * s
#         end
#     end
# end

function gpu_sparse_update!(arr, grad, idxes, lr, lambda)
    # make sure the transferring completes.
    CUDArt.synchronize(CUDArt.null_stream)
    ccall((:sparse_update, vec_lib), Void, (Ptr{Float32}, Ptr{Float32},
        Cint, Cint, Ptr{Cuint}, Float32, Float32, Ptr{Void}, Cint),
        arr.ptr, grad.ptr, size(grad, 1), size(grad, 2),
        idxes.ptr, lr, lambda, null_stream.inner.handle, 256)
end

@generated function proximal_update!(model, ∇αβ, idxes)
    if BACKEND == "GPU"
        quote
            lr_::Float32 = model.lr
            dev = Int(DEVICE)
            siz_αβ = (num_base(model), length(idxes))
            d_∇αβ = CudaArray(model.∇αβ.ptr, siz_αβ, dev)
            copy!(d_∇αβ, ∇αβ)

            ptr_idxes = model.idxes.ptr
            d_idxes = CudaArray(ptr_idxes, (length(idxes),), dev)
            copy!(d_idxes, idxes)
            gpu_sparse_update!(model.α, d_∇αβ, d_idxes, lr_, model.λ)
        end
    else
        quote
            error("Not implemented")
            # lr_::Float32 = model.lr
            # cpu_sparse_update!(model.α, ∇αβ, idxes, lr_, model.λ)
        end
    end
end

function update_sparse_u1!(model, ∇αβ, grad_αβ, uniq_batch)
    idx, inv_idxes = uniq_batch

    ∇αβ = @view ∇αβ[:, 1:length(idx)]
    fill!(∇αβ, 0.0f0)

    # make sure the transferring completes.
    CUDArt.synchronize(CUDArt.null_stream)

    for (i, inv_idx) in enumerate(inv_idxes)
        ∇αβ[:, inv_idx] += grad_αβ[:, i]
    end

    proximal_update!(model, ∇αβ, idx)
    nothing
end

function update_sparse_u2!(model, ∇αβ, grad_αβ, uniq_batch)
    idx, inv_idxes = uniq_batch

    ∇αβ = @view ∇αβ[:, 1:length(idx)]
    fill!(∇αβ, 0.0f0)

    # make sure the transferring completes.
    CUDArt.synchronize(CUDArt.null_stream)

    for (i, inv_idx) in enumerate(inv_idxes)
        ∇αβ[:, inv_idx] += grad_αβ[:, i]
    end
    normalize_columns!(∇αβ)

    proximal_update!(model, ∇αβ, idx)
    nothing
end

# function update_all!(model, grad_cache, grad, uniq_batch)
#     lr_::Float32 = model.lr
#     update_αβ!(model, grad_cache, grad, uniq_batch)
#     axpy!(-lr_, vec(model.∇B), vec(model.B))
#     axpy!(-lr_, vec(model.∇C), vec(model.C))
#     nothing
# end

@generated function sync_params!(cpu_model, model)
    if BACKEND == "CPU"
        return :(cpu_model)
    else
        quote
            for fn in (:α, :B, :C)  #fieldnames(typeof(model))
                src = getfield(model, fn)
                dest = getfield(cpu_model, fn)
                copy!(dest, src)
            end
            CUDArt.synchronize(CUDArt.null_stream)
            cpu_model
        end
    end
end

const _start_time_loginfo = Ref{Float64}(0.0)
const _next_report_time = Ref{Float64}(0.0)
const _sigmas_loginfo = Ref{HostArray{Float32, 2}}()

function log_progress(model, current, progress=progress)
    global _start_time_loginfo
    global _next_report_time
    global _sigmas_loginfo
    global sampleFactory
    factory = get(sampleFactory)

    if current >= _next_report_time[]
        σ = _sigmas_loginfo[]
        copy!(σ, model.σ)
        CUDArt.synchronize(CUDArt.null_stream)
        loss = 0.0f0
        for x in σ
            loss += -log(x+eps(typeof(x)))
        end
        ppl = exp(loss/(model.m*(model.k+1)))

        _next_report_time[] += 3.0
        st = _start_time_loginfo[]
        factory = get(sampleFactory)
        prog = progress(factory) * 100.0
        elapsed = current - st
        etc = (100.0 - prog) / prog * elapsed
        hours = div(etc, 3600)
        rem = etc-hours*3600
        mins = div(rem, 60)
        secs = rem - mins * 60

        mean_speed = factory.consumed[]/elapsed
        info(@sprintf("ETA: %02d:%02d:%02d - loss: %.4f - ppl: %.3f - progress: %5.2f%% - speed: %.1f %s/s - lr: %g", hours, mins, secs, loss, ppl, prog, mean_speed, factory.corpus_type == String ? "lines" : "bytes", model.lr))
    end
    nothing
end

function init_logger(model, start_time)
    global _start_time_loginfo
    global _next_report_time
    global _sigmas_loginfo
    _start_time_loginfo[] = start_time
    _next_report_time[] = start_time
    _sigmas_loginfo[] = HostArray(Float32, (model.k+1), model.m)
end

function train(model; min_lr=1f-4, normalize_sp=false,
        normalize_dict=false, every=500, preratio=0.1,
        saveratio=0.05, part=1.0, save_basename="B")
    if every == 1
        train_f1(model, min_lr=min_lr, normalize_sp=normalize_sp,
        normalize_dict=normalize_dict, preratio=preratio,
        saveratio=saveratio, part=part, save_basename=save_basename)
        return nothing
    end
    global sampleFactory

    if normalize_dict
        global ∇B_cache
        global ∇C_cache
        ∇B_cache[] = HostArray(Float32, (num_vecdim(model), num_base(model)))
        ∇C_cache[] = HostArray(Float32, (num_vecdim(model), num_base(model)))
    end

    cpu_model = model
    model = to_device(cpu_model)

    beta = model.lr0 / min_lr - 1.0f0

    info(@sprintf("Train with lr=%g, normalize_sp=%s, normalize_dict=%s. Update dictionary parameters every %d updates of sparse parameters. Pretrain sparse parameters with %.2f%% amount of training data. Save after every %.2f%% of training. Use %.2f%% of the total training data. Save the trained model to \"%s*.jld\"", model.lr, string(normalize_sp), string(normalize_dict), every, preratio*100, saveratio*100, part*100, save_basename))

    st = time()
    init_logger(model, st)

    grads = HostArray(Float32, (num_base(model), (model.k+2)*model.m))
    grad_cache = HostArray(Float32, (num_base(model), (model.k+2)*model.m))

    factory = get(sampleFactory)
    const producer = factory.producer
    batch_set = Task(() -> producer(factory, model.m, model.k))

    const update_sparse! = normalize_sp ? update_sparse_u2! : update_sparse_u1!
    const update_dict! = normalize_dict ? update_dict_u2! : update_dict_u1!
    const progress_ = factory -> progress(factory)/part

    batch = consume(batch_set)
    old_batch = batch
    compute_grad_sparse(model, old_batch) # this method immediately returns when
                                  # the computations are taking place on a GPU.
                                  # this makes the gpu and cpu runs concurrently
    for batch in batch_set
        uniq_batch = invertable_unique(old_batch[1], old_batch[2])
        fetch_grad_sparse!(grads, model)
        update_sparse!(model, grad_cache, grads, uniq_batch)

        old_batch = batch
        prog = progress_(factory)
        if prog >= preratio
            if preratio > 0.0f0
                info("Pre-train finished.")
            end
            break
        end
        compute_grad_sparse(model, old_batch)
        log_progress(model, time(), progress_)
        model.lr = model.lr0 / (1.0f0 + beta*prog)
    end

    next_save_point = preratio + saveratio
    compute_grad_sparse(model, old_batch)
    nb_batch = 0
    for batch in batch_set
        uniq_batch = invertable_unique(old_batch[1], old_batch[2])
        # synchronize(CUDArt.null_stream)
        # TODO: maybe another stream to handle the data transferring
        fetch_grad_sparse!(grads, model)

        accumulate_grad_dict(model, nb_batch)
        if nb_batch == every - 1
            update_dict!(model)
        end
        update_sparse!(model, grad_cache, grads, uniq_batch)

        old_batch = batch
        compute_grad_sparse(model, old_batch)
        prog = progress_(factory)
        if prog >= next_save_point
            if prog >= 1.0f0
                break
            end
            # synchronize
            sync_params!(cpu_model, model)
            sn = joinpath(project_root,
                @sprintf("%s-checkpoint-%03.0f.jld", save_basename, prog*100.0f0))
            info("Save point reached. saveing to $sn")
            save(cpu_model, sn)
            next_save_point += saveratio
        end
        log_progress(model, time(), progress_)
        model.lr = model.lr0 / (1.0f0 + beta*prog)
        nb_batch = (nb_batch+1) % every
    end

    sync_params!(cpu_model, model)
    sn = joinpath(project_root,
        @sprintf("%s-checkpoint-100.jld", save_basename))
    info("Training finished. Saving to $sn")
    save(cpu_model, sn, save_maps=true)
    nothing
end

function train_f1(model; min_lr=1f-4, normalize_sp=false,
        normalize_dict=false, preratio=0.1,
        saveratio=0.05, part=1.0, save_basename="B")
    global sampleFactory

    if normalize_dict
        global ∇B_cache
        global ∇C_cache
        ∇B_cache[] = HostArray(Float32, (num_vecdim(model), num_base(model)))
        ∇C_cache[] = HostArray(Float32, (num_vecdim(model), num_base(model)))
    end

    cpu_model = model
    model = to_device(cpu_model)

    beta = model.lr0 / min_lr - 1.0f0

    info(@sprintf("Train with lr=%g, normalize_sp=%s, normalize_dict=%s. Update dictionary parameters every 1 updates of sparse parameters. Pretrain sparse parameters with %.2f%% amount of training data. Save after every %.2f%% of training. Use %.2f%% of the total training data. Save the trained model to \"%s*.jld\"", model.lr, string(normalize_sp), string(normalize_dict), preratio*100, saveratio*100, part*100, save_basename))

    st = time()
    init_logger(model, st)

    grads = HostArray(Float32, (num_base(model), (model.k+2)*model.m))
    grad_cache = HostArray(Float32, (num_base(model), (model.k+2)*model.m))

    factory = get(sampleFactory)
    const producer = factory.producer
    batch_set = Task(() -> producer(factory, model.m, model.k))

    const update_sparse! = normalize_sp ? update_sparse_u2! : update_sparse_u1!
    const update_dict! = normalize_dict ? update_dict_u2f1! : update_dict_u1f1!
    const progress_ = factory -> progress(factory)/part

    batch = consume(batch_set)
    old_batch = batch
    compute_grad_sparse(model, old_batch) # this method immediately returns when
                                  # the computations are taking place on a GPU.
                                  # this makes the gpu and cpu runs concurrently
    for batch in batch_set
        uniq_batch = invertable_unique(old_batch[1], old_batch[2])
        fetch_grad_sparse!(grads, model)
        update_sparse!(model, grad_cache, grads, uniq_batch)

        old_batch = batch
        prog = progress_(factory)
        if prog >= preratio
            if preratio > 0.0f0
                info("Pre-train finished.")
            end
            break
        end
        compute_grad_sparse(model, old_batch)
        log_progress(model, time(), progress_)
        model.lr = model.lr0 / (1.0f0 + beta*prog)
    end

    next_save_point = preratio + saveratio
    compute_grad_sparse(model, old_batch)
    for batch in batch_set
        uniq_batch = invertable_unique(old_batch[1], old_batch[2])
        # synchronize(CUDArt.null_stream)
        # TODO: maybe another stream to handle the data transferring
        fetch_grad_sparse!(grads, model)
        update_dict!(model)
        update_sparse!(model, grad_cache, grads, uniq_batch)

        old_batch = batch
        compute_grad_sparse(model, old_batch)
        prog = progress_(factory)
        if prog >= next_save_point
            if prog >= 1.0f0
                break
            end
            # synchronize
            sync_params!(cpu_model, model)
            sn = joinpath(project_root,
                @sprintf("%s-checkpoint-%03.0f.jld", save_basename, prog*100.0f0))
            info("Save point reached. saveing to $sn")
            save(cpu_model, sn)
            next_save_point += saveratio
        end
        log_progress(model, time(), progress_)
        model.lr = model.lr0 / (1.0f0 + beta*prog)
    end

    sync_params!(cpu_model, model)
    sn = joinpath(project_root,
        @sprintf("%s-checkpoint-100.jld", save_basename))
    info("Training finished. Saving to $sn")
    save(cpu_model, sn, save_maps=true)
    nothing
end

end # end of module


function maxent_max_pol(method, model_key)

    ## -------------------------------------------------------------------
    # Monitor
    mon = SimT.OnDiskMonitor(cachedir(iJR, :ME_MONITOR))
    SimT.reset!(mon)

    # Feed jobs
    Ch = Channel(nthreads()) do ch
        cGLCs = Nd.val("cGLC")
        for (exp, cGLC)  in enumerate(cGLCs)
            put!(ch, (exp, cGLC))
        end
    end

    @threads for _ in 1:nthreads()
        thid = threadid()
        for (exp, cGLC) in Ch

            ## -------------------------------------------------------------------
            # handle cache
            is_cached(;method, exp) && continue

            ## -------------------------------------------------------------------
            # Excluded because of unfeasibility or convergence problems
            exp in IGNORE_EXPS && continue

            ## -------------------------------------------------------------------
            # SetUp
            model = iJR.load_model(model_key)
            M, N = size(model)
            biomidx = ChU.rxnindex(model, iJR.BIOMASS_IDER)
            glcidx = ChU.rxnindex(model, iJR.EX_GLC_IDER)
            exp_growth = Nd.val("D", exp) # experimental biom growth rate (equals D)
            # glc per biomass unit supply
            cgD_X = -Nd.cval(:GLC, exp) * Nd.val(:D, exp) / Nd.val(:X, exp)
            biom_beta = 0.0 # current biomass beta
            biom_betas = [biom_beta] # biomass beta round history
            vg_beta = 0.0 # current vg beta
            vg_betas = [vg_beta] # vg beta round history
            vg_avPME = 0.0 # current vg exchange average
            vg_avPME_vgb0 = 0.0 # vg exchange average at vg_beta = 0
            biom_avPME = 0.0 # current biom exchange average
            biom_avPME_vgb0 = 0.0 # biom exchange average at vg_beta = 0
            biom_diff = 0.0 # distance between exp_growth and biom_avPME
            vg_diff = 0.0 # distance between cgD_X and vg_diff
            beta_vec = zeros(N) # ep beta vector
            epouts = Dict() # epout pool for each round (will contain the solution)
            epout = nothing # current epout (used as seed)
            hasvalid_moments = false # a flag that indicate is the momentous are valid
            roundconv = false # global (round) converge flag

            ## ----------------------------------------------------
            # ep params
            me_params = lglob(iJR, :maxent, :params)
            @extract me_params: alpha epsconv epmaxiter=maxiter
            @extract me_params: damp maxvar minvar

            ## ----------------------------------------------------
            # other params
            gdmaxiter = 3000 # maxiter for each gradient descent
            gdth = 0.01  # th of each gradient descend
            roundth = 0.01 # th of the whole simulation
            smooth = 0.1 # gd smooth th

            # After a while without converge, accelerate
            turbo_iter0 = 100 # iter for initing turbo
            turbo_frec = 10 # iter frec for apply turbo
            turbo_factor = 2.0 # turbo strength

            # damp will be reduced when detected
            damp_factor = 0.5 # gd damp penalty factor
            biom_gddamp = 1.0 # biom gd current damp
            vg_gddamp = 1.0 # vg gd current damp

            gdit = -1 # current gd iter
            gderr = -1.0 # current gd error
            last_infotime = -1.0 # time to check if gd needs to update
            upfrec_time = 15 # update info frequency

            # the whole simulation converge as ~log,
            # so I force betas increment by stepping
            beta_scale_rounditer0 = 3 # starting round for beta stepping
            beta_scale_factor = 0.5 # stepping scaling factor

            rounditer = 1 # current round iter
            maxrounds = 50 # max no of rounds

            # monitor
            SimT.record!(mon) do dat
                tdat = get!(dat, exp, Dict())
                tdat[:cgD_X] = cgD_X
                tdat[:method] = method
                tdat[:exp_growth] = exp_growth
            end

            ## -------------------------------------------------------------------
            # closure functions
            function check_roundconv()
                hasvalid_biom_moment = abs(biom_avPME - exp_growth)/abs(exp_growth) <= roundth
                hasvalid_vg_moment = abs(vg_avPME) <= abs(cgD_X) ||
                    abs(vg_avPME - cgD_X)/abs(cgD_X) <= roundth
                hasvalid_moments = hasvalid_biom_moment && hasvalid_vg_moment
                return hasvalid_moments
            end

            ## -------------------------------------------------------------------
            function print_info(msg; varargs...)

                isinfotime = gdit == 1 || abs(last_infotime - time()) > upfrec_time ||
                    epout.status != ChEP.CONVERGED_STATUS

                !isinfotime && return

                lock(WLOCK) do
                    @info(msg,
                        varargs...,
                        epout.status, epout.iter,
                        (biom_avPME_vgb0, biom_avPME, exp_growth), biom_diff,
                        (vg_avPME_vgb0, vg_avPME, cgD_X), vg_diff,
                        (biom_beta, vg_beta),
                        thid
                    ); println()
                    last_infotime = time()
                end
            end

            function check_nan()
                params = [biom_beta, vg_beta, biom_avPME, vg_avPME]
                any(isnan.(params))
            end

            ## -------------------------------------------------------------------
            function gd_core_fun(gdmodel; msg)

                beta_vec[biomidx] = biom_beta
                beta_vec[glcidx] = vg_beta

                # MAXENT
                epout = ChEP.maxent_ep(model;
                    beta_vec,
                    alpha, epsconv,
                    maxvar, minvar, damp,
                    maxiter = epmaxiter,
                    verbose = false,
                    solution = epout
                )

                # update
                biom_avPME = ChU.av(model, epout, iJR.BIOMASS_IDER)
                vg_avPME = ChU.av(model, epout, iJR.EX_GLC_IDER)
                biom_diff = abs(biom_avPME - exp_growth)
                vg_diff = abs(vg_avPME - cgD_X)

                gderr = gdmodel.ϵi
                print_info(msg;
                    exp, rounditer, gdit, gderr,
                    vg_gddamp, biom_gddamp,
                )

                # MONITOR
                SimT.record!(mon) do dat
                    tdat = get!(dat, exp, Dict())
                    tdat[:live_prove] = rand()
                    gddat = get!(tdat, :gd, Dict())
                    SimT.get!push!(gddat;
                        vg_beta, biom_beta,
                        biom_avPME, vg_avPME
                    )
                end

                # RUN OUT OF PATIENT
                (gdit >= turbo_iter0 && rem(gdit, turbo_frec) == 0) &&
                    (gdmodel.maxΔx *= turbo_factor)

                gdit += 1
            end

            while true

                ## -------------------------------------------------------------------
                # BETA SCALING
                scalebeta = rounditer >= beta_scale_rounditer0
                scalebeta && let
                    biom_beta_step = biom_betas[end] - biom_betas[end - 1]
                    biom_beta += biom_beta_step * beta_scale_factor

                    vg_beta_step = vg_betas[end] - vg_betas[end - 1]
                    vg_beta += vg_beta_step * beta_scale_factor
                end

                ## -------------------------------------------------------------------
                # Z GRAD DESCEND: Match biomass momentums
                let
                    target = exp_growth
                    x0 = biom_beta
                    maxΔx = max(abs(biom_beta) * 0.05, 5e2)
                    minΔx = maxΔx * 0.001
                    x1 = x0 + maxΔx * 0.01

                    gdit = 1

                    ## -------------------------------------------------------------------
                    function z_fun(gdmodel)
                        biom_beta = SimT.gd_value(gdmodel)
                        gd_core_fun(gdmodel; msg = "z grad descent... ")
                        return biom_avPME
                    end

                    ## -------------------------------------------------------------------
                    function z_break_cond(gdmodel)
                        roundconv = check_roundconv()
                        zconv = abs(biom_avPME - exp_growth)/abs(exp_growth) <= gdth
                        roundconv || zconv
                    end

                    ## -------------------------------------------------------------------
                    gdmodel = SimT.grad_desc(z_fun;
                        x0, x1, gdth, minΔx, maxΔx, smooth,
                        target, maxiter = gdmaxiter,
                        damp_factor, damp = biom_gddamp,
                        break_cond = z_break_cond,
                        verbose = false
                    )
                    biom_beta = SimT.gd_value(gdmodel)
                    biom_gddamp = gdmodel.damp
                end

                ## -------------------------------------------------------------------
                # AT VG BETA 0 MOMENTS
                isfirstround = (rounditer == 1)
                isfirstround && let
                    biom_avPME_vgb0 = biom_avPME
                    vg_avPME_vgb0 = vg_avPME
                end

                ## -------------------------------------------------------------------
                # FORCE VG BOUNDARY
                let
                    ## -------------------------------------------------------------------
                    # VG GRAD DESCEND: Match biomass momentums
                    target = cgD_X * 0.99 # force to be inside
                    x0 = vg_beta
                    maxΔx = max(abs(vg_beta) * 0.05, 1e2)
                    minΔx = maxΔx * 0.001
                    x1 = x0 + maxΔx * 0.01

                    gdit = 1

                    ## -------------------------------------------------------------------
                    function vg_fun(gdmodel)
                        vg_beta = SimT.gd_value(gdmodel)
                        gd_core_fun(gdmodel; msg = "vg grad descent... ")
                        return vg_avPME
                    end

                    ## -------------------------------------------------------------------
                    function vg_break_cond(epmodel)
                        vgconv = abs(vg_avPME) <= abs(cgD_X)
                        roundconv = check_roundconv()
                        roundconv || vgconv
                    end

                    ## -------------------------------------------------------------------
                    gdmodel = SimT.grad_desc(vg_fun;
                        x0, x1, gdth, minΔx, maxΔx,
                        break_cond = vg_break_cond,
                        damp_factor, damp = vg_gddamp,
                        target, maxiter = gdmaxiter,
                        verbose = false
                    )

                    vg_beta = SimT.gd_value(gdmodel)
                    vg_gddamp = gdmodel.damp
                end

                ## -------------------------------------------------------------------
                # COLLECTING
                push!(biom_betas, biom_beta)
                push!(vg_betas, vg_beta)
                epouts[(biom_beta, vg_beta)] = epout

                ## -------------------------------------------------------------------
                # CHECK NAN
                if check_nan()
                    print_info("Nan detected (BIG PROBLEMS HERE)";
                        exp, rounditer
                    ); break
                end

                ## -------------------------------------------------------------------
                # MONITOR
                SimT.record!(mon) do dat
                    tdat = get!(dat, exp, Dict())
                    tdat[:live_prove] = rand()
                    rdat = get!(tdat, :round, Dict())
                    SimT.get!push!(rdat;
                        vg_beta, biom_beta,
                        biom_avPME, vg_avPME
                    )
                end

                ## -------------------------------------------------------------------
                # PRINT INFO
                print_info("Round Done"; exp, rounditer,
                    hasvalid_moments, roundconv
                )

                ## -------------------------------------------------------------------
                # BREAK
                roundconv && break
                rounditer += 1
                rounditer > maxrounds && break

            end # round while

            ## -------------------------------------------------------------------
            # Storing
            lock(WLOCK) do

                dat = Dict()
                dat[:exp_beta] = (biom_beta, vg_beta)
                dat[:epouts] = epouts
                dat[:model] = model |> ChU.compressed_model

                datfile = dat_file(;method, exp)
                sdat(iJR, dat, datfile)

                print_info("Finished "; exp, rounditer)
            end

        end # for exp, cGLC
    end # for thid

    SimT.reset!(mon)

end


function simulation_init()
    A = Ne
    C = Ms+1
    D = Ms-Ne+1
    return Simulation{A, C, D}(
    zeros(Float64, N+2, 3), #x
    zeros(Float64, N+2, 3), #v
    MVector{3, Float64}(0.0, 0.0, 0.0),    #Δ_no
    zeros(Int64, N, 12), #nn_arr
    zeros(Float64, N+2, 3), #dhp_neutral
    zeros(Float64, N+2, 3), #dhp_ion
    zeros(Float64, N+2, 3), #dhp_coup
    MVector{3, Float64}(0.0, 0.0, 0.0), #hp
    zeros(Float64, N+2, 3), #F
    MVector{A, Int16}(zeros(Int16, Ne)), #surfp
    MVector{C, Int16}(zeros(Int16, Ms+1)), #occnum
    MVector{D, Int16}(zeros(Int16, Ms-Ne+1)), #surfh
    MVector{A, Int16}(zeros(Int16, Ne)), #surfpinit
    MVector{A, Int16}(zeros(Int16, Ne)), #surfpnew
    MVector{1, Float64}(10.0*Å), # trajzmin
    MVector{1, Float64}(0.0), #trajtheta
    zeros(Float64, Ms+1, Ms+1), #H
    zeros(ComplexF64, Ms+1, Ne), #psi
    zeros(ComplexF64, Ms+1, Ne), #phi
    zeros(Float64, Ms+1, Ms+1), #eigvec_H
    MVector{C, Float64}(zeros(Float64, Ms+1)), #eigval_H
    zeros(Float64, Ms+1, Ms+1), #dhdea
    zeros(Float64, Ms+1, Ms+1), #dhdv
    MVector{1, ComplexF64}(zero(ComplexF64)), #akl
    MVector{1, Float64}(0.0), #akk
    zeros(Float64, Ms+1,Ms+1), #DM
    zeros(Float64, Ne, Ms+1-Ne), #blk
    zeros(Float64, Ms+1, Ms+1), #blk2
    zeros(Float64, Ne*(Ms+1-Ne)), #Pb
    zeros(Float64, tsteps), #storage_aop
    zeros(Float64, Ms+1, tsteps), #storage_op,
    zeros(Float64, 5, tsteps), #storage_e
    zeros(Float64, 6, tsteps), #storage_xno
    zeros(Float64, 6, tsteps), #storage_vno
    zeros(Float64, tsteps), #storage_temp
    zeros(Float64, tsteps), #storage_phop
    MVector{1, ComplexF64}(zero(ComplexF64)), #"phipsi"
    MVector{1, Float64}(0.0), #"Pbmaxest"
    zeros(Float64, N+2, 3),#vdot
    zeros(Float64, N+2, 3),#vtemp
    MVector{1, Float64}(0.0), #vscale
    zeros(Float64, tsteps), #KEt
    zeros(Float64, tsteps), #PEt
    zeros(Float64, Int(round(tsteps/2))),
    zeros(Float64, Int(round(tsteps/2))),
    zeros(Float64, Int(round(tsteps/2))),
    zeros(Int64, Int(round(tsteps/2))),
    MVector{1, Int64}(0),
    MVector{1, Int64}(0),
    MVector{1, Int64}(0)
    )
end

function simulation_constructor_x_v!(s::Simulation)

    vvib_no = sqrt(e_vib_i*2.0/(mass_arr[1]*mass_arr[2]/(mass_arr[1] + mass_arr[2])))		# Initial vibrational velocity of NO
    vz_no = sqrt(e_trans_i*2.0/(mass_arr[1] + mass_arr[2]))				# Initial z velocity of NO

    #initial position of NO
    x0no = zeros(Float64, 2, 3)

    xno = rand() #temporary variable
    #xno = 0.1#rand() #temporary variable
    xno = @. xno * [cell[1], cell[2], 1.0*Å]
    #xno[3] = xno[3] + 4.5*Å # Initial z position of NO = 10
    xno[3] = xno[3] + 10.5*Å
    theta_no = rand() # Initial orientation of NO, 0 = O-down, PI = N-down
    #theta_no = 0.1#rand() # Initial orientation of NO, 0 = O-down, PI = N-down
    theta_no = asin(2.0 * theta_no - 1.0) + π/2.0
    phi_no = rand()
    #phi_no = 0.1#rand()
    phi_no = 2.0*π*phi_no

    x0no[1, :] = xno .+ r_no/2.0 * [sin(phi_no)*sin(theta_no), cos(phi_no)*sin(theta_no), cos(theta_no)]
    x0no[2, :] = xno .- r_no/2.0 * [sin(phi_no)*sin(theta_no), cos(phi_no)*sin(theta_no), cos(theta_no)]
    #x0no[1, 3] = x0no[1, 3] +1Å
    xno = x0no[1,:] - x0no[2, :]
    xno = xno / norm(xno)

    #initial velocity of no
    v0no = zeros(Float64, 2, 3)
    phi_no = 0.0
    v0no[1, :] = vz_no * [sin(phi_no)*sin(phi_inc), -cos(phi_no)*sin(phi_inc), -cos(phi_inc)]
    v0no[2, :] = v0no[1, :]
    v0no[1, :] = v0no[1, :] + vvib_no * xno * mass_arr[2] / (mass_arr[1] + mass_arr[2])
    v0no[2, :] = v0no[2, :] - vvib_no * xno * mass_arr[1] / (mass_arr[1] + mass_arr[2])

    #add rotational velocity
    theta_no = rand()
    #theta_no = 0.1#rand()
    theta_no = asin(2.0 * theta_no - 1.0) + π/2.0
    vrot_no = @. [-xno[3], 0.0, xno[1]] * cos(theta_no)/sqrt(xno[1]^2 + xno[3]^2)
    vrot_no = @. vrot_no + [xno[1], -(xno[1]^2 + xno[3]^2)/xno[2], xno[3]] *
                sin(theta_no)/sqrt((xno[1]^2 + xno[3]^2) * (1 + (xno[1]^2 + xno[3]^2)/xno[2]^2))

    vvib_no = sqrt(2.0 * e_rot_i/(mass_arr[1] + mass_arr[2]))
    v0no[1, :] = @. v0no[1, :] + vvib_no * vrot_no * sqrt(mass_arr[2]/mass_arr[1])
    v0no[2, :] = @. v0no[2, :] - vvib_no * vrot_no * sqrt(mass_arr[1]/mass_arr[2])
    #stack information to x_au0 and v_au0
    v0 = [v0no; v_au0]
    x0 = [x0no; x_au0]
    s.x .= x0
    s.v .= v0
    s.Δ_no .= x0[1, :] - x0[2, :]
end

function simulation_constructor_nn!(s::Simulation)
    nn_lattice!(s)
end

function simulation_constructor_energy(s::Simulation)
    s.hp .= get_E_all(s)
end

function simulation_constructor_force(s::Simulation)
    get_F_all(s)
end


#=
This file contains all function used to parse the data and return the data to an
appropriate format to be compute.

Link of the data : http://www.roadef.org/challenge/2005/fr/sujet.php
The data should be in ../data/Instances_Type/name_instance/
There is 3 types of instances labeled A, B and X, so folders
should be named Instances_A/, Instances_B/, and Instances_X/.
=#

# Constant of the problem
const OPTIMISATION_FILE_NAME = "optimization_objectives.txt"
const PAINT_FILE_NAME = "paint_batch_limit.txt"
const RATIO_FILE_NAME = "ratios.txt"
const VEHICLES_FILE_NAME = "vehicles.txt"

# An Instance structure that is used to format data as we want.
struct Instance
    # objectif function
    HPRC_rank::Int
    LPRC_rank::Int # If there is no LPRC, LPRC_rank = -1
    PCB_rank::Int
    # paint limitation
    nb_paint_limitation::Int
    # ratio constraint
    RC_p::Array{Int, 1}
    RC_q::Array{Int, 1}
    nb_HPRC::Int
    nb_LPRC::Int
    # sequence vehicle data
    RC_flag::Array{Bool, 2}
    RC_keys::Array{Int, 1}
    HPRC_keys::Array{Int, 1}
    LPRC_keys::Array{Int, 1}
    same_RC::Dict{Int, Array{Int, 1}}
    same_HPRC::Dict{Int, Array{Int, 1}}
    same_LPRC::Dict{Int, Array{Int, 1}}
    color_code::Array{Int, 1}
    same_color::Dict{Int, Array{Int, 1}}
    # Number of vehicles that were not build the precedent day.
    nb_late_prec_day::Int # Usage 1:nb_late_prec_day give the list of index of those vehicles.
    nb_cars::Int
end

"""
    parser(instance_name::String, instance_type::String, path_folder::String=string(@__DIR__)*"/../data/Instances_")

Returns return the instance of type instance_type parsed form the files in path_folder.
"""
function parser(instance_name::String, instance_type::String, path_folder::String=string(@__DIR__)*"/../data/Instances_")
    path = path_folder * instance_type * "/" * instance_name * "/"

    #Lecture  OPTIMISATION_FILE_NAME
    HPRC_rank = 0
    LPRC_rank = -1
    PCB_rank = 0
    open(path * OPTIMISATION_FILE_NAME) do f
        lines = readlines(f)
        for i in 2:length(lines)
            if lines[i] != ""
                if lines[i][3] == 'h'
                    HPRC_rank = i-1
                elseif lines[i][3] == 'l'
                    LPRC_rank = i-1
                elseif lines[i][3] == 'p'
                    PCB_rank = i-1
                end
            end
        end
    end
    @assert HPRC_rank != 0
    @assert PCB_rank != 0

    #Lecture PAINT_FILE_NAME
    nb_paint_limitation = 0
    open(path * PAINT_FILE_NAME) do f
        lines = readlines(f)
        nb_paint_limitation = parse(Int, lines[2][1:(end-1)])
    end

    #Lecture RATIO_FILE_NAME
    nb_HPRC = 0
    nb_LPRC = 0
    RC_p = Array{Int,1}()
    RC_q = Array{Int,1}()
    open(path * RATIO_FILE_NAME) do f
        lines = readlines(f)
        for line in lines[2:end]
            if line != ""
                values = split(line, ";", limit=2)
                ratio = split(values[1], "/")
                push!(RC_p, parse(Int, ratio[1]))
                push!(RC_q, parse(Int, ratio[2]))
                if values[2][1] == '1'
                    nb_HPRC += 1
                elseif values[2][1] == '0'
                    nb_LPRC += 1
                end
            end
        end
    end

    # TODO:
    # string(Int(RC_flag[car, option])) * RC_value
    # become RC_value *string(Int(RC_flag[car, option]))
    # then:
    # two cars -> difference in RC value is less than 2^nb_LPRC
    # one array not needed but more time to compute difference later
    #Lecture VEHICLES_FILE_NAME
    day = 0
    indice = 0
    nb_RC = nb_HPRC + nb_LPRC
    color_code = Array{Int, 1}()
    RC_keys = Array{Int, 1}()
    HPRC_keys = Array{Int, 1}()
    LPRC_keys = Array{Int, 1}()
    RC_flag = Array{Bool, 2}(undef, 0, nb_RC)
    same_RC = Dict{Int, Array{Int, 1}}()
    same_HPRC = Dict{Int, Array{Int, 1}}()
    same_LPRC = Dict{Int, Array{Int, 1}}()
    same_color = Dict{Int, Array{Int, 1}}()
    nb_cars = 0
    nb_cars_total = 0
    nb_late_prec_day = 0
    open(path * VEHICLES_FILE_NAME) do f
        lines = readlines(f)
        if lines[end] == ""
            nb_cars_total = length(lines) - 2
        else
            nb_cars_total = length(lines) - 1
        end
        RC_flag = Array{Bool, 2}(undef, nb_cars_total, nb_RC)
        color_code = Array{Int, 1}(undef, nb_cars_total)
        RC_keys = Array{Int, 1}(undef, nb_cars_total)
        HPRC_keys = Array{Int, 1}(undef, nb_cars_total)
        LPRC_keys = Array{Int, 1}(undef, nb_cars_total)
        for line in lines[2:end]
            if line != ""
                values = split(line, ";", keepempty = false)
                nb_cars += 1
                car_RC_value = "0"
                for rc in 5:length(values)
                    RC_flag[nb_cars, rc - 4] = (values[rc] == "1")
                    car_RC_value = values[rc] * car_RC_value
                end
                RC_keys[nb_cars] = parse(Int, car_RC_value, base = 2)
                HPRC_keys[nb_cars] = parse(Int, SubString(car_RC_value, nb_LPRC+1), base = 2)
                LPRC_keys[nb_cars] = parse(Int, "0"*SubString(car_RC_value, 1:nb_LPRC), base = 2)
                color_code[nb_cars] = parse(Int, values[4])
                if nb_late_prec_day == 0
                    date = split(values[1], " ")
                    if day == 0
                        day = date[3]
                    elseif date[3] != day
                        day = date[3]
                        nb_late_prec_day = nb_cars-1
                        same_RC[RC_keys[nb_cars]] = Array{Int, 1}([nb_cars])
                        same_HPRC[HPRC_keys[nb_cars]] = Array{Int, 1}([nb_cars])
                        same_LPRC[LPRC_keys[nb_cars]] = Array{Int, 1}([nb_cars])
                        same_color[color_code[nb_cars]] = Array{Int, 1}([nb_cars])
                    end
                else
                    if !haskey(same_RC, RC_keys[nb_cars])
                        same_RC[RC_keys[nb_cars]] = Array{Int, 1}([nb_cars])
                        if !haskey(same_HPRC, HPRC_keys[nb_cars])
                            same_HPRC[HPRC_keys[nb_cars]] = Array{Int, 1}([nb_cars])
                        else
                            push!(same_HPRC[HPRC_keys[nb_cars]], nb_cars)
                        end
                        if !haskey(same_LPRC, LPRC_keys[nb_cars])
                            same_LPRC[LPRC_keys[nb_cars]] = Array{Int, 1}([nb_cars])
                        else
                            push!(same_LPRC[LPRC_keys[nb_cars]], nb_cars)
                        end
                    else
                        push!(same_RC[RC_keys[nb_cars]], nb_cars)
                        push!(same_HPRC[HPRC_keys[nb_cars]], nb_cars)
                        push!(same_LPRC[LPRC_keys[nb_cars]], nb_cars)
                    end
                    if !(color_code[nb_cars] in keys(same_color))
                        same_color[color_code[nb_cars]] = Array{Int, 1}([nb_cars])
                    else
                        push!(same_color[color_code[nb_cars]], nb_cars)
                    end
                end
            end
        end
    end
    @assert nb_cars_total == nb_cars

    # for i in (nb_late_prec_day+1):nb_cars
    #     @assert i in same_RC[RC_keys[i]]
    #     @assert i in same_LPRC[LPRC_keys[i]]
    #     @assert i in same_HPRC[HPRC_keys[i]]
    # end

    return Instance(
            HPRC_rank, LPRC_rank, PCB_rank,                            # objectives file
            nb_paint_limitation,                                       # paint file
            RC_p, RC_q, nb_HPRC, nb_LPRC,                              # ratio file
            RC_flag, RC_keys, HPRC_keys, LPRC_keys,
            same_RC, same_HPRC, same_LPRC,
            color_code, same_color, nb_late_prec_day, nb_cars
        )
end


module BandedMatrices
using Base, FillArrays, LazyArrays

using LinearAlgebra, SparseArrays, Random
using LinearAlgebra.LAPACK
import Base: axes1, getproperty, iterate
import LinearAlgebra: BlasInt,
                    BlasReal,
                    BlasFloat,
                    BlasComplex,
                    axpy!,
                    copy_oftype,
                    checksquare,
                    adjoint,
                    transpose,
                    AdjOrTrans, HermOrSym
import LinearAlgebra.BLAS: libblas
import LinearAlgebra.LAPACK: liblapack, chkuplo, chktrans
import LinearAlgebra: cholesky, cholesky!, norm, diag, eigvals!, eigvals,
            qr, axpy!, ldiv!, mul!, lu, lu!, AbstractTriangular, has_offset_axes,
            chkstride1, kron, lmul!, rmul!
import SparseArrays: sparse

import Base: getindex, setindex!, *, +, -, ==, <, <=, >,
                >=, /, ^, \, transpose, showerror, reindex, checkbounds, @propagate_inbounds

import Base: convert, size, view, unsafe_indices,
                first, last, size, length, unsafe_length, step,
                to_indices, to_index, show, fill!, promote_op,
                MultiplicativeInverses, OneTo, ReshapedArray,
                               similar, copy, convert, promote_rule, rand,
                            IndexStyle, real, imag, Slice, pointer, unsafe_convert, copyto!

import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted, broadcasted

import LazyArrays: MemoryLayout, blasmul!, @lazymul, @lazylmul, @lazyldiv, @blasmatvec, @blasmatmat,
                    AbstractStridedLayout, AbstractColumnMajor, AbstractRowMajor,
                    _copyto!, MatMulVec, MatMulMat, transposelayout, triangulardata,
                    ConjLayout, conjlayout, SymmetricLayout, symmetriclayout, symmetricdata,
                    triangularlayout, InverseLayout, MatMulVec, MatLdivVec, TriangularLayout,
                    ArrayMulArrayStyle, HermitianLayout, hermitianlayout, hermitiandata
import FillArrays: AbstractFill

export BandedMatrix,
       bandrange,
       brand,
       bandwidth,
       BandError,
       band,
       Band,
       BandRange,
       bandwidths,
       colrange,
       rowrange,
       isbanded,
       Zeros,
       Fill,
       Ones,
       Eye

include("blas.jl")
include("lapack.jl")

include("generic/AbstractBandedMatrix.jl")
include("generic/broadcast.jl")
include("generic/matmul.jl")
include("generic/Band.jl")
include("generic/utils.jl")

include("banded/BandedMatrix.jl")
include("banded/BandedLU.jl")
include("banded/BandedQR.jl")
include("banded/gbmm.jl")
include("banded/linalg.jl")

include("symbanded/symbanded.jl")
include("symbanded/BandedCholesky.jl")

include("tribanded.jl")

include("interfaceimpl.jl")

@deprecate setindex!(A::BandedMatrix, v, b::Band) A[b] .= v
@deprecate setindex!(A::BandedMatrix, v, ::BandRangeType, j::Integer) A[BandRange,j] .= v
@deprecate setindex!(A::BandedMatrix, v, kr::Colon, j::Integer) A[:,j] .= v
@deprecate setindex!(A::BandedMatrix, v, kr::AbstractRange, j::Integer) A[kr,j] .= v
@deprecate setindex!(A::BandedMatrix, v, ::Colon, ::Colon) A[:,:] .= v
@deprecate setindex!(A::BandedMatrix, v, ::Colon) A[:] .= v
@deprecate setindex!(A::BandedMatrix, v, k::Integer, ::BandRangeType) A[k,BandRange] .= v
@deprecate setindex!(A::BandedMatrix, v, k::Integer, jr::AbstractRange) A[k,jr] .= v
@deprecate setindex!(A::BandedMatrix, v, kr::AbstractRange, jr::AbstractRange) A[kr,jr] .= v
@deprecate setindex!(A::BandedMatrix, v, ::BandRangeType) BandedMatrices.bandeddata(A) .= v

end #module


__precompile__()
module ObjectTracker

export
    Blob,
    form_blobs,
    is_transient,
    match!,
    Object,
    predict,
    State,
    update

include("Kalman.jl")
include("Blobs.jl")
include("Tracker.jl")
end # module


# import Base: getindex, size, *

struct SFMDistribution{T<:SFFM.Mesh} <: AbstractArray{Float64,2}
    coeffs::Array{Float64,2}
    model::SFFM.Model
    mesh::T
    Fil::SFFM.IndexDict
    function SFMDistribution{T}(
        coeffs::Array{Float64,2},
        model::SFFM.Model,
        mesh::T,
        Fil::SFFM.IndexDict=SFFM.MakeFil(model,mesh.Nodes),
        ) where T<:SFFM.Mesh
        return (size(coeffs,1)==1) ? new(coeffs,model,mesh,Fil) : throw(DimensionMismatch("coeffs must be a row-vector"))
    end
end

size(d::SFMDistribution) = size(d.coeffs)
getindex(d::SFMDistribution,i::Int,j::Int) = d.coeffs[i,j]
setindex!(d::SFMDistribution,x,i::Int,j::Int) = throw(DomainError("inserted value(s) must be Float64/Int"))
setindex!(d::SFMDistribution,x::Float64,i::Int,j::Int) = (d.coeffs[i,j]=x)
setindex!(d::SFMDistribution,x::Int,i::Int,j::Int) = (d.coeffs[i,j]=convert(Float64,x))
*(u::SFMDistribution,B::AbstractArray{Float64,2}) = *(u.coeffs,B)
*(B::AbstractArray{Float64,2},u::SFMDistribution) =
    (size(B,2)==1) ? *(B*u.coeffs) : throw(DimensionMismatch("u is a row-vector and B has more than 1 column"))

function _error_on_nothing(idx)
    (idx===nothing) && throw(DomainError("x is not in the support of the mesh"))
end

function _get_nodes_coeffs_from_index(cell_idx::Int,i::Int,mesh::Mesh,model::Model)
    cell_nodes = CellNodes(mesh)[:,cell_idx]
    N₋ = sum(model.C.<=0)
    coeff_idx = (N₋ + (i-1)*TotalNBases(mesh) + (cell_idx-1)*NBases(mesh)) .+ (1:NBases(mesh))
    return cell_nodes, coeff_idx
end

function _get_coeff_index_pos(x::Float64,i::Int,mesh::Mesh,model::Model)
    cell_idx = findlast(x.>=mesh.Nodes)
    _error_on_nothing(cell_idx)
    cell_nodes, coeff_idx = _get_nodes_coeffs_from_index(cell_idx,i,mesh,model)
    return cell_idx, cell_nodes, coeff_idx
end
function _get_coeff_index_neg(x::Float64,i::Int,mesh::Mesh,model::Model)
    cell_idx = findfirst(x.<=mesh.Nodes) - 1
    _error_on_nothing(cell_idx)
    cell_nodes, coeff_idx = _get_nodes_coeffs_from_index(cell_idx,i,mesh,model)
    return cell_idx, cell_nodes, coeff_idx
end

function _get_point_mass_data_pos(i::Int,mesh::Mesh,model::Model)
    cell_nodes = mesh.Nodes[end]
    N₋ = sum(model.C.<=0)
    coeff_idx = N₋ + TotalNBases(mesh)*NPhases(model) + sum(model.C[1:i].>=0)
    return cell_nodes, coeff_idx
end
function _get_point_mass_data_neg(i::Int,mesh::Mesh,model::Model)
    cell_nodes = mesh.Nodes[1]
    coeff_idx = sum(model.C[1:i].<=0)
    return cell_nodes, coeff_idx
end
_is_left_point_mass(x::Float64,i::Int,mesh::Mesh,model::Model) =
    (x==mesh.Nodes[1])&&(model.C[i]<=0)
_is_right_point_mass(x::Float64,i::Int,mesh::Mesh,model::Model) =
    (x==mesh.Nodes[end])&&(model.C[i]>=0)

function _get_coeffs_index(x::Float64,i::Int,model::Model,mesh::Mesh)
    !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))

    # find which inteval Dₗ,ᵢ x is in
    if _is_left_point_mass(x,i,mesh,model)
        cell_idx = "point mass"
        cell_nodes, coeff_idx = _get_point_mass_data_neg(i,mesh,model)
    elseif _is_right_point_mass(x,i,mesh,model)
        cell_idx = "point mass"
        cell_nodes, coeff_idx = _get_point_mass_data_pos(i,mesh,model)
    else # not a point mass
        if model.C[i]>=0
            cell_idx, cell_nodes, coeff_idx = _get_coeff_index_pos(x,i,mesh,model)
        elseif model.C[i]<0
            cell_idx, cell_nodes, coeff_idx = _get_coeff_index_neg(x,i,mesh,model)
        end
    end

    return cell_idx, cell_nodes, coeff_idx
end

function legendre_to_lagrange(coeffs)
    order = length(coeffs)
    V = vandermonde(NBases(mesh))
    return V.V*coeffs
end

function pdf(d::SFMDistribution{T},model::Model) where T<:Mesh
    throw(DomainError("unknown SFMDistribution{<:Mesh}"))
end

function pdf(d::SFMDistribution{DGMesh},model::Model)
    function f(x::Float64,i::Int) # the PDF
        # check phase is in support
        !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))
        # if x is not in the support return 0.0
        mesh = d.mesh
        if ((x<mesh.Nodes[1])||(x>mesh.Nodes[end]))
            fxi = 0.0
        else
            cell_idx, cell_nodes, coeff_idx = _get_coeffs_index(x,i,model,mesh)
            coeffs = d.coeffs[coeff_idx]
            # if not a point mass, then reconstruct solution
            if !(cell_idx=="point mass")
                if Basis(mesh) == "legendre"
                    coeffs = legendre_to_lagrange(coeffs)
                else
                    V = vandermonde(NBases(mesh))
                    coeffs = (2/(Δ(mesh)[cell_idx]))*(1.0./V.w).*coeffs
                end
                basis_values = lagrange_poly_basis(cell_nodes, x)
                fxi = LinearAlgebra.dot(basis_values,coeffs)
            else
                fxi = coeffs
            end
        end
        return fxi
    end
    return f
end
pdf(d::SFMDistribution{T},model::Model,x,i) where T<:Mesh =
    throw(DomainError("x must be Float64/Int/Array{Float64/Int,1}, i must be Int/Array{Int,1}"))
pdf(d::SFMDistribution{T},model::Model,x::Float64,i::Int) where T<:Mesh = pdf(d,model)(x,i)
pdf(d::SFMDistribution{T},model::Model,x::Int,i::Int) where T<:Mesh = pdf(d,model)(convert(Float64,x),i)
# pdf(d::SFMDistribution{T},model::Model,x::Array{Float64,1},i::Int) where T<:Mesh = pdf(d,model).(x,i)
# pdf(d::SFMDistribution{T},model::Model,x::Array{Float64,1},i::Array{Int,1}) where T<:Mesh = pdf(d,model).(x,i)

function pdf(d::SFMDistribution{FRAPMesh},model::Model)
    function f(x::Float64,i::Int) # the PDF
        # check phase is in support
        !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))
        # if x is not in the support return 0.0
        mesh = d.mesh
        if ((x<mesh.Nodes[1])||(x>mesh.Nodes[end]))
            fxi = 0.0
        else
            cell_idx, cell_nodes, coeff_idx = _get_coeffs_index(x,i,model,mesh)
            coeffs = d.coeffs[coeff_idx]
            # if not a point mass, then reconstruct solution
            if !(cell_idx=="point mass")
                if model.C[i]>0
                    yₖ₊₁ = mesh.Nodes[cell_idx+1]
                    to_go = yₖ₊₁-x
                elseif model.C[i]<0
                    yₖ = mesh.Nodes[cell_idx]
                    to_go = x-yₖ
                end
                me = MakeME(CMEParams[NBases(mesh)], mean = Δ(mesh)[cell_idx])
                fxi = (pdf(Array(coeffs'),me,to_go) + pdf(Array(coeffs'),me,2*Δ(mesh)[cell_idx]-to_go))./cdf(Array(coeffs'),me,2*Δ(mesh)[cell_idx])
            else
                fxi = coeffs
            end
        end
        return fxi
    end
    return f
end

function pdf(d::SFMDistribution{FVMesh},model::Model)
    function f(x::Float64,i::Int) # the PDF
        # check phase is in support
        !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))
        # if x is not in the support return 0.0
        mesh = d.mesh
        if ((x<mesh.Nodes[1])||(x>mesh.Nodes[end]))
            fxi = 0.0
        else
            cell_idx, cell_nodes, coeff_idx = _get_coeffs_index(x,i,model,mesh)
            # if not a point mass, then reconstruct solution
            if !(cell_idx=="point mass")
                ptsLHS = Int(ceil(Order(mesh)/2))
                if cell_idx-ptsLHS < 0
                    nodesIdx = 1:Order(mesh)
                    nodes = CellNodes(mesh)[nodesIdx]
                    poly_vals = lagrange_poly_basis(nodes,x)
                elseif cell_idx-ptsLHS+Order(mesh) > TotalNBases(mesh)
                    nodesIdx = (TotalNBases(mesh)-Order(mesh)+1):TotalNBases(mesh)
                    nodes = CellNodes(mesh)[nodesIdx]
                    poly_vals = lagrange_poly_basis(nodes,x)
                else
                    nodesIdx =  (cell_idx-ptsLHS) .+ (1:Order(mesh))
                    poly_vals = lagrange_poly_basis(CellNodes(mesh)[nodesIdx],x)
                end
                coeff_idx = (sum(model.C.<=0) + (i-1)*TotalNBases(mesh)) .+ nodesIdx
                coeffs = d.coeffs[coeff_idx]#./Δ(mesh)[cell_idx]
                fxi = LinearAlgebra.dot(poly_vals,coeffs)
            else
                coeffs = d.coeffs[coeff_idx]
                fxi = coeffs
            end
        end
        return fxi
    end
    return f
end

############
### CDFs ###
############

function cdf(d::SFMDistribution{T},model::Model) where T<:Mesh
    throw(DomainError("unknown SFMDistribution{<:Mesh}"))
end

"""

    _sum_cells_left(d::SFMDistribution, i::Int, cell_idx::Int, mesh::Mesh, model::Model)

Add up all the probability mass in phase `i` in the cells to the left of `cell_idx`.
"""
function _sum_cells_left(d::SFMDistribution, i::Int, cell_idx::Int, mesh::Mesh, model::Model)
    c = 0.0
    if Basis(mesh) == "legendre"
        for cell in 1:(cell_idx-1)
            # first legendre basis function =1 & has all the mass
            idx = (sum(model.C.<=0) + (i-1)*TotalNBases(mesh) + (cell-1)*NBases(mesh)) .+ 1
            c += d.coeffs[idx]
        end
    else
        for cell in 1:(cell_idx-1)
            idx = (sum(model.C.<=0) + (i-1)*TotalNBases(mesh) + (cell-1)*NBases(mesh)) .+ (1:NBases(mesh))
            c += sum(d.coeffs[idx])
        end
    end
    return c
end

function cdf(d::SFMDistribution{DGMesh},model::Model)
    function F(x::Float64,i::Int) # the PDF
        # check phase is in support
        !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))
        mesh = d.mesh
        Fxi = 0.0
        if (x<mesh.Nodes[1])
            # Fxi = 0.0
        else
            # Fxi = 0.0
            # left pm
            if (x>=mesh.Nodes[1])&&(model.C[i]<=0)
                ~, left_pm_idx = _get_point_mass_data_neg(i,mesh,model)
                left_pm = d.coeffs[left_pm_idx]
                Fxi += left_pm
            end
            # integral over density
            (x.>=mesh.Nodes[end]) ? (xd=mesh.Nodes[end]-sqrt(eps())) : xd = x
            cell_idx, ~, ~ = _get_coeffs_index(xd,i,model,mesh)
            if !(cell_idx=="point mass")
                # add all mass from cells to the left
                Fxi += _sum_cells_left(d, i, cell_idx, mesh, model)

                # integrate up to x in the cell which contains x
                temp_pdf(y) = pdf(d,model)(y,i)
                quad = gauss_lobatto_quadrature(temp_pdf,mesh.Nodes[cell_idx]+sqrt(eps()),xd,NBases(mesh))
                Fxi += quad
            end
            # add the RH point mass if  required
            if (x>=mesh.Nodes[end])&&(model.C[i]>=0)
                ~, ~, right_pm_idx = _get_coeffs_index(mesh.Nodes[end],i,model,mesh)
                right_pm = d.coeffs[right_pm_idx]
                Fxi += right_pm
            end
        end
        return Fxi
    end
    return F
end
cdf(d::SFMDistribution{T},model::Model,x,i) where T<:Mesh =
    throw(DomainError("x must be Float64/Int/Array{Float64/Int,1}, i must be Int/Array{Int,1}"))
cdf(d::SFMDistribution{T},model::Model,x::Float64,i::Int) where T<:Mesh = cdf(d,model)(x,i)
cdf(d::SFMDistribution{T},model::Model,x::Int,i::Int) where T<:Mesh = cdf(d,model)(convert(Float64,x),i)

function cdf(d::SFMDistribution{FRAPMesh},model::Model)
    function F(x::Float64,i::Int) # the PDF
        # check phase is in support
        !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))
        # if x is not in the support return 0.0
        mesh = d.mesh
        if (x<mesh.Nodes[1])
            Fxi = 0.0
        else
            Fxi = 0.0
            # left pm
            if (x>=mesh.Nodes[1])&&(model.C[i]<=0)
                ~, ~, left_pm_idx = _get_coeffs_index(mesh.Nodes[1],i,model,mesh)
                left_pm = d.coeffs[left_pm_idx]
                Fxi += left_pm
            end
            # integral over density
            (x.>=mesh.Nodes[end]) ? (xd=mesh.Nodes[end]-sqrt(eps())) : xd = x
            cell_idx, cell_nodes, coeff_idx = _get_coeffs_index(xd,i,model,mesh)
            coeffs = d.coeffs[coeff_idx]

            if !(cell_idx=="point mass")
                # add all mass from cells to the left
                Fxi += _sum_cells_left(d, i, cell_idx, mesh, model)

                # integrate up to x in the cell which contains x
                me = MakeME(CMEParams[NBases(mesh)], mean = Δ(mesh)[cell_idx])
                a = Array(coeffs')
                if model.C[i]>=0
                    yₖ₊₁ = mesh.Nodes[cell_idx+1]
                    Fxi += (ccdf(a,me,yₖ₊₁-x) - ccdf(a,me,2*Δ(mesh)[cell_idx]-(yₖ₊₁-x)))/(cdf(a,me,2*Δ(mesh)[cell_idx]))
                elseif model.C[i]<0
                    yₖ = mesh.Nodes[cell_idx]
                    Fxi += sum(a) - (ccdf(a,me,x-yₖ) - ccdf(a,me,2*Δ(mesh)[cell_idx]-(x-yₖ)))/(cdf(a,me,2*Δ(mesh)[cell_idx]))
                end
            end
            if (x>=mesh.Nodes[end])&&(model.C[i]>=0)
                ~, ~, right_pm_idx = _get_coeffs_index(mesh.Nodes[end],i,model,mesh)
                right_pm = d.coeffs[right_pm_idx]
                Fxi += right_pm
            end
        end
        return Fxi
    end
    return F
end

function _sum_cells_left_fv(d, i, cell_idx, mesh, model)
    c = 0
    for cell in 1:(cell_idx-1)
        # first legendre basis function =1 & has all the mass
        idx = sum(model.C.<=0) + (i-1)*TotalNBases(mesh) + cell
        c += d.coeffs[idx]*Δ(mesh)[cell]
    end
    return c
end

function cdf(d::SFMDistribution{FVMesh},model::Model)
    function F(x::Float64,i::Int) # the PDF
        # check phase is in support
        !(i∈phases(model)) && throw(DomainError("phase i must be in the support of the model"))
        # if x is not in the support return 0.0
        mesh = d.mesh
        Fxi = 0.0
        if (x<mesh.Nodes[1])
            # Fxi = 0.0
        else
            # Fxi = 0.0
            # left pm
            if (x>=mesh.Nodes[1])&&(model.C[i]<=0)
                ~, ~, left_pm_idx = _get_coeffs_index(mesh.Nodes[1],i,model,mesh)
                left_pm = d.coeffs[left_pm_idx]
                Fxi += left_pm
            end
            # integral over density
            (x.>=mesh.Nodes[end]) ? (xd=mesh.Nodes[end]-sqrt(eps())) : xd = x
            cell_idx, cell_nodes, coeff_idx = _get_coeffs_index(xd,i,model,mesh)
            # if not a point mass, then reconstruct solution
            if !(cell_idx=="point mass")
                # add all mass from cells to the left
                Fxi += _sum_cells_left_fv(d, i, cell_idx, mesh, model)

                # integrate up to x in the cell which contains x
                temp_pdf(y) = pdf(d,model)(y,i)
                quad = gauss_lobatto_quadrature(temp_pdf,mesh.Nodes[cell_idx]+sqrt(eps()),xd,Order(mesh))
                Fxi += quad
            end
            if (x>=mesh.Nodes[end])&&(model.C[i]>=0)
                ~, ~, right_pm_idx = _get_coeffs_index(mesh.Nodes[end],i,model,mesh)
                right_pm = d.coeffs[right_pm_idx]
                Fxi += right_pm
            end
        end
        return Fxi
    end
    return F
end

abstract type SFFMDistribution end

"""

    SFFMDistribution(
        pm::Array{<:Real},
        distribution::Array{<:Real,3},
        x::Array{<:Real},
        type::String,
    )

- `pm::Array{Float64}`: a vector containing the point masses, the first
    `sum(model.C.<=0)` entries are the left hand point masses and the last
    `sum(model.C.>=0)` are the right-hand point masses.
- `distribution::Array{Float64,3}`: "probability" or "density"`
- `x::Array{Float64,2}`:
    - if `type="probability"` is a `1×NIntervals×NPhases` array
        containing the cell centers.
    - if `type="density"` is a `NBases×NIntervals×NPhases` array
        containing the cell nodes at which the denisty is evaluated.
- `type::String`: either `"probability"` or `"density"`. `"cumulative"` is
    not possible.
"""
struct SFFMDensity <: SFFMDistribution
    pm::Array{<:Real}
    distribution::Array{<:Real,3}
    x::Array{<:Real}
end
struct SFFMProbability <: SFFMDistribution
    pm::Array{<:Real}
    distribution::Array{<:Real,3}
    x::Array{<:Real}
end
struct SFFMCDF <: SFFMDistribution
    pm::Array{<:Real}
    distribution::Array{<:Real,3}
    x::Array{<:Real}
end

"""
Convert from a vector of coefficients for the DG system to a distribution.

    Coeffs2Dist(
        model::SFFM.Model,
        mesh::SFFM.Mesh,
        Coeffs;
        type::String = "probability",
    )

# Arguments
- `model`: a Model object
- `mesh`: a Mesh object as output from MakeMesh
- `Coeffs::Array`: a vector of coefficients from the DG method
- `type::String`: an (optional) declaration of what type of distribution you
    want to convert to. Options are `"probability"` to return the probabilities
    ``P(X(t)∈ D_k, φ(t) = i)`` where ``D_k``is the kth cell, `"cumulative"` to
    return the CDF evaluated at cell edges, or `"density"` to return an
    approximation to the density ar at the SFFM.CellNodes(mesh).

# Output
- a tuple with keys
(pm=pm, distribution=yvals, x=xvals, type=type)
    - `pm::Array{Float64}`: a vector containing the point masses, the first
        `sum(model.C.<=0)` entries are the left hand point masses and the last
        `sum(model.C.>=0)` are the right-hand point masses.
    - `distribution::Array{Float64,3}`:
        - if `type="cumulative"` returns a `2×NIntervals×NPhases` array
            containing the CDF evaluated at the cell edges as contained in
            `x` below. i.e. `distribution[1,:,i]` returns the cdf at the
            left-hand edges of the cells in phase `i` and `distribution[2,:,i]`
            at the right hand edges.
        - if `type="probability"` returns a `1×NIntervals×NPhases` array
            containing the probabilities ``P(X(t)∈ D_k, φ(t) = i)`` where ``D_k``
            is the kth cell.
        - if `type="density"` returns a `NBases×NIntervals×NPhases` array
            containing the density function evaluated at the cell nodes as
            contained in `x` below.
    - `x::Array{Float64,2}`:
        - if `type="cumulative"` returns a `2×NIntervals×NPhases` array
            containing the cell edges as contained. i.e. `x[1,:]`
            returns the left-hand edges of the cells and `x[2,:]` at the
            right-hand edges.
        - if `type="probability"` returns a `1×NIntervals×NPhases` array
            containing the cell centers.
        - if `type="density"` returns a `NBases×NIntervals×NPhases` array
            containing the cell nodes.
    - `type`: as input in arguments.
"""
function Coeffs2Dist(
    model::SFFM.Model,
    mesh::SFFM.DGMesh,
    Coeffs::AbstractArray,
    type::Type{T} = SFFMProbability,
    v::Bool = false,
) where {T<:SFFMDistribution}

    V = SFFM.vandermonde(NBases(mesh))
    N₋ = sum(model.C .<= 0)
    N₊ = sum(model.C .>= 0)

    if type == SFFMDensity
        xvals = SFFM.CellNodes(mesh)
        if Basis(mesh) == "legendre"
            yvals = reshape(Coeffs[N₋+1:end-N₊], NBases(mesh), NIntervals(mesh), NPhases(model))
            for i in 1:NPhases(model)
                yvals[:,:,i] = V.V * yvals[:,:,i]
            end
            pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]
        elseif Basis(mesh) == "lagrange"
            yvals =
                Coeffs[N₋+1:end-N₊] .* repeat(1.0 ./ V.w, NIntervals(mesh) * NPhases(model)) .*
                (repeat(2.0 ./ Δ(mesh), 1, NBases(mesh) * NPhases(model))'[:])
            yvals = reshape(yvals, NBases(mesh), NIntervals(mesh), NPhases(model))
            pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]
        end
        if NBases(mesh) == 1
            yvals = [1;1].*yvals
            xvals = [SFFM.CellNodes(mesh)-Δ(mesh)'/2;SFFM.CellNodes(mesh)+Δ(mesh)'/2]
        end
    elseif type == SFFMProbability
        if NBases(mesh) > 1
            xvals = SFFM.CellNodes(mesh)[1, :] + (Δ(mesh) ./ 2)
        else
            xvals = SFFM.CellNodes(mesh)
        end
        if Basis(mesh) == "legendre"
            yvals = (reshape(Coeffs[N₋+1:NBases(mesh):end-N₊], 1, NIntervals(mesh), NPhases(model)).*Δ(mesh)')./sqrt(2)
            pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]
        elseif Basis(mesh) == "lagrange"
            yvals = sum(
                reshape(Coeffs[N₋+1:end-N₊], NBases(mesh), NIntervals(mesh), NPhases(model)),
                dims = 1,
            )
            pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]
        end
    elseif type == SFFMCDF
        if NBases(mesh) > 1
            xvals = SFFM.CellNodes(mesh)[[1;end], :]
        else
            xvals = [SFFM.CellNodes(mesh)-Δ(mesh)'/2;SFFM.CellNodes(mesh)+Δ(mesh)'/2]
        end
        if Basis(mesh) == "legendre"
            tempDist = (reshape(Coeffs[N₋+1:NBases(mesh):end-N₊], 1, NIntervals(mesh), NPhases(model)).*Δ(mesh)')./sqrt(2)
            pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]
        elseif Basis(mesh) == "lagrange"
            tempDist = sum(
                reshape(Coeffs[N₋+1:end-N₊], NBases(mesh), NIntervals(mesh), NPhases(model)),
                dims = 1,
            )
            pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]
        end
        tempDist = cumsum(tempDist,dims=2)
        temppm = zeros(Float64,1,2,NPhases(model))
        temppm[:,1,model.C.<=0] = pm[1:N₋]
        temppm[:,2,model.C.>=0] = pm[N₊+1:end]
        yvals = zeros(Float64,2,NIntervals(mesh),NPhases(model))
        yvals[1,2:end,:] = tempDist[1,1:end-1,:]
        yvals[2,:,:] = tempDist
        yvals = yvals .+ reshape(temppm[1,1,:],1,1,NPhases(model))
        pm[N₋+1:end] = pm[N₋+1:end] + yvals[end,end,model.C.>=0]
    end

    out = type(pm, yvals, xvals)
    v && println("UPDATE: distribution object created with keys ", fieldnames(type))
    return out
end
function Coeffs2Dist(
    model::SFFM.Model,
    mesh::Union{FRAPMesh, FVMesh},
    Coeffs::AbstractArray,
    type::Type{T} = SFFMProbability,
    v::Bool = false,
) where {T<:SFFMDistribution}

    if type != SFFMProbability
        args = [
            model;
            mesh;
            Coeffs;
            type;
        ]
        error("Input Error: no functionality other than 'probability' implemented, yet...")
    end

    N₋ = sum(model.C .<= 0)
    N₊ = sum(model.C .>= 0)

    xvals = SFFM.CellNodes(mesh)

    yvals = sum(
        reshape(Coeffs[N₋+1:end-N₊], NBases(mesh), NIntervals(mesh), NPhases(model)),
        dims = 1,
    )
    pm = [Coeffs[1:N₋]; Coeffs[end-N₊+1:end]]

    out = type(pm, yvals, xvals)
    v && println("UPDATE: distribution object created with keys ", fieldnames(type))
    return out
end

"""
Converts a distribution as output from `Coeffs2Dist()` to a vector of DG
coefficients.

    Dist2Coeffs(
        model::SFFM.Model,
        mesh::SFFM.Mesh,
        Distn::SFFMDistribution,
    )

# Arguments
- `model`: a Model object
- `mesh`: a Mesh object as output from MakeMesh
- `Distn::SFFMDistribution
    - if `type="probability"` is a `1×NIntervals×NPhases` array containing
        the probabilities ``P(X(t)∈ D_k, φ(t) = i)`` where ``D_k``
        is the kth cell.
    - if `type="density"` is a `NBases×NIntervals×NPhases` array containing
        either the density function evaluated at the cell nodes which are in
        `x` below, or, the inner product of the density function against the
        lagrange polynomials.

# Output
- `coeffs` a row vector of coefficient values of length
    `TotalNBases*NPhases + N₋ + N₊` ordered according to LH point masses, RH
    point masses, interior basis functions according to basis function, cell,
    phase. Used to premultiply operators such as B from `MakeB()`
"""
function Dist2Coeffs(
    model::SFFM.Model,
    mesh::SFFM.DGMesh,
    Distn::SFFMDensity,
)
    V = SFFM.vandermonde(NBases(mesh))
    theDistribution =
        zeros(Float64, NBases(mesh), NIntervals(mesh), NPhases(model))
    if Basis(mesh) == "legendre"
        theDistribution = Distn.distribution
        for i = 1:NPhases(model)
            theDistribution[:, :, i] = V.inv * theDistribution[:, :, i]
        end
    elseif Basis(mesh) == "lagrange"
        theDistribution .= Distn.distribution
        # convert to probability coefficients by multiplying by the
        # weights in V.w/2 and cell widths Δ
        theDistribution = ((V.w .* theDistribution).*(Δ(mesh) / 2)')[:]
    end
    # also put the point masses on the ends
    coeffs = [
        Distn.pm[1:sum(model.C .<= 0)]
        theDistribution[:]
        Distn.pm[sum(model.C .<= 0)+1:end]
    ]
    coeffs = Matrix(coeffs[:]')
    return coeffs
end

function Dist2Coeffs(
    model::SFFM.Model,
    mesh::SFFM.DGMesh,
    Distn::SFFMProbability,
)
    V = SFFM.vandermonde(NBases(mesh))
    theDistribution =
        zeros(Float64, NBases(mesh), NIntervals(mesh), NPhases(model))
    if Basis(mesh) == "legendre"
        # for the legendre basis the first basis function is ϕ(x)=Δ√2 and
        # all other basis functions are orthogonal to this. Hence, we map
        # the cell probabilities to the first basis function only.
        theDistribution[1, :, :] = Distn.distribution./Δ(mesh)'.*sqrt(2)
    elseif Basis(mesh) == "lagrange"
        theDistribution .= Distn.distribution
        # convert to probability coefficients by multiplying by the
        # weights in V.w/2
        theDistribution = (V.w .* theDistribution / 2)[:]
    end
    # also put the point masses on the ends
    coeffs = [
        Distn.pm[1:sum(model.C .<= 0)]
        theDistribution[:]
        Distn.pm[sum(model.C .<= 0)+1:end]
    ]
    coeffs = Matrix(coeffs[:]')
    return coeffs
end
function Dist2Coeffs(
    model::SFFM.Model,
    mesh::Union{SFFM.FRAPMesh,SFFM.FVMesh},
    Distn::SFFMDistribution
)

    # also put the point masses on the ends
    coeffs = [
        Distn.pm[1:sum(model.C .<= 0)]
        Distn.distribution[:]
        Distn.pm[sum(model.C .<= 0)+1:end]
    ]

    coeffs = Matrix(coeffs[:]')
    return coeffs
end

"""
Computes the error between distributions.

    starSeminorm(d1::SFFMProbability, d2::SFFMProbability)

# Arguments
- `d1`: a distribution object as output from `Coeffs2Dist`
- `d2`: a distribution object as output from `Coeffs2Dist`
"""
function starSeminorm(d1::SFFMProbability, d2::SFFMProbability)
    e = sum(abs.(d1.pm-d2.pm)) + sum(abs.(d1.distribution-d2.distribution))
    return e
end