################################################################################
#  Copyright 2022, Tom Van Acker (BASF)                                        #
################################################################################
# TFT.jl                                                                       #
# A Julia package for Taylor-Fourier Transform.                                #
# See http://github.com/timmyfaraday/TFT.jl                                    #
################################################################################

# checks
function check_sol(sol, D, H)
    sol.prob.D >= D || Base.error("the required Dth-degree derivative is unavailable")
    H in sol.prob.h || Base.error("the required Hth-harmonic phasor is unavailable")
end

# frequency and angular frequency
F(sol)   = sol.prob.F
F(sol,H) = H * F(sol)
ω(sol,H) = 2.0 * pi * F(sol,H)

# amplitude
"""
    TFT.a(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Shorthand function to obtain the amplitude of the Dth-degree derivative of the
Hth-harmonic phasor, dispatching to `amplitude(sol,D,H)`.
"""
a(sol::AbstractDTFTSolution, D::Int=0, H::Int=1) = amplitude(sol,D,H)

"""
    TFT.amplitude(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Function to obtain the amplitude of the Dth-degree derivative of the
Hth-harmonic phasor.

    ∀ h ∈ {0}:
        a₀⁽ᴰ⁾(t) = ξ₀⁽ᴰ⁾(t)
    ∀ h ∉ {0}:
        aₕ⁽⁰⁾(t) = |2 ⋅ ξₕ⁽⁰⁾(t)| ∈ 𝐑⁺
        aₕ⁽¹⁾(t) = ℜ[2 ⋅ ξₕ⁽¹⁾(t) ⋅ exp(-im ⋅ ϕₕ⁽⁰⁾(t))] ∈ 𝐑
        aₕ⁽²⁾(t) = ℜ[2 ⋅ ξₕ⁽²⁾(t) ⋅ exp(-im ⋅ ϕₕ⁽⁰⁾(t))] + aₕ⁽⁰⁾(t) ⋅ ϕₕ⁽¹⁾(t)² ∈ 𝐑

See: [Assessing Synchrophasor Estimates of an Event Captured by a Phasor
Measurement Unit, pg. 3112](https://ieeexplore.ieee.org/document/9239915)

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `D::Int`                      | degree of the derivative [-], default=0
- `H::Int`                      | harmonic number [-], default=1

Output:
- `a::Vector{<:Real}`           | amplitude aₕ⁽ᴰ⁾(t) [?]
"""
function amplitude(sol::AbstractDTFTSolution, D::Int=0, H::Int=1)
    H == 0 && return ξ(sol,D,H)
    D == 0 && return abs.(2.0 .* ξ(sol,0,H))
    D == 1 && return real.(2.0 .* ξ(sol,1,H) .* exp.(-im .* ϕ(sol,0,H)))
    D == 2 && return real.(2.0 .* ξ(sol,2,H) .* exp.(-im .* ϕ(sol,0,H))) .+
                        a(sol,0,H) .* ϕ(sol,1,H).^2
    return nothing
end

# phase
"""
    TFT.ϕ(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Shorthand function to obtain the phase of the Dth-degree derivative
of the Hth-harmonic phasor, dispatching to `phase(sol, D, H)`.
"""
ϕ(sol::AbstractDTFTSolution, D::Int=0, H::Int=1) = phase(sol, D, H)

"""
    TFT.phase(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Function to obtain the alternative angle of the Dth-degree derivative of the
Hth-harmonic phasor.

    ∀ h ∈ {0}:
        ϕ₀⁽ᴰ⁾(t) = 0.0
    ∀ h ∉ {0}:
        ϕₕ⁽⁰⁾(t) = ∠[pₕ⁽⁰⁾(t)] ∈ [-π,π]
        ϕₕ⁽¹⁾(t) = ℑ[2 ⋅ ξₕ⁽¹⁾(t) ⋅ exp(-im ⋅ ϕₕ⁽⁰⁾(t))] / aₕ⁽⁰⁾(t) ∈ 𝐑
        ϕₕ⁽²⁾(t) = {ℑ[2 ⋅ ξₕ⁽²⁾(t) ⋅ exp(-im ⋅ ϕₕ⁽⁰⁾(t))] - 2 ⋅ aₕ⁽¹⁾(t) ⋅ ϕₕ⁽¹⁾(t)} / aₕ⁽⁰⁾(t) ∈ 𝐑

See: [Assessing Synchrophasor Estimates of an Event Captured by a Phasor
Measurement Unit, pg. 3112](https://ieeexplore.ieee.org/document/9239915)

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `D::Int`                      | degree of the derivative [-], default=0
- `H::Int`                      | harmonic number [-], default=1

Output:
- `ϕ::Vector{<:Real}`           | phase ϕₕ⁽ᴰ⁾(t) [(rad)]
"""
function phase(sol::AbstractDTFTSolution, D::Int=0, H::Int=1)
    H == 0 && return zero(ξ(sol,D,H))
    D == 0 && return Base.angle.(ξ(sol,0,H))
    D == 1 && return imag.(2.0 .* ξ(sol,1,H) .* exp.(-im .* ϕ(sol,0,H))) ./
                        a(sol,0,H)
    D == 2 && return (imag.(2.0 .* ξ(sol,2,H) .* exp.(-im .* ϕ(sol,0,H))) .-
                        2.0 .* a(sol,1,H) .* ϕ(sol,1,H)) ./ a(sol,0,H)
    return nothing
end

# anti-rotating phase
"""
    TFT.φ(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Shorthand function to obtain the anti-rotating phase of the Dth-degree
derivative of the Hth-harmonic phasor, dispatching to `angle(sol,D,H)`.
"""
φ(sol::AbstractDTFTSolution, D::Int=0, H::Int=1) = ar_phase(sol,D,H)

"""
    TFT.ar_phase(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Function to obtain the anti-rotating phase of the Dth-degree derivative of the
Hth-harmonic phasor.

    ∀ h ∈ {0}:
        φ₀⁽ᴰ⁾(t) = 0.0
    ∀ h ∉ {0}:
        φₕ⁽⁰⁾(t) = ∠[ψₕ⁽⁰⁾(t)] ∈ [-π,π]
        φₕ⁽¹⁾(t) = ℑ[2 ⋅ ψₕ⁽¹⁾(t) ⋅ exp(-im ⋅ φₕ⁽⁰⁾(t))] / aₕ⁽⁰⁾(t) ∈ 𝐑
        φₕ⁽²⁾(t) = {ℑ[2 ⋅ ψₕ⁽²⁾(t) ⋅ exp(-im ⋅ φₕ⁽⁰⁾(t))] - 2 ⋅ aₕ⁽¹⁾(t) ⋅ φₕ⁽¹⁾(t)} / aₕ⁽⁰⁾(t) ∈ 𝐑

See: [Assessing Synchrophasor Estimates of an Event Captured by a Phasor
Measurement Unit, pg. 3112](https://ieeexplore.ieee.org/document/9239915)

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `D::Int`                      | degree of the derivative [-], default=0
- `H::Int`                      | harmonic number [-], default=1

Output:
- `φ::Vector{<:Real}`           | anti-rotating phase φₕ⁽ᴰ⁾(t) [(rad)]
"""
function ar_phase(sol::AbstractDTFTSolution, D::Int=0, H::Int=1)
    H == 0 && return zero(ψ(sol,D,H))
    D == 0 && return Base.angle.(ψ(sol,0,H))
    D == 1 && return imag.(2.0 .* ψ(sol,1,H) .* exp.(-im .* φ(sol,0,H))) ./
                        a(sol,0,H)
    D == 2 && return (imag.(2.0 .* ψ(sol,2,H) .* exp.(-im .* φ(sol,0,H))) .-
                        2.0 .* a(sol,1,H) .* φ(sol,1,H)) ./ a(sol,0,H)
    return nothing
end

# frequency
"""
    TFT.f(sol::TFT.AbstractDTFTSolution, H::Int=1)

Shorthand function to obtain the frequency of the zeroth-degree derivative of
the Hth-harmonic phasor, dispatching to `frequency(sol,H)`.
"""
f(sol::AbstractDTFTSolution, H::Int=1) = frequency(sol,H)

"""
    TFT.frequency(sol::TFT.AbstractDTFTSolution, H::Int=1)

Function to obtain the frequency of the zeroth-degree derivative of the
Hth-harmonic phasor.

    fₕ(t) = Fₕ + ϕₕ⁽¹⁾(t) / (2 π) ∈ 𝐑⁺

See: [Fast Taylor-Fourier Transform for Monitoring Modern Power Grids with
Real-Time Dynamic Harmonic Estimation](tbp)

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `H::Int`                      | harmonic number [-], default=1

Output:
- `f::Vector{<:Real}`           | frequency fₕ(t) [Hz]
"""
frequency(sol::AbstractDTFTSolution, H::Int=1) =
    F(sol,H) .+ (ϕ(sol,1,H) ./ (2 * pi))

# rocof
"""
    TFT.r(sol::TFT.AbstractDTFTSolution, H::Int=1)

Shorthand function to obtain the rate-of-change-of-frequency of the
zeroth-degree derivative of the Hth-harmonic phasor, dispatching to
`rocof(sol,H)`.
"""
r(sol::AbstractDTFTSolution, H::Int=1) = rocof(sol,H)

"""
    TFT.rocof(sol::TFT.AbstractDTFTSolution, H::Int=1)

Function to obtain the rate-of-change-of-frequency of the zeroth-degree
derivative of the Hth-harmonic phasor.

    rₕ(t) = ϕₕ⁽²⁾(t) / (2 π)² ∈ 𝐑

See: [Fast Taylor-Fourier Transform for Monitoring Modern Power Grids with
Real-Time Dynamic Harmonic Estimation](tbp)

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `H::Int`                      | harmonic number [-], default=1

Output:
- `r::Vector{<:Real}`           | rocof rₕ(t) [Hz²]
"""
rocof(sol::AbstractDTFTSolution, H::Int=1) =
    ϕ(sol,2,H) ./ (2 * pi)^2

# dynamic phasor
"""
    TFT.ξ(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Shorthand function to obtain the Dth-degree derivative of the Hth-harmonic
dynamic phasor, dispatching to `phasor(sol,D,H)`.
"""
ξ(sol::AbstractDTFTSolution, D::Int=0, H::Int=1) = phasor(sol,D,H)

"""
    TFT.phasor(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Function to obtain the D-th-degree derivative of the Hth-harmonic dynamic
phasor. For the zeroth-harmonic dynamic phasor, only the real part of the
dynamic phasor is returned.

    ∀ h ∈ {0}:
        ξ₀⁽ᴰ⁾(t) ∈ ℝ
    ∀ h ∉ {0}:
        ξₕ⁽ᴰ⁾(t) ∈ ℂ

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `D::Int`                      | degree of the derivative [-], default=0
- `H::Int`                      | harmonic number [-], default=1

Output:
- `ξ::Vector{<:Complex}`        | dynamic phasor ξₕ⁽ᴰ⁾(t) [?]
"""
function phasor(sol::AbstractDTFTSolution, D::Int=0, H::Int=1)
    check_sol(sol, D, H)

    H == 0 && return real(sol.X[H][:,D+1])
    return sol.X[H][:,D+1]
end

# anti-rotating dynamic phasor
"""
    TFT.ψ(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Shorthand function to obtain the Dth-degree derivative of the Hth-harmonic anti-
rotating dynamic phasor, dispatching to `ar_phasor(sol,D,H)`.
"""
ψ(sol::AbstractDTFTSolution, D::Int=0, H::Int=1) = ar_phasor(sol,D,H)

"""
    TFT.ar_phasor(sol::TFT.AbstractDTFTSolution, D::Int=0, H::Int=1)

Function to obtain the D-th-degree derivative of the Hth-harmonic anti-rotating
dynamic phasor.

    ∀ h ∈ {0}:
        ψ₀⁽ᴰ⁾(t) = ξ₀⁽ᴰ⁾(t) ∈ ℝ
    ∀ h ∉ {0}:
        ψₕ⁽ᴰ⁾(t) = ξₕ⁽ᴰ⁾(t) exp(-im ωₕ t) ∈ 𝐂

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `D::Int`                      | degree of the derivative [-], default=0
- `H::Int`                      | harmonic number [-], default=1

Output:
- `ψ::Vector{<:Complex}`        | anti-rotating dynamic phasor ψₕ⁽ᴰ⁾(t) [?]
"""
ar_phasor(sol::AbstractDTFTSolution, D::Int=0, H::Int=1) =
    ifelse(H == 0, ξ(sol,D,H), ξ(sol,D,H) .* exp.(-im .* ω(sol,H) .* sol.prob.t))

# signal
"""
    TFT.signal(sol::TFT.AbstractDTFTSolution)

Function to obtain the overall signal.

    s(t) = ∑ₕ ξₕ⁽⁰⁾(t) + conj(ξₕ⁽⁰⁾(t)) ∈ 𝐑

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]

Output:
- `s::Vector{<:Real}`           | signal s(t) [?]
"""
signal(sol::AbstractDTFTSolution) = sum(signal(sol,nh) for nh in sol.prob.h)

"""
    TFT.signal(sol::TFT.AbstractDTFTSolution, H::Int)

Function to obtain the Hth-harmonic signal.

    ∀ h ∈ {0}:
        s₀(t) = ξ₀⁽⁰⁾(t) ∈ 𝐑
    ∀ h ∉ {0}:
        sₕ(t) = ℜ[ξₕ⁽⁰⁾(t) + conj(ξₕ⁽⁰⁾(t))] ∈ 𝐑

Note: In its essence, the real operator `ℜ[]`` is unnecessary, however, it is
used to convert the complex number `x + j0` to a real number `x`.

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]
- `H::Int`                      | harmonic number [-]

Output:
- `s::Vector{<:Real}`           | signal sₕ(t) [?]
"""
signal(sol::AbstractDTFTSolution, H::Int) =
    ifelse(H == 0, ξ(sol,0,H), real(ξ(sol,0,H) + conj(ξ(sol,0,H))))

# error
"""
    TFT.error(sol::TFT.AbstractDTFTSolution)

Function to obtain the error between the input and computed signal.

Input:
- `sol::AbstractDTFTSolution`   | DTFT solution struct [-]

Output:
- `e::Vector{<:Real}`           | error [?]
"""
error(sol::TFT.AbstractDTFTSolution) = sol.prob.s .- signal(sol)


using JuMP, Mosek, MosekTools, LinearAlgebra, Suppressor, StatsBase, CSV, DataFrames, Ipopt, Compat
include("convexpenalties.jl")
include("alt_min.jl")

gammaMult=5.0 # Need to vary this to generate the four plots

results_template = DataFrame(n=Int[], r=Int[], p=Real[], y=Real[],
    error_GD=Real[], error_AM=Real[], rank_GD=Int[], rank_AM=Int[], LB_GD=Real[], UB_GD=Real[], UB_AM=Real[])

results_basispursuit=similar(results_template, 0)

pRange=collect(0.1:0.1:1.0)
pSize=size(pRange, 1)
yRange=collect(0.1:0.1:1.0)
ySize=size(yRange, 1)

for ARG in ARGS # 1-100
  for iter in 1:5
  n=100

  array_num = parse(Int, ARG)

  y=yRange[(array_num-1)÷ySize+1]
  p=pRange[(array_num-1)%pSize+1]

  r=Int(round(n*(1.0-sqrt(1.0-y*p))))
  if r==0
      r=1
  end
  k=r #assume we picked it via say cross-validation (can only show two dimensions on the plot).
  M_L=randn(n, r)
  M_R=randn(n, r)
  M=M_L*M_R'
  numEntriesSampled=Int(round(p*n^2))
  gamma=gammaMult/p #setting this to be really high as a good proxy for exact constraints (no noise)

  coords=[repeat(1:n, inner=size(1:n,1)) repeat(1:n, outer=size(1:n,1))];
  entries=sample(1:n^2, Int(numEntriesSampled), replace=false)

  ############## Greedy rounding  #####################################################
  @show k
  X_GD,LB_GD, UB_GD,=getMatrixFrobeniusNorm_greedy(coords, entries, M, n, numEntriesSampled, k, gamma, false)

  u,v,w=svd(X_GD);
  @show rnk_GD=sum(v.>1e-2);
  @show err_GD = norm(X_GD.-M)/norm(M);

  ########### Alternating Minimization ###########################################################
  X_AM, UB_AM=alternatingminimization_matrixcompletion(coords, entries, M, n, numEntriesSampled, k, gamma, 50, X_GD)
  u,v,=svd(X_AM)
  @show rnk_AM=sum(v.>1e-2);
  @show err_AM = norm(X_AM.-M)/norm(M);

  ########### Push results #############################################################

  push!(results_basispursuit, [n, r, p, y, err_GD, err_AM, rnk_GD, rnk_AM, LB_GD, UB_GD, UB_AM])
end
  CSV.write("dualitygapplot_gamma5p0onp.csv", results_basispursuit, append=true)

end


module CalibrationErrors

using DataStructures
using Distances
using Parameters
using StatsBase

using LinearAlgebra
using Statistics

# estimation
export calibrationerror

# estimators
export ECE, BiasedSKCE, LinearUnbiasedSKCE, QuadraticUnbiasedSKCE

# binning algorithms
export UniformBinning, MedianVarianceBinning

# scalar-valued kernels
export SquaredExponentialKernel, ExponentialKernel, GaussianKernel, LaplacianKernel

# matrix-valued kernels
export UniformScalingKernel, DiagonalKernel

include("generic.jl")
include("binning/generic.jl")
include("binning/uniform.jl")
include("binning/medianvariance.jl")
include("ece.jl")
include("kernels/generic.jl")
include("kernels/scalar.jl")
include("kernels/matrix.jl")
include("skce/generic.jl")
include("skce/biased.jl")
include("skce/unbiased.jl")

end # module


"Compute pressure levels, vmr, vcd for atmospheric profile, given p_half, T, q"
function compute_atmos_profile_fields(p_half::AbstractArray, T, q; g₀=9.8196)

    # Floating type to use
    FT = eltype(T)

    # Calculate full pressure levels
    p_full = (p_half[2:end] + p_half[1:end - 1]) / 2

    # Dry and wet mass
    dry_mass = 28.9647e-3  / Nₐ  # in kg/molec, weighted average for N2 and O2
    wet_mass = 18.01528e-3 / Nₐ  # just H2O
    ratio = dry_mass / wet_mass
    n_layers = length(T)

    # Also get a VMR vector of H2O (volumetric!)
    vmr_h2o = zeros(FT, n_layers, )
    vcd_dry = zeros(FT, n_layers, )
    vcd_h2o = zeros(FT, n_layers, )

    # Now actually compute the layer VCDs
    for i = 1:n_layers
        Δp = p_half[i + 1] - p_half[i]
        vmr_h2o[i] = q[i] * ratio
        vmr_dry = 1 - vmr_h2o[i]
        M  = vmr_dry * dry_mass + vmr_h2o[i] * wet_mass
        vcd_dry[i] = vmr_dry * Δp / (M * g₀ * 100.0^2)   # includes m2->cm2
        vcd_h2o[i] = vmr_h2o[i] * Δp / (M * g₀ * 100^2)
    end

    return p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o

end

"Validate an input atmospheric profile (YAML only)"
function validate_atmos_profile(params::Dict{Any,Any})

    @assert Set(keys(params)) == Set(["p_surf", "T", "q", "ak", "bk"]) ||
            Set(keys(params)) == Set(["p_half", "T", "q"]) ||
            Set(keys(params)) == Set(["p_half", "T"]) "Set of atmospheric profile fields must be one of (p_surf, T, q, ak, bk), (p_half, T, q), or (p_half, T)"

end

"From a yaml file, get the stored fields (psurf, T, q, ak, bk), calculate derived fields,
and return an AtmosphericProfile object"
function read_atmos_profile(file_path::String)

    # Make sure file is yaml type
    @assert endswith(file_path, ".yaml") "File must be yaml"

    # Read in the data and pass to compute fields
    params_dict = YAML.load_file(file_path)

    # Validate the parameters before doing anything else
    validate_atmos_profile(params_dict)

    T = convert.(Float64, params_dict["T"])

    # Calculate derived fields
    if ("ak" in keys(params_dict))
        psurf = convert(Float64, params_dict["p_surf"])
        q     = convert.(Float64, params_dict["q"])
        ak    = convert.(Float64, params_dict["ak"])
        bk    = convert.(Float64, params_dict["bk"])
        p_half = (ak + bk * psurf)
        p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)
    elseif ("q" in keys(params_dict))
        p_half = convert(Float64, params_dict["p_half"])
        q      = convert.(Float64, params_dict["q"])
        p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)
    else
        p_half = convert(Float64, params_dict["p_half"])
        q = zeros(length(T))
        p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)
    end

    # Return the atmospheric profile struct
    return AtmosphericProfile(nothing, nothing, psurf, T, q, p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o)

end

"Read atmospheric profile from netCDF file (ONLY works for FraLab test file)"
function read_atmos_profile(file::String, lat::Real, lon::Real, time_idx::Int)

    # Time index must be ∈ [1, 2, 3, 4]
    @assert 1 <= time_idx <= 4 "Time index must be ∈ [1, 2, 3, 4]"

    # Make sure file is nc type
    @assert endswith(file, ".nc4") "File must be nc4"

    # Load in the atmospheric profile
    @timeit "loading file" ds = Dataset(file)

    # See how easy it is to actually extract data?
    # Note the [:] in the end reads in ALL the data in one step
    file_lats, file_lons = ds["YDim"][:], ds["XDim"][:]

    # Convert the input lat/lon to right type
    FT = eltype(file_lats)
    lat, lon = FT(lat), FT(lon)

    # Find index (nearest neighbor, one could envision interpolation in space and time!):
    lat_idx, lon_idx = argmin(abs.(file_lats .- lat)), argmin(abs.(file_lons .- lon))

    # Temperature profile
    @timeit "getting T" T = convert(Array{FT,1}, ds["T"][lon_idx, lat_idx,  :, time_idx])

    # Specific humidity profile
    q = convert(Array{FT,1}, ds["QV"][lon_idx, lat_idx, :, time_idx])

    # Surface pressure
    psurf = convert(FT, ds["PS"][lon_idx, lat_idx, time_idx])

    # AK and BK global attributes (important to calculate pressure half-levels)
    ak, bk = ds.attrib["HDF_GLOBAL.ak"][:], ds.attrib["HDF_GLOBAL.bk"][:]

    # Close the file
    close(ds)

    # Calculate derived fields
    p_half = (ak + bk * psurf)
    p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)

    # Return the atmospheric profile struct
    return AtmosphericProfile(lat, lon, psurf, T, q, p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o)
end

"Reduce profile dimensions by re-averaging to near-equidistant pressure grid"
function reduce_profile(n::Int, profile::AtmosphericProfile{FT}) where {FT}

    # Can only reduce the profile, not expand it
    @assert n < length(profile.T)

    # Unpack the profile
    @unpack lat, lon, psurf = profile

    # New rough half levels (boundary points)
    a = range(0, maximum(profile.p), length=n + 1)

    # Matrices to hold new values
    T = zeros(FT, n);
    q = zeros(FT, n);
    p_full = zeros(FT, n);
    p_levels = zeros(FT, n + 1);
    vmr_h2o  = zeros(FT, n);
    vcd_dry  = zeros(FT, n);
    vcd_h2o  = zeros(FT, n);

    # Loop over target number of layers
    for i = 1:n

        # Get the section of the atmosphere with the i'th section pressure values
        ind = findall(a[i] .< profile.p .<= a[i + 1]);

        # Set the pressure levels accordingly
        p_levels[i] = profile.p_levels[ind[1]]
        p_levels[i + 1] = profile.p_levels[ind[end]]

        # Re-average the other parameters to produce new layers
        p_full[i] = mean(profile.p_levels[ind])
        T[i] = mean(profile.T[ind])
        q[i] = mean(profile.q[ind])
        vmr_h2o[i] = mean(profile.vmr_h2o[ind])
        vcd_dry[i] = sum(profile.vcd_dry[ind])
        vcd_h2o[i] = sum(profile.vcd_h2o[ind])
    end

    return AtmosphericProfile(lat, lon, psurf, T, q, p_full, p_levels, vmr_h2o, vcd_dry, vcd_h2o)
end;

"""
$(FUNCTIONNAME)(psurf, λ, depol_fct, vcd_dry)
Returns the Rayleigh optical thickness per layer at reference wavelength `λ` (N₂,O₂ atmosphere, i.e. terrestrial)
Input:
- `psurf` surface pressure in `[hPa]`
- `λ` wavelength in `[μm]`
- `depol_fct` depolarization factor
- `vcd_dry` dry vertical column (no water) per layer
"""
function getRayleighLayerOptProp(psurf, λ, depol_fct, vcd_dry)
    FT = eltype(λ)
    # Total vertical Rayleigh scattering optical thickness
    tau_scat = FT(0.00864) * (psurf / FT(1013.25)) * λ^(-FT(3.916) - FT(0.074) * λ - FT(0.05) / λ)
    tau_scat = tau_scat * (FT(6.0) + FT(3.0) * depol_fct) / (FT(6.0)- FT(7.0) * depol_fct)
    Nz = length(vcd_dry)
    τRayl = zeros(FT,Nz)
    k = tau_scat / sum(vcd_dry)
    for i = 1:Nz
        τRayl[i] = k * vcd_dry[i]
    end

    return convert.(FT, τRayl)
end

"""
$(FUNCTIONNAME)(total_τ, p₀, σp, p_half)
Returns the aerosol optical depths per layer using a Gaussian distribution function with p₀ and σp on a pressure grid
"""
function getAerosolLayerOptProp(total_τ, p₀, σp, p_half)

    # Need to make sure we can also differentiate wrt σp (FT can be Dual!)
    FT = eltype(p₀)
    Nz = length(p_half)
    ρ = zeros(FT,Nz)

    for i = 2:Nz
        dp = p_half[i] - p_half[i - 1]
        ρ[i] = (1 / (σp * sqrt(2π))) * exp(-(p_half[i] - p₀)^2 / (2σp^2)) * dp
    end
    Norm = sum(ρ)
    τAer  =  (total_τ / Norm) * ρ
    return convert.(FT, τAer)
end

"""
$(FUNCTIONNAME)(τRayl, τAer,  aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type)
Computes the composite layer single scattering parameters (τ, ϖ, Z⁺⁺, Z⁻⁺)

Returns:
- `τ`, `ϖ`   : only Rayleigh scattering and aerosol extinction, no gaseous absorption (no wavelength dependence)
- `τ_λ`,`ϖ_λ`: Rayleigh scattering + aerosol extinction + gaseous absorption (wavelength dependent)
- `Z⁺⁺`,`Z⁻⁺`: Composite Phase matrix (weighted average of Rayleigh and aerosols)

# Arguments
- `τRay` layer optical depth for Rayleigh
- `τAer` layer optical depth for Aerosol(s) (vector)
- `aerosol_optics` array of aerosol optics struct
- `Rayl𝐙⁺⁺` Rayleigh 𝐙⁺⁺ phase matrix (2D)
- `Rayl𝐙⁻⁺` Rayleigh 𝐙⁻⁺ phase matrix (2D)
- `Aer𝐙⁺⁺` Aerosol 𝐙⁺⁺ phase matrix (3D)
- `Aer𝐙⁻⁺` Aerosol 𝐙⁻⁺ phase matrix (3D)
- `τ_abs` layer absorption optical depth array (per wavelength) by gaseous absorption
"""
function construct_atm_layer(τRayl, τAer,  aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type)
    FT = eltype(τRayl)
    nAer = length(aerosol_optics)

    # Fix Rayleigh SSA to 1
    ϖRayl = FT(1)

    @assert length(τAer) == nAer "Sizes don't match"

    τ = FT(0)
    ϖ = FT(0)
    A = FT(0)
    Z⁺⁺ = similar(Rayl𝐙⁺⁺);
    Z⁻⁺ = similar(Rayl𝐙⁺⁺);

    if (τRayl + sum(τAer)) < eps(FT)
        fill!(Z⁺⁺, 0); fill!(Z⁻⁺, 0);
        return FT(0), FT(1), Z⁺⁺, Z⁻⁺
    end

    τ += τRayl
    ϖ += τRayl * ϖRayl
    A += τRayl * ϖRayl

    Z⁺⁺ = τRayl * ϖRayl * Rayl𝐙⁺⁺
    Z⁻⁺ = τRayl * ϖRayl * Rayl𝐙⁻⁺

    for i = 1:nAer
        τ   += τAer[i]
        ϖ   += τAer[i] * aerosol_optics[i].ω̃
        A   += τAer[i] * aerosol_optics[i].ω̃ * (1 - aerosol_optics[i].fᵗ)
        Z⁺⁺ += τAer[i] * aerosol_optics[i].ω̃ * (1 - aerosol_optics[i].fᵗ) * Aer𝐙⁺⁺[:,:,i]
        Z⁻⁺ += τAer[i] * aerosol_optics[i].ω̃ * (1 - aerosol_optics[i].fᵗ) * Aer𝐙⁻⁺[:,:,i]
    end

    Z⁺⁺ /= A
    Z⁻⁺ /= A
    A /= ϖ
    ϖ /= τ

    # Rescaling composite SSPs according to Eqs. A.3 of Sanghavi et al. (2013) or Eqs.(8) of Sanghavi & Stephens (2015)
    τ *= (FT(1) - (FT(1) - A) * ϖ)
    ϖ *= A / (FT(1) - (FT(1) - A) * ϖ)#Suniti

    # Adding absorption optical depth / albedo:
    τ_λ = τ_abs .+ τ
    ϖ_λ = (τ .* ϖ) ./ τ_λ

    return Array(τ_λ), Array(ϖ_λ), τ, ϖ, Array(Z⁺⁺), Array(Z⁻⁺)
end

"When performing RT_run, this function pre-calculates properties for all layers, before any Core RT is performed"
function construct_all_atm_layers(FT, nSpec, Nz, NquadN, τRayl, τAer, aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type, qp_μ, μ₀, m)

    # Empty matrices to hold all values
    τ_λ_all   = zeros(FT, nSpec, Nz)
    ϖ_λ_all   = zeros(FT, nSpec, Nz)
    τ_all     = zeros(FT, Nz)
    ϖ_all     = zeros(FT, Nz)
    Z⁺⁺_all   = zeros(FT, NquadN, NquadN, Nz)
    Z⁻⁺_all   = zeros(FT, NquadN, NquadN, Nz)

    dτ_max_all  = zeros(FT, Nz)
    dτ_all      = zeros(FT, Nz)
    ndoubl_all  = zeros(Int64, Nz)
    dτ_λ_all    = zeros(FT, nSpec, Nz)
    expk_all    = zeros(FT, nSpec, Nz)
    scatter_all = zeros(Bool, Nz)

    Threads.@threads for iz=1:Nz

        # Construct atmospheric properties
        τ_λ_all[:, iz], ϖ_λ_all[:, iz], τ_all[iz], ϖ_all[iz], Z⁺⁺_all[:,:,iz], Z⁻⁺_all[:,:,iz] = construct_atm_layer(τRayl[iz], τAer[:,iz], aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs[:,iz], arr_type)

        # Compute doubling number
        dτ_max_all[iz] = minimum([τ_all[iz] * ϖ_all[iz], FT(0.001) * minimum(qp_μ)])
        dτ_all[iz], ndoubl_all[iz] = doubling_number(dτ_max_all[iz], τ_all[iz] * ϖ_all[iz]) #Suniti

        # Compute dτ vector
        dτ_λ_all[:, iz] = (τ_λ_all[:, iz] ./ (FT(2)^ndoubl_all[iz]))
        #@show maximum(dτ_λ_all[:,iz])
        expk_all[:, iz] = exp.(-dτ_λ_all[:, iz] /μ₀) #Suniti

        # Determine whether there is scattering
        scatter_all[iz] = (  sum(τAer[:,iz]) > 1.e-8 ||
                          (( τRayl[iz] > 1.e-8 ) && (m < 3))) ?
                            true : false
    end

    # Compute sum of optical thicknesses of all layers above the current layer
    τ_sum_all = accumulate(+, τ_λ_all, dims=2)

    # First start with all zeros
    # At the bottom of the atmosphere, we have to compute total τ_sum (bottom of lowest layer), for the surface interaction
    τ_sum_all = hcat(zeros(FT, size(τ_sum_all[:,1])), τ_sum_all)

    # Starting scattering interface (None for both added and composite)
    scattering_interface = ScatteringInterface_00()
    scattering_interfaces_all = []

    for iz = 1:Nz
        # Whether there is scattering in the added layer, composite layer, neither or both
        scattering_interface = get_scattering_interface(scattering_interface, scatter_all[iz], iz)
        push!(scattering_interfaces_all, scattering_interface)
    end

    return ComputedAtmosphereProperties(τ_λ_all, ϖ_λ_all, τ_all, ϖ_all, Z⁺⁺_all, Z⁻⁺_all, dτ_max_all, dτ_all, ndoubl_all, dτ_λ_all, expk_all, scatter_all, τ_sum_all, scattering_interfaces_all)
end

"Given the CrossSectionModel, the grid, and the AtmosphericProfile, fill up the τ_abs array with the cross section at each layer
(using pressures/temperatures) from the profile"
function compute_absorption_profile!(τ_abs::Array{FT,2},
                                     model::AbstractCrossSectionModel,
                                     grid,
                                     profile::AtmosphericProfile,
                                     ) where FT <: AbstractFloat

    # The array to store the cross-sections must be same length as number of layers
    @assert size(τ_abs,2) == length(profile.p)

    for iz in 1:length(profile.p)

        # Pa -> hPa
        p = profile.p[iz] / 100
        T = profile.T[iz]

        # Changed index order
        τ_abs[:,iz] = Array(absorption_cross_section(model, grid, p, T)) * profile.vcd_dry[iz] * model.vmr
    end

    return nothing

end


using LSDescent

using Test
using NLPModels, JuMP,  NLPModelsJuMP

using SolverCore, Logging

using LinearAlgebra

# test all solvers with the well known Woods test function
include("woods.jl")
nlp = MathOptNLPModel(woods(40), name="woods")
n = nlp.meta.nvar

include("tests.jl")

#println("Testing bfgs solvers\n\n")
#include("Debugbfgs.jl")

#println("Testing Newton solvers\n\n")
#include("DebugNewton.jl")


using WeibullCount

function test_weibull_count_pdf(x, λ, c; k = 200, transformed=true)
    for T1 ∈ [Int64, Int128, BigInt]
        for T2 ∈ [Float64, BigFloat]
            WeibullCount.cache_clear!()

            print("$T1 and $T2: ")
            println(@time weibull_count_pdf(T1(x), T2(λ), T2(c), T2(1.0); k = T1(k), transform=transformed))
        end
    end
end

# Ensure `weibull_count_pdf_approx` works with different types.
@testset "Parametric types" begin
    x = 2
    c = 0.36787944117144233
    λ = 20.08553692

    for T1 ∈ [Int64, Int128, BigInt]
        for T2 ∈ [Float64, BigFloat]
            WeibullCount.cache_clear!()

            # print("$T1 and $T2: ")
            res, converged = WeibullCount.weibull_count_pdf_approx(T1(x), T2(λ), T2(c), T2(1.0))
            @test isa(res, T2)
        end
    end
end

# Failure of convergence
@testset "DivergedError" begin
    x = 2
    c = 0.36787944117144233
    λ = 20.08553692

    WeibullCount.cache_clear!()
    @test_throws WeibullCount.DivergedError WeibullCount.weibull_count_pdf(x, λ, c)

    WeibullCount.cache_clear!()
    @test WeibullCount.weibull_count_pdf(BigInt(x), BigFloat(λ), BigFloat(c)) ≥ 0
end

# Difficult parameters
@testset "Run-times" begin
    x = 2
    c = 0.36787944117144233
    λ = 20.08553692

    WeibullCount.cache_clear!()
    t1 = @elapsed WeibullCount.weibull_count_pdf_old(BigInt(x), BigFloat(λ), BigFloat(c), BigFloat(1.0); k = BigInt(200))

    WeibullCount.cache_clear!()
    t2 = @elapsed weibull_count_pdf(BigInt(x), BigFloat(λ), BigFloat(c), BigFloat(1.0); k = BigInt(200), increment = BigInt(10), tol = BigFloat(0.001))

    @test t2 < t1
end


@testset "encoding" begin
    # TODO add more examples and types to test

    hdvs = [GradedHDV(10) for i in 1:5]
    alphabet = "STOCK"
    basis = Dict(zip(alphabet,hdvs))

    @testset "NGrams" begin

        fourgrams = compute_4_grams(hdvs)
        fourgrams_chars = compute_4_grams(basis)

        @test fourgrams.d[2][1][4][3] ≈ Π(hdvs[2], 0) * Π(hdvs[1], 1) * Π(hdvs[4], 2) * Π(hdvs[3], 3)
        @test fourgrams.d[2][1][4][3] ≈ fourgrams_chars.d['T']['S']['C']['O']
    end

    @testset "sequence embedding" begin
        sequence = [2, 3, 5, 1, 3, 4]

        emb = sequence_embedding(sequence, hdvs, 3)

        @test typeof(emb) == eltype(hdvs)

        threegrams = compute_3_grams(hdvs)

        @test sequence_embedding(sequence, threegrams) ≈ emb

        threegrams = compute_3_grams(basis)

        @test sequence_embedding([alphabet[i] for i in sequence], threegrams) ≈ emb
    end
end


"""
    CountingTropical{T,CT} <: Number

Counting tropical number type is also a semiring algebra.
It is tropical algebra with one extra field for counting, it is introduced in [arXiv:2008.06888](https://arxiv.org/abs/2008.06888).

Example
-------------------------
```jldoctest; setup=:(using TropicalNumbers)
julia> CountingTropical(1.0, 5.0) + CountingTropical(3.0, 2.0)
(3.0, 2.0)ₜ

julia> CountingTropical(1.0, 5.0) * CountingTropical(3.0, 2.0)
(4.0, 10.0)ₜ

julia> one(CountingTropicalF64)
(0.0, 1.0)ₜ

julia> zero(CountingTropicalF64)
(-Inf, 0.0)ₜ
```
"""
struct CountingTropical{T,CT} <: Number
    n::T
    c::CT
end
CountingTropical(x::T) where T<:Real = CountingTropical(x, one(T))
CountingTropical{T1}(x) where {T1} = CountingTropical{T1,T1}(x)

CountingTropical{T1,CT1}(x) where {T1, CT1} = CountingTropical{T1,CT1}(T1(x), one(CT1))
CountingTropical{T1,CT1}(x::CountingTropical{T1,CT1}) where {T1,CT1} = x
CountingTropical{T1,CT1}(x::CountingTropical) where {T1,CT1} = CountingTropical(T1(x.n), CT1(x.c))
CountingTropical{T1,CT1}(x::Tropical) where {T1,CT1} = CountingTropical{T1,CT1}(T1(TropicalNumbers.content(x)), one(CT1))

Base.:*(a::CountingTropical, b::CountingTropical) = CountingTropical(a.n + b.n, a.c * b.c)
Base.:^(a::CountingTropical, b::Real) = CountingTropical(a.n * b, a.c ^ b)
Base.:^(a::CountingTropical, b::Integer) = CountingTropical(a.n * b, a.c ^ b)
function Base.:+(a::CountingTropical, b::CountingTropical)
    n = max(a.n, b.n)
    if a.n > b.n
        c = a.c
    elseif a.n == b.n
        c =  a.c + b.c
    else
        c = b.c
    end
    CountingTropical(n, c)
end
# inverse and division
Base.inv(x::CountingTropical) = CountingTropical(-x.n, x.c)

Base.zero(::Type{CountingTropical{T}}) where T = zero(CountingTropical{T,T})
Base.zero(::Type{CountingTropical{T,CT}}) where {T<:Integer,CT} = CountingTropical(T(-999999), zero(CT))
Base.zero(::Type{CountingTropical{T,CT}}) where {T<:AbstractFloat,CT} = CountingTropical(typemin(T), zero(CT))
Base.zero(::T) where T<:CountingTropical = zero(T)
Base.one(::Type{CountingTropical{T}}) where T = one(CountingTropical{T,T})
Base.one(::Type{CountingTropical{T,CT}}) where {T<:Integer,CT} = CountingTropical(zero(T), one(CT))
Base.one(::Type{CountingTropical{T,CT}}) where {T<:AbstractFloat,CT} = CountingTropical(zero(T), one(CT))
Base.one(::T) where T<:CountingTropical = one(T)
Base.isapprox(a::CountingTropical, b::CountingTropical; kwargs...) = isapprox(a.n, b.n; kwargs...) && isapprox(a.c, b.c; kwargs...)

Base.show(io::IO, t::CountingTropical) = Base.print(io, "$((t.n, t.c))ₜ")

Base.promote_type(::Type{CountingTropical{T1,CT1}}, b::Type{CountingTropical{T2,CT2}}) where {T1,T2,CT1,CT2} = CountingTropical{promote_type(T1,T2), promote_type(CT1,CT2)}


module ErrorFreeTransforms

typealias IEEE754 Union(Float32, Float64)

include("add.jl")
include("mul.jl")
include("fma.jl")

export
  # add.jl
  err_add,
  err_fast_add,

  # mul.jl
  fsplit,
  err_mul,
  err_fast_mul,

  # fma.jl
  err_fma,
  err_fma_compr,
  err_fma_nearest,
  err_fma_approx

end


export MI

#=
Bias correction techniques

=#

#=
Quadratic extrapolation
Strong et al 1998

I=I_raw + a/n_s + b/n_s^2

Compute MI with number of repetitions (n_s) equal to total, total/2, and total/4
Plot result against 1/n_s, and fit quadratic polynomial by least squares
Intercept is true information
=#

function MI{B<:QE, T<:transformation}(info::Information{B,T},response::Array{Float64,2},stimulus::Array{Float64,1})

    N_s=length(stimulus)

    myinds=1:N_s

    samplepoints=round(Int64,[N_s; N_s/2; N_s/4])

    MI=zeros(Float64,3)

    (conmat, p_s, p_r)=transform(info,response,stimulus)

    MI[1]=mutual_information(conmat,p_s,p_r)

    for i=2:3

        theseinds=sample(myinds, samplepoints[i], replace=false)

        stimulus_2=stimulus[theseinds]
        response_2=response[theseinds,:]

        (conmat2, p_s2, p_r2)=transform(info,response_2,stimulus_2)

        MI[i]=mutual_information(conmat2,p_s2,p_r2)

    end

    I_corrected=polyfit(1./samplepoints, MI, 2)[1]

    (I_corrected, conmat)

end

#=
Panzeri-Treves (PT)
Panzeri and Treves 1996
=#

#=
Nemenman-Shafee-Bialek (NSB)
Nemenman et al 2002
=#

#=
James-Stein shrinkage
Hausser and Strimmer 2008
=#

function mutual_information(p_rs::Array{Float64,2},p_s::Array{Float64,1},p_r::Array{Float64,1})

    MI=0.0

    for i=1:size(p_rs,1)
        for j=1:size(p_rs,2)
            if (p_rs[i,j]==0.0) || (isnan(p_rs[i,j]))
            elseif p_r[j]==0.0
            else
                MI+=p_rs[i,j]*p_s[i]*log2(p_rs[i,j]/(p_r[j]))
            end

        end
    end

    MI

end

function transform{B<:bias,T<:decoder}(info::Information{B,T},response::Array{Float64,2},stimulus::Array{Float64,1})
    predict=project(info.t,response,stimulus)
    (conmat, p_s, p_r)=conmatrix(info.t,predict,stimulus)
end

function polyfit(x::Array{Float64,1}, y::Array{Float64,1}, n::Int64)
  A = [ x[i]^p for i = 1:length(x), p = 0:n ]
  A \ y
end


@inline function add_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(add_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? add_dbdb_db_nonfinite(x,y) : res
end

@inline function sub_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(sub_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? sub_dbdb_db_nonfinite(x,y) : res
end

@inline function mul_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(mul_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? mul_dbdb_db_nonfinite(x,y) : res
end

@inline function dvi_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(dvi_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? dvi_dbdb_db_nonfinite(x,y) : res
end

@inline function add_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = HI(x) + HI(y)
    return DoubleFloat(z, zero(T))
end

@inline function sub_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = HI(x) - HI(y)
    return DoubleFloat(z, zero(T))
end

@inline function mul_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = HI(x) * HI(y)
    return DoubleFloat(z, zero(T))
end

@inline function dvi_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = div(HI(x), HI(y))
    return DoubleFloat(z, zero(T))
end


module CatElements
export ThElements, AbstractElements, Elements, elements

using DataStructures: OrderedDict

using ..CSets, ..FinSets
using ...Present, ...Theories
using ...Theories: Category, ob, hom, dom_nums, codom_nums

@present ThElements(FreeSchema) begin
  (El, Arr, Ob, Hom)::Ob
  Name::AttrType
  src::Hom(Arr, El)
  tgt::Hom(Arr, El)
  dom::Hom(Hom, Ob)
  cod::Hom(Hom, Ob)
  πₑ::Hom(El, Ob)
  πₐ::Hom(Arr, Hom)
  nameo::Attr(Ob, Name)
  nameh::Attr(Hom, Name)

  src⋅πₑ == πₐ⋅dom
  tgt⋅πₑ == πₐ⋅cod
end

@abstract_acset_type AbstractElements
@acset_type Elements(ThElements, index=[:src, :tgt, :πₑ, :πₐ]) <: AbstractElements

"""    elements(X::AbstractACSet)

construct the category of elements from an ACSet. This only correctly handles the CSet part.
This transformation converts an instance of C into a Graph homomorphism. The codomain of the
homomorphism is a graph shaped like the schema. This is one half of the isomorphism between
databases and knowledge graphs.
"""
function elements(X::StructACSet{S}) where S
  Y = Elements{Symbol}()

  obs = ob(S)
  add_parts!(Y, :Ob, length(obs), nameo=obs)
  els = map(enumerate(obs)) do (i,c)
    add_parts!(Y, :El, nparts(X, c), πₑ = i)
  end

  add_parts!(Y, :Hom, length(hom(S)), dom=dom_nums(S), cod=codom_nums(S), nameh=hom(S))
  map(enumerate(zip(hom(S), dom_nums(S), codom_nums(S)))) do (i, (f, ci, di))
    c, d = obs[ci], obs[di]
    nc = nparts(X, c)
    add_parts!(Y, :Arr, nparts(X, c), src=els[ci], tgt=view(els[di], X[f]), πₐ=i)
  end
  return Y
end

"""    presentation(X::AbstractElements)

convert a category of elements into a new schema. This is useful for generating large schemas
that are defined as the category of elements of a specified C-Set. For example, the schema for
Bipartite Graphs is the category of elements of the graph with 2 vertices and 1 edge. The 2 vertices
will get mapped to elements `V_1, V_2` and the one edge will be `E_1` and the source and target maps
will be mapped to `src_E_1, tgt_E_1`.
"""
function presentation(X::AbstractElements)
  P = Presentation(FreeSchema)
  obs = OrderedDict{Tuple{Symbol, Int}, Any}()
  for a in 1:nparts(X, :Ob)
    class = X[a, :nameo]
    for (i,j) in enumerate(incident(X, a, :πₑ))
      ob = Ob(FreeSchema, Symbol("$(class)_$i"))
      obs[(class, j)] = ob
    end
  end
  add_generators!(P, values(obs))
  homs = OrderedDict{Tuple{Symbol, Int}, Any}()
  for f in 1:nparts(X, :Hom)
    class = X[f, :nameh]
    for (i,j) in enumerate(incident(X, f, :πₐ))
      domkey = (X[j, [:src, :πₑ, :nameo]], X[j, :src])
      dom = obs[domkey]
      codkey = (X[j, [:tgt, :πₑ, :nameo]], X[j, :tgt])
      codom = obs[codkey]
      homs[(class, j)] = Hom(Symbol("$(class)_$dom"), dom, codom)
    end
  end
  add_generators!(P, values(homs))
  return P, obs, homs
end
end


# Linear Quadratic Regulator Objective Functions

using SparseArrays

include("QP_Linear_objectives.jl")

abstract type objectiveLQR_abstract <: objectiveFunc end

"""
    LQR_simple(Q, R)

A simple LQR struct where `Q` and `R` are constant for all `x` and `u`.
"""
struct LQR_simple <: objectiveLQR_abstract
    Q
    R
end

"""
    LQR_QP(QR_Full)

An LQR what acts like a QP objective function.
"""
struct LQR_QP <: objectiveLQR_abstract
    QR_Full
end

"""
    LQR_QP_Referenced(lqr_qp::LQR_QP, xRef)

An LQR what acts like a QP objective function, but with a reference
trajectory.
"""
struct LQR_QP_Referenced <: objectiveLQR_abstract
    lqr_qp::LQR_QP
    xRef
end

"""
    LQR_QP_Ref_Terminal(lqr_qp::LQR_QP, xRef)

An LQR what acts like a QP objective function, but with a reference
trajectory.
"""
struct LQR_QP_Ref_Terminal <: objectiveLQR_abstract
    lqr_qp::LQR_QP
    xRef
end

#####################################
#  LQR -> ObjectiveQP Constructors  #
#####################################
@doc raw"""
    makeLQR_TrajSimple(Q, R, NSteps::Int64)

Pseudo-constructor for an LQR objective function (which is really an
[`objectiveQ`](@ref) objective function).

An LQR struct for a trajectory is one that holds the sparse matrix matching:

```math
B_{QR} = [Q 0 0 0 0; 0 R 0 0 0; 0 0 Q 0 0; 0 0 0 R 0; 0 0 0 0 Q]
```

Where size corresponds to the size of the trajectory (or trajectory horizon).

```math
f(y) = \frac{1}{2} \ y^{\top} * B_{QR} * y
```

See also: [`LQR_simple`](@ref) and [`LQR_QP`](@ref)
"""
function makeLQR_TrajSimple(Q, R, NSteps::Int64)
    QSize = size(Q, 1)
    RSize = size(R, 1)
    totSize = (QSize + RSize) * NSteps + QSize

    # Initialize and empty sparse array
    QRFull = spzeros(totSize, totSize)

    rStart = 1

    for k in 1:NSteps
        rEnd = rStart + QSize - 1
        QRFull[rStart:rEnd, rStart:rEnd] = Q

        rStart += QSize
        rEnd = rStart + RSize - 1

        QRFull[rStart:rEnd, rStart:rEnd] = R

        rStart += RSize
    end
    rEnd = rStart + QSize - 1
    QRFull[rStart:rEnd, rStart:rEnd] = Q

    return LQR_QP(QRFull)
end

function makeLQR_TrajSimple(lqr::LQR_simple, NSteps::Int64)
    return makeLQR_TrajSimple(lqr.Q, lqr.R, NSteps)
end

function makeLQR_TrajReferenced(Q, R, NSteps::Int64, xRef)
    baseLQR = makeLQR_TrajSimple(Q, R, NSteps)
    return LQR_QP_Referenced(baseLQR, xRef)
end

function makeLQR_TrajReferenced(lqr::LQR_simple, NSteps::Int64, xRef)
    baseLQR = makeLQR_TrajSimple(lqr.Q, lqr.R, NSteps)
    return LQR_QP_Referenced(baseLQR, xRef)
end

@doc raw"""
    makeLQR_TrajTerminal(Q, R, NSteps::Int64)

Pseudo-constructor for an LQR objective function (which is really an
[`objectiveQ`](@ref) objective function).

An LQR struct for a trajectory is one that holds the sparse matrix matching:

```math
B_{QR} = [0 0 0 0 0; 0 R 0 0 0; 0 0 0 0 0; 0 0 0 R 0; 0 0 0 0 Q]
```

Where size corresponds to the size of the trajectory (or trajectory horizon).

```math
f(y) = \frac{1}{2} \ y^{\top} * B_{QR} * y
```

See also: [`LQR_simple`](@ref) and [`LQR_QP`](@ref)
"""
function makeLQR_TrajTerminal(Q, R, NSteps::Int64)
    QSize = size(Q, 1)
    RSize = size(R, 1)
    totSize = (QSize + RSize) * NSteps + QSize

    # Initialize and empty sparse array
    QRFull = spzeros(totSize, totSize)

    rStart = 1 + QSize

    for k in 1:NSteps
        rEnd = rStart + RSize - 1

        QRFull[rStart:rEnd, rStart:rEnd] = R

        rStart += RSize + QSize
    end
    rStart -= QSize
    rEnd = rStart + QSize - 1
    QRFull[rStart:rEnd, rStart:rEnd] = Q

    return LQR_QP(QRFull)
end

function makeLQR_TrajRefTerminal(Q, R, NSteps::Int64, xRef)
    baseLQR = makeLQR_TrajTerminal(Q, R, NSteps)
    return LQR_QP_Referenced(baseLQR, xRef)
end

function makeLQR_TrajRefTerminal(lqr::LQR_simple, NSteps::Int64, xRef)
    baseLQR = makeLQR_TrajTerminal(lqr.Q, lqr.R, NSteps)
    return LQR_QP_Referenced(baseLQR, xRef)
end

#######################################
###       Evaluate An LQR QP        ###
#######################################

"""
    fObjQP(lqr::LQR_QP, x)

Evaluates a full QP-like LQR. See [`LQR_QP`](@ref).
"""
function fObjQP(lqr::LQR_QP, x)
    return (1/2) * x' * (lqr.QR_Full) * x
end

"""
    fObjQP(lqr::LQR_QP, x, xRef)

Evaluates a full QP-like LQR relative to a reference. (i.e. x -> (x - xRef))

See [`LQR_QP`](@ref).
"""
function fObjQP(lqr::LQR_QP, x, xRef)
    xDiff = x - xRef
    return fObjQP(lqr, xDiff)
end

"""
    fObjQP(lqrRef::LQR_QP_Referenced, x)

Evaluates a full QP-like LQR relative to a reference. (i.e. x -> (x - xRef))

See [`LQR_QP`](@ref), [`LQR_QP_Referenced`](@ref)
"""
function fObjQP(lqrRef::LQR_QP_Referenced, x)
    # println("x = $(size(x))")
    # println("xRef = $(size(lqrRef.xRef))")
    xDiff = x - lqrRef.xRef
    return fObjQP(lqrRef.lqr_qp, xDiff)
end

#######################################
###      Gradients of LQR QP        ###
#######################################
"""
    dfdxQP(lqr::LQR_QP, x)

Evaluates the derivative of a full QP-like LQR at input `x`.

See [`LQR_QP`](@ref), [`dfdxQP`](@ref)
"""
function dfdxQP(lqr::LQR_QP, x)
    return (lqr.QR_Full) * x
end

"""
    dfdxQP(lqr::LQR_QP, x, xRef)

Evaluates the derivative of a full QP-like LQR at input `x` with respect to
a reference `xRef`.

See [`LQR_QP`](@ref), [`dfdxQP`](@ref)
"""
function dfdxQP(lqr::LQR_QP, x, xRef)
    xDiff = x - xRef
    return dfdxQP(lqr, xDiff)
end

"""
    dfdxQP(lqrRef::LQR_QP_Referenced, x)

Evaluates the derivative of a full QP-like LQR at input `x` with respect to
a reference `xRef`.

See [`LQR_QP_Referenced`](@ref), [`LQR_QP`](@ref), [`dfdxQP`](@ref)
"""
function dfdxQP(lqrRef::LQR_QP_Referenced, x)
    return dfdxQP(lqrRef.lqr_qp, x, lqrRef.xRef)
end

#######################################
###      Hessians of LQR QP         ###
#######################################
"""
    hessQP(lqr::LQR_QP)

Evaluates the hessian of a full QP-like LQR, which is simple the base matrix.

See [`LQR_QP`](@ref)
"""
function hessQP(lqr::LQR_QP)
    return lqr.QR_Full
end

"""
    hessQP(lqr::LQR_QP)

Evaluates the hessian of a full QP-like LQR, which is simple the base matrix.

See [`LQR_QP`](@ref)
"""
function hessQP(lqrRef::LQR_QP_Referenced)
    return hessQP(lqrRef.lqr_qp)
end


am_demodulation(y2) = abs.(DSP.Util.hilbert(y2))

function gen_sync_frame(Fs2,sync_frequency)
    nbands = 7
    sync_frame = Vector{Vector{Int}}(undef,length(sync_frequency))

    i = 1;
    pulse_len = round(Int,Fs2/(2*sync_frequency[i]))
    # 7 pulses followed by silence
    sync_frame[i] = vcat(
        fill(-1,pulse_len),
        repeat(vcat(fill(-1,pulse_len),
                    fill(1,pulse_len)),
               nbands),
        fill(-1,5*pulse_len ))

    i = 2;
    pulse_len_on  = round(Int,3/5 * Fs2/(sync_frequency[i]))
    pulse_len_off = round(Int,2/5 * Fs2/(sync_frequency[i]))
    # 7 pulses followed by silence
    sync_frame[i] = vcat(
        fill(-1,pulse_len_off ),
        repeat(vcat(fill(-1,pulse_len_off),
                    fill(1,pulse_len_on)),
               nbands),
        fill(-1,pulse_len_on )) # the last -1 is as long a the previous +1

    return sync_frame
end

function find_sync(y_demod,sync_frame,frame_len)
    # minimum and maximum distance between sync frames
    mindistance = (8*frame_len) ÷ 10
    maxdistance = (12*frame_len) ÷ 10

    conv_sync_full = DSP.conv(y_demod,reverse(sync_frame));

    # skip first incomplete convolution
    conv_sync = @view conv_sync_full[length(sync_frame):end]

    # overall strongest sync frame
    # index0 marks the index of beginning the sync frame
    index0 = findmax(conv_sync)[2];

    # look for all sync frames after the strongest sync frame
    index = index0
    after_index = Int[]
    while index + maxdistance <= length(conv_sync)
        i = findmax(conv_sync[index .+ (mindistance:maxdistance)])[2] + index+mindistance-1
        push!(after_index,i)
        index = i
    end

    # look for all sync frames before the strongest sync frame
    index = index0
    before_index = Int[]
    while index - maxdistance >= 1
        i = findmax(conv_sync[index-maxdistance : index-mindistance])[2] + index-maxdistance-1
        push!(before_index,i)
        index = i
    end

    sync_frame_index = vcat(reverse(before_index),[index0],after_index)
    return sync_frame_index
end

function mark_sync(y_demod,sync_frame,frame_len)
    sync_frame_index = find_sync(y_demod,sync_frame,frame_len)
    tt = zeros(size(y_demod))
    tt[sync_frame_index] .= 1;
    return tt
end

function reshape_signal(s,frame_len)
    nscan = length(s) ÷ frame_len
    return reshape(s[1:frame_len*nscan],(frame_len,nscan))
end

"""
    datatime, (channelA,channelB), data = decode(y,Fs)

Decode the APT image in a time series `y` defined at a frequency `Fs` (in Hz).
`datatime` is the time in seconds counting from the beginning of the recording.
`channelA` and `channelB` are arrays representing the data from the two
different channels (A and B). `data` is raw data as an array including channel A and B,
as well sync markers and telemetry.

# Example
```julia
using FileIO
using APTDecoder

wavname = "gqrx_20190804_141523_137100000.wav"
y,Fs,nbits,opt = load(wavname)
datatime,(channelA,channelB),data = APTDecode.decode(y,Fs)

```
"""
function decode(y,Fs)
    # Fs2 should be a multiple of 4160 Hz and least 8320 Hz
    # 4160 is the least common multiple of 1040 and 832 (the frequency of the
    # sync A and B pulses)

    Fs2 = 3*4160

    # low and high frequency for the band-pass filter (in Hz)
    wpass = (400., 4400.)

    responsetype = DSP.Filters.Bandpass(wpass[1],wpass[2],fs = Fs);
    designmethod = DSP.Filters.Butterworth(6)

    #figure(4); clf(); psd(y[:,1],Fs = Fs, NFFT=1024)
    yf = DSP.filt(DSP.digitalfilter(responsetype, designmethod), y[:,1]);
    #figure(4); clf(); psd(yf[:,1],Fs = Fs, NFFT=1024)

    y2 = DSP.Filters.resample(yf, float(Fs2) / float(Fs) )

    # AM demodulation using the Hilbert Transform
    y_demod = am_demodulation(y2);

    # generate the syncronization frame at the frequency Fs2
    sync_frame = gen_sync_frame(Fs2,sync_frequency)

    # length of a frame
    frame_len = round(Int,Fs2/scans_per_seconds)

    sync_frame_index = find_sync(y_demod,sync_frame[1],frame_len)

    # complete frames
    complete_frame_index = Int[]
    for i = 1:length(sync_frame_index)
        if sync_frame_index[i]+frame_len-1 <= length(y_demod)
            push!(complete_frame_index,sync_frame_index[i])
        end
    end

    data = zeros(length(complete_frame_index),frame_len)
    datatime = (complete_frame_index .- 1) / Fs2

    for i = 1:length(complete_frame_index)
        data[i,:] = y_demod[complete_frame_index[i] : complete_frame_index[i]+frame_len-1]
    end

    # channel A and B
    channels = (view(data,:,259:2985), view(data,:,3380:6103))
    return datatime,channels,data
end


module PooksoftAssetModelingKit

# include -
include("Include.jl")

# export methods -
export price
export evaluate
export sample_underlying_gbm_model
export estimate_single_index_model
export estimate_random_walk_model
export compute_linear_return_array
export compute_log_return_array
export compute_return_volatility

# lattice -
export compute_ternary_price_tree
export compute_binary_price_tree

# technical indicators methods -
export compute_fibonacci_retracement_levels
export compute_exponential_moving_average

# factory methods -
export build_geometric_brownian_motion_model

# export types -
export PSGeometricBrownianMotionModelParameters
export PSHestonAssetPricingModelParameters
export PSSingleIndexModelParameters
export PSRandomWalkModelParameters

end # module


Δt = 0.01
h = 0.5

ρ = 1
μ = 0.001
ν = μ / ρ

# ϕϕ = 1/9
# ϕϕx = 0.0
# ϕxϕx = 4/3

# # LHS * u = RHS
# # LHS = [ϕϕ 0.0;0.0 ϕϕ]/Δt +

# sol = - 2 * μ * ϕxϕx / (ϕϕ/Δt + 2 * μ * ϕxϕx)


# SCE
using UnPack
using Plots
include("CCE.jl")

mutable struct Pop
    position::Vector{Float64}
    cost::Float64
end

# CCE parameters
mutable struct CCE_params
    q::Int64
    alpha::Int64
    beta::Int64
    lb::Vector{Float64}
    ub::Vector{Float64}
    obj_func
end

# SCE parameters
mutable struct SCE_params
    max_iter::Int64
    n_complex::Int64
    n_complex_pop::Int64
    dim::Int64
    lb::Vector{Float64}
    ub::Vector{Float64}
    obj_func
end

import Base.isless
isless(a::Pop, b::Pop) = isless(a.cost, b.cost)

function cost_func(x)
    return sum(x.^2)
end

function uniform_rand(lb::Array{Float64, 1}, ub::Array{Float64, 1})
    dim = length(lb)
    arr = rand(dim) .* (ub .- lb) .+ lb
    return arr
end
# SCE parameters
max_iter = 500
n_complex = 5
n_complex_pop = 10
dim = 10
lb = ones(dim) * -10
ub = ones(dim) * 10
obj_func = cost_func
n_complex_pop = max(n_complex_pop, dim+1) # 如果n_pop_complex小于nVar+1,将n_pop_complex设为nVar+1
# 这个设定来自Nelder-Mead Standard
sce_params = SCE_params(max_iter, n_complex, n_complex_pop, dim, lb, ub, obj_func)

# CCE parameters
cce_q = max(round(Int32, 0.5*n_complex_pop), 2)
cce_alpha = 3
cce_beta = 5

cce_params = CCE_params(cce_q, cce_alpha, cce_beta, lb, ub, obj_func)

function SCE(sce_params, cce_params)
    @unpack max_iter, n_complex, n_complex_pop, dim, lb, ub, obj_func = sce_params

    n_pop = n_complex * n_complex_pop
    I = reshape(1:n_pop, n_complex, :)

    # Step 1. Generate rand_sample
    best_costs = Vector{Float64}(undef, max_iter)

    pops = []
    for i in 1:n_pop
        pop_position = uniform_rand(lb, ub)
        pop_cost = obj_func(pop_position)
        pop = Pop(pop_position, pop_cost)
        push!(pops, pop)
    end
    complex = Array{Pop}(undef, n_complex_pop, n_complex)

    # Step 2. Rank Points
    sort!(pops)
    best_pop = pops[1]

    # Main loop
    for iter in 1:max_iter

        # Step 3. Partion into complexes
        for j in 1:n_complex
            complex[:,j] = deepcopy(pops[I[j,:]])
            # Step 4. Evolve complex, run CCE
            complex[:,j] = CCE(complex[:,j], cce_params)
            pops[I[j,:]] = deepcopy(complex[:,j])
        end
        # Step 5. Shuffle Complexes

        sort!(pops)
        best_pop = pops[1]

        best_costs[iter] = best_pop.cost

        # Show Iteration Information
        println("Iter = ", iter)
        println("The Best Cost is: ", best_costs[iter])
    end
    best_costs
end

best_costs = SCE(sce_params, cce_params)

plot(best_costs, yaxis=:log, label = "cost")

# savefig("Julia.png")


using Test, QSimulator
using LinearAlgebra: diag, diagm, I
const pnt = QSimulator.PERTURBATIVE_NUM_TERMS
const cnt = QSimulator.CHARGE_NUM_TERMS

@testset "DuffingTransmon" begin
	freq = 4.0
	anharm = -0.2
	dims = [3]
	q = DuffingTransmon("q", dims[1], DuffingSpec(freq, anharm))
	h = hamiltonian(q)
	@test isapprox(h, diagm(0 => [0.0, freq, 2 * freq + anharm]))
end

@testset "PerturbativeTransmon" begin
	fmax = 4.0
	anharm_max = -0.2
	fmin = 3.0
	t = fit_transmon(fmax, fmin, anharm_max, PerturbativeTransmon, pnt)
	q = PerturbativeTransmon("q", 3, t)
	h = hamiltonian(q, 0.0)
	@test isapprox(h, diagm(0 => [0.0, fmax, 2 * fmax + anharm_max]))

	ξ = QSimulator.xi_effective(2,160,0,0)
	U = perturbative_transmon_Ueigen(ξ)
	@test size(U) == (5,25)
	Neigen = U*transmon_N(25, ξ)*U'

	σy = 1im*(raising(2) - lowering(2))
	λ = perturbative_transmon_λ(2,160,0,0; num_terms=6)
	Λ = perturbative_transmon_Λ(2,160,0,0; num_terms=6)

	# confirm eq. (24) in arXiv:1706.06566v2
	@test isapprox(Neigen[1:2,1:2], λ*σy/(2*√ξ); rtol=5e-7)
	@test isapprox(Neigen[2:3,2:3], Λ*σy/√(2ξ); rtol=1e-5)
end

@testset "PerturbativeTransmon vs DiagonalChargeBasisTransmon" begin
	t = TransmonSpec(.172, 12.71, 3.69)
	pt = PerturbativeTransmon("q", 3, t, num_terms=pnt)
	ct = DiagonalChargeBasisTransmon("q", 3, t, num_terms=cnt)
	for ϕ in range(0, stop=1.0, length=10)
		ham_pt = hamiltonian(pt, ϕ)
		ham_ct = hamiltonian(ct, ϕ)
	    @test isapprox(ham_pt + ham_ct[1,1] * I, ham_ct, rtol=1e-5)
	end
end

@testset "fit_transmon fixed" begin
	freq = 5.0
	anharm = -0.2
	ham = QSimulator.duffing_hamiltonian(freq, anharm, 3)

	function test_spec(s, model)
		@test s.EJ2 == 0.0
		ham_1 = hamiltonian(model("Test", 3, s))
		@test isapprox(ham + ham_1[1,1] * I, ham_1, rtol=1e-7)
	end

	fit_spec = fit_transmon(freq, freq, anharm, PerturbativeTransmon, pnt)
	test_spec(fit_spec, PerturbativeTransmon)
	test_spec(fit_spec, DiagonalChargeBasisTransmon)
	fit_spec = fit_transmon(freq, freq, anharm, DiagonalChargeBasisTransmon, cnt)
	test_spec(fit_spec, PerturbativeTransmon)
	test_spec(fit_spec, DiagonalChargeBasisTransmon)
end

@testset "fit_transmon tunable" begin
	# tunable transmon
	freq_max = 5.0
	freq_min = 4.0
	anharm_max = -0.2
	ham_max = QSimulator.duffing_hamiltonian(freq_max, anharm_max, 3)
	ham_min = QSimulator.duffing_hamiltonian(freq_max, 0.0, 2)

	function test_spec(s, model)
		@test s.EJ2 != 0.0
		ham_max_1 = hamiltonian(PerturbativeTransmon("Test", 3, s))
		@test isapprox(ham_max + ham_max_1[1,1] * I, ham_max_1, rtol=1e-7)

		ham_min_1 = hamiltonian(PerturbativeTransmon("Test", 2, s))
		@test isapprox(ham_min + ham_min_1[1,1] * I, ham_min_1, rtol=1e-8)
	end

	fit_spec = fit_transmon(freq_max, freq_min, anharm_max, PerturbativeTransmon, pnt)
	test_spec(fit_spec, PerturbativeTransmon)
	test_spec(fit_spec, DiagonalChargeBasisTransmon)
	fit_spec = fit_transmon(freq_max, freq_min, anharm_max, DiagonalChargeBasisTransmon, cnt)
	test_spec(fit_spec, PerturbativeTransmon)
	test_spec(fit_spec, DiagonalChargeBasisTransmon)
end

@testset "DuffingSpec <-> TransmonSpec" begin
	freq = 4.0
	anharm = -0.2
	d = DuffingSpec(freq, anharm)
	t = TransmonSpec(d)
	d1 = DuffingSpec(t)
	@test isapprox(d1.frequency, freq, rtol=1e-8)
	@test isapprox(d1.anharmonicity, anharm, rtol=1e-6)
end


using HPAT
using HPAT.API.NaiveBayes
using MPI

@acc hpat function calcNaiveBayes(file_name)
    num_classes = 20
    points = DataSource(Matrix{Float64},HDF5,"/points", file_name)
    labels = DataSource(Vector{Float64},HDF5,"/labels", file_name)
    coeffs = NaiveBayes(points, labels, num_classes)
    return coeffs
end

c2 = calcNaiveBayes(HPAT.getDefaultDataPath()*"/naivebayes_train_large.hdf5")

if MPI.Comm_rank(MPI.COMM_WORLD)==0	println("Coeffs: ", c2) end


export little_group_elements
export little_group
export little_symmetry
export little_symmetry_iso

"""
    little_group_elements(tsym, psym)

Return elements of `psym` compatible with the translation symmetry `tsym`.
"""
function little_group_elements(tsym::FiniteTranslationSymmetry, psym::PointSymmetry)::Vector{Int}
    return Int[
        i_elem
            for (i_elem, elem) in enumerate(elements(psym))
                if iscompatible(tsym, elem)
    ]
end

"""
    little_group_elements(tsym, tsym_irrep_index, psym)

Return elements of `psym` compatible with the translation symmetry `tsym` at
irrep `tsym_irrep_index`
"""
function little_group_elements(
    tsym::FiniteTranslationSymmetry,
    tsym_irrep_index::Integer,
    psym::PointSymmetry,
)::Vector{Int}
    return Int[
        i_elem
            for (i_elem, pop) in enumerate(elements(psym))
                if iscompatible(tsym, tsym_irrep_index, pop)
    ]
end

"""
    little_group(tsym, tsym_irrep_index, psym)
"""
function little_group(
    tsym::FiniteTranslationSymmetry,
    tsym_irrep_index::Integer,
    psym::PointSymmetry,
)::FiniteGroup
    lg_elems = little_group_elements(tsym, tsym_irrep_index, psym)
    return little_group(tsym, psym, lg_elems)
end

"""
    little_group(tsym, psym, elements)

Generate a little group with given elements.
The elements of the little group, which may be sparse, are compressed into consecutive integers.
"""
function little_group(
    tsym::FiniteTranslationSymmetry,
    psym::PointSymmetry,
    elements::AbstractVector{<:Integer},
)
    element_lookup = Dict(x=>i for (i, x) in enumerate(elements))
    ord_group = length(elements)
    mtab = zeros(Int, (ord_group, ord_group))
    for i in 1:ord_group, j in 1:ord_group
        mtab[i,j] = element_lookup[group_product(psym.group, elements[i], elements[j])]
    end
    return FiniteGroup(mtab)
end

"""
    little_symmetry(tsym, psym, lg_elements)

Generate a little symmetry.
First generate a little group with `lg_elements`, look up `PointSymmetryDatabase` using
the names of the elements, and then checks for group isomorphism.

# Arguments
* `tsym::FiniteTranslationSymmetry`: translation symmetry
* `psym::PointSymmetry`: point symmetry
* `lg_elements::AbstractVector{<:Integer}`: indices of the elements of the little group
"""
function little_symmetry(
    tsym::FiniteTranslationSymmetry,
    psym::PointSymmetry,
    lg_elements::AbstractVector{<:Integer},
)
    (lg_raw, lg_matrep_raw, lg_element_names_raw) = let
        lg_raw = little_group(tsym, psym, lg_elements)
        lg_matrep_raw = psym.matrix_representations[lg_elements]
        lg_element_names_raw = psym.element_names[lg_elements]
        (lg_raw, lg_matrep_raw, lg_element_names_raw)
    end

    psym2 = PointSymmetryDatabase.find(lg_element_names_raw)
    ϕ = group_isomorphism(lg_raw, psym2.group)
    isnothing(ϕ) && error("Group not isomorphic")

    lg_matrep = lg_matrep_raw[ϕ]
    lg_element_names = lg_element_names_raw[ϕ]

    return PointSymmetry(
        psym2.group,
        psym2.generators,
        psym2.conjugacy_classes,
        psym2.character_table,
        psym2.irreps,
        lg_element_names,
        lg_matrep,
        psym2.hermann_mauguin,
        psym2.schoenflies,
    )
end

"""
    little_symmetry(tsym, psym)
"""
function little_symmetry(tsym::FiniteTranslationSymmetry, psym::PointSymmetry)
    return little_symmetry(tsym, psym, little_group_elements(tsym, psym))
end

"""
    little_symmetry(tsym, tsym_irrep_index, psym)
"""
function little_symmetry(tsym::FiniteTranslationSymmetry, tsym_irrep::Integer, psym::PointSymmetry)
    if !iscompatible(tsym, psym)
        throw(ArgumentError("translation and point symmetries are not compatible"))
    end

    tsym_irrep == 1 && return little_symmetry(tsym, psym)

    (lg_raw, lg_matrep_raw, lg_element_names_raw) = let
        lg_elements = little_group_elements(tsym, tsym_irrep, psym)
        lg_raw = little_group(tsym, psym, lg_elements)
        lg_matrep_raw = psym.matrix_representations[lg_elements]
        lg_element_names_raw = psym.element_names[lg_elements]
        (lg_raw, lg_matrep_raw, lg_element_names_raw)
    end

    psym2 = PointSymmetryDatabase.find(lg_element_names_raw)
    ϕ = group_isomorphism(lg_raw, psym2.group)
    isnothing(ϕ) && error("Group not isomorphic")

    lg_matrep = lg_matrep_raw[ϕ]
    lg_element_names = lg_element_names_raw[ϕ]

    return PointSymmetry(
        psym2.group,
        psym2.generators,
        psym2.conjugacy_classes,
        psym2.character_table,
        psym2.irreps,
        lg_element_names,
        lg_matrep,
        psym2.hermann_mauguin,
        psym2.schoenflies,
    )
end

# COV_EXCL_START

"""
    little_symmetry_iso(tsym, tsym_irrep_index, psym)

Find little symmetry using group isomorphism
"""
function little_symmetry_iso(
    tsym::FiniteTranslationSymmetry,
    tsym_irrep_index::Integer,
    psym::PointSymmetry,
)
    @warn "Function little_symmetry_iso deprecated. Correctness not guaranteed." maxlog=1
    if !iscompatible(tsym, psym)
        throw(ArgumentError("translation and point symmetries are not compatible"))
    end
    tsym_irrep_index == 1 && return little_symmetry(tsym, psym)
    (lg_matrep, lg_element_names, lg_irrep) = let
        lg_elements = little_group_elements(tsym, tsym_irrep_index, psym)

        lg_raw = little_group(tsym, psym, lg_elements)
        lg_matrep_raw = psym.matrix_representations[lg_elements]
        lg_element_names_raw = psym.element_names[lg_elements]

        (lg_irrep, ϕ) = IrrepDatabase.find(lg_raw)

        lg_matrep = lg_matrep_raw[ϕ]
        lg_element_names = lg_element_names_raw[ϕ]
        (lg_matrep, lg_element_names, lg_irrep)
    end

    generators = minimal_generating_set(lg_irrep.group)
    hermann_mauguin = join(lg_element_names[generators])
    schoenflies = "unknown"

    simple_element_names = sort(simplify_name.(lg_element_names))
    for i in 1:32
        psym = PointSymmetryDatabase.get(i)
        if ( sort(simplify_name.(psym.element_names)) == simple_element_names
             && !isnothing(group_isomorphism(lg_irrep.group, psym.group)) )
            hermann_mauguin = psym.hermann_mauguin
            break
        end
    end

    return PointSymmetry(
        lg_irrep.group,
        generators,
        lg_irrep.conjugacy_classes,
        lg_irrep.character_table,
        lg_irrep.irreps,
        lg_element_names,
        lg_matrep,
        hermann_mauguin,
        schoenflies,
    )
end
# COV_EXCL_STOP


const Mean = let
    init() = (0, 0.0)
    op((n, sum), val) = n + 1, sum + val
    value((n, sum)) = sum / n
    (; init, op, value)
end

struct ExpectationAnalysis
    options::Dict{Symbol, Any}
end

(e::ExpectationAnalysis)(entry::Entry) = groupreduce(Mean, entry)

"""
    expectation(args...)

Compute the expected value of the last argument conditioned on the preceding ones.
"""
expectation() = Layer((ExpectationAnalysis(Dict{Symbol, Any}()),))


#--------------------------------------------------------------------
# DNSS.jl
# Soham 08-2019
# Test scalar field on Minkowski spacetime
#--------------------------------------------------------------------

using Test, NLsolve, LinearAlgebra, Random, Printf, ForwardDiff
# using LaTeXStrings, Printf, PyPlot

#-----------------------------------------
# Set up initial conditions on u0 and v0
#-----------------------------------------
function psibar(u::Real, v::Real)::Real
    return  exp(-u^2 / σ^2) * sin(pi*u) + exp(-v^2 / σ^2) * sin(pi*v)
end

function computeUboundary(PS::ProductSpace{S1, S2})::NTuple{1, Field{S2}} where {S1, S2}
    return (extractUboundary(Field(PS, psibar), :incoming), )
end

function computeVboundary(PS::ProductSpace{S1, S2})::NTuple{1, Field{S1}} where {S1, S2}
    return (extractVboundary(Field(PS, psibar), :incoming), )
end

#------------------------------------------------------------
# Source term or the change of coordinate functions
#------------------------------------------------------------
function twist(uv::Vector{T})::Vector{T} where {T<:Real}
    u, v = uv
    ω =  (pi/8) * cospi(u/2) * cospi(v/2)
    ũ =  u*cos(ω) + v*sin(ω)
    ṽ = -u*sin(ω) + v*cos(ω)
    return [ũ, ṽ]
end

function dxdX(u::Real, v::Real, component::Symbol)
    if component == :dUdŨ
        return inv(ForwardDiff.jacobian(twist, [u,v]))[1,1]
    elseif component == :dUdṼ
        return inv(ForwardDiff.jacobian(twist, [u,v]))[1,2]
    elseif component == :dVdŨ
        return inv(ForwardDiff.jacobian(twist, [u,v]))[2,1]
    elseif component == :dVdṼ
        return inv(ForwardDiff.jacobian(twist, [u,v]))[2,2]
    end
end

function source(u::Real, v::Real)::Real
    t = u + v
    r = v - u
    return 20 * sinpi(t) * sinpi(r) *  exp(-r^2 / 0.1)
end

#------------------------------------------------------------
# Set up computation on each patch using a non-linear solver
#------------------------------------------------------------
function computePatch(PS::ProductSpace{S1, S2}, ubnd::NTuple{1,Field{S2}}, vbnd::NTuple{1,Field{S1}})::NTuple{1, Field{ProductSpace{S1, S2}}} where {S1, S2}
    if solver == :nonlin
        # Define all the operators we use local to the patch.
        B = incomingboundary(PS)
        I = identity(PS)
        DU, DV = derivative(PS)
        psibnd = combineUVboundary(ubnd, vbnd, :incoming)

        dUdŨ   = Field(PS, (u,v)-> dxdX(u,v, :dUdŨ))
        dVdŨ   = Field(PS, (u,v)-> dxdX(u,v, :dVdŨ))
        dUdṼ   = Field(PS, (u,v)-> dxdX(u,v, :dUdṼ))
        dVdṼ   = Field(PS, (u,v)-> dxdX(u,v, :dVdṼ))

        # Compute the residual. This is slow but works. Sets the residual in the
        # interior as well as the boundary.
        function residual(psi::NTuple{1, Field{S}})::NTuple{1, Field{S}} where {S}
            psi   = first(psi)
            # FIXME: The twisted coordinates do not work
            dpsidŨ    = dUdŨ * (DU * psi) + dVdŨ * (DV * psi)
            ddpsidŨdṼ = dUdṼ * (DU * dpsidŨ) + dVdṼ * (DV * dpsidŨ)
            F = (I - B) * ddpsidŨdṼ  + B * (psi - first(psibnd))
            return (F,)
        end

        # Residual to pass to the non-linear solver in NLsolve
        function f!(fvec::Array{T,1}, x::Array{T,1}) where {T}
            fvec[:]  = reshapeFromTuple(residual(reshapeToTuple(PS, 1, x)))
        end

        # Specify an initial guess
        function initialguess(PS::ProductSpace{S1, S2})::NTuple{1, Field{ProductSpace{S1, S2}}} where {S1, S2}
            return (Field(PS, (u,v)->1), )
        end

        # Compute the initial guess to pass to the solver
        U0 = reshapeFromTuple(initialguess(PS))

        # Call the non-linear solver
        U = nlsolve(f!, U0; method=:trust_region, autodiff=:forward, show_trace=true, ftol=1e-8, iterations=100)
        return reshapeToTuple(PS, 1, U.zero)

    elseif solver == :lin
        # Define all the operators we use local to the patch.
        B = incomingboundary(PS)
        DU, DV = derivative(PS)
        bnd = combineUVboundary(ubnd, vbnd, :incoming)
        L   = 2*DU*DV
        s   = 0 * Field(PS, source)
        b   = first(bnd)

        # Using the linearity of the system to solve the system
        # L u  = s # Wave equation with a source
        # B u =  b # Boundary conditions
        # (L + B) u = b
        U = linsolve(L + B,  s + b)
        return (U, )
    end
end

#-----------------------------------------
# Compare with analytic solution
#-----------------------------------------
function deltapsi(U::NTuple{1, Field})::NTuple{1, Field} where {S}
    # Compute the error at twice the number of points
    psi = project(first(U), prolongate(first(U).space))
    psiexact = Field(psi.space, psibar)
    return (psi - psiexact, )
end

#-----------------------------------------
# h and p convergence functions
#-----------------------------------------
function pconv(range::StepRange, nh::Int=1)
    println("Testing p convergence")
    n_ = collect(range)
    l_ = zeros(size(n_))
    for index in CartesianIndices(n_)
        np = n_[index]
        params = Parameters((np, np), (nh,nh), urange, vrange, 1)
        q = distribute(params, computePatch, computeUboundary, computeVboundary)
        l_[index] = rmse(extract(map(deltapsi, q), 1))
        @printf("  n = %i, rmse = %e\n", n_[index], l_[index])
    end
    return (n_, l_)
end

function hconv(range::StepRange, np::Int=4)
    println("Testing h convergence")
    n_ = collect(range)
    l_ = zeros(size(n_))
    for index in CartesianIndices(n_)
        nh = n_[index] = 2^n_[index]
        params = Parameters((np, np), (nh, nh), urange, vrange, 1)
        l_[index] = rmse(extract(map(deltapsi, distribute(params, computePatch, computeUboundary, computeVboundary)), 1))
        l0 = index[1] > 1 ? l_[index[1] - 1] : 1
        @printf("  n = %3i, rmse[2h] / rmse[h] = %e\n", n_[index], l0 / l_[index])
    end
    return (n_, l_)
end

#-----------------------------------------
# Setup simulation grid parameters
#-----------------------------------------
# TODO: Check if we find convergence for different sizes of domains. What
# happens if the domain size increases? Is the solution unstable for larger
# domain sizes? What we find is that one needs more spacetime elements for
# larger domain sizes to see p-convergence. For a single or fewer elements,
# the p-convergence is either non-existent, or slow.
# FIXME: In any case, we saturate around 1e-10 with all the errors. Is this due
# to the floor set by the condition number of the second derivative matrices?
# What happens if we choose to work with BigFloat?
npoints  = (20, 20)
npatches = (1, 1)
urange   = (-1.0, 1.0)
vrange   = (-1.0, 1.0)
nfields  = 1
Grid     = Parameters(npoints, npatches, urange, vrange, nfields)
solver   = :lin # Use a linear solver instead of the non-linear solver
σ        = 0.5

#-----------------------------------------
# Compute the scalar field over the grid
#-----------------------------------------
# ϕ_ = distribute(Grid, computePatch, computeUboundary, computeVboundary)
# save("output/data/minkowski_2D/sigma_0p5_np_30_nh_4_source/minkowski_2D_sigma_0p5_psi", extract(ϕ_, 1))
# save("output/data/minkowski_2D/sigma_0p5_np_30_nh_4_source/minkowski_2D_sigma_0p5_deltapsi", extract(map(deltapsi,ϕ_), 1))

#-----------------------------------------
# Now test h-p convergence of the error
# in the solution
#-----------------------------------------
np, ep = pconv(2:2:40, 12)
nh, eh = hconv(0:1:9, 4)

#-----------------------------------------
# Save all the data for solution and convergence
#-----------------------------------------
# save("output/data/minkowski_2D_sigma_0p9_psi", extract(ϕ_, 1))
# save("output/data/minkowski_2D_sigma_0p9_deltapsi_pconv", np, ep)
# save("output/data/minkowski_2D_sigma_0p9_deltapsi_hconv", nh, eh)


using Optim, HTTP, GLM, LinearAlgebra, Random, Statistics, DataFrames, CSV, FreqTables,Distributions
include("lgwt.jl")
#data
function PS4()
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS4-mixture/nlsw88t.csv"
df = CSV.read(HTTP.get(url).body)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occ_code;

#Question 1
#from his solution:
function mlogit_with_Z(theta, X, Z, y)

        alpha = theta[1:end-1]
        gamma = theta[end]
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        T = promote_type(eltype(X),eltype(theta))
        num   = zeros(T,N,J)
        dem   = zeros(T,N)
        for j=1:J
            num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)
            dem .+= num[:,j]
        end

        P = num./repeat(dem,1,J)

        loglike = -sum( bigY.*log.(P) )

        return loglike
    end

#values from PS3
startvals=[0.05570767876416688, 0.08342649976722213, -2.344887681361976, 0.04500076157943125, 0.7365771540890512, -3.153244238810631, 0.09264606406280998, -0.08417701777996893, -4.273280002738097, 0.023903455659102114, 0.7230648923377259, -3.749393470343111, 0.03608733246865346, -0.6437658344513095, -4.2796847340030375, 0.0853109465190059, -1.1714299392376775, -6.678677013966667, 0.086620198654063, -0.7978777029320784, -4.969132023685069, -0.0941942241795243];
#look at this to see why you might need this: http://julianlsolvers.github.io/Optim.jl/v0.9.3/algo/autodiff/
td = TwiceDifferentiable(theta -> mlogit_with_Z(theta, X, Z, y), startvals; autodiff = :forward);

theta_hat_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))

#Hessian from ps3
theta_hat_mle_ad = theta_hat_optim_ad.minimizer
println(theta_hat_mle_ad)
H  = Optim.hessian!(td, theta_hat_mle_ad)
theta_hat_mle_ad_se = sqrt.(diag(inv(H)))
println([theta_hat_mle_ad theta_hat_mle_ad_se])

#question 2
# now it is positive which makes more sense because increase in log wage should increase utility

#question 3
#part a

#part b
#1
nodes, weights = lgwt(7,-10,10);
d=Normal(0,2)
sum(weights.*pdf.(d,nodes))
println(sum((nodes.^2).*weights.*pdf.(d,nodes)))

#2
nodes, weights = lgwt(10,-10,10);
sum(weights.*pdf.(d,nodes))
println(sum((nodes.^2).*weights.*pdf.(d,nodes)))

#3
# it is close, with more quadrature points it will converge to 4

#Part C
#1
d=Normal(0,2)
N=1000000
x=rand(Uniform(-10, 10 ),N)

println((20/N)*sum(x.^2 .*(pdf.(d,x))))

#2
println((20/N)*sum(x .*(pdf.(d,x))))

#3
println((20/N)*sum(pdf.(d,x)))

#4
#the difference is around 0.2 which is high in my opinion

#Question 4
function mxlogit_with_quadrature(theta, X, Z, y,rn)

        alpha = theta[1:end-1]
        gamma = theta[end]
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        nodes, weights = lgwt(rn,-4,4);
        d=Normal(0,1)

        T = promote_type(eltype(X),eltype(theta))
        num   = zeros(T,N,J)
        dem   = zeros(T,N)
        for j=1:J
            num[:,j] = sum(exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma).*weights.*pdf.(d,nodes))
            dem .+= num[:,j]
        end

        P = num./repeat(dem,1,J)

        loglike = -sum( log.(P.^bigY) )

        return loglike
    end

#Question 5
function mxlogit_with_simulation(theta, X, Z, y,rn)

        alpha = theta[1:end-1]
        gamma = theta[end]
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        d=Normal(0,1)
        rx=rand(Uniform(0,1 ),rn)

        T = promote_type(eltype(X),eltype(theta))
        num   = zeros(T,N,J)
        dem   = zeros(T,N)
        for j=1:J
            num[:,j] = sum(exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma).*(pdf.(d,x)))
            dem .+= num[:,j]
        end

        P = num./repeat(dem,1,J)

        loglike = -sum( log.(P.^bigY) )

        return loglike
    end

end

PS4()


"""
    divergence(f)
"""
divergence(f) = tr(gradient(f))

function divergence(a::AbstractArray{<:Field})
  ag = gradient(a)
  operate_arrays_of_fields(tr,ag)
end

function symmetric_gradient end

"""
    symmetric_gradient(f)
"""
symmetric_gradient(f) = symmetric_part(gradient(f))

function symmetric_gradient(a::AbstractArray{<:Field})
  ag = gradient(a)
  operate_arrays_of_fields(symmetric_part,ag)
end

"""
    const ε = symmetric_gradient

Alias for the symmetric gradient
"""
const ε = symmetric_gradient

"""
    curl(f)
"""
curl(f) = grad2curl(gradient(f))

function curl(a::AbstractArray{<:Field})
  ag = gradient(a)
  operate_arrays_of_fields(grad2curl,ag)
end

"""
    grad2curl(∇f)
"""
function grad2curl(f)
  @abstractmethod
end

grad2curl(a::GridapType) = operate(grad2curl,a)

grad2curl(f::Field) = operate_fields(grad2curl,f)

@inline function grad2curl(∇u::TensorValue{2})
  ∇u[1,2] - ∇u[2,1]
end

@inline function grad2curl(∇u::TensorValue{3})
  c1 = ∇u[2,3] - ∇u[3,2]
  c2 = ∇u[3,1] - ∇u[1,3]
  c3 = ∇u[1,2] - ∇u[2,1]
  VectorValue(c1,c2,c3)
end

function laplacian end

"""
    const Δ = laplacian

Alias for the `laplacian` function
"""
const Δ = laplacian

"""
    laplacian(f)
"""
function laplacian(f)
  g = gradient(f)
  divergence(g)
end

"""
    ∇⋅f

Equivalent to

    divergence(f)
"""
dot(::typeof(∇),f) = divergence(f)
dot(::typeof(∇),f::GridapType) = divergence(f)

function (*)(::typeof(∇),f)
  msg = "Syntax ∇*f has been removed, use ∇⋅f (\\nabla \\cdot f) instead"
  error(msg)
end

function (*)(::typeof(∇),f::GridapType)
  msg = "Syntax ∇*f has been removed, use ∇⋅f (\\nabla \\cdot f) instead"
  error(msg)
end

"""
    outer(∇,f)

Equivalent to

    gradient(f)
"""
outer(::typeof(∇),f) = gradient(f)
outer(::typeof(∇),f::GridapType) = gradient(f)

"""
    outer(f,∇)

Equivalent to

    transpose(gradient(f))
"""
outer(f,::typeof(∇)) = transpose(gradient(f))
outer(f::GridapType,::typeof(∇)) = transpose(gradient(f))

"""
    cross(∇,f)

Equivalent to

    curl(f)
"""
cross(::typeof(∇),f) = curl(f)
cross(::typeof(∇),f::GridapType) = curl(f)

# Automatic differentiation of functions

function gradient(f::Function)
  function grad_f(x)
    _grad_f(f,x,zero(return_type(f,typeof(x))))
  end
end

function _grad_f(f,x,fx)
  VectorValue(ForwardDiff.gradient(f,get_array(x)))
end

function _grad_f(f,x,fx::VectorValue)
  TensorValue(transpose(ForwardDiff.jacobian(y->get_array(f(y)),get_array(x))))
end

function _grad_f(f,x,fx::MultiValue)
  @notimplemented
end

function divergence(f::Function)
  x -> tr(ForwardDiff.jacobian(y->get_array(f(y)),get_array(x)))
end

function curl(f::Function)
  x -> grad2curl(TensorValue(transpose(ForwardDiff.jacobian(y->get_array(f(y)),get_array(x)))))
end

function laplacian(f::Function)
  function lapl_f(x)
    _lapl_f(f,x,zero(return_type(f,typeof(x))))
  end
end

function _lapl_f(f,x,fx)
  tr(ForwardDiff.jacobian(y->ForwardDiff.gradient(f,y), get_array(x)))
end

function _lapl_f(f,x,fx::VectorValue)
  A = length(x)
  B = length(fx)
  a = ForwardDiff.jacobian(y->transpose(ForwardDiff.jacobian(z->get_array(f(z)),y)), get_array(x))
  tr(ThirdOrderTensorValue{A,A,B}(Tuple(transpose(a))))
end

function _lapl_f(f,x,fx::MultiValue)
  @notimplemented
end

function symmetric_gradient(f::Function)
    x -> symmetric_part(_grad_f(f,x,zero(return_type(f,typeof(x)))))
end


function temperatureasintersection(a::Array{Float64,1}, b::Array{Float64,1}, σ2::Array{Float64,1}, T::Float64, s::Float64)

    x = s*unitgalaxyvector(T)

    p = dot(a, (x-b)) / dot(a,a)

    logpdf(MvNormal(a*p + b, Diagonal(σ2)), x)

end

function temperatureasintersection(a::Array{Float64,1}, b::Array{Float64,1}, σ2::Array{Float64,1}, Trange::AbstractRange{Float64}, Srange::AbstractRange{Float64})

    f(T, s) = temperatureasintersection(a,b,σ2,T,s)

    @showprogress map(P -> f(P[1], P[2]), Iterators.product(Trange, Srange))

end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 4b51367c-09ec-11eb-16d9-3dcd2910078e
begin
	using PyPlot
	#using PyCall #I'll gonna use SpecialFunctions.jl and HypergeometricFunctions.jl
	using LaTeXStrings
	using SpecialFunctions
	using HypergeometricFunctions
end

# ╔═╡ eb325bf2-09e3-11eb-3b88-9b0d15626515
md"## Fractional Fourier-Kravchuk evolution on waveguides"

# ╔═╡ 74a12d8e-0f69-11eb-0650-e79c06fd1080
md"Kravchuk matrix is defined as

$K(i,j)=\sqrt{\frac{\omega(j)}{h(i)}}k(i,j),$
with the functions $\omega$ and $h$ given by

$\omega(j)=\binom{N}{j}p^{j}(1-p)^{N-j},\quad h(i)=\frac{\frac{1-p}{p^{i}}}{\binom{N}{i}},$
and the Kravchuk polynomials obtained by the Hypergeometrical function

$k(i,j)=~_{2}F_{1}(-j,-i,-N,1/p).$

Note: When $p=1/2$, the symmetric Kravchuk polynomials are obtained, henceforth, the symmetric Kravchuk functions follows.
"

# ╔═╡ 7f8aad74-09ec-11eb-3c25-a538f18da2f9
begin
	rc("text", usetex=true)
	ion()
end

# ╔═╡ c28078de-09ec-11eb-0c25-a1bd3c7cdd8b
md" Define the parameters"

# ╔═╡ cb9e6110-09ec-11eb-29bb-b33f49eb026a
begin
	N = 32
	dim = N+1
	const p = 1/2.0
end;

# ╔═╡ bb1e9a12-0fe6-11eb-0551-137403aecef6
function bspecial(n,k)
	return 1/((n+1)*beta(k+1,n-k+1))
end;

# ╔═╡ ee7b88a2-0fe6-11eb-3465-8b70e58f07c1
function kravs(i,j)
	return HypergeometricFunctions._₂F₁maclaurin(-j,-i,-N,1/p)
end;

# ╔═╡ 4f876a76-0fe7-11eb-3f46-0f18d733b446
function omegas(j)
	return bspecial(N,j)*(p^j)*(1-p)^(N-j)
end;

# ╔═╡ 590a100a-0fe7-11eb-017c-4d4ff299ed8a
function hs(i)
	return (((1-p)/p)^i)/(bspecial(N,i))
end;

# ╔═╡ a5f2bf3c-0fe7-11eb-3377-8d5228bb6476
function Ks(i,j)
	return sqrt(omegas(j)/hs(i))*kravs(i,j)
end;

# ╔═╡ d01de678-0fe8-11eb-2831-89f652e45a9a
@timed begin
	KTs = zeros(Float64,(dim,dim))
	for l in 0:N
    	for m in 0:N
        	KTs[l+1,m+1] = Ks(l,m)
    	end
	end
end

# ╔═╡ 4692817e-0fe9-11eb-22e0-dfdab6509141
begin
	fig1 = figure(figsize=(5,5))
	ax1 = gca()
	ax1.imshow(KTs, cmap="gray")
	ax1.set_title(L"\textbf{Kravchuk Matrix}\; K(i,j)\;\textbf{(The trilobite)}")
	ax1.set_xlabel(L"\textbf{Pseudo Energy Level:}\; 0\leq i\leq N")
	ax1.set_ylabel(L"\textbf{Position Index:}\; 0\leq j\leq N")
	ax1.tick_params(direction="out",length=5,width=2,labelsize=10)
	ax1.set_xticks(0:8:N, minor = false)
	ax1.set_yticks(0:8:N, minor = false)

	axins1a = ax1.inset_axes([0, 0.8, 0.2, 0.2])
	axins1a.plot(KTs[1,:],"k:",linewidth=1)
	axins1a.set_xticklabels("")
	axins1a.set_yticklabels("")
	axins1a.text(7.5,0.01,L"K(0,j)",fontsize=9)

	axins2a = ax1.inset_axes([0.8, 0, 0.2, 0.2])
	axins2a.plot(KTs[N+1,:],"k:",linewidth=1)
	axins2a.set_xticklabels("")
	axins2a.set_yticklabels("")
	axins2a.text(15,-0.35,L"K(N,j)",fontsize=9)

	tight_layout()
	#savefig("trilo.png",dpi=300,transparent=true)
end

# ╔═╡ 33089668-09ed-11eb-30db-8f33d969a115
md"
October 16

Status: Kravchuk matrix completed. Changed definitions of functions using native Julia libraries

Next: Write the fractional evolution, and plotting the results
"

# ╔═╡ Cell order:
# ╟─eb325bf2-09e3-11eb-3b88-9b0d15626515
# ╟─74a12d8e-0f69-11eb-0650-e79c06fd1080
# ╠═4b51367c-09ec-11eb-16d9-3dcd2910078e
# ╠═7f8aad74-09ec-11eb-3c25-a538f18da2f9
# ╟─c28078de-09ec-11eb-0c25-a1bd3c7cdd8b
# ╠═cb9e6110-09ec-11eb-29bb-b33f49eb026a
# ╠═bb1e9a12-0fe6-11eb-0551-137403aecef6
# ╠═ee7b88a2-0fe6-11eb-3465-8b70e58f07c1
# ╠═4f876a76-0fe7-11eb-3f46-0f18d733b446
# ╠═590a100a-0fe7-11eb-017c-4d4ff299ed8a
# ╠═a5f2bf3c-0fe7-11eb-3377-8d5228bb6476
# ╠═d01de678-0fe8-11eb-2831-89f652e45a9a
# ╠═4692817e-0fe9-11eb-22e0-dfdab6509141
# ╟─33089668-09ed-11eb-30db-8f33d969a115


function real_eigvectors(A)
    n = size(A, 1)

    vals, vecs = eigen(A)
    ix = imag.(vals) .>= 0

    real.(vals[ix]), real.(vecs[:, ix]), isreal.(vals[ix])
end

function rosenbrok(A, B, C)
    if length(B) == 0
        [
            A
            C
        ]
    elseif length(C) == 0
        [
            A B
        ]
    else
        [
            A B
            C zeros(size(C, 1), size(B, 2))
        ]
    end
end

function fixed_modes(A, B, C)
    ra = eigvals(A)
    rn = length(ra)
    n = size(A, 1)
    N = length(B)

    minsvd = Inf .* ones(rn)

    for i in 1:rn
        for j in 1:2^N
            mask = digits(Bool, j - 1, base=2, pad=N)
            BB = reduce(hcat, B[mask], init=zeros(n, 0))
            CC = reduce(vcat, C[.! mask], init=zeros(0, n))
            s = svdvals(DistributedControl.rosenbrok(A - ra[i] * I, BB, CC))
            if minsvd[i] > s[n]
                minsvd[i] = s[n]
            end
        end
    end
    minsvd
end


@testset "Examples" begin
    @testset "Naive sin" begin
        function sin_naive!(res, x)
            s, t, u = zero(x), zero(x), zero(x)
            tol = one(x)
            Arblib.mul_2exp!(tol, tol, -precision(tol))
            # @show tol
            k = 0
            while true
                Arblib.pow!(t, x, UInt(2k + 1))
                Arblib.fac!(u, UInt(2k + 1))
                Arblib.div!(t, t, u)
                Arblib.abs!(u, t)

                if u ≤ tol
                    # @show u
                    Arblib.add_error!(s, u)
                    break
                end
                if iseven(k)
                    Arblib.add!(s, s, t)
                else
                    Arblib.sub!(s, s, t)
                end
                k += 1
            end
            Arblib.set!(res, s)
            return res
        end

        @test let
            str = ""
            prec = 64
            while true
                x = Arb("2016.1"; prec = prec)
                y = zero(x)
                y = sin_naive!(y, x)
                str *= sprint(print, "Using $(lpad(prec, 5)) bits, sin(x) = ")
                str *= sprint(println, Arblib.string_nice(y, 10))
                y < zero(y) && break
                prec *= 2
            end
            str
        end == """
        Using    64 bits, sin(x) = [+/- 2.67e+859]
        Using   128 bits, sin(x) = [+/- 1.30e+840]
        Using   256 bits, sin(x) = [+/- 3.60e+801]
        Using   512 bits, sin(x) = [+/- 3.01e+724]
        Using  1024 bits, sin(x) = [+/- 2.18e+570]
        Using  2048 bits, sin(x) = [+/- 1.22e+262]
        Using  4096 bits, sin(x) = [-0.7190842207 +/- 1.20e-11]
        """
    end

    @testset "Logistic" begin
        """
            logistic(n[; x0 = 0.5, r = 3.75, digits = 10])
        Compute the nth iterate of the logistic map `x_{n + 1} = r x_n (1 -
        x_n).
        """
        function logistic(n; x0 = Arb(0.5), r = Arb(3.75), digits = 10)
            goalprec = digits * 3.3219280948873623 + 3

            prec = 64
            if typeof(x0) == String
                x = Arb(x0, prec = prec)
            else
                x = Arb(prec = prec)
            end

            t = zero(x)
            while true
                if typeof(x0) == String
                    x = Arb(x0, prec = prec)
                else
                    Arblib.set!(x, x0)
                end

                for i = 1:n
                    Arblib.sub!(t, x, 1)
                    Arblib.neg!(t, t)
                    Arblib.mul!(x, x, t)
                    Arblib.mul!(x, x, r)

                    if Arblib.rel_accuracy_bits(x) < goalprec
                        break
                    end
                end

                if Arblib.rel_accuracy_bits(x) >= goalprec
                    break
                end

                prec *= 2
                x = setprecision(x, prec)
                r = setprecision(r, prec)
                t = setprecision(t, prec)
            end

            return x
        end

        @test string(logistic(10, digits = 10)) == "[0.645367290830930 +/- 3.45e-16]"
        @test string(logistic(10, digits = 20)) ==
              "[0.64536729083093027156146131423635101 +/- 5.18e-36]"
    end
end


using DifferentialEquations
using Roots
using PyCall
pygui(:qt)
using PyPlot
pygui(true)

function condition(u,t,integrator) # Event when event_f(u,t) == 0
 u[1]
end

function affect!(integrator)
end
#cb1 = SavingCallback((u,t,integrator)->(u[4],u[2]), saved_values)

cb2= ContinuousCallback(condition, affect!)
Energy=0.112012;
H=(2*Energy)^2
h=sqrt(H)
n_iter=100
t_end=1000

# H1(p,q,t)= ( (q[2]-q[1])^2+(p[2]-p[1])^2 )*( (q[2]+q[1])^2+(p[2]+p[1])^2 )
# H2(p,q,t)= (p[2]^4+2*p[2]^2*(q[1]^2-1)+(1+q[1]^2)^2 )*(q[2]^4+2*q[2]^2*(p[1]^2+1)+(p[1]^2-1)^2 )
# HAM(p,q,t)=H1(p,q,t)/H2(p,q,t)
Hamil(XX,YY,QQ,PP)=( (QQ-XX)^2+(PP-YY)^2 )*( (QQ+XX)^2+(PP+YY)^2 )/((PP^4+2*PP^2*(XX^2-1)+(1+XX^2)^2 )*(QQ^4+2*QQ^2*(YY^2+1)+(YY^2-1)^2 ))
H_test(u)=( (u[3]-u[1])^2+(u[2]-u[4])^2 )*( (u[3]+u[1])^2+(u[2]+u[4])^2 )/ ((u[2]^4+2*u[2]^2*(u[1]^2-1)+(1+u[1]^2)^2 )*(u[3]^4+2*u[3]^2*(u[4]^2+1)+(u[4]^2-1)^2 ))

ODE1(z,w)=conj(  im * w.*( 1 ./(z.^2-w.^2)+1 ./(1+w.^2) ))
ODE2(z,w)=conj(  im * z.*( 1 ./(w.^2-z.^2)+1 ./(1+z.^2) ))
function Eq_of_M(du,u,p,t)
    du[1]=real(ODE1(u[1]+im*u[2],u[3]+im*u[4]))
    du[2]=imag(ODE1(u[1]+im*u[2],u[3]+im*u[4]))
    du[3]=real(ODE2(u[1]+im*u[2],u[3]+im*u[4]))
    du[4]=imag(ODE2(u[1]+im*u[2],u[3]+im*u[4]))
    return
end
# refernce
# z=X+iP W=Q+iY . Plane Q=0, P=0
# z=u[1]+i u[2] ; w=u[3]+i u[4]

#Q=0, have P find Y
P_start=8e-3
P_end=1e-1
P=range(P_start,stop=P_end,length=n_iter)
Y=zeros(n_iter)
TEST1_0=zeros(n_iter)
TEST1_f=zeros(n_iter)
u1max=zeros(n_iter)
u2max=zeros(n_iter)
u3max=zeros(n_iter)
u4max=zeros(n_iter)
YofP_find(PP)=sqrt((h+PP^2-h*PP^2 )./(1+h-h*PP^2))
axis([ -.2,.2,-.02, .02]);

# figure(dpi=3600)
for k=1:n_iter
    #  u0=zeros(4)
    # Y[k]=YofP_find(P[k])
    Pval=P[k]
    Y_find1(y)=Hamil(0,y,0,Pval)-H
    u0=zeros(4)
    Y[k]=find_zero(Y_find1,.01, maxeval=100, tol=1e-15)
    u0[1]=0 #X
    u0[2]=P[k] #P
    u0[3]=0 #Q
    u0[4]=Y[k] #Y
    TEST1_0[k]=H_test(u0)/H
    prob = ODEProblem(Eq_of_M,u0,(0., t_end))
    t,A=solve(prob, Vern9(),reltol=1e-15,abstol=1e-12,maxiters=1e15, callback=cb2,save_end=false,save_everystep=false)
    plot(A[3,:],A[2,:],"k,")
    plot(A[3,:],-A[2,:],"k,")
    plot(-A[3,:],A[2,:],"k,")
    plot(-A[3,:],-A[2,:],"k,")

    #testing if energy is conserved at end
    uf=zeros(4)
    uf[1]=A[1,end]
    uf[2]=A[2,end]
    uf[3]=A[3,end]
    uf[4]=A[4,end]
    TEST1_f[k]=H_test(uf)/H

    #finding maxes for exit criteria
    u1max[k]=

 end
#P=0, have Q find Y

Q_start=5e-3
Q_end=1e-1
Q=range(Q_start,stop=Q_end,length=n_iter)
Y=zeros(n_iter)
TEST2=zeros(n_iter)
for k=1:n_iter
    Y_find(YY)=Hamil(0,YY,Q[k],0)-H
    u0=zeros(4)
    Y[k]=find_zero(Y_find,.01, maxeval=100,tol =1e-15)
    u0[1]=0 #X
    u0[2]=0 #P
    u0[3]=Q[k] #Q
    u0[4]=Y[k] #Y
    TEST2[k]=H_test(u0)
    prob = ODEProblem(Eq_of_M,u0,(0., t_end))
    t,A=solve(prob, Vern9(),reltol=1e-15,abstol=1e-12,maxiters=1e15, callback=cb2,save_end=false,save_everystep=false)
    plot(A[3,:],A[2,:],"k,")
    plot(A[3,:],-A[2,:],"k,")
    plot(-A[3,:],A[2,:],"k,")
    plot(-A[3,:],-A[2,:],"k,")
end
# savefig("PSS11.png")


"The empty domain with elements of type `T`."
struct EmptySpace{T} <: Domain{T}
end

const AnyEmptySpace = EmptySpace{Any}

EmptySpace() = EmptySpace{Float64}()
EmptySpace(::Type{T}) where {T} = EmptySpace{T}()

similardomain(::EmptySpace, ::Type{T}) where {T} = EmptySpace{T}()

emptyspace(d) = emptyspace(eltype(d))
emptyspace(::Type{T}) where {T} = EmptySpace{T}()

indomain(x::T, d::EmptySpace{T}) where {T} = false
approx_indomain(x, d::EmptySpace, tolerance) = in(x, d)

show(io::IO, d::EmptySpace) = print(io, "{} (empty domain)")

isempty(d::EmptySpace) = true

isopenset(d::EmptySpace) = true
isclosedset(d::EmptySpace) = true

boundary(d::EmptySpace) = d
interior(d::EmptySpace) = d
closure(d::EmptySpace) = d
boundingbox(d::EmptySpace) = d

# Arithmetic operations

issubset1(d1::EmptySpace, d2) = true
issubset2(d1, d2::EmptySpace) = isempty(d1)

# setdiffdomain2(d1, d2::EmptySpace) = d2

map_domain(map::Map{T}, d::EmptySpace{T}) where {T} = d
mapped_domain(map::Map, d::EmptySpace) = EmptySpace{codomaintype(map)}()

==(d1::EmptySpace, d2::EmptySpace) = true
isequal1(d1::EmptySpace, d2) = isempty(d2)
isequal2(d1, d2::EmptySpace) = isempty(d1)

"The full space of elements of type `T`."
struct FullSpace{T} <: Domain{T} end

const AnyFullSpace = FullSpace{Any}

FullSpace() = FullSpace{Float64}()
FullSpace(d) = FullSpace{eltype(d)}()

fullspace(d) = fullspace(eltype(d))
fullspace(::Type{T}) where {T} = FullSpace{T}()

isfullspace(d::FullSpace) = true
isfullspace(d::Domain) = false

similardomain(::FullSpace, ::Type{T}) where {T} = FullSpace{T}()

euclideanspace(n::Val{N}) where {N} = euclideanspace(n, Float64)
euclideanspace(::Val{N}, ::Type{T}) where {N,T} = FullSpace{SVector{N,T}}()

indomain(x::T, d::FullSpace{T}) where {T} = true
approx_indomain(x, d::FullSpace, tolerance) = in(x, d)

show(io::IO, d::FullSpace) = print(io, "{x} (full space)")

# We choose the origin as a point in the full space
point_in_domain(d::FullSpace) = zero(eltype(d))

isempty(::FullSpace) = false

isopenset(d::FullSpace) = true
isclosedset(d::FullSpace) = true

boundary(d::FullSpace{T}) where {T} = EmptySpace{T}()
interior(d::FullSpace) = d
closure(d::FullSpace) = d

# Arithmetic operations

issubset1(d1::FullSpace, d2) = isfullspace(d2)
issubset2(d1, d2::FullSpace) = true

map_domain(m::AbstractAffineMap{T}, d::FullSpace{T}) where {T} = d

==(d1::FullSpace, d2::FullSpace) = true
isequal1(d1::FullSpace, d2) = isfullspace(d2)
isequal2(d1, d2::FullSpace) = isfullspace(d1)

convert(::Type{Domain}, ::Type{T}) where T = FullSpace{T}()
convert(::Type{Domain{S}}, ::Type{T}) where {T,S} = convert(Domain{S}, convert(Domain, T))

infimum(d::FullSpace{T}) where {T} = typemin(T)
supremum(d::FullSpace{T}) where {T} = typemax(T)

# Some convenient complete domains

"The set of integers of type Int (ℤ = \\BbbZ)."
const ℤ = FullSpace{Int}()
"The set of rational numbers of type Rational{Int} (ℚ = \\BbbQ)."
const ℚ = FullSpace{Rational{Int}}()
"The set of reals of type Float64 (ℝ = \\BbbR)."
const ℝ = FullSpace{Float64}()
"The complex plane with Float64 real and imaginar parts (ℂ = \\BbbC)."
const ℂ = FullSpace{Complex{Float64}}()

"The space ℝ^1"
const ℝ1 = FullSpace{SVector{1,Float64}}()
"The space ℝ^2"
const ℝ2 = FullSpace{SVector{2,Float64}}()
"The space ℝ^3"
const ℝ3 = FullSpace{SVector{3,Float64}}()
"The space ℝ^4"
const ℝ4 = FullSpace{SVector{4,Float64}}()


#=

    This is a port of the R2 model TwoCoins.cs

    Output from the R2 model:
    ```
    [0] Mean: 0.339              firstCoin
    [0] Variance: 0.224303
    [1] Mean: 0.309              secondCoin
    [1] Variance: 0.213733
    Number of accepted samples = 847
    ```

    (The R2 model TwoCoinsTuple.cs is basically the same model as TwoCoins, except that
     the return value is a tuple instead of the array [firstCoin,secondCoin], and is
     not ported to Turing.jl.)

    This model:

  parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

   firstCoin    0.3290    0.4701     0.0149    0.0269   487.5706    1.0010      156.9265
  secondCoin    0.3530    0.4781     0.0151    0.0189   636.1597    0.9994      204.7505
   bothHeads    0.0000    0.0000     0.0000    0.0000        NaN       NaN           NaN

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV, DataFrames
include("jl_utils.jl")

#
# Simulate skills and games
#
@model function two_coins()
    firstCoin  ~ flip(0.5)
    secondCoin ~ flip(0.5)
    bothHeads ~ Dirac(firstCoin && secondCoin)

    true ~ Dirac(bothHeads == false)
end

model = two_coins()

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 1_000)
chns = sample(model, PG(5), 1_000)
# chns = sample(model, SMC(), 1_000)
# chns = sample(model, IS(), 1_000)

# chns = sample(model, NUTS(), 1_000)
# chns = sample(model, HMC(0.1,6), 1_000)

display(chns)


using SparseGrids
using Test

@testset "Unique nodes" begin
    @testset "All nodes are different" begin
        nodes = [[0;], [1;]]
        weights = [1; 1]

        un, uw = SparseGrids.uniquenodes(nodes, weights)

        @test un == nodes
        @test uw == weights
    end

    @testset "Two identical nodes" begin
        nodes = [[0;], [0;]]
        weights = [1; 1]

        un, uw = SparseGrids.uniquenodes(nodes, weights)

        @test un == unique(nodes)
        @test uw == [sum(weights)]
    end

    @testset "Two of three nodes identical" begin
        nodes = [[0;], [1;], [0;]]
        weights = [1, 2, 3]

        un, uw = SparseGrids.uniquenodes(nodes, weights)

        @test un == [[0;], [1;]]
        @test uw == [4; 2]
    end
end