@maintain_type struct BicycleControl{T} <: FieldVector{3,T}    # common to all models below
    δ::T    # steering angle
    Fxf::T  # front tire longitudinal force (tire frame)
    Fxr::T  # rear tire longidutinal force (tire frame)
end

struct BicycleModelParams{T}
    # Dimensions
    L::T    # wheelbase
    a::T    # distance from front axle to CG
    b::T    # distance from rear axle to CG
    h::T    # CG height

    # Mass and yaw moment of inertia
    G::T    # gravitational acceleration
    m::T    # vehicle mass
    Izz::T  # yaw moment of inertia

    # Tire model parameters
    μ::T    # coefficient of friction
    Cαf::T  # front tire (pair) cornering stiffness
    Cαr::T  # rear tire (pair) cornering stiffness

    # Longitudinal Drag Force Parameters (FxDrag = Cd0 + Cd1*Ux + Cd2*Ux^2)
    Cd0::T  # rolling resistance
    Cd1::T  # linear drag coefficint
    Cd2::T  # quadratic "aero" drag coefficint
end
BicycleModelParams(vehicle::Dict) = BicycleModelParams((vehicle[n] for n in fieldnames(BicycleModelParams))...)

abstract type AbstractBicycleModel{T} <: DifferentialDynamics end
Base.getproperty(B::AbstractBicycleModel, name::Symbol) = getproperty(getfield(B, :params), name)    # potentially not a great idea...

# Coupled Tire Forces - Simplified Model
@inline function fialatiremodel(α, Cα, μ, Fx, Fz)
    F_max = μ*Fz
    abs(Fx) >= F_max ? zero(F_max) : _fialatiremodel(tan(α), Cα, sqrt(F_max*F_max - Fx*Fx))
end

@inline function _fialatiremodel(tanα, Cα, Fy_max)
    tanα_slide = 3*Fy_max/Cα
    ratio = abs(tanα/tanα_slide)
    if ratio <= 1
        -Cα*tanα*(1 - ratio + ratio*ratio/3)
    else
        -Fy_max*sign(tanα)
    end
end

@inline function invfialatiremodel(Fy, Cα, μ, Fx, Fz)
    F_max = μ*Fz
    Fy_max = sqrt(F_max*F_max - Fx*Fx)
    atan(_invfialatiremodel(Fy, Cα, Fy_max))
end

@inline function _invfialatiremodel(Fy, Cα, Fy_max)    # returns tanα
    if abs(Fy) >= Fy_max
        -(3*Fy_max/Cα)*sign(Fy)
    else
        -(1 + cbrt(abs(Fy)/Fy_max - 1))*sign(Fy)
    end
end

@inline function lateral_tire_forces(B::AbstractBicycleModel, αf, αr, Fxf, Fxr, sδ, cδ, num_iters=3)
    L, a, b, h, m, μ, Cαf, Cαr, G = B.L, B.a, B.b, B.h, B.m, B.μ, B.Cαf, B.Cαr, B.G
    Fyf = zero(Fxf)
    Fx = Fxf*cδ - Fyf*sδ + Fxr
    for i in 1:num_iters
        Fzf = (m*G*b - h*Fx)/L
        Fyf = fialatiremodel(αf, Cαf, μ, Fxf, Fzf)
        Fx = Fxf*cδ - Fyf*sδ + Fxr
    end
    Fzr = (m*G*a + h*Fx)/L
    Fyr = fialatiremodel(αr, Cαr, μ, Fxr, Fzr)
    Fyf, Fyr
end

@inline function lateral_tire_forces(B::AbstractBicycleModel, q::StaticVector{6}, u::StaticVector{3}, num_iters=3)
    Ux, Uy, r = q[4], q[5], q[6]
    δ, Fxf, Fxr = u[1], u[2], u[3]
    a, b = B.a, B.b

    sδ, cδ = sincos(δ)
    αf = atan(Uy + a*r, Ux) - δ
    αr = atan(Uy - b*r, Ux)
    lateral_tire_forces(B, αf, αr, Fxf, Fxr, sδ, cδ, num_iters)
end

# Bicycle Model
struct BicycleModel{T} <: AbstractBicycleModel{T}
    params::BicycleModelParams{T}
end
BicycleModel(vehicle::Dict) = BicycleModel(BicycleModelParams(vehicle))

@maintain_type struct BicycleState{T} <: FieldVector{6,T}
    E::T    # world frame "x" position of CM
    N::T    # world frame "y" position of CM
    ψ::T    # world frame heading of vehicle
    Ux::T   # body frame longitudinal speed
    Uy::T   # body frame lateral speed
    r::T    # yaw rate (dψ/dt)
end

@maintain_type struct LocalRoadGeometry{T} <: FieldVector{4,T}
    ψ::T    # nominal trajectory heading
    κ::T    # nominal trajectory (local) curvature
    θ::T    # nominal trajectory (local) pitch grade
    ϕ::T    # nominal trajectory (local) roll grade
end

function (B::BicycleModel{T})((E, N, ψ, Ux, Uy, r)::StaticVector{6},
                              (δ, Fxf, Fxr)::StaticVector{3},
                              (ψᵣ, κ, θ, ϕ)::StaticVector{4}=zeros(LocalRoadGeometry{T})) where {T}
    a, b, m, Izz, Cd0, Cd1, Cd2 = B.a, B.b, B.m, B.Izz, B.Cd0, B.Cd1, B.Cd2

    sψ, cψ = sincos(ψ)
    sδ, cδ = sincos(δ)
    αf = atan(Uy + a*r, Ux) - δ
    αr = atan(Uy - b*r, Ux)
    Fyf, Fyr = lateral_tire_forces(B, αf, αr, Fxf, Fxr, sδ, cδ)
    Fx_drag = -Cd0 - Ux*(Cd1 + Cd2*Ux)
    Fx_grade = 0    # TODO: figure out how roll/pitch are ordered
    Fy_grade = 0
    F̃xf = Fxf*cδ - Fyf*sδ
    F̃yf = Fyf*cδ + Fxf*sδ
    SVector(
        -Ux*sψ - Uy*cψ, # Ux*cψ - Uy*sψ,  # (ψ measured from N)
        Ux*cψ - Uy*sψ,  # Ux*sψ + Uy*cψ,
        r,
        (F̃xf + Fxr + Fx_drag + Fx_grade)/m + r*Uy,
        (F̃yf + Fyr + Fy_grade)/m - r*Ux,
        (a*F̃yf - b*Fyr)/Izz
    )
end
(B::BicycleModel{T})(q::StaticVector{6}, ur::StaticVector{7}) where {T} = B(q, ur[SUnitRange(1,3)], ur[SUnitRange(4,7)])

# Tracking Bicycle Model
struct TrackingBicycleModel{T} <: AbstractBicycleModel{T}
    params::BicycleModelParams{T}
end
TrackingBicycleModel(vehicle::Dict) = TrackingBicycleModel(BicycleModelParams(vehicle))

@maintain_type struct TrackingBicycleState{T} <: FieldVector{6,T}
    Δs::T   # longitudinal error (w.r.t. nominal trajectory)
    Ux::T   # body frame longitudinal speed
    Uy::T   # body frame lateral speed
    r::T    # yaw rate
    Δψ::T   # heading error (w.r.t. nominal trajectory)
    e::T    # lateral error (w.r.t. nominal trajectory)
end

@maintain_type struct TrackingBicycleParams{T} <: FieldVector{4,T}
    V::T    # nominal trajectory speed
    κ::T    # nominal trajectory (local) curvature
    θ::T    # nominal trajectory (local) pitch grade
    ϕ::T    # nominal trajectory (local) roll grade
end

function (B::TrackingBicycleModel{T})((Δs, Ux, Uy, r, Δψ, e)::StaticVector{6},
                                      (δ, Fxf, Fxr)::StaticVector{3},
                                      (V, κ, θ, ϕ)::StaticVector{4}=zeros(TrackingBicycleParams{T})) where {T}
    a, b, m, Izz, Cd0, Cd1, Cd2 = B.a, B.b, B.m, B.Izz, B.Cd0, B.Cd1, B.Cd2

    sΔψ, cΔψ = sincos(Δψ)
    sδ, cδ = sincos(δ)
    αf = atan(Uy + a*r, Ux) - δ
    αr = atan(Uy - b*r, Ux)
    Fyf, Fyr = lateral_tire_forces(B, αf, αr, Fxf, Fxr, sδ, cδ)
    Fx_drag = -Cd0 - Ux*(Cd1 + Cd2*Ux)
    Fx_grade = 0    # TODO: figure out how roll/pitch are ordered
    Fy_grade = 0
    F̃xf = Fxf*cδ - Fyf*sδ
    F̃yf = Fyf*cδ + Fxf*sδ
    SVector(
        Ux*cΔψ - Uy*sΔψ - V,
        (F̃xf + Fxr + Fx_drag + Fx_grade)/m + r*Uy,
        (F̃yf + Fyr + Fy_grade)/m - r*Ux,
        (a*F̃yf - b*Fyr)/Izz,
        r - (Ux*cΔψ - Uy*sΔψ)*κ,
        Ux*sΔψ + Uy*cΔψ
    )
end
(B::TrackingBicycleModel{T})(q::StaticVector{6}, ur::StaticVector{7}) where {T} = B(q, ur[SUnitRange(1,3)], ur[SUnitRange(4,7)])

# Lateral Tracking Bicycle Model
struct LateralTrackingBicycleModel{T} <: AbstractBicycleModel{T}
    params::BicycleModelParams{T}
end
LateralTrackingBicycleModel(vehicle::Dict) = LateralTrackingBicycleModel(BicycleModelParams(vehicle))

@maintain_type struct LateralTrackingBicycleState{T} <: FieldVector{4,T}
    Uy::T   # body frame lateral speed
    r::T    # yaw rate
    Δψ::T   # heading error (w.r.t. nominal trajectory)
    e::T    # lateral error (w.r.t. nominal trajectory)
end

@maintain_type struct LateralTrackingBicycleParams{T} <: FieldVector{4,T}
    Ux::T   # body frame longitudinal speed
    κ::T    # nominal trajectory (local) curvature
    θ::T    # nominal trajectory (local) pitch grade
    ϕ::T    # nominal trajectory (local) roll grade
end

function (B::LateralTrackingBicycleModel{T})((Uy, r, Δψ, e)::StaticVector{4},
                                             (δ, Fxf, Fxr)::StaticVector{3},
                                             (Ux, κ, θ, ϕ)::StaticVector{4}=zeros(LateralTrackingBicycleParams{T})) where {T}
    a, b, m, Izz = B.a, B.b, B.m, B.Izz

    sΔψ, cΔψ = sincos(Δψ)
    sδ, cδ = sincos(δ)
    αf = atan(Uy + a*r, Ux) - δ
    αr = atan(Uy - b*r, Ux)
    Fyf, Fyr = lateral_tire_forces(B, αf, αr, Fxf, Fxr, sδ, cδ)
    Fy_grade = 0
    F̃yf = Fyf*cδ + Fxf*sδ
    SVector(
        (F̃yf + Fyr + Fy_grade)/m - r*Ux,
        (a*F̃yf - b*Fyr)/Izz,
        r - Ux*κ,
        Ux*sΔψ + Uy*cΔψ
    )
end
(B::LateralTrackingBicycleModel{T})(q::StaticVector{4}, up::StaticVector{7}) where {T} = B(q, up[SUnitRange(1,3)], up[SUnitRange(4,7)])

# Uy/r Stability Envelope
@inline function stable_limits(B::AbstractBicycleModel, Ux, Fxf, Fxr)
    L, a, b, h, m, μ, Cαf, Cαr, G = B.L, B.a, B.b, B.h, B.m, B.μ, B.Cαf, B.Cαr, B.G

    Fx = Fxf + Fxr
    Fy_grade = 0
    Fzf = (m*G*b - h*Fx)/L
    Fzr = (m*G*a + h*Fx)/L
    Ff_max = μ*Fzf
    Fr_max = μ*Fzr
    Fyf_max = abs(Fxf) > Ff_max ? zero(Ff_max) : sqrt(Ff_max*Ff_max - Fxf*Fxf)
    Fyr_max = abs(Fxr) > Fr_max ? zero(Fr_max) : sqrt(Fr_max*Fr_max - Fxr*Fxr)
    tanαf_slide = 3*Fyf_max/Cαf
    tanαr_slide = 3*Fyr_max/Cαr
    αf_slide = atan(tanαf_slide)
    αr_slide = atan(tanαr_slide)

    # https://ddl.stanford.edu/sites/default/files/publications/2012_Thesis_Bobier_A_Phase_Portrait_Approach_to_Vehicle_Stabilization_and_Envelope_Control.pdf
    δ_max = atan(L*( μ*G + Fy_grade/m)/(Ux*Ux) - tanαr_slide) + αf_slide    # corresponds to U̇y = 0; αr < 0; αf < 0; δ > 0; r > 0
    δ_min = atan(L*(-μ*G + Fy_grade/m)/(Ux*Ux) + tanαr_slide) - αf_slide    # corresponds to U̇y = 0; αr > 0; αf > 0; δ < 0; r < 0
    rC  = ( μ*G + Fy_grade/m)/Ux
    UyC = -Ux*tanαr_slide + b*rC
    rD  = Ux/L*(tan( αf_slide + δ_max) - tanαr_slide)     # αr > 0; αf > 0; δ > 0;
    UyD = Ux*tanαr_slide + b*rD
    mCD = (rD - rC)/(UyD - UyC)
    rE  = Ux/L*(tan(-αf_slide + δ_min) + tanαr_slide)    # αr < 0; αf < 0; δ < 0;
    UyE = -Ux*tanαr_slide + b*rE
    rF  = (-μ*G + Fy_grade/m)/Ux
    UyF = Ux*tanαr_slide + b*rF
    mEF = (rF - rE)/(UyF - UyE)

    H_veh = @SMatrix [ 1/Ux -b/Ux;    # β max (≈Uy/Ux)
                      -1/Ux  b/Ux;    # β min (≈Uy/Ux)
                       -mCD     1;    # r max
                        mEF    -1]    # r min
    G_veh = SVector(αr_slide, αr_slide, rC - UyC*mCD, -rF + UyF*mEF)
    δ_min, δ_max, H_veh, G_veh
end

# Bicycle Model with Combined Longitudinal Force Input (Fx = Fxf + Fxr)
@maintain_type struct BicycleControl2{T} <: FieldVector{2,T}
    δ::T    # steering angle
    Fx::T   # combined front/rear tire longitudinal force (sum of each in respective tire frames; approximately, total is applied longitudinally in body frame)
end
BicycleControl2((δ, Fxf, Fxr)::StaticVector{3}) = BicycleControl2(δ, Fxf + Fxr)

struct LongitudinalActuationParams{T}
    fwd_frac::T
    rwd_frac::T
    fwb_frac::T
    rwb_frac::T
end
LongitudinalActuationParams(vehicle::Dict) = LongitudinalActuationParams((vehicle[n] for n in fieldnames(LongitudinalActuationParams))...)
function longitudinal_tire_forces(LP::LongitudinalActuationParams, Fx)
    Fx > 0 ? (Fx*LP.fwd_frac, Fx*LP.rwd_frac) :
             (Fx*LP.fwb_frac, Fx*LP.rwb_frac)
end
BicycleControl(LP::LongitudinalActuationParams, (δ, Fx)::StaticVector{2}) = BicycleControl(δ, longitudinal_tire_forces(LP, Fx)...)

struct ControlLimits{T}
    Fx_max::T
    Fx_min::T
    Px_max::T
    δ_max::T
    κ_max::T
end
ControlLimits(vehicle::Dict) = ControlLimits((vehicle[n] for n in fieldnames(ControlLimits))...)
function apply_control_limits(CL::ControlLimits, (δ, Fx)::StaticVector{2}, Ux)
    Fx_max, Fx_min, Px_max, δ_max = CL.Fx_max, CL.Fx_min, CL.Px_max, CL.δ_max
    Ux = ForwardDiff.value(Ux)    # important for ForwardDiff/linearization, since Ux is technically a state variable
    BicycleControl2(clamp(δ, -δ_max, δ_max),
                    max(min(Fx, Fx_max, Px_max/Ux), Fx_min))
end

struct VehicleModel{T,B<:AbstractBicycleModel{T}} <: DifferentialDynamics
    bicycle_model::B
    longitudinal_params::LongitudinalActuationParams{T}
    control_limits::ControlLimits{T}
end
VehicleModel(vehicle::Dict, B=BicycleModel(vehicle)) = VehicleModel(B, LongitudinalActuationParams(vehicle), ControlLimits(vehicle))

get_Ux(B::BicycleModel, q, u, p) = q[4]
get_Ux(B::TrackingBicycleModel, q, u, p) = q[2]
get_Ux(B::LateralTrackingBicycleModel, q, u, p) = p[1]
function (X::VehicleModel{T})(q, u, p=zeros(SVector{4,T})) where {T}    # TODO: not necessarily always the right-sized p
    Ux = get_Ux(X.bicycle_model, q, u, p)
    X.bicycle_model(q, BicycleControl(X.longitudinal_params, apply_control_limits(X.control_limits, u, Ux)), p)
end
(X::VehicleModel{T,BicycleModel{T}})(q::StaticVector{6}, ur::StaticVector{6}) where {T} = X(q, ur[SUnitRange(1,2)], ur[SUnitRange(3,6)])
(X::VehicleModel{T,TrackingBicycleModel{T}})(q::StaticVector{6}, ur::StaticVector{6}) where {T} = X(q, ur[SUnitRange(1,2)], ur[SUnitRange(3,6)])
(X::VehicleModel{T,LateralTrackingBicycleModel{T}})(q::StaticVector{4}, ur::StaticVector{6}) where {T} = X(q, ur[SUnitRange(1,2)], ur[SUnitRange(3,6)])

# Estimates for Longitudinal Dynamics
function steady_state_estimates(X::VehicleModel{T}, V, A_tan, κ;
                                num_iters=4, r=V*κ, β0=T(0), δ0=T(0), Fyf0=T(0)) where {T}
    # assuming vehicle CM tracks nominal trajectory exactly; constant sideslip w.r.t. nominal trajectory
    B, LP, CL = X.bicycle_model, X.longitudinal_params, X.control_limits
    L, a, b, h, m, Izz, μ, G = B.L, B.a, B.b, B.h, B.m, B.Izz, B.μ, B.G
    Cαf, Cαr, Cd0, Cd1, Cd2 = B.Cαf, B.Cαr, B.Cd0, B.Cd1, B.Cd2
    fwd_frac, rwd_frac, fwb_frac, rwb_frac = LP.fwd_frac, LP.rwd_frac, LP.fwb_frac, LP.rwb_frac
    Fx_max, Px_max = CL.Fx_max, CL.Px_max

    A_rad = V*V*κ
    A_mag = hypot(A_tan, A_rad)
    A_max = μ*G
    if A_mag > A_max             # nominal trajectory total acceleration exceeds friction limit
        if abs(A_rad) > A_max    # nominal trajectory lateral acceleration exceeds friction limit
            A_rad = A_max*sign(A_rad)
            A_tan = zero(A_rad)
            # error("TODO: case when nominal trajectory lateral acceleration exceeds friction limit")
        else    # prioritize radial acceleration for path tracking; reduce intended acceleration along path to compensate
            A_tan = sqrt(A_max*A_max - A_rad*A_rad)*sign(A_tan)
        end
    end
    ṙ = A_tan*κ

    i = 1
    β, δ, Fyf = β0, δ0, Fyf0
    local Ux, Uy, Fxr, Fxf
    while true
        sβ, cβ = sincos(β)
        sδ, cδ = sincos(δ)
        Ux, Uy = V*cβ, V*sβ
        Fx_drag = -Cd0 - Ux*(Cd1 + Cd2*Ux)
        Fx_grade = 0    # TODO: figure out how roll/pitch are ordered
        Fy_grade = 0

        Ax = A_tan*cβ - A_rad*sβ             # Ax = U̇x - r*Uy
        Ay = A_tan*sβ + A_rad*cβ             # Ay = U̇y + r*Ux
        Fx = Ax*m - Fx_drag - Fx_grade       # Fx = F̃xf + Fxr = Fxf*cδ - Fyf*sδ + Fxr
        Fx = min(Fx, min(Fx_max, Px_max/Ux)*(rwd_frac + fwd_frac*cδ) - Fyf*sδ)    # braking force saturated by friction, not engine, limits
        Fzr, Fzf = (m*G*a + h*Fx)/L, (m*G*b - h*Fx)/L
        Fr_max, Ff_max = μ*Fzr, μ*Fzf

        Fxr = clamp((Fx + Fyf*sδ)*ifelse(Fx > 0, rwd_frac/(rwd_frac + fwd_frac*cδ),
                                                 rwb_frac/(rwb_frac + fwb_frac*cδ)),
                    -Fr_max, Fr_max)
        Fyr_max = sqrt(Fr_max*Fr_max - Fxr*Fxr)
        Fyr = (Ay*m - Fy_grade - ṙ*Izz/a)/(1 + b/a)
        Fyr = clamp(Fyr, -Fyr_max, Fyr_max)    # TODO: maybe reduce A_tan in the case that Fyr exceeds Fyr_max
        tanαr = _invfialatiremodel(Fyr, Cαr, Fyr_max)

        F̃xf = clamp(Fx - Fxr, -Ff_max, Ff_max)
        F̃yfmax = sqrt(Ff_max*Ff_max - F̃xf*F̃xf)
        F̃yf = clamp((b*Fyr + ṙ*Izz)/a, -F̃yfmax, F̃yfmax)
        Fxf = F̃xf*cδ + F̃yf*sδ
        Fyf = F̃yf*cδ - F̃xf*sδ
        Fyf_max = sqrt(Ff_max*Ff_max - Fxf*Fxf)
        αf = atan(_invfialatiremodel(Fyf, Cαf, Fyf_max))
        δ = atan(Uy + a*r, Ux) - αf

        if i == num_iters
            Ax = (Fxf*cδ - Fyf*sδ + Fxr + Fx_drag + Fx_grade)/m
            Ay = (Fyf*cδ + Fxf*sδ + Fyr + Fy_grade)/m
            A_tan = Ax*cβ + Ay*sβ
            break
        end
        i = i + 1
        β = atan(tanαr + b*r/Ux)
    end

    sβ, cβ = sincos(β)
    Ux, Uy = V*cβ, V*sβ
    (β=β, Ux=Ux, Uy=Uy, r=r, A=A_tan, δ=δ, Fxf=Fxf, Fxr=Fxr)
end


#=
"""
--------------------------------------------------------------------------------
Function: XP_Bisection
--------------------------------------------------------------------------------
Description: Bisects box in either X or P dimension.
--------------------------------------------------------------------------------
Inputs:
h           function - h(z,p)
hj          function - Jacobian of h(z,p) w.r.t. z
CurrNode    Vector{Interval{Float64}} - Current X node
P           Vector{Interval{Float64}} - Current P node
NextNode    Vector{Interval{Float64}} - New X node
NextNodeP   Vector{Interval{Float64}} - New P node 1
Ptemp       Vector{Interval{Float64}} - New P node 2
PIflag      Bool - Partial Inclusion Flag
Iflag       Bool - Inclusion flag
Eflag       Bool - Exclusion flag
ptol        Float64 - tolerance for P dimension size
pbisect     Bool - bisect based on relative width?
Pstart      Vector{Interval{Float64}} - initial P size
--------------------------------------------------------------------------------
Returns:
bp          Int64 - Cut type flag
CurrNode    Vector{Interval{Float64}} - Current X node
P           Vector{Interval{Float64}} - Current P node
NextNode    Vector{Interval{Float64}} - New X node
NextNodeP   Vector{Interval{Float64}} - New P node 1
Ptemp       Vector{Interval{Float64}} - New P node 2
--------------------------------------------------------------------------------
"""

"""
--------------------------------------------------------------------------------
Function: bisectDirX
--------------------------------------------------------------------------------
Description: Determine direction and position of cut in X dimension.
--------------------------------------------------------------------------------
Inputs:
h        function - h(z,p)
X        Vector{Interval{Float64}} - for state variables
P        Vector{Interval{Float64}} - for decision variables
x        Currently Unused
PIFlag   Bool - Flag for partial inclusion.
--------------------------------------------------------------------------------
Returns:
jmax            Int64 - The dimension to cut X in
xCut            Float64 - Cut value for X
piflagtemp      Bool - Partial Inclusion Flag
--------------------------------------------------------------------------------
"""

"""
--------------------------------------------------------------------------------
Function: bisectDirP
--------------------------------------------------------------------------------
Description: Determines the direction to branch P in based on absolute/relative
width and tolerance for P dimension size.
--------------------------------------------------------------------------------
Inputs:
P         Vector{Interval{Float64}} - P Box to bisect
ptol      Float64 - tolerance for P dimension size
pbisect   Bool - bisect based on relative width?
Pstart    Vector{Interval{Float64}} - initial P size
--------------------------------------------------------------------------------
Returns:
jmax - The dimension to branch P in.
--------------------------------------------------------------------------------
"""

"""
--------------------------------------------------------------------------------
Function:  partIncProc
--------------------------------------------------------------------------------
Description:
Partial inclusion check (uses Newton type operator test with Miranda)
--------------------------------------------------------------------------------
Inputs:
h         function - h(z,p)
hj        function - Jacobian of h(z,p) w.r.t. z
X0        Vector{Interval{Float64}} - state variables
P         Vector{Interval{Float64}} - decision variables
exDim     Int64 - Dimension flag
Iflag     Bool - Inclusion flag
Eflag     Bool - Exclusion flag
PIflag    Bool - Partial Inclusion Flag
--------------------------------------------------------------------------------
Returns:
X         IntervalBox - State variables
Iflag     Bool - Inclusion flag
Eflag     Bool - Exclusion flag
PIflag    Bool - Partial Inclusion Flag
--------------------------------------------------------------------------------
"""
=#

#=
function XP_Bisection(h,hj,CurrNode,P,NextNode,NextNodeP,Ptemp,PIflag,Iflag,Eflag,ptol,pbisect,Pstart)
  #println("CurrNode Top: ", CurrNode)
  pPart = false
  xCut = 0.15
  pCutL = 0.0
  pCutR = 0.0
  IflagL = copy(Iflag)
  IflagR = copy(Iflag)
  EflagL = copy(Eflag)
  EflagR = copy(Eflag)
  XLeft = copy(CurrNode)
  XRight = copy(NextNode)
  PLeft = copy(P)
  PRight = copy(NextNodeP)
  Pmid = copy(Ptemp)
  j = 0
  jp = 0
  bp = 0
  j, xCut,PIflag = bisectDirX(h,CurrNode,P,xCut,PIflag)
  #println("bisectDirX j: ", j)
  #println("bisectDirX xCut: ", xCut)
  #println("bisectDirX PIflag: ", PIflag)
  PIflagL = PIflag
  PIflagR = PIflag
  #println("PIflag:  ", PIflag)
  if (PIflag == false)
    XLeft[j] = Interval(CurrNode[j].lo,xCut)
    XRight[j] = Interval(xCut,CurrNode[j].hi)
    CurrNode = copy(XLeft)
    NextNode = copy(XRight)
    println("bisect trace 1")
    return 1,CurrNode,P,NextNode,NextNodeP,Ptemp
  else
    jp = bisectDirP(P,ptol,pbisect,Pstart)
    if (jp != -1)
      if pPart==true
        bp,jp,pCutL,pCutR = branchP(CurrNode,P,jp,pCutL,pCutR)
      else
        bp = 1
      end
      #println("jp")
      if (bp == 1)
        PLeft[jp] = Interval(P[jp].lo,mid(P[jp]))
        PRight[jp] = Interval(mid(P[jp]),NextNodeP[jp].hi)
        P = copy(PLeft)
        NextNodeP = copy(PRight)
        #println("bisect trace 2")
        return 2,CurrNode,P,NextNode,NextNodeP,Ptemp
      elseif (bp == 3)
        PLeft[jp] = Interval(P[jp].lo,pCutL)
        Pmid[jp] = Interval(pCutL,pCutR)
        PRight[jp] = Interval(pCutR,P[jp].hi)
        P = copy(PLeft)
        Ptemp = copy(Pmid)
        NextNodeP = copy(PRight)
        #println("bisect trace 3")
        return 3,CurrNode,P,NextNode,NextNodeP,Ptemp
      end
    end
  end
  #println("jp = :",jp)
  #println("bp = :",bp)
  if (false) # originally just ((jp == -1) && (bp != 1))
    xCut = xCut*CurrNode[j].lo + (1.0-xCut)*CurreNode[j].hi
    XLeft[j] = Interval(CurrNode[j].lo,xCut)
    XRight[j] = Interval(xCut,CurrNode[j].hi)
    k,rhoL = spectralR(hj,XLeft,P,rhoL)
    k,rhoR = spectralR(hj,XRight,P,rhoR)
    if (rhoL)
      XLeft,Iflag,Eflag,PIflag = partIncProc(h,hj,XLeft,P,j,Iflag,Eflag,PIflag)
    end
    if (rhoR)
      XRight,Iflag,Eflag,PIflag = partIncProc(h,hj,XRight,P,j,Iflag,Eflag,PIflag)
    end
    if ((~PIflagL) && (~PIflagR))
      PIflag = false
    end
  end

  CurrNode = copy(XLeft)
  NextNode = copy(XRight)
  #println("bisect trace 4")
  return bp,CurrNode,P,NextNode,NextNodeP,Ptemp
end

function bisectDirX(h,X,P,x,PIFlag)

  piflagtemp = copy(PIFlag)
  v = 1
  iterm = 4
  jmax = 1
  k = -1
  maxL = 0.0
  disc = 20
  Xold = copy(X)
  Xnew = copy(X)
  xCut = 0.0
  maxL,jmax = findmax(diam.(X))
  if (v == 1)
    rad = diam(X[jmax])/2.0
    for l=0:(disc/2)
      xCut = X[jmax].lo+rad/(iterm) + l*(diam(X[jmax])-2.0*rad/(iterm))/(disc-1)
      Xnew[jmax] = xCut
      k = Miranda(h,Xnew,P)
      if (k>=0)
        return jmax, xCut, piflagtemp
      end
      xCut = X[jmax].hi-rad/(iterm) - l*(diam(X[jmax])-2.0*rad/(iterm))/(disc-1)
      Xnew[jmax] = xCut
      k = Miranda(h,Xnew,P)
      if (k>=0)
        return jmax, xCut, piflagtemp
      else
        Xnew[jmax] = Xold[jmax]
      end
    end
  end

  if (k == -1)
    v = 2
  end

  if (v == 2)
    iterm = iterm*4
    Xnew = copy(Xold)
    disc = 4*disc
    for i = 1:length(X)
        rad = diam(X[i])/2.0
        for l = 0:(disc/2)
          xCut = X[i].lo+rad/(iterm) + l*(diam(X[i])-2.0*rad/(iterm))/(disc-1)
          Xnew[i] = xCut
          k = Miranda(h,Xnew,P)
          if (k>=0)
            return i, xCut, piflagtemp
          end
          xCut = X[i].hi-rad/(iterm) - l*(diam(X[i])-2.0*rad/(iterm))/(disc-1)
          Xnew[i] = xCut
          k = Miranda(h,Xnew,P)
          if (k>=0)
            return i, xCut, piflagtemp
          else
            Xnew[i] = Xold[i]
          end
        end
    end
    if (k == -1)
      piflagtemp = true
    end
  end
  return jmax, xCut, piflagtemp
end

function bisectDirP(P,ptol,pbisect,Pstart)
  jmax = 1
  maxL = 0.0
  if (~pbisect)
    for i=1:length(P)
      if (maxL < diam(P[i]))
        maxL = diam(P[i])
        jmax = i
      end
    end
    if (maxL < ptol)
      return -1
    end
  else
    for i=1:length(P)
      if (maxL < diam(P[i])/diam(Pstart[i]))
        maxL = diam(P[i])/diam(Pstart[i])
        jmax = i
      end
    end
    if (maxL < ptol)
      return -1
    end
  end
  return jmax
end

# Unfinish bisection based on spectral radius
function branchP(X,P,jmax,pCutL,pCutR)

  Plocal = P
  maxD = 0.0
  try
    Pf = [P[j] for j=1:length(P)]
    p = [abs(Pf[j].lo+Pf[j].hi)/2.0 for j=1:length(P)]
    Xf = [X[j] for j=1:length(X)]
    rhok = rho
    fpsecR(Xf,Pf,rho)
  catch
  end

end

function partIncProc(h,hj,X0,P,exDim,Iflag,Eflag,PIflag)

  nx = length(X)
  S1 = Interval(0.0)
  S2 = Interval(0.0)
  k = 0
  incl = -1
  side = 0
  inclusion = [false for i=1:length(X)]
  exclusion = false
  Iflag = false
  Eflag = false

  X = copy(X0)
  Xpart = copy(X0)
  Xpart[exDim] = xCut
  if (X[exDim].lo == xCut)
    side = 1
  end

  x_mid = mid.(X)
  H = h(x_mid,P)
  J = hj(X,P)
  Y = Preconditioner(hj,X,P,jac="User")
  B = Y*H
  M = Y*J
  Xi = X

  for i=1:nx
    S1 = S2 = Interval(0.0)
    for j=1:nx
      if (j<i)
        S1 += M[i,j]*(X[j]-x_mid[j])
      elseif (j>i)
        S2 += M[i,j]*(X[j]-x_mid[j])
      end
    end
    N[i] = x_mid[i] - (B[i]+S1+S2)/M[i,i]
    if N[i] ⊆ X[i]
      inclusion[i] = true
    else
      inclusion[i] = false
    end
    if ~isdisjoint(N[i],X[i])
      if (i == exDim)
        wL = X[i].lo - N[i].lo
        wU = N[i].hi - X[i].hi
        if (side == 0)
          X[i] = Interval(X[i].lo,X[i].hi+0.25*wU)
        else
          X[i] = Interval(X[i].lo-wL*0.25,X[i].hi)
        end
      else
        X[i] = N[i] ∩ X[i]
      end
    else
      exclusion = true
      break
    end
  end
  k += 1

  # checks if all components are included
  if (Iflag == false)
    for i=1:nx
      if (inclusion[i] == true)
        Iflag = true
        continue
      elseif (inclusion[i] == false)
        Iflag = false
        break
      end
    end
  end

  Xpart = X
  if (side == 0)
    Xpart[exDim] = X[exDim].hi
    if (Xpart[exDim]<=xCut)
      PIflag = false
      return X,Iflag,Eflag,PIflag #FIX ME
    end
  else
    Xpart[exDim] = X[exDim].lo
    if (Xpart[exDim]>=xCut)
      PIflag = false
      return X,Iflag,Eflag,PIflag #FIX ME
    end
  end
  incl = Miranda(h,Xpart,P)

  return X,Iflag,Eflag,PIflag #FIX ME
end
=#


using SparseArrays

function numerical_fluxes(flux, mesh, w, scheme, dt=0.0)
    map(i_face -> numerical_flux(flux, mesh, w, scheme, i_face, dt), collect(inner_faces(mesh)))
end

"""Sparse matrix to sum the flux on both side of a face and get the change in a cell.

In 1D: sparse matrix.
For 2D cartesian mesh, it should be a higher order mesh, but since there is no higher dimensional sparse array, the field needs to be reshaped to a vector...
"""
function inner_faces_to_cells_matrix(mesh::AbstractCartesianMesh)
    D = spzeros(FiniteVolumes.nb_cells(mesh), FiniteVolumes.nb_inner_faces(mesh))
    li = LinearIndices(cell_centers(mesh))
    for (i, face) in enumerate(FiniteVolumes.inner_faces(mesh))
        i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, face)
        D[li[i_cell_1], i] -= face_area(mesh, face) / cell_volume(mesh, i_cell_1)
        D[li[i_cell_2], i] += face_area(mesh, face) / cell_volume(mesh, i_cell_1)
    end
    D
end

function update!(w, Φ, mesh, dt=0.0)
    for (i, face) in enumerate(inner_faces(mesh))
        i_cell_1, i_cell_2 = cells_next_to_inner_face(mesh, face)
        w[i_cell_1] -= dt*Φ[i] * face_area(mesh, face) / cell_volume(mesh, i_cell_1)
        w[i_cell_2] += dt*Φ[i] * face_area(mesh, face) / cell_volume(mesh, i_cell_2)
    end
end


# ==============================================================================
# Building model
# See https://easychair.org/publications/open/4cGr
# ==============================================================================

using MAT, Reachability, MathematicalSystems, SX
using SparseArrays, LinearAlgebra

# ==============================
# Load model
# ==============================
file = @relpath "SX/Building_more_decimals.xml"
H = readsxmodel(file, ST=ConstrainedLinearControlContinuousSystem)

# ===================
# Time-varying input
# ===================

n = size(H.modes[1].A, 1)-1 # the sx model has "time" as a state variable
@assert n == 48
A = H.modes[1].A[1:n, 1:n]
B = hcat(H.modes[1].B[1:n, 1])
X = nothing
U = Hyperrectangle(low=[0.8], high=[1.0])
S = ConstrainedLinearControlContinuousSystem(A, B, X, U)

# specify initial states
center_X0 = [fill(0.000225, 10); fill(0.0, 38)]
radius_X0 = [fill(0.000025, 10); fill(0.0, 14); 0.0001; fill(0.0, 23)]
X0 = Hyperrectangle(center_X0, radius_X0)
build_TV = InitialValueProblem(S, X0)

# specifications
pBDS01 = SafeStatesProperty(
    LinearConstraint(sparsevec([25], [1.0], n), 0.0051)  # x25 <= 0.0051
    )

time_horizon = 20.0

# ===================
# Constant input
# ===================

A = Reachability.add_dimension(A) # add an extra zero row and column
S = LinearContinuousSystem(A)
X0 = X0 × U
build_CONST = InitialValueProblem(S, X0)

# specifications
pBLDC01 = SafeStatesProperty(
    LinearConstraint(sparsevec([25], [1.0], n+1), 0.0051)  # x25 <= 0.0051
    )


########################################################################################
# Caches
########################################################################################
abstract type ODEFiltersCache <: OrdinaryDiffEq.OrdinaryDiffEqCache end
mutable struct GaussianODEFilterCache{
    RType,
    ProjType,
    SolProjType,
    PType,
    PIType,
    EType,
    uType,
    duType,
    xType,
    AType,
    QType,
    matType,
    diffusionType,
    diffModelType,
    measType,
    puType,
    llType,
    CType,
    rateType,
    UF,
    JC,
    uNoUnitsType,
} <: ODEFiltersCache
    # Constants
    d::Int                  # Dimension of the problem
    q::Int                  # Order of the prior
    A::AType
    Q::QType
    Ah::AType
    Qh::QType
    diffusionmodel::diffModelType
    R::RType
    Proj::ProjType
    SolProj::SolProjType
    # Also mutable
    P::PType
    PI::PIType
    E0::EType
    E1::EType
    E2::EType
    # Mutable stuff
    u::uType
    u_pred::uType
    u_filt::uType
    tmp::uType
    atmp::uNoUnitsType
    x::xType
    x_pred::xType
    x_filt::xType
    x_tmp::xType
    x_tmp2::xType
    measurement::measType
    m_tmp::measType
    pu_tmp::puType
    H::matType
    du::duType
    ddu::matType
    K1::matType
    K2::matType
    G1::matType
    G2::matType
    covmatcache::matType
    default_diffusion::diffusionType
    local_diffusion::diffusionType
    global_diffusion::diffusionType
    err_tmp::duType
    log_likelihood::llType
    C1::CType
    C2::CType
    du1::rateType
    uf::UF
    jac_config::JC
end

function OrdinaryDiffEq.alg_cache(
    alg::GaussianODEFilter,
    u,
    rate_prototype,
    uEltypeNoUnits,
    uBottomEltypeNoUnits,
    tTypeNoUnits,
    uprev,
    uprev2,
    f,
    t,
    dt,
    reltol,
    p,
    calck,
    ::Val{IIP},
) where {IIP}
    initialize_derivatives = true

    if u isa Number
        error("We currently don't support scalar-valued problems")
    end

    is_secondorder_ode = f isa DynamicalODEFunction
    if is_secondorder_ode
        @warn "Assuming that the given ODE is a SecondOrderODE. If this is not the case, e.g. because it is some other dynamical ODE, the solver will probably run into errors!"
    end

    q = alg.order
    d = is_secondorder_ode ? length(u[1, :]) : length(u)
    D = d * (q + 1)

    u_vec = u[:]
    t0 = t

    uType = typeof(u)
    uElType = eltype(u_vec)
    matType = Matrix{uElType}

    # Projections
    Proj = projection(d, q, uElType)
    E0, E1, E2 = Proj(0), Proj(1), Proj(2)
    @assert f isa AbstractODEFunction
    SolProj = f isa DynamicalODEFunction ? [Proj(1); Proj(0)] : Proj(0)

    # Prior dynamics
    @assert alg.prior == :ibm "Only the ibm prior is implemented so far"
    P, PI = init_preconditioner(d, q, uElType)

    A, Q = ibm(d, q, uElType)

    initial_variance = 1.0 * ones(uElType, D)
    x0 = Gaussian(zeros(uElType, D), SRMatrix(Matrix(Diagonal(sqrt.(initial_variance)))))

    # Measurement model
    R = zeros(uElType, d, d)

    # Pre-allocate a bunch of matrices
    h = zeros(uElType, d)
    H = f isa DynamicalODEFunction ? copy(E2) : copy(E1)
    du = f isa DynamicalODEFunction ? similar(u[2, :]) : similar(u)
    ddu = f isa DynamicalODEFunction ? zeros(uElType, d, 2d) : zeros(uElType, d, d)
    v, S = similar(h), similar(ddu)
    v = similar(h)
    S =
        alg isa EK0 ? SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) :
        SRMatrix(zeros(uElType, d, D), zeros(uElType, d, d))
    measurement = Gaussian(v, S)
    pu_tmp =
        f isa DynamicalODEFunction ?
        Gaussian(zeros(uElType, 2d), SRMatrix(zeros(uElType, 2d, D))) : similar(measurement)
    K = zeros(uElType, D, d)
    G = zeros(uElType, D, D)
    C1 = SRMatrix(zeros(uElType, D, 2D), zeros(uElType, D, D))
    C2 = SRMatrix(zeros(uElType, D, 3D), zeros(uElType, D, D))
    covmatcache = similar(G)

    if alg isa EK1FDB
        H = [E1; E2]
        v = [v; v]
        S = SRMatrix(zeros(uElType, 2d, D), zeros(uElType, 2d, 2d))
        measurement = Gaussian(v, S)
        K = zeros(uElType, D, 2d)
    end

    diffmodel = alg.diffusionmodel
    initdiff = initial_diffusion(diffmodel, d, q, uEltypeNoUnits)
    copy!(x0.Σ, apply_diffusion(x0.Σ, initdiff))

    Ah, Qh = copy(A), copy(Q)
    u_pred = similar(u)
    u_filt = similar(u)
    tmp = similar(u)
    x_pred = copy(x0)
    x_filt = similar(x0)
    x_tmp = similar(x0)
    x_tmp2 = similar(x0)
    m_tmp = similar(measurement)
    K2 = similar(K)
    G2 = similar(G)
    err_tmp = similar(du)

    # Things for calc_J
    uf =
        IIP == true ? OrdinaryDiffEq.UJacobianWrapper(f, t, p) :
        OrdinaryDiffEq.UDerivativeWrapper(f, t, p)
    du1 = similar(rate_prototype)
    dw1 = zero(u)
    jac_config = OrdinaryDiffEq.build_jac_config(alg, f, uf, du1, uprev, u, tmp, dw1)
    atmp = similar(u, uEltypeNoUnits)

    return GaussianODEFilterCache{
        typeof(R),
        typeof(Proj),
        typeof(SolProj),
        typeof(P),
        typeof(PI),
        typeof(E0),
        uType,
        typeof(du),
        typeof(x0),
        typeof(A),
        typeof(Q),
        matType,
        typeof(initdiff),
        typeof(diffmodel),
        typeof(measurement),
        typeof(pu_tmp),
        uEltypeNoUnits,
        typeof(C1),
        typeof(du1),
        typeof(uf),
        typeof(jac_config),
        typeof(atmp),
    }(
        # Constants
        d,
        q,
        A,
        Q,
        Ah,
        Qh,
        diffmodel,
        R,
        Proj,
        SolProj,
        P,
        PI,
        E0,
        E1,
        E2,
        # Mutable stuff
        u,
        u_pred,
        u_filt,
        tmp,
        atmp,
        x0,
        x_pred,
        x_filt,
        x_tmp,
        x_tmp2,
        measurement,
        m_tmp,
        pu_tmp,
        H,
        du,
        ddu,
        K,
        K2,
        G,
        G2,
        covmatcache,
        initdiff,
        initdiff * NaN,
        initdiff * NaN,
        err_tmp,
        zero(uEltypeNoUnits),
        C1,
        C2,
        du1,
        uf,
        jac_config,
    )
end


using LinearAlgebra, FFTW

function create_testfuncs(g::OneDGrid)
    if nx <= 8; error("nx must be > 8"); end

    m = 5
    k0 = g.k[2] # Fundamental wavenumber

    # Test function
    φ = π/3
    f1 = cos.(m*k0*g.x .+ φ)
    f1h = fft(f1)
    f1hr = rfft(f1)

    f1hr_mul = zeros(Complex{eltype(g.x)}, g.nkr)
    mul!(f1hr_mul, g.rfftplan, f1)

    # Theoretical values of the fft and rfft
    f1h_th = zeros(Complex{Float64}, size(f1h))
    f1hr_th = zeros(Complex{Float64}, size(f1hr))

    for i in 1:g.nk
      if abs(real(g.k[i])) == m*k0
        f1h_th[i] = -exp(sign(real(g.k[i]))*im*φ)*g.nx/2
      end
    end

    for i in 1:g.nkr
      if abs(real(g.k[i]))==m*k0
        f1hr_th[i] = -exp(sign(real(g.kr[i]))*im*φ)*g.nx/2
      end
    end

    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th
end

function create_testfuncs(g::TwoDGrid)
    if nx <= 8; error("nx must be > 8"); end

    m, n = 5, 2
    k0 = g.k[2] # fundamental wavenumber
    l0 = g.l[2] # fundamental wavenumber

    # some test functions
    f1 = @. cos(m*k0*g.X) * cos(n*l0*g.Y)
    f2 = @. sin(m*k0*g.X + n*l0*g.Y)

    f1h = fft(f1)
    f2h = fft(f2)
    f1hr = rfft(f1)
    f2hr = rfft(f2)

    f1hr_mul = zeros(Complex{eltype(g.x)}, (g.nkr, g.nl))
    f2hr_mul = zeros(Complex{eltype(g.x)}, (g.nkr, g.nl))
    mul!(f1hr_mul, g.rfftplan, f1)
    mul!(f2hr_mul, g.rfftplan, f2)

    # Theoretical results
    f1h_th = zeros(Complex{eltype(g.x)}, size(f1h))
    f2h_th = zeros(Complex{eltype(g.x)}, size(f2h))
    f1hr_th = zeros(Complex{eltype(g.x)}, size(f1hr))
    f2hr_th = zeros(Complex{eltype(g.x)}, size(f2hr))

    for j in 1:g.nl, i in 1:g.nk
      if ( abs(real(g.K[i, j])) == m*k0 && abs(real(g.L[i, j])) == n*l0 )
        f1h_th[i, j] = - g.nx*g.ny/4
      end
      if ( real(g.K[i, j]) == m*k0 && real(g.L[i, j]) == n*l0 )
        f2h_th[i, j] = -g.nx*g.ny/2
      elseif ( real(g.K[i, j]) == -m*k0 && real(g.L[i, j]) == -n*l0 )
        f2h_th[i, j] = g.nx*g.ny/2
      end
    end
    f2h_th = -im*f2h_th;

    for j in 1:g.nl, i in 1:g.nkr
      if ( abs(g.Kr[i, j])==m*k0 && abs(g.L[i, j])==n*l0 )
        f1hr_th[i, j] = - g.nx*g.ny/4
      end
      if ( real(g.Kr[i, j]) == m*k0 && real(g.L[i, j]) == n*l0 )
        f2hr_th[i, j] = -g.nx*g.ny/2
      elseif ( real(g.Kr[i, j]) == -m*k0 && real(g.L[i, j]) == -n*l0 )
        f2hr_th[i, j] = g.nx*g.ny/2
      end
    end
    f2hr_th = -im*f2hr_th;

    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul,
      f1h_th, f1hr_th, f2h_th, f2hr_th
end

# -----------------------------------------------------------------------------
# FFT's TEST FUNCTIONS

rtol = 1e-12

function test_fft_cosmx(g::OneDGrid)
    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th = create_testfuncs(g)
    isapprox(f1h, f1h_th, rtol=rtol)
end

function test_rfft_cosmx(g::OneDGrid)
    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th = create_testfuncs(g)
    isapprox(f1hr, f1hr_th, rtol=rtol)
end

function test_rfft_mul_cosmx(g::OneDGrid)
    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th = create_testfuncs(g)
    isapprox(f1hr_mul, f1hr_th, rtol=rtol)
end

function test_fft_cosmxcosny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul,
            f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f1h, f1h_th, rtol=rtol)
end

function test_rfft_cosmxcosny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul,
            f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f1hr, f1hr_th, rtol=rtol)
end

function test_rfft_mul_cosmxcosny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul,
            f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f1hr_mul, f1hr_th, rtol=rtol)
end

function test_fft_sinmxny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mulPk, f2hr_mul,
            f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f2h, f2h_th, rtol=rtol)
end

function test_rfft_sinmxny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul,
            f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f2hr, f2hr_th, rtol=rtol)
end

function test_rfft_mul_sinmxny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul,
            f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f2hr_mul, f2hr_th, rtol=rtol)
end

# Test 1D grid
nx = 32             # number of points
Lx = 2π             # Domain width
g1 = OneDGrid(nx, Lx)

# Test 2D rectangular grid
nx, ny = 32, 64     # number of points
Lx, Ly = 2π, 3π     # Domain width
g2 = TwoDGrid(nx, Lx, ny, Ly)

@test test_fft_cosmx(g1)
@test test_rfft_cosmx(g1)
@test test_rfft_mul_cosmx(g1)

@test test_fft_cosmxcosny(g2)
@test test_rfft_cosmxcosny(g2)
@test test_rfft_mul_cosmxcosny(g2)
@test test_fft_sinmxny(g2)
@test test_rfft_sinmxny(g2)
@test test_rfft_mul_sinmxny(g2)


function test_tridiagonalsolve()
    # Test problem:
    # A = [[1 4 0 0]; [3 4 1 0]; [0 2 3 4]; [0 0 1 3]]
    # x = [1; 2; 3; 4]
    # = > A*x = d = [9; 14; 29; 15]

    a = [3.0, 2.0, 1.0]
    b = [1.0, 4.0, 3.0, 3.0]
    c = [4.0, 1.0, 4.0]
    d = [9.0, 14.0, 29.0, 15.0]

    x = zeros(length(b))
    OceanTurb.tridiagonalsolve!(x, d, a, b, c)

    solution = [1.0, 2.0, 3.0, 4.0]
    x == solution
end

function test_ldiv(N)
    x = rand(N)
    A = Tridiagonal(rand(N, N))
    b = A*x
    test_x = zeros(N)

    t_ldiv!(test_x, A, b)

    x ≈ test_x
end

@testset "Solvers" begin
    @test test_tridiagonalsolve()
    @test test_ldiv(10)
end


@testset "chainrules" begin
    Random.seed!(1)

    @testset "general: single input" begin
        for x in (1.0, -1.0, 0.0, 0.5, 10.0, -17.1, 1.5 + 0.7im)
            test_scalar(erf, x)
            test_scalar(erfc, x)
            test_scalar(erfcx, x)
            test_scalar(erfi, x)

            test_scalar(airyai, x)
            test_scalar(airyaiprime, x)
            test_scalar(airybi, x)
            test_scalar(airybiprime, x)

            test_scalar(dawson, x)

            if x isa Real
                test_scalar(logerfc, x)
                test_scalar(logerfcx, x)

                test_scalar(invdigamma, x)
            end

            if x isa Real && 0 < x < 1
                test_scalar(erfinv, x)
                test_scalar(erfcinv, x)
            end

            if x isa Real && x > 0 || x isa Complex
                test_scalar(gamma, x)
                test_scalar(digamma, x)
                test_scalar(trigamma, x)

                if x isa Real
                    test_scalar(airyaix, x)
                    test_scalar(airyaiprimex, x)
                end
            end
        end
    end

    @testset "Bessel functions" begin
        for x in (1.5, 2.5, 10.5, -0.6, -2.6, -3.3, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
            test_scalar(besselj0, x)
            test_scalar(besselj1, x)

            isreal(x) && x < 0 && continue

            test_scalar(bessely0, x)
            test_scalar(bessely1, x)

            for nu in (-1.5, 2.2, 4.0)
                test_frule(besseli, nu, x)
                test_rrule(besseli, nu, x)
                test_frule(besselix, nu, x) # derivative is `NotImplemented`
                test_frule(besselix, nu ⊢ NoTangent(), x) # derivative is a number
                test_rrule(besselix, nu, x)

                test_frule(besselj, nu, x)
                test_rrule(besselj, nu, x)
                test_frule(besseljx, nu, x) # derivative is `NotImplemented`
                test_frule(besseljx, nu ⊢ NoTangent(), x) # derivative is a number
                test_rrule(besseljx, nu, x)

                test_frule(besselk, nu, x)
                test_rrule(besselk, nu, x)
                test_frule(besselkx, nu, x)
                test_rrule(besselkx, nu, x)

                test_frule(bessely, nu, x)
                test_rrule(bessely, nu, x)
                test_frule(besselyx, nu, x) # derivative is `NotImplemented`
                test_frule(besselyx, nu ⊢ NoTangent(), x) # derivative is a number
                test_rrule(besselyx, nu, x)

                test_frule(hankelh1, nu, x)
                test_rrule(hankelh1, nu, x)
                test_frule(hankelh1x, nu, x)
                test_rrule(hankelh1x, nu, x)

                test_frule(hankelh2, nu, x)
                test_rrule(hankelh2, nu, x)
                test_frule(hankelh2x, nu, x)
                test_rrule(hankelh2x, nu, x)
            end
        end
    end

    @testset "erf, beta, and logbeta" begin
        test_points = (1.5, 2.5, 10.5, 1.6 + 1.6im, 1.6 - 1.6im, 4.6 + 1.6im)
        for x in test_points, y in test_points
            test_frule(beta, x, y)
            test_rrule(beta, x, y)

            test_frule(logbeta, x, y)
            test_rrule(logbeta, x, y)

            if x isa Real && y isa Real
                test_frule(erf, x, y)
                test_rrule(erf, x, y)
            end
        end
    end

    @testset "log gamma and co" begin
        # It is important that we have negative numbers with both odd and even integer parts
        test_points = (1.5, 2.5, 10.5, -0.6, -2.6, -3.3, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
        for x in test_points
            for m in (0, 1, 2, 3)
                test_frule(polygamma, m, x; rtol=1e-8)
                test_rrule(polygamma, m, x; rtol=1e-8)
            end

            isreal(x) && x < 0 && continue
            test_scalar(loggamma, x)
            for a in test_points
                test_frule(gamma, a, x; rtol=1e-8)
                test_rrule(gamma, a, x; rtol=1e-8)

                test_frule(loggamma, a, x)
                test_rrule(loggamma, a, x)
            end

            isreal(x) || continue
            test_frule(logabsgamma, x)
            test_rrule(logabsgamma, x; output_tangent=(randn(), randn()))
            for a in test_points
                isreal(a) && a > 0 || continue
                test_frule(gamma_inc, a, x, 0)
                test_rrule(gamma_inc, a, x, 0; output_tangent=(randn(), randn()))
            end
        end
    end

    @testset "exponential integrals" begin
        for x in (1.5, 2.5, 10.5, 1.6 + 1.6im, 1.6 - 1.6im, -4.6 + 1.6im)
            test_scalar(expint, x)
            test_scalar(expintx, x)

            for nu in (-1.5, 2.2, 4.0)
                test_frule(expint, nu, x)
                test_rrule(expint, nu, x)

                test_frule(expintx, nu, x)
                test_rrule(expintx, nu, x)
            end

            isreal(x) || continue
            test_scalar(expinti, x)
            test_scalar(sinint, x)
            test_scalar(cosint, x)
        end
    end

    # https://github.com/JuliaMath/SpecialFunctions.jl/issues/307
    @testset "promotions" begin
        # one argument
        for f in (erf, erfc, logerfc, erfcinv, erfcx, logerfcx, erfi, erfinv, sinint)
            _, ẏ = frule((NoTangent(), 1f0), f, 1f0)
            @test ẏ isa Float32
            _, back = rrule(f, 1f0)
            _, x̄ = back(1f0)
            @test x̄ isa Float32
        end

        # two arguments
        _, ẏ = frule((NoTangent(), 1f0, 1f0), erf, 1f0, 1f0)
        @test ẏ isa Float32
        _, back = rrule(erf, 1f0, 1f0)
        _, x̄ = back(1f0)
        @test x̄ isa Float32
    end
end


# simple program to test affinity propagation

using Distance
using Clustering

x = rand(100, 500)
S = -pairwise(Euclidean(), x, x)

# set diagonal value to median value
S = S - diagm(diag(S)) + median(S)*eye(size(S,1))

result = affinity_propagation(S; max_iter=200, display=:iter)


using Test, QSimulator
using QuadGK: quadgk

EC, EJ1, EJ2 = 0.172, 12.71, 3.69
freq = 0.1
amp = 0.6
harmonics = -50:50
parks = [0.0, 0.25]

@testset "FourierSeries" begin
	times = 0.0:.01:50
	for park in parks
		ϕ(t) = park + amp * sin(2π * freq * t)
		ω(t) = 2π * perturbative_transmon_freq(EC, EJ1, EJ2, ϕ(t))
		fs = FourierSeries(ω, freq, harmonics)
		ans = eval_series(fs, times)
		@test isapprox(ω.(times), ans, rtol=1e-9)
	end
end

@testset "rotating_frame_series" begin
	times = 0.0:.01:1
	for park in parks
		ϕ(t) = park + amp * sin(2π * freq * t)
		ω(t) = 2π * perturbative_transmon_freq(EC, EJ1, EJ2, ϕ(t))
		fs = FourierSeries(ω, freq, harmonics)
		Φ(t) = exp(1im * quadgk(ω, 0.0, t)[1])
		cfs = rotating_frame_series(fs, harmonics)
		# note that we must multiply by the term corresponding to fs.terms[0]
		# as this was ignored in rotating_frame_series
		ans = exp.(1im * fs.terms[0] * times) .* eval_series(cfs, times)
		@test isapprox(Φ.(times), ans, rtol=1e-6)
	end
end


module FringeExplorer

using RadioInterferometry
using Symbolics

export ha2t
export t2ha
export td2wdw
export td2wdw!
export hd2wdw
export hd2wdw!

export telinfo
export antpos
include("telinfo.jl")

export s2c
export c2s
export beamrings
include("spherical.jl")

export MJD0
export radec2hadec
include("erfa_helpers.jl")

"""
Rotaional rate of Earth in radians per second
"""
const Ωe = 7.29211538e-5

"""
    ha2t(ha::Real, Ω=Ωe) -> t

Convert hour angle `ha` in radians to time `t` in seconds. `ha == t == 0` at
transit.  `Ω`, the rate of rotation in radians per second, defaults to `Ωe`, the
rotational rate of Earth.
"""
ha2t(ha::Real, Ω=Ωe) = ha / Ω

"""
    t2ha(t::Real, Ω=Ωe) -> ha

Convert time `t` in seconds to hour angle `ha` in radians. `ha == t == 0` at
transit.  `Ω`, the rate of rotation in radians per second, defaults to `Ωe`, the
rotational rate of Earth.
"""
t2ha(t::Real, Ω=Ωe) = t * Ω

# Symbolic variables
const Ω, t, δ = @variables Ω t δ
const Dt = Differential(t)
const h = Ω * t

# Passing zero for longitude requires rotating XYZ position vectors around Z
# such that positive X intersects the local meridian.
const xyzR3uvw = xyz2uvw(h, δ, 0)

# Substitute value to `Ω`
const xyzR3uvw_e = substitute.(xyzR3uvw, [Ω=>Ωe])

# Delay is w component (use 3:3 to keep as a one row Matrix)
const w = xyzR3uvw_e[3:3,:]
# Delay rate is derivative of w component
const dw = expand_derivatives.(Dt.(w))
# Make 2x3 matrix from w and dw
const wdw_m = [w; dw]

# Create expressions for functions that return wdw for given t and δ
const wdw_expr = build_function(wdw_m, t, δ)

"""
    td2wdw(t, δ) -> Matrix
    tδ2wdw(t, δ) -> Matrix

Returns a 2x3 Matrix that can left multiply a 3 element vector (or a 3 row
matrix) of antenna position(s) to get delay `w` in the first element (or row)
and the delay rate `dw` in the second for a source observed at time `t` and
declination `δ`.  The time `t` should be in seconds, with `t<0` before the
source transits, `t==0` when the source transits, and `t>0` after the source
transits.

# Example
```julia
wdw = td2wdw(0, 0) * antpos
```
"""
td2wdw = eval(wdw_expr[1])
const tδ2wdw = td2wdw

"""
    td2wdw!(dst, t, δ)
    tδ2wdw!(dst, t, δ)

Stores a 2x3 Matrix in `dst` (which should already be a 2x3 Matrix) that can
left multiply a 3 element vector (or a 3 row matrix) of antenna position(s) to
get delay `w` in the first element (or row) and the delay rate `dw` in the
second for a source observed at time `t` and declination `δ`.  The time `t`
should be in seconds, with `t<0` before the source transits, `t==0` when the
source transits, and `t>0` after the source transits.
"""
td2wdw! = eval(wdw_expr[2])
const tδ2wdw! = td2wdw!

"""
    hd2wdw(h, δ, Ω=Ωe) -> Matrix
    hδ2wdw(h, δ, Ω=Ωe) -> Matrix

Returns a 2x3 Matrix that can left multiply a 3 element vector (or a 3 row
matrix) of antenna position(s) to get delay `w` in the first element (or row)
and the delay rate `dw` in the second for a source observed at hour angle `h` and
declination `δ`.  The hour angle `h` should be in radians, with `h<0` before the
source transits, `h==0` when the source transits, and `h>0` after the source
transits.  `Ω` is used to convert `h` to time (see [`ha2t`](@ref)).

# Example
```julia
wdw = hd2wdw(ha, dec) * antpos
```
"""
hd2wdw(h, δ, Ω=Ωe) = td2wdw(ha2t(h, Ω), δ)
hδ2wdw = hd2wdw

"""
    hd2wdw!(dst, t, δ, Ω=Ωe)
    hδ2wdw!(dst, t, δ, Ω=Ωe)

Stores a 2x3 Matrix in `dst` (which should already be a 2x3 Matrix) that can
left multiply a 3 element vector (or a 3 row matrix) of antenna position(s) to
get delay `w` in the first element (or row) and the delay rate `dw` in the
second for a source observed at time `t` and declination `δ`.  The time `t`
should be in seconds, with `t<0` before the source transits, `t==0` when the
source transits, and `t>0` after the source transits.  `Ω` is used to convert
`h` to time (see [`ha2t`](@ref)).
"""
hd2wdw!(dst, h, δ, Ω=Ωe) = td2wdw!(dst, ha2t(h, Ω), δ)
hδ2wdw! = hd2wdw!

end # module FringeExplorer


using PyPlot
x = range(0, stop=1, length=50)
plot(x, sin.(2x).*exp.(-x/3))
savefig(joinpath(@__DIR__, "output", "script2.svg"))


function agm{R<:Real}(a::R, b::R)
    half = isdefined(:parse, (R,)) ? parse(R,"0.5") : (R)(0.5)
    maxIterations = sizeof(R)*8

    a,b = minmax(a,b)
    a0, b0 = deepcopy(a), deepcopy(b)

    for i in 1:maxIterations

        if (eps(a) == eps(b))
           if (significand(a) == significand(b))
              break
           else
              a, b = (abs(a) < abs(b)) ? (a,b) : (b,a)
           end
        end

        a = (a0 + b0) * half
        b = sqrt( a0 * b0 )

        (abs(a-b) < eps(a)) & (eps(a)
           break
        end
    end
end


# coefficients a_0, ..., a_9
let
  pdTableI = [ -0.5000000002 ;
                0.3333333343 ;
               -0.2499998565 ;
                0.1999997049 ;
               -0.1666848753 ;
                0.1428833286 ;
               -0.1241963125 ;
                0.1101687109 ;
               -0.1142650302 ;
                0.1055093006 ]

  # factorials of 0--9
  pdFactorials = [ 1; 1; 2; 6; 24; 120; 720; 5040; 40320; 362880 ]

  invSqrt2Pi = 1/sqrt(2*π)
  inv12 = 1/12
  inv360 = 1/360
  inv1260 = 1/1260

  @inline function procedureF(K::Int64, μ::Float64, s::Float64, ω::Float64,
    c0::Float64, c1::Float64, c2::Float64, c3::Float64)
    ## 1. Poisson probabilities
    px::Float64 = 0.0
    py::Float64 = 0.0
    if K < 10
      px = -μ
      py = μ^K/pdFactorials[K+1]
    else
      KSq::Float64 = K*K
      KCubed::Float64 = K*KSq
      KFived::Float64 = KCubed * KSq
      δ::Float64 = inv12/K - inv360/KCubed + inv1260/KFived
      V::Float64 = (μ-K)/K
      if abs(V) > 0.25
        px = K*log(1+V) - μ + K - δ
      else
        ws::Float64 = 0.0
        Vj::Float64 = 1.0
        for j = 0:9
          ws += pdTableI[j+1]*Vj
          Vj *= V
        end
        px = K*V*V*ws - δ
      end
      py = 1/sqrt(2*π*K)
    end
    ## 2. Discrete normal probabilities
    X::Float64 = (K - μ + 0.5)/s
    XSq::Float64 = X*X
    fx::Float64 = -0.5*XSq ## typo in paper omits minus sign!
    fy::Float64 = ω*(((c3*XSq + c2)*XSq + c1)*XSq + c0)
    return px, py, fx, fy
  end

  inv24 = 1/24
  inv7 = 1/7

  @inline function caseA(μ::Float64, rng::RNG) where RNG <: AbstractRNG
    s::Float64 = sqrt(μ)
    d::Float64 = 6*μ*μ
    L::Int64 = floor(Int64, μ - 1.1484)
    ## Normal sample
    G::Float64 = μ + s*randn(rng)
    U::Float64 = 0.0
    K::Int64 = 0
    if G >= 0
      K = floor(Int64, G)
      ## Immediate acceptance
      if K >= L return K end
      U = rand(rng)
      ## Squeeze acceptance
      if d*U >= (μ - K)^3 return K end
    end
    ## Preparations for Q and H
    ω::Float64 = invSqrt2Pi/s
    b1::Float64 = inv24/μ
    b2::Float64 = 0.3*b1*b1
    c3::Float64 = inv7*b1*b2
    c2::Float64 = b2 - 15*c3
    c1::Float64 = b1 - 6*b2 + 45*c3
    c0::Float64 = 1 - b1 + 3*b2 - 15*c3
    c::Float64 = 0.1069/μ
    px::Float64 = 0.0
    py::Float64 = 0.0
    fx::Float64 = 0.0
    fy::Float64 = 0.0
    if G >= 0
      px, py, fx, fy = procedureF(K, μ, s, ω, c0, c1, c2, c3)
      ## Quotient acceptance
      if fy*(1-U) <= py*exp(px-fx) return K end
    end
    ## Double exponential sample
    while (true)
      E::Float64 = randexp(rng)
      U = 2*rand(rng) - 1
      T::Float64 = 1.8 + E*sign(U)
      if T <= -0.6744
        continue
      end
      K = floor(Int64, μ + s*T)
      px, py, fx, fy = procedureF(K, μ, s, ω, c0, c1, c2, c3)
      ## Hat acceptance
      if c*abs(U) <= py*exp(px+E) - fy*exp(fx+E)
        return K
      end
    end
  end

  ## modification: no table, just inversion.
  @inline function caseB(μ::Float64, rng::RNG) where RNG <: AbstractRNG
    p::Float64 = exp(-μ)
    q::Float64 = p
    U::Float64 = rand(rng)
    K::Int64 = 0
    while U > q
      K += 1
      p *= μ/K
      q += p
    end
    return K
  end

  ## The algorithm from
  ## Ahrens, J.H. and Dieter, U., 1982. Computer generation of Poisson deviates
  ## from modified normal distributions. ACM Transactions on Mathematical
  ## Software (TOMS), 8(2), pp. 163-179.
  ## This is suitable when μ changes every invocation. Modifications or
  ## alternative algorithms should be used for the case where many Poisson
  ## variates are required for the same μ
  global @inline function samplePoisson(μ::Float64,
      rng::RNG = GLOBAL_RNG) where RNG <: AbstractRNG
    if μ >= 10.0
      return caseA(μ, rng)
    else
      return caseB(μ, rng)
    end
  end
end


@inline add_fpfp_dd(x::T, y::T) where {T<:AbstractFloat} = two_sum(x, y)
@inline sub_fpfp_dd(x::T, y::T) where {T<:AbstractFloat} = two_diff(x, y)
@inline mul_fpfp_dd(x::T, y::T) where {T<:AbstractFloat} = two_prod(x, y)
@inline dvi_fpfp_dd(x::T, y::T) where {T<:AbstractFloat} = two_dvi(x, y)
# seems to help
@inline add_fpfp_dd(x::T, y::T) where {T<:IEEEFloat} = two_sum(x, y)
@inline sub_fpfp_dd(x::T, y::T) where {T<:IEEEFloat} = two_diff(x, y)
@inline mul_fpfp_dd(x::T, y::T) where {T<:IEEEFloat} = two_prod(x, y)
@inline dvi_fpfp_dd(x::T, y::T) where {T<:IEEEFloat} = two_dvi(x, y)


module GFIlogisticRegression

import Polyhedra
using CDDLib
import LinearAlgebra
import Distributions
import Optim
import StatsBase
import StatsModels
import DataFrames
import DataFramesMeta

export fidSummary
export fidConfInt
export fidQuantile
export fidProb
export fidSampleLR

function logit(u) # = qlogis
  return log(u / (1.0 - u))
end

function expit(x) # = plogis
  return 1.0 / (1.0 + exp(-x))
end

function rtlogis1(x)
  b = expit(x)
  if b == 0.0
    return x
  end
  return logit(b * rand())
end

function rtlogis2(x)
  a = expit(x)
  if a == 1.0
    return x
  end
  return logit(a + (1.0-a) * rand())
end

function from01(u)
  return map(log, u ./ (1.0 .- u))
end

function from01scalar(u)
  return log(u / (1.0 - u))
end

function to01(x)
  return 1.0 / (1.0 + exp(-x))
end

function dfrom01(u)
  return 1.0 / (u * (1.0 - u))
end

function dlogis(x)
  expminusx = map(exp, -x)
  oneplusexpminusx = 1.0 .+ expminusx
  return expminusx ./ (oneplusexpminusx .* oneplusexpminusx)
end

function ldlogis(x)
  return -x - 2.0 * map(log1p, map(exp, -x))
end

function dldlogis(x)
  return 1.0 .- 2.0 ./ (1.0 .+ map(exp, -x))
end

function forig(x, P, b)
  return prod(dlogis(P * x + b))
end

function f(u, P, b)
  return prod(dlogis(P * from01(u) + b))
end

function logf(u, P, b)
  return sum(ldlogis(P * from01(u) + b))
end

function df(ui, Pi, y1, y2)
  return y1 * dfrom01(ui) * sum(Pi .* y2)
end

function dlogf(ui, Pi, y2)
  return dfrom01(ui) * sum(Pi .* y2)
end

function get_umax(P, b)
  d = size(P, 2)
  fn = function(u)
    return -logf(u, P, b)
  end
  grfn! = function(storage, u)
    y2 = dldlogis(P * from01(u) + b)
    for i in 1:d
      storage[i] = -dlogf(u[i], P[:, i], y2)
    end
  end
  lower = sqrt(eps()) * ones(d)
  upper = 1.0 .- lower
  init = 0.5 * ones(d)
  opt = Optim.optimize(
    fn, grfn!, lower, upper, init, Optim.Fminbox(Optim.LBFGS())
  )
  return (
    mu = from01(opt.minimizer),
    umax = exp(-opt.minimum)^(2.0 / (d + 2))
  )
end

function get_vmin_i(P, b, j, mu)
  d = size(P, 2)
  alpha = 1.0 / (d + 2)
  fn = function(u)
    return f(u, P, b)^alpha * (from01scalar(u[j]) - mu[j])
  end
  grfn! = function(storage, u)
    y1alpha = f(u, P, b)^alpha
    y2 = dldlogis(P * from01(u) + b)
    diff = from01scalar(u[j]) - mu[j]
    storage[j] = y1alpha * dfrom01(u[j]) *
      (alpha * sum(P[:, j] .* y2) * diff + 1.0)
    others = deleteat!(collect(1:d), j)
    for i in others
      storage[i] = y1alpha * dfrom01(u[i]) * alpha * sum(P[:, i] .* y2) * diff
    end
  end
  eta = sqrt(eps())
  lower = eta * ones(d)
  upper = (1.0-eta) * ones(d)
  upper[j] = to01(mu[j])
  init = 0.5 * ones(d)
  init[j] = to01(mu[j]) / 2.0
  opt = Optim.optimize(
    fn, grfn!, lower, upper, init, Optim.Fminbox(Optim.LBFGS())
  )
  return opt.minimum
end

function get_vmin(P, b, mu)
  d = size(P, 2)
  vmin = Vector{Float64}(undef, d)
  for i in 1:d
    vmin[i] = get_vmin_i(P, b, i, mu)
  end
  return vmin
end

function get_vmax_i(P, b, j, mu)
  d = size(P, 2)
  alpha = 1.0 / (d + 2)
  fn = function(u)
    return -f(u, P, b)^alpha * (from01scalar(u[j]) - mu[j])
  end
  grfn! = function(storage, u)
    y1alpha = f(u, P, b)^alpha
    y2 = dldlogis(P * from01(u) + b)
    diff = from01scalar(u[j]) - mu[j]
    storage[j] = -y1alpha * dfrom01(u[j]) *
      (alpha * sum(P[:, j] .* y2) * diff + 1.0)
    others = deleteat!(collect(1:d), j)
    for i in others
      storage[i] = -y1alpha * dfrom01(u[i]) * alpha * sum(P[:, i] .* y2) * diff
    end
  end
  eta = sqrt(eps())
  lower = eta * ones(d)
  lower[j] = to01(mu[j])
  upper = (1.0-eta) * ones(d)
  init = 0.5 * ones(d)
  init[j] = (to01(mu[j]) + 1.0) / 2.0
  opt = Optim.optimize(
    fn, grfn!, lower, upper, init, Optim.Fminbox(Optim.LBFGS())
  )
  return -opt.minimum
end

function get_vmax(P, b, mu)
  d = size(P, 2)
  vmax = Vector{Float64}(undef, d)
  for i in 1:d
    vmax[i] = get_vmax_i(P, b, i, mu)
  end
  return vmax
end

function get_bounds(P, b)
  (mu, umax) = get_umax(P, b)
  return (
    mu = mu,
    umax = umax,
    vmin = get_vmin(P, b, mu),
    vmax = get_vmax(P, b, mu)
  )
end

function rcd(n, P, b)
  d = size(P, 2)
  sims = Vector{Vector{Float64}}(undef, n)
  (mu, umax, vmin, vmax) = get_bounds(P, b)
  if any(vmin .>= vmax)
    error("vmin .>= vmax")
  end
  alpha = 2.0 / (d+2)
  k = 0
  while k < n
    u = umax * rand()
    v = vmin + (vmax - vmin) .* rand(d)
    x = v / sqrt(u) + mu
    if u < forig(x, P, b)^alpha
      k += 1
      sims[k] = x
    end
  end
  return sims
end

"""
    fidSampleLR(formula, data, N[, gmp][, thresh])

Fiducial sampling of the parameters of the logistic regression model.

# Arguments
- `formula`: a formula describing the model
- `data`: data frame in which the variables of the model can be found
- `N`: number of simulations
- `gmp`: whether to use exact arithmetic in the algorithm
- `thresh`: the threshold used in the sequential sampler; the default `N/2` should not be changed

# Example

    using GFIlogisticRegression, DataFrames, StatsModels
    data = DataFrame(
      y = [0, 0, 1, 1, 1],
      x = [-2, -1, 0, 1, 2]
    )
    fidsamples = fidSampleLR(@formula(y ~ x), data, 3000)
"""
function fidSampleLR(formula, data, N, gmp = false, thresh = N/2)
  y = StatsModels.response(formula, data)
  mf = StatsModels.ModelFrame(formula, data)
  mm = StatsModels.ModelMatrix(mf)
  X = mm.m
  (_, colnames) = StatsModels.coefnames(mf.f)
  (n, p) = size(X)
  weight = ones(n, N)
  local WTnorm
  ESS = N .* ones(n)
  T = gmp ? Rational{BigInt} : Float64
  cdd = gmp ? :exact : :float
  CC = Vector{Array{T,2}}(undef, N)
  cc = Vector{Vector{T}}(undef, N)
  # Kstart ####
  Kstart = [1]
  i = 1
  rk = 1
  while rk < p
    i += 1
    Kstart_plus_i = vcat(Kstart, i)
    if LinearAlgebra.rank(X[Kstart_plus_i, :]) == rk + 1
      Kstart = Kstart_plus_i
      rk += 1
    end
  end
  @inbounds Xstart = X[Kstart, :]
  qXstart = convert(Array{T,2}, Xstart)
  ystart = y[Kstart]
  K = setdiff(1:n, Kstart)
  @inbounds XK = X[K, :]
  qXK = convert(Array{T,2}, XK)
  @inbounds yK = y[K]
  # t = 1 to p ####
  At = Array{Float64}(undef, p, N)
  for i in 1:N
    a = map(logit, rand(p))
    @inbounds At[:, i] = a
    C = Array{T,2}(undef, p, p)
    c = Vector{T}(undef, p)
    for j in 1:p
      @inbounds if ystart[j] == 0
        @inbounds C[j, :] = qXstart[j, :]
        @inbounds c[j] = convert(T, a[j])
      else
        @inbounds C[j, :] = -qXstart[j, :]
        @inbounds c[j] = convert(T, -a[j])
      end
    end
    @inbounds CC[i] = C
    @inbounds cc[i] = c
  end
  # t from p+1 to n ####
  for t in 1:(n-p)
    #println(t)
    At = vcat(At, Array{Float64,2}(undef, 1, N))
    @inbounds qXt = qXK[t, :]
    qXt_row = reshape(qXt, 1, :)
    qXtt = LinearAlgebra.transpose(qXt)
    @inbounds if yK[t] == 0
      for i in 1:N
        @inbounds H = Polyhedra.hrep(CC[i], cc[i])
        plyhdrn = Polyhedra.polyhedron(H, CDDLib.Library(cdd))
        pts = collect(Polyhedra.points(plyhdrn))
        MIN = convert(
          Float64, minimum(qXtt * hcat(pts...))
        )
        atilde = rtlogis2(MIN)
        @inbounds weight[t, i] = 1 - expit(MIN)
        @inbounds CC[i] = vcat(CC[i], qXt_row)
        @inbounds cc[i] = vcat(cc[i], convert(T, atilde))
        @inbounds At[p+t, i] = atilde
      end
    else
      for i in 1:N
        @inbounds H = Polyhedra.hrep(CC[i], cc[i])
        plyhdrn = Polyhedra.polyhedron(H, CDDLib.Library(cdd))
        pts = collect(Polyhedra.points(plyhdrn))
        MAX = convert(
          Float64, maximum(qXtt * hcat(pts...))
        )
        atilde = rtlogis1(MAX)
        @inbounds weight[t, i] = expit(MAX)
        @inbounds CC[i] = vcat(CC[i], -qXt_row)
        @inbounds cc[i] = vcat(cc[i], convert(T, -atilde))
        @inbounds At[p+t, i] = atilde
      end
    end
    WT = prod(weight; dims = 1)[1, :]
    WTnorm = WT ./ sum(WT)
    @inbounds ESS[p+t] = 1.0 / sum(WTnorm .* WTnorm)
    @inbounds if ESS[p+t] < thresh || t == n-p
      #println("alteration...")
      Nsons = rand(Distributions.Multinomial(N, WTnorm))
      counter = 1
      At_new = Array{Float64, 2}(undef, p+t, 0)
      @inbounds D = vcat(Xstart, XK[1:t, :])
      QR = LinearAlgebra.qr(D)
      @inbounds P = QR.Q[:, 1:p]#orth(D)
      Pt = LinearAlgebra.transpose(P)
      QQt = LinearAlgebra.I - P*Pt
      M = inv(LinearAlgebra.transpose(D)*D) * LinearAlgebra.transpose(D) * P
      CCtemp = Vector{Array{T,2}}(undef, N)
      cctemp = Vector{Vector{T}}(undef, N)
      for i in 1:N
        @inbounds ncopies = Nsons[i]
        if ncopies >= 1
          @inbounds CCtemp[counter] = CC[i]
          @inbounds cctemp[counter] = cc[i]
          @inbounds At_new = hcat(At_new, At[:, i])
          if ncopies > 1
            @inbounds H = Polyhedra.hrep(CC[i], cc[i])
            plyhdrn = Polyhedra.polyhedron(H, CDDLib.Library(cdd))
            pts = collect(Polyhedra.points(plyhdrn))
            lns = collect(Polyhedra.lines(plyhdrn))
            rys = collect(Polyhedra.rays(plyhdrn))
            @inbounds b = QQt * At[:, i]
            @inbounds B = Pt * At[:, i]
            BTILDES = rcd(ncopies-1, P, b)
            for j in 2:ncopies
              VT_new = Vector{Vector{T}}(undef, length(pts))
              @inbounds Btilde = BTILDES[j-1]
              At_tilde = P * Btilde .+ b
              At_new = hcat(At_new, At_tilde)
              for k in 1:length(pts)
                pt = convert(Vector{Float64}, pts[k])
                VT_new[k] =
                  convert(Vector{T}, pt .- M * (B .- Btilde))
              end
              V = Polyhedra.vrep(VT_new, lns, rys)
              plyhdrn = Polyhedra.polyhedron(V, CDDLib.Library(cdd))
              #HlfSpcs = Polyhedra.hrep(plyhdrn).halfspaces
              #A = map(x -> x.a, HlfSpcs)
              H = Polyhedra.hrep(plyhdrn)
              Hmat = convert(Array{T,2}, CDDLib.extractA(unsafe_load(H.matrix)))
              @inbounds CCtemp[counter + j - 1] = -Hmat[:, 2:end] #LinearAlgebra.transpose(hcat(A...))
              @inbounds cctemp[counter + j - 1] = Hmat[:, 1]#map(x -> x.β, HlfSpcs)
            end
          end
          counter += ncopies
        end
      end
      CC = CCtemp
      cc = cctemp
      At = At_new
      if t < n-p
        weight = ones(n, N)
      end
    end
  end
  Beta = Array{Float64, 2}(undef, N, p)
  for i in 1:N
    @inbounds H = Polyhedra.hrep(CC[i], cc[i])
    plyhdr = Polyhedra.polyhedron(H, CDDLib.Library(cdd))
    vertices = collect(Polyhedra.points(plyhdr))
    pts = hcat(vertices...)
    for j in 1:p
      pt = convert(Vector{Float64}, pts[j, :])
      if rand() < 0.5
        @inbounds Beta[i, j] = minimum(pt)
      else
        @inbounds Beta[i, j] = maximum(pt)
      end
    end
  end
  dfBeta = DataFrames.DataFrame(Beta, Symbol.(colnames))
  return (Beta = dfBeta, Weights = WTnorm)
end # fidSampleLR

"""
    fidSummary(fidsamples)

Summary of the fiducial simulations.

# Argument
- `fidsamples`: an output of `fidSampleLR`

# Example

    using GFIlogisticRegression, DataFrames, StatsModels
    data = DataFrame(
      y = [0, 0, 1, 1, 1],
      x = [-2, -1, 0, 1, 2]
    )
    fidsamples = fidSampleLR(@formula(y ~ x), data, 3000)
    fidSummary(fidsamples)
"""
function fidSummary(fidsamples)
  (_, p) = size(fidsamples.Beta)
  gdf = DataFrames.groupby(DataFrames.stack(fidsamples.Beta, 1:p), :variable)
  wghts = fidsamples.Weights
  function fsummary(x)
    return (
      mean = sum(wghts .* x),
      median = StatsBase.quantile(x, StatsBase.weights(wghts), 0.5),
      lwr = StatsBase.quantile(x, StatsBase.weights(wghts), 0.025),
      upr = StatsBase.quantile(x, StatsBase.weights(wghts), 0.975)
    )
  end
  return DataFrames.combine(gdf, DataFrames.valuecols(gdf) .=> fsummary => DataFrames.AsTable)
end

"""
    fidConfInt(parameter, fidsamples, conf)

Fiducial confidence interval of a parameter of interest.

# Arguments
- `parameter`: an expression of the parameter of interest given as a string; see the example
- `fidsamples`: an output of `fidSampleLR`
- `conf`: confidence level

# Example

    using GFIlogisticRegression, DataFrames, StatsModels
    data = DataFrame(
      y = [0, 0, 1, 1, 1, 1],
      group = ["A", "A", "A", "B", "B", "B"]
    )
    fidsamples = fidSampleLR(@formula(y ~ 0 + group), data, 3000)
    fidConfInt(":\\\"group: A\\\" - :\\\"group: B\\\"", fidsamples, 0.95)
"""
function fidConfInt(parameter, fidsamples, conf = 0.95)
  x = eval(:(DataFramesMeta.@with($(fidsamples.Beta), $(Meta.parse(parameter)))))
  wghts = fidsamples.Weights
  halpha = (1.0 - conf) / 2.0
  qntls = StatsBase.quantile(x, StatsBase.weights(wghts), [halpha; 1.0-halpha])
  return (lower = qntls[1], upper = qntls[2])
end

"""
    fidQuantile(parameter, fidsamples, p)

Fiducial quantile of a parameter of interest.

# Arguments
- `parameter`: an expression of the parameter of interest given as a string; see the example
- `fidsamples`: an output of `fidSampleLR`
- `p`: quantile level, between 0 and 1

# Example

    using GFIlogisticRegression, DataFrames, StatsModels
    data = DataFrame(
      y = [0, 0, 1, 1, 1, 1],
      group = ["A", "A", "A", "B", "B", "B"]
    )
    fidsamples = fidSampleLR(@formula(y ~ 0 + group), data, 3000)
    fidQuantile(":\\\"group: A\\\" ./ :\\\"group: B\\\"", fidsamples, 0.5)
"""
function fidQuantile(parameter, fidsamples, p)
  x = eval(:(DataFramesMeta.@with($(fidsamples.Beta), $(Meta.parse(parameter)))))
  wghts = fidsamples.Weights
  return StatsBase.quantile(x, StatsBase.weights(wghts), p)
end

"""
    fidProb(parameter, fidsamples, q)

Fiducial non-exceedance probability of a parameter of interest.

# Arguments
- `parameter`: an expression of the parameter of interest given as a string; see the example
- `fidsamples`: fiducial simulations, an output of `fidSampleLR`
- `q`: the non-exceedance threshold

# Example

    using GFIlogisticRegression, DataFrames, StatsModels
    data = DataFrame(
      y = [0, 0, 1, 1, 1],
      x = [-2, -1, 0, 1, 2]
    )
    fidsamples = fidSampleLR(@formula(y ~ x), data, 3000)
    fidProb("map(exp, :x)", fidsamples, 1) # this is Pr(exp(x) <= 1)
"""
function fidProb(parameter, fidsamples, q)
  x = eval(:(DataFramesMeta.@with($(fidsamples.Beta), $(Meta.parse(parameter)))))
  wghts = fidsamples.Weights
  return sum(wghts[findall(x .<= q)])
end

end # module

#=
P = [
  -0.82 -0.18;
  -0.41 -0.37;
  0.0   -0.55;
  0.41  -0.73
]
b = [-0.50; 1.58; -1.66; 0.58]
=#

#=
y = [0, 0, 1, 1, 1]
X = [1 -2; 1 -1; 1 0; 1 1; 1 2]
fidsamples = fidSampleLR(y, X, 5)
=#

#=
using DataFrames, StatsModels
data = DataFrame(
  y = [0, 0, 1, 1, 1],
  x = [-2, -1, 0, 1, 2]
)
fidsamples = fidSampleLR(@formula(y ~ x), data, 5)

fidSummary(fidsamples)
fidConfInt("map(exp, :x)", fidsamples)
fidConfInt(":x ./ :\"(Intercept)\"", fidsamples)
=#

#=
data = DataFrame(
  y = [0, 0, 1, 1, 1, 1],
  group = ["A", "A", "A", "B", "B", "B"]
)
fidsamples = fidSampleLR(@formula(y ~ 0 + group), data, 3000)
fidConfInt(":\"group: A\" - :\"group: B\"", fidsamples, 0.95)
=#

#=
    for i in 1:N
      @inbounds H = Polyhedra.hrep(CC[i], cc[i])
      plyhdrn = Polyhedra.polyhedron(H, CDDLib.Library(cdd))
      pts = collect(Polyhedra.points(plyhdrn))
      @inbounds if yK[t] == 0
        MIN = convert(
          Float64, minimum(qXtt * hcat(pts...))
        )
        atilde = rtlogis2(MIN)
        @inbounds weight[t, i] = 1 - expit(MIN)
        @inbounds CC[i] = vcat(CC[i], qXt_row)
        @inbounds cc[i] = vcat(cc[i], convert(T, atilde))
      else
        MAX = convert(
          Float64, maximum(qXtt * hcat(pts...))
        )
        atilde = rtlogis1(MAX)
        @inbounds weight[t, i] = expit(MAX)
        @inbounds CC[i] = vcat(CC[i], -qXt_row)
        @inbounds cc[i] = vcat(cc[i], convert(T, -atilde))
      end
      @inbounds At[p+t, i] = atilde
    end
=#


#%%
using CUTEst
using Projeto2Solvers
using NLPModels
using LinearAlgebra

#%%
# nlp = ADNLPModel(x -> (x[1] - 1)^2 + (x[2] - 2)^2 / 4, zeros(2))
nlp = CUTEstModel("GAUSS2LS")

#%%
print(nlp)
#%%
output = bfgs_rc(nlp)
print(output)

#%%
finalize(nlp)


using transD_GP, Distributed
## make options for the stationary GP
nmin, nmax = 2, 40
pnorm = 2.
K = transD_GP.GP.OrstUhn()
demean = true
fbounds = [-0.5 2.5]
sdev_pos = [0.05abs(diff([extrema(znall)...])[1])]
sdev_prop = 0.05diff(fbounds, dims=2)[:]
xall = permutedims(collect(znall))
xbounds = permutedims([extrema(znall)...])
λ, δ = [2], 0.1

## Initialize a stationary GP using these options
Random.seed!(12)
opt = transD_GP.OptionsStat(nmin = nmin,
                        nmax = nmax,
                        xbounds = xbounds,
                        fbounds = fbounds,
                        xall = xall,
                        λ = λ,
                        δ = δ,
                        demean = demean,
                        sdev_prop = sdev_prop,
                        sdev_pos = sdev_pos,
                        pnorm = pnorm,
                        save_freq = 25,
                        quasimultid = false,
                        K = K,
                        )


using NLsolve

ϕ = 0.5
y = 1
β = 0.996
θ = 0.199
α = 0.4
s = 0.0225

function f!(F, v)
    ω = v[1]
    c = v[2]
    F[1] = (1-(1-ϕ)*α)*ω - ϕ*(y + θ*c)
    F[2] = θ^α*(1-β*(1-s))*c - β*(y-ω)
end

soln = nlsolve(f!, ones(2), autodiff = :forward)

@show soln.zero


function [CG]=DerivativeMatrix(CG)
OrdPolyX=CG.OrdPolyX;
CG.DSX=zeros(OrdPolyX+1,OrdPolyX+1);
for i=1:OrdPolyX+1
  for j=1:OrdPolyX+1
    CG.DSX(i,j)=DLagrange(CG.xwX(i),CG.xwX,j);
  end
end
CG.DWX=-inv(diag(CG.wX))*CG.DSX'*diag(CG.wX);

OrdPolyY=CG.OrdPolyY;
CG.DSY=zeros(OrdPolyY+1,OrdPolyY+1);
for i=1:OrdPolyY+1
  for j=1:OrdPolyY+1
    CG.DSY(i,j)=DLagrange(CG.xwY(i),CG.xwY,j);
  end
end
CG.DWY=-inv(diag(CG.wY))*CG.DSY'*diag(CG.wY);

if strcmp(CG.FeX,'LG')
  OrdPolyX=CG.OrdPolyX;
  CG.DSXLGL=zeros(OrdPolyX+2,OrdPolyX+1);
  CG.IntXLGL=zeros(OrdPolyX+2,OrdPolyX+1);
  CG.IntXLG=zeros(OrdPolyX+1,OrdPolyX+2);
  [CG.wXLGL,CG.xwXLGL]=GaussLobattoQuad(CG.OrdPolyX+1);
  for i=1:OrdPolyX+2
    for j=1:OrdPolyX+1
      CG.DSXLGL(i,j)=DLagrange(CG.xwXLGL(i),CG.xwX,j);
      CG.IntXLGL(i,j)=Lagrange(CG.xwXLGL(i),CG.xwX,j);
      CG.IntXLG(j,i)=Lagrange(CG.xwX(j),CG.xwXLGL,i);
    end
  end
  %CG.IntXLG=inv(diag(CG.wX))*CG.IntXLG*diag(CG.wXLGL);
  CG.DWXLGL=-inv(diag(CG.wX))*CG.DSXLGL'*diag(CG.wXLGL);
end
if strcmp(CG.FeX,'LG') && strcmp(CG.FeY,'LGL')
  DSXLGL=zeros(OrdPolyX+2,OrdPolyX+1);
  for i=1:OrdPolyX+2
    for j=1:OrdPolyX+1
      DSXLGL(i,j)=DLagrange(CG.xwXLGL(i),CG.xwX,j);
    end
  end
  CG.DWXLG=-inv(diag(CG.wX))*DSXLGL'*diag(CG.wXLGL);
end
if strcmp(CG.FeX,'LGL')
  [CG.wXLG,CG.xwXLG]=GaussLegendreQuad(CG.OrdPolyX-1);
  CG.IntXLG=zeros(CG.OrdPolyX,CG.OrdPolyX+1);
  for i=1:CG.OrdPolyX
    for j=1:CG.OrdPolyX+1
      CG.IntXLG(i,j)=Lagrange(CG.xwXLG(i),CG.xwX,j);
      CG.IntXLGL(j,i)=Lagrange(CG.xwX(j),CG.xwXLG,i);
    end
  end
  CG.IntXLGL=eye(CG.OrdPolyX+1);
  CG.IntXLG=eye(CG.OrdPolyX+1);
end

if strcmp(CG.FeY,'LGL')
  [CG.wYLG,CG.xwYLG]=GaussLegendreQuad(CG.OrdPolyY-1);
  CG.IntYLG=zeros(CG.OrdPolyY,CG.OrdPolyY+1);
  for i=1:CG.OrdPolyY
    for j=1:CG.OrdPolyY+1
      CG.IntYLG(i,j)=Lagrange(CG.xwYLG(i),CG.xwY,j);
      CG.IntYLGL(j,i)=Lagrange(CG.xwY(j),CG.xwYLG,i);
    end
  end
  CG.IntYLGL=eye(CG.OrdPolyY+1);
  CG.IntYLG=eye(CG.OrdPolyY+1);
end

if strcmp(CG.FeY,'LG')
  OrdPolyY=CG.OrdPolyY;
  CG.DSYLGL=zeros(OrdPolyY+2,OrdPolyY+1);
  CG.IntYLGL=zeros(OrdPolyY+2,OrdPolyY+1);
  CG.IntYLG=zeros(OrdPolyY+1,OrdPolyY+2);
  [CG.wYLGL,CG.xwYLGL]=GaussLobattoQuad(CG.OrdPolyY+1);
  for i=1:OrdPolyY+2
    for j=1:OrdPolyY+1
      CG.DSYLGL(i,j)=DLagrange(CG.xwYLGL(i),CG.xwY,j);
      CG.IntYLGL(i,j)=Lagrange(CG.xwYLGL(i),CG.xwY,j);
      CG.IntYLG(j,i)=Lagrange(CG.xwY(j),CG.xwYLGL,i);
    end
  end
  %CG.IntYLG=inv(diag(CG.wY))*CG.IntYLG*diag(CG.wYLGL);
  CG.DWYLGL=-inv(diag(CG.wY))*CG.DSYLGL'*diag(CG.wYLGL);
end
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -6.103907844236108 <= q <= 8.378675476959124)

                     add_NL_constraint(m, :(softplus(-0.8297607126794757 + -0.5064858067335005*softplus(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2])) + -0.8619658571743893*softplus(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2])) + 0.8415580445654567*softplus(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))) + softplus(-0.0809505831099564 + 0.6185835236357913*softplus(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2])) + -0.5288114683363174*softplus(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2])) + -0.42070009030061817*softplus(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))) + softplus(0.16011867292674786 + -0.14898607305663658*softplus(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2])) + -0.8220978033005912*softplus(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2])) + -0.8518875556726564*softplus(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using Pkg
Pkg.add("CSV")
Pkg.add("SpecialFunctions")

using CSV
using LinearAlgebra
using SpecialFunctions

X = CSV.read("C:/Users/dhojo/Desktop/X.csv")
y = CSV.read("C:/Users/dhojo/Desktop/y.csv")
n = size(X)[1]

X = convert(Matrix,X)
y = convert(Matrix,y)

r_0 = 0.06
s_0 = 6
Q_0 = [r_0 0;0 s_0]
a_0 = 6
b_0 = 600^2

## the exact value of marginal likelihood

M = X' * X + Q_0
R = I(n) - X * inv(M) * X'
ML = π^(-n/2) * b_0^(a_0/2) * (gamma((n+a_0)/2) / gamma(a_0/2)) *
  (det(Q_0)^(1/2) / det(M)^(1/2)) * ((y'*R*y)[1,1]+b_0)^(-(n+a_0)/2)
log(ML)

logML = -(n/2)*log(pi) + (a_0/2)*log(b_0) + log( gamma((n+a_0)/2)/gamma(a_0/2) ) +
  (1/2)*log( det(Q_0)/det(M) ) - ((n+a_0)/2)*log((y'*R*y)[1,1]+b_0)
logML


using SpecialFunctions, QuadGK, ForwardDiff

function fluence_spatial_freq(z, s, μa, μsp, l)

    #get numerical instabillity at large s
    D = @. 1/(3(μsp + μa))
    α = @. sqrt((D*s^2 + μa)/D)
    A = 0.493
    zb = 2*A*D[1]
    z0 = 1/(μsp[1] + μa[1])
    tolr = 1e-50

    ## equation 11
    if z < z0

        # have to account for exp blowing up to inf
        if exp(α[1]*z0 - α[1]*z) > 1e8
            ϕ = 0.0
        else
            ϕ = D[1]*α[1]*(1 + exp(-2*α[1]*(l - z))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l - z)))
            ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
            ϕ = ϕ*(1 + exp(-2*α[1]*(zb + z0)))/(2*D[1]*α[1])
            ϕ = ϕ*exp(α[1]*z0 - α[1]*z)
            ϕ = ϕ - ((exp(α[1]*(z0 - z))*(1 - exp(-2*α[1]*(z0 - z))))/(2*D[1]*α[1]))
        end

    ## equation 12
    elseif z > z0 && z < l
        ϕ = D[1]*α[1]*(1 + exp(-2*α[1]*(l - z))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l - z)))
        ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
        ϕ = ϕ*(1 + exp(-2*α[1]*(zb + z0)))/(2*D[1]*α[1])
        ϕ = ϕ*exp(α[1]*z0 - α[1]*z)

    #equation 13
    elseif z >= l

        ϕ = exp(α[1]*(z0 - l) + α[2]*(l - z))
        ϕ = ϕ*(1 - exp(-2*α[1]*(zb + z0)))
        ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) + D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
    end

    return ϕ
end

# compute equation 14 integrand
function compute_integral(ρ, z, s, μa, μsp, l)
    return fluence_spatial_freq(z, s, μa, μsp, l)*s*besselj0(ρ*s)
end

# numerically integrate equation 14
function fluence_steadystate(ρ, z, μa, μsp, l)
    return quadgk((s) -> compute_integral(ρ, z, s, μa, μsp, l), 0, Inf, rtol=1e-8)[1]/2π
 end

#compare semi-infinite fluence

function ss_compare(ρ, z, μa, μsp)

    D = 1/(3(μsp + μa))
    A = 0.493
    z0 = 1/(μsp + μa)
    zb = 2*A*D
    μe = sqrt(3*μa*(μa + μsp))

    ϕ = exp(-μe*sqrt((z - z0)^2 + ρ^2))/sqrt((z - z0)^2 + ρ^2)
    ϕ = ϕ - exp(-μe*sqrt((z + z0 + 2*zb)^2 + ρ^2))/sqrt((z + z0 + 2*zb)^2 + ρ^2)
    ϕ = ϕ/(4*π*D)
end

# Eqn 16 to compute reflectance (use auto diff)
function refl_steadystate(ρ, μa, μsp, l)

    f(z) = fluence_steadystate(ρ, z, μa, μsp, l)
    g(z) = ForwardDiff.derivative(z -> f(z), z)

    return 0.118*fluence_steadystate(ρ, 0, μa, μsp, l) + 0.306/(3*(μa[1] + μsp[1]))*g(0)
end


using RDatasets, DataFrames, InteractiveFixedEffectModels, Base.Test
df = dataset("plm", "Cigar")
df[:pState] =  pool(df[:State])
df[:pYear] =  pool(df[:Year])

for method in [:gauss_seidel,  :dogleg, :levenberg_marquardt]
	precision = 1e-1
	println(method)
	result = fit(InteractiveFixedEffectModel(:pState, :pYear, 1), :Sales, df , method =  method, maxiter = 10_000) ;
	@test_approx_eq_eps abs(result.augmentdf[:factors1][1]) 0.18662770198472406 precision
	@test_approx_eq_eps abs(result.augmentdf[:loadings1][1]) 587.2272 precision
	@test_approx_eq_eps result.augmentdf[:residuals][1] -15.6928 precision

	result = fit(InteractiveFixedEffectModel(:pState, :pYear, 2), :Sales,  df, method =  method, maxiter = 10_000) ;
	@test_approx_eq_eps abs(result.augmentdf[:factors1][1])  0.18662770198472406 precision
	@test_approx_eq_eps abs(result.augmentdf[:loadings1][1])  587.227 precision
	@test_approx_eq_eps result.augmentdf[:residuals][1] 2.16611 precision

	result = fit(InteractiveFixedEffectModel(:pState, :pYear, 1), Sales ~ 1 |> pState, df, method = method, maxiter = 10_000) ;
	@test_approx_eq_eps abs(result.augmentdf[:factors1][1])   0.17636 precision
	@test_approx_eq_eps abs(result.augmentdf[:loadings1][1])  20.176432452716522 precision
	@test_approx_eq_eps result.augmentdf[:residuals][1]  -10.0181 precision
	@test_approx_eq_eps result.augmentdf[:pState][1]  107.4766 precision

	result = fit(InteractiveFixedEffectModel(:pState, :pYear, 2), Sales ~ 1 |> pState, df, method =  method, maxiter = 10_000) ;
	@test_approx_eq_eps abs(result.augmentdf[:factors2][1]) 0.244 precision
	@test_approx_eq_eps abs(result.augmentdf[:loadings2][1]) 49.7943 precision
	@test_approx_eq_eps result.augmentdf[:residuals][1]  2.165319 precision
	@test_approx_eq_eps result.augmentdf[:pState][1]  107.47666 precision

	result = fit(InteractiveFixedEffectModel(:pState, :pYear, 2), Sales ~ 1 |> pState, df, method =  method, subset = (df[:State] .<= 30), maxiter = 10_000) ;
	@test size(result.augmentdf, 1) == size(df, 1)
	@test_approx_eq_eps abs(result.augmentdf[:factors1][1])  0.20215 precision
	@test_approx_eq_eps abs(result.augmentdf[:loadings1][1]) 29.546 precision
	@test_approx_eq_eps abs(result.augmentdf[:factors2][1])  0.250768  precision
	@test_approx_eq_eps abs(result.augmentdf[:loadings2][1])   43.578 precision
	@test_approx_eq_eps result.augmentdf[:residuals][1]  2.9448  precision
	@test_approx_eq_eps result.augmentdf[:pState][1] 107.4766 precision

end


"""
Calculation of differential cross sections
from structure functions.

Based on the Rxsecnc.h C++ code provided
by R. Aggarwal.
"""

using QCDNUM

# These globals will later be set by user
ZMass = 91.1876
WMass = 80.398
AlphaEM = 7.297352570e-03
GFermi = 1.16637e-05
TopMass = 171.2
BottomMass = 4.20

Vub = 41.2e-3
Vcb = 3.93e-3

Sin2ThetaW = 0.23127
Sin2ThetaC = 0.05
vu = 0.19164
vd = -0.34582
ve = -0.03746

au = 0.5
ad = -0.5
ae = -0.5

# Variables for Rxsecnc_xq2
sqrt_s = 318.1 # Should be configurable for calculation of y
Lepcharge = 1 # Should be configurable for calculation of Y

export dd_xsecnc_xq2_i
export f2_lo
export xf3_lo
export fl_lo

"""
    f2_lo(x, q2)

Calculate the f2_lo structure function term.
To be run after the evolution of PDFs with QCDNUM.
"""
function f2_lo(x::Float64, q2::Float64)::Float64

    # For weights
    pz = q2 / ((ZMass^2 + q2) * (4*Sin2ThetaW * (1 - Sin2ThetaW)))

    Au = 4.0/9.0 -2*pz*(2.0/3.0)*vu*ve + pz^2*(ve^2 + ae^2)*(vu^2 + au^2)

    Ad = 1.0/9.0 -2*pz*(-1.0/3.0)*vd*ve + pz^2*(ve^2 + ae^2)*(vd^2 + ad^2)

    # As in HERAPDF (top set to 0)
    weights = [0., Ad, Au, Ad, Au, Ad, 0., Ad, Au, Ad, Au, Ad, 0.]

    # Structure function calculation
    output = QCDNUM.zmstfun(2, weights, [x], [q2], 1, 0)

    output[1]
end

"""
    xf3_lo(x, q2)

Calculate the xf3_lo structure function term.
To be run after the evolution of PDFs with QCDNUM.
"""
function xf3_lo(x::Float64, q2::Float64)::Float64

    # For weights
    pz = q2 / ((ZMass^2 + q2) * (4*Sin2ThetaW * (1 - Sin2ThetaW)))

    Bu = -2*(2.0/3.0)*au*ae*pz + 4*au*ae*vu*ve*pz^2

    Bd = -2*(-1.0/3.0)*ad*ae*pz + 4*ad*ae*vd*ve*pz^2

    # weights = [-Bu, -Bd, -Bu, -Bd, -Bu, -Bd, 0.,  Bd, Bu, Bd, Bu, Bd, Bu]

    # As in HERAPDF (top set to 0)
    weights = [0., -Bd, -Bu, -Bd, -Bu, -Bd, 0.,  Bd, Bu, Bd, Bu, Bd, 0.]

    # Structure function calculation
    output = QCDNUM.zmstfun(3, weights, [x], [q2], 1, 0)

    output[1]
end

"""
    fl_lo(x, q2)

Calculate the fl_lo structure function term.
To be run after the evolution of PDFs with QCDNUM.
"""
function fl_lo(x::Float64, q2::Float64)::Float64

    # For weights
    pz = q2 / ((ZMass^2 + q2) * (4*Sin2ThetaW * (1 - Sin2ThetaW)))

    Au = 4.0/9.0 -2*pz*(2.0/3.0)*vu*ve + pz^2*(ve^2 + ae^2)*(vu^2 + au^2)

    Ad = 1.0/9.0 -2*pz*(-1.0/3.0)*vd*ve + pz^2(ve^2 + ae^2)*(vd^2 + ad^2)

    # As in HERAPDF (top set to 0)
    weights = [0., Ad, Au, Ad, Au, Ad, 0.,  Ad, Au, Ad, Au, Ad, 0.]

    output = QCDNUM.zmstfun(1, weights, [x], [q2], 1, 0)

    output[1]
end

"""
    rxsecnc_xq2_i(x, q2)

Reduced cross section for single x, q2.
"""
function rxsecnc_xq2_i(x::Float64, q2::Float64)::Float64

    rxsec = -1.0 #?
    y = 0.04  #?
    y = q2 / sqrt_s / sqrt_s / x

    Y_plus = 1 + (1 - y)^2
    Y_minus = 1 - (1 - y)^2

    F2 = f2_lo(x, q2)
    xF3 = xf3_lo(x, q2)
    FL = fl_lo(x, q2)

    if (Lepcharge == 1)

        rxsec =  F2 - (Y_minus / Y_plus)*xF3 - (y^2 / Y_plus) * FL

    elseif (Lepcharge == -1)

        rxsec =  F2 + (Y_minus / Y_plus)*xF3 - (y^2 / Y_plus) * FL

    end

    rxsec
end

"""
    rxsecnc_xq2(x, q2)

Reduced cross section for all bins.
"""
function rxsecnc_xq2(x_bin_cen::Array{Float64}, q2_bin_cen::Array{Float64})::Array{Float64}

    n_bins = length(x_bin_cen)

    # Check same number of x and q2 vals
    if (n_bins != length(q2_bin_cen))

        throw(DimensionMismatch("x_bin_cen and q2_bin_cen must have same length."))

    end

    # Calculate for each bin
    rxsec = zeros(n_bins)

    for i = 1:n_bins

        rxsec[i] = rxsecnc_xq2_i(x_bin_cen[i], q2_bin_cen[i])

    end

    rxsec
end

"""
    nc_propagator(q2, x)
"""
function nc_propagator(q2::Float64, x::Float64)::Float64

    y = q2 / sqrt_s / sqrt_s / x
    Yplus = 1 + (1 - y)^2
    alpha = AlphaEM
    conversion_factor = 0.3894e9;  # GeV^2 -> pb^2

    conversion_factor*2*π*Yplus*alpha^2 / (x * q2^2)
end

"""
    dd_xsecnc_xq2_i(x, q2)

Double differential cross section for single
x and q2 values.
NB: modifications needed to include pol and order.
"""
function dd_xsecnc_xq2_i(x::Float64, q2::Float64)

    dd_xsec = -1.0

    dd_xsec = nc_propagator(q2, x) * rxsecnc_xq2_i(x, q2)

    dd_xsec
end

"""
    dd_xsecnc_xq2(x_bin_cen, q2_bin_cen)

Double differential cross section for all x and
q2 bins.
NB: modifications needed to include pol and order.
"""
function dd_xsecnc_xq2(x_bin_cen::Array{Float64},
                       q2_bin_cen::Array{Float64})::Array{Float64}

    n_bins = length(x_bin_cen)

    # Check same number of x and q2 vals
    if (n_bins != length(q2_bin_cen))

        throw(DimensionMismatch("x_bin_cen and q2_bin_cen must have same length."))

    end

    xsec = zeros(n_bins)

    for i = 1:n_bins

        xsec[i] = dd_xsecnc_xq2_i(x_bin_cen[i], q2_bin_cen[i])

    end

  xsec
end


function ComputeIntrinsicAlignmentGrid!(CosmologicalGrid::CosmologicalGrid,
    LensingFunction::WLWeightFunction, ConvolvedDensity::AbstractConvolvedDensity,
    BackgroundQuantities::BackgroundQuantities,
    Cosmology::AbstractCosmology)
    c_0 = 2.99792458e5 #TODO: find a package containing the exact value of
                       #physical constants involved in calculations
    input_data = readdlm(joinpath(
        dirname(@__FILE__),"..","input_files","scaledmeanlum-E2Sa.txt"), Float64)
    z = input_data[:,1]
    lum = input_data[:,2]
    spl = Dierckx.Spline1D(z, lum, k = 2)
    lumgrid = spl.(CosmologicalGrid.ZArray)
    for zidx in 1:length(CosmologicalGrid.ZArray)
        for iidx in 1:length(ConvolvedDensity.ZBinArray)-1
            LensingFunction.IntrinsicAlignmentArray[iidx, zidx] =
            - BackgroundQuantities.HZArray[zidx] / c_0 *
            ConvolvedDensity.DensityGridArray[iidx, zidx] *
            LensingFunction.IntrinsicAlignmentModel.𝓐IA *
            LensingFunction.IntrinsicAlignmentModel.𝓒IA * Cosmology.ΩM *
            ( (1 + CosmologicalGrid.ZArray[zidx]) ^
            LensingFunction.IntrinsicAlignmentModel.ηIA ) *
            (lumgrid[zidx] ^ LensingFunction.IntrinsicAlignmentModel.βIA ) /
            BackgroundQuantities.DZArray[zidx]
        end
    end
end


# Everything is experimental! If you want to use it, you need to import it yourself.
# If you find problems, please open an issue.

using Optim
using Turing
using StatsBase
using LinearAlgebra

#      Run like
# using Turing
# @model height(heights) = begin
#     μ ~ Normal(178, 20)
#     σ ~ Uniform(0, 50)
#     heights .~ Normal(μ, σ)
# end
# m = height(d2.height)
# res = quap(m)

# Find the MAP via optimization and get the information matrix at that point.
# Look if your solution converged. Sometimes even solutions that didn't converge
# might be pretty good, on the other hand, just because the solver converged that
# doesn't mean you got the point you are looking for; this isn't even global
# optimization. In any case, trying other optimization methods might help.
function quap(model::Turing.Model, args...; kwargs...)
    opt = optimize(model, MAP(), args...; kwargs...)

    coef = opt.values.array
    var_cov_matrix = informationmatrix(opt)
    sym_var_cov_matrix = Symmetric(var_cov_matrix)  # lest MvNormal complains, loudly
    converged = Optim.converged(opt.optim_result)

    distr = if length(coef) == 1
        Normal(coef[1], √sym_var_cov_matrix[1])  # Normal expects stddev
    else
        MvNormal(coef, sym_var_cov_matrix)       # MvNormal expects variance matrix
    end

    (coef = coef, vcov = sym_var_cov_matrix, converged = converged, distr = distr)
end


function Mstep(d::LogNormal, y)
    mu = mean(log, y)
    sigma = sqrt(mean(x -> (log(x)-mu)^2, y))
    return LogNormal(mu, sigma)
end

function Mstep(d::GeneralizedGamma, y, lr=1)
    shp, scl, pwr = params(d)
    v = [log(shp), log(pwr)]
    g = ForwardDiff.gradient(v -> sum(y-> -logpdf(GeneralizedGamma(exp(v[1]),scl,exp(v[2])),y)), v)
    H = ForwardDiff.hessian(x -> sum(y-> -logpdf(GeneralizedGamma(exp(v[1]),scl,exp(v[2])),y)), v)
    v -= lr*(H \ g)
    b = mean(x -> x^exp(v[2]), y)^exp(-v[2])
    return GeneralizedGamma(exp(v[1]), b, exp(v[2]))
end

function Mstep(d::Gamma, y, lr=1)
    shp, scl = params(d)
    rho = log(shp)
    meanlogy = mean(log,y)
    #f(x,shp) = -shp*log(scl) + (shp-1)*log(x) - (x/scl) -loggamma(shp)
    g = shp*log(scl) - shp*meanlogy + digamma(shp)*shp
    H = shp*log(scl) - shp*meanlogy + trigamma(shp)*shp^2 + digamma(shp)*shp
    rho -= lr*(g/H)
    b = mean(y)/exp(rho)
    return Gamma(exp(rho), b)
end

function Mstep(d::Weibull, y, lr=1)
    shp, scl = params(d)
    rho = log(shp)
    logy = sum(log, y)
    n = length(y)
    A = ((y/scl).^shp)'*(log.(y/scl))
    B = ((y/scl).^shp)'*(log.(y/scl).^2)
    #f(x,shp) = log(shp)-shp*log(scl) + (shp-1)*log(x) - (x/scl)^shp
    g = n*(shp*log(scl)-1) - shp*logy + shp*A
    H = n*(shp*log(scl)-1) - shp*logy + shp*A + (shp^2)*B
    rho -= lr*(g/H)
    b = mean(x -> x^exp(rho), y)^exp(-rho)
    return Weibull(exp(rho), b)
end

# function Mstep(d::Gamma, y)
#     mlog = log(mean(y))
#     logm = mean(log, y)
#     a = shape(d)
#     ia = inv(a)
#     a = inv(ia + (-mlog + logm + log(a) - digamma(a)) / ((a^2) * (ia - trigamma(a))))
#     b = mean(y) / a
#     return Gamma(a,b)
# end

# function Mstep(d::Weibull, y)
#     a, b = params(d)
#     n = length(y)
#     powa = (y.^a)'
#     logy = log.(y)

#     fa = n*(inv(a) - inv(b)) + sum(logy)- n*log(b) - powa*(logy-log(b))/(b^a)
#     dfa = -n*inv(a^2) - (powa*(logy-log(b)).^2)/(b^a)

#     Delta = fx / dfx
#     a += Delta
#     b = mean(x -> x^a, ytilde)^inv(a)
#     return Weibull(a,b)
# end

function Mstep(d::Exponential, y)
    b = mean(y)
    return Exponential(b)
end

#######
#interval censored
function MCEMic(rng, dist, iter, EL, ER, S, lr=1)
    lp = zeros(iter)
    pars = params(dist)
    for it in 1:iter
    ytilde = [rand(rng, truncated(dist,S[i]-ER[i],S[i]-EL[i])) for i in eachindex(S)]
    dist = Mstep(dist, ytilde, lr)
    lp[it] = mean(x-> -logpdf(dist,x), ytilde)
    end
    return dist,lp
end

########
#interval censored with right truncated
function Estep_icrt(rng, dist, EL, ER, S, Tmax)
    N = length(S)
    ys = zeros(N)
    yb = []
    for i in 1:N
        ys[i] = rand(rng, truncated(dist,S[i]-ER[i],S[i]-EL[i]))
        u = rand(rng)
        E = EL[i] + (ER[i]-EL[i])*u
        q = cdf(dist, Tmax - E)
        if 0.0 < q < 1.0
            B = rand(rng, Geometric(q))
            if B > 0
                ysub = [rand(rng, truncated(dist,Tmax-E,Inf)) for i in 1:B]
                yb = [yb;ysub]
            end
        end
    end
    return [ys ; yb]
end

function MCEMicrt(rng, dist, iter, EL, ER, S, Tmax, lr=1)
    lp = zeros(iter)
    pars = params(dist)
    for it in 1:iter
    ytilde = Estep_icrt(rng, dist, EL, ER, S, Tmax)
    dist = Mstep(dist, ytilde, lr)
    lp[it] = mean(x-> -logpdf(dist,x), ytilde)
    end
    return dist,lp
end

######
#doubly interval censored
# function Estep_dic(rng, dist, EL, ER, SL, SR)
#     N = length(EL)
#     ys = zeros(N)
#     for i in 1:N
#         u = rand(rng)
#         E = EL[i] + (ER[i]-EL[i])*u
#         ys[i] = rand(rng, truncated(dist, SL[i]-E, SR[i]-E))
#     end
#     return ys
# end

function Estep_dic(rng, dist, EL, ER, SL, SR)
    N = length(EL)
    ys = zeros(N)
    for i in 1:N
        u = rand(rng)
        S = SL[i] + (SR[i]-SL[i])*u
        ys[i] = rand(rng, truncated(dist, S-ER[i], S-EL[i]))
    end
    return ys
end

function MCEMdic(rng, dist, iter, EL, ER, SL, SR, lr=1)
    lp = zeros(iter)
    pars = params(dist)
    for it in 1:iter
    ytilde = Estep_dic(rng, dist, EL, ER, SL, SR)
    dist = Mstep(dist, ytilde, lr)
    lp[it] = mean(x-> -logpdf(dist, x), ytilde)
    end
    return dist,lp
end