const MOI = MathOptInterface

MOI.Utilities.@model(
    Optimizer,
    (),  # Scalar sets
    (),  # Typed scalar sets
    (MOI.Complements,),  # Vector sets
    (),  # Typed vector sets
    (),  # Scalar functions
    (),  # Typed scalar functions
    (),  # Vector functions
    (MOI.VectorAffineFunction,),  # Typed vector functions
    true,  # is_optimizer
)

function MOI.supports_constraint(
    ::Optimizer, ::Type{MOI.SingleVariable}, ::Type{S}
) where {S <: Union{MOI.Semiinteger, MOI.Semicontinuous, MOI.ZeroOne, MOI.Integer}}
    return false
end

function MOI.supports(
    ::Optimizer, ::MOI.ObjectiveFunction{F}
) where {F<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction}}
    return false
end

struct Solution
    status::MCP_Termination
    x::Vector{Float64}
    info::Information
end

solution(model::Optimizer) = get(model.ext, :solution, nothing)::Union{Nothing, Solution}

"""
    Optimizer()

Define a new PATH optimizer.

Pass options using `MOI.RawParameter`. Common options include:

 - output => "yes"
 - convergence_tolerance => 1e-6
 - time_limit => 3600

A full list of options can be found at http://pages.cs.wisc.edu/~ferris/path/options.pdf.

### Example

    optimizer = PATH.Optimizer()
    MOI.set(optimizer, MOI.RawParameter("output"), "no")
"""
function Optimizer()
    model = Optimizer{Float64}()
    model.ext[:silent] = false
    model.ext[:kwargs] = Dict{Symbol, Any}()
    model.ext[:solution] = nothing
    return model
end

function MOI.set(model::Optimizer, p::MOI.RawParameter, v)
    model.ext[:kwargs][Symbol(p.name)] = v
    return
end

function MOI.get(model::Optimizer, p::MOI.RawParameter)
    return get(model.ext[:kwargs], Symbol(p.name), nothing)
end

MOI.supports(model::Optimizer, ::MOI.Silent) = true

MOI.get(model::Optimizer, ::MOI.Silent) = model.ext[:silent]

function MOI.set(model::Optimizer, ::MOI.Silent, x::Bool)
    model.ext[:silent] = x
    return
end

MOI.get(model::Optimizer, ::MOI.SolverName) = c_api_Path_Version()

function MOI.supports(
    ::Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex}
)
    return true
end

function MOI.get(
    model::Optimizer, ::MOI.VariablePrimalStart, x::MOI.VariableIndex
)
    initial = get(model.ext, :variable_primal_start, nothing)
    if initial === nothing
        return nothing
    end
    return get(initial, x, nothing)
end

function MOI.set(
    model::Optimizer, ::MOI.VariablePrimalStart, x::MOI.VariableIndex, value
)
    initial = get(model.ext, :variable_primal_start, nothing)
    if initial === nothing
        model.ext[:variable_primal_start] = Dict{MOI.VariableIndex, Float64}()
        initial = model.ext[:variable_primal_start]
    end
    if value === nothing
        delete!(initial, x)
    else
        initial[x] = value
    end
    return
end

function _F_linear_operator(model::Optimizer)
    n = MOI.get(model, MOI.NumberOfVariables())
    M = SparseArrays.sparse(Int32[], Int32[], Float64[], n, n)
    q = zeros(n)
    has_term = fill(false, n)
    for index in MOI.get(
        model,
        MOI.ListOfConstraintIndices{
            MOI.VectorAffineFunction{Float64}, MOI.Complements
        }()
    )
        Fi = MOI.get(model, MOI.ConstraintFunction(), index)
        Si = MOI.get(model, MOI.ConstraintSet(), index)

        if 2 * Si.dimension != length(Fi.constants)
            error(
                "Dimension of constant vector $(length(Fi.constants)) does not match the " *
                "required dimension of the complementarity set $(2 * Si.dimension)."
            )
        elseif any(i -> !iszero(Fi.constants[Si.dimension + i]), 1:Si.dimension)
            error(
                "VectorAffineFunction malformed: a constant associated with a " *
                "complemented variable is not zero: $(Fi.constants[Si.dimension+1:end])."
            )
        end

        # First pass: get rows vector and check for invalid functions.
        rows = fill(0, Si.dimension)
        for term in Fi.terms
            if term.output_index <= Si.dimension
                # No-op: leave for second pass.
                continue
            elseif term.output_index > 2 * Si.dimension
                error(
                    "VectorAffineFunction malformed: output_index $(term.output_index) " *
                    "is too large."
                )
            end
            dimension_i = term.output_index - Si.dimension
            row_i = term.scalar_term.variable_index.value
            if rows[dimension_i] != 0 || has_term[row_i]
                error(
                    "The variable $(term.scalar_term.variable_index) appears in more " *
                    "than one complementarity constraint."
                )
            elseif term.scalar_term.coefficient != 1.0
                error(
                    "VectorAffineFunction malformed: variable " *
                    "$(term.scalar_term.variable_index) has a coefficient that is not 1 " *
                    "in row $(term.output_index) of the VectorAffineFunction."
                )
            end
            rows[dimension_i] = row_i
            has_term[row_i] = true
            q[row_i] = Fi.constants[dimension_i]
        end

        # Second pass: add to sparse array
        for term in Fi.terms
            s_term = term.scalar_term
            if term.output_index > Si.dimension || iszero(s_term.coefficient)
                continue
            end
            row_i = rows[term.output_index]
            if iszero(row_i)
                error(
                    "VectorAffineFunction malformed: expected variable in row " *
                    "$(Si.dimension + term.output_index)."
                )
            end
            M[row_i, s_term.variable_index.value] += s_term.coefficient
        end
    end

    return M, q
end

function _bounds_and_starting(model::Optimizer)
    x = MOI.get(model, MOI.ListOfVariableIndices())
    lower = fill(-INFINITY, length(x))
    upper = fill(INFINITY, length(x))
    initial = fill(0.0, length(x))
    for (i, xi) in enumerate(x)
        l, u = MOI.Utilities.get_bounds(model, Float64, xi)
        z = MOI.get(model, MOI.VariablePrimalStart(), xi)
        lower[i] = l
        upper[i] = u
        initial[i] = z !== nothing ? z : 0.5 * (l + u)
    end
    return lower, upper, initial
end

function MOI.optimize!(model::Optimizer)
    model.ext[:solution] = nothing
    lower, upper, initial = _bounds_and_starting(model)
    M, q = _F_linear_operator(model)
    status, x, info = solve_mcp(
        M,
        q,
        lower,
        upper,
        initial;
        silent = model.ext[:silent],
        [k => v for (k, v) in model.ext[:kwargs]]...
    )
    model.ext[:solution] = Solution(status, x, info)
    return
end

const _MCP_TERMINATION_STATUS_MAP =
    Dict{MCP_Termination, Tuple{MOI.TerminationStatusCode, String}}(
        MCP_Solved => (MOI.LOCALLY_SOLVED, "The problem was solved"),
        MCP_NoProgress => (MOI.SLOW_PROGRESS, "A stationary point was found"),
        MCP_MajorIterationLimit => (MOI.ITERATION_LIMIT, "Major iteration limit met"),
        MCP_MinorIterationLimit => (MOI.ITERATION_LIMIT, "Cumulative minor iterlim met"),
        MCP_TimeLimit => (MOI.TIME_LIMIT, "Ran out of time"),
        MCP_UserInterrupt => (MOI.INTERRUPTED, "Control-C, typically"),
        MCP_BoundError => (MOI.INVALID_MODEL, "Problem has a bound error"),
        MCP_DomainError => (MOI.NUMERICAL_ERROR, "Could not find a starting point"),
        MCP_Infeasible => (MOI.INFEASIBLE, "Problem has no solution"),
        MCP_Error => (MOI.OTHER_ERROR, "An error occured within the code"),
        MCP_LicenseError => (MOI.OTHER_ERROR, "License could not be found"),
        MCP_OK => (MOI.OTHER_ERROR, "")
    )

function MOI.get(model::Optimizer, ::MOI.TerminationStatus)
    if solution(model) === nothing
        return MOI.OPTIMIZE_NOT_CALLED
    end
    status, _ = _MCP_TERMINATION_STATUS_MAP[solution(model).status]
    return status
end

function MOI.get(model::Optimizer, ::MOI.RawStatusString)
    if solution(model) === nothing
        return "MOI.optimize! was not called yet"
    end
    _, reason = _MCP_TERMINATION_STATUS_MAP[solution(model).status]
    return reason
end

function MOI.get(model::Optimizer, ::MOI.ResultCount)
    return solution(model) !== nothing ? 1 : 0
end

function MOI.get(
    model::Optimizer, attr::MOI.VariablePrimal, x::MOI.VariableIndex
)
    MOI.check_result_index_bounds(model, attr)
    return solution(model).x[x.value]
end

function MOI.get(model::Optimizer, attr::MOI.PrimalStatus)
    if solution(model) === nothing || attr.N != 1
        return MOI.NO_SOLUTION
    end
    if MOI.get(model, MOI.TerminationStatus()) == MOI.LOCALLY_SOLVED
        return MOI.FEASIBLE_POINT
    end
    return MOI.UNKNOWN_RESULT_STATUS
end


module Turing

##############
# Dependency #
########################################################################
# NOTE: when using anything from external packages,                    #
#       let's keep the practice of explictly writing Package.something #
#       to indicate that's not implemented inside Turing.jl            #
########################################################################

using Requires, Reexport, ForwardDiff
using DistributionsAD, Bijectors, StatsFuns, SpecialFunctions
using Statistics, LinearAlgebra
using Markdown, Libtask
@reexport using Distributions, MCMCChains, Libtask
using Tracker: Tracker

import Base: ~, ==, convert, hash, promote_rule, rand, getindex, setindex!
import DynamicPPL: getspace

const PROGRESS = Ref(true)
function turnprogress(switch::Bool)
    @info "[Turing]: progress logging is $(switch ? "enabled" : "disabled") globally"
    PROGRESS[] = switch
end

const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_TURING", "0")))

# Random probability measures.
include("stdlib/distributions.jl")
include("stdlib/RandomMeasures.jl")
include("utilities/Utilities.jl")
using .Utilities
include("core/Core.jl")
using .Core
include("inference/Inference.jl")  # inference algorithms
using .Inference
include("variational/VariationalInference.jl")
using .Variational

# TODO: re-design `sample` interface in MCMCChains, which unify CmdStan and Turing.
#   Related: https://github.com/TuringLang/Turing.jl/issues/746
#@init @require CmdStan="593b3428-ca2f-500c-ae53-031589ec8ddd" @eval begin
#     @eval Utilities begin
#         using ..Turing.CmdStan: CmdStan, Adapt, Hmc
#         using ..Turing: HMC, HMCDA, NUTS
#         include("utilities/stan-interface.jl")
#     end
# end

@init @require DynamicHMC="bbc10e6e-7c05-544b-b16e-64fede858acb" @eval Inference begin
    import ..DynamicHMC

    if isdefined(DynamicHMC, :mcmc_with_warmup)
        using ..DynamicHMC: mcmc_with_warmup
        include("contrib/inference/dynamichmc.jl")
    else
        error("Please update DynamicHMC, v1.x is no longer supported")
    end
end

###########
# Exports #
###########

# Turing essentials - modelling macros and inference algorithms
export  @model,                 # modelling
        @varname,
        DynamicPPL,

        MH,                     # classic sampling
        RWMH,
        ESS,
        Gibbs,

        HMC,                    # Hamiltonian-like sampling
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,
        DynamicNUTS,
        ANUTS,

        IS,                     # particle-based sampling
        SMC,
        CSMC,
        PG,

        vi,                    # variational inference
        ADVI,

        sample,                 # inference
        psample,
        setchunksize,
        resume,
        @logprob_str,
        @prob_str,

        auto_tune_chunk_size!,  # helper
        setadbackend,
        setadsafe,

        turnprogress,           # debugging

        Flat,
        FlatPos,
        BinomialLogit,
        VecBinomialLogit,
        OrderedLogistic,
        LogPoisson,
        NamedDist,
        filldist,
        arraydist

# Reexports
using AbstractMCMC: sample, psample
export sample, psample

end


module JacProp

export Trajectory, default_activations, ModelTrainer, ADModelTrainer, sample_jacprop, push!, train!, display_modeltrainer, getter

using LTVModelsBase, Parameters, Reexport, Lazy, Juno, FunctionEnsembles
@reexport using LTVModels, Flux, ValueHistories, IterTools, MLDataUtils
using Flux: back!, truncate!, treelike, train!, mse, testmode!, params, jacobian, throttle
using Flux.Optimise: Param, optimiser, RMSProp, expdecay
using DSP, Plots, DataStructures, ReverseDiff#, InteractNext, Observables
export plot, plot!, scatter, scatter!
const RDiff = ReverseDiff
const default_activations = [swish, Flux.sigmoid, tanh, elu]
const IT = IterTools
using Statistics, LinearAlgebra, Random

abstract type AbstractModelTrainer end

"""
ModelTrainer

- `models` A vector of functions
- `opts` A vector of optimizers
- `losses` A vector of loss functions
- `R1` Parameter drift covariance for Kalman smoother
- `R2` Measurement noise covariance
- `P0` Initial covariance matrix for Kalman smoother
- `P` Covariance magnification of NN prior
- `cb` can be either `()->()` or a a Function that creates a closure around `loss` and `data` (loss,data)->(()->())

See also [`LTVModels`](@ref),  [`LTVModels.KalmanModel`](@ref)
"""
@with_kw mutable struct ModelTrainer{cbT,tT} <: AbstractModelTrainer
    models
    opts
    losses
    R1                = I
    R2                = 10_000I
    P0                = 10_000I
    σdivider::Float64 = 10.
    P::Float64        = 1.0
    trajs::Vector{Trajectory} = Trajectory[]
    cb::cbT
    modelhistory = []
    trace::tT = History(Float64)
end

@with_kw mutable struct ADModelTrainer{mT,tT} <: AbstractModelTrainer
    model::mT
    opt
    λ::Float64 = 1.
    testdata::Trajectory
    trajs::Vector{Trajectory} = Trajectory[]
    modelhistory = []
    trace::tT = History(Float64)
    tracev::tT = History(Float64)
    normalizer::Matrix{Float64} = ones(1,1)
end

include("utilities.jl")

to_callback(cb,args...) = methods(cb).mt.max_args > 1 ? cb(args...) : cb

"""
train!(mt::ModelTrainer; epochs=1, jacprop=0, useprior=true, trace = mt.trace)

See also [`ModelTrainer`](@ref)
"""
function Flux.train!(mt::ModelTrainer; epochs=1, jacprop=0, useprior=true, trace = mt.trace)
    @assert !isempty(mt.trajs) "No data in ModelTrainer"
    @unpack models,opts,losses = mt
    data1 = todata(mt)
    jacprop > 0 && (ltvmodels = fit_models(mt, useprior))
    startepoch = length(trace)+1
    @progress for epoch = startepoch:startepoch+epochs-1
        data = if jacprop > 0
            data2 = [todata(sample_jacprop(mt, ltvmodels)) for i = 1:jacprop]
            chain(data1, data2...)
        else
            data1
        end
        for (model,loss, opt) in zip(models,losses,opts)
            train!(loss, data, opt, cb=to_callback(mt.cb, epoch,loss,data1, trace, model, mt))
        end
    end
    push!(mt.modelhistory, deepcopy(models))
    trace
end

function Flux.train!(mt::ADModelTrainer; epochs=1, jacprop=0, trace = mt.trace, cb=()->())
    @assert !isempty(mt.trajs) "No data in ModelTrainer"
    @unpack model,opt,testdata,λ,trace,tracev = mt
    # mt.λ = 0.
    if epochs <= 0
        push!(mt.modelhistory, deepcopy(model))
        return
    end
    w          = model.w
    data       = todata(mt)
    datat      = todata([testdata])
    mt.normalizer = ones(size(data[1][2], 1),size(data[1][1], 1))
    predfun = model isa ADSystem ? pred : predd
    f(x)   = predfun(w,x,model.nx)
    losses = map(data) do d
        x,y        = d
        lossfun    = loss2(w,x,y,mt)
        # lossfun    = (w...)->cost(w,mt.model.nx, x,y)
        # println("Compiling tape ")
        tape = RDiff.GradientTape(lossfun, w) |> RDiff.compile
        # println(" Done")
        increment!(trace, 1, lossfun(w...))
        increment!(tracev, 1, cost(f,datat))
        lossfun, tape
    end
    g          = similar.(w)
    # epochs > 0 && plot(reuse=false)
    startepoch = last(trace)[1]+1
    @progress for epoch = startepoch:startepoch+epochs-1
        # epoch == 500 && (mt.λ = λ)
        # lossfun = loss(w,x,y,mt)
        for (i,(lossfun,tape)) in enumerate(losses)
            RDiff.gradient!(g,tape,w)
            for (opti,gi) in zip(opt,g)
                opti(gi, epoch)
            end
            increment!(tracev, epoch, cost(f,datat))
            increment!(trace, epoch, lossfun(w...))
            if epoch % 5 == 0
                # mt.normalizer = find_normalizer(w,data[i],mt)
                # if i == length(losses) && Distributed.myid() == 1
                #     cb(model)
                #     # plot(trace, reuse=true)
                #     # plot!(tracev)
                #     # gui()
                      # println("Losses: ", last(trace), last(tracev)[2])
                # end
            end
        end
    end
    push!(mt.modelhistory, deepcopy(model))
    trace, tracev
end

function fit_models(mt::ModelTrainer, useprior)
    map(mt.trajs) do traj
        ltvmodel = KalmanModel(mt, traj, useprior=useprior)
    end
end

function sample_jacprop(mt::ModelTrainer, ltvmodels)
    perturbed_trajs = map(mt.trajs, ltvmodels) do traj, ltvmodel
        @unpack x,u,nx,nu = traj
        xa = x .+ std(x,2)/mt.σdivider .* randn(size(x))
        ua = u .+ std(u,2)/mt.σdivider .* randn(size(u))
        ya = LTVModels.predict(ltvmodel, xa, ua)
        Trajectory(xa,ua,ya)
    end
end

function LTVModels.KalmanModel(mt::ModelTrainer, t::Trajectory, ms=mt.models;
    P=mt.P, useprior=!isempty(mt.modelhistory))
    useprior = useprior && !isempty(mt.modelhistory)
    @unpack x,u,nx,nu = t
    @unpack R1,R2,P0  = mt
    T                 = length(t)
    if useprior
        model    = KalmanModel(zeros(nx,nx,T),zeros(nx,nu,T),zeros(1,1,T),false)
        Jm, Js   = jacobians(ms, t)
        Pt       = cat(3,[diagm(Js[:,i].^2 .+ P) for i=1:T]...) # TODO: magic number
        fx       = cat(3,[reshape(Jm[1:nx^2,i], nx,nx) for i=1:T]...)
        fu       = cat(3,[reshape(Jm[nx^2+1:end,i], nx,nu) for i=1:T]...)
        prior    = KalmanModel(fx,fu,Pt,false)
        ltvmodel = KalmanModel(model, prior, x,u,R1,R2,P0, extend = true, printfit=false)
    else # Don't use prior if it's the first time training
        ltvmodel = KalmanModel(x,u,R1,R2,P0, extend = true, printfit=false)
    end
    ltvmodel
end

todata(trajs::Vector{Trajectory}) = [(traj.xu,traj.y) for traj in trajs]
todata(mt::AbstractModelTrainer) = todata(mt.trajs)
function todata2(trajs::Vector{Trajectory})
    xu = hcat(getfield.(trajs,:xu)...)
    y = hcat(getfield.(trajs,:y)...)
    xu, y
end
todata2(mt::AbstractModelTrainer) = todata2(mt.trajs)

Base.push!(mt::AbstractModelTrainer, t::Trajectory) = push!(mt.trajs, t)
Base.push!(mt::AbstractModelTrainer, data::Matrix...) = push!(mt.trajs, Trajectory(data...))

for T in (ModelTrainer, ADModelTrainer)
    @eval (mt::($T))(; kwargs...) = train!(mt; kwargs...)
    @eval function (mt::($T))(t::Trajectory; kwargs...)
        push!(mt,t)
        train!(mt; kwargs...)
    end
    @eval (mt::($T))(data::Matrix...; kwargs...) = mt(Trajectory(data...); kwargs...)
    @eval Lazy.@forward ($T).models predict, simulate, Flux.jacobian, eigvalplot, jacplot
end

include("plots.jl")

end # module


"""Module containing functions for creating and operating on a `RTree`, to quickly store and locate georeferenced polygons."""
module GeoTrees

using Revise
using ArchGDAL
using SpatialIndexing
using JLD2

const agd = ArchGDAL
const si = SpatialIndexing

#=
POTREBBE ESSERE UTILE AVERE DEI VETTORI DI APPOGGIO CHE CONTENGANO TUTTI I POSSIBILI VALORI DEI VARI ATTRIBUTI DEI POLIGONI,
IN QUESTO VERIFICARE SE UN POLIGONO CON UN CERTO VALORE DI UN CERTO ATTRIBUTO ESISTE DIVENTA TRIVIALE, INVECE DI RICHIEDERE
DI RICERCARE UN POLIGONO CON QUELL'ATTRIBUTO IN TUTTO L'ALBERO.
=#
mutable struct CCSTree
    terrain_types::Dict{Float64, String}
    tree::SpatialIndexing.RTree{Float64, 2}

    CCSTree(tree) = new(Dict{Float64, String}(), tree)
end

"""
    Base.show(io::IO, rtree::SpatialIndexing.RTree{T, N})

Write a text reppresentation of `SpatialIndexing.RTree{T, N}` object `rtree` to the output stream `io`.
"""
function Base.show( io::IO, rtree::SpatialIndexing.RTree{T, N} ) where {T, N}
    println(io, typeof(rtree), ":")
    println(io, "\tPolygons: ", rtree.nelems)
    println(io, "\tLevels: ", rtree.root.level)
    println(io, "\tNodes per level:")
    println(io, "\t", reverse(rtree.nnodes_perlevel) )
    println(io, "\tRoot MBR:", rtree.root.mbr)
end

"""
    Base.show(io::IO, ccstree::CCSTree)

Write a text reppresentation of `CCSTree` object `ccstree` to the output stream `io`.
"""
function Base.show( io::IO, ccstree::CCSTree )
    println(io, typeof(ccstree), ":")
    println(io, "\tTerrain typologies: ", length(ccstree.terrain_types))
    println(io, "\tPolygons: ", ccstree.tree.nelems)
    println(io, "\tLevels: ", ccstree.tree.root.level)
    println(io, "\tNodes per level:")
    println(io, "\t", reverse(ccstree.tree.nnodes_perlevel) )
    println(io, "\tRoot MBR:", ccstree.tree.root.mbr)
end

"""
    SpatialIndexing.check( ccstree::CCSTree )

Check wether a `CCSTree` is well formed.
"""
function check( ccstree::CCSTree )
    return si.check(ccstree.tree)
end

"""
    mbrtype( polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} )

Return the type of the mbr returned by function `mbr()` applied to `polygon`
"""
function mbrtype( polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} )
    return SpatialIndexing.HasMBR{ SpatialIndexing.Rect{Float64, 2} }
end

"""
    mbr( polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} )

Return the minimum bounding rectangle of `polygon` as a `SpatialIndexing.Rect{T, N}`
"""
function mbr( polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} )
    boundingbox = agd.envelope(polygon)
    return SpatialIndexing.Rect{Float64, 2}( (boundingbox.MinX, boundingbox.MinY), (boundingbox.MaxX, boundingbox.MaxY) )
end

"""
    printTree( node, n::Int64=0 )

Given a "SpatialIndexing.RTree" or "Node" either("Spatialindexing.Branch" or "Spatialindexing.Leaf") print the content of the tree with `node` as root.
"""
function printTree( node::Union{SpatialIndexing.RTree, SpatialIndexing.Node}, n::Int64=0 )
    if node isa SpatialIndexing.Leaf
        println("LEVEL: 1")
        println("MBR: $(node.mbr)")
        println("ELEMENTS:")
        for (i, el) in enumerate(node.children)
            println("$n.$i) $(el.id)  -  $(el.val[1])  -  $(el.mbr)")
        end
        println("\n")
    elseif node isa SpatialIndexing.Branch
        println("LEVEL: $(node.level)")
        println("MBR: $(node.mbr)")
        println("CHILDREN:")
        for (i, el) in enumerate(node.children)
            println("$n.$i")
            printTree(el, i)
        end
        println("\n")
    else
        printTree(node.root, 0)
    end
end

"""
    distance( element::Union{SpatialIndexing.SpatialElem, SpatialIndexing.Node}, region::SpatialIndexing.Region )

Compute the distance between `element` and `region`.
"""
distance( element::Union{SpatialIndexing.SpatialElem, SpatialIndexing.Node}, region::SpatialIndexing.Region ) = agd.distance(
    region isa SpatialIndexing.Point ? agd.createpoint(region.coord...) :
        let (xl, yl) = region.low, (xh, yh) = region.high
            agd.createpolygon([(xl, yh), (xh, yh), (xh, yl), (xl, yl), (xl, yh)])
        end,
    element isa SpatialIndexing.SpatialElem ? element.val[2] :
        let (xl, yl) = element.mbr.low, (xh, yh) = element.mbr.high
            agd.createpolygon([(xl, yh), (xh, yh), (xh, yl), (xl, yl), (xl, yh)])
        end
)

"""
    createTree( file_path::String, branch_capacity::Int64=7, leaf_capacity::Int64=7 )

Create a `SpatialIndexing.RTree{Float64, 2}` filled with the polygons contained in a shapefile at `file_path`.
"""
function createRTree( file_path::String, branch_capacity::Int64=7, leaf_capacity::Int64=7 )
 # Obtain all the "features" ("polygons" + additional information) of the file
    features = collect(agd.getlayer(agd.read(file_path), 0))
 # Create the RTree (It contains an empty vector as the "terrain_types" field and the empty tree that we just defined as the field "tree")
    tree = RTree{Float64, 2}(Float64, NamedTuple, branch_capacity=branch_capacity, leaf_capacity=leaf_capacity)
 # Insert all of the relevant data of the features in the tree as nodes
    for feature in features
     # Obtain the geometry (polygon) of the feature
        geom = agd.getgeom(feature, 0)
     # Find the MBR of the polygon
        feature_mbr = mbr(geom)
     # Check if the MBR has a valid value
        if si.isvalid(feature_mbr)
         # If so insert a new node in the tree
            si.insert!( tree, feature_mbr, agd.getfield(feature, :objectid), ( code=agd.getfield(feature, :codice_num), geometry=geom ) )
        else
            println("Omitted feature $(agd.getfield(feature, :objectid)), the feature mbr is not valid.")
        end
    end
 # Check that the tree is well-build
    return si.check(tree) ? tree : throw(DomainError("Error in the creation of the tree, check input parameters"))
end

"""
    createCCSTree( file_path::String, branch_capacity::Int64=7, leaf_capacity::Int64=7 )

Create a CCSTree based on a vectorial file at `file_path` containing a series of polygons decribing the terrain composition of an area.
The  created object will contain a field `terrain_types` containing a `Vector{Tuple{Int64, String}}` holding the codes and descriptions of
all the possible kinds of terrain found in the vectorial file.
The "tree" field of the object will instead contain a `SpatialIndexing.RTree{Float64, 2}` with all the polygons of the file.
"""
function createCCSTree( file_path::String, branch_capacity::Int64=7, leaf_capacity::Int64=7 )
 # Obtain all the "features" ("polygons" + additional information) of the file
    features = collect(agd.getlayer(agd.read(file_path), 0))
 # Create the CCSTree from an empty RTree
    ccstree = CCSTree( RTree{Float64, 2}(Float64, NamedTuple, branch_capacity=branch_capacity, leaf_capacity=leaf_capacity) )
  # Insert all of the relevant data of all the features in the tree as nodes
    for feature in features
     # Obtain the geometry (polygon) of the feature
        geom = agd.getgeom(feature, 0)
     # Find the MBR of the polygon
        feature_mbr = mbr(geom)
     # Check if the MBR has a valid value
        if si.isvalid(feature_mbr)
         # Get the code that identifies the kind of landcover of the polygon
            num_code = agd.getfield(feature, :codice_num)
         # Insert the new node corresponding to the feature in the "tree" field of the CCSTree
            si.insert!( ccstree.tree, feature_mbr, agd.getfield(feature, :objectid), ( code=num_code, geometry=geom ) )
         # Insert the landcover type in the dictionary in the field "terrain_types" of the CCSTree along with its description
            push!( ccstree.terrain_types, num_code => agd.getfield(feature, :legenda) )
        else
            println("Omitted feature $(agd.getfield(feature, :objectid)), the feature mbr is not valid.")
        end
    end
  # Check that the tree is well-build
    return check(ccstree) ? ccstree : throw(DomainError("Error in the creation of the tree, check input parameters"))
end

"""
    exists_terrain_type( tree::CCSTree, type_code::Int64 )

Check if the type of terrain corresponding to `type_code` exists in at least one polygon of `tree`.
"""
exists_terrain_type( tree::CCSTree, type_code::Int64 ) = type_code in keys(tree.terrain_types)

"""
    exists_terrain_type( tree::CCSTree, type_description::String )

Check if the type of terrain described by `type_description` exists in at least one polygon of `tree`.
Note: `type_description` must match exactly the description of the terrain type, if the complete description is not
known, consider using function `find_terrain_type_similar` instead.
"""
exists_terrain_type( tree::CCSTree, type_description::String ) = type_description in values(tree.terrain_types)

"""
    find_terrain_type_similar( tree::CCSTree, type::String )

Return a `Dict{Int64, String}` containing all the codes and the respective descriptions in which the latter contains `type` as a substring.
"""
find_terrain_type_similar( tree::CCSTree, type::String ) = filter( (k, v) -> contains(v, type), tree.terrain_types )

# --------------------------------------------------------------- POLYGON SEARCH ---------------------------------------------------------------------------------------
"""
    findPolygons( node::SpatialIndexing.Leaf{T,N}, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} )::Vector{SpatialIndexing.SpatialElem} where {T, N}

Find and return a `Vector{SpatialIndexing.SpatialElem}` that contains all the `SpatialIndexing.SpatialElem`s children of `node` that intersect, contain or are containded in `polygon`.
"""
function findPolygons( node::SpatialIndexing.Leaf{T,N}, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} )::Vector{SpatialIndexing.SpatialElem} where {T, N}
    # Return a Vector of all the children of the node that contain, are contained in, or intersect with `polygon`
    return [
        child
        for child in node.children
        if agd.contains(polygon, child.val.geometry) || agd.contains(child.val.geometry, polygon) || agd.intersects(polygon, child.val.geometry)
    ]
end

"""
    findPolygons( node::SpatialIndexing.Branch{T,N,V}, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} ) where {T, N, V}

Find and return a `Vector{SpatialIndexing.SpatialElem}` that contains all the elements of the subtree rooted in `node` that intersect, contain or are contained in `polygon`.
"""
function findPolygons( node::SpatialIndexing.Branch{T,N,V}, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} ) where {T, N, V}
    res = Vector{SpatialIndexing.SpatialElem}()
    for child in node.children
        poly_mbr = mbr(polygon)
        if si.intersects(child.mbr, poly_mbr) || si.contains(child.mbr, poly_mbr) || si.contains(poly_mbr, child.mbr)
            results = findPolygon(child, polygon)
            if !isempty(results)
                append!(res, results)
            end
        end
    end
    return res
end

"""
    findPolygons( tree::SpatialIndexing.RTree{T,N}, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} ) where {T, N}

Return all the polygons of `tree` intersected by, contained in or containing `polygon`, or the polygon that conains it, or an empty `Vector{SpatialIndexing.SpatialElem}`, if there is no such polygon.
"""
function findPolygons( tree::SpatialIndexing.RTree{T,N}, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} ) where {T, N}
    return findPolygon(tree.root, polygon)
end

"""
    findPolygons( tree::CCSTree, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} ) where {T, N}

Return all the polygons of `tree` intersected by, contained in or containing `polygon`, or the polygon that conains it, or an empty `Vector{SpatialIndexing.SpatialElem}`, if there is no such polygon.
"""
function findPolygons( tree::CCSTree, polygon::ArchGDAL.IGeometry{ArchGDAL.wkbPolygon} ) where {T, N}
    return findPolygon(tree.tree.root, polygon)
end

# ----------------------------------------------------------------- KNN SEARCH -----------------------------------------------------------------------------------------
"""
    findKNN( node::SpatialIndexing.Leaf{T,N}, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N}

Return the `k` `SpatialIndexing.SpatialElem`s children of `node` that are closest to `point` (at most k elements will be returned).
"""
function findKNN( node::SpatialIndexing.Leaf{T,N}, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N}
    agd_point = agd.createpoint(point.coord...)
    results = sort!( # 2) Sort the children in increasing order of distance from `point`
        map( # 1) Map every children (they will all be `SpatialElements`) of the node to the distance between its centroid and `point`
            element -> (
                element,
                agd.distance( # Compute distance between the centroid of the polygon and `point`
                    agd.centroid(element.val.geometry),
                    agd_point
                )
            ),
            node.children
        ),
        lt=(x, y) -> x[2] < y[2]
    )
    # 3) Return the, at most, `k` elements that are closer to `point`
    return results[1:min(k, length(results))]
end

"""
    findKNN( node::SpatialIndexing.Branch{T,N,V}, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N, V}

Return the `k` `SpatialIndexing.SpatialElem`s of the subtree rooted in `node` that are closest to `point` (at most k elements will be returned).
"""
function findKNN( node::SpatialIndexing.Branch{T,N,V}, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N, V}
    # 1) Order the children of the current node by increasing distance from point
    candidates = sort( node.children, lt=(x, y) -> distance(x, point) < distance(y, point) )
    results = sort!( # 5) Sort the resulting vector by increasing distance from `point`
        reduce( # 4) For each child of the current node a Vector of the results will be returned, join the vectors in a single one
            vcat,
            findKNN.( # 3) For each of the children (they could be `Leaf`s or other `Branch`s) find its `k` children closest to `point`
                candidates[1:min(k, length(candidates))], # 2) Take, at most, the first `k` children (if there are less than `k` children, take all of them)
                Ref(point),
                k
            )
        ),
        lt=(x, y) -> x[2] < y[2]
    )
    # 6) Take, at most, the first `k` results (if there are less than `k` children, take all of them)
    return results[1:min(k, length(results))]
end

"""
    findKNN( tree::SpatialIndexing.RTree{T,N}, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N}

Return the `k` `SpatialIndexing.SpatialElem`s contained in `tree` that are closest to `point` (at most k elements will be returned).
"""
function findKNN( tree::SpatialIndexing.RTree{T,N}, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N}
    return findKNN(tree.root, point, k)
end

"""
    findKNN( tree::CCSTree, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N}

Return the `k` `SpatialIndexing.SpatialElem`s contained in `tree` that are closest to `point` (at most k elements will be returned).
"""
function findKNN( tree::CCSTree, point::SpatialIndexing.Point{T,N}, k::Int64 ) where {T, N}
    return findKNN(tree.tree.root, point, k)
end

# ------------------------------------------------------------- DISTANCE BASED SEARCH ----------------------------------------------------------------------------------
"""
    findPolygons( node::SpatialIndexing.Leaf{T,N}, point::SpatialIndexing.Point{T,N} distance::Float64 ) where {T, N}

Return all the `SpatialIndexing.SpatialElem`s children of `node` that are within `distance` from point.
"""
function findPolygons( node::SpatialIndexing.Leaf{T,N}, point::SpatialIndexing.Point{T,N}, distance::Float64 ) where {T, N}
    agd_point = agd.createpoint(point.coord...)
    results = sort!( # 2) Sort the children in increasing order of distance from `point`
        map( # 1) Map every children (they will all be `SpatialElements`) of the node to the distance between its centroid and `point`
            element -> (
                element,
                agd.distance( # Compute distance between the centroid of the polygon and `point`
                    agd.centroid(element.val.geometry),
                    agd_point
                )
            ),
            node.children
        ),
        lt=(x, y) -> x[2] < y[2]
    )
    # 3) Return all the children within `distance` from `point`
    return results[1:something( findfirst(res -> res[2] > distance, result), length(result) )]
end

"""
    findPolygons( node::SpatialIndexing.Branch{T,N,V}, point::SpatialIndexing.Point{T,N} distance::Float64 ) where {T, N, V}

Return all the `SpatialIndexing.SpatialElem`s of the subtree rooted in `node` that are within `distance` from `point`.
"""
function findPolygons( node::SpatialIndexing.Branch{T,N,V}, point::SpatialIndexing.Point{T,N}, distance::Float64 ) where {T, N, V}
    results = sort!( # 5) Sort the results by distance from `point`
        reduce( # 4) For each child of the current node a Vector of the results will be returned, join the vectors in a single one
            vcat,
            findPolygons.( # 3) For each of the children (they could be `Leaf`s or other `Branch`s) find all its children whithin `distance` from `point`
                sort( # 1) Sort the children by increasing distance from `point`
                    node.children,
                    lt=(x, y) -> distance(x, point) < distance(y, point)
                )[1:something( findfirst(res -> res[2] > distance, result), length(result) )], # 2) All children up to the first farther from `point` than `distance`
                Ref(point),
                distance
            )
        ),
        lt=(x, y) -> x[2] < y[2]
    )
    # 6) Return all the results up to the first farther than `distance` from `point`
    return results[1:something( findfirst(res -> res[2] > distance, result), length(result) )]
end

"""
    findPolygons( tree::SpatialIndexing.RTree{T,N}, point::SpatialIndexing.Point{T,N}, distance::Float64 ) where {T, N}

Return all the `SpatialIndexing.SpatialElem`s contained in `tree` that are within `distance` from `point`.
"""
function findPolygons( tree::SpatialIndexing.RTree{T,N}, point::SpatialIndexing.Point{T,N}, distance::Float64 ) where {T, N}
    return findPolygons(tree.root, point, distance)
end

"""
    findPolygons( tree::CCSTree, point::SpatialIndexing.Point{T,N}, distance::Float64 ) where {T, N}

Return all the `SpatialIndexing.SpatialElem`s contained in `tree` that are within `distance` from `point`.
"""
function findPolygons( tree::CCSTree, point::SpatialIndexing.Point{T,N}, distance::Float64 ) where {T, N}
    return findPolygons(tree.tree.root, point, distance)
end

# ----------------------------------------------------------- LOADING AND SAVING -----------------------------------------------------------------------------
"""
    saveTree( tree::Union{SpatialIndexing.RTree{T, N}, CCSTree}, output_file::AbstractString=".\\tree.jld2" ) where {T, N}

Save `tree` as a file in the directory and with the name specified by `output_file` in Julia Data format 2 (JLD2).
"""
function saveTree( tree::Union{SpatialIndexing.RTree{T, N}, CCSTree}, output_file::AbstractString=".\\tree.jld2" ) where {T, N}
    if !occursin(output_file, ".jld2")
        output_file *= ".jld2"
    end
    save_object(tree, output_file)
end

"""
    loadTree( tree_file::AbstractString )

Load a `SpatialIndexing.RTree{T, N}` or a `CCSTree` saved as `tree_file`.
"""
function loadTree( tree_file::AbstractString )
    return load_object(tree_file)
end

#=
# Serve un modo per simulare la possibilità di aggiungere campi in modo da poter creare strutture che si adattino al numero di
# attributi rilevanti di cui si vuole ottenere la lista degli unici.
mutable struct Fields
    fields::Dict{Symbol, Any}

    Fields(fields) = new(fields)
end

mutable struct Cesto
    frutti::Fields

    function Cesto( frutti..., numeri... )
        if length(frutti) != length(numeri)
            throw(DomainError("`frutti` e `numeri` per ogni frutto deve esserci un numero"))
        end
        for (f, n) in zip(frutti, numeri)
            dict = Dict( Symbol(f) => n )
        end
    end
end
=#

using ArchGDAL

const agd = ArchGDAL

ccs_file = "C:\\Users\\Lenovo\\Documents\\GitHub\\Tirocinio\\resources\\Analysis data\\ccs\\ccs.shp"
ccs = agd.read(ccs_file)
features =  collect(agd.getlayer(ccs, 0))
features[1]

ccstree = createCCSTree("C:\\Users\\Lenovo\\Documents\\GitHub\\Tirocinio\\resources\\Analysis data\\ccs\\ccs.shp")

rtree = createRTree("C:\\Users\\Lenovo\\Documents\\GitHub\\Tirocinio\\resources\\Analysis data\\ccs\\ccs.shp")

#= ---------------------------- TREE TESTING [V] --------------------------------

#   ccs_shp = agd.read("C:\\Users\\Lenovo\\Documents\\GitHub\\Tirocinio\\resources\\Analysis data\\ccs\\ccs.shp")
ccs_shp = agd.read("D:\\Z_Tirocinio_Dati\\ccs WGS84\\ccs.shp")
features =  collect(agd.getlayer(ccs_shp, 0))

tree = nothing
GC.gc()
tree = RTree{Float64, 2}(Float64, NamedTuple, branch_capacity=7, leaf_capacity=7)
for feature in features
    geom = agd.getgeom(feature, 0)
    feature_mbr = mbr(geom)
    if si.isvalid(feature_mbr)
        si.insert!( tree, feature_mbr, agd.getfield(feature, :objectid), ( code=agd.getfield(feature, :codice_num), geometry=geom ) )
    end
end
si.check(tree)
tree.nnodes_perlevel

=#
#= ---------------------------- QUERY TESTING [V] --------------------------------

# Multilinea di contorno al poligono
geom = agd.getgeom( agd.getgeom(features[311113]), 0 )
# NUmber of vertexes of the poligon
num_points = agd.ngeom(geom)
# Coordinates of points inside the polygon:
#   Average of each coordinate of the vertexes
xt1 = sum( j -> agd.getx(geom, j), 0:num_points-1 ) / Float64(num_points)
yt1 = sum( j -> agd.gety(geom, j), 0:num_points-1 ) / Float64(num_points)
#   Average of 3 vertexes
xt2 = sum( j -> agd.getx(geom, j), 0:2 ) / 3.0
yt2 = sum( j -> agd.gety(geom, j), 0:2 ) / 3.0
#   In the middle of an edge (mean of 2 vertexes)
xt3 = sum( j -> agd.getx(geom, j), 0:1 ) / 2.0
yt3 = sum( j -> agd.gety(geom, j), 0:1 ) / 2.0
# In the first half of the polygon
xt4 = sum( j -> agd.getx(geom, j), 0:(num_points÷2)-1 ) / (num_points÷2)
yt4 = sum( j -> agd.gety(geom, j), 0:(num_points÷2)-1 ) / (num_points÷2)

#= Results for each ponint
    res = findPolygon(tree, si.Point((xt1, yt1)))
    res = findPolygon(tree, si.Point((xt2, yt2)))
    res = findPolygon(tree, si.Point((xt3, yt3)))
    res = findPolygon(tree, si.Point((xt4, yt4)))
=#

res1 = knn( tree, si.Point((xt1, yt1)), 3 )
res2 = knn( tree, si.Point((xt2, yt2)), 3 )
res3 = knn( tree, si.Point((xt3, yt3)), 3 )
res4 = knn( tree, si.Point((xt4, yt4)), 3 )

sat_shp = agd.read(split( @__DIR__ , "\\Porting\\")[1] * "\\Mappe\\sat WGS84\\sette_sorelle.shp")
sat = agd.getgeom(collect(agd.getlayer(sat_shp, 0))[1])
res = findPolygon(tree, sat)
# Get the result in decreasing order based on the surface of the polygon
sort!( res, lt=(x, y) -> agd.geomarea(x.val.geometry) < agd.geomarea(y.val.geometry), rev=true )

# Id di tutti i poligoni intersecati dal poligono "sette_sorelle" (ottenuti con QGIS)
objectids = [
    192271, 192307, 197591, 214336, 214356, 214422,
    214518, 214606, 214777, 214840, 214967, 214999,
    215264, 215314, 215391, 215555, 215918, 215944,
    221275, 221300, 221335, 221415, 221641, 222633,
    223271, 231330, 232459, 232478, 232492, 232541,
    240004, 240325, 240533, 240803, 240903, 241012,
    241266, 241369, 241738, 249519, 249534, 249542,
    249588, 249619, 249622, 252658, 252687, 252737,
    252754, 252931, 258981, 261396, 261404, 261417,
    262501, 263548, 263567, 263578, 263584, 263602,
    263625, 265696, 265697, 265699, 265700, 265701,
    265702, 265703, 265704, 265705, 265706, 265708,
    265709, 265711, 265712, 265714, 265715, 267016,
    267021, 267029, 267043, 267064, 395922, 396242,
    396281
]

# Check that all the returned polygons are correct
all(r -> r.id in objectids, res)
resids = [ Int64(r.id) for r in res ]
# Check whether there are any missing polygon
all(id -> id in resids, objectids)

#=  PER SALVARE GLI ID TROVATI
    path = "C:\\Users\\DAVIDE-FAVARO\\Desktop\\Dati\\res_ids.txt"
    path = "D:\\Documents and Settings\\DAVIDE-FAVARO\\My Documents\\GitHub\\Tirocinio\\Mappe\\polygon_ids.txt"
    open( path, "w" ) do io
        for id in objectids
            write(io, "$id\n")
        end
    end
=#

=#
#= ---------------------------- SAVE TESTING [V] --------------------------------

path = "D:\\Z_Tirocinio_Dati\\tree_test.jld2"
save_object(path, tree)
loaded_tree = load_object(path);
si.check(loaded_tree)
l_res = findPolygon(loaded_tree, si.Point((xt1, yt1)))
l_res = findPolygon(loaded_tree, si.Point((xt2, yt2)))
l_res = findPolygon(loaded_tree, si.Point((xt3, yt3)))

=#
# ---------------------------- PLOT TESTING [V] --------------------------------
# PRESO DA "plot_utils.jl" CONTENUTO IN "https://github.com/alyst/SpatialIndexing.jl/tree/master/examples"

using PlotlyJS
using Printf: @sprintf

const SI = SpatialIndexing

# webcolor palette for R-tree levels
const LevelsPalette = Dict(
    1 => "#228B22", # Forest Green
    2 => "#DAA520", # Goldenrod
    3 => "#FF6347", # Tomato
    4 => "#B22222", # Firebrick
    5 => "#4B0082", # Indigo
    6 => "#800080", # Purple
    7 => "#4169E1", # Royal blue
    8 => "#008080", # Teal
    9 => "#00FFFF", # Cyan
)

# create plotly trace for a single R-tree `node` (rectangle edges)
function node_trace(node::SI.Node{<:Any,2}, ix::Union{Integer, Nothing};
                    showlegend::Bool)
    nbr = SI.mbr(node)
    res = scatter(x = [nbr.low[1], nbr.high[1], nbr.high[1], nbr.low[1], nbr.low[1]],
                  y = [nbr.low[2], nbr.low[2], nbr.high[2], nbr.high[2], nbr.low[2]],
                  mode=:lines, line_color=get(LevelsPalette, SI.level(node), "#708090"),
                  line_width=SI.level(node),
                  hoverinfo=:text, hoveron=:points,
                  text="lev=$(SI.level(node)) ix=$(ix !== nothing ? ix : "none")<br>nchildren=$(length(node)) nelems=$(SI.nelements(node)) area=$(@sprintf("%.2f", SI.area(SI.mbr(node))))",
                  name="level $(SI.level(node))",
                  legendgroup="level $(SI.level(node))", showlegend=showlegend)
    return res
end

# create plotly trace for a single R-tree `node` (cube edges)
function node_trace(node::SI.Node, ix::Union{Integer, Nothing};
                    showlegend::Bool)
    nbr = SI.mbr(node)
    res = scatter3d(x = [nbr.low[1],  nbr.high[1], nbr.high[1], nbr.low[1], nbr.low[1],
                         nbr.low[1],  nbr.high[1], nbr.high[1], nbr.low[1], nbr.low[1],
                         NaN, nbr.low[1], nbr.low[1],
                         NaN, nbr.high[1], nbr.high[1],
                         NaN, nbr.high[1], nbr.high[1]],
                    y = [nbr.low[2],  nbr.low[2],  nbr.high[2], nbr.high[2], nbr.low[2],
                         nbr.low[2],  nbr.low[2],  nbr.high[2], nbr.high[2], nbr.low[2],
                         NaN, nbr.high[2], nbr.high[2],
                         NaN, nbr.high[2], nbr.high[2],
                         NaN, nbr.low[2], nbr.low[2]],
                    z = [nbr.low[3],  nbr.low[3],  nbr.low[3],  nbr.low[3], nbr.low[3],
                         nbr.high[3], nbr.high[3], nbr.high[3], nbr.high[3], nbr.high[3],
                         NaN, nbr.low[3], nbr.high[3],
                         NaN, nbr.low[3], nbr.high[3],
                         NaN, nbr.low[3], nbr.high[3]],
                  mode=:lines, line_color=get(LevelsPalette, SI.level(node), "#708090"),
                  line_width=SI.level(node),
                  hoverinfo=:text, hoveron=:points,
                  text="lev=$(SI.level(node)) ix=$(ix !== nothing ? ix : "none")<br>nchildren=$(length(node)) nelems=$(SI.nelements(node)) area=$(@sprintf("%.2f", SI.area(SI.mbr(node))))",
                  name="level $(SI.level(node))",
                  legendgroup="level $(SI.level(node))", showlegend=showlegend)
    return res
end

# create plotly traces for the nodes in a subtree with the `node` root
# and append them to `node_traces`
function append_subtree_traces!(node_traces::Vector{PlotlyBase.AbstractTrace},
                                node::SI.Node, ix::Union{Integer, Nothing},
                                levels::Set{Int})
    push!(node_traces, node_trace(node, ix, showlegend = SI.level(node) ∉ levels))
    push!(levels, SI.level(node)) # show each level once in legend
    if node isa SI.Branch
        for (i, child) in enumerate(SI.children(node))
            append_subtree_traces!(node_traces, child, i, levels)
        end
    end
    return nothing
end

data_trace(tree::RTree{<:Any, 2}) =
    scatter(mode=:markers, name=:data, marker_color = "#333333", marker_size = 2,
            x=[SI.center(SI.mbr(x)).coord[1] for x in tree],
            y=[SI.center(SI.mbr(x)).coord[2] for x in tree],
            text=["id=$(SI.id(x))" for x in tree], hoverinfo=:text)

data_trace(tree::RTree) =
    scatter3d(mode=:markers, name=:data, marker_color = "#333333", marker_size = 2,
              x=[SI.center(SI.mbr(x)).coord[1] for x in tree],
              y=[SI.center(SI.mbr(x)).coord[2] for x in tree],
              z=[SI.center(SI.mbr(x)).coord[3] for x in tree],
              text=["id=$(SI.id(x))" for x in tree], hoverinfo=:text)

# create Plotly plot of the given tree
function PlotlyJS.plot(tree::RTree)
    ndims(tree) == 1 && throw(ArgumentError("1-D R-trees not supported"))
    ndims(tree) > 3 && @warn("Only 1st-3rd dimensions would be shown for $(ndims(tree))-D trees")
    node_traces = Vector{PlotlyBase.AbstractTrace}()
    append_subtree_traces!(node_traces, tree.root, nothing, Set{Int}())
    PlotlyJS.plot([node_traces; [data_trace(tree)]], Layout(hovermode=:closest))
end

# MOLTO PESANTE A LIVELLO DI PRESTAZIONI
tree_plot = plot(tree);
open( "C:\\Users\\DAVIDE-FAVARO\\Desktop\\Dati\\tree_plot.png", "w" ) do io
    PlotlyJS.savefig(io, tree_plot, width=900, height=1000)
end

end # module


"""
    ESS

Elliptical slice sampling algorithm.

# Examples
```jldoctest; setup = :(Random.seed!(1))
julia> @model gdemo(x) = begin
           m ~ Normal()
           x ~ Normal(m, 0.5)
       end
gdemo (generic function with 2 methods)

julia> sample(gdemo(1.0), ESS(), 1_000) |> mean
Mean

│ Row │ parameters │ mean     │
│     │ Symbol     │ Float64  │
├─────┼────────────┼──────────┤
│ 1   │ m          │ 0.824853 │
```
"""
struct ESS{space} <: InferenceAlgorithm end

ESS() = ESS{()}()
ESS(space::Symbol) = ESS{(space,)}()

mutable struct ESSState{V<:VarInfo} <: AbstractSamplerState
    vi::V
end

function Sampler(alg::ESS, model::Model, s::Selector)
    # sanity check
    vi = VarInfo(model)
    space = getspace(alg)
    vns = _getvns(vi, s, Val(space))
    length(vns) == 1 ||
        error("[ESS] does only support one variable ($(length(vns)) variables specified)")
    for vn in vns[1]
        dist = getdist(vi, vn)
        EllipticalSliceSampling.isgaussian(typeof(dist)) ||
            error("[ESS] only supports Gaussian prior distributions")
    end

    state = ESSState(vi)
    info = Dict{Symbol, Any}()

    return Sampler(alg, info, s, state)
end

# always accept in the first step
function AbstractMCMC.step!(
    ::AbstractRNG,
    model::Model,
    spl::Sampler{<:ESS},
    ::Integer,
    ::Nothing;
    kwargs...
)
    return Transition(spl)
end

function AbstractMCMC.step!(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:ESS},
    ::Integer,
    transition;
    kwargs...
)
    # obtain random variable that should be sampled
    vi = spl.state.vi

    # obtain previous sample
    f = vi[spl]

    # recompute log-likelihood in logp
    if spl.selector.tag !== :default
        runmodel!(model, vi, spl)
    end

    # define previous sampler state
    oldstate = EllipticalSliceSampling.EllipticalSliceSamplerState(f, getlogp(vi))

    # compute next state
    state = AbstractMCMC.step!(rng, ESSModel(model, spl),
                               EllipticalSliceSampling.EllipticalSliceSampler(), 1, oldstate;
                               progress = false)

    # update sample and log-likelihood
    vi[spl] = state.sample
    setlogp!(vi, state.loglikelihood)

    return Transition(spl)
end

struct ESSModel{M<:Model,S<:Sampler{<:ESS},T} <: AbstractMCMC.AbstractModel
    model::M
    spl::S
    μ::T
end

function ESSModel(model::Model, spl::Sampler{<:ESS})
    vi = spl.state.vi
    vns = _getvns(vi, spl)
    μ = mapreduce(vcat, vns[1]) do vn
        dist = getdist(vi, vn)
        vectorize(dist, mean(dist))
    end

    ESSModel(model, spl, μ)
end

# sample from the prior
function EllipticalSliceSampling.sample_prior(rng::Random.AbstractRNG, model::ESSModel)
    spl = model.spl
    vi = spl.state.vi
    vns = _getvns(vi, spl)
    set_flag!(vi, vns[1][1], "del")
    runmodel!(model.model, vi, spl)
    return vi[spl]
end

# compute proposal and apply correction for distributions with nonzero mean
function EllipticalSliceSampling.proposal(model::ESSModel, f, ν, θ)
    sinθ, cosθ = sincos(θ)
    a = 1 - (sinθ + cosθ)
    return @. cosθ * f + sinθ * ν + a * model.μ
end

function EllipticalSliceSampling.proposal!(out, model::ESSModel, f, ν, θ)
    sinθ, cosθ = sincos(θ)
    a = 1 - (sinθ + cosθ)
    @. out = cosθ * f + sinθ * ν + a * model.μ
    return out
end

# evaluate log-likelihood
function Distributions.loglikelihood(model::ESSModel, f)
    spl = model.spl
    vi = spl.state.vi
    vi[spl] = f
    runmodel!(model.model, vi, spl)
    getlogp(vi)
end

function DynamicPPL.tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, vn::VarName, inds, vi)
    if vn in getspace(sampler)
        return DynamicPPL.tilde(LikelihoodContext(), SampleFromPrior(), right, vn, inds, vi)
    else
        return DynamicPPL.tilde(ctx, SampleFromPrior(), right, vn, inds, vi)
    end
end

function DynamicPPL.tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vi)
    return DynamicPPL.tilde(ctx, SampleFromPrior(), right, left, vi)
end

function DynamicPPL.dot_tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vn::VarName, inds, vi)
    if vn in getspace(sampler)
        return DynamicPPL.dot_tilde(LikelihoodContext(), SampleFromPrior(), right, left, vn, inds, vi)
    else
        return DynamicPPL.dot_tilde(ctx, SampleFromPrior(), right, left, vn, inds, vi)
    end
end

function DynamicPPL.dot_tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vi)
    return DynamicPPL.dot_tilde(ctx, SampleFromPrior(), right, left, vi)
end


#####
##### Fields
#####

##### Auxiliary fields

# Face & Center
aux_vars_ref_state(FT) = (; ref_state = (ρ0 = FT(0), α0 = FT(0), p0 = FT(0)))

# Center only
cent_aux_vars_gm(FT) = (;
    tke = FT(0),
    Hvar = FT(0),
    QTvar = FT(0),
    HQTcov = FT(0),
    q_liq = FT(0),
    q_ice = FT(0),
    RH = FT(0),
    s = FT(0),
    T = FT(0),
    buoy = FT(0),
    cloud_fraction = FT(0),
    H_third_m = FT(0),
    W_third_m = FT(0),
    QT_third_m = FT(0),
    # From RadiationBase
    dTdt_rad = FT(0), # horizontal advection temperature tendency
    dqtdt_rad = FT(0), # horizontal advection moisture tendency
    # From ForcingBase
    subsidence = FT(0), #Large-scale subsidence
    dTdt = FT(0), #Large-scale temperature tendency
    dqtdt = FT(0), #Large-scale moisture tendency
    dTdt_hadv = FT(0), #Horizontal advection of temperature
    H_nudge = FT(0), #Reference H profile for relaxation tendency
    dTdt_fluc = FT(0), #Vertical turbulent advection of temperature
    dqtdt_hadv = FT(0), #Horizontal advection of moisture
    qt_nudge = FT(0), #Reference qt profile for relaxation tendency
    dqtdt_fluc = FT(0), #Vertical turbulent advection of moisture
    u_nudge = FT(0), #Reference u profile for relaxation tendency
    v_nudge = FT(0), #Reference v profile for relaxation tendency
    ug = FT(0), #Geostrophic u velocity
    vg = FT(0), #Geostrophic v velocity
    ∇θ_liq_ice_gm = FT(0),
    ∇q_tot_gm = FT(0),
)
cent_aux_vars(FT, n_up) = (; aux_vars_ref_state(FT)..., cent_aux_vars_gm(FT)..., TC.cent_aux_vars_edmf(FT, n_up)...)

# Face only
face_aux_vars_gm(FT) = (; massflux_s = FT(0), diffusive_flux_s = FT(0), total_flux_s = FT(0), f_rad = FT(0))
face_aux_vars(FT, n_up) = (; aux_vars_ref_state(FT)..., face_aux_vars_gm(FT)..., TC.face_aux_vars_edmf(FT, n_up)...)

##### Diagnostic fields

# Center only
cent_diagnostic_vars_gm(FT) = ()
cent_diagnostic_vars(FT, n_up) = (; cent_diagnostic_vars_gm(FT)..., TC.cent_diagnostic_vars_edmf(FT, n_up)...)

# Face only
face_diagnostic_vars_gm(FT) = ()
face_diagnostic_vars(FT, n_up) = (; face_diagnostic_vars_gm(FT)..., TC.face_diagnostic_vars_edmf(FT, n_up)...)

##### Prognostic fields

# Center only
cent_prognostic_vars(FT, n_up) = (; cent_prognostic_vars_gm(FT)..., TC.cent_prognostic_vars_edmf(FT, n_up)...)
cent_prognostic_vars_gm(FT) = (; u = FT(0), v = FT(0), θ_liq_ice = FT(0), q_tot = FT(0))

# Face only
face_prognostic_vars(FT, n_up) = (; w = FT(0), TC.face_prognostic_vars_edmf(FT, n_up)...)
# TC.face_prognostic_vars_edmf(FT, n_up) = (;) # could also use this for empty model

#####
##### Methods
#####

function satadjust(gm::TC.GridMeanVariables, grid, state)
    p0_c = TC.center_ref_state(state).p0
    ρ0_c = TC.center_ref_state(state).ρ0
    aux_gm = TC.center_aux_grid_mean(state)
    prog_gm = TC.center_prog_grid_mean(state)
    param_set = TC.parameter_set(gm)
    @inbounds for k in TC.real_center_indices(grid)
        θ_liq_ice = prog_gm.θ_liq_ice[k]
        q_tot = prog_gm.q_tot[k]
        ts = TC.thermo_state_pθq(param_set, p0_c[k], θ_liq_ice, q_tot)
        aux_gm.q_liq[k] = TD.liquid_specific_humidity(ts)
        aux_gm.q_ice[k] = TD.ice_specific_humidity(ts)
        aux_gm.T[k] = TD.air_temperature(ts)
        ρ = TD.air_density(ts)
        aux_gm.buoy[k] = TC.buoyancy_c(param_set, ρ0_c[k], ρ)
        aux_gm.RH[k] = TD.relative_humidity(ts)
    end
    return
end


import Base

"""
Struct for WilsonFermion
"""
struct WilsonFermion_2D_wing{NC,NDW} <: WilsonFermion_2D{NC} #AbstractFermionfields_2D{NC}

    f::Array{ComplexF64,4}
    NC::Int64
    NX::Int64
    #NY::Int64
    #NZ::Int64
    NT::Int64
    NG::Int64
    NDW::Int64
    Dirac_operator::String
    #BoundaryCondition::Vector{Int8}

    function WilsonFermion_2D_wing(NC::T,NX::T,NT::T) where T<: Integer
        NG = 2
        NDW = 1
        #@assert NDW == 1 "only NDW = 1 is supported. Now NDW = $NDW"
        f = zeros(ComplexF64,NC,NX+2NDW,NT+2NDW,NG)
        Dirac_operator = "Wilson"
        return new{NC,NDW}(f,NC,NX,NT,NG,NDW,Dirac_operator)
    end

    function WilsonFermion_2D_wing{NC}(NX::T,NT::T) where {T<: Integer,NC}
        NG = 2
        NDW = 1
        #@assert NDW == 1 "only NDW = 1 is supported. Now NDW = $NDW"
        f = zeros(ComplexF64,NC,NX+2NDW,NT+2NDW,NG)
        Dirac_operator = "Wilson"
        return new{NC,NDW}(f,NC,NX,NT,NG,NDW,Dirac_operator)
    end

end

function Base.length(x::WilsonFermion_2D_wing{NC,NDW}) where {NC,NDW}
    return NC*x.NX*x.NT*x.NG
end

function Base.setindex!(x::WilsonFermion_2D_wing{NC,NDW},v,i1,i2,i5,i6)  where {NC,NDW}
    @inbounds x.f[i1,i2 + NDW,i5 + NDW,i6] = v
end

function Base.getindex(x::WilsonFermion_2D_wing{NC,NDW},i1,i2,i5,i6) where {NC,NDW}
    @inbounds return x.f[i1,i2 .+ NDW,i5 .+ NDW,i6]
end

function Base.getindex(x::WilsonFermion_2D_wing{NC,NDW},i1::N,i2::N,i5::N,i6::N) where {NC,NDW,N<: Integer}
    @inbounds return  x.f[i1,i2 + NDW,i5 + NDW,i6]
end

function Base.similar(x::T) where T <: WilsonFermion_2D_wing
    return WilsonFermion_2D_wing(x.NC,x.NX,x.NT)
end

function set_wing_fermion!(a::WilsonFermion_2D_wing{NC,NDW},boundarycondition) where {NC,NDW}
    NT = a.NT
    #NZ = a.NZ
    #NY = a.NY
    NX = a.NX
    NG = a.NG

    #!  X-direction
    for ialpha=1:NG
        for it=1:NT
            #for iz = 1:NZ
                #for iy=1:NY
                    @simd for k=1:NC
                        a[k,0, it,ialpha] = boundarycondition[1]*a[k,NX, it,ialpha]
                    end
                #end
            #end
        end
    end

    for ialpha=1:NG
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for k=1:NC
                        a[k,NX+1, it,ialpha] = boundarycondition[1]*a[k,1, it,ialpha]
                    end
                #end
            #end
        end
    end

    for ialpha=1:NG

        #T-direction
        #for iz=1:NZ
            #for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix, 0,ialpha] = boundarycondition[2]*a[k,ix, NT,ialpha]
                        a[k,ix, NT+1,ialpha] =boundarycondition[2]*a[k,ix, 1,ialpha]
                    end
                end
            #end
        #end

    end

end

function set_wing_fermion!(a::WilsonFermion_2D_wing{NC,NDW},boundarycondition,iseven::Bool) where {NC,NDW}
    NT = a.NT
    #NZ = a.NZ
    #NY = a.NY
    NX = a.NX
    NG = a.NG

    #!  X-direction
    @inbounds for ialpha=1:NG
        for it=1:NT
            #for iz = 1:NZ
                #for iy=1:NY
                    ix = NX
                    evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,0, it,ialpha] = boundarycondition[1]*a[k,NX, it,ialpha]
                        end
                    end
                #end
            #end
        end
    end

    @inbounds for ialpha=1:NG
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    ix = 1
                    evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,NX+1, it,ialpha] = boundarycondition[1]*a[k,1, it,ialpha]
                        end
                    end
                #end
            #end
        end
    end

    @inbounds for ialpha=1:NG

        #T-direction
        #for iz=1:NZ
            #for iy=1:NY
                for ix=1:NX
                    it = NT
                    evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,ix, 0,ialpha] = boundarycondition[2]*a[k,ix, NT,ialpha]
                        end
                    end

                    it = 1
                    evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,ix, NT+1,ialpha] =boundarycondition[2]*a[k,ix, 1,ialpha]
                        end
                    end

                end
            #end
        #end

    end

end

function add_fermion!(c::WilsonFermion_2D_wing{NC,NDW},α::Number,a::T1,β::Number,b::T2,iseven) where {NC,NDW,T1 <: Abstractfermion,T2 <: Abstractfermion}#c += alpha*a + beta*b
    n1,n2,n5,n6 = size(c.f)

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5 -NDW
            #for i4=1:n4
                #iz = i4 -NDW
                #for i3=1:n3
                    #iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse(( ix + it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                #println(a.f[i1,i2, i5,i6],"\t",b.f[i1,i2, i5,i6] )
                                c.f[i1,i2, i5,i6] += α*a.f[i1,i2, i5,i6] + β*b.f[i1,i2, i5,i6]
                            end
                        end
                    end
                #end
            #end
        end
    end
    return
end

function add_fermion!(c::WilsonFermion_2D_wing{NC,NDW},α::Number,a::T1,iseven::Bool) where {NC,NDW,T1 <: Abstractfermion,T2 <: Abstractfermion}#c += alpha*a + beta*b
    n1,n2, n5,n6 = size(c.f)

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5 -NDW
            #for i4=1:n4
                #iz = i4 -NDW
                #for i3=1:n3
                    #iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse(( ix + it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                #println(a.f[i1,i2, i5,i6],"\t",b.f[i1,i2, i5,i6] )
                                c.f[i1,i2, i5,i6] += α*a.f[i1,i2, i5,i6]
                            end
                        end
                    end
                #end
            #end
        end
    end
    return
end

function WWx!(xout::T,U::Array{G,1},x::T,A)  where  {T <: WilsonFermion_2D_wing,G <: AbstractGaugefields} #(1 - K^2 Teo Toe) xe
    iseven = true
    isodd = false
    temp = A._temporary_fermi[7]#temps[4]
    temp2 = A._temporary_fermi[6]#temps[4]

    #println("Wx")
    #@time Wx!(xout,U,x,A)
    clear_fermion!(xout,iseven)

    #Tx!(temp,U,x,A)
    Toex!(temp,U,x,A,iseven) #Toe
    #Tx!(temp2,U,temp,A)
    Toex!(temp2,U,temp,A,isodd) #Teo

    #set_wing_fermion!(temp,A.boundarycondition)
    #add_fermion!(xout,1,x,-1,temp2)
    add_fermion!(xout,1,x,-1,temp2,iseven)
    set_wing_fermion!(xout,A.boundarycondition,iseven)

    #Wx!(xout,U,x,A)
    return

end

function WWdagx!(xout::T,U::Array{G,1},x::T,A)  where  {T <: WilsonFermion_2D_wing,G <: AbstractGaugefields} #(1 - K^2 Teo Toe) xe
    iseven = true
    isodd = false
    temp = A._temporary_fermi[7]#temps[4]
    temp2 = A._temporary_fermi[6]#temps[4]

    clear_fermion!(xout)

    #Tx!(temp,U,x,A)
    Tdagoex!(temp,U,x,A,iseven) #Toe
    #Tx!(temp2,U,temp,A)
    Tdagoex!(temp2,U,temp,A,isodd) #Teo

    #set_wing_fermion!(temp,A.boundarycondition)
    #add_fermion!(xout,1,x,-1,temp2)
    add_fermion!(xout,1,x,-1,temp2,iseven)
    set_wing_fermion!(xout,A.boundarycondition)

    return
end

function clear_fermion!(a::WilsonFermion_2D_wing{NC,NDW} ,iseven) where {NC,NDW}
    n1,n2, n5,n6 = size(a.f)
    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5-NDW
            #for i4=1:n4
                #iz = i4-NDW
                #for i3=1:n3
                    #iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                a.f[i1,i2, i5,i6]= 0
                            end
                        end
                    end
                #end
            #end
        end
    end
end

function LinearAlgebra.mul!(x::WilsonFermion_2D_wing{NC,NDW},A::TA) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    #NY = x.NY
    #NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            x[ic,ix, it,1] = A[1,1]*e1+A[1,2]*e2#+A[1,3]*e3+A[1,4]*e4
                            x[ic,ix, it,2] = A[2,1]*e1+A[2,2]*e2#+A[2,3]*e3+A[2,4]*e4
                            #x[ic,ix, it,3] = A[3,1]*e1+A[3,2]*e2#+A[3,3]*e3+A[3,4]*e4
                            #sx[ic,ix, it,4] = A[4,1]*e1+A[4,2]*e2#+A[4,3]*e3+A[4,4]*e4

                    end
                #end
            #end
        end
    end

end

function LinearAlgebra.mul!(x::WilsonFermion_2D_wing{NC,NDW},A::TA,iseven::Bool) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    #NY = x.NY
    #NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            x[ic,ix, it,1] = A[1,1]*e1+A[1,2]*e2#+A[1,3]*e3+A[1,4]*e4
                            x[ic,ix, it,2] = A[2,1]*e1+A[2,2]*e2#+A[2,3]*e3+A[2,4]*e4
                            #x[ic,ix, it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            #x[ic,ix, it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4
                        end
                    end
                #end
            #end
        end
    end

end

#function LinearAlgebra.mul!(xout::WilsonFermion_2D_wing{NC,NDW},A::TA,x::WilsonFermion_2D_wing{NC}) where {TA <: AbstractMatrix, NC,NDW}
function LinearAlgebra.mul!(xout::WilsonFermion_2D_wing{NC,NDW},A::TA,x::Abstractfermion) where {TA <: AbstractMatrix, NC,NDW}

    NX = xout.NX
    #NY = xout.NY
    #NZ = xout.NZ
    NT = xout.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            xout[ic,ix, it,1] = A[1,1]*e1+A[1,2]*e2#+A[1,3]*e3+A[1,4]*e4
                            xout[ic,ix, it,2] = A[2,1]*e1+A[2,2]*e2#+A[2,3]*e3+A[2,4]*e4
                            #xout[ic,ix, it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            #xout[ic,ix, it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4

                    end
                #end
            #end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_2D_wing{NC,NDW},A::TA,x::WilsonFermion_2D_wing{NC},iseven::Bool) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    #NY = x.NY
    #NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            xout[ic,ix, it,1] = A[1,1]*e1+A[1,2]*e2#+A[1,3]*e3+A[1,4]*e4
                            xout[ic,ix, it,2] = A[2,1]*e1+A[2,2]*e2#+A[2,3]*e3+A[2,4]*e4
                            #xout[ic,ix, it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            #xout[ic,ix, it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4
                        end

                    end
                #end
            #end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_2D_wing{NC,NDW},x::WilsonFermion_2D_wing{NC},A::TA) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    #NY = x.NY
    #NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            xout[ic,ix, it,1] = A[1,1]*e1+A[2,1]*e2#+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix, it,2] = A[1,2]*e1+A[2,2]*e2#+A[3,2]*e3+A[4,2]*e4
                            #xout[ic,ix, it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            #xout[ic,ix, it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4

                    end
                #end
            #end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_2D_wing{NC,NDW},x::AbstractFermionfields_2D{NC},A::TA) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    #NY = x.NY
    #NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            xout[ic,ix, it,1] = A[1,1]*e1+A[2,1]*e2#+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix, it,2] = A[1,2]*e1+A[2,2]*e2#+A[3,2]*e3+A[4,2]*e4
                            #xout[ic,ix, it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            #xout[ic,ix, it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4

                    end
                #end
            #end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_2D_wing{NC,NDW},x::WilsonFermion_2D_wing{NC},A::TA,iseven::Bool) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    #NY = x.NY
    #NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix, it,1]
                            e2 = x[ic,ix, it,2]
                            #e3 = x[ic,ix, it,3]
                            #e4 = x[ic,ix, it,4]

                            xout[ic,ix, it,1] = A[1,1]*e1+A[2,1]*e2#+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix, it,2] = A[1,2]*e1+A[2,2]*e2#+A[3,2]*e3+A[4,2]*e4
                            #xout[ic,ix, it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            #xout[ic,ix, it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4
                        end
                    end
                #end
            #end
        end
    end

end

function LinearAlgebra.dot(a::WilsonFermion_2D_wing{NC,NDW},b::WilsonFermion_2D_wing{NC,NDW},iseven::Bool) where {NC,NDW}
    NT = a.NT
    #NZ = a.NZ
    #NY = a.NY
    NX = a.NX
    NG = a.NG

    c = 0.0im
    @inbounds for α=1:NG
        for it=1:NT
            #for iz=1:NZ
                #for iy=1:NY
                    for ix=1:NX
                        evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for ic=1:NC
                                c+= conj(a[ic,ix, it,α])*b[ic,ix, it,α]
                            end
                        end
                    end
                #end
            #end
        end
    end
    return c
end

#Overwrite Y with X*a + Y*b, where a and b are scalars. Return Y.
function LinearAlgebra.axpby!(a::Number, X::T, b::Number, Y::WilsonFermion_2D_wing{NC,NDW},iseven::Bool) where {NC,NDW,T <: AbstractFermionfields_2D}
    n1,n2, n5,n6 = size(Y.f)
    #println("axpby")

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5+NDW
            #for i4=1:n4
                #iz = i4+NDW
                #for i3=1:n3
                    #iy = i3+NDW
                    for i2=1:n2
                        ix = i2+NDW
                        evenodd = ifelse(( ix+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                Y.f[i1,i2, i5,i6] = a*X.f[i1,i2, i5,i6] + b*Y.f[i1,i2, i5,i6]
                            end
                        end
                    end
                #end
            #end
        end
    end
end

"""
c--------------------------------------------------------------------------c
c     y = sigma_3 * x
c     here
c                  ( 1       )
c        SIGMA3 =  (   -1     )
c--------------------------------------------------------------------------c
    """
    function mul_γ5x!(y::WilsonFermion_2D_wing{NC},x::WilsonFermion_2D_wing{NC}) where NC
        n1,n2, n5,n6 = size(x.f)
        @inbounds for i6=1:n6
            for i5=1:n5
                #it = i5+NDW
                #for i4=1:n4
                    #iz = i4+NDW
                    #for i3=1:n3
                        #iy = i3+NDW
                        for i2=1:n2
                            #ix = i2+NDW
                            @simd for ic=1:NC
                                y.f[i1,i2, i5,i6]=x.f[i1,i2, i5,i6]*ifelse(i6 == 2,-1,1)
                            end
                        end
                    #end
                #end
            end
        end

    end

    function apply_γ5!(x::WilsonFermion_2D_wing{NC}) where NC
        n1,n2, n5,n6 = size(x.f)
        #println("axpby")

        @inbounds for i6=1:n6
            for i5=1:n5
                #it = i5+NDW
                #for i4=1:n4
                    #iz = i4+NDW
                    #for i3=1:n3
                        #iy = i3+NDW
                        for i2=1:n2
                            #ix = i2+NDW
                            @simd for ic=1:NC
                                x.f[i1,i2, i5,i6]=x.f[i1,i2, i5,i6]*ifelse(i6 == 2,-1,1)
                            end
                        end
                    #end
                #end
            end
        end

    end

    #

    function mul_1plusγ5x!(y::WilsonFermion_2D_wing{NC},x::WilsonFermion_2D_wing{NC})  where NC#(1+gamma_5)/2
        NX = x.NX
        #NY = x.NY
        #NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                #for iz=1:NZ
                    #for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix, it,1] = x[ic,ix, it,1]
                            y[ic,ix, it,2] = 0#-1*x[ic,ix, it,2]
                            #y[ic,ix, it,3] = x[ic,ix, it,3]
                            #y[ic,ix, it,4] = x[ic,ix, it,4]
                        end
                    #end
                #end
            end
        end
    end

    function mul_1plusγ5x_add!(y::WilsonFermion_2D_wing{NC},x::WilsonFermion_2D_wing{NC},factor) where NC#x = x +(1+gamma_5)/2
        NX = x.NX
        #NY = x.NY
        #NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                #for iz=1:NZ
                    #for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix, it,1] += factor*x[ic,ix, it,1]
                            #y[ic,ix, it,4] += factor*x[ic,ix, it,4]
                        end
                    #end
                #end
            end
        end
    end

    function mul_1minusγ5x!(y::WilsonFermion_2D_wing{NC},x::WilsonFermion_2D_wing{NC}) where NC#(1-gamma_5)/2
        NX = x.NX
        #NY = x.NY
        #NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                #for iz=1:NZ
                    #for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix, it,1] = 0#x[ic,ix, it,1]
                            y[ic,ix, it,2] = x[ic,ix, it,2]
                            #y[ic,ix, it,3] = 0#x[ic,ix, it,3]
                            #y[ic,ix, it,4] = 0#x[ic,ix, it,4]
                        end
                    #end
                #end
            end
        end
    end

    function mul_1minusγ5x_add!(y::WilsonFermion_2D_wing{NC},x::WilsonFermion_2D_wing{NC},factor) where NC#+(1-gamma_5)/2
        NX = x.NX
        #NY = x.NY
        #NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                #for iz=1:NZ
                    #for iy=1:NY
                        @simd for ix=1:NX
                            #y[ic,ix, it,1] += factor*x[ic,ix, it,1]
                            y[ic,ix, it,2] += factor*x[ic,ix, it,2]
                        end
                    #end
                #end
            end
        end
    end


module BQCESubroutine

export broutine!, @broutine

using MLStyle
using ExprTools
using YaoLocations: YaoLocations, plain, Locations, CtrlLocations, AbstractLocations, merge_locations
using LinearAlgebra
using LoopVectorization
using PaddedMatrices

"""
    broutine!(st, op, locs[, ctrl, args...])

Qubit-based quantum circuit subroutine.

# Args

- `st`: the raw storage type.
- `op`: the operator
- `locs`: a `Locations`, location of the operator
- `ctrl`: optional. A `CtrlLocations`, location of the control qubits.
- `args`: optional. parameters of the operator.
"""
function broutine! end

include("utils.jl")
include("threading.jl")

using .Utilities

include("codegen/utils.jl")
include("codegen/broutine.jl")
include("mul.jl")
include("statevector.jl")

end


###############################################################################
#
#  NfRelOrd
#
###############################################################################

# I don't like that we have to drag S around
# Still hoping for julia/#18466

mutable struct NfRelOrdSet{T}
  nf::NumField{T}

  function NfRelOrdSet{T}(K::NumField{T}) where {T}
    a = new(K)
    return a
  end
end

mutable struct NfRelOrd{T, S, U} <: Ring
  nf::NumField{T}
  basis_nf::Vector{U}
  basis_matrix::Generic.MatSpaceElem{T}
  basis_mat_inv::Generic.MatSpaceElem{T}
  basis_pmatrix::PMat{T, S}
  pseudo_basis::Vector{Tuple{U, S}}

  disc_abs::NfOrdIdl # used if T == nf_elem
  disc_rel#::NfRelOrdIdl{T} # used otherwise; is a forward declaration
  parent::NfRelOrdSet{T}

  isequation_order::Bool

  ismaximal::Int                   # 0 Not known
                                   # 1 Known to be maximal
                                   # 2 Known to not be maximal

  trace_mat::Generic.MatSpaceElem{T}

  index::NfOrdIdl #Only if the base field is AnticNumberField

  inv_coeff_ideals::Vector{S}
  index_div
  @declare_other

  function NfRelOrd{T, S, U}(K::NumField{T}) where {T, S, U}
    z = new{T, S, U}()
    z.nf = K
    z.parent = NfRelOrdSet{T}(K)
    z.isequation_order = false
    z.ismaximal = 0
    z.index_div = Dict{ideal_type(order_type(base_field(K))), Vector{Tuple{ideal_type(z), Int}}}()
    return z
  end

  function NfRelOrd{T, S, U}(K::NumField{T}, M::PMat{T, S}) where {T, S, U}
    z = NfRelOrd{T, S, U}(K)
    z.nf = K
    z.parent = NfRelOrdSet{T}(K)
    z.basis_pmatrix = M
    z.basis_matrix = M.matrix
    z.index_div = Dict{ideal_type(order_type(base_field(K))), Vector{Tuple{ideal_type(z), Int}}}()
    return z
  end

  function NfRelOrd{T, S, U}(K::NumField{T}, M::Generic.MatSpaceElem{T}) where {T, S, U}
    z = NfRelOrd{T, S, U}(K)
    z.nf = K
    z.parent = NfRelOrdSet{T}(K)
    z.basis_matrix = M
    z.basis_pmatrix = pseudo_matrix(M)
    z.index_div = Dict{ideal_type(order_type(base_field(K))), Vector{Tuple{ideal_type(z), Int}}}()
    return z
  end
end

###############################################################################
#
#  NfRelOrdElem
#
###############################################################################

mutable struct NfRelOrdElem{T, U} <: RingElem
  parent#::NfRelOrd{T, S} # I don't want to drag the S around
  elem_in_nf::U
  coordinates::Vector{T}
  has_coord::Bool

  function NfRelOrdElem{T, U}(O::NfRelOrd{T}) where {T, U}
    z = new{T, U}()
    z.parent = O
    z.elem_in_nf = zero(nf(O))
    z.coordinates = Vector{T}(undef, degree(O))
    z.has_coord = false
    return z
  end

  function NfRelOrdElem{T, U}(O::NfRelOrd{T}, a::U) where {T, U}
    z = new{T, U}()
    z.parent = O
    z.elem_in_nf = a
    z.coordinates = Vector{T}(undef, degree(O))
    z.has_coord = false
    return z
  end

  function NfRelOrdElem{T, U}(O::NfRelOrd{T}, a::U, arr::Vector{T}) where {T, U}
    z = new{T, U}()
    z.parent = O
    z.elem_in_nf = a
    z.coordinates = arr
    z.has_coord = true
    return z
  end
end

###############################################################################
#
#  NfRelOrdFracIdl
#
###############################################################################

mutable struct NfRelOrdFracIdlSet{T, S, U}
  order::NfRelOrd{T, S, U}

  function NfRelOrdFracIdlSet{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    a = new(O)
    return a
  end
end

mutable struct NfRelOrdFracIdl{T, S, U}
  order::NfRelOrd{T, S, U}
  parent::NfRelOrdFracIdlSet{T, S, U}
  basis_pmatrix::PMat{T, S}
  pseudo_basis::Vector{Tuple{U, S}}
  basis_matrix::Generic.MatSpaceElem{T}
  basis_mat_inv::Generic.MatSpaceElem{T}
  den::fmpz

  norm
  has_norm::Bool

  function NfRelOrdFracIdl{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    z = new{T, S, U}()
    z.order = O
    z.parent = NfRelOrdFracIdlSet{T, S, U}(O)
    z.has_norm = false
    return z
  end

  function NfRelOrdFracIdl{T, S, U}(O::NfRelOrd{T, S, U}, M::PMat{T, S}) where {T, S, U}
    z = NfRelOrdFracIdl{T, S, U}(O)
    z.basis_pmatrix = M
    z.basis_matrix = M.matrix
    return z
  end
end

###############################################################################
#
#  NfRelOrdIdl
#
###############################################################################

mutable struct NfRelOrdIdlSet{T, S, U}
  order::NfRelOrd{T, S, U}

  function NfRelOrdIdlSet{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    a = new(O)
    return a
  end
end

mutable struct NfRelOrdIdl{T, S, U}
  order::NfRelOrd{T, S, U}
  parent::NfRelOrdIdlSet{T, S, U}
  basis_pmatrix::PMat{T, S}
  pseudo_basis::Vector{Tuple{U, S}}
  basis_matrix::Generic.MatSpaceElem{T}
  basis_mat_inv::Generic.MatSpaceElem{T}

  norm
  has_norm::Bool
  is_prime::Int            # 0: don't know
                           # 1 known to be prime
                           # 2 known to be not prime
  splitting_type::Tuple{Int, Int}

  minimum
  non_index_div_poly::fq_poly # only used if the ideal is a prime ideal not dividing the index
  p_uniformizer::NfRelOrdElem{T, U}
  anti_uniformizer::U

  function NfRelOrdIdl{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    z = new{T, S, U}()
    z.order = O
    z.parent = NfRelOrdIdlSet{T, S, U}(O)
    z.has_norm = false
    z.is_prime = 0
    z.splitting_type = (0,0)
    return z
  end

  function NfRelOrdIdl{T, S, U}(O::NfRelOrd{T, S, U}, M::PMat{T, S}) where {T, S, U}
    z = NfRelOrdIdl{T, S, U}(O)
    z.basis_pmatrix = M
    z.basis_matrix = M.matrix
    return z
  end
end

################################################################################
#
#  NfRelNS / NfRelNSElem
#
################################################################################

mutable struct NfRelNS{T} <: NonSimpleNumField{T}
  base_ring::Nemo.Field
  pol::Array{Nemo.Generic.MPoly{T}, 1}
  abs_pol::Array{Generic.Poly{T}, 1}
  S::Array{Symbol, 1}
  basis_traces::Vector{Vector{T}}
  auxilliary_data::Array{Any, 1}
  @declare_other

  function NfRelNS(abs_pol::Array{Generic.Poly{T}}, f::Array{Nemo.Generic.MPoly{T}, 1}, S::Array{Symbol, 1}; cached::Bool = false) where T
    r = new{T}()
    r.pol = f
    r.abs_pol = abs_pol
    r.base_ring = base_ring(f[1])
    r.S = S
    r.auxilliary_data = Array{Any}(undef, 5)
    return r
  end
end

mutable struct NfRelNSElem{T} <: NonSimpleNumFieldElem{T}
  data::Nemo.Generic.MPoly{T}
  parent::NfRelNS{T}

  NfRelNSElem{T}(g::Generic.MPoly{T}) where {T} = new{T}(g)
end

################################################################################
#
#  Quotient rings of orders
#
################################################################################

# T1 is the type of the order, T2 the type of the ideal,
# T3 the type of the basis pseudo matrix of the ideal
mutable struct RelOrdQuoRing{T1, T2, T3} <: Ring
  base_ring::T1
  ideal::T2
  basis_pmatrix::T3

  function RelOrdQuoRing{T1, T2, T3}(O::T1, I::T2) where { T1, T2, T3 }
    z = new{T1, T2, T3}()
    z.base_ring = O
    z.ideal = I
    z.basis_pmatrix = basis_pmatrix_wrt(I, O)
    return z
  end
end

function RelOrdQuoRing(O::T1, I::T2) where { T1, T2 }
  @assert T2 == ideal_type(O)
  return RelOrdQuoRing{T1, T2, typeof(basis_pmatrix(I, copy = false))}(O, I)
end

# T1, T2 and T3 as for RelOrdQuoRing, S is the elem_type of the order
mutable struct RelOrdQuoRingElem{T1, T2, T3, S} <: RingElem
  elem::S
  parent::RelOrdQuoRing{T1, T2, T3}

  function RelOrdQuoRingElem{T1, T2, T3, S}(Q::RelOrdQuoRing{T1, T2, T3}, x::S) where { T1, T2, T3, S }
    z = new{T1, T2, T3, S}()
    z.elem = mod(x, Q)
    z.parent = Q
    return z
  end
end

function RelOrdQuoRingElem(Q::RelOrdQuoRing{T1, T2, T3}, x::S) where { T1, T2, T3, S }
  return RelOrdQuoRingElem{T1, T2, T3, S}(Q, x)
end


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# src/quadexpr.jl
# Defines all types relating to expressions with a quadratic and affine part
# - GenericQuadExpr             ∑qᵢⱼ xᵢⱼ  +  ∑ aᵢ xᵢ  +  c
#   - QuadExpr                  Alias for (Float64, Variable)
# - GenericQuadConstraint       ∑qᵢⱼ xᵢⱼ  +  ∑ aᵢ xᵢ  +  c  [≤,≥]  0
#   - QuadConstraint            Alias for (Float64, Variable)
# Operator overloads in src/operators.jl
#############################################################################

#############################################################################
# GenericQuadExpr
# ∑qᵢⱼ xᵢⱼ  +  ∑ aᵢ xᵢ  +  c
type GenericQuadExpr{CoefType,VarType} <: AbstractJuMPScalar
    qvars1::Vector{VarType}
    qvars2::Vector{VarType}
    qcoeffs::Vector{CoefType}
    aff::GenericAffExpr{CoefType,VarType}
end
coeftype{C,V}(::GenericQuadExpr{C,V}) = C

Base.isempty(q::GenericQuadExpr) = (length(q.qvars1) == 0 && isempty(q.aff))
Base.zero{C,V}(::Type{GenericQuadExpr{C,V}}) = GenericQuadExpr(V[], V[], C[], zero(GenericAffExpr{C,V}))
Base.one{C,V}(::Type{GenericQuadExpr{C,V}})  = GenericQuadExpr(V[], V[], C[],  one(GenericAffExpr{C,V}))
Base.zero(q::GenericQuadExpr) = zero(typeof(q))
Base.one(q::GenericQuadExpr)  =  one(typeof(q))
Base.copy(q::GenericQuadExpr) = GenericQuadExpr(copy(q.qvars1),copy(q.qvars2),copy(q.qcoeffs),copy(q.aff))

function Base.append!{T,S}(q::GenericQuadExpr{T,S}, other::GenericQuadExpr{T,S})
    append!(q.qvars1, other.qvars1)
    append!(q.qvars2, other.qvars2)
    append!(q.qcoeffs, other.qcoeffs)
    append!(q.aff, other.aff)
    q
end

function assert_isfinite(q::GenericQuadExpr)
    assert_isfinite(q.aff)
    for i in 1:length(q.qcoeffs)
        isfinite(q.qcoeffs[i]) || error("Invalid coefficient $(q.qcoeffs[i]) on quadratic term $(q.qvars1[i])*$(q.qvars2[i])")
    end
end

function Base.isequal{T,S}(q::GenericQuadExpr{T,S},other::GenericQuadExpr{T,S})
    isequal(q.aff,other.aff)   || return false
    length(q.qvars1) == length(other.qvars1) || return false
    for i in 1:length(q.qvars1)
        isequal(q.qvars1[i],  other.qvars1[i]) || return false
        isequal(q.qvars2[i],  other.qvars2[i]) || return false
        isequal(q.qcoeffs[i],other.qcoeffs[i]) || return false
    end
    return true
end

# Alias for (Float64, Variable)
typealias QuadExpr GenericQuadExpr{Float64,Variable}
Base.convert(::Type{QuadExpr}, v::Union{Real,Variable,AffExpr}) = QuadExpr(Variable[], Variable[], Float64[], AffExpr(v))
QuadExpr() = zero(QuadExpr)

function setobjective(m::Model, sense::Symbol, q::QuadExpr)
    m.obj = q
    if m.internalModelLoaded
        if method_exists(MathProgBase.setquadobjterms!, (typeof(m.internalModel), Vector{Cint}, Vector{Cint}, Vector{Float64}))
            verify_ownership(m, m.obj.qvars1)
            verify_ownership(m, m.obj.qvars2)
            MathProgBase.setquadobjterms!(m.internalModel, Cint[v.col for v in m.obj.qvars1], Cint[v.col for v in m.obj.qvars2], m.obj.qcoeffs)
        else
            isa(m.internalModel, MathProgBase.AbstractLinearQuadraticModel) && Base.warn_once("Solver does not support adding a quadratic objective to an existing model. Hot-start is disabled.")
            m.internalModelLoaded = false
        end
    end
    setobjectivesense(m, sense)
end

# Copy a quadratic expression to a new model by converting all the
# variables to the new model's variables
function Base.copy(q::QuadExpr, new_model::Model)
    QuadExpr(copy(q.qvars1, new_model), copy(q.qvars2, new_model),
                copy(q.qcoeffs), copy(q.aff, new_model))
end

function getvalue(a::QuadExpr)
    ret = getvalue(a.aff)
    for it in 1:length(a.qvars1)
        ret += a.qcoeffs[it] * getvalue(a.qvars1[it]) * getvalue(a.qvars2[it])
    end
    return ret
end
getvalue(arr::Array{QuadExpr}) = map(getvalue, arr)

##########################################################################
# GenericQuadConstraint
# ∑qᵢⱼ xᵢⱼ  +  ∑ aᵢ xᵢ  +  c  [≤,≥]  0
# As RHS is implicitly taken to be zero, we store only LHS and sense
type GenericQuadConstraint{QuadType} <: AbstractConstraint
    terms::QuadType
    sense::Symbol
end
Base.copy{CON<:GenericQuadConstraint}(c::CON, new_model::Model) = CON(copy(c.terms, new_model), c.sense)

# Alias for (Float64, Variable)
typealias QuadConstraint GenericQuadConstraint{QuadExpr}

function Base.copy(c::QuadConstraint, new_model::Model)
    return QuadConstraint(copy(c.terms, new_model), c.sense)
end

function addconstraint(m::Model, c::QuadConstraint)
    push!(m.quadconstr,c)
    if m.internalModelLoaded
        if method_exists(MathProgBase.addquadconstr!, (typeof(m.internalModel),
                                                       Vector{Cint},
                                                       Vector{Float64},
                                                       Vector{Cint},
                                                       Vector{Cint},
                                                       Vector{Float64},
                                                       Char,
                                                       Float64))
            if !((s = string(c.sense)[1]) in ['<', '>', '='])
                error("Invalid sense for quadratic constraint")
            end
            terms = c.terms
            verify_ownership(m, terms.qvars1)
            verify_ownership(m, terms.qvars2)
            MathProgBase.addquadconstr!(m.internalModel,
                                        Cint[v.col for v in c.terms.aff.vars],
                                        c.terms.aff.coeffs,
                                        Cint[v.col for v in c.terms.qvars1],
                                        Cint[v.col for v in c.terms.qvars2],
                                        c.terms.qcoeffs,
                                        s,
                                        -c.terms.aff.constant)
        else
            Base.warn_once("Solver does not appear to support adding quadratic constraints to an existing model. Hot-start is disabled.")
            m.internalModelLoaded = false
        end
    end
    return ConstraintRef{Model,QuadConstraint}(m,length(m.quadconstr))
end
addconstraint(m::Model, c::Array{QuadConstraint}) =
    error("Vectorized constraint added without elementwise comparisons. Try using one of (.<=,.>=,.==).")

function addVectorizedConstraint(m::Model, v::Array{QuadConstraint})
    ret = Array(ConstraintRef{Model,QuadConstraint}, size(v))
    for I in eachindex(v)
        ret[I] = addconstraint(m, v[I])
    end
    ret
end


module Clapeyron
using StaticArrays
using LinearAlgebra
using SparseArrays
#for the assoc solver and the sparse packed VofV
import PackedVectorsOfVectors
const PackedVofV = PackedVectorsOfVectors.PackedVectorOfVectors

#for non allocating vectors of zeros and ones
using FillArrays: FillArrays

using Roots: Roots
using NLSolvers
import Metaheuristics
using DiffResults, ForwardDiff
using Scratch

include("solvers/Solvers.jl")
using .Solvers
using .Solvers: log, sqrt
∂Tag = Solvers.∂Tag
using Unitful
import LogExpFunctions

#The Base of Clapeyron: EoSModel and eos(model,V,T,z)
include("base/constants.jl")
include("base/EoSModel.jl")
include("base/eosshow.jl")

#EoSParam, ClapeyronParam
include("params/paramvectors.jl")
include("params/ClapeyronParam.jl")
include("params/combiningrules.jl")

using CSV, Tables
#getparams machinery
include("database/ParamOptions.jl")
include("database/database.jl")
include("database/UserReader.jl")

#macros, used for defining models
include("utils/macros.jl")

#splitting models, useful for methods.
include("utils/split_model.jl")

#Clapeyron methods (AD, property solvers, etc)
include("methods/methods.jl")

#=
the dependency chain is the following:

base --> params -|-> database ----------------|
                 |-> split_model --> methods -|-> models
                 |-> macros ------------------|

=#

#Clapeyron EoS collection
include("models/types.jl") #type hierarchy
include("models/ideal/ideal.jl")
include("models/ideal/BasicIdeal.jl")
include("models/ideal/MonomerIdeal.jl")
include("models/ideal/ReidIdeal.jl")
include("models/ideal/WalkerIdeal.jl")
include("models/ideal/JobackIdeal.jl")

include("models/SAFT/PCSAFT/PCSAFT.jl")
include("models/SAFT/PCSAFT/variants/sPCSAFT.jl")
include("models/SAFT/ogSAFT/ogSAFT.jl")
include("models/SAFT/CPA/CPA.jl")
include("models/SAFT/CPA/variants/sCPA.jl")
include("models/SAFT/SAFTVRSW/SAFTVRSW.jl")
include("models/SAFT/LJSAFT/LJSAFT.jl")
include("models/SAFT/softSAFT/softSAFT.jl")
include("models/SAFT/SAFTVRMie/SAFTVRMie.jl")
include("models/SAFT/SAFTVRMie/variants/SAFTVRQMie.jl")
include("models/SAFT/SAFTgammaMie/SAFTgammaMie.jl")
include("models/SAFT/CKSAFT/CKSAFT.jl")
include("models/SAFT/CKSAFT/variants/sCKSAFT.jl")
include("models/SAFT/BACKSAFT/BACKSAFT.jl")
include("models/SAFT/equations.jl")

include("models/cubic/equations.jl")
include("models/cubic/vdW.jl")
include("models/cubic/RK/RK.jl")
include("models/cubic/PR/PR.jl")

include("models/Activity/Wilson/Wilson.jl")
include("models/Activity/NRTL/NRTL.jl")
include("models/Activity/UNIQUAC/UNIQUAC.jl")
include("models/Activity/UNIFAC/UNIFAC.jl")
include("models/Activity/UNIFAC/variants/ogUNIFAC.jl")
include("models/Activity/UNIFAC/variants/PSRK.jl")
include("models/Activity/UNIFAC/variants/VTPR.jl")
include("models/Activity/equations.jl")

include("models/Activity/COSMOSAC/utils.jl")
include("models/Activity/COSMOSAC/COSMOSAC02.jl")
include("models/Activity/COSMOSAC/COSMOSAC10.jl")
include("models/Activity/COSMOSAC/COSMOSACdsp.jl")

include("models/cubic/alphas/alphas.jl")
include("models/cubic/mixing/mixing.jl")
include("models/cubic/translation/translation.jl")

include("models/cubic/RK/variants/SRK.jl")
include("models/cubic/RK/variants/PSRK.jl")
include("models/cubic/PR/variants/PR78.jl")
include("models/cubic/PR/variants/VTPR.jl")
include("models/cubic/PR/variants/UMRPR.jl")

include("models/EmpiricHelmholtz/IAPWS95.jl")
include("models/EmpiricHelmholtz/PropaneRef.jl")
include("models/EmpiricHelmholtz/LJRef.jl")
include("models/EmpiricHelmholtz/MultiFluid/multifluid.jl")

include("models/LatticeFluid/SanchezLacombe/SanchezLacombe.jl")

include("models/SPUNG/SPUNG.jl")
include("models/UFTheory/UFTheory.jl")

include("models/cached/CachedEoS.jl")

end # module