using AbstractAlgebra
using Random

@testset "Input-output AbstractAlgebra" begin

    rng = Random.MersenneTwister(42)

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:lex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :lex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:deglex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :deglex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:deglex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :deglex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3//4, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:lex)
    fs = [x^2*y + 3//4, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :lex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:deglex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :deglex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:deglex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y, x*y - y^2, x*y - x^2]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    meta = Groebner.set_metaparameters(ring, :deglex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs
end

@testset "Input-output generic :hasparent" begin

    rng = Random.MersenneTwister(42)

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:lex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    ring.origring = :hasparent
    meta = Groebner.set_metaparameters(ring, :lex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(GF(2^31 - 1), ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    ring.origring = :hasparent
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:degrevlex)
    fs = [x^2*y + 3//4, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    ring.origring = :hasparent
    meta = Groebner.set_metaparameters(ring, :degrevlex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs

    R, (x, y) = PolynomialRing(QQ, ["x", "y"], ordering=:lex)
    fs = [x^2*y + 3//4, (2^31 - 5)*x - (2^31 - 4)*y]
    ring, exps, cfs = Groebner.convert_to_internal(fs, :input)
    ring.origring = :hasparent
    meta = Groebner.set_metaparameters(ring, :lex, false, false, :exact, rng)
    fsfs = Groebner.convert_to_output(ring, fs, exps, cfs, meta)
    @test fsfs == fs
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMBeam.jl/blob/master/LICENSE

"""
    get_beam_mass_matrix_2d(X1, X2, A, ro)

Function integrates mass matrix for 6 DOF Euler-Bernoulli beam element in 2d.

X1 = beams left node coordinates
X2 = beams right node coordinates
A = Cross section area
ro = Density
"""
function get_beam_mass_matrix_2d(X1,X2,A,ro)
    le=norm(X2-X1)                      # Lenght of element
    a=atan((X2[2]-X1[2])/(X2[1]-X1[1])) # Rotation angle of the element
    nn=2                            # Number of nodes
    nd=3*1+3                           # Number of DOFs
    Gp=[-1/3*sqrt(5+2*sqrt(10/7)) -1/3*sqrt(5-2*sqrt(10/7)) 0 1/3*sqrt(5-2*sqrt(10/7)) 1/3*sqrt(5+2*sqrt(10/7))] # Five Gauss integration points
    w=[(322-13*sqrt(70))/900 (322+13*sqrt(70))/900 128/225 (322+13*sqrt(70))/900 (322-13*sqrt(70))/900] # Five Gauss integration weights
    # Rotation matrix
    B=[cos(a) sin(a) 0  0     0       0;
      -sin(a) cos(a) 0  0     0       0;
       0      0      1  0     0       0;
       0      0      0  cos(a) sin(a) 0;
       0      0      0 -sin(a) cos(a) 0;
       0      0      0  0      0      1]
    # Integration of the truss mass matrix tm
    detJ=le/2
    function tmint(wi::Float64, xi::Float64)
        N1s=1+(-1/2)*(1+xi)
        N2s=(1/2)*(1 + xi)
        Ns=[N1s N2s]
        return wi*Ns'*Ns
    end
    tm=zeros(2,2)
    for i = 1:size(Gp,2)
        tm +=tmint(w[i],Gp[i])
    end
    tm=ro*A*tm*detJ
    # Integration of the beam mass matrix bm
    detJ=(le/2)
    function bmint(wi::Float64, xi::Float64)
        N1=1/4*(1-xi)^2*(2+xi)
        N2=le/8*(1-xi)^2*(xi+1)
        N3=1/4*(1+xi)^2*(2-xi)
        N4=le/8*(1+xi)^2*(xi-1)
        N=[N1 N2 N3 N4]
        return wi*N'*N
    end
    bm=zeros(4,4)
    for i = 1:size(Gp,2)
        bm +=bmint(w[i],Gp[i])
    end
    bm=ro*A*bm*detJ
    # Assembly of the 6 DOF truss-beam mass matrix m
    m=zeros(6,6)
    m[1,1]=tm[1,1];m[1,4]=tm[1,2];m[4,1]=tm[2,1];m[4,4]=tm[2,2]
    m[2,2]=bm[1,1];m[2,3]=bm[1,2];m[2,5]=bm[1,3];m[2,6]=bm[1,4]
    m[3,2]=bm[2,1];m[3,3]=bm[2,2];m[3,5]=bm[2,3];m[3,6]=bm[2,4]
    m[5,2]=bm[3,1];m[5,3]=bm[3,2];m[5,5]=bm[3,3];m[5,6]=bm[3,4]
    m[6,2]=bm[4,1];m[6,3]=bm[4,2];m[6,5]=bm[4,3];m[6,6]=bm[4,4]
    # Rotation
    m=B'*m*B
    return m
end


include("eggify.jl")
using Metatheory.Library
using Metatheory.EGraphs.Schedulers

or_alg = @theory begin
  ((p ∨ q) ∨ r) == (p ∨ (q ∨ r))
  (p ∨ q) == (q ∨ p)
  (p ∨ p)     => p
  (p ∨ true)  => true
  (p ∨ false) => p
end

and_alg = @theory begin
  ((p ∧ q) ∧ r) == (p ∧ (q ∧ r))
  (p ∧ q) == (q ∧ p)
  (p ∧ p)     => p
  (p ∧ true)  => p
  (p ∧ false) => false
end

comb = @theory begin
  # DeMorgan
  ¬(p ∨ q) == (¬p ∧ ¬q)
  ¬(p ∧ q) == (¬p ∨ ¬q)
  # distrib
  (p ∧ (q ∨ r)) == ((p ∧ q) ∨ (p ∧ r))
  (p ∨ (q ∧ r)) == ((p ∨ q) ∧ (p ∨ r))
  # absorb
  (p ∧ (p ∨ q)) => p
  (p ∨ (p ∧ q)) => p
  # complement
  (p ∧ (¬p ∨ q)) => p ∧ q
  (p ∨ (¬p ∧ q)) => p ∨ q
end

negt = @theory begin
  (p ∧ ¬p)   => false
  (p ∨ ¬(p)) => true
  ¬(¬p) == p
end

impl = @theory begin
  (p == ¬p) => false
  (p == p)  => true
  (p == q)  => (¬p ∨ q) ∧ (¬q ∨ p)
  (p => q)  => (¬p ∨ q)
end

fold = @theory begin
  (true == false)  => false
  (false == true)  => false
  (true == true)   => true
  (false == false) => true
  (true ∨ false)   => true
  (false ∨ true)   => true
  (true ∨ true)    => true
  (false ∨ false)  => false
  (true ∧ true)    => true
  (false ∧ true)   => false
  (true ∧ false)   => false
  (false ∧ false)  => false
  ¬(true)          => false
  ¬(false)         => true
end

theory = or_alg ∪ and_alg ∪ comb ∪ negt ∪ impl ∪ fold

query = :(¬(((¬p ∨ q) ∧ (¬r ∨ s)) ∧ (p ∨ r)) ∨ (q ∨ s))

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

params = SaturationParams(timeout = 22, eclasslimit = 3051, scheduler = ScoredScheduler)#, schedulerparams=(1000,5, Schedulers.exprsize))

for i in 1:2
  G = EGraph(query)
  report = saturate!(G, theory, params)
  ex = extract!(G, astsize)
  println("Best found: $ex")
  println(report)
end

open("src/main.rs", "w") do f
  write(f, rust_code(theory, query, params))
end


function cr3bpEomControlMTKSys(μ::AbstractFloat)
    @parameters atx aty atz
    @variables t x(t) y(t) z(t) vx(t) vy(t) vz(t)
    D = Differential(t)

    # Compute r1 and r2
    r1 = sqrt((x+μ)^2 + y^2 + z^2)
    r2 = sqrt((x+μ-1)^2 + y^2 + z^2)

    # Compute g
    gx = x - (1 - μ)*(x + μ)/r1^3 - μ*(x + μ - 1)/r2^3
    gy = y - (1 - μ)*y/r1^3 - μ*y/r2^3
    gz = -(1 - μ)*z/r1^3 - μ*z/r2^3

    # Compute h
    hx = 2*vy
    hy = -2*vx

    eqs = [D(x) ~ vx,
           D(y) ~ vy,
           D(z) ~ vz,
           D(vx) ~ gx + hx + atx,
           D(vy) ~ gy + hy + aty,
           D(vz) ~ gz + atz]

    return ODESystem(eqs)
end

function cr3bpEomNoControlMTKSys(μ::AbstractFloat)
    _cr3bpEomWControl = cr3bpEomControlMTKSys(μ)
    (atx, aty, atz) = parameters(_cr3bpEomWControl)
    eqs = [eq.lhs ~ simplify(
        substitute(substitute(substitute(
            eq.rhs, atx => 0), aty => 0), atz => 0)) for
            eq in equations(_cr3bpEomWControl)]
    return ODESystem(eqs)
end

function cr3bpEOMIndirect(μ::AbstractFloat,    TU::AbstractFloat,
                          LU::AbstractFloat,   MU::AbstractFloat,
                          tMax::AbstractFloat, c::AbstractFloat)
    @parameters u
    @variables t x(t) y(t) z(t) vx(t) vy(t) vz(t) m(t)
    @variables λx(t) λy(t) λz(t) λvx(t) λvy(t) λvz(t) λm(t) H(t)
    D    = Differential(t)

    # Compute r1 and r2
    r1 = sqrt((x+μ)^2 + y^2 + z^2)
    r2 = sqrt((x+μ-1)^2 + y^2 + z^2)

    # Compute g
    gx = x - (1 - μ)*(x + μ)/r1^3 - μ*(x + μ - 1)/r2^3
    gy = y - (1 - μ)*y/r1^3 - μ*y/r2^3
    gz = -(1 - μ)*z/r1^3 - μ*z/r2^3

    # Compute h
    hx = 2*vy
    hy = -2*vx

    # Compute G
    G11     = 1 - (1 - μ)/r1^3 + 3*(1 - μ)*(x + μ)^2/r1^5 -
                μ/r2^3 + 3*μ*(x + μ - 1)^2/r2^5;
    G22     = 1 - (1 - μ)/r1^3 + 3*(1 - μ)*y^2/r1^5 -
                μ/r2^3 + 3*μ*y^2/r2^5;
    G33     = -(1 - μ)/r1^3 + 3*(1 - μ)*z^2/r1^5 -
                μ/r2^3 + 3*μ*z^2/r2^5;
    G12     = 3*(1 - μ)*(x + μ)*y/r1^5 +
                3*μ*(x + μ - 1)*y/r2^5;
    G13     = 3*(1 - μ)*(x + μ)*z/r1^5 +
                3*μ*(x + μ - 1)*z/r2^5;
    G23     = 3*(1 - μ)*y*z/r1^5 +
                3*μ*y*z/r2^5;

    # Compute scaled max thrust and exhaust velocity
    tMaxSc = tMax * TU * TU / (MU*LU*1000.0)
    cSc = c * TU / (LU*1000.0)

    # Compute thrust direction
    λv = sqrt(λvx^2 + λvy^2 + λvz^2)
    αx = -λvx / λv
    αy = -λvy / λv
    αz = -λvz / λv

    # Compute thrust acceleration
    atx = u*tMaxSc*αx / m
    aty = u*tMaxSc*αy / m
    atz = u*tMaxSc*αz / m

    # Dynamics
    dx   = vx
    dy   = vy
    dz   = vz
    ddx  = gx + hx + atx
    ddy  = gy + hy + aty
    ddz  = gz + atz
    dm   = -u*tMaxSc / cSc
    dλx  = -G11*λvx - G12*λvy - G13*λvz
    dλy  = -G12*λvx - G22*λvy - G23*λvz
    dλz  = -G13*λvx - G23*λvy - G33*λvz
    dλvx = -λx + 2.0*λvy
    dλvy = -λy - 2.0*λvx
    dλvz = -λz
    dλm  = -λv*u*tMaxSc / (m^2)

    # Equations
    eqs = [D(x)   ~ dx,
           D(y)   ~ dy,
           D(z)   ~ dz,
           D(vx)  ~ ddx,
           D(vy)  ~ ddy,
           D(vz)  ~ ddz,
           D(m)   ~ dm,
           D(λx)  ~ dλx,
           D(λy)  ~ dλy,
           D(λz)  ~ dλz,
           D(λvx) ~ dλvx,
           D(λvy) ~ dλvy,
           D(λvz) ~ dλvz,
           D(λm)  ~ dλm]

    return ODESystem(eqs, t)
end


@diffusion_process FitzHughDiffusionAltAuxLin begin
    :dimensions
    process --> 2
    wiener --> 1

    :parameters
    (ϵ, s, γ, β, σ, t, T) --> Float64
    (u, v) --> ℝ{2}

    :additional
    constdiff --> true
    linear --> true
end

function B(t, P::FitzHughDiffusionAltAuxLin)
    @SMatrix [
        0.0  1.0;
        (1.0-P.γ-3.0*P.v[1]^2-6*P.v[1]*P.v[2])/P.ϵ (1.0-P.ϵ-3.0*P.v[1]^2)/P.ϵ
    ]
end

function β(t, P::FitzHughDiffusionAltAuxLin)
    ℝ{2}(
        0.0,
        (2*P.v[1]^3+P.s-P.β+6*P.v[1]^2*P.v[2])/P.ϵ
    )
end

σ(t, P::FitzHughDiffusionAltAuxLin) = ℝ{2}(0.0, P.σ/P.ϵ)


using TensorTrains

w=2
L=5
h,V = hubbard_2D(w,L;U=4,w_pbc=false,L_pbc=false)
H_tto = hV_to_mpo(h,V,tol=1e-5)
ψ_0 = half_filling(w*L)
E_als,ψ_r = als_eigsolv(H_tto,ψ_0,sweep_schedule=[2,5],rmax_schedule=[16,32])
E_mals,ψ_tt,r = mals_eigsolv(H_tto,ψ_r,tol=1e-6,sweep_schedule=[2,5],rmax_schedule=[32,64])
#Fiedler order
σ_fiedler = fiedler_order(ψ_r)
h_f = h[σ_fiedler,σ_fiedler]
V_f = V[σ_fiedler,σ_fiedler,σ_fiedler,σ_fiedler]
H_tto = hV_to_mpo(h_f,V_f,tol=1e-5)
E_fiedler,ψ_fiedler,r = mals_eigsolv(H_tto,ψ_r,tol=1e-6,sweep_schedule=[3],rmax_schedule=[64])
#BWPO order
σ_bwpo = bwpo_order(ψ_r)
h_b = h[σ_bwpo,σ_bwpo]
V_b = V[σ_bwpo,σ_bwpo,σ_bwpo,σ_bwpo]
H_tto = hV_to_mpo(h_b,V_b,tol=1e-5)
E_bwpo,ψ_bwpo,r = mals_eigsolv(H_tto,ψ_r,tol=1e-6,sweep_schedule=[3],rmax_schedule=[64])


using TensorOperations

function propagate(LHS, X, Y)
    @tensor R[6,7,8] := LHS[1,2,3]*X[1,4,6]*Y[2,5,7,4]*X[3,5,8]
end


"""
    SinusoidSource(eltype, samplerate, freqs)

SinusoidSource is a single-channel sine-tone signal generator. `freqs` can be an
array of frequencies for a multi-frequency source, or a single frequency for a
single sinusoid source.

Inputs
------
* `samplerate` specifies the sample rate of the signal.
* `freqs` sinusoid frequencies to generate.

Output
------
* SampleSource object

Example
-------
```julia
source_object = SinusoidSource(Float64, 48u"kHz", 200:200:2400)
cn = read(source_object, 50u"ms")     # Generate 50 ms of harmonic stack audio
```
"""
mutable struct SinusoidSource{T} <: SampleSource
    samplerate::Float64
    freqs::Vector{Float64} # in radians/sample
    phases::Vector{Float64}

    function SinusoidSource(eltype, samplerate::Number, freqs::Array)
        # convert frequencies from cycles/sec to rad/sample
        radfreqs = map(f->2pi*f/samplerate, freqs)
        new{eltype}(Float64(samplerate), radfreqs, zeros(length(freqs)))
    end
    SinusoidSource(eltype, samplerate, freq::StepRange) = SinusoidSource(eltype, samplerate, collect(freq))
    SinusoidSource(eltype, samplerate::Number, freq::Real) = SinusoidSource(eltype, samplerate, [freq])
    SinusoidSource(eltype, samplerate::Unitful.Frequency, freq::Real) = SinusoidSource(eltype, samplerate |> u"Hz" |> ustrip, [freq])
    SinusoidSource(eltype, samplerate::Unitful.Frequency, freq::Array) = SinusoidSource(eltype, samplerate |> u"Hz" |> ustrip, freq)
end

Base.eltype(::SinusoidSource{T}) where T = T
nchannels(source::SinusoidSource) = 1
samplerate(source::SinusoidSource) = source.samplerate

function unsafe_read!(source::SinusoidSource, buf::Array, frameoffset, framecount)
    inc = 2pi / samplerate(source)

    for i in 1:framecount
        buf[i+frameoffset, 1] = 0
    end

    for tone_idx in 1:length(source.freqs)
        tone_freq = source.freqs[tone_idx]
        tone_phas = source.phases[tone_idx]

        for i in 1:framecount
            buf[i+frameoffset, 1] += sin.(tone_phas)
            tone_phas += tone_freq
        end
        source.phases[tone_idx] = tone_phas

    end

    framecount
end


using EconPDEs, Distributions

Base.@kwdef struct HaddadModel
  # consumption process parameters
  μbar::Float64 = 0.018
  vbar::Float64 = 0.00052
  κμ::Float64 = 0.3
  νμ::Float64 = 0.456
  κv::Float64 = 0.012
  νv::Float64 = 0.00472

  # active capital
  αbar::Float64 = 1.2
  λ::Float64 = 0.018

  # utility parameters
  ρ::Float64 = 0.0132
  γ::Float64 = 10.0
  ψ::Float64 = 1.5
end

function initialize_stategrid(m::HaddadModel; μn = 30, vn = 30)
  μbar = m.μbar; vbar = m.vbar ; κμ = m.κμ ; νμ = m.νμ ; κv = m.κv ; νv = m.νv ; αbar = m.αbar ; λ = m.λ ; ρ = m.ρ ; γ = m.γ ; ψ = m.ψ

  σ = sqrt(νμ^2 * vbar / (2 * κμ))
  μmin = quantile(Normal(μbar, σ), 0.001)
  μmax = quantile(Normal(μbar, σ), 0.999)
  μs = range(μmin,  μmax, length = μn)

  α = 2 * κv * vbar / νv^2
  β = νv^2 / (2 * κv)
  vmin = quantile(Gamma(α, β), 0.001)
  vmax = quantile(Gamma(α, β), 0.999)
  vs = range(vmin,  vmax, length = vn)

  OrderedDict(:μ => μs, :v => vs)
end

function (m::HaddadModel)(state::NamedTuple, y::NamedTuple)
  (; μbar, vbar, κμ, νμ, κv, νv, αbar, λ, ρ, γ, ψ) = m
  (; μ, v) = state
  (; p, pμ_up, pμ_down, pv_up, pv_down, pμμ, pμv, pvv) = y

  # drift and volatility of μ, ν, p
  μc = μ
  σc = sqrt(v)
  μμ = κμ * (μbar - μ)
  σμ = νμ * sqrt(v)
  μv = κv * (vbar - v)

  pμ = (μμ >= 0) ? pμ_up : pμ_down
  pv = (νμ >= 0) ? pv_up : pv_down
  σv = νv * sqrt(v)
  σp_Zμ = pμ / p * σμ
  σp_Zv = pv / p * σv
  σp2 = σp_Zμ^2 + σp_Zv^2
  μp = pμ / p * μμ + pv / p * μv + 0.5 * pμμ / p * σμ^2 + 0.5 * pvv / p * σv^2

  # Market Price of Risk κ
  αstar = αbar - sqrt(λ * αbar / (0.5 * γ * (σc^2 + σp_Zμ^2 + σp_Zv^2)))
  αstar = clamp(αstar, 0.0, 1.0)
  κ_Zc = αstar * γ * σc
  κ_Zμ = (αstar - (1 / γ - 1) / (ψ - 1)) * γ * σp_Zμ
  κ_Zv = (αstar - (1 / γ - 1) / (ψ - 1)) * γ * σp_Zv
  κ2 = κ_Zc^2 + κ_Zμ^2 + κ_Zv^2

  # Interest rate r
  #r = ρ + μc / ψ - (1 - 1 / ψ) / (2 * γ) * κ2 - (1 / ψ - γ) / (2 * γ * (ψ - 1)) * σp2 - (1 - γ) / (γ * ψ) * (κμ * σμ + κv * σv) + 1 / ψ  * (κ_Zc * σc + κμ * σμ + κv * σv)
  #out = p * (1 / p + μc + μp - r - κ_Zc * σc - κ_Zμ * σp_Zμ - κ_Zv * σp_Zv)

  pt = - p * (1 / p - ρ + (1 - 1 / ψ) * (μc - 0.5 * γ * σc^2 * (1 - (αstar - 1)^2)) + μp + (0.5 * (1 / ψ - γ) / (1 - 1 / ψ) + 0.5 * γ * (1 - 1 / ψ) * (αstar - 1)^2) * σp2)
  return (; pt)
end

m = HaddadModel()
μn = 30
vn = 30
stategrid = initialize_stategrid(m; μn = μn, vn = vn)
yend = OrderedDict(:p => ones(μn, vn))
result = pdesolve(m, stategrid, yend)
@assert result.residual_norm <= 1e-5


## Cauchy

# pseudocauchy does not normalize at ∞
pseudostieltjes(S::Space{<:Arc},f,z) = stieltjes(setcanonicaldomain(S),f,mobius(S,z))
pseudohilbert(S::Space{<:Arc},f,z) = hilbert(setdomain(S,ChebyshevInterval()),f,mobius(S,z))

stieltjes(S::Space{<:Arc},f,z) =
    stieltjes(setcanonicaldomain(S),f,mobius(S,z))-stieltjes(setcanonicaldomain(S),f,mobius(S,Inf))
function hilbert(S::Space{<:Arc},f,z)
    y = mobius(S,z)
    abs(imag(y)) < 10E-15 || throw(ArgumentError())
    hilbert(setcanonicaldomain(S),f,real(y))+(1/π)*stieltjes(setcanonicaldomain(S),f,mobius(S,Inf))
end

function PseudoHilbert(sp::JacobiWeight{LS,RR},k::Integer) where {LS,RR<:Arc}
    @assert k==1
    csp=setcanonicaldomain(sp)
    H=Hilbert(csp)
    PseudoHilbertWrapper(SpaceOperator(H,setdomain(domainspace(H),domain(sp)),setdomain(rangespace(H),domain(sp))),k)
end

function Hilbert(sp::JacobiWeight{LS,RR},k::Integer) where {LS,RR<:Arc}
    @assert k==1
    csp=setcanonicaldomain(sp)
    St = Stieltjes(csp,mobius(sp,Inf))
    H=Hilbert(csp)+(1/π)*SpaceOperator(St,csp,ConstantSpace(eltype(St)))
    HilbertWrapper(SpaceOperator(H,setdomain(domainspace(H),domain(sp)),setdomain(rangespace(H),domain(sp))),k)
end

## stieltjesintegral

function stieltjesintegral(sp::Space{<:Arc},w,z)
    g=Fun(setcanonicaldomain(sp),w)*fromcanonicalD(sp)
    stieltjesintegral(g,mobius(sp,z))-
        stieltjesintegral(g,mobius(sp,Inf))+
        sum(Fun(sp,w))*log(z-fromcanonical(sp,Inf))
end

function linesumstieltjesintegral(sp::Space{<:Arc},w,z)
    g=Fun(setcanonicaldomain(sp),w)*abs(fromcanonicalD(sp))
    stieltjesintegral(g,mobius(sp,z))-
        stieltjesintegral(g,mobius(sp,Inf))+
        linesum(Fun(sp,w))*log(z-fromcanonical(sp,Inf))
end

function logkernel(sp::Space{<:Arc},w,z)
    g=Fun(setcanonicaldomain(sp),w)*abs(fromcanonicalD(sp))
    logkernel(g,mobius(sp,z))-
        logkernel(g,mobius(sp,Inf))+
        linesum(Fun(sp,w))*logabs(z-fromcanonical(sp,Inf))/π
end

function SingularIntegral(S::JacobiWeight{JW,RR},k::Integer) where {JW,RR<:Arc}
    d=domain(S)
    if k==0
        tol=1E-15
        # the mapped logkernel
        csp=setcanonicaldomain(S)
        Σ=SingularIntegral(csp,0)
        M=Multiplication(abs(fromcanonicalD(d,Fun(identity,csp))),csp)

        z∞=mobius(d,Inf)
        cnst=Vector{Float64}()
        for j = 1:10000
            push!(cnst,logkernel(Fun(csp,[zeros(j-1);1.]),z∞))
            if j ≥ 3 && norm(cnst[end-2:end]) < tol
                break
            end
        end
        L∞=FiniteOperator(cnst',csp,ConstantSpace(eltype(cnst)))

        x=Fun(identity,S)
        SpaceOperator((Σ-L∞)*M,S,setdomain(rangespace(Σ),d)) +
            (logabs(x-fromcanonical(d,Inf))/π)*DefiniteLineIntegral(S)
    else
        # multiply by abs(M')/M' to change to dz to ds
        Mp=fromcanonicalD(d)
        Hilbert(S,k)[setdomain(abs(Mp)/Mp,d)]
    end
end


# module NeighbourOperators

# using Random: randperm
# using Distributions: NegativeBinomial
# using StatsBase: sample, ProbabilityWeights
# using Solutions: Solution, swap!, permu!

# export small_shift!,
#        swap_close_planes!,
#        swap_close_costly_planes!,
#        permutation_operator!

function small_shift!(sol::Solution; p=0.5)
    n = sol.inst.nb_planes
    shifted_plane = sample(1:n)
    new_position = find_shift_destination(shifted_plane, p, n)
    shift!(sol, shifted_plane, new_position)
end

function shift_costly_plane!(sol::Solution; p=0.5)
    shifted_plane = cost_weighted_sample(sol)
    new_position = find_shift_destination(shifted_plane, p,
                                          sol.inst.nb_planes)
    shift!(sol, shifted_plane, new_position)
end

"""Given a position, returns a random but valid destination.

The size of the shift is a random variable 1 + X,
where X follows a Negative Binomial of parameters (1, p)
"""
function find_shift_destination(current_position::Int, p::Float64, n_planes::Int)
    negative_binomial = NegativeBinomial(1, p)
    direction = sample([-1, 1])
    norm = 1 + rand(negative_binomial)
    new_position = current_position + direction * norm
    new_position = min(max(new_position, 1), n_planes)
    return new_position
end

"""Swaps a random plane and one of its neighbour."""
function swap_close_planes!(sol::Solution)
    n = sol.inst.nb_planes
    swapped_plane_1_id = sample(1:n)
    if swapped_plane_1_id == n
        swapped_plane_2_id = n-1
    elseif swapped_plane_1_id == 1
        swapped_plane_2_id = 2
    else
        swapped_plane_2_id = swapped_plane_1_id + sample([-1, 1])
    end
    swap!(sol, swapped_plane_1_id, swapped_plane_2_id)
end

"""Swaps a random plane according to the cost distribution and one of its neighbour."""
function swap_close_costly_planes!(sol::Solution)
    n = sol.inst.nb_planes
    swapped_plane_1_id = cost_weighted_sample(sol)
    if swapped_plane_1_id == n
        swapped_plane_2_id = n-1
    elseif swapped_plane_1_id == 1
        swapped_plane_2_id = 2
    else
        swapped_plane_2_id = swapped_plane_1_id + sample([-1, 1])
    end
    swap!(sol, swapped_plane_1_id, swapped_plane_2_id)
end

function permutation_operator!(sol::Solution; permuted_ratio=0.1)
    n_permuted_planes = Int(round(permuted_ratio * sol.inst.nb_planes))
    println(n_permuted_planes)
    if n_permuted_planes <= 2
        swap!(sol)
    else
        permuted_planes = sample(1:sol.inst.nb_planes, n_permuted_planes,
                                replace=false)
        permu!(sol, permuted_planes,
               permuted_planes[randperm(length(permuted_planes))])
    end
end

function cost_weighted_sample(sol::Solution)
    sample(1:sol.inst.nb_planes, ProbabilityWeights(sol.costs))
end

# end


include("common.jl")

mmesh = Mesh(joinpath(PDATA,"rectangle.stl"))
c1=(1.0, 0.0)
c2=(1.0, 1.0)
fault_idx, fault_coords, fault_normal_vec = get_fault_info(mmesh, c1, c2)

mmesh, conn = fix_mesh(mmesh, fault_idx, rectangle_redblack)
close("all")
visualize_mesh(mmesh)
v = 0.05*(rand(mmesh.nnode).-0.5)
u = 0.05*(rand(mmesh.nnode).-0.5)
scatter(mmesh.nodes[:,1] + u, mmesh.nodes[:,2] + v, s = 2)
savefig("mesh.png")

bdnode = bcnode((x, y)->(x<1e-5 || x>2-1e-5), mmesh)

Conn = compute_slip_boundary_condition( conn, mmesh)
NT = 100
Δt = 0.004
λ = 1.0
μ = 1.0
invη = ones(get_ngauss(mmesh))
sn = size(conn,1)
slipvec = [zeros(NT, sn) ones(NT, sn)] # acceleration
gravity = 0.0

U, Sigma, Varepsilon = subduction_solver(invη, gravity, Conn, slipvec, bdnode, mmesh)

sess = Session(); init(sess)
Uval, Stress, Strain = run(sess, [U, Sigma, Varepsilon])
p = visualize_displacement(Uval[1:5:end,:], mmesh)
saveanim(p, "subduction_rectangle.gif")

close("all")
u = Uval[end,:]
visualize_scalar_on_fem_points(u[mmesh.ndof+1:2mmesh.ndof],mmesh)
savefig("y_displacement.png")


# Closed-form formulas for the inner products between different kinds of components

# step x step
function GMStepStep(m::Int, M::Int, N::Int, μm::Float64, σm::Float64, μM::Float64, σM::Float64)::Float64

  return  (m*μm*μM-(M-m)*(1-μm)*μM+(N-M)*(1-μm)*(1-μM)) :: Float64

end
function GMStepStep(i::Int, j::Int, d, IT)::Float64
  	if i>j
    	return GMStepStep(j, i, IT, d.μ[STEP][j], d.σ[STEP][j], d.μ[STEP][i], d.σ[STEP][i]) :: Float64
    else
    	return GMStepStep(i, j, IT, d.μ[STEP][i], d.σ[STEP][i], d.μ[STEP][j], d.σ[STEP][j]) :: Float64
    end
    #return 0.0 :: Float64
end

# step x slope
function GMStepSlope(
    t::Int, l::Int, N::Int, μSTEP::Float64, σSTEP::Float64,
    μSLOPE::Float64, σSLOPE::Float64
    )::Float64

  if l > t
    return (
      t*μSTEP*μSLOPE - (l-t)*(1-μSTEP)*μSLOPE - (N+1)*l-μSLOPE*(N+1)
      + (N+1)*l*μSTEP + μSLOPE*(N+1)*μSTEP + (1/2)*(N+1)^2-(1/2)*N
      - (1/2)*μSTEP*(N+1)^2 + (1/2)*μSTEP*(N+1) + (l+1)*l+μSLOPE*(l+1)
      - (l+1)*l*μSTEP - μSLOPE*(l+1)*μSTEP - (1/2)*(l+1)^2+(1/2)*l
      + (1/2)*μSTEP*(l+1)^2 - (1/2)*μSTEP*(l+1)
      ) :: Float64
  else
    return (
      - (l+1)*l*μSTEP+(1/2)*t - (1/2)*N+(1/2)*(N+1)^2 + μSLOPE*(N+1)*μSTEP
      - μSLOPE*(l+1)*μSTEP - μSLOPE*(N+1) - (1/2)*μSTEP*(N+1)^2
      + (1/2)*μSTEP*(N+1) + (1/2)*μSTEP*(l+1)^2 - (1/2)*μSTEP*(l+1)
      + μSLOPE*(t+1) + l*μSTEP*μSLOPE - (1/2)*(t+1)^2 + (N+1)*l*μSTEP
      - (N+1)*l+(t+1)*l
      ) :: Float64
  end

  #return 0.0 :: Float64
end
function GMStepSlope(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMStepSlope(i, j, IT, d.μ[STEP][i], d.σ[STEP][i], d.μ[SLOPE][j], d.σ[SLOPE][j]) :: Float64
end
function GMSlopeStep(i::Int, j::Int, d, IT)::Float64
    return GMStepSlope(j, i, IT, d.μ[STEP][j], d.σ[STEP][j], d.μ[SLOPE][i], d.σ[SLOPE][i]) :: Float64
end

# step x spike
function GMStepSpike(
    i::Int, j::Int, N::Int, μSTEP::Float64, σSTEP::Float64,
    μSPIKE::Float64, σSPIKE::Float64
    )::Float64

  if j > i
    return (
      i*μSPIKE*μSTEP - (N-i)*(1-μSTEP)*μSPIKE + (1-μSTEP)*μSPIKE + (1-μSPIKE)*(1-μSTEP)
      )::Float64
  else
    return (
      i*μSPIKE*μSTEP - μSTEP*μSPIKE - μSTEP*(1-μSPIKE) - (N-i)*(1-μSTEP)*μSPIKE
      )::Float64
  end

  return 0.0 :: Float64
end
function GMStepSpike(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMStepSpike(i, j, IT, d.μ[STEP][i], d.σ[STEP][i], d.μ[SPIKE][j], d.σ[SPIKE][j])::Float64
end
function GMSpikeStep(i::Int, j::Int, d, IT)::Float64
    return GMStepSpike(j, i, IT, d.μ[STEP][j], d.σ[STEP][j], d.μ[SPIKE][i], d.σ[SPIKE][i])::Float64
end

# step x sine
function GMStepSin(
    t::Int, s::Int, N::Int, μSTEP::Float64, σSTEP::Float64,
    μSIN::Float64, σSIN::Float64,ω::Float64
    )::Float64

  return (
    μSTEP*μSIN*(t+1)-(1/2)*μSTEP*sin(ω)*cos((t+1)*ω)/(cos(ω)-1)
    +(1/2)*μSTEP*sin((t+1)*ω)-μSTEP*μSIN+(1/2)*μSTEP*sin(ω)*cos(ω)/(cos(ω)-1)
    -(1/2)*μSTEP*sin(ω)+(μSTEP*μSIN-μSIN)*(N+1)-(1/2)*(-1+μSTEP)*sin(ω)*cos((N+1)*ω)/(cos(ω)-1)
    +(-1/2+(1/2)*μSTEP)*sin((N+1)*ω)-(μSTEP*μSIN-μSIN)*(t+1)
    +(1/2)*(-1+μSTEP)*sin(ω)*cos((t+1)*ω)/(cos(ω)-1)-(-1/2+(1/2)*μSTEP)*sin((t+1)*ω)
    )::Float64

  #return GM::Float64
end
function GMStepSin(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMStepSin(i, j, IT, d.μ[STEP][i], d.σ[STEP][i], d.μ[SIN][j], d.σ[SIN][j],d.fs[j]) :: Float64
end
function GMSinStep(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMStepSin(j, i, IT, d.μ[STEP][j], d.σ[STEP][j], d.μ[SIN][i], d.σ[SIN][i],d.fs[i]) :: Float64
end

# step x cosine
function GMStepCos(
    t::Int, c::Int, N::Int, μSTEP::Float64, σSTEP::Float64,
    μCOS::Float64, σCOS::Float64,ω::Float64
    )::Float64

    return (
    μSTEP*μCOS*(t+1)+(1/2)*μSTEP*cos((t+1)*ω)
    +(1/2)*μSTEP*sin(ω)*sin((t+1)*ω)/(cos(ω)-1)
    -μSTEP*μCOS-(1/2)*μSTEP*cos(ω)-(1/2)*μSTEP*sin(ω)^2/(cos(ω)-1)
    +(μSTEP*μCOS-μCOS)*(N+1)+(-1/2+(1/2)*μSTEP)*cos((N+1)*ω)
    -(1/2)*(1-μSTEP)*sin(ω)*sin((N+1)*ω)/(cos(ω)-1)
    -(μSTEP*μCOS-μCOS)*(t+1)-(-1/2+(1/2)*μSTEP)*cos((t+1)*ω)
    +(1/2)*(1-μSTEP)*sin(ω)*sin((t+1)*ω)/(cos(ω)-1)
    )::Float64
end
function GMStepCos(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMStepCos(i, j, IT, d.μ[STEP][i], d.σ[STEP][i], d.μ[COS][j], d.σ[COS][j],d.fc[j]) :: Float64
end
function GMCosStep(i::Int, j::Int, d, IT)::Float64
    return GMStepCos(j, i, IT, d.μ[STEP][j], d.σ[STEP][j], d.μ[COS][i], d.σ[COS][i],d.fc[i]) :: Float64
end

# slope x slope
function GMSlopeSlope(m::Int, M::Int, N::Int, μm::Float64, σm::Float64, μM::Float64, σM::Float64)::Float64

  return (
    (1/6)*N - (1/6)*M + m*μM*μM - μM*(m+1)*μM - μM*(m+1)*m
    - (1/2)*(N+1)^2 + (1/3)*(N+1)^3 + (1/2)*(M+1)^2 - (1/3)*(M+1)^3
    + (1/2)*μM*(m+1)^2 - (1/2)*μM*(m+1) + (1/2)*m*(M+1)^2
    - (1/2)*m*(M+1) + (1/2)μM*(M+1)^2 - (1/2)*μM*(M+1)
    + (1/2)*M*(M+1)^2 - (1/2)*M*(M+1) - (1/2)*m*(N+1)^2 + (1/2)*m*(N+1)
    - (1/2)*μM*(N+1)^2 + (1/2)*μM*(N+1) - (1/2)*M*(N+1)^2
    + (1/2)*M*(N+1) - (1/2)*μM*(N+1)^2 - (M+1)*m*M
    - (M+1)*μM*M + (N+1)*m*M + μM*(N+1)*m + (N+1)*μM*M
    + μM*(N+1)*μM + (1/2)*μM*(N+1)
    ) :: Float64
end
function GMSlopeSlope(i::Int, j::Int, d::dataCD, IT::Int)::Float64
  	if j>i
    	return GMSlopeSlope(i, j, IT, d.μ[SLOPE][i], d.σ[SLOPE][i], d.μ[SLOPE][j], d.σ[SLOPE][j]) :: Float64
	else
		return GMSlopeSlope(j, i, IT, d.μ[SLOPE][j], d.σ[SLOPE][j], d.μ[SLOPE][i], d.σ[SLOPE][i]) :: Float64
	end
	return 0.0::Float64
end

# slope x spike
function GMSlopeSpike(
    l::Int, p::Int, N::Int, μSLOPE::Float64, σSLOPE::Float64,
                            μSPIKE::Float64, σSPIKE::Float64,
    )::Float64

  if p > l
    return (
      l*μSLOPE*μSPIKE + μSPIKE*(N+1)*l + μSPIKE*(N+1)*μSLOPE
      - (1/2)*μSPIKE*(N+1)^2 + (1/2)*μSPIKE*(N+1) - μSPIKE*(l+1)*l
      - μSPIKE*(l+1)*μSLOPE + (1/2)*μSPIKE*(l+1)^2
      - (1/2)*μSPIKE*(l+1) + (p-l-μSLOPE)*μSPIKE
      + (p-l-μSLOPE)*(1-μSPIKE)
      )::Float64
  else
    return (
      l*μSLOPE*μSPIKE + μSPIKE*(N+1)*l + μSPIKE*(N+1)*μSLOPE
      - (1/2)*μSPIKE*(N+1)^2 + (1/2)*μSPIKE*(N+1) - μSPIKE*(l+1)*l
      - μSPIKE*(l+1)*μSLOPE + (1/2)*μSPIKE*(l+1)^2
      - (1/2)*μSPIKE*(l+1) - μSLOPE*μSPIKE - μSLOPE*(1-μSPIKE)
      )::Float64
  end

  return 0.0::Float64
end
function GMSlopeSpike(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSlopeSpike(i, j, IT, d.μ[SLOPE][i], d.σ[SLOPE][i], d.μ[SPIKE][j], d.σ[SPIKE][j]) :: Float64
end
function GMSpikeSlope(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSlopeSpike(j, i, IT, d.μ[SLOPE][j], d.σ[SLOPE][j], d.μ[SPIKE][i], d.σ[SPIKE][i]) :: Float64
end

# slope x sine
function GMSlopeSin(
    l::Int, s::Int, N::Int, μSLOPE::Float64, σSLOPE::Float64,
    μSIN::Float64, σSIN::Float64,ω::Float64
    )::Float64

  return (
    μSLOPE*μSIN*(l+1)-(1/2)*μSLOPE*sin(ω)*cos((l+1)*ω)/(cos(ω)-1)
    +(1/2)*μSLOPE*sin((l+1)*ω)-μSLOPE*μSIN+(1/2)*μSLOPE*sin(ω)*cos(ω)/(cos(ω)-1)
    -(1/2)*μSLOPE*sin(ω)+(l*μSIN+(1/2)*μSIN+μSLOPE*μSIN)*(N+1)
    -(1/2)*μSIN*(N+1)^2+(1/2)*sin(ω)*(N+1)*cos((N+1)*ω)/(cos(ω)-1)-(1/2*(N+1))*sin((N+1)*ω)
    -(1/2)*(μSLOPE+l)*sin(ω)*cos((N+1)*ω)/(cos(ω)-1)
    +(1/2)*(l*cos(ω)+μSLOPE*cos(ω)-l-μSLOPE-1)*sin((N+1)*ω)/(cos(ω)-1)
    -(l*μSIN+(1/2)*μSIN+μSLOPE*μSIN)*(l+1)
    +(1/2)*μSIN*(l+1)^2-(1/2)*sin(ω)*(l+1)*cos((l+1)*ω)/(cos(ω)-1)
    +(1/2*(l+1))*sin((l+1)*ω)
    +(1/2)*(μSLOPE+l)*sin(ω)*cos((l+1)*ω)/(cos(ω)-1)
    -(1/2)*(l*cos(ω)+μSLOPE*cos(ω)-l-μSLOPE-1)*sin((l+1)*ω)/(cos(ω)-1)
    )::Float64
end
function GMSlopeSin(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSlopeSin(i, j, IT, d.μ[SLOPE][i], d.σ[SLOPE][i], d.μ[SIN][j], d.σ[SIN][j], d.fs[j]) :: Float64
end
function GMSinSlope(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSlopeSin(j, i, IT, d.μ[SLOPE][j], d.σ[SLOPE][j], d.μ[SIN][i], d.σ[SIN][i], d.fs[i]) :: Float64
end

# slope x cosine
@inline function GMSlopeCos(
    l::Int, c::Int, N::Int, μSLOPE::Float64, σSLOPE::Float64,
    μCOS::Float64, σCOS::Float64, ω::Float64
    )::Float64

  return (
    μSLOPE*μCOS*(l+1)+(1/2)*μSLOPE*cos((l+1)*ω)
    +(1/2)*μSLOPE*sin(ω)*sin((l+1)*ω)/(cos(ω)-1)
    -μSLOPE*μCOS-(1/2)*μSLOPE*cos(ω)-(1/2)*μSLOPE*sin(ω)^2/(cos(ω)-1)
    +(μSLOPE*μCOS+l*μCOS+(1/2)*μCOS)*(N+1)
    -(1/2)*μCOS*(N+1)^2-(1/2*(N+1))*cos((N+1)*ω)
    -(1/2)*sin(ω)*(N+1)*sin((N+1)*ω)/(cos(ω)-1)
    +(1/2)*(l*cos(ω)+μSLOPE*cos(ω)-l-μSLOPE-1)*cos((N+1)*ω)/(cos(ω)-1)
    +(1/2)*(l+μSLOPE)*sin(ω)*sin((N+1)*ω)/(cos(ω)-1)
    -(μSLOPE*μCOS+l*μCOS+(1/2)*μCOS)*(l+1)
    +(1/2)*μCOS*(l+1)^2+(1/2*(l+1))*cos((l+1)*ω)
    +(1/2)*sin(ω)*(l+1)*sin((l+1)*ω)/(cos(ω)-1)
    -(1/2)*(l*cos(ω)+μSLOPE*cos(ω)-l-μSLOPE-1)*cos((l+1)*ω)/(cos(ω)-1)
    -(1/2)*(l+μSLOPE)*sin(ω)*sin((l+1)*ω)/(cos(ω)-1)
    )::Float64
end
function GMSlopeCos(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSlopeCos(i, j, IT, d.μ[SLOPE][i], d.σ[SLOPE][i], d.μ[COS][j], d.σ[COS][j], d.fc[j])::Float64
end
function GMCosSlope(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSlopeCos(j, i, IT, d.μ[SLOPE][j], d.σ[SLOPE][j], d.μ[COS][i], d.σ[COS][i], d.fc[i])::Float64
end

# spike x spike
@inline function GMSpikeSpike(i::Int, j::Int, N::Int, μSPIKEi::Float64, σSPIKEi::Float64)::Float64

  if i != j
    return (
      (N-2)*μSPIKEi^2 - 2*μSPIKEi*(1-μSPIKEi)
    )::Float64
  else
    return float(N)::Float64
  end

  return 0.0::Float64
end
function GMSpikeSpike(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSpikeSpike(i, j, IT, d.μ[SPIKE][i], d.σ[SPIKE][i])::Float64
end

# spike x sine
@inline function GMSpikeSin(
    p::Int, s::Int, N::Int, μSPIKE::Float64, σSPIKE::Float64,
    μSIN::Float64, σSIN::Float64, ω::Float64
    )::Float64

   return (
    μSPIKE*μSIN*(N+1)-(1/2)*μSPIKE*sin(ω)*cos((N+1)*ω)/(cos(ω)-1)
    +(1/2)*μSPIKE*sin((N+1)*ω)-μSPIKE*μSIN+(1/2)*μSPIKE*sin(ω)*cos(f[s])/(cos(f[s])-1)
    -(1/2)*μSPIKE*sin(ω)+μSPIKE*(sin(p*ω)-μSIN)+(1-μSPIKE)*(sin(p*ω)-μSIN)
    ) :: Float64
end
function GMSpikeSin(i::Int, j::Int, d::dataCD, IT::Int)
    return GMSpikeSin(i, j, IT, d.μ[SPIKE][i], d.σ[SPIKE][i], d.μ[SIN][j], d.σ[SIN][j], d.fs[j]) :: Float64
end
function GMSinSpike(i::Int, j::Int, d::dataCD, IT::Int)
    return GMSpikeSin(j, i, IT, d.μ[SPIKE][j], d.σ[SPIKE][j], d.μ[SIN][i], d.σ[SIN][i], d.fs[i]) :: Float64
end

# spike x cosine
@inline function GMSpikeCos(
    p::Int, c::Int, N::Int, μSPIKE::Float64, σSPIKE::Float64,
    μCOS::Float64, σCOS::Float64, ω::Float64
    )

  return (
    μSPIKE*μCOS*(N+1)+(1/2)*μSPIKE*cos((N+1)*ω)
    +(1/2)*μSPIKE*sin(ω)*sin((N+1)*ω)/(cos(ω)-1)
    -μSPIKE*μCOS-(1/2)*μSPIKE*cos(ω)-(1/2)*μSPIKE*sin(ω)^2/(cos(ω)-1)
    +μSPIKE*(cos(p*ω)-μCOS)+(1-μSPIKE)*(cos(p*ω)-μCOS)
    )::Float64
end
function GMSpikeCos(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSpikeCos(i, j, IT, d.μ[SPIKE][i], d.σ[SPIKE][i], d.μ[COS][j], d.σ[COS][j],d.fc[j]) :: Float64
end
function GMCosSpike(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSpikeCos(j, i, IT, d.μ[SPIKE][j], d.σ[SPIKE][j], d.μ[COS][i], d.σ[COS][i],d.fc[i]) :: Float64
end

# sine x sine
@inline function GMSinSin(i::Int, j::Int, N::Int, μSINi::Float64, σSINi::Float64, μSINj::Float64, σSINj::Float64, ωi::Float64, ωj::Float64)::Float64

  #ATENTION HERE
  #there still may be a ptoblem from cos(ωi) = cos(ωj) with i != j
  if i == j
    return float(N)::Float64
  elseif abs( cos(ωj)-cos(ωi) ) <= 1e-8
    #cancelation that need to happen for finiteness
    return (
      μSINj*μSINi*(N+1)
      -(1/2)*sin((N+1)*ωj)*sin((N+1)*ωi)-(1/2)*μSINj*sin(ωi)*cos((N+1)*ωi)/(-1+cos(ωi))
      -(1/2)*μSINi*sin(ωj)*cos((N+1)*ωj)/(cos(ωj)-1)+(1/2)*μSINj*sin((N+1)*ωi)
      +(1/2)*sin((N+1)*ωj)*μSINi-μSINj*μSINi
      +(1/2)*sin(ωi)*sin(ωj)
      +(1/2)*μSINj*sin(ωi)*cos(ωi)/(-1+cos(ωi))+(1/2)*μSINi*sin(ωj)*cos(ωj)/(cos(ωj)-1)
      -(1/2)*μSINj*sin(ωi)-(1/2)*μSINi*sin(ωj)
      )::Float64
  else
    return (
      μSINj*μSINi*(N+1)+(1/2)*sin(ωj)*cos((N+1)*ωj)*sin((N+1)*ωi)/(cos(ωj)-cos(ωi))
      -(1/2)*sin(ωi)*sin((N+1)*ωj)*cos((N+1)*ωi)/(cos(ωj)-cos(ωi))
      -(1/2)*sin((N+1)*ωj)*sin((N+1)*ωi)-(1/2)*μSINj*sin(ωi)*cos((N+1)*ωi)/(-1+cos(ωi))
      -(1/2)*μSINi*sin(ωj)*cos((N+1)*ωj)/(cos(ωj)-1)+(1/2)*μSINj*sin((N+1)*ωi)
      +(1/2)*sin((N+1)*ωj)*μSINi-μSINj*μSINi-(1/2)*sin(ωj)*cos(ωj)*sin(ωi)/(cos(ωj)-cos(ωi))
      +(1/2)*sin(ωi)*sin(ωj)*cos(ωi)/(cos(ωj)-cos(ωi))+(1/2)*sin(ωi)*sin(ωj)
      +(1/2)*μSINj*sin(ωi)*cos(ωi)/(-1+cos(ωi))+(1/2)*μSINi*sin(ωj)*cos(ωj)/(cos(ωj)-1)
      -(1/2)*μSINj*sin(ωi)-(1/2)*μSINi*sin(ωj)
      )::Float64
  end

  return 0.0::Float64
end
function GMSinSin(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMSinSin(i, j, IT, d.μ[SIN][i], d.σ[SIN][i], d.μ[SIN][j], d.σ[SIN][j], d.fs[i], d.fs[j]) :: Float64
end

# sine x cosine
@inline function GMSinCos(
    s::Int, c::Int, N::Int, μSIN::Float64, σSIN::Float64,
    μCOS::Float64, σCOS::Float64, ωs::Float64, ωc::Float64
    )::Float64

  if abs(cos(ωc)-cos(ωs)) >= 1e-5

    return (
      μSIN*μCOS*(N+1)+(1/2)*sin(ωs)*cos((N+1)*ωs)*cos((N+1)*ωc)/(cos(ωs)-cos(ωc))
      -(1/2)*sin((N+1)*ωs)*cos((N+1)*ωc)+(1/2)*sin(ωc)*sin((N+1)*ωs)*sin((N+1)*ωc)/(cos(ωs)-cos(ωc))
      +(1/2)*μSIN*cos((N+1)*ωc)-(1/2)*μCOS*sin(ωs)*cos((N+1)*ωs)/(cos(ωs)-1)
      +(1/2)*μSIN*sin(ωc)*sin((N+1)*ωc)/(cos(ωc)-1)+(1/2)*sin((N+1)*ωs)*μCOS-μSIN*μCOS
      -(1/2)*sin(ωs)*cos(ωs)*cos(ωc)/(cos(ωs)-cos(ωc))+(1/2)*sin(ωs)*cos(ωc)
      -(1/2)*sin(ωc)^2*sin(ωs)/(cos(ωs)-cos(ωc))-(1/2)*μSIN*cos(ωc)+(1/2)*μCOS*sin(ωs)*cos(ωs)/(cos(ωs)-1)
      -(1/2)*μSIN*sin(ωc)^2/(cos(ωc)-1)-(1/2)*sin(ωs)*μCOS
      )::Float64

  else
    return (
      μSIN*μCOS*(N+1)-(1/2)*cos(ωs)*cos((N+1)*ωs)^2/sin(ωs)-(1/2)*sin((N+1)*ωs)*cos((N+1)*ωs)
      +(1/2)*(cos(ωs)*μCOS+μSIN*sin(ωs)+μCOS)*cos((N+1)*ωs)/sin(ωs)
      +(1/2)*(-cos(ωs)*μSIN+μCOS*sin(ωs)-μSIN)*sin((N+1)*ωs)/sin(ωs)-μSIN*μCOS
      +(1/2)*cos(ωs)^3/sin(ωs)+(1/2)*cos(ωs)*sin(ωs)-(1/2)*(cos(ωs)*μCOS+μSIN*sin(ωs)+μCOS)*cos(ωs)/sin(ωs)
      +(1/2)*cos(ωs)*μSIN-(1/2)*μCOS*sin(ωs)+(1/2)*μSIN
      )::Float64

  end

  return 0.0::Float64
end
function GMSinCos(i::Int, j::Int, d::dataCD, IT::Int)
    return GMSinCos(i, j, IT, d.μ[SIN][i], d.σ[SIN][i], d.μ[COS][j], d.σ[COS][j],d.fs[i],d.fc[j]) :: Float64
end
function GMCosSin(i::Int, j::Int, d::dataCD, IT::Int)
    return GMSinCos(j, i, IT, d.μ[SIN][j], d.σ[SIN][j], d.μ[COS][i], d.σ[COS][i],d.fs[j],d.fc[i]) :: Float64
end

# cosine x cosine
@inline function GMCosCos(i::Int, j::Int, N::Int, μCOSi::Float64, σCOSi::Float64, μCOSj::Float64, σCOSj::Float64, ωi::Float64, ωj::Float64)::Float64

  if i == j
    return float(N)::Float64
  elseif abs(cos(ωi)-cos(ωj) ) >= 1e-8
    return (
      μCOSj*μCOSi*(N+1)-(1/2)*cos((N+1)*ωj)*cos((N+1)*ωi)
      +(1/2)*sin(ωi)*cos((N+1)*ωj)*sin((N+1)*ωi)/(cos(ωj)-cos(ωi))
      -(1/2)*sin(ωj)*sin((N+1)*ωj)*cos((N+1)*ωi)/(cos(ωj)-cos(ωi))
      +(1/2)*μCOSj*cos((N+1)*ωi)+(1/2)*cos((N+1)*ωj)*μCOSi
      +(1/2)*μCOSj*sin(ωi)*sin((N+1)*ωi)/(-1+cos(ωi))
      +(1/2)*μCOSi*sin(ωj)*sin((N+1)*ωj)/(cos(ωj)-1)-μCOSj*μCOSi+(1/2)*cos(ωj)*cos(ωi)
      -(1/2)*sin(ωi)^2*cos(ωj)/(cos(ωj)-cos(ωi))
      +(1/2)*sin(ωj)^2*cos(ωi)/(cos(ωj)-cos(ωi))-(1/2)*μCOSj*cos(ωi)
      -(1/2)*cos(ωj)*μCOSi-(1/2)*μCOSj*sin(ωi)^2/(-1+cos(ωi))-(1/2)*μCOSi*sin(ωj)^2/(cos(ωj)-1)
      )::Float64
  else
    return (
      μCOSj*μCOSi*(N+1)-(1/2)*cos((N+1)*ωj)*cos((N+1)*ωi)
      +(1/2)*μCOSj*cos((N+1)*ωi)+(1/2)*cos((N+1)*ωj)*μCOSi
      +(1/2)*μCOSj*sin(ωi)*sin((N+1)*ωi)/(-1+cos(ωi))
      +(1/2)*μCOSi*sin(ωj)*sin((N+1)*ωj)/(cos(ωj)-1)
      -μCOSj*μCOSi+(1/2)*cos(ωj)*cos(ωi)
      -(1/2)*μCOSj*cos(ωi)-(1/2)*cos(ωj)*μCOSi
      -(1/2)*μCOSj*sin(ωi)^2/(-1+cos(ωi))
      -(1/2)*μCOSi*sin(ωj)^2/(cos(ωj)-1)
      )::Float64
  end
  return 0.0::Float64
end
function GMCosCos(i::Int, j::Int, d::dataCD, IT::Int)::Float64
    return GMCosCos(i, j, IT, d.μ[COS][i], d.σ[COS][i],d.μ[COS][j], d.σ[COS][j],d.fc[i],d.fc[j])::Float64
end

# GM = Matrix{Function}(TOTALCOMPONENTS, TOTALCOMPONENTS)
# GM[STEP, STEP]   = GMStepStep
# GM[STEP, SPIKE]  = GMStepSpike
# GM[SPIKE, STEP]  = GMSpikeStep
# GM[STEP, SLOPE]  = GMStepSlope
# GM[SLOPE, STEP]  = GMSlopeStep
# GM[STEP, SIN]    = GMStepSin
# GM[SIN, STEP]    = GMSinStep
# GM[STEP, COS]    = GMStepCos
# GM[COS, STEP]    = GMCosStep
# GM[SPIKE, SPIKE] = GMSpikeSpike
# GM[SPIKE, SLOPE] = GMSpikeSlope
# GM[SLOPE, SPIKE] = GMSlopeSpike
# GM[SPIKE, SIN]   = GMSpikeSin
# GM[SIN, SPIKE]   = GMSinSpike
# GM[SPIKE, COS]   = GMSpikeCos
# GM[COS, SPIKE]   = GMCosSpike
# GM[SLOPE, SLOPE] = GMSlopeSlope
# GM[SLOPE, SIN]   = GMSlopeSin
# GM[SIN, SLOPE]   = GMSinSlope
# GM[SLOPE, COS]   = GMSlopeCos
# GM[COS, SLOPE]   = GMCosSlope
# GM[SIN, SIN]     = GMSinSin
# GM[SIN, COS]     = GMSinCos
# GM[COS, SIN]     = GMCosSin
# GM[COS, COS]     = GMCosCos

function GM2(c1,c2,i::Int, j::Int, d::dataCD, IT::Int)::Float64
  if c1 == STEP && c2 == STEP
    return GMStepStep(i::Int, j::Int, d, IT)::Float64
  elseif c1 == STEP && c2 == SPIKE
    return GMStepSpike(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SPIKE && c2 == STEP
    return GMSpikeStep(i::Int, j::Int, d, IT)::Float64
  elseif c1 == STEP && c2 == SLOPE
    return GMStepSlope(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SLOPE && c2 == STEP
    return GMSlopeStep(i::Int, j::Int, d, IT)::Float64
  elseif c1 == STEP && c2 == SIN
    return GMStepSin(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SIN && c2 == STEP
    return GMSinStep(i::Int, j::Int, d, IT)::Float64
  elseif c1 == STEP && c2 == COS
    return GMStepCos(i::Int, j::Int, d, IT)::Float64
  elseif c1 == COS && c2 == STEP
    return GMCosStep(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SPIKE && c2 == SPIKE
    return GMSpikeSpike(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SPIKE && c2 == SLOPE
    return GMSpikeSlope(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SLOPE && c2 == SPIKE
    return GMSlopeSpike(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SPIKE && c2 == SIN
    return GMSpikeSin(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SIN && c2 == SPIKE
    return GMSinSpike(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SPIKE && c2 == COS
    return GMSpikeCos(i::Int, j::Int, d, IT)::Float64
  elseif c1 == COS && c2 == SPIKE
    return GMCosSpike(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SLOPE && c2 == SLOPE
    return GMSlopeSlope(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SLOPE && c2 == SIN
    return GMSlopeSin(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SIN && c2 == SLOPE
    return GMSinSlope(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SLOPE && c2 == COS
    return GMSlopeCos(i::Int, j::Int, d, IT)::Float64
  elseif c1 == COS && c2 == SLOPE
    return GMCosSlope(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SIN && c2 == SIN
    return GMSinSin(i::Int, j::Int, d, IT)::Float64
  elseif c1 == SIN && c2 == COS
    return GMSinCos(i::Int, j::Int, d, IT)::Float64
  elseif c1 == COS && c2 == SIN
    return GMCosSin(i::Int, j::Int, d, IT)::Float64
  elseif c1 == COS && c2 == COS
    return GMCosCos(i::Int, j::Int, d, IT)::Float64
  else
    return 0.0::Float64
  end
end


########## reward ##########
#POMDPs.reward(DP::DrivePOMDP, Ss::Sts, acc_ego::Float64) = Reward(DP, Ss, acc_ego)
POMDPs.reward(DP::DrivePOMDP, Ss::Sts, Aego::Symbol, Ssnext::Sts) = Reward(DP, Ss, Aego, Ssnext)

# based directly on next state; used in DiscreteValueIteration Solver
function R_crash(DP::DrivePOMDP, Ssnext::Sts, acc_ego::Float64)
    Ego_nextPos = GetGlobalPosition(DP.Routes[Ssnext.Ego.r], min(Ssnext.Ego.s, DP.sV[end]))
    Other_nextPos = GetGlobalPosition(DP.Routes[Ssnext.Other.r], min(Ssnext.Other.s, DP.sV[end]))
    if haskey(DP.Routes[Ssnext.Ego.r].intersect_Infos, Ssnext.Other.r) # if two routes intesect with each other
        if Distance2D(Ego_nextPos, Other_nextPos) < 2 # crash happens
            return -100.0 - 1.5*Ssnext.Ego.v^2 - 2.6*acc_ego*abs(acc_ego)
        end
    end
    return 0.0
end

# include the situation that junction point is far away from every state-point but collision actually happens
function R_crash(DP::DrivePOMDP, Ss::Sts, Ssnext::Sts)
    if haskey(DP.Routes[Ss.Ego.r].intersect_Infos, Ss.Other.r) # if two routes intesect with each other
        for (i, PtDist) in enumerate(DP.Routes[Ss.Ego.r].intersect_Infos[Ss.Other.r])
            InterDistEgo = PtDist[3] # distance from start point to junction point
            InterDistOther = DP.Routes[Ss.Other.r].intersect_Infos[Ss.Ego.r][i][3]
            if Ss.Ego.s < InterDistEgo < Ssnext.Ego.s && Ss.Other.s < InterDistOther < Ssnext.Other.s
                return -150.0
            end
        end
    end
    return 0.0
end

"""
if Vego != Vref, R_v should always be less than R_acc, so that Ego will get a better reward if it chooses to accelerate/decelerate itself.
R_v and R_acc should have a relation between each other.
"""

function R_v(DP::DrivePOMDP, Ss::Sts, acc_ego::Float64)
    Kvp = 1.5
    Kvn = 1.0
    Vref = DP.Routes[Ss.Ego.r].Vref[UInt16(Ss.Ego.s/DP.Δs)+1] # m/s, should be defined as array according to Route/Geometries
    Vego = Ss.Ego.v + acc_ego*DP.Δt
    if Vego > Vref
        #return -Kvp*(Vref - Vego)^2
        return -Kvp*((Vref - Vego)*3.6)^2
    else
        return -Kvn*(Vref - Vego)*3.6
    end
end

function R_acc(DP::DrivePOMDP, Ss::Sts, acc_ego::Float64) #
    Factor = -0.77
    Δa = acc_ego - DP.Routes[Ss.Ego.r].Aref[UInt16(Ss.Ego.s/DP.Δs)+1]
    return Factor*(Δa^2+2.3*abs(Δa))
    #return Factor*(acc_ego^2+2.3*abs(acc_ego))
end

function Reward(DP::DrivePOMDP, Ss::Sts, Aego::Symbol, Ssnext::Sts)
    acc_ego = 0.0
    if Ssnext.Other.s != Ss.Other.s
        acc_k = (Ssnext.Other.v^2 - Ss.Other.v^2)/(2*(Ssnext.Other.s - Ss.Other.s))
        Ego, Other, accs = IDMtransit(DP, Ss, Aego, acc_k)
        if Ego.s != Ss.Ego.s
            acc_ego = (Ego.v^2 - Ss.Ego.v^2)/(2*(Ego.s - Ss.Ego.s))
        end
        #acc_ego = (Ssnext.Ego.v^2 - Ss.Ego.v^2)/(2*(Ssnext.Ego.s - Ss.Ego.s))
    end
    reward = R_crash(DP, Ss, Ssnext) + R_crash(DP, Ss, acc_ego) + R_crash(DP, Ssnext, acc_ego) + R_v(DP, Ss, acc_ego) + R_acc(DP, Ss, acc_ego)
    reward = trunc(reward; digits=4)
    #println("## Realtime Reward:$reward")
    return reward
end


__precompile__(true)

module Simplices
using LinearAlgebra
using Printf

using Distributions

include("simplexsplitting/SimplexSplitting.jl")
include("barycentric-coordinates.jl")
include("Binary.jl")
include("Circumsphere.jl")
include("complementary.jl")
include("embed.jl")
include("geometry.jl")
include("heaviside.jl")
include("intersection-of-boundaries.jl")
include("intersection-of-boundaries-loop.jl")
include("init_functions.jl")
include("NullSpace.jl")
include("polytope-generating-vertices.jl")
include("QR.jl")
include("shared-vertices.jl")
include("sharing-a-face.jl")
include("some-vertex-in-circumsphere.jl")
include("simplexoperations.jl")
include("simplex_sampling.jl")
include("TriangulationPolytopeFaces.jl")
include("TriangulationNonSimplicialFaces.jl")
include("Update.jl")
include("volume-computation.jl")

include("simplexintersect.jl")
include("intersecting_vertices.jl")

export BarycentricCoordinates,
        centroid,
        Circumsphere,
        childsimplex,
        childpoint,
        childpoints,
        Delaunay,
        even_sampling, even_sampling_rules, evenly_sample,
        heaviside,
        heaviside0,
	insidepoints,
        intersecting_simplices,
        intersectingvertices,
        IntersectionOfBoundaries,
        IntersectionOfBoundaries_NoStorage,
        issingular,
        nontrivially_intersecting_simplices,
        nontrivial_intersection,
        outsidepoint,
        orientation,
        PolytopeGeneratingVertices,
        radius,
        SharedVertices,
        SharedFaceVolume,
        SharedFaceVertices,
        shared_vertex_intersection,
	SimplexSplitting,
        simplices_sharing_vertices,
        simplexintersection,
        SomeVertexInCircumsphere,
        tensordecomp,
        volume,
        VolumeComputation,
	outsidepoint,
	outsidepoints

end


function update_weights(data, e, same, other, weights, weights_mult, neigh_weights_same, neigh_weights_other, m, max_f_vals, min_f_vals)

	for t = 1:size(data, 2)

		# Penalty term
		penalty = sum(neigh_weights_same.*(abs.(e[t] .- same[:, t])/(max_f_vals[t] .- min_f_vals[t] .+ eps(Float64))))

		# Reward term
		reward = sum(neigh_weights_other.*(weights_mult .* (abs.(e[t] .- other[:, t])/(max_f_vals[t] .- min_f_vals[t] .+ eps(Float64)))))

		# Weights update
		weights[t] = weights[t] - penalty/(m*size(same, 1) + eps(Float64)) +
			reward/(m*size(other, 1) + eps(Float64))
	end

	# Return updated weights.
	return weights
end


module jPCA

export fit, JPCA

using MultivariateStats,Optim

include("reshape_skew.jl")
include("getRealVs.jl")

import MultivariateStats.fit

type JPCA
    pca
    M
    Mskew
    proj
end

"""
    fit(jpca::Type{JPCA}, X, k=10 or num features)
NB: X should be a time X features array.
"""
function fit(jpca::Type{JPCA}, X, k=size(X,2) < 10 ? size(X,2) : 10)
    size(X,2) > size(X,1) && error("Must have more rows than columns")
    k > size(X,2) && error("k must be <= the number of columns in X")
	isodd(k) && error("k must be even")

    pca   = fit(PCA, X, pratio=1., maxoutdim=k)
    X2    = pca.proj[:,1:k]

    T     = (1:size(X2,1))[2:end]
    Tpre  = (1:size(X2,1))[1:end-1]
    dX2 = X2[T,:].-X2[Tpre,:]
    M = dX2/X2[Tpre,:]

    Mskew = skewsymregress(dX2, X2[Tpre,:])
    D,V   = eig(Mskew)
    V     .= V[:,sortperm(abs.(D))]

    jPCs = zeros(size(V)...)
    for pair in 1:div(k,2)
		vi1 = 1+2(pair-1)
		vi2 = 2pair

		VconjPair  = V[:, [vi1,vi2]]
		evConjPair = D[[vi1, vi2]]
		jPCs[:,vi1] .= getRealVs(VconjPair, evConjPair)[1]
		jPCs[:,vi2] .= getRealVs(VconjPair, evConjPair)[2]
    end
	proj = pca.proj * jPCs

    JPCA(pca, M, Mskew, proj)
end

function skewsymregress(dX,X)
    M0 = X\dX
    M0 .= 0.5*(M0 - M0')
    m0 = reshape_skew(M0)

    function skewsymderiv(m)
        # Evaluate objective function
        f = vecnorm( dX - X*reshape_skew(m))^2;

        # Evaluate derivative
        # D = (dX - X*reshape_skew(m))'*X;
        # df = 2*reshape_skew(D - D')
        f
    end

    M = optimize(skewsymderiv, m0)
    return reshape_skew(M.minimizer)
end

end


include("../../external/mbh2dfmm/julia/modbhrouts.jl")

using FINUFFT

## FREESPACE SOLVER WITH NUFFT
"
fs_modstokes(F1, F2, L, lambda, xnu, ynu; precomp=GR)

Solves modified Stokes equation on NxN grid of length L with rhs with
freespace boundary conditions and compactly supported rhs (F1, F2).

GR is precomputed truncated Green's function to modified biharmonic:
GR = fs_modbh_precomp(N, L, lambda)
"
function fs_modstokes(F1, F2, L, lambda, xnu=[], ynu=[]; precomp=[])
    # Check input
    M = size(F1)
    @assert length(M)>1 && M[1]==M[2] "Grid must be square"
    @assert size(F1)==size(F2)
    N = M[1]
    # Get truncated Green's function
    if isempty(precomp)
        GRpre = fs_modbh_precomp(N, L, lambda)
    else
        GRpre = precomp
    end
    # Pad input to double size
    sf = 2
    N2 = sf*N
    info("FFT size is $(N2)x$(N2)")
    Fpad1 = zeros(Complex128, N2, N2)
    Fpad2 = zeros(Complex128, N2, N2)
    Fpad1[1:N, 1:N] = F1
    Fpad2[1:N, 1:N] = F2
    # Do FFTs in-place
    fft!(Fpad1)
    fft!(Fpad2)
    Fhat1, Fhat2 = Fpad1, Fpad2
    # Convolve with truncated greens function
    Uhat1, Uhat2 = Fhat1, Fhat2 # Operate in-place
    k1, k2 = k_vectors([N2, N2], [L*sf, L*sf])
    k1 = ifftshift(k1)
    k2 = ifftshift(k2)
    k1sq, k2sq = k1.^2, k2.^2
    for j=1:N2
        for i=1:N2
            K1 = k1[i]
            K2 = k2[j]
            Ksq = k1sq[i] + k2sq[j]
            KdotFhat = K1*Fhat1[i,j] + K2*Fhat2[i,j]
            Uhat1[i,j] = GRpre[i,j]*(Ksq*Fhat1[i,j] - K1*KdotFhat)
            Uhat2[i,j] = GRpre[i,j]*(Ksq*Fhat2[i,j] - K2*KdotFhat)
        end
    end
    if length(xnu) > 0
        # NUFFT for off-grid points
        @assert length(xnu)==length(ynu)
        # TODO: Get this scaling right so we can use modeord=0
        origin = pi/sf
        scale = 2*pi/(sf*L)
        xn = origin + xnu*scale
        yn = origin + ynu*scale
        Nnu = length(xnu)
        unu1 = Array{Complex128}(Nnu)
        unu2 = Array{Complex128}(Nnu)
        opts = finufft_default_opts()
        opts.modeord = 1
        nufft2d2!(xn, yn, unu1, 1, 1e-15, Uhat1, opts)
        nufft2d2!(xn, yn, unu2, 1, 1e-15, Uhat2, opts)
        uscale = 1/(N*sf)^2
        unu1 = real(unu1)*uscale
        unu2 = real(unu2)*uscale
    else
        unu1, unu2 = [], []
    end
    # Transform back
    ifft!(Uhat1)
    ifft!(Uhat2)
    # Truncate results to original grid size
    U1 = zeros(N,N)
    U2 = zeros(N,N)
    for j=1:N
        for i=1:N
            U1[i,j] = real(Uhat1[i,j])
            U2[i,j] = real(Uhat2[i,j])
        end
    end
    return U1, U2, unu1, unu2
end

"Precompute truncated Green's function to modified biharmonic equation"
function fs_modbh_precomp(N, L, lambda)
    sf = 1+sqrt(2) # Minimum oversampling in 2D
    Nf = ceil(Int, sf*N)
    if mod(Nf,2) != 0
        Nf+=1 # Even grids are always faster (right?)
    end
    info("Precompute FFT size is $(Nf)x$(Nf)")
    sf = Nf/N # Effective oversampling
    R = sqrt(2)*L
    logR = log(R)
    rscale = 1.0
    n = 1
    qn = modbh_qn([R],lambda,rscale,n)
    Q0 = qn[1,1]
    Q1 = qn[2,1]
    lam2 = lambda*lambda
    # Compute
    GR = Array{Complex128}(Nf, Nf)
    k1, k2 = k_vectors([Nf, Nf], [L*sf, L*sf])
    k1 = ifftshift(k1)
    k2 = ifftshift(k2)
    k1sq, k2sq = k1.^2, k2.^2
    for j=1:Nf
        for i=1:Nf
            K1 = k1[i]
            K2 = k2[j]
            Ksq = k1sq[i] + k2sq[j]
            K = sqrt(Ksq)
            J1 = besselj(1, K*R)
            J0M1 = besselj0m1(K*R)
            J0 = J0M1+1
            GR[i,j] = lam2*(J0M1 + K*R*logR*J1)/(Ksq*(lam2+Ksq))
            GR[i,j] += (Q0*K*R*J1-lambda*R*Q1*J0)/(lam2+Ksq)
        end
    end
    # Correct at k=0
    mid = 1 # grid is even
    GR[mid, mid] = R*R*(-0.25+0.5*logR)- R*Q1/lambda
    GR .*= 1/lam2
    ifft!(GR)
    # GR now in real space and has rubbish in the center,
    # Truncate by picking out corner values
    GRtrunc = Array{Complex128}(2*N, 2*N)
    idx = 1:N
    offset = Nf - N
    for i=idx
        for j=idx
            GRtrunc[i,  j  ] = GR[i,       j       ]
            GRtrunc[i+N,j  ] = GR[i+offset,j       ]
            GRtrunc[i,  j+N] = GR[i,       j+offset]
            GRtrunc[i+N,j+N] = GR[i+offset,j+offset]
        end
    end
    # Return Fourier transform
    fft!(GRtrunc)
    return GRtrunc
end

## PJ FUNCTIONS
function GR_purejulia(K, R)
    J0 = besselj(0, K*R)
    J1 = besselj(1, K*R)
    K0 = besselk(0, lambda*R)
    K1 = besselk(1, lambda*R)
    GR = -( J0*(lambda*K^2*R*K1 - K^2 - lambda^2) -
            K*R*J1*((K^2+lambda^2)*log(R) + K^2*K0) + lambda^2
            ) / (K^2*(lambda^2+K^2))
    return GR
end

function GR0_purejulia(R)
    GR = -(lambda^2*R^2 - 2*lambda^2*R^2*log(R) +
           4*lambda*R*besselk(1, R*lambda) - 4)/(4*lambda^2)
    return GR
end

## HELPERS
function k_vectors(M,box)
    function k_vec(M,L)
        if mod(M,2)==0
            MM = M/2;
            k = (2*pi/L)*collect(-MM:(MM-1));
        elseif mod(M-1,2)==0
            MM = (M-1)/2;
            k = (2*pi/L)*collect(-MM:MM);
        else
            error("k-vectors not computed");
        end
        return k
    end
    k1 = k_vec(M[1], box[1])
    k2 = k_vec(M[2], box[2])
    return k1, k2
end


function steadystate!(m::HetDSGEGovDebt;
                      βlo::S = 0.5*exp(m[:γ].scaledvalue)/(1 + m[:r].scaledvalue),
                      βhi::S = exp(m[:γ].scaledvalue)/(1 + m[:r].scaledvalue),
                      excess::S = 5000.,
                      tol::S = 1e-4,
                      maxit::Int64 = 20,
                      βband::S = 1e-2) where {S<:AbstractFloat}
    # If we have already solved for βstar (i.e. it's not NaN) and we only want to
    # estimate the non steady state parameters, there's no need to recompute
    # zlo/zhi, etc.
    if !isnan(m[:βstar].value) && get_setting(m, :estimate_only_non_steady_state_parameters)
        return
    else
        reset_grids!(m)
        nx = get_setting(m, :nx)
        ns = get_setting(m, :ns)
        ni = get_setting(m, :ni)
        nz = get_setting(m, :nz)

        # Determines whether random or not
        us, zs = if get_setting(m, :fix_random_matrices)
            get_setting(m, :us), get_setting(m, :zs)
        else
            generate_us_and_zs(ni, nz)
        end

        # This is a test setting; TODO: Remove because now we can just fix the randomness
        if get_setting(m, :steady_state_only)
            find_steadystate!(m; βlo = βlo, βhi = βhi, excess = excess, tol = tol, maxit = maxit,
                              βband = βband)
            return
        else
            # Do you want to calibrate for matching the income moments?
            if get_setting(m, :calibrate_income_targets)

                target = get_setting(m, :calibration_targets)
                lower  = get_setting(m, :calibration_targets_lb)
                upper  = get_setting(m, :calibration_targets_ub)
                m[:sH_over_sL], m[:zlo], _, _ = best_fit(m[:pLH].value, m[:pHL].value,
                                                         target, lower, upper, us, zs)
            else
                m[:varlinc], m[:vardlinc] = skill_moments(m[:sH_over_sL].value, m[:zlo].value,
                                                          m[:pLH].value, m[:pHL].value,
                                                          us, zs, ni)
            end
            # Whether one calibrates or not, zhi needs to be updated here
            m[:zhi] = 2.0 - m[:zlo].value

            # Construct Markov transition matrix for skill
            f, sgrid, swts, sscale = persistent_skill_process(m[:sH_over_sL].value, m[:pLH].value,
                                                              m[:pHL].value, get_setting(m, :ns))
            m.grids[:sgrid] = Grid(sgrid, swts, sscale)
            m.grids[:fgrid] = f

            xgrid, xwts, xlo, xhi, xscale = cash_grid(sgrid, m[:ωstar].value, m[:H].value,
                                                      m[:r].scaledvalue, m[:η].value,
                                                      m[:γ].scaledvalue,
                                                      m[:Tstar].value, m[:zlo].value, nx)

            m.grids[:xgrid] = Grid(uniform_quadrature(xscale), xlo, xhi, nx, scale = xscale)
            m <= Setting(:xlo, xlo)
            m <= Setting(:xhi, xhi)
            m <= Setting(:xscale, xscale)
            xswts = kron(swts, xwts)
            m.grids[:weights_total] = xswts

            # Once have updated grids, can call steady state and compute other two moments
            find_steadystate!(m; βlo = βlo, βhi = βhi,
                              excess = excess, tol = tol, maxit = maxit,
                              βband = βband)
            m[:mpc] = ave_mpc(m[:μstar].value,   m[:cstar].value, xgrid, xswts, nx, ns)
            m[:pc0] = frac_zero(m[:μstar].value, m[:cstar].value, xgrid, xswts, ns)
        end
    end
end

function find_steadystate!(m::HetDSGEGovDebt;
                           βlo::S = 0.5*exp(m[:γ].scaledvalue)/(1 + m[:r].scaledvalue),
                           βhi::S = exp(m[:γ].scaledvalue)/(1 + m[:r].scaledvalue),
                           excess::S = 5000.,
                           tol::S = 1e-4,
                           maxit::Int64 = 20,
                           βband::S = 1e-2) where {S<:AbstractFloat}
    # Load settings
    nx = get_setting(m, :nx)
    ns = get_setting(m, :ns)

    xlo = get_setting(m, :xlo)
    xhi = get_setting(m, :xhi)

    zlo = m[:zlo].value
    zhi = m[:zhi].value

    # Load parameters
    R  = 1 + m[:r].scaledvalue
    H  = m[:H].value
    η  = m[:η].value
    γ  = m[:γ].scaledvalue
    ω  = m[:ωstar].value
    T  = m[:Tstar].value
    bg = m[:bg].value

    # Load grids
    xgrid = m.grids[:xgrid].points
    sgrid = m.grids[:sgrid].points
    xswts = m.grids[:weights_total]
    f     = m.grids[:fgrid]

    βlo_temp = βlo
    βhi_temp = βhi

    reject  = false
    counter = 1

    n  = ns*nx
    c  = zeros(n)
    bp = zeros(n)
    KF = zeros(n,n)
    μ  = zeros(n)

    # Initial guess
    β   = 1.0
    Win = Vector{Float64}(undef, n)
    Win_guess = ones(n)

    if get_setting(m, :use_last_βstar) && !isnan(m[:βstar].value)

        # This short-circuits computation of the policy function
        βlo = βhi = m[:βstar].value
    elseif !isnan(m[:βstar].value)

        # If one has computed β* before, we first bisect into a neighborhood around it
        βlo_temp = m[:βstar].value - βband
        βhi_temp = m[:βstar].value + βband

        c, bp, Win, KF, reject = policy_hetdsgegovdebt(nx, ns, βlo_temp, R, ω, H, η, T, γ,
                                                       zhi, zlo, xgrid, sgrid, xswts, Win_guess,
                                                       f, damp = get_setting(m, :policy_damp), maxit = get_setting(m, :policy_maxit))
        excess_lo, μ = compute_excess(xswts, KF, bp, bg)

        if excess_lo < 0 && abs(excess_lo) > tol
            βlo = βlo_temp
            c, bp, Win, KF, reject = policy_hetdsgegovdebt(nx, ns, βhi_temp, R, ω, H, η, T, γ, zhi,
                                                           zlo, xgrid, sgrid, xswts, Win_guess, f,
                                                           damp = get_setting(m, :policy_damp), maxit = get_setting(m, :policy_maxit))
            excess_hi, μ = compute_excess(xswts, KF, bp, bg)

            if excess_hi > 0
                βhi = βhi_temp
            end
        elseif excess_lo >= 0
            βhi = βlo_temp
        end
    end

    while abs(excess) > tol && counter < maxit # clearing markets
        β = (βlo + βhi) / 2.0
        c, bp, Win, KF, reject = policy_hetdsgegovdebt(nx, ns, β, R, ω, H, η, T, γ, zhi, zlo,
                                                       xgrid, sgrid, xswts, Win_guess, f,
                                                       damp = get_setting(m, :policy_damp), maxit = get_setting(m, :policy_maxit))
        excess, μ = compute_excess(xswts, KF, bp, bg)
        # bisection
        if excess > 0
            βhi = β
        elseif excess < 0
            βlo = β
        end
        counter += 1
        if get_setting(m, :use_last_βstar) && !isnan(m[:βstar].value)
            break
        end
    end

  #=  @show β
    This all was for testing when we were trying to compare the old steadystate to the new steadystate
    GR.inline("png")
    p = Plots.plot(xgrid, μ[1:300], label = "lowskill")
    Plots.plot!(p, xgrid, μ[301:600], label = "highskill")
    Plots.savefig(p, "agrid_vs_D_together_old.png")
    p = Plots.plot(xgrid, μ[1:300], label = "lowskill")
    Plots.savefig(p, "agrid_vs_D_lowskill_old.png")
    p = Plots.plot(xgrid, μ[301:600], label = "highskill")
    Plots.savefig(p, "agrid_vs_D_highskill_old.png")

    Plots.plot(μ)
    Plots.savefig("mu_old.png")

    p = Plots.plot(xgrid, c[1:300], label = "lowskill")
    Plots.plot!(p, xgrid, c[301:600], label = "highskill")
    Plots.savefig(p, "agrid_vs_cpol_old.png")

    p = Plots.plot(xgrid, xgrid - c[1:300], label = "lowskill")
    Plots.plot!(p, xgrid, xgrid - c[301:600], label = "highskill")
    Plots.savefig(p, "agrid_minus_cpol_old.png")

    p = Plots.plot()
    for mult = 1:6
        for i = 1:mult*100
            Plots.plot!(p, KF[:, i])
        end
        Plots.savefig(p,"KF_old_mult=$mult.png")
    end
    for i = 1:600
        p = Plots.plot(KF[:, i])
        Plots.savefig(p, "KF_dir/KF_old_$i.png")
    end=#

    # If policy function does not converge, we signal to likelihood that should reject
    m <= Setting(:auto_reject, reject)
    m[:lstar]  = Win
    m[:cstar]  = c
    m[:μstar]  = μ
    m[:βstar]  = β
    m[:β_save] = β
    nothing
end

# These are for computing moments to compare against targets (the next few functions)
function ave_mpc(m::AbstractArray, c::AbstractArray, agrid::AbstractArray,
                 aswts::AbstractArray, na::Int, ns::Int)
	mpc = 0.
	for iss=1:ns
		i = na*(iss-1)+1
		mpc += aswts[i]*m[i]*(c[i+1] - c[i])/(agrid[2] - agrid[1])
		for ia=2:na
			i = na*(iss-1)+ia
			mpc += aswts[i]*m[i]*(c[i] - c[i-1])/(agrid[ia] - agrid[ia-1])
		end
	end
	return mpc
end

function frac_zero(m::AbstractArray, c::AbstractArray, agrid::AbstractArray,
                   aswts::AbstractArray, ns::Int)
	return sum(aswts .* m .* (c .== repeat(agrid, ns)))
end

function ssample(us::Matrix{S}, P::Matrix{S}, πss::AbstractArray,
                 ni::Int) where {S<:AbstractFloat}
	shist = ones(Int,ni,8)
	for i=1:ni
		if us[i,1] > πss[1]
			shist[i,1] = 2
		end
		for t=2:8
			if us[i,t] > P[shist[i,t-1],1]
				shist[i,t] = 2
			end
		end
	end
	return shist
end

function ln_annual_inc(zhist::Matrix{S}, us::Matrix{S}, zlo::S, P::Matrix{S},
                       πss::AbstractArray, sgrid::AbstractArray,
                       ni::Int) where {S<:AbstractFloat}
  	s_inds = ssample(us,P,πss,ni)
	linc1 = zeros(ni)
	linc2 = zeros(ni)
	for i=1:ni
		inc1 = 0.
		for t=1:4
			zshock = 1. + (1. - zlo)*(zhist[i,t]-1.)
			sshock = (sgrid[1] + (sgrid[2] - sgrid[1])*(s_inds[i,t] - 1.))
			inc1 += zshock*sshock
		end
		inc2 = 0.
		for t=5:8
			zshock = 1. + (1. - zlo)*(zhist[i,t]-1.)
			sshock = (sgrid[1] + (sgrid[2] - sgrid[1])*(s_inds[i,t] - 1.))
			inc2 += zshock*sshock
		end
		linc1[i] += log(inc1)
     	linc2[i] += log(inc2)
	end
	return linc1, linc2
end

function skill_moments(sH_over_sL::Real, zlo::Real, pLH::S, pHL::S, us::Matrix{S},
                       zs::Matrix{S}, ni::Int = 10000) where {S<:AbstractFloat}
	πL    = pHL / (pLH + pHL)
	πss   = [πL; 1.0-πL]
	P     = [[1.0-pLH pLH]; [pHL 1.0-pHL]]
	slo   = 1.0 / (πL + (1-πL) * sH_over_sL)
	shi   = sH_over_sL * slo
	sgrid = [slo; shi]
	linc1, linc2 = ln_annual_inc(zs, us, zlo, P, πss, sgrid, ni)
	return var(linc1), var(linc2 - linc1)
end

"""
```
function loss(x::Vector{S}, target::Vector{S}) where {S<:AbstractFloat}
```

Computes for given vector and target vector, the sum of absolute loss.
"""
loss(x::Vector{S}, target::Vector{S}) where {S<:AbstractFloat} = sum(abs.(x-target))

"""
```
function best_fit(pLH::S, pHL::S, target::Vector{S}, lower::Vector{S}, upper::Vector{S},
                  us::Matrix{S}, zs::Matrix{S}, max_iter::Int = 20,
                  initial_guess::Vector{S} = [6.3, 0.03]) where {S<:AbstractFloat}
```

Uses Nelder-Mead (gradient descent algorithm) to optimize for income moments.
"""
function best_fit(pLH::S, pHL::S, target::Vector{S}, lower::Vector{S}, upper::Vector{S},
                  us::Matrix{S}, zs::Matrix{S}, max_iter::Int = 20,
                  initial_guess::Vector{S} = [6.3, 0.03]) where {S<:AbstractFloat}

    skill_moments_f(x) = loss(collect(skill_moments(x[1], x[2], pLH, pHL, us, zs)), target)

    res = optimize(skill_moments_f, lower, upper, initial_guess, Fminbox(NelderMead()),
                   Optim.Options(f_calls_limit = max_iter))

    sH_over_sL_argmin, zlo_argmin = Optim.minimizer(res)
    min_varlinc, min_vardlinc = skill_moments(sH_over_sL_argmin, zlo_argmin, pLH,
                                              pHL, us, zs)
    return sH_over_sL_argmin, zlo_argmin, min_varlinc, min_vardlinc
end

function zsample(uz::Matrix{S}, zgrid::AbstractArray, zcdf::AbstractArray,
                 ni::Int, nz::Int) where {S<:AbstractFloat}
    zave = 0.5*zgrid[1:nz-1]+0.5*zgrid[2:nz]
	zs = zeros(ni,8)
	for i=1:ni
		for t=1:8
			for iz=1:nz-1
				if zcdf[iz] < uz[i,t] <= zcdf[iz+1]
					zs[i,t] = zave[iz]
				end
			end
		end
	end
	return zs
end

# This function computes excess savings. The KF that pops out of steadystate is nx*ns by nx*ns but we want to reduce that to a vector of length nx*ns which is the stationary distribution (this is what the eigenstuff is doing). Stationary because no matter where you are now, you want to have the same probability of going to all the different gridpoints tomorrow.
@inline function compute_excess(xswts::Vector{S}, KF::Matrix{S}, bp::Vector{S},
                                bg::S; print_warning::Bool = false,
                                tol::S = 2e-1) where {S<:Float64}
    LPMKF = xswts[1] * KF

    # Find eigenvalue closest to 1
    (D,V) = (eigen(LPMKF)...,)
    max_D = argmax(abs.(D))
    D = D[max_D]

    if abs(D - 1) > tol && print_warning
        @warn "Your eigenvalue is too far from 1, something is wrong."
    end

    # Pick eigenvector associated w/ largest eigenvalue and moving it back to values
    μ = real(V[:,max_D])
    # @show μ
    μ = μ ./ dot(xswts, μ) # Scale of eigenvectors not determinate: rescale to integrate to 1
    excess = dot(xswts, (μ .* bp)) - bg # Compute excess supply of savings, which is a fn of w
    return excess, μ
end

# policy function (consumption and ell = 1/consumption are the "policies" being decied on here)
function policy_hetdsgegovdebt(nx::Int, ns::Int, β::S, R::S, ω::S, H::S, η::S,
                               T::S, γ::S, zhi::S, zlo::S, xgrid::Vector{S},
                               sgrid::Vector{S}, xswts::Vector{S}, Win::Vector{S},
                               f::Matrix{S}, dist::S = 1., tol::S = 1e-4;
                               maxit::Int64 = 500, damp::S = 0.5) where {S<:AbstractFloat}
    n    = nx*ns
    c    = zeros(n)                  # consumption
    bp   = Vector{Float64}(undef, n) # savings
    Wout = Vector{Float64}(undef, length(Win))
    counter = 1
    reject = false
    qfunction(x::Float64) = mollifier_hetdsgegovdebt(x, zhi, zlo)

    while dist>tol && counter<maxit
        # compute c(w) given guess for Win = β*R*E[u'(c_{t+1})]
        for iss in 1:ns
            for ia in 1:nx
                c[nx*(iss-1)+ia] = min(1/Win[nx*(iss-1)+ia],xgrid[ia]+η)
            end
        end

        bp = R*(exp(-γ))*(repeat(xgrid, ns) - c)  # compute bp(w) given guess for Win
        # Wout and Win are the ell function (in the paper!). The point is we're trying to get the policy function (ell) to converge (which is the business with Win and Wout because you need to compare the difference between Win and Wout to see if it converged)
        Wout = parameterized_expectations_hetdsgegovdebt(nx, ns, β, R, ω, H, T, γ,
                                                         qfunction, xgrid,
                                                         sgrid, xswts, c, bp, f)
        dist = maximum(abs.(Wout - Win))
        Win  = damp*Wout + (1.0-damp) * Win
        counter += 1
    end
    if counter == maxit
        @warn "Euler iteration did not converge"
        reject = true
        return c, bp, Wout, zeros(n,n), reject
    end
    tr = kolmogorov_fwd_hetdsgegovdebt(nx, ns, ω, H, T, R, γ, qfunction, xgrid, sgrid, bp, f)
    return c, bp, Wout, tr, reject
end

# This function is the formula for ell_t in the middle top of page 7 of the paper
@inline function parameterized_expectations_hetdsgegovdebt(nx::Int, ns::Int, β::S,
                                                           R::S, ω::S, H::S, T::S, γ::S,
                                                           qfunc::Function, xgrid::Vector{S},
                                                           sgrid::Vector{S}, xswts::Vector{S},
                                                           c::Vector{S}, bp::Vector{S},
                                                           f::Matrix{S}) where {S<:AbstractFloat}
    l_out = zeros(nx*ns)
    for iss=1:ns     # loop off s (only 2 values)
        for ia=1:nx # loop over a
            sumn = 0.0
            for isp=1:ns # loop over s' (only 2 values)
                for iap=1:nx # loop over a'
                    # Name in code                                                 : notation in paper
                    #xswts                                                         : weights for the integral approximation
                    #c[nx*(isp-1)+iap]                                             : c_{t+1}(a', s')
                    # qfunc                                                        : g(⋅)
                    # qfunc((xgrid[iap] - bp[nx*(iss-1)+ia] - T) / (ω*H*sgrid[isp]): (a'-(1+r)exp(-γ)(a-c_t(a,s) - T))/ w*s'*H
                    # f[iss, isp]/ sgrid[isp]                                      : p(s'|s)/s'
                    # β*R*exp(-γ))/ω*H  (line 429)                                 β*(1+r)*exp(-z)/ω_{t+1}*H{t+1}
                    # When you're trying to match code to paper, be creative in the ordering of stuff as the variables in the paper are commutatively multipled in different orders
                    sumn += (xswts[nx*(isp-1)+iap]/c[nx*(isp-1)+iap]) *
                        qfunc((xgrid[iap] - bp[nx*(iss-1)+ia] - T) /
                              (ω*H*sgrid[isp])) * f[iss,isp] ./ sgrid[isp]
                end
            end
            l_out[nx*(iss-1)+ia] = (β*R*(exp(-γ))/ω*H)*sumn
        end
    end
    return l_out
end

# This function is the exprression for m_{t+1} in the middle of page 7 of the paper. See the key in lines 417-424 of this file to help you in matching stuff up
@inline function kolmogorov_fwd_hetdsgegovdebt(nx::Int, ns::Int, ω::S, H::S, T::S, R::S,
                                               γ::S, qfunc::Function, xgrid::Vector{S},
                                               sgrid::Vector{S}, bp::Vector{S},
                                               f::Matrix{S}) where {S<:AbstractFloat}
    tr = zeros(nx*ns,nx*ns)
    for iss=1:ns
        for ia=1:nx
            for isp=1:ns
                for iap=1:nx
                    tr[nx*(isp-1)+iap, nx*(iss-1)+ia] = qfunc((xgrid[iap] -
                                                               bp[nx*(iss-1)+ia] - T) /
                                                              (ω*H*sgrid[isp])) * f[iss,isp] ./
                                                              (ω * H * sgrid[isp])
                end
            end
        end
    end
    return tr
end

# also known as qfunction of g(⋅)
@inline function mollifier_hetdsgegovdebt(z::S, ehi::S, elo::S) where {S<:AbstractFloat}
    In = 0.443993816237631
    if z<ehi && z>elo
        temp = -1.0 + 2.0 * (z - elo) / (ehi - elo)
        return (2.0 / (ehi - elo)) * exp(-1.0 / (1.0 - temp^2)) / In
    end
    return 0.0
end
# also known as qfunction of g(⋅) but the derivative of it
@inline function dmollifier_hetdsgegovdebt(x::S, ehi::S, elo::S) where {S<:AbstractFloat}
    In = 0.443993816237631
    if x<ehi && x>elo
        temp = (-1.0 + 2.0*(x-elo)/(ehi-elo))
        out  = -(2*temp./((1 - temp.^2).^2)).*(2/(ehi-elo)) .*
            mollifier_hetdsgegovdebt(x, ehi, elo)
    else
        out = 0.0
    end
    return out
end

function calibrate_pLH_pHL(m::HetDSGEGovDebt)
    target_mpc, target_pc0 = get_setting(m, :targets)
    σt_mpc, σt_pc0         = get_setting(m, :target_σt)

    function minimize_penalty(m, pLHpHL::Vector{Float64})
        m[:pLH] = pLHpHL[1]
        m[:pHL] = pLHpHL[2]
        steadystate!(m)
        if get_setting(m, :auto_reject)
            m <= Setting(:auto_reject, false)
            @show pLHpHL, Inf
            return Inf
        end
        @show pLHpHL, m[:mpc].value, m[:pc0].value
        println(0.5*((log(m[:mpc])-log(target_mpc))^2/σt_mpc +
                    (log(m[:pc0].value)-log(target_pc0))^2/σt_pc0^2))
        return 0.5*((log(m[:mpc])-log(target_mpc))^2/σt_mpc +
                    (log(m[:pc0].value)-log(target_pc0))^2/σt_pc0^2)
    end
    minimize_me(x) = minimize_penalty(m, x)
    res = optimize(minimize_me, [0.005, 0.005], [0.095, 0.095], [0.01125, 0.03],
                   Fminbox(Optim.NelderMead()), Optim.Options(f_calls_limit=300))
    println(res)
    return res
end


function hat(ω)
    # skew-symmetric
    return [0 -ω[3] ω[2];
            ω[3] 0 -ω[1];
            -ω[2] ω[1] 0]
end

function L(Q)
    [Q[1] -Q[2:4]'; Q[2:4] Q[1]*I + hat(Q[2:4])]
end

function R(Q)
    [Q[1] -Q[2:4]'; Q[2:4] Q[1]*I - hat(Q[2:4])]
end

H = [zeros(1,3); I];

T = Diagonal([1.0; -1.0; -1.0; -1.0])

function G(Q)
    return L(Q)*H  # 4x3
end

function Ḡ(q)
    Qb = q[4:7]
    Q0 = q[11:14]
    Q1 = q[18:21]
    Q2 = q[25:28]
    Q3 = q[32:35]

    # 35x30
    return blockdiag(sparse(I, 3, 3), sparse(G(Qb)),
                     sparse(I, 3, 3), sparse(G(Q0)),
                     sparse(I, 3, 3), sparse(G(Q1)),
                     sparse(I, 3, 3), sparse(G(Q2)),
                     sparse(I, 3, 3), sparse(G(Q3)),)
end

function Expq(ϕ)

    # The quaternion exponential map ϕ → q
    Q = zeros(4)
    θ = norm(ϕ)
    Q = [cos(θ/2); 0.5*ϕ*sinc(θ/(2*pi))]

    return Q
end

function rotate(Q, p)
    # Rotate a position vector p by a quaternion Q
    # H'L(Q)*R(Q)'*H*p
    return H'*R(Q)'*L(Q)*H*p
end

function anglesolve(Q)
    # convert arbitrary quaternion to unsigned angle
    2 * atan(norm(Q[2:4]),Q[1])
end
#=
function anglebt(Q1, Q2)
    # get smallest angle b/t quaternions
    acos((Q1'*Q2) / (norm(Q1)*norm(Q2)))
end
=#

function angle_y(Q1, Q2)
    # signed angle about y axis of y-axis-constrained quaternions
    Q12 = L(Q1)'*Q2
    Q12 = Q12/(norm(Q12))
    return 2 * asin(Q12[3])
end

function a_joint(q)
    # convert quaternions to relative joint angles b/t links

    Qb = q[4:7]
    Q0 = q[11:14]
    Q1 = q[18:21]
    Q2 = q[25:28]
    Q3 = q[32:35]

    a0 = -angle_y(Qb, Q0)
    a1 = -angle_y(Q0, Q1)
    a2 = -angle_y(Qb, Q2)
    a3 = -angle_y(Q2, Q3)

    return [a0; a1; a2; a3]

end

function a_act(q)
    # convert quaternions to relative joint angles b/t links
    # only actuated joints
    Qb = q[4:7]
    Q0 = q[11:14]
    Q2 = q[25:28]

    a0 = -angle_y(Qb, Q0)
    a2 = -angle_y(Qb, Q2)

    return [a0; a2]

end

function a_vel(a, a_prev, dt)
    # get joint velocities from current and previous joint angle
    return (a .- a_prev)/dt
end

function smoothsqrt(x)
    ϵ = 1e-7
    return sqrt(x.+ϵ*ϵ) .- ϵ
end


@doc raw"""
`IDAE`: Implicit Differential Algebraic Equation

Defines a partitioned differential algebraic initial value problem
```math
\begin{align*}
\dot{q} (t) &= v(t) + u(t, q(t), p(t), \lambda(t)) , & q(t_{0}) &= q_{0} , \\
\dot{p} (t) &= f(t, q(t), v(t)) + r(t, q(t), p(t), \lambda(t)) , & p(t_{0}) &= p_{0} , \\
p(t) &= p(t, q(t), v(t)) , && \
0 &= \phi (t, q(t), p(t), \lambda(t)) , & \lambda(t_{0}) &= \lambda_{0} ,
\end{align*}
```
with vector field ``f``, the momentum defined by ``p``, projection ``u`` and ``r``,
algebraic constraint ``\phi=0``,
conditions ``(q_{0}, p_{0})`` and ``\lambda_{0}``, the dynamical variables
``(q,p)`` taking values in ``\mathbb{R}^{d} \times \mathbb{R}^{d}`` and
the algebraic variable ``\lambda`` taking values in ``\mathbb{R}^{n}``.

### Fields

* `d`: dimension of dynamical variables ``q`` and ``p`` as well as the vector fields ``f`` and ``p``
* `m`: dimension of algebraic variable ``\lambda`` and the constraint ``\phi``
* `n`: number of initial conditions
* `ϑ`: function determining the momentum
* `f`: function computing the vector field ``f``
* `u`: function computing the projection
* `g`: function computing the projection
* `ϕ`: algebraic constraint
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q``
* `p₀`: initial condition for dynamical variable ``p``
* `λ₀`: initial condition for algebraic variable ``\lambda``

"""
struct IDAE{dType <: Number, tType <: Number,
            ϑType <: Function, fType <: Function,
            uType <: Function, gType <: Function,
            ϕType <: Function, vType <: Union{Function,Nothing},
            pType <: Union{Tuple,Nothing}, N} <: AbstractEquationPDAE{dType, tType}

    d::Int
    m::Int
    n::Int
    ϑ::ϑType
    f::fType
    u::uType
    g::gType
    ϕ::ϕType
    v::vType
    t₀::tType
    q₀::Array{dType, N}
    p₀::Array{dType, N}
    λ₀::Array{dType, N}
    parameters::pType
    periodicity::Vector{dType}

    function IDAE(DT::DataType, N::Int, d::Int, m::Int, n::Int,
                  ϑ::ϑType, f::fType, u::uType, g::gType, ϕ::ϕType, t₀::tType,
                  q₀::AbstractArray{dType}, p₀::AbstractArray{dType}, λ₀::AbstractArray{dType};
                  v=nothing, parameters=nothing, periodicity=zeros(DT,d)) where {
                        dType <: Number, tType <: Number,
                        ϑType <: Function, fType <: Function,
                        uType <: Function, gType <: Function,
                        ϕType <: Function}

        @assert d == size(q₀,1) == size(p₀,1)
        @assert m == size(λ₀,1)
        @assert n == size(q₀,2) == size(p₀,2) == size(λ₀,2)
        @assert 2d ≥ m
        @assert ndims(q₀) == ndims(p₀) == ndims(λ₀) == N ∈ (1,2)

        new{DT, tType, ϑType, fType, uType, gType, ϕType, typeof(v), typeof(parameters), N}(d, m, n, ϑ, f, u, g, ϕ, v, t₀,
                convert(Array{DT}, q₀), convert(Array{DT}, p₀), convert(Array{DT}, λ₀),
                parameters, periodicity)
    end
end

function IDAE(ϑ, f, u, g, ϕ, t₀::Number, q₀::AbstractArray{DT}, p₀::AbstractArray{DT}, λ₀::AbstractArray{DT}=zeros(q₀); kwargs...) where {DT}
    IDAE(DT, ndims(q₀), size(q₀,1), size(λ₀,1), size(q₀,2), ϑ, f, u, g, ϕ, t₀, q₀, p₀, λ₀; kwargs...)
end

function IDAE(ϑ, f, u, g, ϕ, q₀::AbstractArray, p₀::AbstractArray, λ₀::AbstractArray=zeros(q₀); kwargs...)
    IDAE(ϑ, f, u, g, ϕ, zero(eltype(q₀)), q₀, p₀, λ₀; kwargs...)
end

Base.hash(dae::IDAE, h::UInt) = hash(dae.d, hash(dae.m, hash(dae.n,
        hash(dae.ϑ, hash(dae.f, hash(dae.u, hash(dae.g, hash(dae.ϕ, hash(dae.v,
        hash(dae.t₀, hash(dae.q₀, hash(dae.p₀, hash(dae.λ₀,
        hash(dae.periodicity, hash(dae.parameters, h)))))))))))))))

Base.:(==)(dae1::IDAE, dae2::IDAE) = (
                                dae1.d == dae2.d
                             && dae1.m == dae2.m
                             && dae1.n == dae2.n
                             && dae1.ϑ == dae2.ϑ
                             && dae1.f == dae2.f
                             && dae1.u == dae2.u
                             && dae1.g == dae2.g
                             && dae1.ϕ == dae2.ϕ
                             && dae1.v == dae2.v
                             && dae1.t₀ == dae2.t₀
                             && dae1.q₀ == dae2.q₀
                             && dae1.p₀ == dae2.p₀
                             && dae1.λ₀ == dae2.λ₀
                             && dae1.parameters == dae2.parameters
                             && dae1.periodicity == dae2.periodicity)

function Base.similar(dae::IDAE, q₀, p₀, λ₀=get_λ₀(q₀, dae.λ₀); kwargs...)
    similar(dae, dae.t₀, q₀, p₀, λ₀; kwargs...)
end

function Base.similar(dae::IDAE, t₀::TT, q₀::AbstractArray{DT}, p₀::AbstractArray{DT}, λ₀::AbstractArray{DT}=get_λ₀(q₀, dae.λ₀);
                      v=dae.v, parameters=dae.parameters, periodicity=dae.periodicity) where {DT  <: Number, TT <: Number}
    @assert dae.d == size(q₀,1) == size(p₀,1)
    @assert dae.m == size(λ₀,1)
    IDAE(dae.ϑ, dae.f, dae.u, dae.g, dae.ϕ, t₀, q₀, p₀, λ₀; v=v, parameters=parameters, periodicity=periodicity)
end

@inline Base.ndims(dae::IDAE) = dae.d

@inline periodicity(equation::IDAE) = equation.periodicity


using JuMP
using Ipopt

function Solve_Linearized_Dist_SOC(buses, lines, generators, fp0, fq0, l0)
info("Starting Linearized DistSOC Model")

fp0 = vcat([0], fp0)
fq0 = vcat([0], fq0)
l0 = vcat([0], l0)

println(fp0)

bus_set = collect(keys(buses))
line_set = collect(keys(lines))
gen_set = collect(keys(generators))
gen_bus = []
for g in keys(generators)
    push!(gen_bus, generators[g].bus_idx)
end
root_bus = 0
for b in keys(buses)
    if buses[b].is_root
        root_bus = b
    end
end

lines_to = Dict()
for l in keys(lines)
    lines_to[lines[l].to_node] = lines[l]
end

v_root = 1

m = Model(solver=IpoptSolver())

@variable(m, v[bus_set] >= 0) #variable for voltage square
@variable(m, fp[bus_set]) #variable for active power flow
@variable(m, fq[bus_set]) # variable for reactive power flow
@variable(m, gp[bus_set]) #variable for active power generation
@variable(m, gq[bus_set]) #variable for reactive power generation
@variable(m, l[bus_set] >= 0) # variable for quadratic current

@objective(m, Min, sum((gp[b])*buses[b].generator.cost for b in gen_bus))

for b in gen_bus
# All buses with generation
    @constraint(m, gp[b] <= buses[b].generator.g_P_max) #Upper limit constraint for active power generation
    @constraint(m, gp[b] >= 0) #Lower limit constraint for active power generation
    @constraint(m, gq[b] <= buses[b].generator.g_Q_max) #Upper limit constraint for reactive power generation
    @constraint(m, gq[b] >= -buses[b].generator.g_Q_max) #Lower limit constraint for reactive power generation
end

for b in setdiff(bus_set, gen_bus)
# All buses without generation
    @constraint(m, gp[b] == 0)
    @constraint(m, gq[b] == 0)
end

for b in bus_set
# All buses
    @constraint(m, buses[b].d_P - gp[b] + sum((fp[k] + l[k]*lines_to[k].r) for k in buses[b].children) == fp[b])
    @constraint(m, buses[b].d_Q - gq[b] + sum((fq[k] + l[k]*lines_to[k].x) for k in buses[b].children) == fq[b])

    @constraint(m, v[b] <= buses[b].v_max)
    @constraint(m, v[b] >= buses[b].v_min)
end

for b in setdiff(bus_set, [root_bus])
# All buses without root
    b_ancestor = buses[b].ancestor[1]
    R = lines_to[b].r
    X = lines_to[b].x
    term_p = (fp0[b]+l0[b]*R)
    term_q = (fq0[b]+l0[b]*X)

    @constraint(m, -fp0[b]^2 - fq0[b]^2 + 2fp0[b]*fp[b] + 2fq0[b]*fq[b] <= lines_to[b].s_max^2)
    @constraint(m, term_p^2 + term_q^2 +
                     2*(fp[b] - fp0[b])*term_p +
                     2*(fq[b] - fq0[b])*term_q +
                     2*(l[b] - l0[b])*(term_p*R + term_q*X) <= lines_to[b].s_max^2)
    # Same as above but rearranged
    # @constraint(m, -1*(fp0[b] + l0[b]*R)^2 - (fq0[b] + l0[b]*X)^2 +
    #                 2*(fp0[b] + l0[b]*R)*fp[b] + 2*(fq0[b] + l0[b]*X)*fq[b] +
    #                 2*(l0[b]*R^2 + fp0[b]*R + l0[b]*X^2 + fq0[b]*X)*l[b] <= lines_to[b].s_max^2)

    @constraint(m, v[b_ancestor] - 2*(R*fp[b] + X*fq[b]) - l[b]*(R^2 + X^2) == v[b])
    @constraint(m, ((2*fp0[b])/l0[b]) * fp[b] + ((2*fq0[b])/l0[b]) * fq[b] - ((fp0[b]^2 + fq0[b]^2)/l0[b]^2) * l[b] <= v[b])
end

#Voltage and flow constraints for root node
@constraint(m, v[root_bus] == v_root)
@constraint(m, l[root_bus] == 0)
@constraint(m, fp[root_bus] == 0)
@constraint(m, fq[root_bus] == 0)

solve(m)

# Prepare Results
bus_results = DataFrame(bus = Any[], d_P = Any[], d_Q = Any[], g_P = Any[], g_Q = Any[], v_squared = Any[])
line_results = DataFrame(line = Any[], from = Any[], to = Any[], f_P = Any[], f_q = Any[], a_squared = Any[])
for b in bus_set
    res = [b, buses[b].d_P, buses[b].d_Q, getvalue(gp[b]), getvalue(gq[b]), getvalue(v[b])]
    push!(bus_results, res)

    # Calc square current on line
    a_res = 0
    fp_res = getvalue(fp[b])
    fq_res = getvalue(fq[b])
    if b != root_bus
        a_res = getvalue(l[b])
        lres = [lines_to[b].index, buses[b].ancestor[1], b, fp_res, fq_res, a_res]
        push!(line_results, lres)
    end
end

sort!(bus_results, cols=:bus)
sort!(line_results, cols=:to)

return getobjectivevalue(m), bus_results, line_results

end


using LinearAlgebra;
using BivMatFun;
using SpecialFunctions;
using Printf;
using MAT;
using BenchmarkTools;
using DelimitedFiles;

include("experiments_common.jl");

function run_test(use_grcar = false)

    success = true;

    a = 1/2;
    f = (x,y,i,j) -> (x+y).^(-a-i-j) * gamma(a+i+j) ./ gamma(a) * (-1)^(i+j);

    ntests = 5;

    # n, errD, errT, nblocksDA, nblocksDB, nblocksTA, nblocksTB, digitsD, digitsT, errest, timeD, timeT
    data = zeros(2*ntests, 12);

    for i = 1 : ntests
        n = 32 * i;

        vars = matread("$(data_directory)/experiment1_$(n)_random.mat")
        A = vars["A"]; B = vars["B"]; C = vars["C"]; Y = vars["Y"]; errest1 = vars["errest"]

        @printf("n = %d, A,B,C random\n", n);

        XD, infod = fun2m(f, A, B, C, method = BivMatFun.Diag);
        if ! is_ci_test()
            tdiag = @benchmark fun2m($f, $A, $B, $C, method = $BivMatFun.Diag);
            tdiag = median(tdiag.times) / 1e9
        else
            tdiag = 0.0
        end

        @printf("  Diag: err = %e, nblocks = %d %d, max digits = %d, time = %f, errest = %e\n", norm(XD - Y) / norm(Y), infod.nblocksA, infod.nblocksB, infod.digits, tdiag, errest1);

        XT, infot = fun2m(f, A, B, C, method = BivMatFun.Taylor);

        if ! is_ci_test()
            tt = @benchmark fun2m($f, $A, $B, $C, method = $BivMatFun.Taylor);
            tt = median(tt.times) / 1e9;
        else
            tt = 0.0
        end

        @printf("  Taylor: err = %e, nblocks = %d %d, max deg = %d, time = %f, errest = %e\n", norm(XT - Y) / norm(Y), infot.nblocksA, infot.nblocksB, infot.digits, tt, errest1);

        success &= norm(XD - Y) / norm(Y) < errest1 * 100;
        success &= norm(XT - Y) / norm(Y) < errest1 * 100;

        # GRCAR
        vars = matread("$(data_directory)/experiment1_$(n)_grcar.mat")
        A = vars["A"]; B = vars["B"]; C = vars["C"]; Y2 = vars["Y"]; errest2 = vars["errest"]

        @printf("n = %d, A grcar, B,C random\n", n);

        XD2, infod2 = fun2m(f, A, B, C, method = BivMatFun.Diag);

        if ! is_ci_test()
            tdiag2 = @benchmark fun2m($f, $A, $B, $C, method = $BivMatFun.Diag);
            tdiag2 = median(tdiag2.times) / 1e9;
        else
            tdiag2 = 0.0
        end

        @printf("  Diag: err = %e, nblocks = %d %d, max digits = %d, time = %f, errest = %e\n", Float64(norm(XD2 - Y2) / norm(Y2)), infod2.nblocksA, infod2.nblocksB, infod2.digits, tdiag2, errest2);

        BivMatFun.warnings(false);
        XT2, infot2 = fun2m(f, A, B, C, method = BivMatFun.Taylor);
        if ! is_ci_test()
            tt2 = @benchmark fun2m($f, $A, $B, $C, method = $BivMatFun.Taylor);
            tt2 = median(tt2.times) / 1e9;
        else
            tt2 = 0.0
        end

        BivMatFun.warnings(true);

        @printf("  Taylor: err = %e, nblocks = %d %d, max deg = %d, time = %f, errest = %e\n", Float64(norm(XT2 - Y2) / norm(Y2)), infot2.nblocksA, infot2.nblocksB, infot2.digits, tt2, errest2);

        # For the GRCAR matrix, which involves a-posteriori re-blocking due to a
        # highly non-normal eigenvector matrix, we expect Taylor to fail, hence
        # we do not check it.
        success &= norm(XD2 - Y2) / norm(Y) < errest2 * 100;

        #
        # % n, errD, errT, nblocksDA, nblocksDB, nblocksTA, nblocksTB, digitsD, digitsT, errest, timeD, timeT
        data[2*i-1, 1] = n;
        data[2*i-1, 2] = norm(XD - Y) / norm(Y);
        data[2*i-1, 3] = norm(XT - Y) / norm(Y);
        data[2*i-1, 4] = infod.nblocksA;
        data[2*i-1, 5] = infod.nblocksB;
        data[2*i-1, 6] = infot.nblocksA;
        data[2*i-1, 7] = infot.nblocksB;
        data[2*i-1, 8] = infod.digits;
        data[2*i-1, 9] = infot.digits;
        data[2*i-1, 10] = errest1;
        data[2*i-1, 11] = tdiag;
        data[2*i-1, 12] = tt;
        #
        data[2*i, 1] = n;
        data[2*i, 2] = norm(XD2 - Y2) / norm(Y2);
        data[2*i, 3] = norm(XT2 - Y2) / norm(Y2);
        data[2*i, 4] = infod2.nblocksA;
        data[2*i, 5] = infod2.nblocksB;
        data[2*i, 6] = infot2.nblocksA;
        data[2*i, 7] = infot2.nblocksB;
        data[2*i, 8] = infod2.digits;
        data[2*i, 9] = infot2.digits;
        data[2*i, 10] = errest2;
        data[2*i, 11] = tdiag2;
        data[2*i, 12] = tt2;

        writedlm("exp1.dat", data, '\t');
    end

    return success;

end

run_test(true);


function normcol_equal(matin)
# solve the proximal problem
# matout = argmin||matout-matin||_F^2, s.t. matout(:,i)=1
    matin ./ repmat(sqrt(sum(matin .^ 2, 1) + eps()), size(matin, 1), 1)
end


using LinearAlgebra

function proj_dual_iso2d!(p, temp, weight)
    p1, p2 = view(p,:,:,1), view(p,:,:,2)
    temp .= 1 ./ max.(1.0, sqrt.( p1 .^2 + p2 .^ 2 ) ./ (weight+1e-8) )
    p1 .*= temp
    p2 .*= temp
end

"Project to dual l1 norm"
function proj_dual_l1!(x, weight)
    x ./= max.(1.0, abs.(x) / weight)
end

function div2d!(divp, p)
    p1, p2 = view(p,:,:,1), view(p,:,:,2)
    p1_x, p2_y = view(divp,:,:,1), view(divp,:,:,2)

    # p2_y is temp variable to save memory
    p1_x         .=  p1 - circshift!(p2_y, p1, [0, 1])
    p1_x[:, end] .= -p1[:, end-1]
    p1_x[:,   1] .=  p1[:, 1]

    circshift!(p2_y, p2, [1, 0])
    p2_y         .= p2 - p2_y
    p2_y[end, :] .= -p2[end-1, :]
    p2_y[1,   :] .=  p2[1, :]

    p1_x .+= p2_y
    return p1_x
end

function grad2d!(du, u) where {T<:AbstractFloat}
    du1, du2 = view(du, :, :, 1), view(du, :, :, 2)
    du1 .= circshift!(du1, u, [0, -1]) .- u
    du2 .= circshift!(du2, u, [-1, 0]) .- u
    du1[:, end] .= 0.0
    du2[end, :] .= 0.0
end

"Compute opeartor norm based on the Power method"
function compute_opnorm_block(As,shape, niter=3)
    λ = 0
    size,count = shape
    x = rand(size*count)

    for i=1:niter
        x_next = zeros(size*count)
        for j=1:count
            A = As[j]
            Ax = A*x[((j-1)*size)+1:((j-1)*size)+size]
            x_next_j = A'*Ax
            x_next[((j-1)*size)+1:((j-1)*size)+size] = x_next_j
        end
        λ_prev = λ
        λ = sqrt(sum(x_next.^2))

        if (abs(λ_prev-λ) < 1e-9)
            println("break at iteration: ", i)
            break
        end
        x = x_next / λ
    end
    return sqrt(λ)
end

function _recon2d_tv_primaldual_flow(As,A_norm,W_list,W_norm,u0s,bs,w_tv,w_flow,c,niter,p1,p2,p3)
    #A variational reconstruction method for undersampled dynamic x-ray tomography based on physical motion models (Burger)#
    height, width, frames = size(u0s)

    u = u0s
    u_prev = similar(u)
    ubar = deepcopy(u)

    p_adjoint  = similar(p3)
    p3_adjoint = similar(p3)
    p_adjoint_prev = similar(p3)
    opsnorm = [A_norm, sqrt(8), W_norm]
    tau = c /sum(opsnorm)

    data1 = similar(bs)
    u_grad = similar(p2)
    u_temp = similar(u)
    divp2 = similar(p2)

    Wus = similar(u)

    sigmas = map(n -> 1.0/(n*c), opsnorm)
    #@info sigmas

    for it=1:niter
        u_prev .= u
        p_adjoint_prev .= p_adjoint
        Threads.@threads for t=1:frames
            A = As[t]
            #sigmas = map(n-> 1.0/(n*c), opsnorm)

            @views data1[:,t] .= mul!(data1[:,t], A, vec(ubar[:,:,t])) .- bs[:,t]
            @views p1[:,t] .= (p1[:,t] .+ sigmas[1] * data1[:,t]) / (sigmas[1] + 1.0)
            @views mul!( view(p_adjoint[:,:,t], :), A', p1[:,t])

            @views grad2d!(u_grad[:,:,:,t], ubar[:,:,t])
            @views p2[:,:,:,t] .+= sigmas[2] .* u_grad[:,:,:,t]
            @views proj_dual_iso2d!(p2[:,:,:,t], u_temp[:,:,t], w_tv)

            @views p_adjoint[:,:,t] .-= div2d!(divp2[:,:,:,t], p2[:,:,:,t])

            if t < frames
                _p3 = view(p3, :, :, t)
                Wu = view(Wus, :, :, t)
                Wuv = vec(Wu)
                @views Wuv .= mul!(Wuv, W_list[t], vec(ubar[:,:,t+1])) .- vec(ubar[:,:,t])

                @views _p3 .+= sigmas[3] .* Wu
                proj_dual_l1!(_p3, w_flow)
                p3_adj = view(p3_adjoint[:,:,t], :)

                mul!(p3_adj, W_list[t]', vec(_p3))
                # @views p_adjoint[:,:,t+1] .+= p3_adjoint[:,:,t] # race condition
                @views p_adjoint[:,:,t] .-= _p3
            end
        end

        # add p3_adjoint to prevent race condition
        @views p_adjoint[:,:,2:end] .+= p3_adjoint[:,:,1:end-1]

        # primal update
        u .-= tau * p_adjoint
        u .= max.(u, 0.0) # positivity constraint

        # acceleration
        ubar .= 2*u .- u_prev

        if it % 50 == 0
            #du = u_prev - u
            primal_gap = sum(abs.(-p_adjoint+p_adjoint_prev + (u_prev-u)/tau)) / length(p_adjoint)
            #@info "primal gap:" primal_gap
            @info it primal_gap
        end
    end
    return u
end

"""
    recon2d_tv_primaldual_flow(As, bs, u0s, niter, w_tv, w_flow, c=1.0)

Reconstruct a 2d image by total variation model using Primal Dual optimization method and optical flow estimation

# Args
As : List of forward opeartors
bs : List of projection data
u0s: Initial guess of images
niter: number of iterations
w_tv: weight for TV term
W_flow: weight for flow term
c : See 61 page in 2016_Chambolle,Pock_An_introduction_to_continuous_optimization_for_imagingActa_Numerica

"""
function recon2d_tv_primaldual_flow(A_list, bs::Array{R, 2}, u0s::Array{R, 3}, niter1::Int, niter2::Int, w_tv::R, w_flow::R, c=10.0) where {R <: AbstractFloat}

    height,width,frames = size(u0s)
    v = zeros(height, width, 2, frames)
    u = u0s
    shape = (height*width, frames)
    A_norm = compute_opnorm_block(A_list, shape)

    p1 = zeros(size(bs))
    p2 = zeros(height, width, 2,frames)
    p3 = zeros(height, width, frames)
    for i=1:niter1
        @info i
        u_prev = u
        v_prev = v
        W_list = mapslices(f -> compute_warping_operator(f), v,dims=[1,2,3])
        # check if W_list is correct
        println(sum(abs.(W_list[1]*vec(u[:,:,2]) - vec(u[:,:,1]))))
        @views @info "check if W_list is correct. data term after and before warping" sum(abs.(W_list[1]*vec(u[:,:,2]) - vec(u[:,:,1])))/(height*width) sum(abs.(u[:,:,2] - u[:,:,1]))/(height*width)

        W_norm = compute_opnorm_block(W_list, shape)
        u = _recon2d_tv_primaldual_flow(A_list,A_norm,W_list,W_norm,u,bs,w_tv,w_flow,c,niter2,p1,p2,p3)
        v = get_flows(u)
    end
    return u
end