module MochaClassification

using Compat
using Mocha
using MochaClassificationWrapper
using MLUtilities
using Utilities
using DataModel

export MochaClassificationModel

"""
Neural Network Classification using Mocha

fields:
    - backend: Mocha Backend (CPU or GPU)
    - mochaNet: Mocha wrapper
    - labels
    - ntrain: number of data items
    - batchsize
"""
type MochaClassificationModel <: AbstractDataModel
    backend::Backend
    mochaNet::MochaClassificationWrapper.MWrap
    labels::Array{Int64,1}
    ntrain::Float64 # training set size
    batchsize::Int64
end

function MochaClassificationModel(
    datax,datac,
    modelfactory::Function,
    backend::Backend;
    ntrain::Int = length(datac),
    batchsize::Int = 100,
    do_shuffle::Bool = true,
    do_accuracy::Bool = true,
    do_predprob::Bool = false
    )

    data_layer = MemoryDataLayer(name = "data",
                               data = Array[datax,datac],
                               batch_size = batchsize,
                               shuffle = do_shuffle)

    mochaNet = MochaClassificationWrapper.MWrap(data_layer,
                                 modelfactory,
                                 "MochaClassificationNet",
                                 do_accuracy,
                                 do_predprob,
                                 backend)
    MochaClassificationModel(backend,mochaNet,datac[:],ntrain,batchsize)
end

"""
`DataModel.getllik(dms::MochaClassificationModel)` returns a function that provides the full loglikelihood
"""
function DataModel.getllik(dms::MochaClassificationModel)
    function llik(x)
        (accuracy, loglikelihood) = MochaClassificationWrapper.evaluateTestNN(dms.mochaNet, x, dms.batchsize)
        return loglikelihood
    end
end

"""
`DataModel.getgrad(dms::MochaClassificationModel)` returns function returning the stochastic gradient function.
"""
function DataModel.getgrad(dms::MochaClassificationModel)
    function grad(x)
        (llik, grad) = MochaClassificationWrapper.forwardbackward(dms.mochaNet, x, regu_coef = 0.)
        return grad * dms.ntrain
    end
end

"""
`getparams(dms::MochaClassificationModel)` reads out the current weights.
"""
function getparams(dms::MochaClassificationModel)
    MochaClassificationWrapper.getparams(dms.mochaNet)
end

"""
`setparams!(dms::MochaClassificationModel,para::Array{Float64,1` sets the weights to `para`.
"""
function setparams!(dms::MochaClassificationModel,para::Array{Float64,1})
    MochaClassificationWrapper.setparams!(dms.mochaNet,para)
end
"""
`fetchnparams(dms::MochaClassificationModel)` returns the number of parameters of the model.
"""
function fetchnparams(dms::MochaClassificationModel)
    MochaClassificationWrapper.getnparams(dms.mochaNet)
end

#-------------------
# Initialisation functions
#-------------------
"""
Xavier Initialisation
"""
function init_xavier(dms::MochaClassificationModel)
    MochaClassificationWrapper.init_xavier(dms.mochaNet)
end
"""
Initialisation based on fan-in.
"""
function init_simple_fanin(dms::MochaClassificationModel)
    MochaClassificationWrapper.init_simple_fanin(dms.mochaNet)
end
"""
Gaussian Initialisation
"""
function init_gaussian(dms::MochaClassificationModel, initvar::Float64)
    MochaClassificationWrapper.init_gaussian(dms.mochaNet, initvar)
end
"""
Uniform Initialisation
"""
function init_uniform(dms::MochaClassificationModel, initvar::Float64)
    MochaClassificationWrapper.init_uniform(dms.mochaNet, initvar)
end

#-------------------
# evaluation functions
#-------------------
"""
Evaluate accuracy based on predictive probabilities
"""
function evaluate(dms::MochaClassificationModel,
    predprobs::Array{Float64,2})
    prediction =  (findmax(predprobs,1)[2]-1.) % size(predprobs)[1]
    accuracy = sum(dms.labels .== prediction[:])./(length(dms.labels)+0.0)
    loglikelihood = sum(log([predprobs[i,j] for (i,j) in zip(dms.labels+1,1:round(Int,dms.ntrain))]))
    @dict(accuracy, loglikelihood)
end

"""
Evaluate accuracy at weights `x`.
"""
function evaluate(dms::MochaClassificationModel,
    x::Vector{Float64})
    @assert length(x) == fetchnparams(dms)
    (accuracy, loglikelihood) = MochaClassificationWrapper.evaluateTestNN(dms.mochaNet, x, dms.batchsize)
    @dict(accuracy, loglikelihood)
end

"""
Get predictive probabilities given weights `x`
"""
function getpredprobs(dms::MochaClassificationModel,
    x::Vector{Float64})
    predictiveprobs = MochaClassificationWrapper.getpredprobs(dms.mochaNet, x, dms.batchsize)
    @dict(predictiveprobs)
end

end


#= = # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   fluxtable.txt
#
#   This file stores the F10.7 in different formats.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # = =#

struct _fluxtable_Structure
    F10obs::_space_indices_itp_constant{Float64}
    F10adj::_space_indices_itp_constant{Float64}
end

# Remote file: fluxtable.txt
_fluxtable = @RemoteFile(
    "ftp://ftp.seismo.nrcan.gc.ca/spaceweather/solar_flux/daily_flux_values/fluxtable.txt",
    file = "fluxtable.txt",
    dir = "$path_si_root",
    updates = :daily
)

# Optional variable that will store the `fluxtable.txt` data.
@OptionalData(_fluxtable_data, _fluxtable_Structure,
              "Run `init_space_indices()` with `:fluxtable` in `enabled_files` array to initialize required data.")

################################################################################
#                              Private Functions
################################################################################

"""
    _init_fluxtable(;force_download = false, local_path = nothing)

Initialize the data in the file `fluxtable.txt` by creating `_fluxtable_data`.
The initialization process is composed of:

1. Download the file, if it is necessary;
2. Parse the file;
3. Create the interpolations and the structures.

If the keyword `force_download` is `true`, then the file will always be
downloaded.

The user can also specify a location for the file using the keyword
`local_path`. If it is `nothing`, which is the default, then the file will be
downloaded.

"""
function _init_fluxtable(;force_download=false, local_path=nothing)
    # Update the remote files if no path is given.
    if local_path == nothing
        download(_fluxtable; force=force_download)
        local_path = path(_fluxtable)
    end

    push!(_fluxtable_data,   _parse_fluxtable(local_path))

    nothing
end

"""
    _parse_fluxtable(path::AbstractString)

Parse the `fluxtable.txt` file in `path` and return an instance of the structure
`_fluxtable_Structure` with the initialize interpolations.

"""
function _parse_fluxtable(path::AbstractString)
    # Allocate raw data.
    DT      = DateTime[]
    F10obs  = Float64[]
    F10adj  = Float64[]

    open(path) do file
        line_num = 0

        for ln in eachline(file)
            line_num += 1

            # Jump comments at the beginning of the file.
            (line_num in [1,2]) && continue

            # Get the tokens.
            tokens   = split(ln)
            num_toks = length(tokens)

            # Check if the line is valid.
            (num_toks != 7) && @error "Error parsing the line $line_num of fluxtable.txt"

            # We will save only the flux obtained at 20h. This is what
            # NRLMSISE-00 and JB2008 uses.
            fluxtime = parse(Int, tokens[2])
            (fluxtime != 200000) && continue

            # The DT of the data will be computed at noon. Hence, we will be
            # able to use the nearest-neighbor algorithm in the interpolations.
            year  = parse(Int, tokens[1][1:4])
            month = parse(Int, tokens[1][5:6])
            day   = parse(Int, tokens[1][7:8])
            DT_k  = DateTime(year, month, day, 12, 0, 0)

            # Get the raw data.

            if isempty(DT) || DT_k != DT[end]
                push!(DT,      DT_k)
                push!(F10obs,  parse(Float64, tokens[5]))
                push!(F10adj,  parse(Float64, tokens[6]))
            else

                DT[end] =  DT_k
                F10obs[end] = parse(Float64, tokens[5])
                F10adj[end] = parse(Float64, tokens[6])
            end
        end
    end

    # Create the interpolations for each parameter.
    knots = (datetime2julian.(DT),)
    # @show knots
    # @show F10obs
    # throw(error("test"))
    itp_F10obs  = interpolate(knots, F10obs,  Gridded(Constant()))
    itp_F10adj  = interpolate(knots, F10adj,  Gridded(Constant()))

    _fluxtable_Structure(itp_F10obs, itp_F10adj)
end


"""

    TransData{FT<:AbstractFloat}

Store information relevant to the particular geometrical configuration, not intrinsic to
molecules themselves, involving rigid geometrical transformations (translations and
rotations); this depends on the number of molecules `nummols` from
`MolSystem{T<:OneMol{<:AbstractFloat}}`.

`numTrans` is the number of transformations ``t`` (an integer).

`transArray` is a `3*nummols`-by-`numTrans` real array where each 3-element vector
`transArray[3*nn-2:3*nn, tt]` is the center-of-mass translation vector (relative to the
baseline) for molecule `nn` in transformation labeled `tt`.

`rotArray` is a `3*nummols`-by-`3*numTrans` real array where each 3-by-3 block
`rotArray[3*nn-2:3*nn, 3*tt-2:3*tt]` is the center-of-mass rotation matrix (relative to the
baseline) for molecule `nn` in transformation labeled `tt`.

`transLabels` is a `numTrans`-element vector of strings labeling each transformation `tt`.

`changedFromBefore` is an `nummols`-by-`numTrans` `BitArray` such that
`changedFromBefore[nn, tt]` is `true` if molecule `nn` is transformed in transformation `tt`
relative to transformation `tt - 1`, `false` otherwise.

`allAtomPos` is the horizontal concatenation of atomic coordinates (initially unmodified,
later modified, by transformations) across all molecules.

`baselineTrans` is a `3*nummols`-element real vector where each 3-element vector
`baselineTrans[3*nn-2:3*nn]` is the baseline center of mass coordinate for molecule `nn`.

baselineRot is a `3*nummols`-by-3 real matrix where each 3-by-3 block
`baselineRot[3*nn-2:3*nn,:]` is the baseline center of mass rotation matrix for the
orientation of molecule `nn` (using the alibi convention).
"""
struct TransData{FT<:AbstractFloat}

    transArray::Array{FT, 2}
    rotArray::Array{FT, 2}
    transLabels::Array{<:AbstractString, 1}
    changedFromBefore::BitArray{2}
    numTrans::Integer
    allAtomPos::Array{FT, 2}
    baselineTrans::Array{FT, 1}
    baselineRot::Array{FT, 2}

end

"""

    readTransData(myMolData::MolSystem{<:OneMol{FT}},
                  translistfilename::AbstractString="") where FT<:AbstractFloat

Construct TransData from existing molecular data `myMolData` and a file of name
`translistfilename` containing information about geometric transformations; if no file is
specified, produce a default set of transformations (no displacements or rotations at all).

See also: [`MolSystem{<:OneMol{FT}}`](@ref)

"""
function readTransData(myMolData::MolSystem{<:OneMol{FT}},
                       translistfilename::AbstractString="") where FT<:AbstractFloat

    ## default: no transformation file specified
    if (translistfilename == "")
        numtrans = 1;
        allmoltrans = zeros(FT, numtrans, 6*myMolData.nummols);
        allmoltranslabels = ["DEFAULT"];
        changedfrombefore = falses(myMolData.nummols, numtrans);
        transArray = zeros(FT, 3*myMolData.nummols, numtrans);
        rotArray = repeat(Matrix{FT}(I, 3, 3), myMolData.nummols, numtrans);

        allAtomPos = Array{FT,2}(undef, 3, myMolData.cumnumatomslist[end]);
        concatenateAtomPos!(myMolData, allAtomPos);

        baselineTrans = zeros(FT, 3*myMolData.nummols);
        baselineRot = repeat(Matrix{FT}(I, 3, 3), myMolData.nummols, 1);

        myTransData = TransData{FT}(transArray, rotArray, allmoltranslabels,
                                    changedfrombefore, numtrans, allAtomPos,
                                    baselineTrans, baselineRot);
        initializeAtomPos!(myTransData, myMolData);
        return myTransData;
    end

    # read file and format each line to get rid of extraneous spaces and things like that
    transstrarray = strip.(chomp.(readlines(translistfilename)));
    transstartinds = findall(map(u -> startswith(u, "TRANS "), transstrarray));
    transendinds = findall(transstrarray .== "ENDTRANS");
    numtrans = length(transendinds);
    if (length(transstartinds) != numtrans)
        error(string(translistfilename, ": number of TRANS and ENDTRANS statements",
                     " should be the same"));
    end
    transbothinds = vec(transpose(hcat(transstartinds, transendinds)));
    if (transbothinds != sort(transbothinds))
        error(string(translistfilename, ": TRANS and ENDTRANS statements should",
                     " alternate in ordered occurrence"));
    end

    # populate the lists of transformation labels, translation vectors, and Euler angles
    allmoltranslabels = Array{String, 1}(undef, numtrans);
    allmoltrans = zeros(FT, 6*myMolData.nummols + 1, numtrans);
    for tt=1:numtrans
        allmoltranslabels[tt] = replace(strip(chomp(transstrarray[transstartinds[tt]])),
                                        "TRANS " => "");
        currtransstrarray = strip.(chomp.(transstrarray[transstartinds[tt]+1:transendinds[tt]-1]));
        for currtransstr in currtransstrarray
            currtransstrsplit = split(currtransstr);
            currmollabel = currtransstrsplit[1];
            currmolind = findfirst(myMolData.molLabels .== currmollabel);

            if (currmolind != nothing)
                ## assemble transformation list in the following order for each row in a transformation:
                ## [phi_1, theta_1, psi_1, ..., phi_N, theta_N, psi_N, x_1, y_1, z_1, ..., x_N, y_N, z_N]
                ## Note: all Euler angles should be given in the transformation file as ALIBI transformations in DEGREES
                allmoltrans[3*myMolData.nummols+3*currmolind-2,
                            tt] = parse(FT,currtransstrsplit[2]);
                allmoltrans[3*myMolData.nummols+3*currmolind-1,
                            tt] = parse(FT, currtransstrsplit[3]);
                allmoltrans[3*myMolData.nummols+3*currmolind,
                            tt] = parse(FT, currtransstrsplit[4]);
                if (length(currtransstrsplit) == 7)
                    allmoltrans[3*currmolind-2, tt] = parse(FT, currtransstrsplit[5]);
                    allmoltrans[3*currmolind-1, tt] = parse(FT, currtransstrsplit[6]);
                    allmoltrans[3*currmolind, tt] = parse(FT, currtransstrsplit[7]);
                end
            end
        end
        allmoltrans[end, tt] = tt;
    end

    # populate the lists of baseline translation vectors and Euler angles
    baselinestartind = findfirst(map(u -> startswith(u, "BASELINE"), transstrarray));
    baselineendind = findfirst(map(u -> startswith(u, "ENDBASELINE"), transstrarray));
    baselinetrans = zeros(FT, 6*myMolData.nummols);

    if (baselinestartind != nothing && baselineendind != nothing)
        baselinestrarray = strip.(chomp.(transstrarray[(baselinestartind+
                                                        1):(baselineendind-1)]));
        if (length(baselinestrarray) > 0)
            for currbaselinestr in baselinestrarray
                currbaselinestrsplit = split(currbaselinestr);
                currmollabel = currbaselinestrsplit[1];
                currmolind = findfirst(myMolData.molLabels .== currmollabel);

                if (currmolind != nothing)
                    ## Note: all Euler angles should be given in the transformation file in DEGREES
                    baselinetrans[3*myMolData.nummols+3*currmolind-2] = parse(FT, currbaselinestrsplit[2]);
                    baselinetrans[3*myMolData.nummols+3*currmolind-1] = parse(FT, currbaselinestrsplit[3]);
                    baselinetrans[3*myMolData.nummols+3*currmolind] = parse(FT, currbaselinestrsplit[4]);
                    if (length(currbaselinestrsplit) == 7)
                        baselinetrans[3*currmolind-2] = parse(FT, currbaselinestrsplit[5]);
                        baselinetrans[3*currmolind-1] = parse(FT, currbaselinestrsplit[6]);
                        baselinetrans[3*currmolind] = parse(FT, currbaselinestrsplit[7]);
                    end
                end
            end
        end
    end

    # sort such that the displacement of the last molecule changes first
    # between transformation labels, then the displacement of the
    # next-to-last molecule, and so on, then the rotation of the last
    # molecule, and so on
    allmoltranssorted = sortslices(allmoltrans, dims=2);
    sortindvec = round.(Int, vec(allmoltranssorted[end, :]));
    permute!(allmoltranslabels, sortindvec);
    allmoltrans = allmoltranssorted[1:end-1, :];

    # convert Euler angles into (alibi) rotation matrices
    allmoltranslations = deepcopy(allmoltrans[3*myMolData.nummols+1:end, :]);
    allmolrotations = Array{FT, 2}(undef, 3*myMolData.nummols, 3*numtrans);
    for tt=1:numtrans
        for nn=1:myMolData.nummols
            if (allmoltrans[3*nn-2:3*nn, tt] == zeros(FT, 3))
                allmolrotations[3*nn-2:3*nn, 3*tt-2:3*tt] = Matrix{FT}(I, 3, 3);
            else
                cosphi = cosd(allmoltrans[3*nn-2, tt]);
                sinphi = sind(allmoltrans[3*nn-2, tt]);
                costheta = cosd(allmoltrans[3*nn-1, tt]);
                sintheta = sind(allmoltrans[3*nn-1, tt]);
                cospsi = cosd(allmoltrans[3*nn, tt]);
                sinpsi = sind(allmoltrans[3*nn, tt]);

                allmolrotations[3*nn-2, 3*tt-2] = costheta*cospsi;
                allmolrotations[3*nn-2, 3*tt-1] = costheta*sinpsi;
                allmolrotations[3*nn-2, 3*tt] = -1.0*sintheta;
                allmolrotations[3*nn-1, 3*tt-2] = sinphi*sintheta*cospsi - cosphi*sinpsi;
                allmolrotations[3*nn-1, 3*tt-1] = sinphi*sintheta*sinpsi + cosphi*cospsi;
                allmolrotations[3*nn-1, 3*tt] = sinphi*costheta;
                allmolrotations[3*nn, 3*tt-2] = cosphi*sintheta*cospsi + sinphi*sinpsi;
                allmolrotations[3*nn, 3*tt-1] = cosphi*sintheta*sinpsi - sinphi*cospsi;
                allmolrotations[3*nn, 3*tt] = cosphi*costheta;
            end
        end
    end

    baselinetranslations = deepcopy(baselinetrans[3*myMolData.nummols+1:end]);
    baselinerotations = Array{FT, 2}(undef, 3*myMolData.nummols, 3);
    for nn=1:myMolData.nummols
        if (baselinetrans[3*nn-2:3*nn] == zeros(FT, 3))
            baselinerotations[3*nn-2:3*nn, :] = Matrix{FT}(I, 3, 3);
        else
            cosphi = cosd(baselinetrans[3*nn-2]);
            sinphi = sind(baselinetrans[3*nn-2]);
            costheta = cosd(baselinetrans[3*nn-1]);
            sintheta = sind(baselinetrans[3*nn-1]);
            cospsi = cosd(baselinetrans[3*nn]);
            sinpsi = sind(baselinetrans[3*nn]);

            baselinerotations[3*nn-2, 1] = costheta*cospsi;
            baselinerotations[3*nn-2, 2] = costheta*sinpsi;
            baselinerotations[3*nn-2, 3] = -1.0*sintheta;
            baselinerotations[3*nn-1, 1] = sinphi*sintheta*cospsi - cosphi*sinpsi;
            baselinerotations[3*nn-1, 2] = sinphi*sintheta*sinpsi + cosphi*cospsi;
            baselinerotations[3*nn-1, 3] = sinphi*costheta;
            baselinerotations[3*nn, 1] = cosphi*sintheta*cospsi + sinphi*sinpsi;
            baselinerotations[3*nn, 2] = cosphi*sintheta*sinpsi - sinphi*cospsi;
            baselinerotations[3*nn, 3] = cosphi*costheta;
        end
    end

    # after sorting, determine if each successive transformation affects a given molecule
    changedfrombefore = trues(myMolData.nummols, numtrans);
    for tt=2:numtrans
        for currmolind=1:myMolData.nummols
            if (allmoltrans[3*myMolData.nummols+3*currmolind-2, tt] ==
                allmoltrans[3*myMolData.nummols+3*currmolind-2, tt-1] &&
                allmoltrans[3*myMolData.nummols+3*currmolind-1, tt] ==
                allmoltrans[3*myMolData.nummols+3*currmolind-1, tt-1] &&
                allmoltrans[3*myMolData.nummols+3*currmolind, tt] ==
                allmoltrans[3*myMolData.nummols+3*currmolind, tt-1] &&
                allmoltrans[3*currmolind-2, tt] == allmoltrans[3*currmolind-2, tt-1] &&
                allmoltrans[3*currmolind-1, tt] == allmoltrans[3*currmolind-1, tt-1] &&
                allmoltrans[3*currmolind, tt] == allmoltrans[3*currmolind, tt-1])

                changedfrombefore[currmolind, tt] = false;

            end
        end
    end

    # horizontally concatenate all atomic coordinates over all (not just unique) molecules
    allAtomPos = Array{FT, 2}(undef, 3, myMolData.cumnumatomslist[end]);

    allAtomPos[:, 1:myMolData.numatomslist[1]] .=
    myMolData.uniqueMolDataArray[1].atompos0[:, :];

    for nn=2:myMolData.nummols

        currmolDupeInd = myMolData.dupemolinds[nn];
        if (currmolDupeInd == -1)
            allAtomPos[:,
                       myMolData.cumnumatomslist[nn-1]+1:myMolData.cumnumatomslist[nn]] .=
            myMolData.uniqueMolDataArray[myMolData.uniquemolinds[nn]].atompos0[:, :];
        else
            allAtomPos[:,
                       myMolData.cumnumatomslist[nn-1]+1:myMolData.cumnumatomslist[nn]] .=
            myMolData.uniqueMolDataArray[myMolData.dupemolinds[nn]].atompos0[:, :];
        end

    end

    # create and return the TransData
    myTransData = TransData{FT}(allmoltranslations, allmolrotations,
                                allmoltranslabels, changedfrombefore,
                                numtrans, allAtomPos, baselinetranslations,
                                baselinerotations);
    initializeAtomPos!(myTransData, myMolData);
    return myTransData;

end

"""

    initializeAtomPos!(myTransData::TransData{FT},
                       myMolData::MolSystem{<:OneMol{FT}}) where FT<:AbstractFloat

Initialize atomic coordinates across all molecules by first shifting coordinates of all
molecules to have centers of mass of 0, then rotating all molecular coordinates according to
baseline orientations of each, then translating to desired baseline centers of mass,
modifying `myTransData` in-place.

See also: [`transformAtomPos!`](@ref), [`untransformAtomPos!`](@ref)

"""
function initializeAtomPos!(myTransData::TransData{FT},
                            myMolData::MolSystem{<:OneMol{FT}}) where FT<:AbstractFloat
    startatom = 1;
    endatom = myMolData.numatomslist[1];
    for nn=1:myMolData.nummols
        if (nn > 1)
            startatom = myMolData.cumnumatomslist[nn-1] + 1;
            endatom = myMolData.cumnumatomslist[nn];
        end
        makeCOM0!(myTransData.allAtomPos, startatom, endatom);
        if (myTransData.baselineRot[3*nn-2:3*nn, :] != Matrix{FT}(I, 3, 3))
            rotatecoords!(myTransData.allAtomPos, myTransData.baselineRot[3*nn-2:3*nn, :],
                          startatom, endatom);
        end
        if (myTransData.baselineTrans[3*nn-2:3*nn] != zeros(FT, 3))
            translateCOM!(myTransData.allAtomPos, myTransData.baselineTrans[3*nn-2:3*nn],
                          startatom, endatom);
        end
    end
end

"""

    transformAtomPos!(myTransData::TransData{FT},
                      myMolData::MolSystem{<:OneMol{FT}},
                      tt::Integer) where FT<:AbstractFloat

Transform atomic coordinates across all molecules by first shifting coordinates of all
molecules to have centers of mass of 0, then rotating all molecular coordinates according to
orientations corresponding to transformation `tt` relative to the baseline, then translating
to desired centers of mass corresponding to transformation `tt` relative to the baseline,
modifying `myTransData` in-place.

See also: [`initializeAtomPos!`](@ref), [`untransformAtomPos!`](@ref)

"""
function transformAtomPos!(myTransData::TransData{FT},
                           myMolData::MolSystem{<:OneMol{FT}},
                           tt::Integer) where FT<:AbstractFloat
    startatom = 1;
    endatom = myMolData.numatomslist[1];
    for nn=1:myMolData.nummols
        if (nn > 1)
            startatom = myMolData.cumnumatomslist[nn-1] + 1;
            endatom = myMolData.cumnumatomslist[nn];
        end
        if (myTransData.baselineTrans[3*nn-2:3*nn] != zeros(FT, 3))
            makeCOM0!(myTransData.allAtomPos, startatom, endatom);
        end
        if (myTransData.rotArray[3*nn-2:3*nn, 3*tt-2:3*tt] != Matrix{FT}(I, 3, 3))
            rotatecoords!(myTransData.allAtomPos,
                          myTransData.rotArray[3*nn-2:3*nn, 3*tt-2:3*tt],
                          startatom, endatom);
        end
        if (myTransData.baselineTrans[3*nn-2:3*nn] .+
            myTransData.transArray[3*nn-2:3*nn, tt] != zeros(FT, 3))
            translateCOM!(myTransData.allAtomPos,
                          myTransData.baselineTrans[3*nn-2:3*nn] .+
                          myTransData.transArray[3*nn-2:3*nn, tt],
                          startatom, endatom);
        end
    end
end

"""

    untransformAtomPos!(myTransData::TransData{FT},
                        myMolData::MolSystem{<:OneMol{FT}},
                        tt::Integer) where FT<:AbstractFloat

Undo transformation of atomic coordinates across all molecules by first shifting coordinates
of all molecules to have centers of mass of 0, then rotating all molecular coordinates
according to orientations corresponding to the opposites of transformation `tt` to return to
the baseline, then translating back to the desired baseline centers of mass, modifying
`myTransData` in-place.

See also: [`initializeAtomPos!`](@ref), [`transformAtomPos!`](@ref)

"""
function untransformAtomPos!(myTransData::TransData{FT},
                             myMolData::MolSystem{<:OneMol{FT}},
                             tt::Integer) where FT<:AbstractFloat
    startatom = 1;
    endatom = myMolData.numatomslist[1];
    for nn=1:myMolData.nummols
        if (nn > 1)
            startatom = myMolData.cumnumatomslist[nn-1] + 1;
            endatom = myMolData.cumnumatomslist[nn];
        end
        if (myTransData.baselineTrans[3*nn-2:3*nn] .+
            myTransData.transArray[3*nn-2:3*nn, tt] != zeros(FT, 3))
            makeCOM0!(myTransData.allAtomPos, startatom, endatom);
        end
        if (myTransData.rotArray[3*nn-2:3*nn, 3*tt-2:3*tt] != Matrix{FT}(I, 3, 3))
            rotatecoords!(myTransData.allAtomPos,
                          transpose(myTransData.rotArray[3*nn-2:3*nn,3*tt-2:3*tt]),
                          startatom, endatom);
        end
        if (myTransData.baselineTrans[3*nn-2:3*nn] != zeros(FT, 3))
            translateCOM!(myTransData.allAtomPos,
                          myTransData.baselineTrans[3*nn-2:3*nn],
                          startatom, endatom);
        end
    end
end


module CylinderModel

using GridapGmsh
using Gridap.Io
to_json_file(GmshDiscreteModel("mesh.msh"),"mesh.json")

end


# using CSV, Dates, HDF5, Statistics, Parameters

hydrooptions() = Dict(
    :gisregion => "Europe8",                    # "Europe8", "Eurasia38", "Scand3"

    :costclasses_min => [ 0,  50,  100],         # US $/MWh
    :costclasses_max => [50, 100, 9999],

    :storageclasses_min => [   0, 1e-6,  12],   # weeks (discharge time)
    :storageclasses_max => [1e-6,   12, 9e9]
)

mutable struct HydroOptions
    gisregion               ::String
    costclasses_min         ::Vector{Float64}
    costclasses_max         ::Vector{Float64}
    storageclasses_min      ::Vector{Float64}
    storageclasses_max      ::Vector{Float64}
end

HydroOptions() = HydroOptions("",[],[],[],[])

function HydroOptions(d::Dict{Symbol,Any})
    options = HydroOptions()
    for (key,val) in d
        setproperty!(options, key, val)
    end
    return options
end

function getregion(lon, lat, regions, lonrange=1:36000, latrange=1:18000)
    res = 0.01
    res2 = res/2
    lons = (-180+res2:res:180-res2)[lonrange]         # longitude values (pixel center)
    lats = (90-res2:-res:-90+res2)[latrange]          # latitude values (pixel center)
    flon = findfirst(round(lon-res2, digits=5) .< lons .<= round(lon+res2, digits=5))
    flat = findfirst(round(lat-res2, digits=5) .< lats .<= round(lat+res2, digits=5))
    (flon == nothing || flat == nothing) && return 0
    # println(lons[flon],", ",lats[flat])
    # println(flon, " ", flat)
    if regions[flon, flat] > 0
        return regions[flon, flat]
    else
        return maximum(regions[max(1,flon-2):min(length(lons),flon+2), max(1,flat-2):min(length(lats),flat+2)])
    end
end

function GIShydro(; optionlist...)
    options = HydroOptions(merge(hydrooptions(), optionlist))

    potential, existing, hydroplants, WECcapacity, WECpotential = readhydrodatabases()

    dayspermonth, nordicprofile, eleccost, capacity, reservoirenergy, dischargetime, monthlyinflow =
        potentialhydrovars(potential)

    countries, countrynames, regions, regionlist, lonrange, latrange = read_countries_and_regions(options)

    existingcapac, existinginflowcf =
        existinginflow(existing, countries, countrynames, regions, regionlist, lonrange, latrange, hydroplants, WECcapacity,
                        dayspermonth)

    potentialcapac, potentialinflowcf, potentialmeancost, potentialmeandischargetime =
        potentialinflow(options, potential, capacity, regions, regionlist, lonrange, latrange, monthlyinflow, eleccost,
                        dischargetime, dayspermonth)

    println("\nSaving...")

    @unpack gisregion = options
    mkpath(in_datafolder("output"))
    matopen(in_datafolder("output", "GISdata_hydro_$gisregion.mat"), "w") do file
        write(file, "existingcapac", existingcapac)
        write(file, "existinginflowcf", existinginflowcf)
        write(file, "potentialcapac", potentialcapac)
        write(file, "potentialinflowcf", potentialinflowcf)
        write(file, "potentialmeancost", potentialmeancost)
        write(file, "potentialmeandischargetime", potentialmeandischargetime)
    end
    nothing
end

function readhydrodatabases()
    println("\nReading hydro databases...")

    # lat,lon,COE,Production_GWh,Lake_surface_m2,Lake_volume_m3,
    #   Qm1,Qm2,Qm3,Qm4,Qm5,Qm6,Qm7,Qm8,Qm9,Qm10,Qm11,Qm12,Qm13,ContID,BasinID,SysID,CapCost
    # original headers:
    # lat,lon,COE ($/kWh),Production (GWh),Lake surface (m2),Lake volume (m3),
    #   Qm1,Qm2,Qm3,Qm4,Qm5,Qm6,Qm7,Qm8,Qm9,Qm10,Qm11,Qm12,Qm13,ContID,BasinID,
    #   SysID (1=DiversionalCanalPower/2=RiverPower), CapCost ($/kW)
    potential = DataFrame(CSV.File(in_datafolder("Hydro database (Gernaat) - potential.csv")))

    # GrandID,lat,lon,Production_kWh,m1,m2,m3,m4,m5,m6,m7,m8,m9,m10,m11,m12,m13
    # original headers:
    # GrandID,lat,lon,Production_GWh,m1,m2,m3,m4,m5,m6,m7,m8,m9,m10,m11,m12,m13
    existing = DataFrame(CSV.File(in_datafolder("Hydro database (Gernaat) - existing (GRanD).csv")))

    # country,country_long,name,gppd_idnr,capacity_mw,latitude,longitude,fuel1,fuel2,fuel3,fuel4,commissioning_year,owner,source,url,geolocation_source,year_of_capacity_data,generation_gwh_2013,generation_gwh_2014,generation_gwh_2015,generation_gwh_2016,estimated_generation_gwh
    elecplants = DataFrame(CSV.File(in_datafolder("WRI - Global Power Plant Database v1.10", "global_power_plant_database.csv")), copycols=true)

    # clean up wrong coordinates
    wrong = findall((elecplants.fuel1 .== "Hydro") .& (
            (elecplants.latitude.>90) .| (elecplants.latitude.<-90) .| (elecplants.longitude.>180) .| (elecplants.longitude.<-180)
        ))
    # elecplants.url[wrong]     # 8 results, check urls online
    # 1: DDMMSS in lat and lon
    # 2-8: lat & lon flipped
    w = wrong[1]
    elecplants.latitude[w], elecplants.longitude[w] = dms2deg(elecplants.latitude[w]), dms2deg(elecplants.longitude[w])
    w = wrong[2:8]
    elecplants.latitude[w], elecplants.longitude[w] = elecplants.longitude[w], elecplants.latitude[w]
    wrong = findall((elecplants.fuel1 .== "Hydro") .& (
            (elecplants.latitude.>90) .| (elecplants.latitude.<-90) .| (elecplants.longitude.>180) .| (elecplants.longitude.<-180)
        ))
    if !isempty(wrong)
        error("Cleanup of power plant database coordinates failed.")
    end

    # capacity_mw,latitude,longitude
    hydroplants = filter(r -> !ismissing(r.fuel1) && r.fuel1 == "Hydro", elecplants)[:,5:7]

    # Country, Capacity_MW, Pumped_MW, Other_MW, Generation_GWh, Generation_BP_GWh   % [0 = no data]
    # original headers:
    # Country, Total Hydropower Capacity (MW) in 2015, Pumped Storage Capacity (MW) in 2015, Excluding Pumped Storage (MW) in 2015, Estimated Net Hydropower Generation (GWh) in 2015, Consumption of Hydroelectricity (GWh) in 2015 (BP 2016)
    WECcapacity = CSV.File(in_datafolder("WEC hydro capacity 2015.csv")) |> DataFrame

    # Country, Undeveloped_GWh, Potential_GWh, Utilisation
    # original headers:
    # Country, Undeveloped (GWh/year), Total Potential (GWh/year), Current Utilisation (%)
    WECpotential = CSV.File(in_datafolder("WEC hydro potentials.csv")) |> DataFrame

    return potential, existing, hydroplants, WECcapacity, WECpotential
end

function potentialhydrovars(potential)
    dayspermonth = [Dates.daysinmonth(Date("2018-$m")) for m = 1:12]

    # No data on Nordic sites (hydro data limited to 60 degrees north), so take Nordic monthly inflow from a European model instead.
    nordicprofile = [2224, 1746, 2032, 4089, 15163, 23469, 16126, 9740, 9026, 9061, 5438, 3544]

    production = potential[:, :Production_GWh]              # GWh
    profile = Matrix(potential[:, 7:18])                    # m3/s  (columns Qm1-Qm12)
    eleccost = potential[:, :COE]                           # $/kWh
    capcost = potential[:, :CapCost]                        # $/kW
    type = potential[:, :SysID]
    reservoirsize = potential[:, :Lake_volume_m3] / 1e6     # Mm3

    cfriver = mean(profile, dims=2) ./ maximum(profile, dims=2)
    crf = CRF(0.1, 40)                                      # Gernaat assumes 10% discount rate and 40 year lifetime
    cf = capcost * crf/8760 ./ eleccost
    effic = 0.9*(type .== 1) .+ 0.7*(type .== 2)
    capacity = production ./ cf * 1000/8760                 # MW
    sortedflow = sort(profile, dims=2)                      # m3/s
    designflow = sortedflow[:, 9]                           # m3/s
    waterdensity = 997                                      # kg/m3
    grav = 9.81                                             # m/s2

    # My estimate of fallheight using eq 2-3 in the Gernaat paper.
    # Maybe ask for a new version of the database with this field included.
    fallheight = production*1e9/8760 ./ cf ./ effic ./ (designflow*waterdensity*grav)     # m

    # remove outliers
    f = findall(fallheight .> 2000)
    fallheight[f] .= 2000
    production[f] = fallheight[f] .* cf[f] .* effic[f] .* designflow[f] * waterdensity*grav*8760/1e9
    capacity[f] = production[f] ./ cf[f] * 1000/8760
    eleccost[f] = eleccost[f] .* potential[f, :Production_GWh] ./ production[f]
    capcost[f] = capcost[f] .* potential[f, :Production_GWh] ./ production[f]

    # J = kg m2/s2
    # m * m3/s * kg/m3 * m/s2 = J/s = W
    energyprofile = fallheight * dayspermonth' * 24/1e9 .* profile * waterdensity*grav    # GWh/month
    waterenergy = sum(energyprofile, dims=2)    # GWh/year

    # m * Mm3 * kg/m3 * m/s2 = MJ
    reservoirenergy = fallheight .* reservoirsize * waterdensity*grav / 3600 / 1000     # GWh
    dischargetime = reservoirenergy ./ capacity * 1000     # h
    # Matlab plots:  figure; hist(log10(reservoirenergy(reservoirenergy>0)),1000)
    # hist(log10(dischargetime(reservoirenergy>0)),1000)

    # Try calculating a monthly cf based on water inflow and turbine size.
    # Maybe not always accurate for large reservoirs since they can store between months,
    # but should work for run-of-river. Annual production and CF should still be fine.
    # (But why is my production estimate sometimes much higher than Gernaat's?)
    monthlyinflow = energyprofile .* effic
    monthlyprod = min.(monthlyinflow, capacity * dayspermonth' * 24/1000)   # GWh/month

    return dayspermonth, nordicprofile, eleccost, capacity, reservoirenergy, dischargetime, monthlyinflow
end

function read_countries_and_regions(options)
    println("Reading Global GADM country-level database...")
    countries, _, countrynames, _, _ = loadregions("Global_GADM0")
    numcountries = length(countrynames)

    @unpack gisregion = options
    println("Reading regions for $gisregion...")
    regions, offshoreregions, regionlist, lonrange, latrange = loadregions(gisregion)
    # bbox = [-90 -180; 90 180]

    return countries, countrynames, regions, regionlist, lonrange, latrange
end

function existinginflow(existing, countries, countrynames, regions, regionlist, lonrange, latrange, hydroplants, WECcapacity,
                        dayspermonth)
    println("\nCalculate monthly inflow for existing hydro...");

    existingprofile = Matrix(existing[:, 5:16])             # m3/s  (columns m1-m12)

    numcountries = length(countrynames)
    capacwri = zeros(numcountries,1)        # GW
    annualcf = zeros(numcountries,1)
    scalefactor_capacwri = zeros(numcountries,1)

    for i = 1:size(hydroplants,1)
        country = getregion(hydroplants[i,:longitude], hydroplants[i,:latitude], countries)
        i == 4867 && continue   # weird data point in the Indian Ocean
        if country > 0 && country != NOREGION
            capacwri[country] += hydroplants[i,:capacity_mw] / 1e3    # GW
        else
            println("Can't identify country: $(hydroplants[i,:])")
        end
    end

    # calculate scale factor for national hydro capacity (WRI capac/WEC capac)
    # calculate annual cf for each country from WEC data
    for i = 1:size(WECcapacity,1)-1     # last row is "World"
        cty = findfirst(countrynames .== Symbol(WECcapacity[i,:Country]))
        cty == nothing && error("Can't find country name: $(Symbol(WECcapacity[i,:Country]))")
        scalefactor_capacwri[cty] = capacwri[cty] / WECcapacity[i,:Capacity_MW] * 1000
        annualcf[cty] = WECcapacity[i,:Generation_GWh] / WECcapacity[i,:Capacity_MW] * 1000/8760
    end

    numreg = length(regionlist)
    existingcapac = zeros(numreg,1)         # GW
    existinginflow = zeros(numreg,12)       # GWh/month

    # calculate annual generation at each site using capacity and scale factor & CF from above
    # then get monthly inflow profile from nearest site in Gernaat/Grand
    # finally calculate monthly inflow for existing sites
    for i = 1:size(hydroplants,1)
        cty = getregion(hydroplants[i,:longitude], hydroplants[i,:latitude], countries)
        reg = getregion(hydroplants[i,:longitude], hydroplants[i,:latitude], regions, lonrange, latrange)
        (reg == 0 || reg == NOREGION || cty == 0 || cty == NOREGION) && continue
        scalefactor = scalefactor_capacwri[cty] > 0 ? scalefactor_capacwri[cty] : 1.0
        existingcapacity = hydroplants[i,:capacity_mw] / 1e3 / scalefactor      # GW
        annualgeneration = existingcapacity * annualcf[cty] * 8760              # GWh/year
        existingcapac[reg] += existingcapacity

        # approx distance in km to all existing plants in Gernaat/Grand
        dist = 111*sqrt.((existing[:,:lon] .- hydroplants[i,:longitude]).^2 + (existing[:,:lat] .- hydroplants[i,:latitude]).^2)
        d, ndx = findmin(dist)
        if regionlist[reg] == "NOR"
            inflowprofile = nordicprofile
        else
            inflowprofile = existingprofile[ndx,:]
        end
        monthlyprofile = inflowprofile .* dayspermonth
        monthlyprofile = monthlyprofile / sum(monthlyprofile)

        monthlygeneration = annualgeneration * monthlyprofile       # GWh/month
        existinginflow[reg,:] += monthlygeneration
    end

    existinginflowcf = existinginflow ./ existingcapac ./ dayspermonth' / 24

    return existingcapac, existinginflowcf
end

function potentialinflow(options, potential, capacity, regions, regionlist, lonrange, latrange, monthlyinflow, eleccost,
            dischargetime, dayspermonth)
    println("Calculate monthly inflow for potential hydro...")

    @unpack costclasses_min, costclasses_max, storageclasses_min, storageclasses_max, gisregion = options

    ncostclasses = length(costclasses_min)
    nstorageclasses = length(storageclasses_min)

    numreg = length(regionlist)
    potentialcapac = zeros(numreg,ncostclasses,nstorageclasses)        # GW
    potentialinflow = zeros(numreg,ncostclasses,nstorageclasses,12)    # GWh/month

    potentialmeancost = zeros(numreg,ncostclasses,nstorageclasses)             # $/kWh
    potentialmeandischargetime = zeros(numreg,ncostclasses,nstorageclasses)    # hours
    nobservations = zeros(numreg,ncostclasses,nstorageclasses)

    for i = 1:size(potential,1)
        reg = getregion(potential[i,:lon], potential[i,:lat], regions, lonrange, latrange)
        (reg == 0 || reg == NOREGION) && continue

        weeks = dischargetime[i] / 168
        storageclass = findfirst((weeks .>= storageclasses_min) .& (weeks .< storageclasses_max))
        cost = eleccost[i]*1000     # US $/MWh
        costclass = findfirst((cost .>= costclasses_min) .& (cost .< costclasses_max))

        if regionlist[reg] == "NOR"
            inflowprofile = nordicprofile/sum(nordicprofile) * sum(monthlyinflow[i,:])
        else
            inflowprofile = monthlyinflow[i,:]
        end

        potentialcapac[reg,costclass,storageclass] += capacity[i]/1000
        potentialinflow[reg,costclass,storageclass,:] += inflowprofile

        potentialmeancost[reg,costclass,storageclass] += eleccost[i]
        potentialmeandischargetime[reg,costclass,storageclass] += min(8760, dischargetime[i])
        nobservations[reg,costclass,storageclass] += 1
    end

    potentialinflowcf = potentialinflow ./ potentialcapac ./ reshape(dayspermonth, (1,1,1,12)) / 24
    potentialmeancost = potentialmeancost ./ nobservations
    potentialmeandischargetime = potentialmeandischargetime ./ nobservations

    return potentialcapac, potentialinflowcf, potentialmeancost, potentialmeandischargetime
end

# function mapurl(latlon1,latlon2)
#     url = "https://www.google.com/maps/dir/?api=1&origin=$(latlon1[1]),$(latlon1[2])" *
#             "&destination=$(latlon2[1]),$(latlon2[2])&travelmode=walking"
# end


"""
    module Interfaces

This module contains interfaces that can be used to extend and modify the algorithms and behaviours of `Rimu`.

# Interfaces
Follow the links for the definitions of the interfaces!
* [`AbstractHamiltonian`](@ref) for defining [`Hamiltonians`](@ref)
* [`AbstractDVec`](@ref) for defining data structures for `Rimu` as in [`DictVectors`](@ref)
* [`StochasticStyle`](@ref) for controlling the stochastic algorithms used by [`lomc!`](@ref)
  as implemented in [`StochasticStyles`](@ref)

# Additonal exports

## Interface functions for[`AbstractHamiltonian`](@ref)s:
* [`diagonal_element`](@ref)
* [`num_offdiagonals`](@ref)
* [`get_offdiagonal`](@ref)
* [`offdiagonals`](@ref).
* [`random_offdiagonal`](@ref)
* [`starting_address`](@ref)
* [`LOStructure`](@ref)

## working with  [`AbstractDVec`](@ref)s and [`StochasticStyle`](@ref)
* [`deposit!`](@ref)
* [`default_style`](@ref)
* [`CompressionStrategy`](@ref)

## Functions Rimu.jl uses to do FCIQMC:

* [`fciqmc_col!`](@ref)
* [`step_stats`](@ref)
* [`update_dvec!`](@ref)
"""
module Interfaces

using LinearAlgebra
using StaticArrays

using ..ConsistentRNG
import OrderedCollections: freeze

export
    StochasticStyle, default_style, StyleUnknown, fciqmc_col!, step_stats, update_dvec!,
    CompressionStrategy, NoCompression, compress!
export
    AbstractDVec, deposit!, storage, localpart, freeze
export
    AbstractHamiltonian, diagonal_element, num_offdiagonals, get_offdiagonal, offdiagonals,
    random_offdiagonal, starting_address,
    LOStructure, IsHermitian, AdjointKnown, AdjointUnknown, has_adjoint

include("stochasticstyles.jl")
include("dictvectors.jl")
include("hamiltonians.jl")

end


__precompile__(true)
module Nothings

using Compat

export allownothing, disallownothing, isnothing, nothings,
       NothingException, levels, coalesce

if VERSION >= v"0.7.0-DEV.2762"
    using Base: Nothing, nothing
else
    export Nothing, nothing

    """
        Nothing

    A type with no fields whose singleton instance [`nothing`](@ref) is used
    to represent nothing values.
    """
    struct Nothing end

    """
        nothing

    The singleton instance of type [`Nothing`](@ref) representing a nothing value.
    """
    const nothing = Nothing()

    Base.show(io::IO, x::Nothing) = print(io, "nothing")

end

"""
 NothingException(msg)

Exception thrown when a [`nothing`](@ref) value is encountered in a situation
where it is not supported. The error message, in the `msg` field
may provide more specific details.
"""
struct NothingException <: Exception
    msg::AbstractString
end

import Base: ==, !=, <, *, <=, !, +, -, ^, /, &, |, xor

Base.showerror(io::IO, ex::NothingException) =
    print(io, "NothingException: ", ex.msg)

isnothing(::Any) = false
isnothing(::Nothing) = true

Base.promote_rule(::Type{T}, ::Type{Nothing}) where {T} = Union{T, Nothing}
Base.promote_rule(::Type{T}, ::Type{Union{S,Nothing}}) where {T,S} = Union{promote_type(T, S), Nothing}
Base.promote_rule(::Type{T}, ::Type{Any}) where {T} = Any
Base.promote_rule(::Type{Any}, ::Type{Nothing}) = Any
Base.promote_rule(::Type{Nothing}, ::Type{Any}) = Any
Base.promote_rule(::Type{Nothing}, ::Type{Nothing}) = Nothing

Base.convert(::Type{Union{T, Nothing}}, x) where {T} = convert(T, x)

# Comparison operators
==(::Nothing, ::Nothing) = nothing
==(::Nothing, b) = nothing
==(a, ::Nothing) = nothing
# != must be defined explicitly since fallback expects a Bool
!=(::Nothing, ::Nothing) = nothing
!=(::Nothing, b) = nothing
!=(a, ::Nothing) = nothing
Base.isequal(::Nothing, ::Nothing) = true
Base.isequal(::Nothing, b) = false
Base.isequal(a, ::Nothing) = false
<(::Nothing, ::Nothing) = nothing
<(::Nothing, b) = nothing
<(a, ::Nothing) = nothing
Base.isless(::Nothing, ::Nothing) = false
Base.isless(::Nothing, b) = false
Base.isless(a, ::Nothing) = true
if VERSION < v"0.7.0-DEV.300"
    <=(::Nothing, ::Nothing) = nothing
    <=(::Nothing, b) = nothing
    <=(a, ::Nothing) = nothing
end

# Unary operators/functions
for f in (:(!), :(+), :(-), :(Base.identity), :(Base.zero), :(Base.one), :(Base.oneunit),
        :(Base.abs), :(Base.abs2), :(Base.sign),
        :(Base.acos), :(Base.acosh), :(Base.asin), :(Base.asinh), :(Base.atan), :(Base.atanh),
        :(Base.sin), :(Base.sinh), :(Base.cos), :(Base.cosh), :(Base.tan), :(Base.tanh),
        :(Base.exp), :(Base.exp2), :(Base.expm1), :(Base.log), :(Base.log10), :(Base.log1p),
        :(Base.log2), :(Base.exponent), :(Base.sqrt), :(Base.gamma), :(Base.lgamma),
        :(Base.iseven), :(Base.ispow2), :(Base.isfinite), :(Base.isinf), :(Base.isodd),
        :(Base.isinteger), :(Base.isreal), :(Base.isimag), :(Base.isnan), :(Base.isempty),
        :(Base.iszero), :(Base.transpose), :(Base.float))
    @eval $(f)(d::Nothing) = nothing
end

for f in (:(Base.zero), :(Base.one), :(Base.oneunit))
    @eval function $(f)(::Type{Union{T, Nothing}}) where T
        T === Any && throw(MethodError($f, (Any,)))
        $f(T)
    end
end

# Binary operators/functions

for f in (:(+), :(-), :(*), :(/), :(^),
        :(Base.div), :(Base.mod), :(Base.fld), :(Base.rem), :(Base.min), :(Base.max))
    @eval ($f)(::Nothing, ::Nothing) = nothing
end

for f in (:(+), :(*), :(Base.min), :(Base.max))
    @eval begin
        ($f)(a::Number, o::Nothing) = a
        ($f)(o::Nothing, a::Number) = a
    end
end

Base.:(-)(a::Number, o::Nothing) =  a
Base.:(-)(o::Nothing, a::Number) = -a
Base.:(^)(a::Number, o::Nothing) =  a
Base.:(^)(o::Nothing, a::Number) = nothing

for f in (:(/), :(Base.div), :(Base.fld), :(Base.cld))
    @eval begin
        ($f)(a::Number, o::Nothing) = a
        ($f)(o::Nothing, a::Number) = inv(a)
    end
end

for f in (:(Base.mod), :(Base.rem))
    @eval begin
        ($f)(a::Number, o::Nothing) = a
        ($f)(o::Nothing, a::Number) = nothing
    end
end

# to avoid ambiguity warnings
(^)(::Nothing, ::Integer) = nothing

# Bit operators
(&)(::Nothing, ::Nothing) = nothing
(&)(a::Nothing, b::Bool) = b
(&)(b::Bool, a::Nothing) = b
(&)(::Nothing, b::Integer) = b
(&)(b::Integer, ::Nothing) = b
(|)(::Nothing, ::Nothing) = nothing
(|)(a::Nothing, b::Bool) = b
(|)(b::Bool, a::Nothing) = b
(|)(::Nothing, b::Integer) = b
(|)(b::Integer, ::Nothing) = b
xor(::Nothing, ::Nothing) = nothing
xor(a::Nothing, b::Bool) = b
xor(b::Bool, a::Nothing) = b
xor(a::Nothing, b::Integer) = b
xor(b::Integer, a::Nothing) = b

# String functions
*(d::Nothing, x::AbstractString) = x
*(d::AbstractString, x::Nothing) = x

end # module Nothings


push!(LOAD_PATH, "/Users/megorov/Desktop/projects/stanford/sniper/egorov/sniperRepo/new-version")

using UrbanMap_
using SniperMDP_
using ValueIterations
using Simulator
using SimServers

# create a map and the MDP
map = UrbanMap(1, [0.5], [0.5], [0.42], [0.22], 10, 10)
mdp = SniperMDP(map, nSnipers = 1, nMonitors = 2, targets=[(2.0,8.0), (2.0,1.0)])

# load the policy
policy = load_policy("policies/mmdp-10-iter-1.policy")

# set up the server
socket = 4444
server = SniperServer(socket)

# create the simulation type
init = 470872
sim  = MDPSimulation(mdp, init, 100)

# begin simulation and send results to client
#simAndSend(server, sim, mdp, policy)

# begin batch sim
n = 500
rewards = zeros(n)
shot    = zeros(Int64, n)
obs     = zeros(Int64, n)
for i = 1:n
    res = simulate(policy, mdp, 100, initState=init)
    rewards[i] = sum(res.rewards)
    shot[i]    = res.nKills
    obs[i]     = res.nObs
end


######### sensis EUR du 11 oct 2017
using CSV, DataFrames
using FileIO
using Optim
using VegaLite
using Rotolo
using Base.Markdown

global_theme = """
	text-align: justify;
	border-style:solid;
	padding: 5px
  """

methods(Media.render, (Rotolo.Juno, Any))

methods(Media.render, (Atom.Inline, Void))

methods(Media.render, (Atom.Editor, Any))

methods(Media.render, (Atom.Inline, Any))

@session MontantsEquiv6 global_theme
# @endsession

@redirect Markdown.MD
sleep(3)

@redirect Katex

Rotolo.currentSession.redirected_types

@container header style=>"text-align: center;"

@redirect Base.Markdown.MD

md"# Montants équivalents"

typeof(md"# Montants équivalents")

methods(Juno.render, (Juno.Editor, Base.Markdown.MD))
methods(Atom.render, (Atom.Editor, Base.Markdown.MD))

@style md"*From Wikipedia, the free encyclopedia*" style=>"font-size:x-small"

Katex("f(x,y)=(x^2+y-11)^2+(x+y^2-7)^2")

5+6

@redirect Int

5+65

# delete!(Rotolo.currentSession.redirected_types, Int)

"abcd"

@redirect String

"abcd"

Rotolo.loadmsg(Katex)
Katex("f(x,y)=(x^2+y-11)^2+(x+y^2-7)^2")

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

fvar(s) = s' * covmat * s
dfvar(s) = covmat * s + (s' * covmat)'

function minfunc(initial::Vector{Float64}, vidx::Int64, h::Float64)
  nv = copy(initial)
  nv[vidx] += h
  fvar(nv)
end

function stephedge1(minfunc::Function, initial::Vector, mask::BitArray)
  ovar = round(Int64, minfunc(initial, 1, 0.))

  minzer = zeros(length(initial))
  minmum = [ 1e20 for i in 1:length(initial) ]
  for i in collect(1:length(initial))[mask]
    res = optimize(x -> minfunc(initial, i, x), -1e6, 1e6 )
    minzer[i] = Optim.minimizer(res)
    minmum[i] = Optim.minimum(res)
  end
  vidx = indmin(minmum)
  amnt = round(Int64, minzer[vidx])
  nvar = round(Int64, minmum[vidx])
  println("old var = $ovar, mat $(sensimat[vidx]) ($vidx), amount = $amnt, new var. = $nvar")

  newvec = copy(initial)
  newvec[vidx] += minzer[vidx]
  newmask = copy(mask)
  newmask[vidx] = false
  vidx, newvec, newmask
end

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

sensis = CSV.read(IOBuffer("""
  Mats	EUR.EIB.6M	USD.LIB.3M
  3M	-1202	4589
  6M	-8415	4684
  1Y	-1626	-13340
  18M	-15253	17372
  2Y	-9917	15681
  3Y	5901	-11285
  4Y	3657	-15135
  5Y	65133	24432
  6Y	-30523	-16872
  7Y	-41440	-12484
  8Y	49300	-1001
  9Y	-5135	-1847
  10Y	-21788	6180
  11Y	11166	-9559
  12Y	17216	-4027
  13Y	-2389	1057
  14Y	-11172	2661
  15Y	2677	1759
  16Y	-6815	2324
  17Y	-14232	5430
  18Y	4067	5638
  19Y	23327	645
  20Y	-18821	-118
  21Y	-21295	0
  22Y	-6	0
  23Y	-16188	0
  24Y	-2756	0
  25Y	50119	0
  26Y	4388	0
  27Y	-364	0
  28Y	-19	0
  """), delim='\t')
sensimat = Vector(String.(sensis[:,1]))

###########  EUR  #########################################

dat = load("c:/temp/histovarEUR.jld")
matdelta = dat["matvar"]
matmatu  = dat["maturities"]

midx = indexin(sensimat, String.(matmatu))
sum(midx .== 0)  # =0, ok on a toutes les mat qu'il faut

covmat = cov(matdelta[:,midx])
extrema(covmat) # 1.86e-9 , 3.61e-7

sensivec = Vector{Float64}(sensis[Symbol("EUR.EIB.6M")])
initial, mask = copy(sensivec), trues(length(sensivec))
_, initial, mask = stephedge1(minfunc, initial, mask)
# old var = 50, mat 5Y (8), amount = -31042, new var. = 12
# old var = 12, mat 18M (4), amount = 20089, new var. = 6
# old var = 6, mat 6M (2), amount = 7675, new var. = 5
# old var = 5, mat 4Y (7), amount = -5172, new var. = 4
# old var = 4, mat 2Y (5), amount = 3887, new var. = 4
# old var = 4, mat 28Y (31), amount = -3262, new var. = 4
# old var = 4, mat 7Y (10), amount = 1876, new var. = 4
# old var = 4, mat 3Y (6), amount = -2555, new var. = 3
# old var = 3, mat 17Y (20), amount = 1735, new var. = 3
# old var = 3, mat 27Y (30), amount = -1825, new var. = 3
# old var = 3, mat 16Y (19), amount = 1494, new var. = 3
# old var = 3, mat 25Y (28), amount = -1546, new var. = 3
# old var = 3, mat 8Y (11), amount = 1103, new var. = 3

std₀ = sqrt(50) * 10000 # 71k€
sqrt(12) * 10000        # 35k€
sqrt(6) * 10000         # 25k€

-31042 * 10000 / 5 / 0.98 # 5Y  prêt de 63 M€
20089 * 10000 / 1.5       # 18M emprunt de 134 M€
7675 * 10000 / 0.5       # 6M emprunt de 154 M€

using VegaLite

pd = DataFrame(mats = sensimat, std= sqrt.(diag(covmat)) * 10000)
pd |> markbar() |>
  encoding(xordinal(field=:mats, vlscale(domain=sensimat)),
  yquantitative(field=:std))

###########  USD  #########################################

dat = load("c:/temp/histovarUSD.jld")
matdelta = dat["matvar"]
matmatu  = dat["maturities"]

cidx = isnull.(sensis[3]) .== false
sensimat = Vector{String}(sensis[cidx,1])

midx = indexin(sensimat, String.(matmatu))
sum(midx .== 0)  # =0, ok on a toutes les mat qu'il faut

covmat = cov(matdelta[:,midx])
extrema(covmat) # 1.86e-9 , 3.61e-7

pd = DataFrame(mats = sensimat, std= sqrt.(diag(covmat)) * 10000)
pd |> markbar() |>
  encoding(xordinal(field=:mats, vlscale(domain=sensimat)),
  yquantitative(field=:std))

sensivec = Vector{Float64}(sensis[cidx, Symbol("USD.LIB.3M")])
initial, mask = copy(sensivec), trues(length(sensivec))
_, initial, mask = stephedge1(minfunc, initial, mask)

# old var = 32, mat 18M (4), amount = -13699, new var. = 25
# old var = 25, mat 7Y (10), amount = 11879, new var. = 16
# old var = 16, mat 2Y (5), amount = -6602, new var. = 14
# old var = 14, mat 1Y (3), amount = 6651, new var. = 12
# old var = 12, mat 3M (1), amount = -5328, new var. = 9
# old var = 9, mat 3Y (6), amount = 4410, new var. = 8
# old var = 8, mat 6M (2), amount = -3390, new var. = 7
# old var = 7, mat 20Y (23), amount = -3719, new var. = 6
# old var = 6, mat 8Y (11), amount = 3072, new var. = 5
# old var = 5, mat 21Y (24), amount = -2938, new var. = 5
# old var = 5, mat 4Y (7), amount = 2501, new var. = 5

std₀ = sqrt(32) * 10000 # 57k€
sqrt(25) * 10000        # 50k€
sqrt(16) * 10000         # 40k€
sqrt(14) * 10000         # 37k€

-13699 * 10000 / 1.5  # 18M  prêt de 91M€
11879 * 10000 / 7     # 7Y   emprunt de 17 M€
-6602 * 10000 / 2     # 2Y   prêt de 33 M€

#############  find simplified equivalment pose   ############################

# subset of sensis to minimize
sensimat = Vector{String}(sensis[1])
sensimat2 = String["3M", "6M", "1Y", "3Y", "7Y", "12Y"]
mmask = indexin(sensimat2, sensimat)
sensi₀ = Vector{Float64}(sensis[3])

dat = load("c:/temp/histovarUSD.jld")
matdelta = dat["matvar"]
matmatu  = dat["maturities"]
midx = indexin(sensimat, String.(matmatu))
sum(midx .== 0)  # =0, ok on a toutes les mat qu'il faut

covmat = cov(matdelta[:,midx])
extrema(covmat)

sensi₁ = copy(sensi₀)
function minfunc(x::Vector{Float64})
  sensi₁[mmask] = x
  sensi₁' * covmat * sensi₁
end

result = optimize(minfunc, zeros(length(mmask)), BFGS())
Optim.minimizer(result)
Optim.minimum(result)

sensi₁ - sensi₀

sqrt(sensi₀' * covmat * sensi₀) * 10000 # 56 k€
sqrt(sensi₁' * covmat * sensi₁) * 10000 # 37 k€

pdtfr = vcat(DataFrame(mat=sensimat,  sensi=sensi₀, src = "orig"),
             DataFrame(mat=sensimat2, sensi=(sensi₀-sensi₁)[mmask], src="simplified"))

pdtfr |>
  markbar() |>
  encoding(columnordinal(field=:mat),
           xordinal(field=:src),
           colornominal(field=:src),
           yquantitative(field=:sensi))


export DDPGPolicy

mutable struct DDPGPolicy{
    BA<:NeuralNetworkApproximator,
    BC<:NeuralNetworkApproximator,
    TA<:NeuralNetworkApproximator,
    TC<:NeuralNetworkApproximator,
    P,
    R<:AbstractRNG,
} <: AbstractPolicy

    behavior_actor::BA
    behavior_critic::BC
    target_actor::TA
    target_critic::TC
    γ::Float32
    ρ::Float32
    batch_size::Int
    start_steps::Int
    start_policy::P
    update_after::Int
    update_every::Int
    act_limit::Float64
    act_noise::Float64
    step::Int
    rng::R
    # for logging
    actor_loss::Float32
    critic_loss::Float32
end

Flux.functor(x::DDPGPolicy) = (
    ba = x.behavior_actor,
    bc = x.behavior_critic,
    ta = x.target_actor,
    tc = x.target_critic,
),
y -> begin
    x = @set x.behavior_actor = y.ba
    x = @set x.behavior_critic = y.bc
    x = @set x.target_actor = y.ta
    x = @set x.target_critic = y.tc
    x
end

"""
    DDPGPolicy(;kwargs...)

# Keyword arguments

- `behavior_actor`,
- `behavior_critic`,
- `target_actor`,
- `target_critic`,
- `start_policy`,
- `γ = 0.99f0`,
- `ρ = 0.995f0`,
- `batch_size = 32`,
- `start_steps = 10000`,
- `update_after = 1000`,
- `update_every = 50`,
- `act_limit = 1.0`,
- `act_noise = 0.1`,
- `step = 0`,
- `rng = Random.GLOBAL_RNG`,
"""
function DDPGPolicy(;
    behavior_actor,
    behavior_critic,
    target_actor,
    target_critic,
    start_policy,
    γ = 0.99f0,
    ρ = 0.995f0,
    batch_size = 32,
    start_steps = 10000,
    update_after = 1000,
    update_every = 50,
    act_limit = 1.0,
    act_noise = 0.1,
    step = 0,
    rng = Random.GLOBAL_RNG,
)
    copyto!(behavior_actor, target_actor)  # force sync
    copyto!(behavior_critic, target_critic)  # force sync
    DDPGPolicy(
        behavior_actor,
        behavior_critic,
        target_actor,
        target_critic,
        γ,
        ρ,
        batch_size,
        start_steps,
        start_policy,
        update_after,
        update_every,
        act_limit,
        act_noise,
        step,
        rng,
        0.0f0,
        0.0f0,
    )
end

# TODO: handle Training/Testing mode
function (p::DDPGPolicy)(env)
    p.step += 1

    if p.step <= p.start_steps
        p.start_policy(env)
    else
        D = device(p.behavior_actor)
        s = state(env)
        s = Flux.unsqueeze(s, ndims(s) + 1)
        action = p.behavior_actor(send_to_device(D, s)) |> vec |> send_to_host
        clamp(action[] + randn(p.rng) * p.act_noise, -p.act_limit, p.act_limit)
    end
end

function RLBase.update!(p::DDPGPolicy, traj::CircularArraySARTTrajectory)
    length(traj) > p.update_after || return
    p.step % p.update_every == 0 || return
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(p.batch_size))
    update!(p, batch)
end

function RLBase.update!(p::DDPGPolicy, batch::NamedTuple{SARTS})
    s, a, r, t, s′ = send_to_device(device(p), batch)

    A = p.behavior_actor
    C = p.behavior_critic
    Aₜ = p.target_actor
    Cₜ = p.target_critic

    γ = p.γ
    ρ = p.ρ

    # !!! we have several assumptions here, need revisit when we have more complex environments
    # state is vector
    # action is scalar
    a′ = Aₜ(s′)
    qₜ = Cₜ(vcat(s′, a′)) |> vec
    y = r .+ γ .* (1 .- t) .* qₜ
    a = Flux.unsqueeze(a, 1)

    gs1 = gradient(Flux.params(C)) do
        q = C(vcat(s, a)) |> vec
        loss = mean((y .- q) .^ 2)
        ignore() do
            p.critic_loss = loss
        end
        loss
    end

    update!(C, gs1)

    gs2 = gradient(Flux.params(A)) do
        loss = -mean(C(vcat(s, A(s))))
        ignore() do
            p.actor_loss = loss
        end
        loss
    end

    update!(A, gs2)

    # polyak averaging
    for (dest, src) in zip(Flux.params([Aₜ, Cₜ]), Flux.params([A, C]))
        dest .= ρ .* dest .+ (1 - ρ) .* src
    end
end


# MOI backend AbstractModel
MOIU.@model(ParametronMOIModel, # modelname
    (MOI.ZeroOne, MOI.Integer), # scalarsets
    (MOI.LessThan, MOI.GreaterThan, MOI.EqualTo), # typedscalarsets
    (MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives), # vectorsets
    (), # typedvectorsets
    (MOI.SingleVariable,), # scalarfunctions
    (MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction), # typedscalarfunctions
    (), # vectorfunctions
    (MOI.VectorAffineFunction,) # typedvectorfunctions
)

# Variable conversion
Functions.Variable(v::MOI.VariableIndex) = Variable(v.value)
MOI.VariableIndex(v::Variable) = MOI.VariableIndex(v.index)

# Sense conversion
function MOI.OptimizationSense(sense::Sense)
    if sense == Minimize
        MOI.MIN_SENSE
    elseif sense == Maximize
        MOI.MAX_SENSE
    else
        error()
    end
end

# update!
struct IdentityVarMap end
Base.getindex(::IdentityVarMap, i) = MOI.VariableIndex(i)

function update!(moi_f::MOI.ScalarAffineFunction, f::AffineFunction, varmap = IdentityVarMap())
    moi_f.constant = f.constant[]
    resize!(moi_f.terms, length(f.linear))
    @inbounds for i in eachindex(f.linear)
        term = f.linear[i]
        moi_f.terms[i] = MOI.ScalarAffineTerm(term.coeff, varmap[term.var.index])
    end
    moi_f
end

function update!(moi_f::MOI.ScalarQuadraticFunction, f::QuadraticFunction, varmap = IdentityVarMap())
    affine = f.affine
    moi_f.constant = affine.constant[]
    resize!(moi_f.affine_terms, length(affine.linear))
    @inbounds for i in eachindex(affine.linear)
        term = affine.linear[i]
        moi_f.affine_terms[i] = MOI.ScalarAffineTerm(term.coeff, varmap[term.var.index])
    end
    resize!(moi_f.quadratic_terms, length(f.quadratic))
    @inbounds for i in eachindex(f.quadratic)
        term = f.quadratic[i]
        var_index_1 = varmap[term.rowvar.index]
        var_index_2 = varmap[term.colvar.index]
        coeff = ifelse(term.rowvar == term.colvar, 2 * term.coeff, term.coeff)
        moi_f.quadratic_terms[i] = MOI.ScalarQuadraticTerm(coeff, var_index_1, var_index_2)
    end
    moi_f
end

function update!(moi_f::MOI.VectorAffineFunction, fs::Vector{<:AffineFunction}, varmap = IdentityVarMap())
    resize!(moi_f.constants, length(fs))
    num_linear_terms = 0
    for f in fs
        num_linear_terms += length(f.linear)
    end
    resize!(moi_f.terms, num_linear_terms)
    i = 1
    @inbounds for row in eachindex(fs)
        f = fs[row]
        for term in f.linear
            moi_f.terms[i] = MOI.VectorAffineTerm(Int64(row), MOI.ScalarAffineTerm(term.coeff, varmap[term.var.index]))
            i += 1
        end
        moi_f.constants[row] = f.constant[]
    end
    moi_f
end

update!(moi_f::MOI.AbstractFunction, f::Nothing, varmap = IdentityVarMap()) = moi_f

# make_moi_equivalent
make_moi_equivalent(::Type{AffineFunction{T}}) where {T} =
    MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{T}[], zero(T))
make_moi_equivalent(::Type{QuadraticFunction{T}}) where {T} =
    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{T}[], MOI.ScalarQuadraticTerm{T}[], zero(T))
make_moi_equivalent(::Type{<:AbstractVector{AffineFunction{T}}}) where {T} =
    MOI.VectorAffineFunction(MOI.VectorAffineTerm{T}[], T[])

# canonical_function_type
canonical_function_type(::Type{Variable}, ::Type{T}) where {T} = AffineFunction{T}
canonical_function_type(::Type{<:LinearTerm}, ::Type{T}) where {T} = AffineFunction{T}
canonical_function_type(::Type{<:AffineFunction}, ::Type{T}) where {T} = AffineFunction{T}
canonical_function_type(::Type{<:AbstractVector{<:AffineFunction}}, ::Type{T}) where {T} = Vector{AffineFunction{T}}
canonical_function_type(::Type{<:QuadraticTerm}, ::Type{T}) where {T} = QuadraticFunction{T}
canonical_function_type(::Type{<:QuadraticFunction}, ::Type{T}) where {T} = QuadraticFunction{T}

# moi_to_native_type
moi_to_native_type(::Type{MOI.ScalarAffineFunction}) = AffineFunction{T} where T
moi_to_native_type(::Type{MOI.ScalarQuadraticFunction}) = QuadraticFunction{T} where T
moi_to_native_type(::Type{MOI.VectorAffineFunction}) = Vector{AffineFunction{T}} where T
moi_to_native_type(::Type{MOI.VectorQuadraticFunction}) = Vector{QuadraticFunction{T}} where T
moi_to_native_type(::Type{MOI.SingleVariable}) = Nothing

# Objective
struct Objective{E, F}
    expr::WrappedExpression{E}
    f::F
    isconstant::Bool
end

function Objective(::Type{T}, expr) where T
    val = evalarg(expr)
    E = canonical_function_type(typeof(val), T)
    converted = @expression convert(E, expr)
    isconstant = converted isa E # i.e., it's just a value; not a LazyExpression or a Parameter
    wrapped = convert(WrappedExpression{E}, converted)
    f = make_moi_equivalent(E)
    F = typeof(f)
    update!(f, wrapped())
    Objective{E, F}(wrapped, f, isconstant)
end

function update!(objective::Objective{E, F}, optimizer::MOI.AbstractOptimizer, varmap) where {E, F}
    if !objective.isconstant
        update!(objective.f, objective.expr(), varmap)
        MOI.set(optimizer, MOI.ObjectiveFunction{F}(), objective.f)
    end
    nothing
end

# Constraint
mutable struct Constraint{E, F<:MOI.AbstractFunction, S<:MOI.AbstractSet}
    expr::WrappedExpression{E}
    f::F
    set::S
    isconstant::Bool
    modelindex::MOI.ConstraintIndex{F, S}
    optimizerindex::MOI.ConstraintIndex{F, S}

    function Constraint(::Type{T}, expr, set::S) where {T, S<:MOI.AbstractSet}
        val = evalarg(expr)
        E = canonical_function_type(typeof(val), T)
        converted = @expression convert(E, expr)
        isconstant = converted isa E # i.e., it's just a value; not a LazyExpression or a Parameter
        wrapped = convert(WrappedExpression{E}, converted)
        f = make_moi_equivalent(E)
        F = typeof(f)
        update!(f, wrapped())
        new{E, F, S}(wrapped, f, set, isconstant)
    end

    function Constraint(f::F, set::S) where {F<:MOI.AbstractFunction, S<:MOI.AbstractSet}
        E = Nothing
        wrapped = convert(WrappedExpression{E}, @expression nothing)
        new{E, F, S}(wrapped, f, set, true)
    end
end

function update!(constraint::Constraint, optimizer::MOI.AbstractOptimizer, varmap)
    if !constraint.isconstant
        update!(constraint.f, constraint.expr(), varmap)
        MOI.set(optimizer, MOI.ConstraintFunction(), constraint.optimizerindex, constraint.f)
    end
    nothing
end
update!(constraint::Constraint{Nothing}, optimizer::MOI.AbstractOptimizer, varmap) = nothing

# Constraints
let
    constraint_specs = [(MOI.ScalarAffineFunction{T} where T, MOI.GreaterThan{T} where T),
                        (MOI.ScalarAffineFunction{T} where T, MOI.LessThan{T} where T),
                        (MOI.ScalarAffineFunction{T} where T, MOI.EqualTo{T} where T),
                        (MOI.VectorAffineFunction{T} where T, MOI.Nonnegatives),
                        (MOI.VectorAffineFunction{T} where T, MOI.Nonpositives),
                        (MOI.VectorAffineFunction{T} where T, MOI.Zeros),
                        (MOI.ScalarQuadraticFunction{T} where T, MOI.GreaterThan{T} where T),
                        (MOI.ScalarQuadraticFunction{T} where T, MOI.LessThan{T} where T),
                        (MOI.ScalarQuadraticFunction{T} where T, MOI.EqualTo{T} where T),
                        (MOI.VectorQuadraticFunction{T} where T, MOI.Nonnegatives),
                        (MOI.VectorQuadraticFunction{T} where T, MOI.Nonpositives),
                        (MOI.VectorQuadraticFunction{T} where T, MOI.Zeros),
                        (MOI.SingleVariable, MOI.Integer),
                        (MOI.SingleVariable, MOI.ZeroOne)]
    fieldnames = Symbol[]
    constrainttypes = []
    for (F, S) in constraint_specs
        E = moi_to_native_type(F)
        funcsym = F isa UnionAll ? F.body.name.name : F.name.name
        setsym = S isa UnionAll ? S.body.name.name : S.name.name
        ctype = if F isa UnionAll && S isa UnionAll
            Constraint{E{T}, F{T}, S{T}} where T
        elseif F isa UnionAll
            Constraint{E{T}, F{T}, S} where T
        elseif S isa UnionAll
            Constraint{E, F, S{T}} where T
        else
            Constraint{E, F, S}
        end
        push!(fieldnames, Symbol(lowercase(String(funcsym)) * "_in_" * lowercase(String(setsym))))
        push!(constrainttypes, ctype)
    end

    # Constraints struct
    @eval begin
        struct Constraints{T}
            $([:($name::Vector{$(ctype isa UnionAll ? :($ctype{T}) : ctype)}) for (name, ctype) in zip(fieldnames, constrainttypes)]...)
        end

        function Constraints{T}() where {T}
            Constraints{T}($([:(Vector{$(ctype isa UnionAll ? :($ctype{T}) : ctype)}()) for ctype in constrainttypes]...))
        end
    end

    # Base.push! methods
    for ((F, S), name, ctype) in zip(constraint_specs, fieldnames, constrainttypes)
        if ctype isa UnionAll
            @eval function Base.push!(constraints::Constraints{T}, constraint::$(:($ctype{T}))) where T
                push!(constraints.$name, constraint)
            end
        else
            @eval Base.push!(constraints::Constraints, constraint::$(ctype)) = push!(constraints.$name, constraint)
        end
    end

    # update!
    @eval begin
        function update!(constraints::Constraints, optimizer::MOI.AbstractOptimizer, varmap)
            $(map(fieldnames) do fieldname
                quote
                    for constraint in constraints.$fieldname
                        update!(constraint, optimizer, varmap)
                    end
                end
            end...)
            nothing
        end
    end

    # mapindices!
    @eval begin
        function mapindices!(constraints::Constraints, idxmap)
            $(map(fieldnames) do fieldname
                quote
                    for constraint in constraints.$fieldname
                        constraint.optimizerindex = idxmap[constraint.modelindex]
                    end
                end
            end...)
            nothing
        end
    end
end


const FAKE_CONCEPTNET = """
5 5
/c/en/this 1 2 3 4 5
/c/en/is -1 -2 -3 -4 -5
/c/en/# 1 1 1 1 1
/c/fr/etre 6 7 8 9 10
/c/fr/# 2 2 2 2 2
"""
@testset "Conceptnet-2-WordVectors" begin
    filepath, io = mktemp()
    write(io, FAKE_CONCEPTNET)
    close(io)
    cptnet = load_embeddings(filepath)
    for lang in keys(cptnet.embeddings)
        wv = conceptnet2wv(cptnet, lang)
        _, _, data_type = typeof(cptnet).parameters
        @test wv isa WordVectors{String, data_type, Int}
        @test length(wv.vocab) == length(cptnet.embeddings[lang])
        @test size(wv.vectors, 1) == cptnet.width
        for word in keys(cptnet.embeddings[lang])
            @test cptnet[lang, word] == get_vector(wv, word)
        end
    end
end


module TikhonovEnsembleKalmanSampling

using Distributed
using Statistics
using LinearAlgebra
using Optim

include("sampling.jl")

export teks_update, whteks_update
export teks, whteks

end # module


using Test, MathOptInterface

const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIU = MOI.Utilities

import Cbc
const OPTIMIZER = Cbc.Optimizer()
MOI.set(OPTIMIZER, MOI.Silent(), true)

@testset "SolverName" begin
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
end

@testset "supports_default_copy_to" begin
    @test !MOIU.supports_allocate_load(OPTIMIZER, false)
    @test !MOIU.supports_allocate_load(OPTIMIZER, true)
    @test !MOIU.supports_default_copy_to(OPTIMIZER, false)
    @test !MOIU.supports_default_copy_to(OPTIMIZER, true)
end

const CACHE = MOIU.UniversalFallback(MOIU.Model{Float64}())
const CACHED = MOIU.CachingOptimizer(CACHE, OPTIMIZER)
const BRIDGED = MOI.Bridges.full_bridge_optimizer(CACHED, Float64)

const CONFIG = MOIT.TestConfig(duals = false, infeas_certificates = false)

@testset "basic_constraint_tests" begin
    MOIT.basic_constraint_tests(CACHED, CONFIG)
end

@testset "Unit" begin
    MOIT.unittest(BRIDGED, CONFIG, [
        "number_threads", # FIXME implement `MOI.NumberOfThreads`
        "solve_time", # FIXME implement `MOI.SolveTime`
        "solve_unbounded_model", # INFEASIBLE_OR_UNBOUNDED instead of DUAL_INFEASIBLE
        "delete_soc_variables", "solve_qcp_edge_cases", "solve_qp_edge_cases"  # No quadratics
    ])
end

@testset "ModelLike" begin
    @test MOI.get(CACHED, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
    @testset "default_objective_test" begin
         MOIT.default_objective_test(CACHED)
     end
     @testset "default_status_test" begin
         MOIT.default_status_test(CACHED)
     end
    @testset "nametest" begin
        MOIT.nametest(CACHED)
    end
    @testset "validtest" begin
        MOIT.validtest(CACHED)
    end
    @testset "emptytest" begin
        # Requires VectorOfVariables
        # MOIT.emptytest(CACHED)
    end
    @testset "orderedindicestest" begin
        MOIT.orderedindicestest(CACHED)
    end
    @testset "copytest" begin
        # Requires VectorOfVariables
        # MOIT.copytest(CACHED, MOIU.CachingOptimizer(
        #     ModelForCachingOptimizer{Float64}(),
        #     Cbc.Optimizer()
        # ))
    end
end

@testset "Continuous Linear" begin
    MOIT.contlineartest(BRIDGED, CONFIG,
                        # FIXME starting values
                        ["partial_start"])
end

@testset "Integer Linear" begin
    MOIT.intlineartest(BRIDGED, CONFIG, [
        "indicator1", "indicator2", "indicator3", "indicator4",
        # FIXME starting values
        "knapsack"
    ])
end

@testset "Test params" begin
    knapsack_model = MOIU.Model{Float64}()
    MOIU.loadfromstring!(knapsack_model, """
        variables: x, y
        maxobjective: x + y
        c1: x in ZeroOne()
        c2: y in ZeroOne()
        c3: x + y <= 1.0
    """)
    model = Cbc.Optimizer(maxNodes = 0, presolve = "off", cuts = "off",
                          heur = "off", logLevel = 0)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.NODE_LIMIT
    # We also check that options are not destroyed on `empty!`.
    MOI.empty!(model)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.NODE_LIMIT
end


abstract type AbstractEmitter{T} end

@inline get_doa(emitter::AbstractEmitter) = get_doa(emitter.doa)
@inline get_existence(emitter::AbstractEmitter) = get_existence(emitter.exists)

# Can be removed once https://github.com/JuliaLang/julia/pull/32968 is merged
# Julia v1.4
filteremitters(f, xs::Tuple) = Base.afoldl((ys, x) -> f(x) ? (ys..., x) : ys, (), xs...)

function get_steer_vec(manifold::AbstractManifold, emitter::AbstractEmitter, attitude)
    get_steer_vec(manifold, get_doa(emitter), attitude)
end

function get_steer_vec(manifold::IdealManifold{1}, emitter::AbstractEmitter, attitude)
    1.0
end


################################################################################
# same with LLL
################################################################################

function single_env{T}(c::ClassGrpCtx{T}, I::Hecke.SmallLLLRelationsCtx, nb::Int, rat::Float64, max_good::Int = 2)
  bad_norm = 0
  rk = rank(c.M)
  good = 0
  while true
    e = class_group_small_lll_elements_relation_next(I)
    n = norm_div(e, norm(I.A), nb)
    if I.cnt > length(I.b)^2
      break
    end
    if nbits(numerator(n)) > nb - 25
      bad_norm += 1
      if I.cnt > 100  && bad_norm / I.cnt > 0.1
        @vprint :ClassGroup 2 "norm too large, $(I.cnt) has size $(nbits(numerator(n))) should be <= $(nb - 20) $bad_norm $(I.cnt)\n"
        break
      end
      continue
    end
    fl = class_group_add_relation(c, e, n, norm(I.A), integral = true)
    if !fl  && I.cnt/(good+1) > 2*c.expect
      @vprint :ClassGroup 2 "not enough progress $(I.cnt) $(c.expect) $good\n"
      break
    end
    if fl
      good += 1
    end
    if fl && max_good > -1
      if max_good < good
        @vprint :ClassGroup 2 "found enough $(I.cnt)\n"
        break
      end
    end
    if false && fl && good > 0 && (rank(c.M)- rk+1)/good < rat
      @vprint :ClassGroup 2 "rank too slow $(I.cnt) $good $(rank(c.M)) $rk\n"
      break
    end
  end
end

function class_group_via_lll{T}(c::ClassGrpCtx{T}, rat::Float64 = 0.2)
  O = order(c.FB.ideals[1])
  nb = nbits(abs(discriminant(O)))
  nb = div(nb, 2) + 30

  rt = time_ns()
  I = class_group_small_lll_elements_relation_start(c, O)
  single_env(c, I, nb, rat/10, -1)
  @vprint :ClassGroup 1 "search in order:  $((time_ns()-rt)*1e-9) rel mat:  $(c.M.bas_gens)\nin $(c.M)"

  @vtime :ClassGroup 1 h, piv = class_group_get_pivot_info(c)
  if h == 1 return c; end

#  for p = piv
#    I = class_group_small_lll_elements_relation_start(c, c.FB.ideals[p])
#    single_env(c, I, nb, rat, 1)
#  end
#
#  @vprint :ClassGroup 1 "search in ideals:  $((time_ns()-rt)*1e-9) rel mat:  $(c.M.bas_gens)\n"

#  @vtime :ClassGroup 1 h, piv = class_group_get_pivot_info(c)
#  if h == 1 return c; end

#  @vprint :ClassGroup 1 "Now with random...\n"
#  @vprint :ClassGroup 1 "length(piv) = $(length(piv)) and h = $h\n"
#  @vprint :ClassGroup 1 "$(piv)\n"

  class_group_new_relations_via_lll(c, rat, extra = -1)

  return c
end

function class_group_new_relations_via_lll{T}(c::ClassGrpCtx{T}, rat::Float64 = 0.2; extra::Int = 5, rand_exp::Int = 1)

  O = order(c.FB.ideals[1])
  nb = nbits(abs(discriminant(O)))
  nb = div(nb, 2) + 30

  st = c.rel_cnt

  @vtime :ClassGroup 1 h, piv = class_group_get_pivot_info(c)

  @vprint :ClassGroup 1 "Now with random...\n"
  @vprint :ClassGroup 1 "length(piv) = $(length(piv)) and h = $h\n"
  @vprint :ClassGroup 1 "$(piv)\n"
  if length(piv) == 0
    for i=1:5
      push!(piv, rand(1:length(c.FB.ideals)))
    end
    @vprint :ClassGroup 1 "piv was empty, supplemented it to\n"
    @vprint :ClassGroup 1 "$(piv)\n"
  end

  start = max(1, max(div(length(c.FB.ideals), 2)+1, length(c.FB.ideals)-10*(1+div(rand_exp, 3))))
  stop = length(c.FB.ideals)
  if isdefined(c, :randomClsEnv)
    rand_env = c.randomClsEnv
    random_extend(rand_env, 2.0)
    @vprint :ClassGroup 1 "re-using and extending random: $(nbits(norm(rand_env.rand))) and $(rand_env.exp)\n"
  else
    @vprint :ClassGroup 1 "want $(stop-start) primes for random. Try distinct rational primes...\n"
    JJ = [c.FB.ideals[stop]]
    start += length(c.FB.ideals) - stop
    stop  += length(c.FB.ideals) - stop
    i = stop
    while i>1 && length(JJ) < stop - start
      i -= 1
      if degree(c.FB.ideals[i]) > 1
        continue
      end
      if minimum(c.FB.ideals[i]) in [minimum(x) for x = JJ]
        continue
      else
        push!(JJ, c.FB.ideals[i])
      end
    end
    while length(JJ) < stop - start
      AA = rand(c.FB.ideals)
      if !(AA in JJ)
        push!(JJ, AA)
      end
    end
    rand_env = random_init(JJ, lb = root(abs(discriminant(O)), 2)^1, ub = abs(discriminant(O))^1, reduce = false)
    c.randomClsEnv = rand_env
  end

  if h > 0
    rand_exp += 1
    while gcd(h, rand_exp) > 1
      rand_exp += 1
    end
  end

  while true
    st = c.rel_cnt
    while (c.rel_cnt - st < 2)
      for p = piv
        @vprint :ClassGroup 1 "p: $p $rand_exp $(length(rand_env.base))\n"
        @vtime :ClassGroup 3 J = random_get(rand_env, reduce = false)
  #      @show nbits(norm(J)), rand_env.exp, rand_exp
        @vtime :ClassGroup 3 J *= c.FB.ideals[p]^rand_exp
        @vtime :ClassGroup 3 I = class_group_small_lll_elements_relation_start(c, J)
        @vtime :ClassGroup 3 single_env(c, I, nb, 0.8, -1)
        if h > 0 && c.rel_cnt - st > 2
          return
        else
          if h == 0 && c.rel_cnt - st > length(piv) + 2
            break
          end
        end
      end
    end

    @vprint :ClassGroup 1 "eval info\n"
    @vtime :ClassGroup 1 h, piv_new = class_group_get_pivot_info(c)
    @vprint :ClassGroup 1 "length(piv) = $(length(piv_new)) and h = $h\n"
    @vprint :ClassGroup 1 "$(piv_new)\n"

    if piv_new == piv
      if h > 0
        extra = 5
#        J = [rand(c.FB.ideals) for x=1:10]
#        println("extending rand")
#        random_extend(rand_env, J)
        random_extend(rand_env, root(abs(discriminant(O)), 2))
      end
      rand_exp += 1
      rand_exp = min(rand_exp, 13)
      if h>0
        while gcd(rand_exp, h) > 1
          rand_exp += 1
        end
      end

      if rand_exp % 3 == 0
        start = max(start -10, 1)
      end
    end
    piv = piv_new
    if h == 1 return end
  end
end


for T in (:Int8, :Int16, :Int32, :Int64, :Int128, :Float16, :Float32, :Float64)
    @eval begin
        promote_type(::Type{ArbComplex}, ::Type{$T}) = ArbComplex
        promote_type(::Type{ArbReal}, ::Type{$T}) = ArbReal
        promote_type(::Type{ArbFloat}, ::Type{$T}) = ArbFloat

        promote_type(::Type{ArbComplex{P}}, ::Type{$T}) where {P} = ArbComplex{P}
        promote_type(::Type{ArbReal{P}}, ::Type{$T}) where {P} = ArbReal{P}
        promote_type(::Type{ArbFloat{P}}, ::Type{$T}) where {P} = ArbFloat{P}

        promote_rule(::Type{ArbComplex}, ::Type{$T}) = ArbComplex
        promote_rule(::Type{ArbReal}, ::Type{$T}) = ArbReal
        promote_rule(::Type{ArbFloat}, ::Type{$T}) = ArbFloat

        promote_rule(::Type{ArbComplex{P}}, ::Type{$T}) where {P} = ArbComplex{P}
        promote_rule(::Type{ArbReal{P}}, ::Type{$T}) where {P} = ArbReal{P}
        promote_rule(::Type{ArbFloat{P}}, ::Type{$T}) where {P} = ArbFloat{P}
    end
end

promote_type(::Type{ArbComplex}, ::Type{ArbReal}) = ArbComplex
promote_type(::Type{ArbComplex}, ::Type{ArbFloat}) = ArbComplex
promote_type(::Type{ArbReal}, ::Type{ArbFloat}) = ArbReal

promote_rule(::Type{ArbComplex}, ::Type{ArbReal}) = ArbComplex
promote_rule(::Type{ArbComplex}, ::Type{ArbFloat}) = ArbComplex
promote_rule(::Type{ArbReal}, ::Type{ArbFloat}) = ArbReal

promote_type(::Type{ArbComplex{P}}, ::Type{ArbReal{P}}) where {P} = ArbComplex{P}
promote_type(::Type{ArbComplex{P}}, ::Type{ArbFloat{P}}) where {P} = ArbComplex{P}
promote_type(::Type{ArbReal{P}}, ::Type{ArbFloat{P}}) where {P} = ArbReal{P}

promote_rule(::Type{ArbComplex{P}}, ::Type{ArbReal{P}}) where {P} = ArbComplex{P}
promote_rule(::Type{ArbComplex{P}}, ::Type{ArbFloat{P}}) where {P} = ArbComplex{P}
promote_rule(::Type{ArbReal{P}}, ::Type{ArbFloat{P}}) where {P} = ArbReal{P}

promote_rule(::Type{ArbComplex{P}}, ::Type{Mag}) where {P} = ArbComplex{P}
promote_rule(::Type{ArbReal{P}}, ::Type{Mag}) where {P} = ArbReal{P}
promote_rule(::Type{ArbFloat{P}}, ::Type{Mag}) where {P} = ArbFloat{P}

promote_type(::Type{ArbFloat{P}}, ::Type{BigInt}) where {P} = ArbFloat{P}
promote_type(::Type{ArbReal{P}}, ::Type{BigInt}) where {P} = ArbReal{P}
promote_type(::Type{ArbComplex{P}}, ::Type{BigInt}) where {P} = ArbComplex{P}

promote_type(::Type{ArbFloat{P}}, ::Type{Rational{I}}) where {P,I} = ArbFloat{P}
promote_type(::Type{ArbReal{P}}, ::Type{Rational{I}}) where {P,I} = ArbReal{P}
promote_type(::Type{ArbComplex{P}}, ::Type{Rational{I}}) where {P,I} = ArbComplex{P}

promote_rule(::Type{ArbFloat{P}}, ::Type{BigInt}) where {P} = promote_type(ArbFloat{P}, BigInt)
promote_rule(::Type{ArbReal{P}}, ::Type{BigInt}) where {P} = promote_type(ArbReal{P}, BigInt)
promote_rule(::Type{ArbComplex{P}}, ::Type{BigInt}) where {P} = promote_type(ArbComplex{P}, BigInt)

promote_rule(::Type{ArbFloat{P}}, ::Type{Rational{I}}) where {P,I} = promote_type(ArbFloat{P}, Rational{I})
promote_rule(::Type{ArbReal{P}}, ::Type{Rational{I}}) where {P,I} = promote_type(ArbReal{P}, Rational{I})
promote_rule(::Type{ArbComplex{P}}, ::Type{Rational{I}}) where {P,I} = promote_type(ArbComplex{P}, Rational{I})

promote_type(::Type{BigFloat},::Type{ArbFloat{P}}) where {P} = ArbFloat{P}
promote_type(::Type{BigFloat},::Type{ArbReal{P}}) where {P} = ArbReal{P}
promote_type(::Type{BigFloat},::Type{ArbComplex{P}}) where {P} = ArbComplex{P}
promote_type(::Type{ArbFloat{P}}, ::Type{BigFloat}) where {P} = ArbFloat{P}
promote_type(::Type{ArbReal{P}}, ::Type{BigFloat}) where {P} = ArbReal{P}
promote_type(::Type{ArbComplex{P}}, ::Type{BigFloat}) where {P} = ArbComplex{P}

promote_rule(::Type{ArbFloat{P}}, ::Type{BigFloat}) where {P} = promote_type(ArbFloat{P}, BigFloat)
promote_rule(::Type{ArbReal{P}}, ::Type{BigFloat}) where {P} = promote_type(ArbReal{P}, BigFloat)
promote_rule(::Type{ArbComplex{P}}, ::Type{BigFloat}) where {P} = promote_type(ArbComplex{P}, BigFloat)

promote_rule(::Type{ArbFloat{P}}, ::Type{I}) where {P, I<:Integer} = ArbFloat{P}
promote_rule(::Type{ArbReal{P}}, ::Type{I}) where {P, I<:Integer} = ArbReal{P}
promote_rule(::Type{ArbComplex{P}}, ::Type{I}) where {P, I<:Integer} = ArbComplex{P}

promote_rule(::Type{ArbFloat{P}}, ::Type{F}) where {P, F<:AbstractFloat} = ArbFloat{P}
promote_rule(::Type{ArbReal{P}}, ::Type{F}) where {P, F<:AbstractFloat} = ArbReal{P}
promote_rule(::Type{ArbComplex{P}}, ::Type{F}) where {P, F<:AbstractFloat} = ArbComplex{P}
promote_rule(::Type{ArbComplex{P}}, ::Type{C}) where {P, C<:Complex} = ArbComplex{P}

promote_rule(::Type{ArbFloat{P}}, x::ArbFloat{Q}) where {P,Q} =
    P>Q ? ArbFloat{P} : ArbFloat{Q}
promote_rule(::Type{ArbReal{P}}, x::ArbReal{Q}) where {P,Q} =
    P>Q ? ArbReal{P} : ArbReal{Q}
promote_rule(::Type{ArbComplex{P}}, x::ArbComplex{Q}) where {P,Q} =
    P>Q ? ArbComplex{P} : ArbComplex{Q}
promote_rule(::Type{ArbFloat{P}}, x::ArbReal{Q}) where {P,Q} =
    P>Q ? ArbReal{P} : ArbReal{Q}
promote_rule(::Type{ArbFloat{P}}, x::ArbComplex{Q}) where {P,Q} =
    P>Q ? ArbComplex{P} : ArbComplex{Q}
promote_rule(::Type{ArbReal{P}}, x::ArbComplex{Q}) where {P,Q} =
    P>Q ? ArbComplex{P} : ArbComplex{Q}

promote_rule(::Type{ArbFloat{P}},::Type{ArbFloat{Q}}) where {P,Q} = P<Q ? ArbFloat{P} : ArbFloat{Q}
promote_rule(::Type{ArbReal{P}},::Type{ArbReal{Q}}) where {P,Q} = P<Q ? ArbReal{P} : ArbReal{Q}
promote_rule(::Type{ArbComplex{P}},::Type{ArbComplex{Q}}) where {P,Q} = P<Q ? ArbComplex{P} : ArbComplex{Q}

promote_rule(::Type{<:AbstractIrrational}, ::Type{ArbFloat{P}}) where {P} = ArbFloat{P}
promote_rule(::Type{<:AbstractIrrational}, ::Type{ArbReal{P}}) where {P} = ArbReal{P}
promote_rule(::Type{<:AbstractIrrational}, ::Type{ArbComplex{P}}) where {P} = ArbComplex{P}

convert(::Type{ArbFloat}, x::ArbFloat{P}) where {P} = x
convert(::Type{ArbReal}, x::ArbReal{P}) where {P} = x
convert(::Type{ArbComplex}, x::ArbComplex{P}) where {P} = x
convert(::Type{ArbFloat}, x::ArbReal{P}) where {P} = ArbFloat{P}(x)
convert(::Type{ArbFloat}, x::ArbComplex{P}) where {P} = ArbFloat{P}(real(x))
convert(::Type{ArbReal}, x::ArbComplex{P}) where {P} = real(x)
convert(::Type{ArbReal}, x::ArbFloat{P}) where {P} = ArbReal{P}(x)
convert(::Type{ArbComplex}, x::ArbFloat{P}) where {P} = ArbComplex{P}(x)

convert(::Type{ArbReal{P}}, x::ArbComplex{P}) where {P} = real(x)
convert(::Type{ArbFloat{P}}, x::ArbComplex{P}) where {P} = midpoint(real(x), ArbFloat{P})
convert(::Type{Mag}, x::ArbComplex{P}) where {P} = radius(real(x), Mag)

convert(::Type{ArbComplex{P}}, x::ArbReal{P}) where {P} = ArbComplex{P}(x)
convert(::Type{ArbComplex{P}}, x::ArbFloat{P}) where {P} = ArbComplex{P}(ArbReal{P}(x))
convert(::Type{ArbComplex{P}}, x::Mag) where {P} = ArbComplex{P}(ArbFloat{P}(x))

convert(::Type{ArbFloat{P}}, x::ArbReal{P}) where {P} = midpoint(x, ArbFloat{P})
convert(::Type{Mag}, x::ArbReal{P}) where {P} = radius(x, Mag)

convert(::Type{ArbReal{P}}, x::ArbFloat{P}) where {P} = ArbReal{P}(x)
convert(::Type{ArbReal{P}}, x::Mag) where {P} = ArbReal{P}(ArbFloat{P}(x))

convert(::Type{Mag}, x::ArbFloat{P}) where {P} = Mag(x)
convert(::Type{ArbFloat{P}}, x::Mag) where {P} = ArbFloat{P}(x)

convert(::Type{ArbFloat{P}}, x::BigFloat) where {P} = ArbFloat{P}(x)
convert(::Type{ArbReal{P}}, x::BigFloat) where {P} = ArbReal{P}(x)
convert(::Type{ArbComplex{P}}, x::BigFloat) where {P} = ArbComplex{P}(x)

convert(::Type{ArbFloat{P}}, x::BigInt) where {P} = ArbFloat{P}(x)
convert(::Type{ArbReal{P}}, x::BigInt) where {P} = ArbReal{P}(x)
convert(::Type{ArbComplex{P}}, x::BigInt) where {P} = ArbComplex{P}(x)

convert(::Type{ArbFloat{P}}, x::I) where {P, I<:Integer} = ArbFloat{P}(x)
convert(::Type{ArbReal{P}}, x::I) where {P, I<:Integer} = ArbReal{P}(x)
convert(::Type{ArbComplex{P}}, x::I) where {P, I<:Integer} = ArbComplex{P}(x)

convert(::Type{ArbFloat{P}}, x::Rational{I}) where {P, I} = ArbFloat{P}(x.num) / ArbFloat{P}(x.den)
convert(::Type{ArbReal{P}}, x::Rational{I}) where {P, I} = ArbReal{P}(x.num) / ArbReal{P}(x.den)
convert(::Type{ArbComplex{P}}, x::Rational{I}) where {P, I} = ArbComplex{P}(ArbReal{P}(x))

function rationalize(::Type{I}, x::ArbFloat{P}) where {P,I<:Integer}
    oldprec = precision(BigFloat)
    setprecision(BigFloat, P)
    y = BigFloat(x)
    z = rationalize(I, y)
    setprecision(BigFloat, oldprec)
    return z
end

rationalize(::Type{I}, x::ArbReal{P}) where {P, I<:Integer} = rationalize(I, ArbFloat{P}(x))
rationalize(::Type{I}, x::ArbComplex{P}) where {P, I<:Integer} = Complex{I}(rationalize(I, ArbFloat{P}(real(x))), rationalize(I, ArbFloat{P}(imag(x))))

rationalize(x::ArbFloat{P}) where {P} = rationalize(Int64, x)
rationalize(x::ArbReal{P}) where {P} = rationalize(Int64, ArbFloat{P}(x))
rationalize(x::ArbComplex{P}) where {P} = Complex{Int64}(rationalize(Int64, ArbFloat{P}(real(x))), rationalize(Int64, ArbFloat{P}(imag(x))))

convert(::Type{Rational{I}}, x::ArbFloat{P}) where {P, I} = rationalize(I, x)
convert(::Type{Rational{I}}, x::ArbReal{P}) where {P, I} = rationalize(I, x)
convert(::Type{Rational{I}}, x::ArbComplex{P}) where {P, I} = rationalize(I, x)

convert(::Type{ArbFloat{P}}, x::AbstractFloat) where {P} = ArbFloat{P}(x)
convert(::Type{ArbReal{P}}, x::AbstractFloat) where {P} = ArbReal{P}(x)
convert(::Type{ArbComplex{P}}, x::AbstractFloat) where {P} = ArbComplex{P}(x)
convert(::Type{ArbComplex{P}}, x::Complex) where {P} = ArbComplex{P}(x)


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

"""
    SemidiscretizationEulerAcoustics(semi_acoustics::SemiAcoustics, semi_euler::SemiEuler;
                                     source_region=x->true, weights=x->1.0)

!!! warning "Experimental code"
    This semidiscretization is experimental and may change in any future release.

Construct a semidiscretization of the acoustic perturbation equations that is coupled with
the compressible Euler equations via source terms for the perturbed velocity. Both
semidiscretizations have to use the same mesh and solvers with a shared basis. The coupling region
is described by a function `source_region` that maps the coordinates of a single node to `true` or
`false` depending on whether the point lies within the coupling region or not. A weighting function
`weights` that maps coordinates to weights is applied to the acoustic source terms.
Note that this semidiscretization should be used in conjunction with
[`EulerAcousticsCouplingCallback`](@ref) and only works in two dimensions.
"""
struct SemidiscretizationEulerAcoustics{SemiAcoustics, SemiEuler, Cache} <: AbstractSemidiscretization
  semi_acoustics::SemiAcoustics
  semi_euler::SemiEuler
  performance_counter::PerformanceCounter
  cache::Cache

  function SemidiscretizationEulerAcoustics{SemiAcoustics, SemiEuler, Cache}(
      semi_acoustics, semi_euler, cache) where {SemiAcoustics, SemiEuler, Cache}

    # Currently both semidiscretizations need to use a shared mesh
    @assert semi_acoustics.mesh == semi_euler.mesh

    # Check if both solvers use the same polynomial basis
    @assert semi_acoustics.solver.basis == semi_euler.solver.basis

    performance_counter = PerformanceCounter()
    new(semi_acoustics, semi_euler, performance_counter, cache)
  end
end

function SemidiscretizationEulerAcoustics(semi_acoustics::SemiAcoustics, semi_euler::SemiEuler;
                                          source_region=x->true, weights=x->1.0) where
    {Mesh, SemiAcoustics<:SemidiscretizationHyperbolic{Mesh, <:AbstractAcousticPerturbationEquations},
     SemiEuler<:SemidiscretizationHyperbolic{Mesh, <:AbstractCompressibleEulerEquations}}

  cache = create_cache(SemidiscretizationEulerAcoustics, source_region, weights,
                       mesh_equations_solver_cache(semi_acoustics)...)

  return SemidiscretizationEulerAcoustics{typeof(semi_acoustics), typeof(semi_euler), typeof(cache)}(
    semi_acoustics, semi_euler, cache)
end

function create_cache(::Type{SemidiscretizationEulerAcoustics}, source_region, weights,
                      mesh, equations::AcousticPerturbationEquations2D, dg::DGSEM, cache)

  coupled_element_ids = get_coupled_element_ids(source_region, equations, dg, cache)

  acoustic_source_terms = zeros(eltype(cache.elements), (ndims(equations), nnodes(dg), nnodes(dg),
                                                         length(coupled_element_ids)))

  acoustic_source_weights = precompute_weights(source_region, weights, coupled_element_ids,
                                               equations, dg, cache)

  return (; acoustic_source_terms, acoustic_source_weights, coupled_element_ids)
end

function get_coupled_element_ids(source_region, equations, dg::DGSEM, cache)
  coupled_element_ids = Vector{Int}(undef, 0)

  for element in eachelement(dg, cache)
    for j in eachnode(dg), i in eachnode(dg)
      x = get_node_coords(cache.elements.node_coordinates, equations, dg, i, j, element)
      if source_region(x)
        push!(coupled_element_ids, element)
        break
      end
    end
  end

  return coupled_element_ids
end

function precompute_weights(source_region, weights, coupled_element_ids, equations, dg::DGSEM, cache)
  acoustic_source_weights = zeros(eltype(cache.elements),
                                  (nnodes(dg), nnodes(dg), length(coupled_element_ids)))

  @threaded for k in 1:length(coupled_element_ids)
    element = coupled_element_ids[k]
    for j in eachnode(dg), i in eachnode(dg)
      x = get_node_coords(cache.elements.node_coordinates, equations, dg, i, j, element)
      acoustic_source_weights[i, j, k] = source_region(x) ? weights(x) : zero(weights(x))
    end
  end

  return acoustic_source_weights
end

function Base.show(io::IO, semi::SemidiscretizationEulerAcoustics)
  @nospecialize semi # reduce precompilation time

  print(io, "SemidiscretizationEulerAcoustics(")
  print(io,       semi.semi_acoustics)
  print(io, ", ", semi.semi_euler)
  print(io, ", cache(")
  for (idx, key) in enumerate(keys(semi.cache))
    idx > 1 && print(io, " ")
    print(io, key)
  end
  print(io, "))")
end

function Base.show(io::IO, mime::MIME"text/plain", semi::SemidiscretizationEulerAcoustics)
  @nospecialize semi # reduce precompilation time

  if get(io, :compact, false)
    show(io, semi)
  else
    summary_header(io, "SemidiscretizationEulerAcoustics")
    summary_line(io, "semidiscretization Euler", semi.semi_euler |> typeof |> nameof)
    show(increment_indent(io), mime, semi.semi_euler)
    summary_line(io, "semidiscretization acoustics", semi.semi_acoustics |> typeof |> nameof)
    show(increment_indent(io), mime, semi.semi_acoustics)
    summary_footer(io)
  end
end

# The acoustics semidiscretization is the main semidiscretization.
@inline function mesh_equations_solver_cache(semi::SemidiscretizationEulerAcoustics)
  return mesh_equations_solver_cache(semi.semi_acoustics)
end

@inline Base.ndims(semi::SemidiscretizationEulerAcoustics) = ndims(semi.semi_acoustics)
@inline Base.real(semi::SemidiscretizationEulerAcoustics) = real(semi.semi_acoustics)

# Computes the coefficients of the initial condition
@inline function compute_coefficients(t, semi::SemidiscretizationEulerAcoustics)
  compute_coefficients(t, semi.semi_acoustics)
end

@inline function compute_coefficients!(u_ode, t, semi::SemidiscretizationEulerAcoustics)
  compute_coefficients!(u_ode, t, semi.semi_acoustics)
end

@inline function calc_error_norms(func, u, t, analyzer, semi::SemidiscretizationEulerAcoustics,
                                  cache_analysis)
  calc_error_norms(func, u, t, analyzer, semi.semi_acoustics, cache_analysis)
end

function rhs!(du_ode, u_ode, semi::SemidiscretizationEulerAcoustics, t)
  @unpack semi_acoustics, cache = semi
  @unpack acoustic_source_terms, acoustic_source_weights, coupled_element_ids = cache

  du_acoustics = wrap_array(du_ode, semi_acoustics)

  time_start = time_ns()

  @trixi_timeit timer() "acoustics rhs!" rhs!(du_ode, u_ode, semi_acoustics, t)

  @trixi_timeit timer() "add acoustic source terms" add_acoustic_source_terms!(
    du_acoustics, acoustic_source_terms, acoustic_source_weights, coupled_element_ids,
    mesh_equations_solver_cache(semi_acoustics)...)

  runtime = time_ns() - time_start
  put!(semi.performance_counter, runtime)

  return nothing
end

function add_acoustic_source_terms!(du_acoustics, acoustic_source_terms, source_weights,
                                    coupled_element_ids, mesh::TreeMesh{2}, equations, dg::DGSEM,
                                    cache)

  @threaded for k in 1:length(coupled_element_ids)
    element = coupled_element_ids[k]

    for j in eachnode(dg), i in eachnode(dg)
      du_acoustics[1, i, j, element] += source_weights[i, j, k] * acoustic_source_terms[1, i, j, k]
      du_acoustics[2, i, j, element] += source_weights[i, j, k] * acoustic_source_terms[2, i, j, k]
    end
  end

  return nothing
end

end # @muladd