# == THETA ========================================================================================================================================= #
function θ(contracts::Array{Y,1}; number_of_levels::Int64=2, T::Float64=(1 / 365),
    σ::Float64=0.15, Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational)::Array{Float64,1} where {Y<:AbstractDerivativeContractModel}

    # value array -
    value_array = Array{Float64,1}()

    # compute -
    for contract ∈ contracts
        value = θ(contract; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ, choice=choice)
        push!(value_array, value)
    end

    # return -
    return value_array
end

function θ(contract::Y; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational)::Float64 where {Y<:AbstractDerivativeContractModel}

    # θ : 1 day diff
    Tₒ = T
    T₁ = Tₒ - (1 / 365)

    # build a binary tree with N levels -
    mₒ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=Tₒ, μ=μ)
    m₁ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T₁, μ=μ)

    # compute -
    Pₒ = premium(contract, mₒ; choice=choice)
    P₁ = premium(contract, m₁; choice=choice)

    # compute theta -
    θ_value = (P₁ - Pₒ)

    # return the value -
    return θ_value
end

function θ(model::CRRJITContractPremiumLatticeModel, point::PQContractPremiumLatticePoint)::Float64

    # get stuff from the lattice model -
    contractType = model.contractType

    # get our operating point from the point model -
    underlying_value = model.underlying[point.s]
    iv_value = model.iv[point.j]
    K_value = model.strike[point.i]
    dte_value = model.dte[point.d]
    number_of_levels = model.number_of_levels
    risk_free_rate = model.risk_free_rate

    # build an empty contract -
    contract = build(contractType, Dict{String,Any}())
    contract.number_of_contracts = 1
    contract.direction = 1
    contract.strike_price = K_value

    # compute -
    return θ(contract; number_of_levels = number_of_levels, T = (dte_value/365.0), σ = iv_value,
        Sₒ = underlying_value, μ = risk_free_rate);
end
# ================================================================================================================================================== #

# == DELTA ========================================================================================================================================= #
function δ(contracts::Array{Y,1}; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational)::Array{Float64,1} where {Y<:AbstractDerivativeContractModel}

    # value array -
    value_array = Array{Float64,1}()

    # compute -
    for contract ∈ contracts
        value = δ(contract; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ, choice=choice)
        push!(value_array, value)
    end

    # return -
    return value_array
end

function δ(contract::Y; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational)::Float64 where {Y<:AbstractDerivativeContractModel}

    # advance base price by 1 -
    S₁ = Sₒ + 1

    # build models -
    mₒ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ)
    m₁ = build(CRRLatticeModel; Sₒ=S₁, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ)

    # compute -
    Pₒ = premium(contract, mₒ; choice=choice)
    P₁ = premium(contract, m₁; choice=choice)

    # compute theta -
    δ_value = (P₁ - Pₒ)

    # return the value -
    return δ_value
end

function δ(model::CRRJITContractPremiumLatticeModel, point::PQContractPremiumLatticePoint)::Float64

    # get stuff from the lattice model -
    contractType = model.contractType

    # get our operating point from the point model -
    underlying_value = model.underlying[point.s]
    iv_value = model.iv[point.j]
    K_value = model.strike[point.i]
    dte_value = model.dte[point.d]
    number_of_levels = model.number_of_levels
    risk_free_rate = model.risk_free_rate

    # build an empty contract -
    contract = build(contractType, Dict{String,Any}())
    contract.number_of_contracts = 1
    contract.direction = 1
    contract.strike_price = K_value

    # call the original δ function -
    return δ(contract; number_of_levels = number_of_levels, T = (dte_value/365.0), σ = iv_value,
        Sₒ = underlying_value, μ = risk_free_rate)
end

function δ(model::CRRJITContractPremiumLatticeModel, points::Array{PQContractPremiumLatticePoint,1})::Array{Float64,1}

    # initialize -
    delta_array = Array{Float64,1}()

    # process the points -
    for point ∈ points
        value = δ(model, point)
        push!(delta_array, value)
    end

    # return -
    return delta_array
end
# ================================================================================================================================================== #

# == GAMMA ========================================================================================================================================= #
function γ(contract::Y; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational) where {Y<:AbstractDerivativeContractModel}

    # advance base price by 1 -
    S₁ = Sₒ + 1

    # compute -
    δₒ = δ(contract; number_of_levels=number_of_levels, T=T, σ=σ, Sₒ=Sₒ, μ=μ, choice=choice)
    δ₁ = δ(contract; number_of_levels=number_of_levels, T=T, σ=σ, Sₒ=S₁, μ=μ, choice=choice)

    @show (δₒ,δ₁, T)

    # compute γ -
    γ_value = (δ₁ - δₒ)

    # return -
    return γ_value
end

function γ(contracts::Array{Y,1}; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational) where {Y<:AbstractDerivativeContractModel}

    # initialize -
    value_array = Array{Float64,1}()

    # compute -
    for contract ∈ contracts
        value = γ(contract; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ, choice=choice)
        push!(value_array, value)
    end

    # return -
    return value_array
end

function γ(model::CRRJITContractPremiumLatticeModel, point::PQContractPremiumLatticePoint;
    choice::Function=_rational)::Float64

    # get stuff from the lattice model -
    contractType = model.contractType

    # get our operating point from the point model -
    underlying_value = model.underlying[point.s]
    iv_value = model.iv[point.j]
    K_value = model.strike[point.i]
    dte_value = model.dte[point.d]
    number_of_levels = model.number_of_levels
    risk_free_rate = model.risk_free_rate

    # build an empty contract -
    contract = build(contractType, Dict{String,Any}())
    contract.number_of_contracts = 1
    contract.direction = 1
    contract.strike_price = K_value

    # compute -
    return γ(contract; number_of_levels = number_of_levels, T = (dte_value/365.0), σ = iv_value,
        Sₒ = underlying_value, μ = risk_free_rate, choice=choice);
end
# ================================================================================================================================================== #

# == VEGA ========================================================================================================================================== #
function vega(contracts::Array{Y,1}; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational) where {Y<:AbstractDerivativeContractModel}

    # initialize -
    value_array = Array{Float64,1}()

    # compute -
    for contract ∈ contracts
        value = vega(contract; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ, choice=choice)
        push!(value_array, value)
    end

    # return -
    return value_array
end

function vega(contract::Y; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational) where {Y<:AbstractDerivativeContractModel}

    # setup the calculation -
    σₒ = σ
    σ₁ = σ + 0.01

    # build models -
    mₒ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σₒ, T=T, μ=μ)
    m₁ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ₁, T=T, μ=μ)

    # compute -
    Pₒ = premium(contract, mₒ; choice=choice)
    P₁ = premium(contract, m₁; choice=choice)

    # compute theta -
    vega_value = (P₁ - Pₒ)

    # return the value -
    return vega_value
end

function vega(model::CRRJITContractPremiumLatticeModel, point::PQContractPremiumLatticePoint)::Float64

    # get stuff from the lattice model -
    contractType = model.contractType

    # get our operating point from the point model -
    underlying_value = model.underlying[point.s]
    iv_value = model.iv[point.j]
    K_value = model.strike[point.i]
    dte_value = model.dte[point.d]
    number_of_levels = model.number_of_levels
    risk_free_rate = model.risk_free_rate

    # build an empty contract -
    contract = build(contractType, Dict{String,Any}())
    contract.number_of_contracts = 1
    contract.direction = 1
    contract.strike_price = K_value

    # compute -
    return vega(contract; number_of_levels = number_of_levels, T = (dte_value/365.0), σ = iv_value,
        Sₒ = underlying_value, μ = risk_free_rate);
end
# ================================================================================================================================================== #

# == RHO =========================================================================================================================================== #
function ρ(contract::Y; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational) where {Y<:AbstractDerivativeContractModel}

    # setup mu -
    μₒ = μ
    μ₁ = μ + 0.001

    # build models -
    mₒ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μₒ)
    m₁ = build(CRRLatticeModel; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ₁)

    # compute -
    Pₒ = premium(contract, mₒ; choice=choice)
    P₁ = premium(contract, m₁; choice=choice)

    # compute theta -
    ρ_value = (P₁ - Pₒ)

    # return the value -
    return ρ_value
end

function ρ(contracts::Array{Y,1}; number_of_levels::Int64=2, T::Float64=(1 / 365), σ::Float64=0.15,
    Sₒ::Float64=1.0, μ::Float64=0.0015, choice::Function=_rational) where {Y<:AbstractDerivativeContractModel}

    # initialize -
    value_array = Array{Float64,1}()

    # compute -
    for contract ∈ contracts
        value = ρ(contract; Sₒ=Sₒ, number_of_levels=number_of_levels, σ=σ, T=T, μ=μ, choice=choice)
        push!(value_array, value)
    end

    # return -
    return value_array
end
# ================================================================================================================================================== #


# Example - BVAR
using Plots, CSV, StatsBase
include(joinpath(Pkg.dir("VectorAutoregressions"),"src","bvar.jl"))
plotly()

y = CSV.read(joinpath(Pkg.dir("VectorAutoregressions"),"test","bvar_data.csv"), header = false)
y = convert(Array,y[:,1:3])
const λ = 0.1:0.1:1
const τ = 10*(0.1:0.1:1)
const ε = 0.0001
const p = 4
const H = 10
const reps = 10000
const burnin = 8000
const max_try = 1000
const update = 1000
prior = Hyperparameter(λ,τ,ε,p,H,reps,burnin,max_try,update)

mForecast = fit_bvar(y,prior)
pFore = plot(layout = grid(1,3));
for j in 1:size(mForecast,3)
plot!(pFore,size(y,1)-H:size(y,1)+H, [y[size(y,1)-H:end,j]; median(mForecast[:,:,j],1)'], legend = false,subplot = j);
plot!(pFore,size(y,1)+1:size(y,1)+H, [[percentile(mForecast[:,i,j],16) for i in 1:H] median(mForecast[:,:,j],1)' [percentile(mForecast[:,i,j],84) for i in 1:H]], color = "red", line = [:dash :solid :dash], subplot = j)
end

gui(pFore)


module SparseVolterraExamples

    using ApproxFun, NLsolve, BandedMatrices, BlockBandedMatrices, LinearAlgebra, BlockArrays, MultivariateOrthogonalPolynomials, Test, Plots

    include("triIntegroDiff.jl")
    include("triQEgen.jl")
    include("triVolterraFullKernel.jl")
    include("triVolterraMonKernel.jl")

    export triNonLinearVolterraObjective,
        triNonLinearIntegroDiffVolterraObjective,
        pochhammer,
        alternI,
        EvalJPat1mat,
        DirectEvalLHSP10at0,
        DirectEvalLHSP10atPrime0,
        DirectEvalLHSP10atPPrime0,
        DirectEvalLHSP10atPPPrime0,
        DirectEvalLHSP10atPPPPrime0,
        triQgen,
        triEygenP01,
        triQEygenP01,
        WLoweringP01P00,
        reflectPabtoPba,
        WLoweringP01Pm11,
        ReflectorP01,
        triEygenP00,
        triEygenP11,
        triExgenP01,
        triVolterraEQ1FullKernelSolver,
        triVolterraEQ2FullKernelSolver,
        triVolterraFullKernelFunP01,
        triVolterraFullKernelOpP01,
        triVolterraFullKernelOpP01,
        triBandedSgen,
        triVolterraEQ1Solver,
        triVolterraEQ2Solver,
        triVolterraFunP01,
        triVolterraOpP01,
        triJQEJbasislexP01,
        triJQEJbasisgridgenP01,
        triNoKernelDirect

end # module


module SSLTestMosaic
using SimpleSDMLayers
using Test
using Statistics

n1 = fill(1.0f0, 5, 5)
s1 = SimpleSDMResponse(n1, 0., 2., 0., 2.)

n2 = fill(2.0, 5, 5)
s2 = SimpleSDMResponse(n2, 1., 3., 1., 3.)

n3 = fill(3.0, 10, 10)
s3 = SimpleSDMResponse(n3, -1.8, 2.2, 1.4, 5.4)

n4 = fill(4.0, 5, 3)
s4 = SimpleSDMResponse(n4, 3.0, 4.2, 3.0, 5.0)

# Test
L = mosaic(mean, [s1, s2, s3, s4])

@test typeof(L) == SimpleSDMResponse{Float32}
@test extrema(L) == (1.0f0, 4.0f0)

end


"""
Reproduction of the empirical results concerning an empirical application
in the prediction of transportation mode choice using the best subset
maximum score approach of Chen and Lee (2017).
The work-trip mode choice dataset of Horowitz (1993) is included in the package.
"""
module transportationMode

using DataFrames, Compat
using MathProgBase, Gurobi
using ..BSBP

""" Best subset maximum score approach on the work-trip mode choice dataset of Horowitz.
Args:
    warm_start (int,optional):  Set to 1 for warm start strategy. Default = 1.
    tau (float,optional):       Tuning paramater to construct the refined bound used in the warm start approach. Default = 1.5.
    mio (int,optional):         Set to 1 for using Method 1 for the MIO formulation. Set to 2 for Method 2. Default to 1.
    q (int, optional):          Value of the variable selection bound. Default = 1.
    series_exp (int, optional): Set to 1 to use quadratic expansion terms as covariates. Default = 0.
    b (int, optional):          Bound value. Default = 1.
"""
function TransportationMode(;warm_start=0,tau=1.5,mio=2,series_exp=0,beta0=1,b=10,q=1,time_limit=86400)
  # load data
  data = readdlm("./data_horowitz.csv",',')

  bhat = []
  score = 0
  gap = 0
  rtime = 0
  ncount = 0

  println("estimation based on full sample")
  Y_tr=data[:,1]
  temp=data[:,2:end]

  n_tr=length(Y_tr)
  # [DCOST CARS DOVTT DIVTT]
  x_std=(temp.-repmat(mean(temp,1),n_tr,1))./repmat(std(temp,1),n_tr,1)
  x_foc = [x_std[:,1] ones(n_tr,1)] # [DCOST Intercept]

  if series_exp == 1
      z2 = x_std[:,2]
      z3 = x_std[:,3]
      z4 = x_std[:,4]
      x_aux1 = [z2, z3, z4] # linear terms
      x_aux2 = [z2.*z3, z3.*z4, z2.*z4]
      x_aux3 = [z2.*z2,z3.*z3, z4.*z4]
      x_aux = [x_aux1, x_aux2, x_aux3]
  else
      x_aux = x_std[:,2:4] # [CARS DOVTT DIVTT]
  end

  k=size(x_foc)[2]
  d=size(x_aux)[2]

  bnd= hcat(-b*ones(k-1+d,1),b*ones(k-1+d,1)) # set the initial parameter bounds

  tol = floor(sqrt(log(n_tr)*n_tr)/2) # set the tolerance level value
  println("tolerance level: ", tol/n_tr)

  if warm_start == 1 # warm start MIO
      bhat,score,gap,rtime,ncount  = WarmStartMaxScore(Y_tr,x_foc,x_aux,beta0,q,time_limit,tol,bnd,mio,tau)
  else # cold start MIO
      bhat,score,gap,rtime,ncount  = MaxScore(Y_tr,x_foc,x_aux,beta0,q,time_limit,tol,bnd,mio)
  end

  println("parameter estimates: ", bhat)
  println("avg_score gap time node_count:", score, " , ", gap, " , ", rtime, " , ", ncount)
  return score,gap,rtime,ncount
end

export TransportationMode
end


using Test, MathematicalPredicates

@time @testset "Basic functionality" begin include("basic.jl") end
@time @testset "LazySets support" begin include("LazySets.jl") end


# slab_rhs!
using ClimaCore

struct ThermalSlabParameters# <: CLIMAParameters.AbstractEarthParameterSet{F}
    h::FT
    ρ::FT
    c::FT
    T_init::FT
end

# domain
function ShellDomain(; radius = 6371e3, Nel = 8, Nq = 2)
    domain = ClimaCore.Domains.SphereDomain(radius)
    mesh = ClimaCore.Meshes.EquiangularCubedSphere(domain, Nel)
    topology = ClimaCore.Topologies.Topology2D(mesh)
    quad = ClimaCore.Spaces.Quadratures.GLL{Nq}()
    space = ClimaCore.Spaces.SpectralElementSpace2D(topology, quad)
end

# init simulation
function slab_space_init(::Type{FT}, space, params) where {FT}

    coords = ClimaCore.Fields.coordinate_field(space)

    # initial condition
    T_sfc = map(coords) do coord
        T_sfc_0 = params.T_init
        anom_ampl = FT(2)
        radlat = coord.lat / FT(180) * pi
        lat_0 = FT(60) / FT(180) * pi
        lon_0 = FT(-90) / FT(180) * pi
        radlon = coord.long / FT(180) * pi
        stdev = FT(5) / FT(180) * pi
        anom = anom_ampl * exp(-((radlat - lat_0)^2 / 2stdev^2 + (radlon - lon_0)^2 / 2stdev^2))
        T_sfc = T_sfc_0 + anom
    end

    # prognostic variable
    Y = ClimaCore.Fields.FieldVector(T_sfc = T_sfc)

    return Y, space
end

# ode
function slab_rhs!(dY, Y, Ya, t)
    """
    Slab ocean:
    ∂_t T_sfc = F_sfc + G
    """
    p, F_sfc = Ya

    @. dY.T_sfc = (F_sfc) / (p.h * p.ρ * p.c)
end

struct SlabSimulation{P, Y, D, I}
    params::P
    Y_init::Y
    domain::D
    integrator::I
end

function slab_init(
    ::Type{FT},
    tspan;
    stepper = Euler(),
    nelements = 6,
    npolynomial = 4,
    dt = 0.02,
    saveat = 1.0e10,
    space = nothing,
) where {FT}

    params = ThermalSlabParameters(FT(0.5), FT(1500.0), FT(800.0), FT(281.0))

    Y, space = slab_space_init(FT, space, params)
    Ya = (params = params, F_sfc = ClimaCore.Fields.zeros(space)) #auxiliary
    problem = OrdinaryDiffEq.ODEProblem(slab_rhs!, Y, tspan, Ya)
    integrator = OrdinaryDiffEq.init(problem, stepper, dt = dt, saveat = saveat)

    SlabSimulation(params, Y, space, integrator)
end

get_slab_energy(slab_sim, T_sfc) = slab_sim.params.ρ .* slab_sim.params.c .* T_sfc .* slab_sim.params.h


include("auto_encoder_task.jl")

# ::: NETWORK
networks = Dict(
    "full_sigmoid" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
   "partial_sigmoid1" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "partial_sigmoid2" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
       ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "partial_sigmoid3" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
       ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]))

# ::: TRAINING
runtime_plotting = false
plot_intervall = 20
begin
    full_sigmoid_tm = [NaN for _ in 1:TRAIN_EPISODES];
    partial_sigmoid_tm1 = [NaN for _ in 1:TRAIN_EPISODES];
    partial_sigmoid_tm2 = [NaN for _ in 1:TRAIN_EPISODES];
    partial_sigmoid_tm3 = [NaN for _ in 1:TRAIN_EPISODES];

    println("---"); @time for e in 1:TRAIN_EPISODES
        t = time()

        z1 = forward!(networks["full_sigmoid"][1], train_data_x[e])
        z2 = forward!(networks["partial_sigmoid1"][1], train_data_x[e])
        z3 = forward!(networks["partial_sigmoid2"][1], train_data_x[e])
        z4 = forward!(networks["partial_sigmoid3"][1], train_data_x[e])

        y1 = forward!(networks["full_sigmoid"][2], z1)
        y2 = forward!(networks["partial_sigmoid1"][2], z2)
        y3 = forward!(networks["partial_sigmoid2"][2], z3)
        y4 = forward!(networks["partial_sigmoid3"][2], z4)

        dloss1 = -Δmse(train_data_x[e], y1)
        dloss2 = -Δmse(train_data_x[e], y2)
        dloss3 = -Δmse(train_data_x[e], y3)
        dloss4 = -Δmse(train_data_x[e], y4)

        dz1 = backward!(reverse(networks["full_sigmoid"][2]), dloss1, opt)
        _ = backward!(reverse(networks["full_sigmoid"][1]), dz1, opt)
        dz2 = backward!(reverse(networks["partial_sigmoid1"][2]), dloss2, opt)
        _ = backward!(reverse(networks["partial_sigmoid1"][1]), dz2, opt)
        dz3 = backward!(reverse(networks["partial_sigmoid2"][2]), dloss3, opt)
        _ = backward!(reverse(networks["partial_sigmoid2"][1]), dz3, opt)
        dz4 = backward!(reverse(networks["partial_sigmoid3"][2]), dloss4, opt)
        _ = backward!(reverse(networks["partial_sigmoid3"][1]), dz4, opt)

        full_sigmoid_tm[e] = mse(train_data_x[e], y1)
        partial_sigmoid_tm1[e] = mse(train_data_x[e], y2)
        partial_sigmoid_tm2[e] = mse(train_data_x[e], y3)
        partial_sigmoid_tm3[e] = mse(train_data_x[e], y4)

        # if runtime_plotting && e % plot_intervall == 0
        #     p = plot(FC_train_metric)
        #     display(p)
        # end

        # println("episode = $e : loss = $loss : dt = $(time() - t)")
    end
end

begin
    gr = 30
    plot([sum(full_sigmoid_tm[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (full diminish): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
    plot!([sum(partial_sigmoid_tm1[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (Relu [1] + diminish [3]): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
    plot!([sum(partial_sigmoid_tm2[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (Relu [2] + diminish [2]): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
    plot!([sum(partial_sigmoid_tm3[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (Relu [3] + diminish [1]): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
end

generate_images(10, "/Users/BBM2H16AHO/Documents/programming/DNN-Framework/auto_encoder/images/relu_sigmoid_test")


"""
    oddsratio(a, b, c, d; alpha=.05, exact=false)
    oddsratio(A, B; alpha=.05, exact=false)	# todo

Calculate odds ratio of 2x2 table (Case-Control Studies).
If A, B are vectors, a two-way table is generated using `crossTable` and OR is calculated.

Exact method is not implemented.

|             | case | non-cases |         |
|:------------|:-----|:----------|:--------|
| exposed     | a    | b         | a+b     |
| non-exposed | c    | d         | c+d     |
|             | a+c  | b+d       | a+b+c+d |

## Example outputs

```julia
B = 2×2 Named Array{Int64,2}
Exposed ╲ Case │  +   -
───────────────┼───────
+              │ 25  12
-              │  9  69
1×3 Named Array{Float64,2}
      ╲  │      OR    lower    upper
─────────┼──────────────────────────
OR 95%CI │ 15.9722  6.00717  42.4679
```
"""
function oddsratio(a, b, c, d; alpha=0.05, exact=false)
		if any(i==0 for i in [a b c d]) @error("Zero cell(s), cannot compute OR's variance") end
		B = NamedArray([a b; c d], ( ["+", "-"], ["+", "-"] ), ("Exposed", "Case"))
		rowname = string("OR ", string(100*(1-alpha)), "%CI")
		A = NamedArray([.0 .0 .0], ( [rowname], ["OR", "lower", "upper"] ), ("", ""))
		A[1] = (a * d) / (b * c)
		v = sum([1/i for i in [a, b, c, d]])
		z = quantile(Normal(), 1-alpha/2)
		A[2], A[3] = A[1] * [exp(-z*sqrt(v)) exp(z*sqrt(v))]
		@show B
		A
end

"""
    riskratio(a, b, c, d; alpha=.05)
    riskratio(A, B; alpha=.05, exact=false)	# todo

Calculate risk ratio (RR) of 2x2 table; this can also be used for calculating prevalence ratio (PR).

If A, B are vectors, a two-way table is generated using `crossTable` and RR is calculated.

|             | case | non-cases |         |
|:------------|:-----|:----------|:--------|
| exposed     | a    | b         | a+b     |
| non-exposed | c    | d         | c+d     |
|             | a+c  | b+d       | a+b+c+d |

## Example outputs

```julia
julia> riskratio(122,355,22,965, alpha=0.01)
B = 2×2 Named Array{Int64,2}
Exposed ╲ Case │   +    -
───────────────┼─────────
+              │ 122  355
-              │  22  965
1×3 Named Array{Float64,2}
        ╲  │      RR    lower    upper
───────────┼──────────────────────────
RR 99.0%CI │ 11.4746  6.43047  20.4752
```
"""
function riskratio(a, b, c, d; alpha=0.05)
		if any(i==0 for i in [a c]) error("Zero cell(s), cannot compute RR's variance") end
		B = NamedArray([a b; c d], ( ["+", "-"], ["+", "-"] ), ("Exposed", "Case"))
		rowname = string("RR ", string(100*(1-alpha)), "%CI")
		A = NamedArray([.0 .0 .0], ( [rowname], ["RR", "lower", "upper"] ), ("", ""))
		p11, p12 = (a/(a+b), c/(c+d))
		A[1] = p11 / p12
		v = (1-p11)/a + (1-p12)/c
		z = quantile(Normal(), 1-alpha/2)
		A[2], A[3] = A[1] * [exp(-z*sqrt(v)) exp(z*sqrt(v))]
		@show B
		A
end


function computeRelativityCorrection(x::Float64,y::Float64,z::Float64,
    vx::Float64,vy::Float64,vz::Float64,)
    relcorr = -2.0 * (x * vx + y * vy + z * vz) / C_MPS^2
    return relcorr
end

function getSidTime(dt::DateTime)
    jd = datetime2julian(DateTime(year(dt),
                            month(dt), day(dt)))
    jc = (jd - 2451545.0) / 36525.0
    sid = 24110.54841 + jc * (8640184.812866 +
                            jc * (0.093104 -
                                jc * 0.0000062 ) )
    sid /= 3600.0
    sid = sid % 24.0
    if sid < 0.0 sid += 24.0 end
    return sid
end

function secondofday(dt::DateTime)
    return hour(dt) * 3600.0 + minute(dt) * 60.0 +
                second(dt) + millisecond(dt) / 100000.0
end

function derivative(inState::Array{Float64,1},
                    accel::Array{Float64,1},
                    j20::Float64=PZ90Ellipsoid.j20,
                    mu::Float64=PZ90Ellipsoid.gm_km,
                    ae::Float64=PZ90Ellipsoid.a_km,)
    x  = inState[1]
    y  = inState[3]
    z  = inState[5]
    r² = x^2 + y^2 + z^2
    r  = sqrt(r²)
    xmu = mu / r²
    rho = ae / r
    xr  = x / r
    yr  = y / r
    zr  = z / r
    zr2 = zr^2
    k1  = j20 * xmu * 1.5 * rho^2
    cm  = k1 * (1.0 - 5.0 * zr2)
    cmz = k1 * (3.0 - 5.0 * zr2)
    k2  = cm - xmu

    gloAx = k2 * xr + accel[1]
    gloAy = k2 * yr + accel[2]
    gloAz = (cmz - xmu) * zr + accel[3]

    dxt = zeros(6)

    dxt[1] = inState[2]
    dxt[2] = gloAx
    dxt[3] = inState[4]
    dxt[4] = gloAy
    dxt[5] = inState[6]
    dxt[6] = gloAz
    dxt
end

function calculate_Fit_pos(epoch::DateTime,
        ephTime::DateTime,
        clkbias::Float64, # tau
        clkdrift::Float64, # gamma
        ix::Float64,iy::Float64,iz::Float64,
        ivx::Float64,ivy::Float64,ivz::Float64,
        iax::Float64,iay::Float64,iaz::Float64,
        health::Float64=-1.0,
        step::Float64=1.0,
        we::Float64=PZ90Ellipsoid.ge.angVelocity,
)
    elapte = (epoch - ephTime).value / 1000.0 # TODO  second?
    # Values to be returned will be stored here
    # If the exact epoch is found, let's return the values
    if epoch == ephTime        # exact match for epoch
        oxx = ix * 1.e3;   # m
        oxy = iy * 1.e3;   # m
        oxz = iz * 1.e3;   # m
        ovx = ivx * 1.e3;  # m/sec
        ovy = ivy * 1.e3;  # m/sec
        ovz = ivz * 1.e3;  # m/sec
    # In the GLONASS system, 'clkbias' already includes the
    # relativistic correction, therefore we must substract the late
    # from the former.
        relcorr = computeRelativityCorrection(oxx, oxy, oxz, ovx, ovy, ovz)
        clkbias = clkbias + clkdrift * elapte - relcorr;
        clkdrift = clkdrift;
        frame = fmPZ90;

    # We are done, let's return
        return Xvt(ECEF(oxx, oxy, oxz), (ovx, ovy, ovz),
    relcorr,clkbias,clkdrift,frame,
    health == 1.0 ? hsHealthy : hsUnhealthy)
    end

    # Get the data out of the GloRecord structure
    px = ix   # X coordinate (km)
    vx = ivx   # X velocity   (km/s)
    ax = iax   # X acceleration (km/s^2)

    py = iy   # Y coordinate
    vy = ivy   # Y velocity
    ay = iay   # Y acceleration

    pz = iz   # Z coordinate
    vz = ivz   # Z velocity
    az = iaz   # Z acceleration

    # We will need some PZ-90 ellipsoid parameters

    # Get sidereal time at Greenwich at 0 hours UT

    gst = getSidTime(ephTime)
    s0 =  gst * π / 12.0
    numSeconds = secondofday(ephTime)

    s =  s0 + we * numSeconds
    cs =  cos(s)
    ss =  sin(s)

    # Initial state matrix
    initialState = zeros(6)
    accel = zeros(3)
    # dxt1  = zeros(6)
    # dxt2  = zeros(6)
    # dxt3  = zeros(6)
    # dxt4  = zeros(6)
    tempRes = zeros(6)

    # Get the reference state out of GloEphemeris object data. Values
    # must be rotated from PZ-90 to an absolute coordinate system
    # Initial x coordinate (m)
    initialState[1]  = (px * cs - py * ss);
    # Initial y coordinate
    initialState[3]  = (px * ss + py * cs);
    # Initial z coordinate
    initialState[5]  = pz;

    # Initial x velocity   (m/s)
    initialState[2]  = (vx * cs - vy * ss - we * initialState[3] );
    # Initial y velocity
    initialState[4]  = (vx * ss + vy * cs + we * initialState[1] );
    # Initial z velocity
    initialState[6]  = vz;

    # Integrate satellite state to desired epoch using the given step
    rkStep = step

    if epoch < ephTime
        rkStep = step * (-1.0)
    end

    rkStepNs = Nanosecond(rkStep * 1e9)
    workEpoch = TimeDate(ephTime)
    tolerance = Nanosecond(1) # 1e-9

    while true
    # If we are about to overstep, change the stepsize appropriately
    # to hit our target final time.
        if rkStep > 0.0
            if (workEpoch + rkStepNs) > epoch
                rkStepNs = epoch - workEpoch
                rkStep = rkStepNs.value / 1e9
            end
        else
            if (workEpoch + rkStepNs) < epoch
                rkStepNs = (epoch - workEpoch)
                rkStep = rkStepNs.value / 1e9
            end
        end

        numSeconds += rkStep;
        s = s0 + we * ( numSeconds );
        cs = cos(s);
        ss = sin(s);

        # Accelerations are computed once per iteration
        accel[1] = ax * cs - ay * ss
        accel[2] = ax * ss + ay * cs
        accel[3] = az
        dxt1 = derivative(initialState, accel)
        for j = 1:6
            tempRes[j] = initialState[j] + rkStep * dxt1[j] / 2.0
        end
        dxt2 = derivative(tempRes, accel)
        for j = 1:6
            tempRes[j] = initialState[j] + rkStep * dxt2[j] / 2.0
        end
        dxt3 = derivative(tempRes, accel)
        for j = 1:6
            tempRes[j] = initialState[j] + rkStep * dxt3[j];
        end
        dxt4 = derivative(tempRes, accel)
        for j = 1:6
            initialState[j] = initialState[j] + rkStep *
                            ( dxt1[j] +
                                2.0 * ( dxt2[j] + dxt3[j] ) +
                                dxt4[j] ) / 6.0;
        end
        # @bp
        # If we are within tolerance of the target time, we are done.
        workEpoch += rkStepNs
        if ( abs(epoch - workEpoch) < tolerance )
            break
        end
    end
    # End of 'while (!done)...'
    px = initialState[1]
    py = initialState[3]
    pz = initialState[5]
    vx = initialState[2]
    vy = initialState[4]
    vz = initialState[6]
    # @bp

    oxx = 1000.0 * ( px * cs + py * ss );         # X coordinate
    oxy = 1000.0 * (-px * ss + py * cs);          # Y coordinate
    oxz = 1000.0 * pz;                        # Z coordinate

    ovx = 1000.0 * ( vx * cs + vy * ss + we * (oxy / 1000.0) ) # X velocity
    ovy = 1000.0 * (-vx * ss + vy * cs - we * (oxx / 1000.0) ) # Y velocity
    ovz = 1000.0 * vz;                        # Z velocity

    # In the GLONASS system, 'clkbias' already includes the relativistic
    # correction, therefore we must substract the late from the former.
    ox = ECEF(oxx, oxy, oxz)

    relcorr = computeRelativityCorrection(oxx, oxy, oxz, ovx, ovy, ovz)
    clkbias = clkbias + clkdrift * elapte - relcorr
    clkdrift = clkdrift;
    frame = fmPZ90;

    # We are done, let's return
    return Xvt(ECEF(oxx, oxy, oxz), (ovx, ovy, ovz),
                    relcorr,clkbias,clkdrift,frame,
                    health == 1.0 ? hsHealthy : hsUnhealthy)

end

function calculate_Fit_pos(ees::EphEpochStore, dt::DateTime,
                            vaild_d::Minute=Minute(15),)
    k = ees.keys
    if length(k) == 0
        return ECEF{Float64}(NaN, NaN, NaN)
    end
    d = k .- dt
    i = argmin(abs.(d))

    md = d[i]
    if md > vaild_d || md < -vaild_d
        return ECEF{Float64}(NaN, NaN, NaN)
    end

    mk = k[i]
    me = ees.eph[mk].d

    xvt = calculate_Fit_pos(dt, mk,
                        me[1], me[2],

                        me[4], me[8],  me[12],
                        me[5], me[9],  me[13],
                        me[6], me[10], me[14],

                        me[7],)
    return xvt.x
end


@testset "SOCP Atoms: " for solver in solvers
    if can_solve_socp(solver)
        @testset "norm 2 atom" begin
            x = Variable(2, 1)
            A = [1 2; 2 1; 3 4]
            b = [2; 3; 4]
            p = Problem{BigFloat}(:minimize, norm2(A * x + b))
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 0.64888 atol=TOL
            @test evaluate(norm2(A * x + b)) ≈ 0.64888 atol=TOL

            x = Variable(2, 1)
            A = [1 2; 2 1; 3 4]
            b = [2; 3; 4]
            lambda = 1
            p = Problem{BigFloat}(:minimize, norm2(A * x + b) + lambda * norm2(x), x >= 1)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 14.9049 atol=TOL
            @test evaluate(norm2(A * x + b) + lambda * norm2(x)) ≈ 14.9049 atol=TOL

            x = Variable(2)

            p = Problem{BigFloat}(:minimize, norm2([x[1] + 2x[2] + 2; 2x[1] + x[2] + 3; 3x[1]+4x[2] + 4]) + lambda * norm2(x), x >= 1)
            @test vexity(p) == ConvexVexity()

            solve!(p, solver)
            @test p.optval ≈ 14.9049 atol=TOL
            @test evaluate(norm2(A * x + b) + lambda * norm2(x)) ≈ 14.9049 atol=TOL

            x = Variable(2, 1)
            A = [1 2; 2 1; 3 4]
            b = [2; 3; 4]
            lambda = 1
            p = Problem{BigFloat}(:minimize, norm2(A * x + b) + lambda * norm_1(x), x >= 1)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 15.4907 atol=TOL
            @test evaluate(norm2(A * x + b) + lambda * norm_1(x)) ≈ 15.4907 atol=TOL
        end

        @testset "frobenius norm atom" begin
            m = Variable(4, 5)
            c = [m[3, 3] == 4, m >= 1]
            p = Problem{BigFloat}(:minimize, norm(vec(m), 2), c)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ sqrt(35) atol=TOL
            @test evaluate(norm(vec(m), 2)) ≈ sqrt(35) atol=TOL
        end

        @testset "quad over lin atom" begin
            x = Variable(3, 1)
            A = [2 -3 5; -2 9 -3; 5 -8 3]
            b = [-3; 9; 5]
            c = [3 2 4]
            d = -3
            p = Problem{BigFloat}(:minimize, quadoverlin(A*x + b, c*x + d))
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 17.7831 atol=TOL
            @test (evaluate(quadoverlin(A * x + b, c * x + d)))[1] ≈ 17.7831 atol=TOL
        end

        @testset "sum squares atom" begin
            x = Variable(2, 1)
            A = [1 2; 2 1; 3 4]
            b = [2; 3; 4]
            p = Problem{BigFloat}(:minimize, sumsquares(A*x + b))
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 0.42105 atol=TOL
            @test (evaluate(sumsquares(A * x + b)))[1] ≈ 0.42105 atol=TOL
        end

        @testset "square atom" begin
            x = Variable(2, 1)
            A = [1 2; 2 1; 3 4]
            b = [2; 3; 4]
            p = Problem{BigFloat}(:minimize, sum(square(A*x + b)))
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 0.42105 atol=TOL
            @test evaluate(sum(square(A * x + b))) ≈ 0.42105 atol=TOL

            x = Variable(2, 1)
            A = [1 2; 2 1; 3 4]
            b = [2; 3; 4]
            expr = A * x + b
            p = Problem{BigFloat}(:minimize, sum(dot(^)(expr,2))) # elementwise ^
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 0.42105 atol=TOL
            @test evaluate(sum(broadcast(^, expr, 2))) ≈ 0.42105 atol=TOL

            p = Problem{BigFloat}(:minimize, sum(dot(*)(expr, expr))) # elementwise *
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 0.42105 atol=TOL
            @test evaluate(sum((dot(*))(expr, expr))) ≈ 0.42105 atol=TOL
        end

        @testset "inv pos atom" begin
            x = Variable(4)
            p = Problem{BigFloat}(:minimize, sum(invpos(x)), invpos(x) < 2, x > 1, x == 2, 2 == x)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 2 atol=TOL
            @test evaluate(sum(invpos(x))) ≈ 2 atol=TOL

            x = Variable(3)
            p = Problem{BigFloat}(:minimize, sum(dot(/)([3,6,9], x)), x<=3)
            solve!(p, solver)
            @test x.value ≈ fill(3.0, (3, 1)) atol=TOL
            @test p.optval ≈ 6 atol=TOL
            @test evaluate(sum((dot(/))([3, 6, 9], x))) ≈ 6 atol=TOL

            x = Variable()
            p = Problem{BigFloat}(:minimize, sum([3,6,9]/x), x<=3)
            solve!(p, solver)
            @test x.value ≈ 3 atol=TOL
            @test p.optval ≈ 6 atol=TOL
            @test evaluate(sum([3, 6, 9] / x)) ≈ 6 atol=TOL
        end

        @testset "geo mean atom" begin
            x = Variable(2)
            y = Variable(2)
            p = Problem{BigFloat}(:minimize, geomean(x, y), x >= 1, y >= 2)
            # not DCP compliant
            @test vexity(p) == ConcaveVexity()
            p = Problem{BigFloat}(:maximize, geomean(x, y), 1 < x, x < 2, y < 2)
            # Just gave it a vector as an objective, not okay
            @test_throws Exception solve!(p, solver)

            p = Problem{BigFloat}(:maximize, sum(geomean(x, y)), 1 < x, x < 2, y < 2)
            solve!(p, solver)
            @test p.optval ≈ 4 atol=TOL
            @test evaluate(sum(geomean(x, y))) ≈ 4 atol=TOL
        end

        @testset "sqrt atom" begin
            x = Variable()
            p = Problem{BigFloat}(:maximize, sqrt(x), 1 >= x)
        end

        @testset "quad form atom" begin
            x = Variable(3, 1)
            A = [0.8608 0.3131 0.5458; 0.3131 0.8584 0.5836; 0.5458 0.5836 1.5422]
            p = Problem{BigFloat}(:minimize, quadform(x, A), [x >= 1])
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 6.1464 atol=TOL
            @test (evaluate(quadform(x, A)))[1] ≈ 6.1464 atol=TOL

            x = Variable(3, 1)
            A = -1.0*[0.8608 0.3131 0.5458; 0.3131 0.8584 0.5836; 0.5458 0.5836 1.5422]
            c = [3 2 4]
            p = Problem{BigFloat}(:maximize, c*x , [quadform(x, A) >= -1])
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 3.7713 atol=TOL
            @test (evaluate(quadform(x, A)))[1] ≈ -1 atol=TOL
        end

        @testset "huber atom" begin
            x = Variable(3)
            p = Problem{BigFloat}(:minimize, sum(huber(x, 1)), x >= 2)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 9 atol=TOL
            @test evaluate(sum(huber(x, 1))) ≈ 9 atol=TOL
        end

        @testset "rational norm atom" begin
            A = [1 2 3; -1 2 3]
            b = A * ones(3)
            x = Variable(3)
            p = Problem{BigFloat}(:minimize, norm(x, 4.5), [A * x == b])
            @test vexity(p) == ConvexVexity()
            # Solution is approximately x = [1, .93138, 1.04575]
            solve!(p, solver)
            @test p.optval ≈ 1.2717 atol=TOL
            @test evaluate(norm(x, 4.5)) ≈ 1.2717 atol=TOL
        end

        @testset "rational norm dual norm" begin
            v = [0.463339, 0.0216084, -2.07914, 0.99581, 0.889391]
            x = Variable(5)
            q = 1.379;  # q norm constraint that generates many inequalities
            qs = q / (q - 1);  # Conjugate to q
            p = Problem{BigFloat}(:minimize, x' * v)
            p.constraints += (norm(x, q) <= 1)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)  # Solution is -norm(v, q / (q - 1))
            @test p.optval ≈ -2.144087 atol=TOL
            @test sum(evaluate(x' * v)) ≈ -2.144087 atol=TOL
            @test evaluate(norm(x, q)) ≈ 1 atol=TOL
            @test sum(evaluate(x' * v)) ≈ -(sum(abs.(v) .^ qs) ^ (1 / qs)) atol=TOL
        end

        @testset "rational norm atom sum" begin
            A = [-0.719255  -0.229089
                -1.33632   -1.37121
                0.703447  -1.4482]
            b = [-1.82041, -1.67516, -0.866884]
            q = 1.5
            xvar = Variable(2)
            p = Problem{BigFloat}(:minimize, .5 * sumsquares(xvar) + norm(A * xvar - b, q))
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            # Compute gradient, check it is zero(ish)
            x_opt = xvar.value
            margins = A * x_opt - b
            qs = q / (q - 1);  # Conjugate
            denom = sum(abs.(margins).^q)^(1/qs)
            g = x_opt + A' * (abs.(margins).^(q-1) .* sign.(margins)) / denom
            @test p.optval ≈ 1.7227 atol=TOL
            @test norm(g, 2) ^ 2 ≈ 0 atol=TOL
        end

        @testset "norm consistent with Base for matrix variables" begin
            A = randn(4, 4)
            x = Variable(4, 4)
            x.value = A
            # Matrix norm
            @test evaluate(opnorm(x)) ≈ opnorm(A) atol=TOL
            @test evaluate(opnorm(x, 1)) ≈ opnorm(A, 1) atol=TOL
            @test evaluate(opnorm(x, 2)) ≈ opnorm(A, 2) atol=TOL
            @test evaluate(opnorm(x, Inf)) ≈ opnorm(A, Inf) atol=TOL
            # Vector norm
            # TODO: Once the deprecation for norm on matrices is removed, remove the `vec` calls
            @test evaluate(norm(vec(x), 1)) ≈ norm(vec(A), 1) atol=TOL
            @test evaluate(norm(vec(x), 2)) ≈ norm(vec(A), 2) atol=TOL
            @test evaluate(norm(vec(x), 7)) ≈ norm(vec(A), 7) atol=TOL
            @test evaluate(norm(vec(x), Inf)) ≈ norm(vec(A), Inf) atol=TOL
        end

        @testset "Fixed and freed variables" begin
            @testset "fix and free addition" begin
                x = Variable()
                y = Variable()

                p = Problem{BigFloat}(:minimize, x+y, x>=0, y>=0)
                solve!(p, solver)
                @test p.optval ≈ 0 atol=TOL

                y.value = 4
                fix!(y)
                solve!(p, solver)
                @test p.optval ≈ 4 atol=TOL

                free!(y)
                solve!(p, solver)
                @test p.optval ≈ 0 atol=TOL
            end

            @testset "fix multiplication" begin
                a = [1,2,3,2,1]
                x = Variable(length(a))
                gamma = Variable(Positive())
                fix!(gamma, 0.7)

                p = Problem{BigFloat}(:minimize, norm(x-a) + gamma*norm(x[1:end-1] - x[2:end]))
                solve!(p, solver)
                o1 = p.optval
                # x should be very close to a
                @test o1 ≈ 0.7 * norm(a[1:end - 1] - a[2:end]) atol=TOL
                # increase regularization
                fix!(gamma, 1.0)
                solve!(p, solver)
                o2 = p.optval
                # x should be very close to mean(a)
                @test o2 ≈ norm(a .- mean(a)) atol=TOL

                @test o1 <= o2
            end
        end
    end
end


println("Welcome to Julia, Julian!")

arithmetic_mean(x) = sum(x) / length(x)

# Arithmetic Mean
avg(x) = (ones(length(x))'x) / length(x)

# Norm of a vector
norm(x) = sqrt(x'x)

# Standard Deviation
std(x) = norm(x .- avg(x)) / sqrt(length(x))

# Quadratic Mean
rms(x) = sqrt(avg(x)^2 + std(x)^2)

# Distance function
distance(x, y) = norm(x - y)

head = first
tail = last


using DynamicIterators
using Distributions

D = MvNormal([1.0, 0.5], [1.0 0.5; 0.5 1.5] )
struct Move{T}
    x::T
    σ::Float64
    i::Int
end
m1(x) = Move(x, 0.1, 1)
m2(x) = Move(x, 0.1, 2)
Base.rand(M::Move) = M.x + M.σ*randn()*[M.i-1, 2-M.i]
Distributions.logpdf(M::Move, x) = logpdf(Normal(M.x[M.i], M.σ), x[M.i])
MH1 = MetropolisHastings(D, m1, logpdf)
MH2 = MetropolisHastings(D, m2, logpdf)

I = Evolve(i->rand(1:2))

MH = mixture(I, (MH1, MH2))

X = values(trace(MH, 1=>(1, [0.0, 0.0]), endtime(2000)))


# Solve a SAA of the resource allocation model given scenarios of the demands

function solveSAAModel(demand_scen::Array{Float64,2})

	numScenarios::Int64 = size(demand_scen,1)

	mod = Model(solver=GurobiSolver(gurobi_env,Presolve=0,OutputFlag=0))

	@variable(mod, z[1:numResources] >= 0)
	@variable(mod, v[1:numResources,1:numCustomers,1:numScenarios] >= 0)
	@variable(mod, w[1:numCustomers,1:numScenarios] >= 0)

	@objective(mod, Min, sum(costVector[i]*z[i] for i=1:numResources) + (1.0/numScenarios)*sum(recourseCostCoeff[j]*w[j,s] for j=1:numCustomers for s=1:numScenarios))

	@constraint(mod, [i=1:numResources, s=1:numScenarios], sum(v[i,j,s] for j = 1:numCustomers) <= rho[i]*z[i])
	@constraint(mod, [j=1:numCustomers, s=1:numScenarios], sum(mu[i,j]*v[i,j,s] for i = 1:numResources) + w[j,s] >= demand_scen[s,j])

	status = solve(mod)
	z_soln = getvalue(z)
	objValue_soln::Float64 = getobjectivevalue(mod)

	return z_soln, objValue_soln
end


"Problem 5 in the Hock-Schittkowski suite"
function hs5()

  nlp = Model()

  l = [-1.5, -3]
  u = [4, 3]
  @variable(nlp, l[i] ≤ x[i=1:2] ≤ u[i], start=0.0)

  @NLobjective(
    nlp,
    Min,
    sin(x[1] + x[2]) + (x[1] - x[2])^2 - 1.5 * x[1] + 2.5 * x[2] + 1
  )

  return nlp
end


using Catlab.CategoricalAlgebra, Catlab.Graphs, Catlab.Present, Catlab.Graphics, Catlab.Theories, Catlab.Present
using Catlab.CategoricalAlgebra.FinCats: FinCatGraphEq

@present ThSIR(FreeSchema) begin
    (S,I,R)::Ob
end

"""This code will work for any choice of 'type graph'"""
function update_state(state::StructACSet{S}) where S
    rules = []
    while !isempty(rules)
        state = apply_rule(state, pop!(rules))
    end
end

@acset_type SIR(ThSIR)

# infect = Rule("name", L, R)

I = @acset SIR begin I=1 end
I2 = @acset SIR begin I=2 end
SI = @acset SIR begin S=1; I=1 end
# L_infect =  homomorphism(I, SI) # equivalent to below b/c only one morphism
L_infect =  ACSetTransformation(I, SI; I=[1])
R_infect = ACSetTransformation(I, I2; I=[1])

ex_state = @acset SIR begin S=10; I=3; R=1 end

# rewrite(L_infect, R_infect, ex_state) # execute immediately

match = homomorphism(SI, ex_state)

matches = [match_1, match_2] # pointing to ex_state

"""
L    I    R
|
G <- K -> H (updated G)
   kg  kh

"""

_, kg, _, kh = rewrite_match_maps(L_infect, R_infect, match)

# pretend state_2 is a homomorphism from state_1 -> state_2

matches = [match for match in matches]

state_3 = rewrite_match(L_infect, R_infect, match_2)

# # type 2 --------------------------------------------------
# @present ThSIR2(FreeSchema) begin
#     State::Ob
#     Agents::Ob
#     state::Hom(Agents, State)
# end

# @acset_type SIR2(ThSIR2)

# state0 = vcat(fill.(1:3, [5,4,1])...)

# SIR2_state = @acset SIR2 begin State=3; Agents=10; state=state0 end

# I = @acset SIR2 begin Agents=1; State=3; state=2 end
# I2 = @acset SIR2 begin Agents=2; State=3; state=[2,2] end
# SI = @acset SIR2 begin Agents=2; State=3; state=[1,2] end

# # L_infect =  homomorphism(I, SI) # equivalent to below b/c only one morphism
# homomorphism(I, SI)
# homomorphism(I, I2)

# L_infect =  ACSetTransformation(I, SI; Agents=[1])
# R_infect = ACSetTransformation(I, I2)


"""
    Apply Laplace method for a given decision matrix

"""
function laplace(decisionMatrix::DataFrame)::LaplaceResult

    n, p = size(decisionMatrix)

    probs = [1.0 / p for i in 1:p]

    expecteds = zeros(Float64, n)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    for i in 1:n
        expecteds[i] = (m[i,:] .* probs)  |> sum
    end

    bestIndex = sortperm(expecteds) |> last

    result = LaplaceResult(
        expecteds,
        bestIndex
    )

    return result
end

"""
    Apply maximin method for a given decision matrix

"""
function maximin(decisionMatrix::DataFrame)::MaximinResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmins = rowmins(decisionMatrix)

    bestIndex = sortperm(rmins) |> last

    result = MaximinResult(
        rmins,
        bestIndex
    )

    return result
end

"""
    Apply maximax method for a given decision matrix

"""
function maximax(decisionMatrix::DataFrame)::MaximaxResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmaxs = rowmaxs(decisionMatrix)

    bestIndex = sortperm(rmaxs) |> last

    result = MaximaxResult(
        rmaxs,
        bestIndex
    )

    return result
end

"""
    Apply minimax method for a given decision matrix

"""
function minimax(decisionMatrix::DataFrame)::MinimaxResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmaxs = rowmaxs(decisionMatrix)

    bestIndex = sortperm(rmaxs) |> first

    result = MinimaxResult(
        rmaxs,
        bestIndex
    )

    return result
end

"""
    Apply minimin method for a given decision matrix

"""
function minimin(decisionMatrix::DataFrame)::MiniminResult

    n, p = size(decisionMatrix)

    # m = convert(Array{Float64,2}, decisionMatrix)
    m = Matrix{Float64}(decisionMatrix)

    rmins = rowmins(decisionMatrix)

    bestIndex = sortperm(rmins) |> first

    result = MiniminResult(
        rmins,
        bestIndex
    )

    return result
end

"""
    Apply Savage method for a given decision matrix

"""
function savage(decisionMatrix::DataFrame)::SavageResult

    n, p = size(decisionMatrix)

    cmaxs = colmaxs(decisionMatrix)

    newDecisionMatrix = similar(decisionMatrix)

    @inbounds for i in 1:p
        newDecisionMatrix[:, i] = cmaxs[i] .- decisionMatrix[:, i]
    end

    rmaxs = scores = rowmaxs(newDecisionMatrix)

    bestIndex = rmaxs |> sortperm |> first

    result = SavageResult(
        newDecisionMatrix,
        scores,
        bestIndex
    )

    return result
end

"""
    Apply Hurwicz method for a given decision matrix

"""
function hurwicz(decisionMatrix::DataFrame; alpha::Float64=0.5)::HurwiczResult

    n, p = size(decisionMatrix)

    rmaxs = rowmaxs(decisionMatrix)
    rmins = rowmins(decisionMatrix)

    scores = alpha .* rmaxs .+ (1.0 - alpha) .* rmins

    bestIndex = scores |> sortperm |>  last

    result = HurwiczResult(
        scores,
        bestIndex
    )

    return result

end

"""
    Apply MLE (Maximum Likelihood) method for a given decision matrix

"""
function mle(decisionMatrix::DataFrame, weights::Array{Float64,1})::MLEResult

    w = unitize(weights)

    weightedMatrix = w * decisionMatrix

    scores = rowsums(weightedMatrix)

    bestIndex = scores |> sortperm |> last

    result = MLEResult(
        scores,
        bestIndex
    )

    return result

end

"""
    Apply Expected Regret method for a given decision matrix

"""
function expectedregret(decisionMatrix::DataFrame, weights::Array{Float64,1})

    _, p = size(decisionMatrix)

    w = unitize(weights)

    cmaxs = colmaxs(decisionMatrix)

    regretmat = similar(decisionMatrix)

    for i in 1:p
        regretmat[:, i] = cmaxs[i] .- decisionMatrix[:, i]
    end

    weightedregret = w * regretmat

    scores = rowsums(weightedregret)

    bestIndex = scores |> sortperm |> first

    result = ExpectedRegretResult(
        scores,
        bestIndex
    )

    return result
end


# Now, you need to figure out how to pilot this thing.
#
# It seems like the submarine can take a series of commands like forward 1,
# down 2, or up 3:
#
# forward X increases the horizontal position by X units.
# down X increases the depth by X units.
# up X decreases the depth by X units.
#
# Note that since you're on a submarine, down and up affect your depth, and so
# they have the opposite result of what you might expect.
#
# The submarine seems to already have a planned course (your puzzle input). You
# should probably figure out where it's going. For example:
#
# forward 5
# down 5
# forward 8
# up 3
# down 8
# forward 2
#
# Your horizontal position and depth both start at 0. The steps above would
# then modify them as follows:
#
# forward 5 adds 5 to your horizontal position, a total of 5.
# down 5 adds 5 to your depth, resulting in a value of 5.
# forward 8 adds 8 to your horizontal position, a total of 13.
# up 3 decreases your depth by 3, resulting in a value of 2.
# down 8 adds 8 to your depth, resulting in a value of 10.
# forward 2 adds 2 to your horizontal position, a total of 15.
#
# After following these instructions, you would have a horizontal position of
# 15 and a depth of 10. (Multiplying these together produces 150.)
#
# Calculate the horizontal position and depth you would have after following
# the planned course. What do you get if you multiply your final horizontal
# position by your final depth?

dirs = map(eachline("input.txt")) do l
  (dir, amt) = split(l, ' ')
  amt = parse(Int, amt)
  if dir == "forward"
    (0, amt)
  elseif dir == "up"
    (-amt, 0)
  else
    (amt, 0)
  end
end
(depth, horz) = reduce(.+, dirs)
println("Result: ", depth * horz)


import SpecialFunctions: digamma, gamma
using Neighborhood: Theiler, KDTree, bulksearch

abstract type NearestNeighborEntropyEstimator <: EntropyEstimator end

function maximum_neighbor_distances(pts, est::EntropyEstimator)
    theiler = Theiler(est.w)
    est.w >= 0 || error("w, the number of neighbors to exclude, must be >= 0")
    tree = KDTree(pts) # Euclidean metric forced
    idxs, dists = bulksearch(tree, pts.data, NeighborNumber(est.k), theiler)
    # Distance to k-th nearest neighbor for each of the points
    return [d[end] for d in dists]
end

"Volume of a unit ball in R^d."
ball_volume(d::Int) = π^(d/2)/gamma((d/2)+1)

include("KozachenkoLeonenko.jl")
include("Kraskov.jl")


#=
Find the unique positive integer whose square has the form 1_2_3_4_5_6_7_8_9_0,
where each “_” is a single digit.
=#

#TODO: Try iterating by calculating the squares as permutations then grabbing sqrts.
function calc()
  # This is kinda icky.
  for i in 1238902700:3238902700
    s = i^2
    if length("$s") != 19
      continue
    end
    if occursin(r"1\d2\d3\d4\d5\d6\d7\d8\d9\d0", "$s")
      return round(Int64, sqrt(s))
    end
  end
end
@time println(calc())


using Bolt
using ForwardDiff
using PyPlot
using BenchmarkTools

# Cₗ as a function of baryon density
function clb(Ω_b::DT, ells) where DT
    par = CosmoParams{DT}(Ω_b=Ω_b)
    bg = Background(par)
    ih = IonizationHistory(Peebles(), par, bg)
    k_grid = quadratic_k(0.1bg.H₀, 1000bg.H₀, 500)
    sf = source_grid(par, bg, ih, k_grid, BasicNewtonian())
    return cltt(ells, par, bg, ih, sf)
end

ells = 10:2:1200
f(Ω_b) = clb(Ω_b, ells)
@time cl = f(0.046)
@time ∂cl = ForwardDiff.derivative(f, 0.046)  # you can just ForwardDiff the whole thing

##
clf()
plt.figure(figsize=(10,5))
plot(ells, cl .* ells.^2, "-", label=raw"$C_{\ell}$")
plot(ells, ∂cl .* ells.^2 / 10, "-",
    label=raw"$\partial C_{\ell}/\partial\Omega_b / 10$")

# uncomment to see finite differences
# Δ = 1e-3
# @time finitediff_∂cl = (f(0.046 + Δ) .- f(0.046 - Δ)) ./ 2Δ
# plt.plot(ells, finitediff_∂cl .* ells.^2 / 10, "--",
#     label=raw"$\Delta C_{\ell}/\Delta\Omega_b / 10$")

ylabel(raw"$\ell^2 C_{\ell}^{TT}$")
xlabel(raw"$\ell$")
legend()
ylim(-0.3, 0.5)
tight_layout()
# savefig("docs/assets/example_spectrum.png")
gcf()

##
sf

##


# Test variational information

using Test
using Clustering

@testset "varinfo() (variational information)" begin

a1 = [1, 1, 1, 2, 2, 2, 3, 3, 3, 3]
a2 = [1, 1, 1, 1, 2, 2, 2, 2, 2, 2]

@test varinfo(a1, a1) ≈ 0.0 atol=1.0e-12
@test varinfo(a2, a2) ≈ 0.0 atol=1.0e-12

v = varinfo(a1, a2)
v_ = varinfo(a2, a1)
@test 0.0 < v < log(3)
@test v ≈ v_

a1 = [1, 2, 3, 4, 5]
a2 = [1, 1, 1, 1, 1]

@test varinfo(a1, a2) ≈ log(5)
@test varinfo(a2, a1) ≈ log(5)

a1 = [1, 1, 1, 2, 2, 2]
a2 = [2, 2, 2, 1, 1, 1]
@test varinfo(a1, a2) ≈ 0.0 atol=1.0e-12

end


module PointwiseKDEs

export PointwiseKDE

using Distributions
using LinearAlgebra
using StatsFuns
using Random

struct PointwiseKDE{T <: Real} <: ContinuousMultivariateDistribution
    pts::Matrix{T}
    CF::Cholesky{T, Matrix{T}}
end

"""
    PointwiseKDE(points::Matrix{T <: Real})

Construct a Gaussian KDE from the given matrix of points.

`points` should have size `(ndim, npts)`.  Bandwidth selection is Scott's rule.

The result is a `MultivariateDistribution` object suitable for use with `rand`
or `logpdf` or ....
"""
function PointwiseKDE(pts::Matrix{T}) where T <: Real
    nd, np = size(pts)
    C = cov(pts, dims=2) ./ np^(2/(4+nd))
    CF = cholesky(C)

    PointwiseKDE(pts, CF)
end

function Distributions.length(d::PointwiseKDE)
    size(d.pts, 1)
end

function Distributions.sampler(d::PointwiseKDE)
    d
end

function Distributions.eltype(d::PointwiseKDE{T}) where T <: Number
    T
end

function Distributions._rand!(rng::AbstractRNG, d::PointwiseKDE{T}, x::AbstractArray{T}) where T <: Real
    nd, np = size(d.pts)
    i = rand(rng, 1:np)

    pt = vec(d.pts[:,i])
    cpt = d.CF.L*randn(rng, nd)

    for i in eachindex(x)
        x[i] = pt[i] + cpt[i]
    end
    x
end

function Distributions._rand!(rng::AbstractRNG, d::PointwiseKDE{T}, x::AbstractMatrix{T}) where T <: Real
    nd, np = size(x)

    for j in 1:np
        x[:,j] = Distributions._rand!(rng, d, x[:,j])
    end

    x
end

function Distributions._logpdf(d::PointwiseKDE{T}, x::AbstractArray{T}) where T <: Real
    nd, np = size(d.pts)

    logp::T = zero(T) - Inf
    for i in 1:np
        r = x .- vec(d.pts[:,i])
        ru = d.CF.L \ r

        logp = logaddexp(logp, -0.5*dot(ru, ru))
    end

    logp - sum(0.5*log(2*pi) .+ log.(diag(d.CF.L))) - log(np)
end

end # module


"""
calculated this by hand so may not be right,
all kernels (both smoothing and base process) in EQ form, see eqn 2
"""
function threeEQs(t::AbstractFloat, tp::AbstractFloat,
				  Gp1::AbstractArray, Gp2::AbstractArray, bp::AbstractArray)::AbstractFloat
	a, b, c = Gp1[1], Gp2[1], bp[1]
	coeff = 1 / (sqrt(2 * π) * sqrt(abs(a^2 + b^2 + c^2)))
	ans =  exp(-0.5 * (1 / (a^2 + b^2 + c^2)) * (t - tp)^2) * coeff
	ans
end

function threeEQs(ts::AbstractVector, # like ϕ[i, j] := scaledEQs(ts[i], ts[j], Gs[i, :], Gs[j, :], u[:])
	Gs::AbstractMatrix, bp::AbstractVector)
    @tullio phi[i,j] := exp(-0.5 * (inv(Gs[i,1]^2 + Gs[j,1]^2 + bp[1]^2) * (ts[i] - ts[j])^2)) * 1 / (sqrt(2 * π * (Gs[i,1]^2 + Gs[j,1]^2 + bp[1]^2)))
end

function scaledEQs(t::Real, tp::Real,
				   Gp1::AbstractArray, Gp2::AbstractArray, bp::AbstractArray)::Real
	sq = Gp1[1]^2 + Gp2[1]^2 + bp[1]^2
	exp(-0.5 * (1 / (sq) * (t - tp)^2)) * (Gp1[2] * Gp2[2]) / (sqrt(2 * π * sq))
end

function scaledEQs(ts::AbstractVector, # like ϕ[i, j] := scaledEQs(ts[i], ts[j], Gs[i, :], Gs[j, :], u[:])
	Gs::AbstractMatrix, bp::AbstractVector)
    @tullio phi[i,j] := exp(-0.5 * (inv(Gs[i,1]^2 + Gs[j,1]^2 + bp[1]^2) * (ts[i] - ts[j])^2)) * (Gs[i,2] * Gs[j,2]) / (sqrt(2 * π * (Gs[i,1]^2 + Gs[j,1]^2 + bp[1]^2)))
end

function scaledEQs_adj(t::AbstractFloat, tp::AbstractFloat,
					   Gp1::AbstractArray, Gp2::AbstractArray, bp::AbstractArray)
	sq = Gp1[1]^2 + Gp2[1]^2 + bp[1]^2
	f_val = scaledEQs(t, tp, Gp1, Gp2, bp)
	∇sq =  (t - tp)^2 / (2 * sq^2) * f_val - 0.5 * (1 / sq) * f_val
	(nothing, nothing, [2 * Gp1[1] * ∇sq, f_val / Gp1[2]], [2 * Gp2[1] * ∇sq, f_val / Gp2[2]], [2 * bp[1] * ∇sq])
end

@adjoint function scaledEQs(t::AbstractFloat, tp::AbstractFloat,
	Gp1::AbstractArray, Gp2::AbstractArray, bp::AbstractArray)
	scaledEQs(t, tp, Gp1, Gp2, bp), _ -> scaledEQs_adj(t, tp, Gp1, Gp2, bp)
end


module Mods

import Base: (==), (+), (-), (*), (inv), (/), (//), (^), hash, show
import Base: rand, conj, iszero

export Mod, modulus, value, AbstractMod
export is_invertible
export CRT

abstract type AbstractMod <: Number end

"""
`Mod{m}(v)` creates a modular number in mod `m` with value `mod(v,m)`.
"""
struct Mod{N,T} <: AbstractMod
    val::T
end

# safe constructors (slower)
function Mod{N}(x::T) where {T<:Union{Integer,Complex{<:Integer}},N}
    @assert N isa Integer && N > 1 "modulus must be at least 2"
    Mod{N,T}(x)
end

# type casting
Mod{N,T}(x::Mod{N,T2}) where {T,N,T2} = Mod{N,T}(T(x.val))
Mod{N,T}(x::Mod{N,T}) where {T,N} = x

show(io::IO, z::Mod{N}) where {N} = print(io, "Mod{$N}($(value(z)))")
show(io::IO, ::MIME"text/plain", z::Mod{N}) where {N} = show(io, z)

"""
`modulus(a::Mod)` returns the modulus of this `Mod` number.
```
julia> a = Mod{13}(11);

julia> modulus(a)
13
```
"""
modulus(a::Mod{N}) where {N} = N

"""
`value(a::Mod)` returns the value of this `Mod` number.
```
julia> a = Mod{13}(11);

julia> value(a)
11
```
"""
value(a::Mod{N}) where {N} = mod(a.val, N)

Base.abs(a::Mod{N,<:Real} where {N}) = abs(value(a))

function hash(x::Mod, h::UInt64 = UInt64(0))
    v = value(x)
    m = modulus(x)
    return hash(v, hash(m, h))
end

# Test for equality
iszero(x::Mod{N,T}) where {N,T} = iszero(mod(x.val, N))
==(x::Mod{N,T1}, y::Mod{M,T2}) where {M,N,T1,T2} = false
==(x::Mod{N,T1}, y::Mod{N,T2}) where {N,T1,T2} = iszero(value(x - y))

# Easy arithmetic
@inline function +(x::Mod{N,T}, y::Mod{N,T}) where {N,T}
    s, flag = Base.add_with_overflow(x.val, y.val)
    if !flag
        return Mod{N,T}(s)
    end
    t = widen(x.val) + widen(y.val)    # add with added precision
    return Mod{N,T}(mod(t, N))
end

function -(x::Mod{M,T}) where {M,T<:Signed}
    if (x.val isa BigInt || x.val == typemin(T))
        return Mod{M,T}(-mod(x.val, M))
    else
        return Mod{M,T}(-x.val)
    end
end

function -(x::Mod{M,T}) where {M,T<:Unsigned}
    return Mod{M,T}(M - value(x))
end

-(x::Mod, y::Mod) = x + (-y)

@inline function *(x::Mod{N,T}, y::Mod{N,T}) where {N,T}
    p, flag = Base.mul_with_overflow(x.val, y.val)
    if !flag
        return Mod{N,T}(p)
    else
        q = widemul(x.val, y.val)         # multipy with added precision
        return Mod{N,T}(mod(q, N)) # return with proper type
    end
end

# Division stuff
"""
`is_invertible(x::Mod)` determines if `x` is invertible.
"""
function is_invertible(x::Mod{M})::Bool where {M}
    return gcd(x.val, M) == 1
end

"""
`inv(x::Mod)` gives the multiplicative inverse of `x`.
"""
@inline function inv(x::Mod{M,T}) where {M,T}
    Mod{M,T}(_invmod(x.val, M))
end
_invmod(x::Unsigned, m::Unsigned) = invmod(x, m)
# faster version of `Base.invmod`, only works for for signed types
@inline function _invmod(x::Signed, m::Signed)
    (g, v, _) = gcdx(x, m)
    if g != 1
        error("$x (mod $m) is not invertible")
    end
    return v
end

function /(x::Mod{N,T}, y::Mod{N,T}) where {N,T}
    return x * inv(y)
end

(//)(x::Mod, y::Mod) = x / y
(//)(x::Number, y::Mod{N}) where {N} = x / y
(//)(x::Mod{N}, y::Number) where {N} = x / y

Base.promote_rule(::Type{Mod{M,T1}}, ::Type{Mod{N,T2}}) where {M,N,T1,T2<:Number} =
    error("can not promote types `Mod{$M,$T1}`` and `Mod{$N,$T2}`")
Base.promote_rule(::Type{Mod{M,T1}}, ::Type{Mod{M,T2}}) where {M,T1,T2<:Number} =
    Mod{M,promote_type(T1, T2)}
Base.promote_rule(::Type{Mod{M,T1}}, ::Type{T2}) where {M,T1,T2<:Number} =
    Mod{M,promote_type(T1, T2)}
Base.promote_rule(::Type{Mod{M,T1}}, ::Type{Rational{T2}}) where {M,T1,T2} =
    Mod{M,promote_type(T1, T2)}

# Operations with rational numbers
Mod{N}(k::Rational) where {N} = Mod{N}(numerator(k)) / Mod{N}(denominator(k))
Mod{N,T}(k::Rational{T2}) where {N,T,T2} = Mod{N,T}(numerator(k)) / Mod{N,T}(denominator(k))

# Random
rand(::Type{Mod{N}}, args::Integer...) where {N} = rand(Mod{N,Int}, args...)
rand(::Type{Mod{N,T}}) where {N,T} = Mod{N}(rand(T))
rand(::Type{Mod{N,T}}, dims::Integer...) where {N,T} = Mod{N}.(rand(T, dims...))

# Chinese remainder theorem functions
"""
    `CRT([T=BigInt, ]m1, m2,...)`

Chinese Remainder Theorem.

```
julia> CRT(Int, Mod{11}(4), Mod{14}(814))
92

julia> 92%11
4

julia> 92%14
8

julia> CRT(Mod{9223372036854775783}(9223372036854775782), Mod{9223372036854775643}(9223372036854775642))
85070591730234614113402964855534653468
```

!!! note

    `CRT` uses `BigInt` by default to prevent potential integer overflow.
    If you are confident that numbers do not overflow in your application,
    please specify an optional type parameter as the first argument.
"""
function CRT(::Type{T}, remainders, primes) where {T}
    length(remainders) == length(primes) || error("size mismatch")
    isempty(remainders) && return zero(T)
    primes = convert.(T, primes)
    M = prod(primes)
    Ms = M .÷ primes
    ti = _invmod.(Ms, primes)
    mod(sum(convert.(T, remainders) .* ti .* Ms), M)
end

function CRT(::Type{T}, rs::Mod...) where {T}
    CRT(T, value.(rs), modulus.(rs))
end
CRT(rs::Mod...) = CRT(BigInt, rs...)

include("GaussMods.jl")
include("iterate.jl")

end # end of module Mods


#############################################################################
# quantum_relative_entropy returns tr(A*(log(A)-log(B))) where A and B
# are positive semidefinite matrices.  Note this function uses logarithm
# base e, not base 2, so return value is in units of nats, not bits.
#
# Quantum relative entropy is convex (jointly) in (A,B). This function
# implements the semidefinite programming approximation given in the reference
# below.  Parameters m and k control the accuracy of this approximation: m is
# the number of quadrature nodes to use and k the number of square-roots to
# take. See reference for more details.
#
# Implementation uses the expression
#    D(A||B) = e'*D_{op} (A \otimes I || I \otimes B) )*e
# where D_{op} is the operator relative entropy and e = vec(Matrix(I, n, n)).
#
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#
#REFERENCE
#   Ported from CVXQUAD which is based on the paper: "Lieb's concavity
#   theorem, matrix geometric means and semidefinite optimization" by Hamza
#   Fawzi and James Saunderson (arXiv:1512.03401)
#############################################################################

const MatrixOrConstant = Union{AbstractMatrix,Constant}

struct QuantumRelativeEntropy1 <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr,AbstractExpr}
    size::Tuple{Int,Int}
    m::Integer
    k::Integer

    function QuantumRelativeEntropy1(
        A::AbstractExpr,
        B::AbstractExpr,
        m::Integer,
        k::Integer,
    )
        children = (A, B)
        if size(A) != size(B)
            throw(DimensionMismatch("A and B must be the same size"))
        end
        n = size(A)[1]
        if size(A) != (n, n)
            throw(DimensionMismatch("A and B must be square"))
        end
        return new(
            :quantum_relative_entropy,
            hash(children),
            children,
            (1, 1),
            m,
            k,
        )
    end
end

struct QuantumRelativeEntropy2 <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int,Int}
    m::Integer
    k::Integer
    B::AbstractMatrix
    J::AbstractMatrix
    K::AbstractMatrix

    function QuantumRelativeEntropy2(
        A::AbstractExpr,
        B::AbstractMatrix,
        m::Integer,
        k::Integer,
        nullspace_tol::Real,
    )
        children = (A,)

        if size(A) != size(B)
            throw(DimensionMismatch("A and B must be the same size"))
        end
        n = size(A)[1]
        if size(A) != (n, n)
            throw(DimensionMismatch("A and B must be square"))
        end
        if norm(B - B') > nullspace_tol
            throw(DomainError(B, "B must be Hermitian"))
        end

        # nullspace of A must contain nullspace of B
        v, U = eigen(Hermitian(B))
        if any(v .< -nullspace_tol)
            throw(DomainError(B, "B must be positive semidefinite"))
        end
        J = U'[v.>nullspace_tol, :]
        K = U'[v.<nullspace_tol, :]

        return new(
            :quantum_relative_entropy,
            hash(children),
            children,
            (1, 1),
            m,
            k,
            B,
            J,
            K,
        )
    end
end

sign(atom::Union{QuantumRelativeEntropy1,QuantumRelativeEntropy2}) = Positive()

function monotonicity(atom::QuantumRelativeEntropy1)
    return (NoMonotonicity(), NoMonotonicity())
end
monotonicity(atom::QuantumRelativeEntropy2) = (NoMonotonicity(),)

function curvature(atom::Union{QuantumRelativeEntropy1,QuantumRelativeEntropy2})
    return ConvexVexity()
end

function evaluate(atom::QuantumRelativeEntropy1)
    A = evaluate(atom.children[1])
    B = evaluate(atom.children[2])
    return quantum_relative_entropy(A, B)
end

function evaluate(atom::QuantumRelativeEntropy2)
    A = evaluate(atom.children[1])
    return quantum_relative_entropy(A, atom.B)
end

function quantum_relative_entropy(
    A::AbstractExpr,
    B::AbstractExpr,
    m::Integer = 3,
    k::Integer = 3,
)
    #println("quantum_relative_entropy general case")
    return QuantumRelativeEntropy1(A, B, m, k)
end

function quantum_relative_entropy(
    A::AbstractExpr,
    B::MatrixOrConstant,
    m::Integer = 3,
    k::Integer = 3,
    nullspace_tol::Real = 1e-6,
)
    #println("quantum_relative_entropy general case")
    return QuantumRelativeEntropy2(A, evaluate(B), m, k, nullspace_tol)
end

function quantum_relative_entropy(
    A::MatrixOrConstant,
    B::AbstractExpr,
    m::Integer = 3,
    k::Integer = 3,
)
    #println("quantum_relative_entropy constant A")
    A = evaluate(A)
    return -quantum_entropy(A, m, k) - trace_logm(B, A, m, k)
end

function quantum_relative_entropy(
    A::MatrixOrConstant,
    B::MatrixOrConstant,
    m::Integer = 0,
    k::Integer = 0,
    nullspace_tol::Real = 1e-6,
)
    #println("quantum_relative_entropy constant A and B")
    A = evaluate(A)
    B = evaluate(B)

    if size(A) != size(B)
        throw(DimensionMismatch("A and B must be the same size"))
    end
    if size(A) != (size(A)[1], size(A)[1])
        throw(DimensionMismatch("A and B must be square"))
    end
    if norm(A - A') > nullspace_tol
        throw(DomainError(A, "A must be Hermitian"))
    end
    if norm(B - B') > nullspace_tol
        throw(DomainError(B, "B must be Hermitian"))
    end

    # need to project down to support of A
    v, U = eigen(Hermitian(A))
    if any(v .< -nullspace_tol)
        throw(DomainError(A, "A must be positive semidefinite"))
    end
    if any(eigvals(Hermitian(B)) .< -nullspace_tol)
        throw(DomainError(B, "B must be positive semidefinite"))
    end
    J = U'[v.>nullspace_tol, :]
    Ap = Hermitian(J * A * J')
    Bp = Hermitian(J * B * J')

    if any(eigvals(Bp) .< nullspace_tol)
        return Inf
    end

    return real(tr(Ap * (log(Ap) - log(Bp))))
end

function conic_form!(atom::QuantumRelativeEntropy1, unique_conic_forms)
    #println("conic_form QuantumRelativeEntropy1")
    if !has_conic_form(unique_conic_forms, atom)
        A = atom.children[1]
        B = atom.children[2]
        m = atom.m
        k = atom.k
        n = size(A)[1]
        eye = Matrix(1.0 * I, n, n)
        e = vec(eye)

        conic_form!(A ⪰ 0, unique_conic_forms)
        conic_form!(B ⪰ 0, unique_conic_forms)

        τ = Variable()
        conic_form!(
            τ in
            RelativeEntropyEpiCone(kron(A, eye), kron(eye, conj(B)), m, k, e),
            unique_conic_forms,
        )

        cache_conic_form!(unique_conic_forms, atom, minimize(τ))
    end
    return get_conic_form(unique_conic_forms, atom)
end

function conic_form!(atom::QuantumRelativeEntropy2, unique_conic_forms)
    #println("conic_form QuantumRelativeEntropy2")
    if !has_conic_form(unique_conic_forms, atom)
        A = atom.children[1]
        B = atom.B
        J = atom.J
        K = atom.K
        m = atom.m
        k = atom.k

        conic_form!(A ⪰ 0, unique_conic_forms)

        if length(K) > 0
            conic_form!(K * A * K' == 0, unique_conic_forms)
            Ap = J * A * J'
            Bp = Hermitian(J * B * J')
            τ = -quantum_entropy(Ap, m, k) - real(tr(Ap * log(Bp)))
        else
            τ = -quantum_entropy(A, m, k) - real(tr(A * log(B)))
        end

        cache_conic_form!(unique_conic_forms, atom, minimize(τ))
    end
    return get_conic_form(unique_conic_forms, atom)
end


"""
`CALCULATE_AFFINE` - Compute left-hand affine transform from two Nxd point sets

```
tform = calculate_affine(moving_pts, fixed_pts)
```

* moving_pts: Nxd array, each row being homogeneous coords of a point, of the
    point set that is moving to fit the fixed point set.
* fixed_pts: Nxd array, each row being homogeneous coords of a point, of the
    reference point set.
* tform: affine transformation that produces a mapping of the moving_pts to
    fixed_pts.

```
fixed_pts = moving_pts * tform
```
"""
function calculate_affine(moving_pts, fixed_pts)
  return moving_pts \ fixed_pts
end

"""
`CALCULATE_RIGID` - Compute right-hand rigid transform from two Nxd point sets

```
tform = calculate_rigid(moving_pts, fixed_pts)
```

* moving_pts: Nxd array, each row being homogeneous coords of a point, of the
    point set that is moving to fit the fixed point set.
* fixed_pts: Nxd array, each row being homogeneous coords of a point, of the
    reference point set.
* tform: rigid transformation that produces a mapping of the moving_pts to
    fixed_pts.

```
fixed_pts = moving_pts * tform
```
"""
function calculate_rigid(moving_pts, fixed_pts)
  moving_pts = moving_pts[:,1:end-1]
  fixed_pts = fixed_pts[:,1:end-1]
  n, dim = size(moving_pts)
  moving_bar = mean(moving_pts, 1)
  fixed_bar = mean(fixed_pts, 1)
  moving_centered = moving_pts .- moving_bar
  fixed_centered = fixed_pts .- fixed_bar
  C = fixed_centered.' * moving_centered / n
  U,s,V = svd(C)
  # Rotation R in least squares sense:
  # moving_pts - moving_bar = (fixed_pts - fixed_bar)*R
  R = (U * diagm(vcat(ones(dim-1), det(U*V.'))) * V.' ).'
  t = fixed_bar - moving_bar*R
  return [R zeros(dim); t 1]
end

"""
`CALCULATE_TRANSLATION` - Compute left-hand translation from two Nxd point sets

```
tform = calculate_translation(moving_pts, fixed_pts)
```

* moving_pts: Nxd array, each row being nonhomogeneous coords of a point, of the
    point set that is moving to fit the fixed point set.
* fixed_pts: Nxd array, each row being nonhomogeneous coords of a point, of the
    reference point set.
* tform: translation transformation that produces a mapping of the moving_pts to
    fixed_pts.

```
fixed_pts = moving_pts * tform
```
"""
function calculate_translation(moving_pts, fixed_pts)
  n, dim = size(moving_pts)
  moving_bar = mean(moving_pts, 1)
  fixed_bar = mean(fixed_pts, 1)
  t = fixed_bar - moving_bar
  return [eye(2) zeros(dim); t 1]
end

function calculate_affine(matches::Matches)
  return calculate_affine(matches.src_points, matches.dst_points)
end

function calculate_rigid(matches::Matches)
  return calculate_rigid(matches.src_points, matches.dst_points)
end

function calculate_translation(matches::Matches)
  return calculate_translation(matches.src_points, matches.dst_points)
end

function calculate_affine(mesh::Mesh)
  return calculate_affine(mesh.src_nodes, mesh.dst_nodes)
end

function calculate_rigid(mesh::Mesh)
  return calculate_rigid(mesh.src_nodes, mesh.dst_nodes)
end

function calculate_translation(mesh::Mesh)
  return calculate_translation(mesh.src_nodes, mesh.dst_nodes)
end

"""
Matrix for rotation in degrees (clockwise rotation with left-handed Cartesian)
"""
function make_rotation_matrix(theta, image_size = nothing)
  angle = deg2rad(theta)
  tform = [cos(angle) sin(angle) 0; -sin(angle) cos(angle) 0; 0 0 1]
  tform[abs(tform) .< eps] = 0
  if image_size != nothing
    	rotated_bb = snap_bb(tform_bb(sz_to_bb([image_size[1], image_size[2]]), tform))
    	rotation_offset = [rotated_bb.i, rotated_bb.j]
	translation_matrix = make_translation_matrix(-rotation_offset);
	tform = tform * translation_matrix
  end
  return tform
end

function make_translation_matrix(offset)
  return [1 0 0; 0 1 0; offset[1] offset[2] 1]
end

function make_scale_matrix(scale)
  return [scale 0 0; 0 scale 0; 0 0 1]
end


# # Hydraulic functions

# This tutorial shows how to specify the hydraulic functions
# used in Richard's equation. In particular,
# we show how to choose the formalism for matric potential and hydraulic
# conductivity, and how to make the hydraulic conductivity account for
# the presence of ice as well as the temperature dependence of the
# viscosity of liquid water.

# # Preliminary setup

# External modules
using Plots

# ClimateMachine modules
using ClimateMachine
using ClimateMachine.Land
using ClimateMachine.Land.SoilWaterParameterizations

FT = Float32;
# # Specifying a hydraulics model
# ClimateMachine's Land model allows the user to pick between two hydraulics models,
# that of van Genuchten [vanGenuchten1980](@cite) or that of Brooks and Corey, see [BrooksCorey1964](@cite) or [Corey1977](@cite). The
# same model is consistently used for the matric potential
# and hydraulic conductivity.

# The van Genuchten model requires two free parameters, `α` and `n`.
# A third parameter, `m`, is computed from `n`. Of these, only `α` carries
# units, of inverse meters. The Brooks and Corey model also uses
# two free parameters, `ψ_b`, the magnitude of the matric potential at saturation,
#  and a constant `M`. `ψ_b` carries units of meters. These parameter sets are stored in
# either the [`vanGenuchten`](@ref ClimateMachine.Land.SoilWaterParameterizations.vanGenuchten) or the
# [`BrooksCorey`](@ref ClimateMachine.Land.SoilWaterParameterizations.BrooksCorey)
# hydraulics model structures (more details below). These parameters are enough to compute the matric potential.

# The hydraulic conductivity
# requires an additional parameter, `Ksat` (m/s), which is the hydraulic conductivity
# in saturated soil. This parameter is
# not stored in the hydraulics model, but rather as part of the
# [`WaterParamFunctions`](@ref ClimateMachine.Land.WaterParamFunctions), which stores
# other parameters needed for the soil water modeling.

# Below we show how to create two concrete examples of these hydraulics models,
# for sandy loam ([Bonan19a](@cite)). Note that the parameters chosen are a function of soil type,
# and that the parameters are converted to type `FT` internally.
vg_α = 7.5 # m^-1
vg_n = 1.89
hydraulics = vanGenuchten(FT; α = vg_α, n = vg_n);

ψ_sat = 0.218 # m
Mval = 0.2041
hydraulics_bc = BrooksCorey(FT; ψb = ψ_sat, m = Mval);
# # Matric Potential
# The matric potential `ψ` reflects the negative pressure of water
# in unsaturated soil. The negative pressure (suction) of water arises
# because of adhesive forces between water and soil.

# The van Genuchten expression for matric potential is
# ``
# ψ = -\frac{1}{α} S_l^{-1/(nm)}\times (1-S_l^{1/m})^{1/n},
# ``

# and the Brooks and Corey expression is
# ``
# ψ = -ψ_b S_l^{-M}.
# ``

# Here `S_l` is the effective saturation of liquid water, `θ_l/ν`, where `ν` is
# porosity of the soil. We generally neglect the residual pore space in the CliMA model,
# but the user can set the parameter in the
# [`WaterParamFunctions`](@ref ClimateMachine.Land.WaterParamFunctions) structure if it is
# desired.

# In the CliMA code, we use [multiple dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch).
# With multiple dispatch, a function can have many
# ways of executing (called methods), depending on the *type* of the
# variables passed in. A simple example of multiple dispatch is the division operation.
# Integer division takes two numbers as input, and returns an integer - ignoring the decimal.
# Float division takes two numbers as input, and returns a floating point number, including the decimal.
# In Julia, we might write these as:

# ```julia
# function division(a::Int, b::Int)
#      return floor(Int, a/b)
# end
# ```
# ```julia
# function division(a::Float64, b::Float64)
#      return a/b
# end
# ```

# We can see that `division` is now a function with two methods.

# ```julia
# julia> division
# division (generic function with 2 methods)
# ```

# Now, using the same function signature, we can carry out integer
# division or floating point division, depending on the types of the
# arguments:

# ```julia
# julia> division(1,2)
# 0
#
# julia> division(1.0,2.0)
# 0.5
# ```

# Here are more pertinent examples:
# Based on our choice of `FT = Float32`,

# ```julia
# julia> typeof(hydraulics)
# vanGenuchten{Float32,Float32,Float32,Float32}
# ```

# but meanwhile,

# ```julia
# julia> typeof(hydraulics_bc)
# BrooksCorey{Float32,Float32,Float32}
# ```

# The function `matric_potential` will execute different methods
# depending on if we pass a hydraulics model of type `vanGenuchten` or
# `BrooksCorey`. In both cases, it will return the correct value
# for `ψ`.

# Let's plot the matric potential as a function of the effective saturation `S_l = θ_l/ν`,
# which can range from zero to one.
S_l = FT.(0.01:0.01:0.99)
ψ = matric_potential.(Ref(hydraulics), S_l)
ψ_bc = matric_potential.(Ref(hydraulics_bc), S_l)
plot(
    S_l,
    log10.(-ψ),
    xlabel = "effective saturation",
    ylabel = "Log10(|ψ|)",
    label = "van Genuchten",
)
plot!(S_l, log10.(-ψ_bc), label = "Brooks and Corey")
savefig("bc_vg_matric_potential.png")
# ![](bc_vg_matric_potential.png)

# The steep slope in
# `ψ` near saturated and completely dry soil are part of the reason
# why Richard's equation is such a challenging numerical problem.

# # Hydraulic conductivity
# The hydraulic conductivity is a more complex function than the matric potential,
# as it depends on the temperature of the water, the volumetric ice fraction, and
# the volumetric liquid water fraction. It also depends on the hydraulics model
# chosen.

# We represent the hydraulic conductivity `K` as the product of four factors:
# `Ksat`, an impedance factor (which accounts for the effect of ice on conductivity)
# a viscosity factor (which accounts for the effect of temperature on the
# viscosity of liquid water, and how that in turn affects conductivity)
# and a moisture factor (which accounts for the effect of liquid water, and is determined by the hydraulics model).
# We are going to calculate `K = Ksat × viscosity factor × impedance factor × moisture factor`.
# In the code, each of these factors is
# computed by a function with multiple methods, except for `Ksat`.
# Like we defined new type
# classes for `vanGenuchten` and `BrooksCorey`, we also created new type classes
# for the impedance choice, the viscosity choice, and the moisture choice.

# The function [`viscosity_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.viscosity_factor)
# takes as arguments the temperature of the soil and the
# viscosity model desired, and returns the factor `k_v` by which the hydraulic conductivity is scaled.
# One option is to account for this effect:

# ``
# k_v = e^{γ (T-T_{\rm ref})}
# ``

# where γ = 0.0264/K and ``T_{\rm ref}`` = 288K.

# For example, at the freezing point of water, using the default values
# for γ and T_ref, viscosity reduces the conductivity by a third:
viscous_effect_model = TemperatureDependentViscosity{FT}();
viscosity_factor(viscous_effect_model, FT(273.15))

# The other option is to ignore this effect:

# ``
# k_v = 1
# ``

# This is the default approach.
no_viscous_effect_model = ConstantViscosity{FT}();
viscosity_factor(no_viscous_effect_model, FT(273.15))

# Very similarly, the function
# [`impedance_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.impedance_factor)
# takes as arguments the liquid water and ice
# volumetric fractions in the soil, as well as the impedance model being used, and returns
# the factor `k_i` by which the hydraulic conductivity is scaled.
# One option is to account for this effect:

# ``
# k_i = 10^{-Ω f_i},
# ``

# where `Ω = 7` is an empirical factor and
# `f_i` is the ratio of the volumetric
# ice fraction to total volumetric water fraction  ([Lundin1990](@cite)).

# For example, with ``\theta_i = \theta_l``, or f_i = 0.5, ice reduces the conductivity by over 1000x.
impedance_effect_model = IceImpedance{FT}();
impedance_factor(impedance_effect_model, FT(0.5))

# The other option is to ignore this effect:

# ``
# k_i = 1
# ``

# This is the default approach.
no_impedance_effect_model = NoImpedance{FT}();
impedance_factor(no_impedance_effect_model, FT(0.5))

# As for the moisture dependence of hydraulic conductivity, it can also be either
# independent of moisture, or dependent on moisture. If it is dependent on moisture,
# the specific function evaluated is dictated by the hydraulics model.
# The [`moisture_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.moisture_factor)
# for the van Genuchten model is (denoting it as ``k_m``)

# ``
#  k_m = \sqrt{S_l}[1-(1-S_l^{1/m})^m]^2,
# ``

# for ``S_l < 1``,

# and for the Brooks and Corey model it is

# ``
# k_m = S_l^{2M+3},
# ``

# also for ``S_l<1``. When ``S_l\geq 1``, ``k_m = 1`` for each model.

# Let's put all these factors together now. Below
# we choose additional parameters, consistent with the hydraulics parameters
# for sandy loam ([Bonan19a](@cite)), and show how hydraulic conductivity varies with
# liquid water content, in the case without ice impedance or temperature effects.
Ksat = FT(4.42 / (3600 * 100))
T = FT(0.0)
f_i = FT(0.0)
K =
    hydraulic_conductivity.(
        Ref(Ksat),
        Ref(impedance_factor(NoImpedance{FT}(), f_i)),
        Ref(viscosity_factor(ConstantViscosity{FT}(), T)),
        moisture_factor.(Ref(MoistureDependent{FT}()), Ref(hydraulics), S_l),
    );

# Let's also compute `K` when we include the effects of temperature
# and ice on the hydraulic conductivity.
# In the cases where a
# [`TemperatureDependentViscosity`](@ref ClimateMachine.Land.SoilWaterParameterizations.TemperatureDependentViscosity)
# or
# [`IceImpedance`](@ref ClimateMachine.Land.SoilWaterParameterizations.IceImpedance)
# type is passed, the correct factors are calculated,
# based on the temperature `T` and volumetric ice fraction `θ_i`.

T = FT(273.15)
S_i = FT(0.1); # = θ_i/ν
# The total volumetric water fraction cannot
# exceed unity, so the effective liquid water saturation
# should have a max of 1-S_i.
S_l_accounting_for_ice = FT.(0.01:0.01:(0.99 - S_i))
f_i = S_i ./ (S_l_accounting_for_ice .+ S_i)
K_w_factors =
    hydraulic_conductivity.(
        Ref(Ksat),
        impedance_factor.(Ref(NoImpedance{FT}()), f_i),
        Ref(viscosity_factor(ConstantViscosity{FT}(), T)),
        moisture_factor.(
            Ref(MoistureDependent{FT}()),
            Ref(hydraulics),
            S_l_accounting_for_ice,
        ),
    );
plot(
    S_l,
    log10.(K),
    xlabel = "total effective saturation, (θ_i+θ_l)/ν",
    ylabel = "Log10(K)",
    label = "Base case",
    legend = :bottomright,
)
plot!(
    S_l_accounting_for_ice .+ S_i,
    log10.(K_w_factors),
    label = "θ_i = 0.1, T = 273.15",
)
savefig("T_ice_K.png")
# ![](T_ice_K.png)

# If the user is not considering phase transitions
# and does not add in Freeze/Thaw source terms, the default is for zero
# ice in the model, for all time and space. In this case the ice impedance
# factor evaluates to 1 regardless of which type is passed.

# # Other features
# The user also has the choice of making the conductivity constant by choosing
# [`MoistureIndependent`](@ref ClimateMachine.Land.SoilWaterParameterizations.MoistureIndependent)
# along with
# [`ConstantViscosity`](@ref ClimateMachine.Land.SoilWaterParameterizations.ConstantViscosity)
# and
# [`NoImpedance`](@ref ClimateMachine.Land.SoilWaterParameterizations.NoImpedance).
# This is useful for debugging!
no_moisture_dependence = MoistureIndependent{FT}()
K_constant =
    hydraulic_conductivity.(
        Ref(Ksat),
        Ref(FT(1.0)),
        Ref(FT(1.0)),
        moisture_factor.(Ref(no_moisture_dependence), Ref(hydraulics), S_l),
    );
# ```julia
# julia> unique(K_constant)
# 1-element Array{Float32,1}:
#  1.2277777f-5
# ```

# Note that choosing this option does not mean the matric potential
# is constant, as a hydraulics model is still required and employed.

# And, lastly, you might also find it helpful in debugging
# to be able to turn off the flow of water by setting `Ksat = 0`.

# # References
# - [vanGenuchten1980](@cite)
# - [BrooksCorey1964](@cite)
# - [Corey1977](@cite)
# - [Lundin1990](@cite)
# - [Bonan19a](@cite)


include(joinpath(@__DIR__, "..", "src", "RBFBase.jl"))
using .RBF
using ForwardDiff: gradient, hessian, jacobian     # to automatically differentiate the model outputs for comparison
using Test
import FiniteDiff;
fd = FiniteDiff;
m = RBF;

# RBFModel with 2 input and 2 outputs

f1(x) = 1 .+ sum(x.^2);
f2(x) = sin(sum(x));
f(x) = [f1(x), f2(x)];

sites = vec( [[x;y] for x = range( -2, 2, length = 5),
                    y = range( -2, 2, length = 5) ] )

sites = [rand(2) for i = 1 : 5]

vals_1 = f1.(sites);
vals_2 = f2.(sites);
vals = f.(sites);

shape_parameters = [ 1 ]
kernels = [:multiquadric, :exp, :cubic, :thin_plate_spline]

println("Looping through different RBF settings.")
for poly_deg ∈ [ -1, 0, 1 ]
    for kernel ∈ kernels
        for sp ∈ shape_parameters

            A = 1e-3    # somewhat loose tolerance because of possibly bad conditioning
            println("Kernel $(kernel), sp = $sp, poly_deg = $poly_deg")
            rbf = m.RBFModel(
                training_sites = sites,
                training_values = vals,
                kernel = kernel,
                shape_parameter = sp,
                polynomial_degree = poly_deg,
            );
            m.train!(rbf);

            # Test interpolation properties
            rbf1 = x -> m.output( rbf, 1, x );
            rbf2 = x -> m.output( rbf, 2, x );

            evals_1 = rbf1.( sites )
            @test all( isapprox.( evals_1 , vals_1,  atol=A ) )

            evals_2 = rbf2.( sites )
            @test all( isapprox.( evals_2 , vals_2,  atol=A ) )

            rbf_output = x -> m.output( rbf, x );
            @test rbf_output( sites[1] ) == rbf.function_handle( sites[1] )

            # Test gradient calculations
            d_rbf1 = x -> m.grad( rbf, 1, x);
            d_rbf2 = x -> m.grad( rbf, 2, x);

            fwd_rbf1 = χ -> gradient( x -> m.output( rbf, 1, x ), χ);
            fwd_rbf2 = χ -> gradient(x -> m.output( rbf, 2, x ), χ);

            dvals_1 = d_rbf1.(sites)
            dvals_2 = d_rbf2.(sites)
            fwdvals_1 = fwd_rbf1.(sites)
            fwdvals_2 = fwd_rbf2.(sites)

            @test all( isapprox.( dvals_1, fwdvals_1, atol = A) )
            @test all( isapprox.( dvals_2, fwdvals_2, atol = A) )

            jacobian_handle = x -> m.jac( rbf, x )
            @test jacobian_handle(sites[1]) ≈ [d_rbf1(sites[1])'; d_rbf2(sites[1])']

            # test hessian correctness
            h_rbf1 = x -> m.hess( rbf, 1, x )

            H_fwd1 = χ -> fd.finite_difference_hessian( x -> m.output( rbf, 1, x), χ )

            X = rand(2)
            @test h_rbf1( X ) ≈ H_fwd1( X ) atol = A

        end
    end
end


@use GADT

import Base: convert

struct Fun{T,R}
    fn::Function
end

function (typed_fn::Fun{T,R})(arg::T)::R where {T,R}
    typed_fn.fn(arg)
end

function convert(::Type{Fun{T,R}}, fn::Function) where {T,R}
    Fun{T,R}(fn)
end

function convert(::Type{Fun{T,R}}, fn::Fun{C,D}) where {T,R,C<:T,D<:R}
    Fun{T,R}(fn.fn)
end

⇒(::Type{A}, ::Type{B}) where {A,B} = Fun{A,B}

@data public Exp{T} begin
    Sym{T}::Symbol => Exp{T}
    Val{A}::A => Exp{A}
    App{A,B}::(Exp{Fun{A,B}}, Exp{<:A}) => Exp{B}
    Lam{A,B}::(Symbol, Exp{B}) => Exp{Fun{A,B}}
    If{A}::(Exp{Bool}, Exp{A}, Exp{A}) => Exp{A}
end

function substitute(template::Exp{T}, pair::Tuple{Symbol,Exp{G}}) where {T,G<:T}
    (sym, exp) = pair
    @match template begin
        Sym(&sym) => exp
        Val(_) => template
        App(f, a) => App(substitute(f, pair), substitute(a, pair))
        Lam(&sym, exp) => template
        Lam(sym, exp) => Lam(sym, substitute(exp, pair))
        If(cond, exp1, exp2) => begin
            (cond, exp1, exp2) =
                substitute(cond, pair), substitute(exp1, pair), substitute(exp2, pair)
            If(cond, exp1, exp2)
        end
    end
end

function eval_exp(exp::Exp{T_}, ctx::Dict{Symbol,Any}) where {T_}
    let T = T_ # fix static parameter issues for Julia 1.0 and 1.1
    @match exp begin
        Sym(a) => (ctx[a]::T, ctx)
        Val(a::T) => (a, ctx)
        App{A,T}(f::Exp{Fun{A,T}}, arg::Exp{<:A}) where {A, T} =>
            let (f, ctx) = eval_exp(f, ctx), (arg, ctx) = eval_exp(arg, ctx)
                (f(arg), ctx)
            end
        Lam{A,B}(sym, exp::Exp{B}) where {A,B} =>
            let f(x) = begin
                    eval_exp(substitute(exp, sym => Val(x)), ctx)[1]
                end
                (f, ctx)
            end
        If(cond, exp1, exp2) => let (cond, ctx) = eval_exp(cond, ctx)
            eval_exp(cond ? exp1 : exp2, ctx)
        end
        _ => error(exp)
    end
    end
end

add = Val{Number ⇒ Number ⇒ Number}(x -> y -> x + y)
sub = Val{Number ⇒ Number ⇒ Number}(x -> y -> x - y)
gt = Val{Number ⇒ Number ⇒ Bool}(x -> y -> x > y)
ctx = Dict{Symbol,Any}()

@assert 3 == eval_exp(App(App(add, Val(1)), Val(2)), ctx)[1]
@assert -1 == eval_exp(App(App(sub, Val(1)), Val(2)), ctx)[1]
@assert 1 == eval_exp(
    If(
        App(App(gt, Sym{Int}(:x)), Sym{Int}(:y)),
        App(App(sub, Sym{Int}(:x)), Sym{Int}(:y)),
        App(App(sub, Sym{Int}(:y)), Sym{Int}(:x)),
    ),
    Dict{Symbol,Any}(:x => 1, :y => 2),
)[1]