# Elements of Chemical Reaction Engineering
# Fifth Edition
# H. SCOTT FOGLER
# Chapter 13: Unsteady-State Nonisothermal Reactor Design
# Section 13.5: Nonisothermal Multiple Reactions
# Example 13–5 Multiple Reactions in a Semibatch Reactor
# p. 658

function semibatch_reactor(; store_results=false::Bool)
    datadir = nothing
    if store_results
        datadir = generate_data_subdir(@__FILE__)
    end

    # initial conditions and timepoints
    t0 = 0.0f0
    tf = 0.4f0  # Bradfoard uses 0.4
    Δt = 0.01f0
    tspan = (t0, tf)

    # state: CA, CB, CC, T, Vol
    u0 = [1.0f0, 0.0f0, 0.0f0, 290.0f0, 100.0f0]

    # control constraints
    # F = volumetric flow rate
    # V = exchanger temperature
    # F = 240 & V = 298 in Fogler's book
    # F ∈ (0, 250) & V ∈ (200, 500) in Bradford 2017
    control_ranges = [(100.0f0, 700.0f0), (0.0f0, 400.0f0)]

    # state constraints
    # T ∈ (0, 420]
    # Vol ∈ (0, 200]
    T_up = 380.0f0
    V_up = 100.0f0

    system_params = (
        CpA=30.0f0,
        CpB=60.0f0,
        CpC=20.0f0,
        CpH2SO4=35.0f0,
        N0H2S04=100.0f0,
        T0=305.0f0,
        CA0=4.0f0,
        HRA=-6500.0f0,
        HRB=8000.0f0,
        E1A=9500.0f0 / 1.987f0,
        E2B=7000.0f0 / 1.987f0,
        A1=1.25f0,
        A2=0.08f0,
        UA=35000.0f0,  # 45000,  Bradford value
        Tr1=320.0f0,  # 420, Bradford value
        Tr2=290.0f0,  # 400, Bradford value
    )

    function system!(du, u, p, t, controller; input=:state)
        @argcheck input in (:state, :time)

        CpA, CpB, CpC, CpH2SO4, N0H2S04, T0, CA0, HRA, HRB, E1A, E2B, A1, A2, UA, Tr1, Tr2 = values(
            system_params
        )

        # neural network outputs controls taken by the system
        CA, CB, CC, T, Vol = u  # states
        # controls
        if input == :state
            c_F, c_T = controller(u, p)
        elseif input == :time
            c_F, c_T = controller(t, p)
        end
        k1A = A1 * exp(E1A * ((1.0f0 / Tr1) - (1.0f0 / T)))
        k2B = A2 * exp(E2B * ((1.0f0 / Tr2) - (1.0f0 / T)))

        k1CA = k1A * CA
        k2CB = k2B * CB
        F_Vol = c_F / Vol

        ra = -k1CA
        rb = 0.5f0 * k1CA - k2CB
        rc = 3.0f0 * k2CB

        num =
            UA * (c_T - T) - CA0 * c_F * CpA * (T - T0) +
            (HRA * (-k1CA) + HRB * (-k2CB)) * Vol
        den = (CA * CpA + CpB * CB + CpC * CC) * Vol + N0H2S04 * CpH2SO4

        # dynamics of the controlled system
        dCA = ra + (CA0 - CA) * F_Vol
        dCB = rb - CB * F_Vol
        dCC = rc - CC * F_Vol
        dT = num / den
        dVol = c_F

        # update in-place
        @inbounds begin
            du[1] = dCA
            du[2] = dCB
            du[3] = dCC
            du[4] = dT
            du[5] = dVol
        end
    end

    # simulate the system with constant controls as in Fogler's
    # to reproduce his results and verify correctness
    fogler_ref = [240.0f0, 298.0f0]  # reference values in Fogler
    fogler_timespan = (0.0f0, 1.5f0)
    fixed_controlODE = ControlODE((u, p) -> fogler_ref, system!, u0, fogler_timespan)
    @info "Fogler's case: final time state" solve(fixed_controlODE, nothing).u[end]
    plot_simulation(
        fixed_controlODE,
        nothing;
        only=:states,
        vars=[1, 2, 3],
    )
    plot_simulation(
        fixed_controlODE,
        nothing;
        only=:states,
        vars=[4, 5],
    )

    function collocation(
        u0;
        num_supports::Integer=length(controlODE.tsteps),
        nodes_per_element::Integer=3,
        constrain_states::Bool=false,
    )
        optimizer = optimizer_with_attributes(
            Ipopt.Optimizer, "print_level" => 0, "check_derivatives_for_naninf" => "yes"
        )
        model = InfiniteModel(optimizer)
        method = OrthogonalCollocation(nodes_per_element)
        @infinite_parameter(
            model, t in [t0, tf], num_supports = num_supports, derivative_method = method
        )

        @variables(
            model,
            begin
                # state variables
                x[1:5], Infinite(t)
                # control variables
                c[1:2], Infinite(t)
            end
        )

        # tricks to make IPOPT work...
        almost_zero = 1.0f-3
        close_to_infinity = 1.0f3

        set_start_value_function(x[1], t -> 0.5f0)
        set_start_value_function(x[2], t -> 0.5f0)
        set_start_value_function(x[3], t -> 0.5f0)
        set_start_value_function(x[4], t -> 2.0f2)
        set_start_value_function(x[5], t -> 1.0f2)

        set_start_value_function(c[1], t -> 2.0f2)
        set_start_value_function(c[2], t -> 1.0f2)

        @constraints(
            model,
            begin
                almost_zero <= x[1] <= close_to_infinity
                almost_zero <= x[2] <= close_to_infinity
                almost_zero <= x[3] <= close_to_infinity
                almost_zero <= x[4] <= close_to_infinity
                almost_zero <= x[5] <= close_to_infinity
            end
        )

        # fixed_parameters
        CpA, CpB, CpC, CpH2SO4, N0H2S04, T0, CA0, HRA, HRB, E1A, E2B, A1, A2, UA, Tr1, Tr2 = values(
            system_params
        )

        initial_conditions = @constraints(
            model,
            begin
                x[1](0) == u0[1]
                x[2](0) == u0[2] + almost_zero
                x[3](0) == u0[3] + almost_zero
            end
        )

        control_constraints = @constraints(
            model,
            begin
                control_ranges[1][1] <= c[1] <= control_ranges[1][2]
                control_ranges[2][1] <= c[2] <= control_ranges[2][2]
            end
        )

        if constrain_states
            state_constraints = @constraints(
                model,
                begin
                    x[4] <= T_up
                    x[5] <= V_up
                end
            )
        end

        dynamic_constraints = @constraints(
            model,
            begin
                ∂(x[1], t) ==
                -(A1 * exp(E1A * ((1.0f0 / Tr1) - (1.0f0 / x[4]))) * x[1]) +
                (CA0 - x[1]) * c[1] / x[5]
                ∂(x[2], t) ==
                0.5f0 * (A1 * exp(E1A * ((1.0f0 / Tr1) - (1.0f0 / x[4]))) * x[1]) -
                (A2 * exp(E2B * ((1.0f0 / Tr2) - (1.0f0 / x[4]))) * x[2]) -
                x[2] * c[1] / x[5]
                ∂(x[3], t) ==
                3.0f0 * (A2 * exp(E2B * ((1.0f0 / Tr2) - (1.0f0 / x[4]))) * x[2]) -
                x[3] * c[1] / x[5]
                ∂(x[4], t) ==
                (
                    UA * (c[2] - x[4]) - CA0 * c[1] * CpA * (x[4] - T0) +
                    (
                        HRA * (-(A1 * exp(E1A * ((1.0f0 / Tr1) - (1.0f0 / x[4]))) * x[1])) +
                        HRB * (-(A2 * exp(E2B * ((1.0f0 / Tr2) - (1.0f0 / x[4]))) * x[2]))
                    ) * x[5]
                ) / ((x[1] * CpA + CpB * x[2] + CpC * x[3]) * x[5] + N0H2S04 * CpH2SO4)
                ∂(x[5], t) == c[1]
            end
        )

        @objective(model, Max, x[2](tf))

        optimize!(model)
        jump_model = optimizer_model(model)

        # list possible termination status: model |> termination_status |> typeof
        @info solution_summary(jump_model; verbose=false)
        if Int(termination_status(jump_model)) ∉ (1, 4)  # OPTIMAL = 1, LOCALLY_SOLVED = 4
            @warn raw_status(jump_model) termination_status(jump_model)
        end

        states = hcat(value.(x)...) |> permutedims
        controls = hcat(value.(c)...) |> permutedims

        return model, supports(t), states, controls
    end

    # set arquitecture of neural network controller
    controller = FastChain(
        (x, p) -> [1f2x[1], 1f2x[2], 1f2x[3], x[4], x[5]],
        FastDense(5, 8, tanh_fast),
        FastDense(8, 8, tanh_fast),
        FastDense(8, 2),
        # (x, p) -> [240f0, 298f0],
        scaled_sigmoids(control_ranges),
    )

    control_profile, infopt_model, times_collocation, states_collocation, controls_collocation = collocation_preconditioner(
        u0,
        collocation;
        plot=false,
        nodes_per_element=2,
        constrain_states=false,
    )
    _, _, _, states_collocation_constrained, controls_collocation_constrained = collocation_preconditioner(
        u0,
        collocation;
        plot=false,
        nodes_per_element=2,
        constrain_states=true,
    )
    #=
    plt.figure()
    plt.plot(times_collocation, states_collocation[1, :]; label="s1")
    plt.plot(times_collocation, states_collocation[2, :]; label="s2")
    plt.plot(times_collocation, states_collocation[3, :]; label="s3")
    plt.plot(
        times_collocation, states_collocation_constrained[1, :];
        label="s1_constrained", color=plt.gca().lines[1].get_color(), ls="dashdot"
    )
    plt.plot(
        times_collocation, states_collocation_constrained[2, :];
        label="s2_constrained", color=plt.gca().lines[2].get_color(), ls="dashdot"
    )
    plt.plot(
        times_collocation, states_collocation_constrained[3, :];
        label="s3_constrained", color=plt.gca().lines[3].get_color(), ls="dashdot"
    )
    plt.legend()
    plt.show()

    plt.figure()
    plt.plot(times_collocation, states_collocation[4, :]; label="s4")
    plt.plot(times_collocation, states_collocation[5, :]; label="s5")
    plt.plot(
        times_collocation, states_collocation_constrained[4, :];
        label="s4_constrained", color=plt.gca().lines[1].get_color(), ls="dashdot"
    )
    plt.plot(
        times_collocation, states_collocation_constrained[5, :];
        label="s5_constrained", color=plt.gca().lines[2].get_color(), ls="dashdot"
    )
    plt.legend()
    plt.show()

    plt.figure()
    plt.plot(times_collocation, controls_collocation[1, :]; label="c1")
    plt.plot(times_collocation, controls_collocation[2, :]; label="c2")
    plt.plot(
        times_collocation, controls_collocation_constrained[1, :];
        label="c1_constrained", color=plt.gca().lines[1].get_color(), ls="dashdot"
    )
    plt.plot(
        times_collocation, controls_collocation_constrained[2, :];
        label="c2_constrained", color=plt.gca().lines[2].get_color(), ls="dashdot"
    )
    plt.legend()
    plt.show()
    =#

    controlODE = ControlODE(controller, system!, u0, tspan; tsteps)

    θ = preconditioner(
        controlODE,
        control_profile;
        progressbar=true,
        plot_progress=false,
        # control_range_scaling=[range[end] - range[1] for range in control_ranges],
    )

    plot_simulation(controlODE, θ; only=:states)
    plot_simulation(controlODE, θ; only=:controls)
    store_simulation("precondition", controlODE, θ; datadir)

    # objective function splitted componenets to optimize
    function losses(controlODE, params; α=1.0f-3, δ=1.0f1, ρ=1.0f-2, kwargs...)

        # integrate ODE system
        sol_array = Array(solve(controlODE, params; kwargs...))

        # running cost
        out_temp = map(x -> relaxed_log_barrier(T_up - x; δ), sol_array[4, 1:end])
        out_vols = map(x -> relaxed_log_barrier(V_up - x; δ), sol_array[5, 1:end])

        # terminal cost
        # L = - (100 x₁ - x₂) + penalty  # Bradford
        objective = -sol_array[2, end]

        # integral penalty
        state_penalty = α * Δt * (sum(out_temp) + sum(out_vols))
        control_penalty = 0.0f0
        regularization = ρ * mean(abs2, θ)
        return objective, state_penalty, control_penalty, regularization
    end

    # α: penalty coefficient
    # δ: barrier relaxation coefficient
    α0, δ0 = 1.0f-5, 1.0f1
    barrier_iterations = 0:8
    αs = [α0 for _ in barrier_iterations]
    δs = [δ0 * 0.8f0^i for i in barrier_iterations]
    θ = constrained_training(
        controlODE,
        losses;
        starting_params=θ,
        αs,
        δs,
        show_progressbar=true,
        # plots_callback,
        datadir,
    )

    @info "Final states"
    plot_simulation(controlODE, θ; only=:states, vars=[1, 2, 3])
    plot_simulation(
        controlODE, θ; only=:states, vars=[4, 5], yrefs=[T_up, V_up]
    )

    @info "Final controls"
    plot_simulation(controlODE, θ; only=:controls)#  only=:states, vars=[1,2,3])

    @info "Final loss" losses(controlODE, θ; δ=δs[end])

    final_objective, final_state_penalty, final_control_penalty, final_regularization = losses(
        controlODE, θ; δ=δs[end]
    )

    return store_simulation(
        "constrained",
        controlODE,
        θ;
        datadir,
        metadata=Dict(
            :objective => final_objective,
            :state_penalty => final_state_penalty,
            :cotrol_penalty => final_control_penalty,
            :regularization => final_regularization,
            :num_params => length(θ),
            :layers => controller_shape(controller),
            :deltas => δs,
            :t0 => t0,
            :tf => tf,
            :Δt => Δt,
        ),
    )
end  # function wrapper


#fdtd.jl
using DelimitedFiles

#
# single 1d fdtd step in-place
#
function FDTDStep1D!(e::Array{T}, h::Array{T}, invϵ::Array{T}, Δx::T, Δt::T) where T
    ne = length(e)
    nh = length(h)
    dtdxFactor = Δt / Δx

    #e-field pass
    e[1] = e[end] = zero(T) #ensure boundary condition
    @inbounds for i = 2:(ne - 1) #spare first and last element
        e[i] += dtdxFactor * invϵ[i] * (h[i - 1] - h[i])
    end
    #h-field pass
    @inbounds for i = 1:nh
        h[i] += dtdxFactor * (e[i] - e[i + 1])
    end
end

#
# fdtd 1d simulation
#
function FDTD1D(steps::Int, snapshots::Array{Int}, n::Int, Δx::Float64, Δt::Float64, ϵBorder::Float64 = Inf)

    # spatial coordinates (not needed for calculation)
    x0 = (n - 1.0) / 2.0 * Δx # artificial midpoint
    x_e = collect(range(-x0 / 2.0         , length = n,     step = Δx))
    x_h = collect(range(-x0 / 2.0 + Δx / 2, length = n - 1, step = Δx))

    #initial fields: gaussian packet
    e = exp.(-0.5e2 * (n - 1) * Δx .* x_e .^ 2)
    h = exp.(-0.5e2 * (n - 2) * Δx .* x_h .^ 2)
    invϵ = map(x -> if (x > ϵBorder) 1.0 / 4.0 else 1.0 end, x_e)

    #initialize the result variables and snapshot-track-variables
    snapshotIndex = 1
    nextSnapshotIndex = snapshots[1]
    eResult = Array{Float64}(undef, length(snapshots), n)
    hResult = Array{Float64}(undef, length(snapshots), n - 1)

    #lets go
    for i = 1:steps
        FDTDStep1D!(e, h, invϵ, Δx, Δt) #magic

        if i == nextSnapshotIndex #if we reached a snapshot point:
            #copy the current fields
            for m = 1:n
                eResult[snapshotIndex, m] = e[m]
            end
            for m = 1:(n - 1)
                hResult[snapshotIndex, m] = h[m]
            end
            if snapshotIndex < length(snapshots) #and advance the snapshot-track-variables
                snapshotIndex += 1
                nextSnapshotIndex = snapshots[snapshotIndex]
            end
        end
    end

    #return out the results
    (efield = eResult, xe = x_e, hfield = hResult, xh = x_h, snaps = snapshots)
end

#
# save result in csv file
#
function SaveFDTDResult(resultFile, result)
    snapshots = result.snaps
    open(resultFile, "w") do io
        writedlm(io, [result.snaps, result.xe, result.xh], ',')
        writedlm(io, result.efield, ',')
        writedlm(io, result.hfield, ',')
    end
end


function run()
  println("running...")

  # initial step
  evaluate()

  # main loop
  while t < t_end
    dt = timestep()

    std_output(dt)

    integrate(dt)

    output(t, t_output)

    update_time(t, dt)
  end
end

function std_output(dt)
  diagnose()
  Printf.@printf("t = %5f, dt = %.5e, p_err = %.8e, E_err = %.8e, L1_err = %.8e\n", t, dt, p_error,
                   E_error, L1_error)
  log_history(dt)
end

function diagnose()
  global L1_error = measure_L1error()
  p_sum = 0
  p_sum_abs = 0
  E_sum = 0
  E_sum_abs = 0
  for i in 1:length(sph)
    p_sum += sph[i].p
    p_sum_abs += abs(sph[i].p)
    E_sum += sph[i].U + 0.5 * sph[i].p^2 / sph[i].m
    E_sum_abs += abs(sph[i].U + 0.5 * sph[i].p^2 / sph[i].m)
  end
  global p_error = abs(p_sum - p_initial) / p_sum_abs
  global E_error = abs(E_sum - E_initial) / E_sum_abs
end

function measure_L1error()
  # exact solution of Riemann problem
  step_riemann = lbox / Nsample_riemann
  x_riemann = x_min + step_riemann/2:step_riemann:x_max - step_riemann/2
  rho_riemann_array = map(x -> rho_riemann(x, t, v_s, P_s, x_center), x_riemann)

  # calculate mean density in each bin
  numbin = ceil(Int, (x_max - x_min)/L1_binsize)
  rhobar_exact = zeros(numbin)
  count_exact = zeros(numbin)
  rhobar_sim = zeros(numbin)
  count_sim = zeros(numbin)
  for i in 1:length(x_riemann)
    index = ceil(Int, (x_riemann[i] - x_min)/L1_binsize)
    rhobar_exact[index] += rho_riemann_array[i]
    count_exact[index] += 1
  end
  for i in 1:length(sph)
    index = ceil(Int, (sph[i].x - x_min)/L1_binsize)
    rhobar_sim[index] += sph[i].rho
    count_sim[index] += 1
  end
  for i in 1:numbin
    rhobar_exact[i] = count_exact[i] > 0 ? rhobar_exact[i]/count_exact[i] : 0
    rhobar_sim[i] = count_sim[i] > 0 ? rhobar_sim[i]/count_sim[i] : 0
  end

  # measure L1 error of [x_min + (1/4)lbox, x_min + (3/4)lbox]
  L1error = 0
  bin25 = ceil(Int, 0.25*lbox/L1_binsize)
  bin75 = ceil(Int, 0.75*lbox/L1_binsize)
  for i in bin25:bin75
    L1error += abs(rhobar_exact[i] - rhobar_sim[i])/(bin75 - bin25)
  end
  return L1error
end

function log_history(dt)
  open(hstfilename, "a") do io
    println(io, "$t $dt $p_error $E_error $L1_error")
  end
end


push!(LOAD_PATH, abspath(@__DIR__,".."))
using Plots
using AttractorNetworksBase ; const A = AttractorNetworksBase
using Calculus
using BenchmarkTools, Cthulhu
using Random
using Test
using Statistics,LinearAlgebra
Random.seed!(0)

using Plots ; theme(:dark)
function plotvs(x::AbstractArray{<:Real},y::AbstractArray{<:Real})
  x,y=x[:],y[:]
  @info """
  The max differences between the two are $(extrema(x .-y ))
  """
  plt=plot()
  scatter!(plt,x,y;leg=false,ratio=1,color=:white)
  lm=xlims()
  plot!(plt,identity,range(lm...;length=3);linestyle=:dash,color=:yellow)
  return plt
end
function test_Jgradient(utest::Vector{R},ntw::A.BaseNetwork) where R
    Jalloc = A.JAlloc(ntw)
    gradW=similar(ntw.weights)
    gradu=similar(ntw.weights)
    J=similar(gradW)
    # analytic gradient
    A.jacobian!(J,gradW,gradu,utest,ntw,Jalloc)
    # numerical gradient weights
    grad_numW = similar(gradW)
    ntw2=copy(ntw)
    for ij in eachindex(ntw2.weights)
        function gradfun(w)
            ntw2.weights[ij] = w
            A.jacobian!(J,nothing,nothing,utest,ntw2,Jalloc)
            return J[ij]
        end
        grad_numW[ij] = Calculus.gradient(gradfun,ntw2.weights[ij])
    end
    # numerical gradient u
    grad_numu = similar(gradu)
    for lj in CartesianIndices(gradu)
        (j,l)=Tuple(lj)
        function gradfun(u)
            u2 = copy(utest)
            u2[l]=u
            A.jacobian!(J,nothing,nothing,u2,ntw,Jalloc)
            return J[lj]
        end
        grad_numu[lj] = Calculus.gradient(gradfun,utest[l])
    end
    return gradW,grad_numW,gradu,grad_numu
end

## linear stable dynamcs goes to zero
ne,ni = 59,33
ntot = ne+ni
ntw = A.BaseNetwork(ne,ni ; gfun=A.IOId{Float64}())
ntw.weights .= 0.3.*randn(ntot)
B = ntw.weights - I
ntw.external_input .= 0.0
ntw.membrane_taus .= 1.0
r_start = 30.0.*randn(ntot)

# check velocity
A.velocity(r_start,ntw)

u_end,r_end=A.run_network_to_convergence(ntw,r_start;veltol=1E-5)

vel_test=similar(u_end)
A.velocity!(vel_test,u_end,ntw.iofunction.(u_end),ntw)

t,rdyn,_=A.run_network(ntw,r_end,4.0)

plot(t,rdyn[4,:];leg=false)

##
fill!(ntw.weights,0.0)
u_start = randn(A.ndims(ntw))
r_start = ntw.iofunction.(u_start)

u_end,r_end=A.run_network_to_convergence(ntw,r_start)
@test all(isapprox.(u_end,ntw.external_input;atol=0.05))

ne,ni = 20,23
ntot = ne+ni
ntw = A.BaseNetwork(ne,ni ; gfun=A.IOQuad(0.02))
u_start = randn(A.ndims(ntw))
r_start = ntw.iofunction.(u_start)
u_end,r_end=A.run_network_to_convergence(ntw,r_start)
t,ut,rt=A.run_network(ntw,r_end,1.0)
@test all(isapprox.(u_end,ut[:,end];atol=0.05))

##

##
Wtest = A.make_wmat(80,20,A.WDalian)
heatmap(Wtest ; ratio=1)

##
pe,pi = A.populations(80,20)
mah = rand(100,100)

testee = pe*pe'
testie = pi*pe'
all( view(testee,pe,pe))
all( view(testie,pi,pe))

view(mah,pe*pe')

heatmap(testie ; ratio=1)
heatmap(pe*pi' ; ratio = 1)
heatmap(pi*pe' ; ratio = 1)

mah[pe,pe] .- mah[1:80,1:80]

##
ne,ni = 34,44
wmeanf = [ 0.333  -1.123
            12.122 -0.06432]
W = A.make_wmat(ne,ni,wmeanf;noautapses=true)

pe,pi = A.populations(ne,ni)

sum(W[pe,pe]; dims=2)
#
# whatevs = A.run_network(10.0 .* rand(ntot), 0.5,ntw)
#
# using Plots; theme(:dark)
# _ = let t=whatevs.t,
#   u12 = [ uu[1:2] for uu in whatevs.u]
#   u12=hcat(u12...) |> permutedims
#   plot(t,u12;leg=false,linewidth=4)
# end

##

ne,ni = 4,3
ntot = ne+ni
ntw = A.BaseNetwork(ne,ni ; gfun=A.IOQuad(0.02))

utest = randn(ntot)

grad_ana,grad_num = test_Jgradient_weights(utest,ntw)
grad_ana,grad_num = test_Jgradient_u(utest,ntw)

grad_ana ./ grad_num

grad_num

grad_ana


# Nonequivalent Groups : Linear methods
function ObservableStats(F::NEAT)
    x = F.rawX; v = F.rawV
    μx = mean(x); σx = std(x)
    μv = mean(v); σv = std(v)
    covxv = cov(x, v); corxv = cor(x, v)
    return μx, σx, μv, σv, covxv, corxv
end

function SummaryStats(F::NEAT)
    r(x) = round2(x; digits = 2)
    x = F.rawX; v = F.rawV
    μx = mean(x); σx = std(x); kx = kurtosis(x); sx = skewness(x);
    μv = mean(v); σv = std(v); kv = kurtosis(v); sv = skewness(v);
    df = DataFrame(
        test = ["X", "V"], mean = [μx, μv], sigma = [σx, σv], kurtosis = [kx, kv], skewness = [sx, sv],
        min = [minimum(x), minimum(v)], max = [maximum(x), maximum(v)], N = [length(x), length(v)]
    )
    println(df)
    return df;
end

mutable struct NEATEquateResult <: NEATEquateMethod
    method
    table
    synthetic
    estimates
    data
end

"""
    Tucker(X::NEAT, Y::NEAT; w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁)

Tucker observed score method.
"""
function Tucker(X::NEAT, Y::NEAT; w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁)
    W = w₁ + w₂
    w₁ = w₁ / W; w₂ = w₂ / W
    # test score
    x = X.rawX; xv = X.rawV
    y = Y.rawX; yv = Y.rawV
    # ObservableStats
    μx, σx, μxv, σxv, covxv, corxv = ObservableStats(X)
    μy, σy, μyv, σyv, covyv, coryv = ObservableStats(Y)
    # regression slope
    γ₁ = covxv / σxv^2
    γ₂ = covyv / σyv^2
    # synthetic mean and var
    μsX = μx - w₂*γ₁*(μxv-μyv)
    μsY = μy + w₁*γ₂*(μxv-μyv)
    σ²sX = σx^2 - w₂*γ₁^2*(σxv^2-σyv^2) + w₁*w₂*γ₁^2*(μxv-μyv)^2
    σ²sY = σy^2 + w₁*γ₂^2*(σxv^2-σyv^2) + w₁*w₂*γ₂^2*(μxv-μyv)^2
    # transformation
    slope = sqrt(σ²sY)/sqrt(σ²sX); intercept = μsY - slope*μsX
    lYx = @. X.tableX.scale * slope + intercept
    tbl = DataFrame(scaleX = X.tableX.scale, lYx = lYx)
    return NEATEquateResult(
        :Tucker,
        tbl,
        DataFrame(Group = [1, 2], μ = [μsX, μsY], σ = [sqrt(σ²sX), sqrt(σ²sY)], γ = [γ₁, γ₂], w = [w₁, w₂]),
        (slope = slope, intercept = intercept),
        (X = X, Y = Y)
    )
end
# Nonequivalent Groups : Levine under a classical congeneric model
"""
    LevineTrueScore(X::NEAT, Y::NEAT; common = :external)
Levine true score method applied to observed scores. This algorithm assumes Classical Congeneric model, which dosen't need to estimate true scores.
"""
function LevineTrueScore(X::NEAT, Y::NEAT; common = :external)
    # test score
    x = X.rawX; xv = X.rawV
    y = Y.rawX; yv = Y.rawV
    # ObservableStats
    μx, σx, μxv, σxv, covxv, corxv = ObservableStats(X)
    μy, σy, μyv, σyv, covyv, coryv = ObservableStats(Y)
    # regression slope
    if common == :internal
        γ₁ = σx^2 / covxv
        γ₂ = σy^2 / covyv
    elseif common == :external
        γ₁ = (σx^2 + covxv) / (σxv^2 + covxv)
        γ₂ = (σy^2 + covyv) / (σyv^2 + covyv)
    else
        println("`common` argument can recieve only :internal or :external")
    end
    # doesn't use the synthetic population.
    # transformation
    slope = γ₂/γ₁ ; intercept = μy + γ₂ * (μxv - μyv) - slope*μx
    lYx = @. X.tableX.scale * slope + intercept
    tbl = DataFrame(scaleX = X.tableX.scale, lYx = lYx)
    return NEATEquateResult(
        :LevineTrueScore,
        tbl,
        DataFrame(Group = [1, 2], μ = [μx, μy], σ = [σx, σy], γ = [γ₁, γ₂], w = [missing, missing]),
        (slope = slope, intercept = intercept),
        (X = X, Y = Y)
    )
end

"""

Levine observed score method. This algorithm assumes Classical Congeneric model, which dosen't need to estimate true scores.
"""
function LevineObservedScore(X::NEAT, Y::NEAT; common = :external, w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁)
    W = w₁ + w₂
    w₁ = w₁ / W; w₂ = w₂ / W
    # test score
    x = X.rawX; xv = X.rawV
    y = Y.rawX; yv = Y.rawV
    # ObservableStats
    μx, σx, μxv, σxv, covxv, corxv = ObservableStats(X)
    μy, σy, μyv, σyv, covyv, coryv = ObservableStats(Y)
    # regression slope
    if common == :internal
        γ₁ = σx^2 / covxv
        γ₂ = σy^2 / covyv
    elseif common == :external
        γ₁ = (σx^2 + covxv) / (σxv^2 + covxv)
        γ₂ = (σy^2 + covyv) / (σyv^2 + covyv)
    else
        println("`common` argument can recieve only :internal or :external")
    end
    # synthetic mean and var
    μsX = μx - w₂*γ₁*(μxv-μyv)
    μsY = μy + w₁*γ₂*(μxv-μyv)
    σ²sX = σx^2 - w₂*γ₁^2*(σxv^2-σyv^2) + w₁*w₂*γ₁^2*(μxv-μyv)^2
    σ²sY = σy^2 + w₁*γ₂^2*(σxv^2-σyv^2) + w₁*w₂*γ₂^2*(μxv-μyv)^2
    # transformation
    slope = sqrt(σ²sY)/sqrt(σ²sX); intercept = μsY - slope*μsX
    lYx = @. X.tableX.scale * slope + intercept
    tbl = DataFrame(scaleX = X.tableX.scale, lYx = lYx)
    return NEATEquateResult(
        :LevineObservedScore,
        tbl,
        DataFrame(Group = [1, 2], μ = [μsX, μsY], σ = [sqrt(σ²sX), sqrt(σ²sY)], γ = [γ₁, γ₂], w = [w₁, w₂]),
        (slope = slope, intercept = intercept),
        (X = X, Y = Y)
    )
end

# Nonequivalent Groups : Chained linear Observed Score Equating
"""
    ChainedLinear(X::NEAT, Y::NEAT)

Equate 2 tests in a chain through the common test scores. Now supports only 2 form equating.

# Basic Algorithm

1. put X on the scale of V -call lV(x);
2. put V on the scale of Y - call lY(v);
3. obtain Y-equivalent as lY(x).

"""
function ChainedLinear(X::NEAT, Y::NEAT)
    # test score
    x = X.rawX; xv = X.rawV
    y = Y.rawX; yv = Y.rawV
    # ObservableStats
    μx, σx, μxv, σxv, covxv, corxv = ObservableStats(X)
    μy, σy, μyv, σyv, covyv, coryv = ObservableStats(Y)
    # regression slope
    γ₁ = σx / σxv
    γ₂ = σy / σyv
    # estimate
    slope = γ₂/γ₁
    intercept = μy + γ₂ *(μxv - μyv) - slope * μx
    lYx = @. X.tableX.scale * slope + intercept
    tbl =  DataFrame(scaleX = X.tableX.scale, lYx = lYx)
    return NEATEquateResult(
        :ChainedLinear,
        tbl,
        DataFrame(Group = [1,2], γ = [γ₁, γ₂]),
        (slope = slope, intercept = intercept),
        (X = X, Y = Y)
    )
end

"""
    ChainedMean(X::NEAT, Y::NEAT)
Chained mean equating, which results are identical for Levine's observed score and true score (not implemented) when w₁ = 1.0 (there is no option for the population weights in the mean equating).
"""
function ChainedMean(X::NEAT, Y::NEAT)
    # test score
    x = X.rawX; xv = X.rawV
    y = Y.rawX; yv = Y.rawV
    # ObservableStats
    μx, σx, μxv, σxv, covxv, corxv = ObservableStats(X)
    μy, σy, μyv, σyv, covyv, coryv = ObservableStats(Y)
    γ₂ = σy / σyv
    # estimate
    slope = 1.0
    intercept = μy + γ₂ *(μxv - μyv) - slope * μx
    lYx = @. X.tableX.scale * slope + intercept
    tbl =  DataFrame(scaleX = X.tableX.scale, lYx = lYx)
    return NEATEquateResult(
        :ChainedMean,
        tbl,
        DataFrame(Group = [1,2], γ = [γ₁, γ₂]),
        (slope = slope, intercept = intercept),
        (X = X, Y = Y)
    )
end
# Nonequivalent Groups : Braun-Holland Linear Method
"""
    BraunHolland(X::NEAT, Y::NEAT; w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁)

Conduct Braun-Holland Linear equating, which use means and standard deviations estimated by the synsetic population distributions.

# Arguments

- `X` frequency table generated by `freqtab(X, V)`, `V` is the test score of common item part
- `Y` same as above.
- `w₁`, `w₂` The synsetic population weights, that are constrained to be w₁ + w₂ = 1.
"""
function BraunHolland(X::NEAT, Y::NEAT; w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁)
    # synthetic weight
    W = w₁ + w₂
    w₁ = w₁ / W; w₂ = w₂ / W
    # prior (the weights from common part)
    J = length(X.tableV.freq)
    h₁ = X.tableV.prob
    h₂ = Y.tableV.prob
    # Conditional distribution
    conditionalX = X.marginal ./ sum(X.marginal)
    conditionalY = Y.marginal ./ sum(Y.marginal)
    for v in 1:length(X.tableV.scale)
        conditionalX[:, v] ./= h₁[v]
    end
    for v in 1:length(Y.tableV.scale)
        conditionalY[:, v] ./= h₂[v]
    end
    # Marginal out
    f₂x = zeros(Float64, length(X.tableX.freq))
    g₁y = zeros(Float64, length(Y.tableX.freq))
    for j in 1:length(X.tableX.scale)
        f₂x[j] = conditionalX[j,:]'h₂
    end
    for j in 1:length(Y.tableX.scale)
        g₁y[j] = conditionalY[j,:]'h₁
    end
    # normalize
    f₂x = f₂x / sum(f₂x)
    g₁y = g₁y / sum(g₁y)
    # synthetic population
    fsx = @. w₁ * X.tableX.prob + w₂ * f₂x; fsx = fsx ./ sum(fsx)
    fsy = @. w₁ * g₁y + w₂ * Y.tableX.prob; fsy = fsy ./ sum(fsy)
    # @show f₂x, g₁y
    # synthetic pupulation parameter
    μsx = fsx' * X.tableX.scale
    σsx = sqrt(fsx' * (X.tableX.scale .- μsx) .^2)
    μsy = fsy' * Y.tableX.scale
    σsy = sqrt(fsy' * (Y.tableX.scale .- μsy) .^2)
    # external regression parameter
    slope = σsy / σsx; intercept = μsy - slope * μsx
    lYx = @. X.tableX.scale * slope + intercept
    tbl = DataFrame(scaleX = X.tableX.scale, lYx = lYx)
    return NEATEquateResult(
        :BraunHolland,
        tbl,
        DataFrame(Group = [1, 2], μ = [μsx, μsy], σ = [σsx, σsy], w = [w₁, w₂]),
        (slope = slope, intercept = intercept),
        (X = X, Y = Y)
    )
end

# Nonequivalent Goups : Frequency Estimation
"""
    FrequencyEstimation(X::NEAT, Y::NEAT; w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁)

Conduct frequency estimation equipercentile equating, which assumes, for both form, the conditional distribution of total score given each common part score is the same in both populations.

"""
function FrequencyEstimation(X::NEAT, Y::NEAT; w₁ = length(X.rawX) / (length(X.rawX) + length(Y.rawX)), w₂ = 1.0 - w₁, case = :middle)
    # synthetic weight
    W = w₁ + w₂
    w₁ = w₁ / W; w₂ = w₂ / W
    # prior (the weights from common part)
    J = length(X.tableV.freq)
    h₁ = X.tableV.prob
    h₂ = Y.tableV.prob
    # Conditional distribution
    conditionalX = X.marginal ./ sum(X.marginal)
    conditionalY = Y.marginal ./ sum(Y.marginal)
    for v in 1:length(X.tableV.scale)
        conditionalX[:, v] ./= h₁[v]
    end
    for v in 1:length(Y.tableV.scale)
        conditionalY[:, v] ./= h₂[v]
    end
    # Marginal out
    f₂x = zeros(Float64, length(X.tableX.freq))
    g₁y = zeros(Float64, length(Y.tableX.freq))
    for j in 1:length(X.tableX.scale)
        f₂x[j] = conditionalX[j,:]'h₂
    end
    for j in 1:length(Y.tableX.scale)
        g₁y[j] = conditionalY[j,:]'h₁
    end
    # normalize
    f₂x = f₂x / sum(f₂x)
    g₁y = g₁y / sum(g₁y)
    # synthetic population
    fsx = @. w₁ * X.tableX.freq + w₂ * f₂x
    fsy = @. w₁ * g₁y + w₂ * Y.tableX.freq
    # Equipercentile Equating
    ftX = FreqTab(DataFrame(scale = X.tableX.scale, freq = fsx, cumprob = cumsum(fsx) ./ sum(fsx)),
                  X.rawX, X.intervalX, (; dummy = "dummy"))
    ftY = FreqTab(DataFrame(scale = Y.tableX.scale, freq = fsy, cumprob = cumsum(fsy) ./ sum(fsy)),
                  Y.rawX, Y.intervalX, (; dummy = "dummy"))
    tbl = Equipercentile(ftX, ftY; case = case)
    return NEATEquateResult(
        :FrequencyEstimation,
        tbl.table,
        DataFrame(Group = [1, 2], μ = [mean(fsx), mean(fsy)], σ = [std(fsx), std(fsy)], w = [w₁, w₂]),
        nothing,
        (X = X, Y = Y)
    )
end

# Nonequivalent Groups : Chained Equipercentile Method
"""
    ChainedEquipercentile(X::NEAT, Y::NEAT; case = :middle)

# Algorithm
Chained equipercentile equating is, in short, to conduct equipercentile equating consecutively.
First, find equipercentile relationship for score X to scores V on the first population.
This function is referred to as 𝑒V1(𝑥).
Second, fubd the equipercentile relationship for converting scores on the common items(V) to scores on the form Y based on examinees from the population 2.
Refer to the resulting function as 𝑒Y2(𝑣).
"""
function ChainedEquipercentile(X::NEAT, Y::NEAT; case = :lower)
    # Find equipercentile relationship for score V on the first population
    ftX = freqtab(X.rawX; interval = X.intervalX, scale = X.tableX.scale)
    ftXV = freqtab(X.rawV; interval = X.intervalV, scale = X.tableV.scale)
    ftYV = freqtab(Y.rawV; interval = Y.intervalV, scale = Y.tableV.scale)
    ftY = freqtab(Y.rawX; interval = Y.intervalX, scale = Y.tableX.scale)
    eV₁x = Equipercentile(ftX, ftXV; case = case)
    # Search percentile of score V on scale Y
    eYxu = zeros(Float64, length(eV₁x.table.scaleX))
    eYxl = zeros(Float64, length(eV₁x.table.scaleX))
    for (i,v) in enumerate(eV₁x.table.eYx)
        P = PRF(v, ftYV)
        eYxu[i] = PFu(P, ftY)
        eYxl[i] = PFl(P, ftY)
    end
    if case == :upper
        eYx = eYxu
    elseif case == :lower
        eYx = eYxl
    elseif case == :both
        eYx = string.(eYxu, "_", eYxl)
    elseif case == :middle
        eYx = (eYxu .+ eYxl) ./ 2.0
    end
    tbl = DataFrame(scaleX = eV₁x.table.scaleX, eYx = eYx)
    return NEATEquateResult(
        :ChainedEquipercentile,
        tbl,
        nothing,
        nothing,
        (X = X, Y = Y)
    )
end
#-----------------


using Test
using jInvVis
using jInv.Mesh

# tests for regular mesh
domain = [0 1.1 0 1.0 0  1.1]
n      = [8 5 3]
Mr    = getRegularMesh(domain,n)
Mt    = getTensorMesh3D(rand(n[1]),rand(n[2]),rand(n[3]))
I     = randn(tuple(n...))
println("=== test viewOrthoSlices  ===")
figure(2)
viewOrthoSlices2D(I,Mr)
figure(1)
viewOrthoSlices2D(I,Mt)
println("==== all passed === ")


#--------------------------------------------------------------------------------------------------
# Tuning the constant stepsize for SGD and SMD - E-NET VERSION
#--------------------------------------------------------------------------------------------------

using Random
using PyPlot
using LaTeXStrings, Printf

include("solve_enet.jl");

#-------------------------------------
#   Parameters
#-------------------------------------
maxIter = 7000;
r = 5;  # rank
d = 100;
stoch_err = sqrt(10);  # standard deviation of errors in stochastic measurements b
init_radius = 1.0;  # 2-norm measure of relative deviation between Xtrue and Xinit

num_tests = 5;
subgrad_stepsizes = [1e-6, 2e-6, 3e-6, 4e-6, 5e-6];
mirror_stepsizes = [0.3, 0.4, 0.5, 0.6, 0.7];

Random.seed!(321);  # for reproducibility

#-------------------------------------
#   Initialize plots
#-------------------------------------
subgrad_plot = figure(figsize=[10,6]);
xlabel(L"Iteration $k$");
ylabel(L"Relative distance to solution set $\min_U \; \frac{\| X_k U - X_{true} \|_2^2 }{\| X_{true} \|_2^2}$");
title(@sprintf("SGD for e-net covariance estimation (r=%i, d=%i)", r,d));
xlim(0,maxIter)
ylim(1e-4, 1e1)

subgrad_fun_plot = figure(figsize=[10,6]);
xlabel(L"Iteration $k$");
ylabel("Function Error");
title(@sprintf("SGD for e-net covariance estimation (r=%i, d=%i)", r,d));
xlim(0,maxIter)

mirror_plot = figure(figsize=[10,6]);
xlabel(L"Iteration $k$");
ylabel(L"Relative distance to solution set $\min_U \; \frac{\| X_k U - X_{true} \|_2^2 }{\| X_{true} \|_2^2}$");
title(@sprintf("SMD for e-net covariance estimation (r=%i, d=%i)", r,d));
xlim(0,maxIter)
ylim(1e-4, 1e1)

mirror_fun_plot = figure(figsize=[10,6]);
xlabel(L"Iteration $k$");
ylabel("Function Error");
title(@sprintf("SMD for e-net covariance estimation (r=%i, d=%i)", r,d));
xlim(0,maxIter)

#-----------------------------------------------------------------------------------------
#   Run method(s) for each dimension
#-----------------------------------------------------------------------------------------
for i=1:num_tests
    ηsubgrad = subgrad_stepsizes[i];
    ηmirror = mirror_stepsizes[i];

    # Generate true matrix we are searching for
    Xtrue = randn(d,r);

    # Generate initial point
    pert = randn(d,r);
    Xinit = Xtrue + init_radius * (norm(Xtrue, 2) / norm(pert, 2)) * pert;

    # Run the two methods
    (subgrad_err_hist, subgrad_fun_hist) = solve_enet(Xinit, Xtrue, stoch_err,  maxIter,
                                                            fill(ηsubgrad, maxIter), method="subgradient",
                                                            noiseType="sparse")
    (mirror_err_hist, mirror_fun_hist) = solve_enet(Xinit, Xtrue, stoch_err,  maxIter,
                                                                fill(ηmirror, maxIter), method="mirror",
                                                                noiseType="sparse")

    # Add errors to plots
    plt[:figure](subgrad_plot[:number])
    semilogy(subgrad_err_hist, #color=step_colors[i],
                    label=@sprintf("step=%1.2e", ηsubgrad));

    plt[:figure](subgrad_fun_plot[:number])
    semilogy(subgrad_fun_hist, #color=step_colors[i],
                    label=@sprintf("step=%1.2e", ηsubgrad));

    plt[:figure](mirror_plot[:number])
    semilogy(mirror_err_hist, #color=step_colors[i],
                    label=@sprintf("step=%1.2e", ηmirror));

    plt[:figure](mirror_fun_plot[:number])
    semilogy(mirror_fun_hist, #color=step_colors[i],
                    label=@sprintf("step=%1.2e", ηmirror));
end

# Add legends to plots and save final plots
plt[:figure](subgrad_plot[:number])
legend(loc="upper right")
savefig("plots/sgd_enet_stepsize_test.pdf");

plt[:figure](subgrad_fun_plot[:number])
legend(loc="upper right")
savefig("plots/sgd_enet_stepsize_test_fun_vals.pdf");

plt[:figure](mirror_plot[:number])
legend(loc="upper right")
savefig("plots/smd_enet_stepsize_test.pdf");

plt[:figure](mirror_fun_plot[:number])
legend(loc="upper right")
savefig("plots/smd_enet_stepsize_test_fun_vals.pdf");


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.4
#     language: julia
#     name: julia-1.6
# ---

# %%
# https://twitter.com/iitenki_moruten/status/1434070751884177415

using Plots
using StatsBase
using Random
using Printf
using Distributions
using BenchmarkTools
using QuadGK
using SpecialFunctions
using Zygote

# Generating Gaussian distribution (1-dim) by using HMC method
# 改良後

v(x) = 0.5*x^2
u(x) = x^2*(x-3)^2/2

function S(x,n_size)
    sum(u(x[i]) for i=1:n_size)
end

function dSdx(x,k)
    return u'(x)
end

function H(x,p,n_size)
    K = sum(p[i]^2 for i=1:n_size)
    return S(x,n_size) + 0.5*K
end

#１変数更新のHMC法
function HMC_1dim(x,parameter)
        n_size,n_tau,τ = parameter
        #初期化
        p_ini = rand(Normal(0,1),n_size)
        p_fin = zeros(n_size)
        x_ini = zeros(n_size)
        x_fin = zeros(n_size)

        #初期配位を保存
        for i=1:n_size x_ini[i] = x[i] end

        #初期配位のハミルトニアンを計算
        H_ini = H(x_ini,p_ini,n_size)

        #リープフロッグで時間発展
        leapfrog!(x,p_ini,parameter)

        #時間発展後の配位を保存
        for i=1:n_size x_fin[i] = x[i]; p_fin[i] = p_ini[i] end

        #時間発展後のハミルトニアンを計算
        H_fin = H(x_fin,p_fin,n_size)

        #Metropolis check
        r=rand()
        if r<exp(H_ini - H_fin)
            for i=1:n_size x[i] = x_fin[i] end
        else
            for i=1:n_size x[i] = x_ini[i] end
        end
end

#x,pをリープフロッグ法で更新する．
function leapfrog!(x,p,parameter)
    n_size,n_tau,τ = parameter
    Δτ  = τ/n_tau
    p0  = 0.0
    p1  = 0.0
    x05 = 0.0
    x15 = 0.0
    for k=1:n_size
        #initial setting : τ0
        p0  = p[k]
        x05 = x[k] + p0 * 0.5*Δτ

        #(n-1)step : τ0 -> τn-1
        for i=1:n_tau-1
            p1  = p0 - dSdx(x05,k)*Δτ
            x15 = x05 + p1*Δτ
            p0  = p1
            x05 = x15
        end

        #final step : τn-1 -> τn
        p_end = p1 - dSdx(x15,k)*Δτ
        x_end = x15 + p_end*0.5*Δτ

        x[k]  = x_end
        p[k]  = p_end
    end
end

function gen_configs(Ω,n_con,parameter)
    n_size,n_tau,τ = parameter
    x = zeros(n_size)
    for j=1:n_con
        HMC_1dim(x,parameter)
        Ω[j] = x[1]
    end
end

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

"""q = [n_size,n_tau,τ]"""

#パラメータベクトル
n_config = 1_000_000          # # of configurations
n_size   = 1
n_tau    = 100
τ        = 5
q 　　　  = [n_size,n_tau,τ]
Ω        = zeros(n_config)

#result = @benchmark gen_configs(n_config,q)
@time gen_configs(Ω,n_config,q)

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

###plot###
gauss(x) = exp(-0.5*x^2)/sqrt(2*pi)
pdf_u(x) = exp(-u(x))
Z0,error = quadgk(pdf_u,-Inf,Inf)
normed_u(x) = pdf_u(x)/Z0

xr = range(-3,6,length=1000)

plot(Ω,st=:histogram,nbins=50,norm=:pdf,alpha=0.3,label="HMC's sample",legend=:topleft)
plot!(xr,normed_u.(xr),label="pdf_ana")

# %%
##########################################################

"""q = [n_size,n_tau,τ]"""

#パラメータベクトル
n_config = 1_000_000          # # of configurations
n_size   = 1
n_tau    = 100
τ        = 5
q 　　　  = [n_size,n_tau,τ]
Ω        = zeros(n_config)

#result = @benchmark gen_configs(n_config,q)
@time gen_configs(Ω,n_config,q)

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

###plot###
gauss(x) = exp(-0.5*x^2)/sqrt(2*pi)
pdf_u(x) = exp(-u(x))
Z0,error = quadgk(pdf_u,-Inf,Inf)
normed_u(x) = pdf_u(x)/Z0

xr = range(-3,6,length=1000)

plot(Ω,st=:histogram,nbins=50,norm=:pdf,alpha=0.3,label="HMC's sample",legend=:topleft)
plot!(xr,normed_u.(xr),label="pdf_ana")

# %% [markdown]
# 以下のコードでの主な変更点
#
# * `r < exp(H_ini - H_fin)` → `r ≤ min(1, exp(H_ini - H_fin))` (Metropolis-Hastings化)
# * `n_size` を函数内から削除 (必要なら `n_size = length(x)` とできる)
# * `parameter` に作業用配列 `p_ini`, `p_fin`, `x_ini`, `x_fin` を含めて、メモリアロケーションを削減。
# * `parameter` にポテンシャル函数 `u` も含め、函数 `H` などに渡すようにした。
# * Zygote.jl ではなく、ForwardDiff.jl を使うようにして、初回実行時の遅延の短縮。
#
# References
#
# * https://docs.julialang.org/en/v1/manual/performance-tips/#Pre-allocating-outputs
# * https://docs.julialang.org/en/v1/manual/performance-tips/#More-dots:-Fuse-vectorized-operations
# * https://github.com/genkuroki/public/blob/main/0018/HMC%20leapfrog.ipynb

# %%
using Plots
# using StatsBase
# using Random
# using Printf
using Distributions
using BenchmarkTools
using QuadGK
# using SpecialFunctions

module Revised

using Random
using Distributions
using ForwardDiff

# Generating Gaussian distribution (1-dim) by using HMC method
# 改良後

S(x, parameter) = ((u,) = parameter; sum(u, x))
dSdx(x, k, parameter) = ((u,) = parameter; ForwardDiff.derivative(u, x))
square(x) = x^2
H(x, p, parameter) = S(x, parameter) + 0.5sum(square, p)

#１変数更新のHMC法
function HMC_1dim!(x, parameter)
        _, n_tau, τ, p_ini, p_fin, x_ini, x_fin = parameter

        #初期化
        rand!(Normal(0,1), p_ini)

        #初期配位を保存
        x_ini .= x

        #初期配位のハミルトニアンを計算
        H_ini = H(x_ini, p_ini, parameter)

        #リープフロッグで時間発展
        leapfrog!(x, p_ini, parameter)

        #時間発展後の配位を保存
        x_fin .= x
        p_fin .= p_ini

        #時間発展後のハミルトニアンを計算
        H_fin = H(x_fin, p_fin, parameter)

        #Metropolis check
        r = rand()
        if r ≤ min(1, exp(H_ini - H_fin))
            x .= x_fin
        else
            x .= x_ini
        end
end

#x,pをリープフロッグ法で更新する．
function leapfrog!(x, p, parameter)
    n_size = length(x)
    _, n_tau, τ = parameter
    Δτ  = τ/n_tau
    for k=1:n_size
        #initial setting : τ0
        p0  = p[k]
        x05 = x[k] + p0 * 0.5*Δτ

        #(n-1)step : τ0 -> τn-1
        local p1, x15
        for i=1:n_tau-1
            p1  = p0 - dSdx(x05, k, parameter)*Δτ
            x15 = x05 + p1*Δτ
            p0  = p1
            x05 = x15
        end

        #final step : τn-1 -> τn
        p_end = p1 - dSdx(x15, k, parameter)*Δτ
        x_end = x15 + p_end*0.5*Δτ

        x[k]  = x_end
        p[k]  = p_end
    end
end

function gen_configs(Ω, x, parameter)
    n_tau, τ = parameter
    for j in eachindex(Ω)
        HMC_1dim!(x, parameter)
        Ω[j] = x[1]
    end
end

end

v(x) = 0.5x^2
u(x) = x^2*(x-3)^2/2

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

n_config = 10^6          # # of configurations
n_size   = 1
n_tau    = 100
τ        = 5

Ω        = zeros(n_config)
x        = zeros(n_size)
p_ini    = similar(x)
p_fin    = similar(x)
x_ini    = similar(x)
x_fin    = similar(x)

"""q = (u, n_tau, τ, p_ini, p_fin, x_ini, x_fin)"""
q 　　　  = (u, n_tau, τ, p_ini, p_fin, x_ini, x_fin)

#result = @benchmark gen_configs(n_config,q)
@time Revised.gen_configs(Ω, x, q)

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

###plot###
gauss(x) = exp(-0.5*x^2)/sqrt(2*pi)
pdf_u(x) = exp(-u(x))
Z0,error = quadgk(pdf_u,-Inf,Inf)
normed_u(x) = pdf_u(x)/Z0

xr = range(-3,6,length=1000)

plot(Ω,st=:histogram,nbins=50,norm=:pdf,alpha=0.3,label="HMC's sample",legend=:topleft)
plot!(xr,normed_u.(xr),label="pdf_ana")

# %%
##########################################################

n_config = 10^6          # # of configurations
n_size   = 1
n_tau    = 100
τ        = 5

Ω        = zeros(n_config)
x        = zeros(n_size)
p_ini    = similar(x)
p_fin    = similar(x)
x_ini    = similar(x)
x_fin    = similar(x)

"""q = (u, n_tau, τ, p_ini, p_fin, x_ini, x_fin)"""
q 　　　  = (u, n_tau, τ, p_ini, p_fin, x_ini, x_fin)

#result = @benchmark gen_configs(n_config,q)
@time Revised.gen_configs(Ω, x, q)

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

###plot###
gauss(x) = exp(-0.5*x^2)/sqrt(2*pi)
pdf_u(x) = exp(-u(x))
Z0,error = quadgk(pdf_u,-Inf,Inf)
normed_u(x) = pdf_u(x)/Z0

xr = range(-3,6,length=1000)

plot(Ω,st=:histogram,nbins=50,norm=:pdf,alpha=0.3,label="HMC's sample",legend=:topleft)
plot!(xr,normed_u.(xr),label="pdf_ana")

# %%


function secular(A::SymArrow{T}, λ::T) where T
    a, b, c = A.a, A.b, A.c
    ret = λ-c
    @inbounds for i = 1:length(a)
        ret += b[i]^2/(a[i]-λ)
    end
    ret
end

function secular_derivative(A::SymArrow{T}, λ::T) where T
    a, b, c = A.a, A.b, A.c
    ret = one(T)
    @inbounds for i = 1:length(a)
        ret += (b[i]/(a[i]-λ))^2
    end
    ret
end

function secular_second_derivative(A::SymArrow{T}, λ::T) where T
    a, b, c = A.a, A.b, A.c
    ret = zero(T)
    @inbounds for i = 1:length(a)
        ret += b[i]^2/(a[i]-λ)^3
    end
    2ret
end

function secular_over_derivative(A::SymArrow{T}, λ::T) where T
    a, b, c = A.a, A.b, A.c
    num = λ-c
    den = one(T)
    @inbounds for i = 1:length(a)
        bi2 = b[i]^2
        aiλ = a[i]-λ
        bi2daiλ = bi2/aiλ
        num += bi2daiλ
        den += bi2daiλ/aiλ
    end
    num/den
end

function pick_parameter(A::SymArrow{T}, j::Integer, λ::T) where T
    f = secular(A, λ)
    fp = secular_derivative(A, λ)
    fpp = secular_second_derivative(A, λ)
    c1 = inv(A.a[j-1]-λ)
    c2 = inv(A.a[j]-λ)
    # α, β, γ
    # γ = (fpp - 2c1*fp)/(2*(c2-c1)*c2^2)
    # β = (fp-c2^2*γ)/c1^2
    # α = f - c1*β - c2*γ
    c2γ = (fpp - 2c1*fp)/(2*(c2-c1)*c2)
    c1β = (fp-c2*c2γ)/c1
    return f - c1β - c2γ
end

function pick_zero_update(A::SymArrow{T}, j::Integer, α::T, λ::T) where T
    f = secular(A, λ)
    fp = secular_derivative(A, λ)
    a = α/((A.a[j-1]-λ)*(λ-A.a[j]))
    b = fp - (inv(A.a[j-1]-λ)+inv(A.a[j]-λ))*f
    c = -f
    return 2c/(b+sqrt(b^2-4a*c))
end

function pick_zero_update(A::SymArrow{T}, j::Integer, λ::T) where T
    f = secular(A, λ)
    fp = secular_derivative(A, λ)
    fpp = secular_second_derivative(A, λ)
    c1 = inv(A.a[j-1]-λ)
    c2 = inv(A.a[j]-λ)
    # α, β, γ
    # γ = (fpp - 2c1*fp)/(2*(c2-c1)*c2^2)
    # β = (fp-c2^2*γ)/c1^2
    # α = f - c1*β - c2*γ
    c2γ = (fpp - 2c1*fp)/(2*(c2-c1)*c2)
    c1β = (fp-c2*c2γ)/c1
    α = f - c1β - c2γ
    a = α/((A.a[j-1]-λ)*(λ-A.a[j]))
    b = fp - (inv(A.a[j-1]-λ)+inv(A.a[j]-λ))*f
    c = -f
    return 2c/(b+sqrt(b^2-4a*c))
end

function first_initial_guess(a1::T, nrmb2::T, c::T) where T
    ret = a1
    if c < a1
        amc2 = (a1-c)/2
        ret -= amc2+sqrt(amc2^2+nrmb2)
    else
        cma2 = (c-a1)/2
        ret -= nrmb2/(cma2+sqrt(cma2^2+nrmb2))
    end
    ret
end

function first_pick_zero_update(A::SymArrow{T}, λ::T) where T
    f = secular(A, λ)
    fp = secular_derivative(A, λ)
    n = length(A.a)
    a1 = A.a[1]
    α = one(T)
    @inbounds for i = 2:n
        α += (A.b[i]/(A.a[i]-λ))^2*(a1-A.a[i])/(λ-A.a[i])
    end
    a = -α/(λ-a1)
    b = fp + f/(λ-a1)
    c = -f
    return 2c/(b+sqrt(b^2-4a*c))
end

function last_initial_guess(an::T, nrmb2::T, c::T) where T
    ret = an
    if c < an
        amc2 = (an-c)/2
        ret += nrmb2/(amc2+sqrt(amc2^2+nrmb2))
    else
        cma2 = (c-an)/2
        ret += cma2+sqrt(cma2^2+nrmb2)
    end
    ret
end

function last_pick_zero_update(A::SymArrow{T}, λ::T) where T
    f = secular(A, λ)
    fp = secular_derivative(A, λ)
    n = length(A.a)
    an = A.a[n]
    α = one(T)
    @inbounds for i = 1:n-1
        α += (A.b[i]/(A.a[i]-λ))^2*(an-A.a[i])/(λ-A.a[i])
    end
    a = -α/(λ-an)
    b = fp + f/(λ-an)
    c = -f
    return 2c/(b+sqrt(b^2-4a*c))
end

# Fast Multipole Method used to accelerate the eigensolve

@inline cauchykernel{T}(x::T, y::T) = inv(x-y)
@inline coulombkernel{T}(x::T, y::T) = inv((x-y)^2)
@inline coulombprimekernel{T}(x::T, y::T) = inv((x-y)^3)
@inline logkernel{T}(x::T, y::T) = log(abs(x-y))

function secular(A::SymArrow{T}, λ::Vector{T}) where T
    a, b, c = A.a, A.b, A.c
    ret = λ-c
    for i = 1:length(λ)
        @inbounds for j = 1:length(a)
            ret[i] += b[j]^2/(a[j]-λ[i])
        end
    end
    ret
end

function secular_FMM(A::SymArrow{T}, λ::Vector{T}) where T
    a, b, c = A.a, A.b, A.c
    amin, amax = extrema(a)
    λmin, λmax = extrema(λ)
    K = KernelMatrix(cauchykernel, reverse(λ), reverse(a), λmax, λmin, amax, amin)
    #ret = λ - c + K*(b.^2)
    ret = λ - c - reverse!(K*(b.^2))
    ret
end

function secular_derivative(A::SymArrow{T}, λ::Vector{T}) where T
    a, b, c = A.a, A.b, A.c
    ret = ones(T, length(λ))
    for i = 1:length(λ)
        @inbounds for j = 1:length(a)
            ret[i] += (b[j]/(a[j]-λ[i]))^2
        end
    end
    ret
end

function secular_derivative_FMM(A::SymArrow{T}, λ::Vector{T}) where T
    a, b, c = A.a, A.b, A.c
    amin, amax = extrema(a)
    λmin, λmax = extrema(λ)
    K = KernelMatrix(coulombkernel, reverse(λ), reverse(a), λmax, λmin, amax, amin)
    ret = ones(T, length(λ)) + K*(b.^2)
    reverse!(ret)
end

function secular_second_derivative(A::SymArrow{T}, λ::Vector{T}) where T
    a, b, c = A.a, A.b, A.c
    ret = zeros(T, length(λ))
    for i = 1:length(λ)
        @inbounds for j = 1:length(a)
            ret[i] += 2b[j]^2/(a[j]-λ[i])^3
        end
    end
    ret
end

function secular_second_derivative_FMM(A::SymArrow{T}, λ::Vector{T}) where T
    a, b, c = A.a, A.b, A.c
    amin, amax = extrema(a)
    λmin, λmax = extrema(λ)
    K = KernelMatrix(coulombprimekernel, reverse(λ), reverse(a), λmax, λmin, amax, amin)
    ret = K*(-2b.^2)
    reverse!(ret)
end

# For all 2 ≤ j ≤ n-1
function pick_zero_update(A::SymArrow{T}, λ::Vector{T}) where T
    n = size(A, 1)
    f = secular(A, λ)
    fp = secular_derivative(A, λ)
    fpp = secular_second_derivative(A, λ)
    c1 = T[inv(A.a[j-1]-λ[j-1]) for j in 2:n-1]
    c2 = T[inv(A.a[j]-λ[j-1]) for j in 2:n-1]
    c2γ = (fpp - 2c1.*fp)./(2*(c2-c1).*c2)
    c1β = (fp-c2.*c2γ)./c1
    α = f - c1β - c2γ
    a = T[α[j-1]/((A.a[j-1]-λ[j-1])*(λ[j-1]-A.a[j])) for j in 2:n-1]
    b = T[fp[j-1] - (inv(A.a[j-1]-λ[j-1])+inv(A.a[j]-λ[j-1]))*f[j-1] for j in 2:n-1]
    c = -f
    return 2c./(b+sqrt.(b.^2-4a.*c))
end

function pick_zero_update_FMM(A::SymArrow{T}, λ::Vector{T}) where T
    n = size(A, 1)
    f = secular_FMM(A, λ)
    fp = secular_derivative_FMM(A, λ)
    fpp = secular_second_derivative_FMM(A, λ)
    c1 = T[inv(A.a[j-1]-λ[j-1]) for j in 2:n-1]
    c2 = T[inv(A.a[j]-λ[j-1]) for j in 2:n-1]
    c2γ = (fpp - 2c1.*fp)./(2*(c2-c1).*c2)
    c1β = (fp-c2.*c2γ)./c1
    α = f - c1β - c2γ
    a = T[α[j-1]/((A.a[j-1]-λ[j-1])*(λ[j-1]-A.a[j])) for j in 2:n-1]
    b = T[fp[j-1] - (inv(A.a[j-1]-λ[j-1])+inv(A.a[j]-λ[j-1]))*f[j-1] for j in 2:n-1]
    c = -f
    return 2c./(b+sqrt.(b.^2-4a.*c))
end


# some model
mgrid = [range(0.0, stop=10.,step=0.05), range(0.0, stop=10.,step=0.05)];

model = Medium(mgrid);
vpb=[2100.,2200.];rhob=[2100., 2300.]
update!(model, [:vp,:rho], [vpb, rhob])
fill!(model)

nznx=prod(length.(model.mgrid))

for parameterization in [[:χKI, :χrhoI, :null],[:χKI, :null, :null],[:χvp, :χrho, :null],[:χvp, :null, :null]]
	δx1=randn(count(parameterization .≠ :null)*nznx)
	δxout1=zero(δx1)

	copyto!(δxout1, δx1, model, parameterization)

	δxout2=randn(size(δxout1));
	δx2=zero(δx1)
	GeoPhyInv.pert_chainrule!(δx2, δxout2, model, parameterization)

	@test LinearAlgebra.dot(δxout2,δxout1) ≈ LinearAlgebra.dot(δx1,δx2)
end


"""
Given parEst, an estimate the ergm parameters, sample from that ergm nSample times and re-estimate the same ergm on the obtained sample of observations.

Return a distribution of estimated parameter, aka parametric bootstrap following "Exponential Random Graph Models with Big Networks: Maximum Pseudolikelihood Estimation and the Parametric Bootstrap" of Christian S. Schmid, Bruce A. Desmarais.
"""
function get_par_boot_ergm_distrib(model, parEst, N; nSample = 100)

    mat_sample = StaticNets.sample_ergm(model, N, parEst, nSample)

    estimates = zeros(length(parEst), nSample)

   for i in 1:nSample
        estimates[:,i] = StaticNets.estimate(model,  mat_sample[:,:,i])
    end
    return estimates
end
export get_par_boot_ergm_distrib

"""
Get the confidence intervals from parametric bootstrap, as described in "Exponential Random Graph Models with Big Networks: Maximum Pseudolikelihood Estimation and the Parametric Bootstrap" of Christian S. Schmid, Bruce A. Desmarais.
"""
function get_conf_int_ergm_par_boot(model, N, parEst, quantilesVals; nSample = 100)

    par_b_dist = get_par_boot_ergm_distrib(model, parEst, N; nSample = nSample)

    confInt = reduce(vcat, [quantile(par_b_dist[p,:], quantilesVals)' for p in 1:model.nErgmPar])
end
export get_conf_int_ergm_par_boot

is_between(x, interval) = interval[1] <= x <= interval[2]

function get_coverage_conf_int_par_boot_ergm(model, N, parDgp, quantilesVals; nSampleDgp=100, nSampleBands = 100)

    # sample from dgp
    matsSample = StaticNets.sample_ergm(model, N, parDgp, nSampleDgp)

    estimates = reduce(hcat,[StaticNets.estimate(model, A) for A in eachslice(matsSample, dims=3)])

    coverFlags =falses(size(estimates))

    for n in 1:nSampleDgp
        confInt = get_conf_int_ergm_par_boot(model, N, estimates[:, n], quantilesVals; nSample=nSampleBands)

        for p in 1:model.nErgmPar
            coverFlags[p,n] = is_between(parDgp[p], confInt[p,:])
        end
    end

    return coverFlags
end
export get_coverage_conf_int_par_boot_ergm


function fastparse(str::ASCIIString)
    dy = 0
    for i = 1:4
        c = str[i]
        dy = 10*dy + (c - '0')
    end

    dm = 0
    for i = 6:7
        c = str[i]
        dm = 10*dm + (c - '0')
    end

    dd = 0
    for i = 9:10
        c = str[i]
        dd = 10*dd + (c - '0')
    end

    th = 0
    for i = 12:13
        c = str[i]
        th = 10*th + (c - '0')
    end

    tm = 0
    for i = 15:16
        c = str[i]
        tm = 10*tm + (c - '0')
    end

    ts = 0
    for i = 18:19
        c = str[i]
        ts = 10*ts + (c - '0')
    end

    tms = 0
    for i = 21:23
        if i <= length(str)
            c = str[i]
        else
            c = '0'
        end
        tms = 10*tms + (c - '0')
    end
    DateTime(dy,dm,dd,th,tm,ts,tms)
end

function fastparse{S<:AbstractString}(ads::Vector{S})
    l = length(ads)
    a = Array(DateTime, l)
    for i in 1:l
        a[i] = fastparse(ads[i])
    end
    a
end


module Probs
include("./samplesGenerator.jl")
export generateChanges
# export ifDetected, notDetected, extractFps, extractTps, belowH, aboveH

function ifDetected(cde::Int64, changes::Vector{Int64}, delay::Int64)
    # 1) If change detection event is TP
    for chp in changes
        if abs(chp - cde) <= delay
            return true
        end
    end
    return false
end

function notDetected(cde::Int64, changes::Vector{Int64}, delay::Int64)
    # 2) If change detection event is not TP
    return ! ifDetected(cde, changes, delay)
end

function extractFps(detections::Vector{Int64},
                    trueChanges::Vector{Int64},
                    delay::Int64)
    # 3) Extract False Positives from a sequence of CDEs
    filter(x->notDetected(x, trueChanges, delay), detections)
end

function extractTps(detections::Vector{Int64},
                    trueChanges::Vector{Int64},
                    delay::Int64)
    # 4) Extract True Positives from a sequence of CDEs
    filter(x->ifDetected(x, trueChanges, delay), detections)
end

function ifPredictedTp(cde::Int64, probs::Vector{Float64}, h::Float64)
    # 5) Check if an individual TP is correctly predicted
    probs[cde] >= h
end

function ifPredictedFp(cde::Int64, probs::Vector{Float64}, h::Float64)
    # 6) Check if an individual FP is correctly ignored
    probs[cde] < h
end

function correctlyPredictedTps(probs::Vector{Float64},
                               trueChanges::Vector{Int64},
                               h::Float64)
    # 7) Given probs vector count those provided true changes for
    # which corresponding probs >= h. Time locations must be integers
    filter(x->ifPredictedTp(x, probs, h), trueChanges)
end

function correctlyPredictedFps(probs::Vector{Float64},
                               fPositives::Vector{Int64},
                               h::Float64)
    # 8) Given probs vector count those provided false alarms for
    # which corresponding probs < h. Time locations must be integers
    filter(x->ifPredictedFp(x, probs, h), fPositives)
end

function belowH(h::Float64, probs::Vector{Float64}, delay::Int64)
    # 9) Returns _indices_ given a vector of probs and threshold
    # function. Returs indices of those elements in probs vector which
    # are below threshold value
    ret :: Vector{Int64}
    subprobs = probs[1 : min(delay, length(probs))]
    ret = find( x -> x <= h, subprobs)
end

function aboveH(h, probs, delayConfirm)
    # 10)
    subprobs = probs[1 : min(delayConfirm, length(probs))]
    find(x->(x>h), subprobs)
end

#function initProbs(mu, sigma)
#    # 11)
#    0
#end
#
#function whichBin(num, edges)
#    # 12) Edges:[1,2,3] Bins:[(-inf,1) (1,2) (2,3) (3, inf)] Bins nums: 1,2,3,4
#    n = length(edges)
#    if num <= edges[1]
#        return 1
#    elseif num > edges[end]
#        return n+1
#    else
#        for i in 2:n
#            if edges[i-1] < num <= edges[i]
#                return i
#            end
#        end
#    end
#end

function gaussPdf(x, mu, sigma)
    k = 1/(sigma * sqrt(2*pi))
    p = k * exp(-((x-mu)^2/(2*sigma^2)))
end

# 13, 14 - don't need
function generateChanges(n, mu, sigma, width)
    # 15) Generate sequence of changes
    changes::Array{Int,1}
    values = 1:width
    weights = map(x->gaussPdf(x, mu, sigma), 1:width)
    samples = genSamples(values, weights, n)
    changes = cumsum(samples)
end

function dynamicH(probs,threshold,hmin,hmax,tmax)
    dynamich::Vector{Float64}
    dynamich=map(probs) do p
        if p>=threshold
            hmin
        else
            hmax
        end
    end
    dynamich[1:tmax]
end

end
#------------------ End module Probs-------------------------#

#------------------------------------------------------------#
# module Pmf
# export gaussPmf
#
# using Distributions
# using Bisect
#
# type GaussPmf
#     # Prototype is in thinkbayes.py/def MakeGaussianPmf(mu, sigma, num_sigmas, n=201):
#     mu         :: Float64
#     sigma      :: Float64
#     num_sigmas :: Int
#     n          :: Int
#     edges      :: Vector{Float64}
#     probs      :: Vector{Float64}
#     function GaussPmf(mu::Float64, sigma::Float64)
#         n_default          = 251
#         num_sigmas_default = 3
#         low                = mu - num_sigmas_default * sigma
#         high               = mu + num_sigmas_default * sigma
#         d                  = Normal(mu, sigma)
#         edges              = linspace(low, high, n_default)
#         probs              = pdf(d, edges)
#         new(mu, sigma, num_sigmas_default, n_default, edges, probs)
#     end
# end
#
# function gaussPmf(gp::GaussPmf, val)
#      Bisect.idx_of_closest(gp.probs, float(val))
# end
#
#end
#----------------- End module Pmf ------------------------------#


include("../src/hypergraph-affinity-functions.jl")
names = ["Clothing, Shoes & Accessories"
"Electronics and Office"
"Home, Furniture & Appliances"
"Home Improvement & Patio"
"Baby"
"Toys, Games, and Video Games"
"Food, Household & Pets"
"Pharmacy, Health & Beauty"
"Sports, Fitness & Outdoors"
"Auto, Tires & Industrial"
"Other"]
dataset = "walmart-trips"
r = 25
H, labels = read_hypergraph_data(dataset,r)

## Categories 1 and 7 make up over 50% of the dataset.
# At a high level, this suggests that a lot of shopping can be broadly
# categorized into grocery shopping or clothes shopping. We take the subset
# of the data corresponding to these two (almost balanced) categories, to obtain
# a two-class hypergraph
CF = findall(x->in(x,[1,7]),labels)
labels = labels[CF]
classes = round.(Int64,(labels .- 1) ./6)
n = length(CF)

H2 = H[:,CF]
order = vec(sum(H2,dims = 2))
nonempty = findall(x->x>0,order)
H = H2[nonempty,:]
# class 1 is food

## Bootstraps
alpha, B1, B2, H1, H2, tR1, tR2, N = Hypergraph_to_Scores(H,classes,25)

for k = 4:4:12
B = 100
TR1 = tR1[k,1:k]  # True affinity scores
TR0 = tR2[k,1:k]

TA1 = H1[k,1:k]  # True affinity scores
TA0 = H2[k,1:k]

# Bootstraps
R0, R1, MR1, MR0, SR1, SR0, A0, A1, MA1, MA0, SA1, SA0 = Bootstrap_Affinities(H,k,B,classes)
## Plot
sc = 2
lw = 2.5
ms = 9
gfs = 15
tfs = 15
titlesize = 14
s1 = 400
s2 = 300
color1 = :purple
color2 = RGB(218/255,165/255,35/255)
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity / Baseline"
p = Plots.plot([1; k],[1; 1],fontfamily = "helvetica",linewidth = 2,yscale = :log10,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = k/4:k/4:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)

plot!(p,1:k,[TR0 TR0], fillrange=[MR0-sc*SR0 MR0+sc*SR0], fillalpha=0.3,color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)
plot!(p,1:k,[TR1 TR1], fillrange=[MR1-sc*SR1 MR1+sc*SR1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)

savefig("Figures/Ratios_$(k)_Walmart_Food_Clothing.pdf")

## Next
s1 = 400
s2 = 300
title = "$k-size hyperedge"
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity"
p = Plots.plot(yscale = :log10,fontfamily = "helvetica",linewidth = 2,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = k/4:k/4:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)
plot!(p,1:k,[TA1 TA1], fillrange=[MA1-sc*SA1 MA1+sc*SA1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)
plot!(p,1:k,[TA0 TA0], fillrange=[MA0-sc*SA0 MA0+sc*SA0], fillalpha=0.3,color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)

savefig("Figures/Affinities_$(k)_Walmart_Food_Clothing_log.pdf")
end
## Group Homophily Index, print usable LaTeX code for building table
# include("../src/hypergraph-affinity-functions.jl")
# S1,S2  = GHI(H,classes,14)


function sum_rows(data, n, m)
    s = zero(eltype(data))
    for i in 1:n
        for j in 1:m
            s += data[i,j]
        end
    end
    return s
end

function sum_columns(data, n, m)
    s = zero(eltype(data))
    for j in 1:m
        for i in 1:n
            s += data[i,j]
        end
    end
    return s
end

using Random

Random.seed!(0)
n, m = 500, 10^6
data = rand(n,m)

#dry run for precompilation
sum_columns(data, n, m)
sum_rows(data, n, m)

println("Summing with the column iteration moving fastest")
@time sum_columns(data, n, m)
println("Summing with the row iteration moving fastest")
@time sum_rows(data, n, m);


@testset "Geometry" begin
    @testset "Boosts" begin
        β0 = Vec3(0, 0, 0)
        u1 = Vec4(1.2, 0.2, 0.2, 1.)
        @test boost(u1, β0) ≈ u1
        β1 = Vec3(0, 0, 0.85)
        v1 = boost(u1, β1)
        @test u1⋅u1 ≈ v1⋅v1
        β2 = u1 / u1.t
        v2 = boost(u1, β2)
        @test v2.t^2 ≈ u1⋅u1
        ε = 1e-16
        @test norm(Vec3(v2)) < ε
        β3 = Vec3(0.55, -0.37, 0.11)
        v3 = boost(u1, β3)
        u3 = boost(v3, -β3)
        @test u1 ≈ u3
    end

    @testset "Rotations" begin
        x̂ = Vec3(1, 0, 0)
        ŷ = Vec3(0, 1, 0)
        ẑ = Vec3(0, 0, 1)
        @test rotate(x̂, ẑ, π/2) ≈ ŷ
        @test rotate(ŷ, x̂, π/2) ≈ ẑ
        @test rotate(ẑ, ŷ, π/2) ≈ x̂
        @test rotate(x̂, ẑ, π) ≈ -x̂
        @test rotate(x̂, ŷ, π) ≈ -x̂
        @test rotate(ŷ, ẑ, π) ≈ -ŷ
        @test rotate(ŷ, x̂, π) ≈ -ŷ
        @test rotate(ẑ, x̂, π) ≈ -ẑ
        @test rotate(ẑ, ŷ, π) ≈ -ẑ
        u = Vec3(1.1, -0.3, 0.9)
        k̂ = normalize(Vec3(-0.6, 0.8, 1.3))
        @test rotate(u, k̂, -π/3) ≈ rotate(u, -k̂, π/3)
        @test norm(rotate(u, k̂, 0.6)) ≈ norm(u)
        @test rotate(u, normalize(u), 1.3) ≈ u
        v̂ = normalize(k̂ × u)
        @test rotate(u, v̂, π) ≈ -u
    end
end


# Solves the Lorenz problem over different periods, and writes time average
module LorenzMod

using DelimitedFiles, DifferentialEquations, GeometricIntegratorsDiffEq

include("gregory.jl")
using .Gregory

function calcZAverage(sol, ibeg::Int, iend::Int,
    num_steps::Int, order::Int)::Float64
    @assert( iend - ibeg == num_steps)
    #data = zeros(num_steps+1)
    #for i = 1:num_steps+1
    #    data[i] = u[ibeg+i-1][3]
    #end
    numbnd = order-1
    if 2*numbnd > num_steps
        numbnd = div(num_steps, 2)
        order = numbnd + 1
        println("WARNING: had to adopt order ", order, " quadrature.")
    end
    val = integrate(sol[3,ibeg:iend], numbnd, order, h = 1/num_steps)
    return val
end

"""
Print3DArray(file, data)

Print a 3D array to file using column-major ordering.
"""
function Print3DArray(file, arr::AbstractArray{Float64,3})
    for k = 1:size(arr,3)
        for j = 1:size(arr,2)
            for i = 1:size(arr,1)
                print(file, arr[i,j,k], " ")
            end
            println(file)
        end
    end
end

# define the Lorenz problem
function lorenz!(du,u,p,t)
    du[1] = 10.0*(u[2]-u[1])
    du[2] = u[1]*(28.0-u[3]) - u[2]
    du[3] = u[1]*u[2] - (8/3)*u[3]
end

"""
gatherStatsLorenz()

Runs the Lorenz problem with different discretization orders, steps sizes and
integration periods, and writes results to file.
"""
function gatherStatsLorenz()

    u0 = [-5.276131687990789; -8.274832825254473; 16.36301350713599]
    time = 1000.0
    probLorenz = ODEProblem(lorenz!, u0, (0, time))

    time_fac = 10
    #num_steps = [6250; 12500; 25000; 50000; 100000]
    base_time = time/time_fac^3
    base_steps = [10; 20; 40; 80; 160]
    order = [2, 4, 8]
    #file_prefix = ["Heun_u_", "RK4_u_", "GLRK(4)_u_"]
    method = [Heun(), RK4(), Vern8() ]

    # loop over the different order methods
    for p = 1:size(order,1)
        println("Gathering statistics for order ",order[p])
        # each method gets its own file

        dt = zeros(Float64, (time_fac^3, size(base_steps,1), 3) )
        z_avg = zeros(size(dt))
        cputime = zeros(size(dt))

        println("Initial compile...")
        cpu_full = @elapsed sol = solve(probLorenz, method[p],
                                        dt=base_time/base_steps[1],
                                        adaptive=false)

        for n = 1:size(base_steps,1)
            dt[:,n,:] .= base_time/base_steps[n]
            println("\tSolve using num steps = ",time_fac^3*base_steps[n])
            cpu_full = @elapsed sol = solve(probLorenz, method[p],
                                            dt=base_time/base_steps[n],
                                            adaptive=false)
            #if any(isnan, sol)
            if sol.retcode != :Success
                println("\t\tLorenz failed to solve due to instability.")
                continue
            end

            # We partition the simulation into smaller periods for statistics
            # gathering.  For example, if the simulation time is 1000 units
            # then we have 10 samples for 100, 100 samples for 10 units, and
            # 1000 samples for 1 unit.
            # The statistics for the medium and long periods are duplicated to
            # simplify the file I/O (NumPy loadtxt is easier)
            ptr = 1
            for s1 = 1:time_fac
                ibeg = base_steps[n]*(s1-1)*time_fac*time_fac + 1
                iend = base_steps[n]*s1*time_fac*time_fac + 1
                z_avg_long = calcZAverage(sol, ibeg, iend,
                                          time_fac^2*base_steps[n],
                order[p])
                for s2 = 1:time_fac
                    ibeg = base_steps[n]*((s1-1)*time_fac + s2 - 1)*time_fac + 1
                    iend = base_steps[n]*((s1-1)*time_fac + s2)*time_fac + 1
                    z_avg_med = calcZAverage(sol, ibeg, iend,
                    time_fac*base_steps[n], order[p])
                    for s3 = 1:time_fac
                        # get statistics for the shortest period
                        ibeg = base_steps[n]*(
                        ( (s1-1)*time_fac + s2-1)*time_fac + s3 - 1) + 1
                        iend = base_steps[n]*(
                        ( (s1-1)*time_fac + s2-1)*time_fac + s3) + 1
                        z_avg[ptr,n,1] = calcZAverage(sol, ibeg, iend,
                        base_steps[n], order[p])
                        cputime[ptr,n,1] = cpu_full/(time_fac^3)

                        # get statistics for the medium period
                        ibeg = base_steps[n]*((s1-1)*time_fac+s2-1)*time_fac + 1
                        iend = base_steps[n]*((s1-1)*time_fac + s2)*time_fac + 1
                        z_avg[ptr,n,2] = z_avg_med
                        cputime[ptr,n,2] = cpu_full/(time_fac^2)

                        # get statistics for the long period
                        z_avg[ptr,n,3] = z_avg_long
                        cputime[ptr,n,3] = cpu_full/time_fac

                        ptr += 1
                    end
                end
            end
        end

        # write the data
        data = open("raw_statistics_lorenz_order$(order[p]).dat", "w")
        Print3DArray(data, dt)
        Print3DArray(data, z_avg)
        Print3DArray(data, cputime)
        close(data)
    end
end

end # LorenzMod


using StatisticalRethinking, CmdStan, MCMCChain
gr(size=(500,500));

ProjDir = rel_path("..", "scripts", "04")
cd(ProjDir)

howell1 = CSV.read(rel_path("..", "data", "Howell1.csv"), delim=';')
df = convert(DataFrame, howell1);

df2 = filter(row -> row[:age] >= 18, df)
first(df2, 5)

heightsmodel = "
// Inferring a Rate
data {
  int N;
  real<lower=0> h[N];
}
parameters {
  real<lower=0> sigma;
  real<lower=0,upper=250> mu;
}
model {
  // Priors for mu and sigma
  mu ~ normal(178, 20);
  sigma ~ uniform( 0 , 50 );

  // Observed heights
  h ~ normal(mu, sigma);
}
";

stanmodel = Stanmodel(name="heights", model=heightsmodel);

heightsdata = Dict("N" => length(df2[:height]), "h" => df2[:height]);

rc, a3d, cnames = stan(stanmodel, heightsdata, ProjDir, diagnostics=false,
  CmdStanDir=CMDSTAN_HOME);

pi = filter(p -> length(p) > 2 && p[end-1:end] == "__", cnames)
p = filter(p -> !(p in  pi), cnames)

chn = Chains(a3d,
  Symbol.(cnames),
  Dict(
    :parameters => Symbol.(p),
    :internals => Symbol.(pi)
  )
)

describe(chn)
describe(chn, section=:internals)


function find_joltage_differences(adapters)
    pushfirst!(adapters, 0)
    device = maximum(adapters) + 3
    push!(adapters, device)
    sort!(adapters)
    one = 0
    three = 0
    for (x, y) in zip(adapters, adapters[2:end])
        diff = y - x
        if diff == 1
            one += 1
        elseif diff == 3
            three += 1
        end
    end
    one * three
end

function main()
    adapters = map(x -> parse(Int64, x), readlines("../resources/input_10.txt"))
    result = find_joltage_differences(adapters)
    println(result)
end

main()


@testset "Test composition order" begin
  Qx, x = FlintQQ["x"]
  K, a = NumberField(x^6 + 108, "a")
  A = automorphisms(K)
  for f in A
    for g in A
      @test g(f(a)) == compose(f, g)(a)
      @test g(f(a)) == (f * g)(a)
    end
  end
end

@testset "Induce image" begin
  Qx, x = FlintQQ["x"]
  K, a = NumberField(x^6 + 108, "a")
  A = automorphisms(K)
  OK = maximal_order(K)
  lP = prime_decomposition(OK, 7)
  lP1 = [x[1] for x in lP]
  for i = 1:length(lP)
    I_new = A[2](lP1[1])
    id = findfirst(isequal(I_new), lP1)
    @test id != nothing
  end
  f = hom(K, K, a^4//12+a//2)
  E = EquationOrder(K)
  I = ideal(E, E(a))
  @test_throws ErrorException Hecke.induce_image(f, I)

end


using PyPlot
using QuantEcon: meshgrid

fig = figure()
ax = fig[:gca](projection="3d")

x_min, x_max = -5, 5
y_min, y_max = -5, 5

alpha, beta = 0.2, 0.1

ax[:set_xlim]((x_min, x_max))
ax[:set_ylim]((x_min, x_max))
ax[:set_zlim]((x_min, x_max))

# Axes
ax[:set_xticks]((0,))
ax[:set_yticks]((0,))
ax[:set_zticks]((0,))
gs = 3
z = linspace(x_min, x_max, gs)
x = zeros(gs)
y = zeros(gs)
ax[:plot](x, y, z, "k-", lw=2, alpha=0.5)
ax[:plot](z, x, y, "k-", lw=2, alpha=0.5)
ax[:plot](y, z, x, "k-", lw=2, alpha=0.5)

# Fixed linear function, to generate a plane
f(x, y) = alpha .* x + beta .* y

# Vector locations, by coordinate
x_coords = [3, 3]
y_coords = [4, -4]
z = f(x_coords, y_coords)
for i=1:2
    ax[:text](x_coords[i], y_coords[i], z[i], LaTeXString("\$a_{$i}\$"), fontsize=14)
end

# Lines to vectors
for i=1:2
    x = (0, x_coords[i])
    y = (0, y_coords[i])
    z = (0, f(x_coords[i], y_coords[i]))
    ax[:plot](x, y, z, "b-", lw=1.5, alpha=0.6)
end

# Draw the plane
grid_size = 20
xr2 = linspace(x_min, x_max, grid_size)
yr2 = linspace(y_min, y_max, grid_size)
x2, y2 = meshgrid(xr2, yr2)
z2 = f(x2, y2)
ax[:plot_surface](x2, y2, z2, rstride=1, cstride=1, cmap=ColorMap("jet"),
        linewidth=0, antialiased=true, alpha=0.2)
plt.show()


# Load Julia packages (libraries) needed  for the snippets in chapter 0

using DynamicHMCModels

# CmdStan uses a tmp directory to store the output of cmdstan

ProjDir = @__DIR__
cd(ProjDir)

# Read the milk data

df = CSV.read(joinpath("..", "..", "data", "milk.csv"), DataFrame)
df = filter(row -> !(row[:neocortex_perc] == "NA"), df)
#df[:, :kcal_per_g] = convert(Vector{Float64}, df[:, :kcal_per_g])
df[!, :log_mass] = log.(convert(Vector{Float64}, df[:, :mass]))

# Define the model struct

Base.@kwdef mutable struct MilkModel{Ty <: AbstractVector, Tx <: AbstractMatrix}
    "Observations."
    y::Ty
    "Covariates"
    x::Tx
end

# Write a function to return properly dimensioned transformation.

function make_transformation(model::MilkModel)
  as((β = as(Array, size(model.x, 2)), σ = asℝ₊))
end

# Instantiate the model with data and inits.

x = hcat(ones(size(df, 1)), df[:, :log_mass]);
model = MilkModel(;y=df[:, :kcal_per_g], x=x)

# Make the type callable with the parameters *as a single argument*.

function (model::MilkModel)(θ)
    @unpack y, x, = model   # extract the data
    @unpack β, σ = θ            # works on the named tuple too
    ll = 0.0
    ll += logpdf(Normal(0, 100), x[1])
    ll += logpdf(Normal(0, 1), x[2])
    ll += logpdf(TDist(1.0), σ)
    ll += loglikelihood(Normal(0, σ), y .- x*β)
    ll
end

println()
model((β = [1.0, 2.0], σ = 1.0))
println()

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

p = as_particles(posterior)
display(p)

stan_result = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
          Mean         SD        Naive SE       MCSE      ESS
    a  0.70472876 0.057040655 0.00090189195 0.0011398893 1000
   bm -0.03150330 0.023642759 0.00037382484 0.0004712342 1000
sigma  0.18378372 0.039212805 0.00062000888 0.0011395979 1000

Quantiles:
          2.5%       25.0%       50.0%        75.0%       97.5%
    a  0.59112968  0.66848775  0.70444950  0.741410500 0.81915225
   bm -0.07729257 -0.04708425 -0.03104865 -0.015942925 0.01424901
sigma  0.12638780  0.15605950  0.17800600  0.204319250 0.27993590
";

# End of `05/5.6d.jl`


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ bbfa7498-30bf-4264-b4d9-ee0c238c9984
using Plots, Combinatorics, PlutoTest

# ╔═╡ f0d31c21-f866-4d51-86ec-d473147f1677
using LightGraphs, SimpleWeightedGraphs, GraphPlot

# ╔═╡ f20ff74b-16f0-4a9d-9969-2d240ba18177
plotly()

# ╔═╡ 56cedfe0-b02c-11eb-1f56-3557d6c170bc
md"""
# [Day 9: All in a Single Night](https://adventofcode.com/2015/day/9)

Every year, Santa manages to deliver all of his presents in a single night.

This year, however, he has some new locations to visit; his elves have provided him the distances between every pair of locations. He can start and end at any two (different) locations he wants, but he must visit each location exactly once. What is the shortest distance he can travel to achieve this?

For example, given the following distances:

```
London to Dublin = 464
London to Belfast = 518
Dublin to Belfast = 141

```

The possible routes are therefore:

```
Dublin -> London -> Belfast = 982
London -> Dublin -> Belfast = 605
London -> Belfast -> Dublin = 659
Dublin -> Belfast -> London = 659
Belfast -> Dublin -> London = 605
Belfast -> London -> Dublin = 982

```

The shortest of these is `London -> Dublin -> Belfast = 605`, and so the answer is `605` in this example.

**What is the distance of the shortest route?**
"""

# ╔═╡ 1609ee1f-1d90-4889-b39d-5fe0fb6de8c4
puzzle_input = "AlphaCentauri to Snowdin = 66
AlphaCentauri to Tambi = 28
AlphaCentauri to Faerun = 60
AlphaCentauri to Norrath = 34
AlphaCentauri to Straylight = 34
AlphaCentauri to Tristram = 3
AlphaCentauri to Arbre = 108
Snowdin to Tambi = 22
Snowdin to Faerun = 12
Snowdin to Norrath = 91
Snowdin to Straylight = 121
Snowdin to Tristram = 111
Snowdin to Arbre = 71
Tambi to Faerun = 39
Tambi to Norrath = 113
Tambi to Straylight = 130
Tambi to Tristram = 35
Tambi to Arbre = 40
Faerun to Norrath = 63
Faerun to Straylight = 21
Faerun to Tristram = 57
Faerun to Arbre = 83
Norrath to Straylight = 9
Norrath to Tristram = 50
Norrath to Arbre = 60
Straylight to Tristram = 27
Straylight to Arbre = 81
Tristram to Arbre = 90"

# ╔═╡ 3ffc6481-2267-4d97-8065-d2fd78c637e9
function build_graph(paths)
	edges = []
	for line in split(paths, "\n")
		m = match(r"(\w+) to (\w+) = (\d+)", line)
		push!(edges, (m[1], m[2], parse(Int, m[3])))
	end

	nodes = Dict()
	next_id = 1
	for (from, to, weight) in edges
		if !haskey(nodes, from)
			nodes[from] = next_id
			next_id += 1
		end
		if !haskey(nodes, to)
			nodes[to] = next_id
			next_id += 1
		end
	end

	g = SimpleWeightedGraph(length(nodes));

	node_labels = []
	edge_labels = []

	for (from, to, weight) in edges
		push!(edge_labels, string(weight))

		if from ∉ node_labels
			push!(node_labels, from)
		end
		if to ∉ node_labels
			push!(node_labels, to)
		end

		add_edge!(g, nodes[from], nodes[to], weight);
	end

	g, node_labels, edge_labels
end

# ╔═╡ 78d4f47f-96b6-44f9-b9ab-2cb84bbb19b0
example_graph, example_nodes, example_edges = build_graph("""London to Dublin = 464
London to Belfast = 518
Dublin to Belfast = 141""")

# ╔═╡ 7e17c6f5-3972-4777-9611-f4c321108144
gplot(example_graph, nodelabel=example_nodes, edgelabel=example_edges)

# ╔═╡ 8e78e6fe-cb4a-46f7-839b-42aa35f7b71a
puzzle_graph, puzzle_nodes, puzzle_edges = build_graph(puzzle_input)

# ╔═╡ 8dde0e83-121f-4720-80b6-00020bd33a6b
gplot(puzzle_graph, nodelabel=puzzle_nodes, edgelabel=puzzle_edges)

# ╔═╡ 76183faa-4f76-4da6-bcff-255e7d70b730
function read(paths)
	nodes = []
	distances = Dict()

	for line in split(paths, "\n")
		m = match(r"(\w+) to (\w+) = (\d+)", line)
		distances[m[1] * m[2]] = distances[m[2] * m[1]] = parse(Int, m[3])

		if m[1] ∉ nodes
			push!(nodes, m[1])
		end
		if m[2] ∉ nodes
			push!(nodes, m[2])
		end
	end

	nodes, distances
end

# ╔═╡ 8abefd65-5a57-4717-8bf4-4e92ebe9d562
function smallest((nodes, distances),)
	smallest_distance = Inf

	for c in permutations(nodes)
		d = 0
		for idx in 1:length(c)-1
			m = c[idx]
			n = c[idx+1]
			name = m * n
			d += distances[name]
		end

		if d < smallest_distance
			smallest_distance = d
		end
	end

	smallest_distance
end

# ╔═╡ 684af9bd-4282-40ea-85d8-2d2d89c258c0
@assert smallest(read("""London to Dublin = 464
London to Belfast = 518
Dublin to Belfast = 141""")) == 605

# ╔═╡ be574368-3f0d-4863-b420-ad14106b5175
part1 = smallest(read(puzzle_input))

# ╔═╡ a1dafb48-4984-431a-82de-d5280f30af21
@test part1 == 141

# ╔═╡ 1d1f08c8-7092-4ea9-8884-43e28062c036
md"Your puzzle answer was `141`."

# ╔═╡ 64028478-8bc7-412a-82be-de768aaa1fe6
md"""
# Part Two

The next year, just to show off, Santa decides to take the route with the longest distance instead.

He can still start and end at any two (different) locations he wants, and he still must visit each location exactly once.

For example, given the distances above, the longest route would be `982` via (for example) `Dublin -> London -> Belfast`.

**What is the distance of the longest route?**
"""

# ╔═╡ ccd1c5e7-1c13-4608-8293-e4e0a76cca37
function biggest((nodes, distances),)
	biggest_distance = 0

	for c in permutations(nodes)
		d = 0
		for idx in 1:length(c)-1
			m = c[idx]
			n = c[idx+1]
			name = m * n
			d += distances[name]
		end

		if d > biggest_distance
			biggest_distance = d
		end
	end

	biggest_distance
end

# ╔═╡ a548a1da-187a-4aa7-a338-c6c0b4114177
@assert biggest(read("""London to Dublin = 464
London to Belfast = 518
Dublin to Belfast = 141""")) == 982

# ╔═╡ dbdedef1-f657-460f-93a7-e27a2563fcae
part2 = biggest(read(puzzle_input))

# ╔═╡ 7d51e278-0442-4967-9bba-0cbc1518c693
@test part2 == 736

# ╔═╡ 0d5b2494-5454-41ec-9a3d-e627e4c27f51
md"Your puzzle answer was `736`."

# ╔═╡ Cell order:
# ╠═bbfa7498-30bf-4264-b4d9-ee0c238c9984
# ╠═f0d31c21-f866-4d51-86ec-d473147f1677
# ╠═f20ff74b-16f0-4a9d-9969-2d240ba18177
# ╟─56cedfe0-b02c-11eb-1f56-3557d6c170bc
# ╟─1609ee1f-1d90-4889-b39d-5fe0fb6de8c4
# ╠═3ffc6481-2267-4d97-8065-d2fd78c637e9
# ╠═78d4f47f-96b6-44f9-b9ab-2cb84bbb19b0
# ╠═7e17c6f5-3972-4777-9611-f4c321108144
# ╠═8e78e6fe-cb4a-46f7-839b-42aa35f7b71a
# ╠═8dde0e83-121f-4720-80b6-00020bd33a6b
# ╠═76183faa-4f76-4da6-bcff-255e7d70b730
# ╠═8abefd65-5a57-4717-8bf4-4e92ebe9d562
# ╠═684af9bd-4282-40ea-85d8-2d2d89c258c0
# ╠═be574368-3f0d-4863-b420-ad14106b5175
# ╠═a1dafb48-4984-431a-82de-d5280f30af21
# ╟─1d1f08c8-7092-4ea9-8884-43e28062c036
# ╟─64028478-8bc7-412a-82be-de768aaa1fe6
# ╠═ccd1c5e7-1c13-4608-8293-e4e0a76cca37
# ╠═a548a1da-187a-4aa7-a338-c6c0b4114177
# ╠═dbdedef1-f657-460f-93a7-e27a2563fcae
# ╠═7d51e278-0442-4967-9bba-0cbc1518c693
# ╟─0d5b2494-5454-41ec-9a3d-e627e4c27f51


"""
    newton(x0::T,xL::T,xU::T,f::Function,df::Function,envp1::T,envp2::T)

Defines a local 1D newton method to solve for the root of `f` between the bounds
`xL` and `xU` using `x0` as a starting point. The derivative of `f` is `df`. The
inputs `envp1` and `envp2` are the envelope calculation parameters.
"""
@inline function newton(x0::T,xL::T,xU::T,f::Function,df::Function,envp1::T,envp2::T) where {T<:Real}
  dfk::T = zero(x0)

  xk::T = max(xL,min(x0,xU))
  fk::T = f(xk,envp1,envp2)

  for i=1:MC_param.env_max_int
    dfk = df(xk,envp1,envp2)
    if (abs(fk)<MC_param.env_tol)
      return xk
    end
    if (dfk == zero(x0))
      error("NEWTON EXCEPTION")
    elseif (xk==xL && fk/dfk>zero(x0))
      return xk
    elseif (xk==xU && fk/dfk<zero(x0))
      return xk
    end
    xk = max(xL,min(xU,xk-fk/dfk))
    fk = f(xk,envp1,envp2)
  end
  error("NEWTON EXCEPTION")
end

"""
    secant(x0::T,x1::T,xL::T,xU::T,f::Function,envp1::T,envp2::T)  where {T<:Real}

Defines a local 1D secant method to solve for the root of `f` between
the bounds `xL` and `xU` using `x0` and `x1` as a starting points. The inputs
`envp1` and `envp2` are the envelope calculation parameters.
"""
@inline function secant(x0::T,x1::T,xL::T,xU::T,f::Function,envp1::T,envp2::T) where {T<:Real}
  xkm::T = max(xL,min(xU,x0))
  xk::T = max(xL,min(xU,x1))
  fkm::T = f(xkm,envp1,envp2)

  for i=1:MC_param.env_max_int
    fk = f(xk,envp1,envp2)
    Bk::T = (fk-fkm)/(xk-xkm)
    if (abs(fk)<MC_param.env_tol)
      return xk
    end
    if (Bk == zero(x0))
      error("SECANT EXCEPTION")
    elseif ((xk==xL)&(fk/Bk>zero(x0)))
      return xk
    elseif ((xk==xU)&(fk/Bkzero(x0)))
      return xk
    end
    xkm = xk
    fkm = fk
    xk = max(xL,min(xU,xk-fk/Bk))
  end
  error("SECANT EXCEPTION")
end

"""
    golden_section(xL::T,xU::T,f::Function,envp1::T,envp2::T) where {T<:Real}

Defines a local 1D golden section method to solve for the root of `f` between
the bounds `xL` and `xU` using `x0` as a starting point. Define iteration used
in golden section method. The inputs `envp1` and `envp2` are the envelope
calculation parameters.
"""
@inline function golden_section(xL::T,xU::T,f::Function,envp1::T,envp2::T) where {T<:Real}
  fL::T = f(xL,envp1,envp2)
  fU::T = f(xU,envp1,envp2)

  if (fL*fU > zero(fL))
    error("GOLDEN EXCEPTION")
  end
  xm::T = xU-(2.0-golden)*(xU-xL)
  fm::T = f(xm,envp1,envp2)
  return golden_section_it(1,xL,fL,xm,fm,xU,fU,f,envp1,envp2)
end
"""
    golden_section_it(init::Int64,a::T,fa::T,b::T,fb::T,c::T,
                      fc::T,f::Function,envp1::T,envp2::T) where {T<:Real}

Define iteration used in golden section method. The inputs `fa`,`fb`, and `fc`,
are the function `f` evaluated at `a`,`b`, and `c` respectively. The inputs
`envp1` and `envp2` are the envelope calculation parameters. The value `init` is
the iteration number of the golden section method.
"""
@inline function golden_section_it(init::Int64,a::T,fa::T,b::T,fb::T,c::T,
                                   fc::T,f::Function,envp1::T,envp2::T) where {T<:Real}
  b_t_x::Bool = (c-b > b-a)
  if (b_t_x)
    x::T = b + (2.0-golden)*(c-b)
  else
    x = b - (2.0-golden)*(b-a)
  end
  itr::Int64 = init
  if (abs(c-a)<MC_param.env_tol*(abs(b)+abs(x))||(itr>MC_param.env_max_int))
    return (c+a)/2.0
  end
  itr += 1
  fx::T = f(x,envp1,envp2)
  if (b_t_x)
    if (fa*fx<zero(fa))
      golden_section_it(itr,a,fa,b,fb,x,fx,f,envp1,envp2)
    else
      golden_section_it(itr,b,fb,x,fx,c,fc,f,envp1,envp2)
    end
  else
    if (fa*fb<(fa))
      golden_section_it(itr,a,fa,x,fx,b,fb,f,envp1,envp2)
    else
      golden_section_it(itr,x,fx,b,fb,c,fc,f,envp1,envp2)
    end
  end
end


export reflect

"""
    reflect(P::LazySet{N}) where {N<:Real}

Concrete reflection of a convex set `P`, resulting in the reflected set `-P`.

### Note

This function requires that the list of constraints of the set `P` is
available, i.e. such that it can be written as
``P = \\{z ∈ ℝⁿ: ⋂ sᵢᵀz ≤ rᵢ, i = 1, ..., N\\}.``

This function can be used to implement the alternative definition of the
Minkowski Difference, which writes as
```math
A ⊖ B = \\{a − b | a ∈ A, b ∈ B\\} = A ⊕ (-B)
```
by calling `minkowski_sum(A, reflect(B))`.
"""
function reflect(P::LazySet)

    @assert applicable(constraints_list, P)  "this function " *
        "requires that the list of constraints is available, but it is not; " *
        "if the set is bounded, try overapproximating with an `HPolytope` first"

    F,g = tosimplehrep(P)
    T = isbounded(P) ? HPolytope : HPolyhedron
    return T(-F, g)
end