using LinearAlgebra
using Test

using SetProg
using Polyhedra

using JuMP

@testset "Algebraic" begin
    □ = polyhedron(HalfSpace([1, 0], 1.0) ∩ HalfSpace([-1, 0], 1) ∩ HalfSpace([0, 1], 1) ∩ HalfSpace([0, -1], 1))
    @testset "Ellipsoid" begin
        @testset "John" begin
            mock = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
            model = JuMP.direct_model(mock);
            # Q = [1 0
            #      0 1]
            # t = √det(Q) = 1                                                                  Q11  Q12  Q22  t
            MOI.Utilities.set_mock_optimize!(mock, mock -> MOI.Utilities.mock_optimize!(mock, [1.0, 0.0, 1.0, 1.0]));

            @variable(model, ◯, Ellipsoid(symmetric=true, dimension=2))
            cref = @constraint(model, ◯ ⊆ □)
            @objective(model, Max, nth_root(volume(◯)))

            SetProg.optimize!(model)
            @test JuMP.termination_status(model) == MOI.OPTIMAL
            @test JuMP.objective_sense(model) == MOI.MAX_SENSE
            @test JuMP.objective_value(model) == 1.0
            @test JuMP.value(◯) isa SetProg.Sets.PolarEllipsoid
            @test JuMP.value(◯).Q == Symmetric([1.0 0.0; 0.0 1.0])
        end
        @testset "Löwner" begin
            mock = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
            model = JuMP.direct_model(mock);
            # Q = [√2  0
            #       0 √2]
            # t = √det(Q) = 2                                                                  Q11  Q12  Q22  t
            MOI.Utilities.set_mock_optimize!(mock, mock -> MOI.Utilities.mock_optimize!(mock, [0.5, 0.0, 0.5, 0.5]));

            @variable(model, ◯, Ellipsoid(symmetric=true, dimension=2))
            cref = @constraint(model, □ ⊆ ◯)
            @objective(model, Min, nth_root(volume(◯)))

            SetProg.optimize!(model)
            @test JuMP.termination_status(model) == MOI.OPTIMAL
            @test JuMP.objective_sense(model) == MOI.MAX_SENSE
            @test JuMP.objective_value(model) == 0.5
            @test JuMP.value(◯) isa SetProg.Sets.Ellipsoid
            @test JuMP.value(◯).Q == Symmetric([0.5 0.0; 0.0 0.5])
        end
    end
    @testset "Quartic" begin
        @testset "Inner" begin
            mock = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
            model = JuMP.direct_model(mock);
            # The PSD matrix for the variable  is 3 x 3 so 3 * (3+1) / 2 = 6
            # The PSD matrix for the convexity is 6 x 6 so 6 * (6+1) / 2 = 21
            # entries
            # 1 variable for t
            # hence 28 variables
            MOI.Utilities.set_mock_optimize!(mock, mock -> MOI.Utilities.mock_optimize!(mock, ones(28)))

            @variable(model, ◯, PolySet(symmetric=true, degree=4, dimension=2,
                                        convex=true))
            cref = @constraint(model, ◯ ⊆ □)
            @objective(model, Max, nth_root(volume(◯)))

            SetProg.optimize!(model)
            @test JuMP.termination_status(model) == MOI.OPTIMAL
            @test JuMP.objective_sense(model) == MOI.MAX_SENSE
            @test JuMP.objective_value(model) == 1.0
            @test JuMP.value(◯) isa SetProg.Sets.PolarPolySet{Float64}
            @test JuMP.value(◯).degree == 4
            x, y = SetProg.data(model).polyvars
            @test polynomial(JuMP.value(◯).p) == x^4 + 2x^3*y + 3x^2*y^2 + 2x*y^3 + y^4
            @test JuMP.value(◯).convexity_proof.n == 6
            @test JuMP.value(◯).convexity_proof.Q == ones(21)
            @test JuMP.objective_value(model) == 1.0
        end
        @testset "Outer" begin
            mock = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
            model = JuMP.direct_model(mock);
            MOI.Utilities.set_mock_optimize!(mock, mock -> MOI.Utilities.mock_optimize!(mock, ones(28)))

            @variable(model, ◯, PolySet(symmetric=true, degree=4, dimension=2,
                                        convex=true))
            cref = @constraint(model, □ ⊆ ◯)
            @objective(model, Min, nth_root(volume(◯)))

            SetProg.optimize!(model)
            @test JuMP.termination_status(model) == MOI.OPTIMAL
            @test JuMP.objective_sense(model) == MOI.MAX_SENSE
            @test JuMP.objective_value(model) == 1.0
            @test JuMP.value(◯) isa SetProg.Sets.PolySet{Float64}
            @test JuMP.value(◯).degree == 4
            x, y = SetProg.data(model).polyvars
            @test polynomial(JuMP.value(◯).p) == x^4 + 2x^3*y + 3x^2*y^2 + 2x*y^3 + y^4
            @test JuMP.value(◯).convexity_proof.n == 6
            @test JuMP.value(◯).convexity_proof.Q == ones(21)
        end
    end
end


##############################################################################################
module IEEE39bus
##############################################################################################
# imports and exports
using Random, Distributions
using LinearAlgebra
export dx_dt

##############################################################################################
##############################################################################################
# Type union declarations for multiple dispatch and type aliases

# vectors and ensemble members of sample
VecA = Union{Vector{Float64}, SubArray{Float64, 1}}

# dictionary for model parameters
ParamDict = Union{Dict{String, Array{Float64}}, Dict{String, Vector{Float64}}}

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

function dx_dt(x::VecA, t::Float64, dx_params::ParamDict)

    """Time derivative of the phase and fequency of the effective-network swing eq model

    Input x is a 2 n_g vector of the phase and fequency at each of the n_g generator buses.
    The input dx_params is a ParamDict of all system parameters to be passed to the
    integration scheme.  The system is currenty defined autonomously to be run as an SDE,
    noise perturbed steady state."""

    # unpack the system parameters effective network of
    # Nishikawa, T., & Motter, A. E. (2015). Comparative analysis of existing
    # models for power-grid synchronization.
    A = dx_params["A"]::Array{Float64}
    D = dx_params["D"]::Array{Float64}
    H = dx_params["H"]::Array{Float64}
    K = dx_params["K"]::Array{Float64}
    γ = dx_params["γ"]::Array{Float64}
    ω = dx_params["ω"]::Array{Float64}

    # convert the effective bus coupling and passive injection to contain the change
    # of variable terms
    K = ω[1] * K / 2.0
    A = ω[1] * A / 2.0

    # unpack the phase and frequency at the n_g buses, with all phases listed first, then all
    # fequencies in the order of the bus index
    n_g = convert(Int, length(x) / 2)
    δ_1 = @view x[1:n_g]
    δ_2 = @view x[n_g+1:end]

    # define the vector of the derivatives
    dx = zeros(2 * n_g)

    # derivative of the phase equals frequency
    dx[1:n_g] .= δ_2

    # compute the derivative of the inertia normalized frequencies
    # entry j is defined as
    # A_j * ω/2 - D_j /2 * δ_2 - Σ_{i!=j} K * ω/2 * sin(δ_j - δ_i - γ_ij)
    for j in 1:n_g
        for i in 1:n_g
            if j != i
                # K is symmetric, we loop over the columns for faster memory access
                # with the same variable j as in the row index of the derivative
                dx[n_g + j] += -K[i, j] * sin(δ_1[j] - δ_1[i] - γ[i, j])
            end
        end
        # finally apply the remaining terms
        dx[n_g + j] += A[j] - δ_2[j] * D[j] / 2.0
    end
    # to compute the derivative of the frequencies, we finally
    # divide back out by the inertia
    dx[n_g + 1 : end] = dx[n_g + 1: end] ./ H
    return dx
end

##############################################################################################
# end module

end


# DE

# Import Packages

using PyCall
using PyPlot
@pyimport numpy as np

pygui(true)

# Drop Wave Function

f(x) = -(1 + cos(12 * x)) / (0.5 * x ^ 2 + 2)

limit = [-5, 5]
x = np.linspace(limit[1], limit[2], 500)
y = f.(x)

plot(x, y)

# Rats

N = 50
iter = 100
AC = 0.5
CR = 0.8

rat = zeros(N, 2)

for ii = 1: N

    global rat

    rat[ii, 1] = rand() * 2 * limit[2] + limit[1]
    rat[ii, 2] = f.(rat[ii, 1])

end

scatter(rat[:, 1], rat[:, 2])

globalBest = zeros(1, 2)
globalBest[1, 1] = rand() * 2 + 0.5
globalBest[1, 2] = f(globalBest[1, 1])

# Iteration

for jj = 1: iter

    for kk = 1: N

        global rat, globalBest

        ratCR = zeros(3, 1)

        while true

            ratCR = Int64.(floor.(rand(3, 1) * 30) .+ 1)

            if ratCR[1] ≠ ratCR[2] ≠ ratCR[3] ≠ kk

                break

            end

        end

        ratNew = zeros(1, 2)

        if rand() > CR

            ratNew[1, 1] = rat[ratCR[1]] + AC * (rat[ratCR[2]] - rat[ratCR[3]])

            if limit[1] ≤ ratNew[1, 1] ≤ limit[2]

                ratNew[1, 2] = f(ratNew[1, 1])

            else

                ratNew[1, 2] = 0

            end

        else

            ratNew[1, :] = rat[kk, :]

        end

        if ratNew[1, 2] ≤ rat[kk, 2]

            rat[kk, :] = ratNew[1, :]

            if rat[kk, 2] ≤ globalBest[1, 2]

                globalBest[1, :] = rat[kk, :]

            end

        end

    end

end

scatter(globalBest[1, 1], globalBest[1, 2])
println("min = $globalBest")


using COAST
using JuMP
using Test

@testset "COAST.jl" begin
    @test loaded_COAST()==true
end

@testset "he_forward.jl" begin
    ## input
    # rdaam params
    # eqs, fanning curvilinear, ketcham et al 07
    alpha = 0.04672
    c0 = 0.39528
    c1 = 0.01073
    c2 = -65.12969
    c3 =  -7.91715
    rmr0 = 0.79
    eta_q = 0.91
    L_dist = 8.1*1e-4 # cm (!)

    # general params
    U238 = 4*1e-6
    R = 1.9872*1e-3
    L = 60*1e-4 # m
    n_T_pts = 100
    n_iter = 60.0

    # preprocess concentrations
    (U238_V,U235_V,Th232_V,U238_mol,U235_mol,Th232_mol) = conc_to_atoms_per_volume(U238,0.0,density=3.20)

    """
    check if RDAAM reproduces vanilla Durango
    from farley et al, 00
    """
    L1 = 90*1e-4
    logD0_a2 = log10(10^1.5/(L1^2)) # cm^2/s
    Ea = 32.9

    T = collect(LinRange(120.0,0.01,ceil(Int,n_T_pts)).+273.15)
    times = collect(LinRange(120.0,0.01,ceil(Int,n_T_pts)+1).*3.1558e7*1e6)
    mass_he_t1 = rdaam_forward_diffusion(alpha,c0,c1,c2,c3,rmr0,eta_q,L_dist,0.0,0.0,0.0,R,Ea,logD0_a2,n_iter,
                  U238_mol,U238_V,U235_mol,U235_V,Th232_mol,Th232_V,L1,times...,T...)
    print("mass t1 Farley, 00 is: $mass_he_t1 \n")
    pre_he_t1 = (8*(U238_mol*exp(60*1e6*sec_in_yrs/τ38)-U238_mol)+
                 7*(U235_mol*exp(60*1e6*sec_in_yrs/τ35)-U235_mol)) # Dodson value
    print("mass t1 Farley, 00 predicted is: $pre_he_t1 \n")
    @test isapprox(mass_he_t1/pre_he_t1,1.0; rtol = 1e-2)

    """
    calculate erho
    rdaam params in diffn
    test fig 2, 2009 RDAAM
    """
    Rjoules = 8.314
    psi = 1e-13
    omega = 1e-22
    Etrap = 34*1e3 # J/mol
    E_L = 122.3*1e3 # J/mol
    L2 = 60*1e-4
    log10D0L_a2_rdaam = log10(exp(9.733)*L^2/L2^2)
    T = collect(LinRange(120.0,0.0,ceil(Int,n_T_pts)).+273.15)
    times = collect(LinRange(120.0,0.01,ceil(Int,n_T_pts)+1).*3.1558e7*1e6)
    mass_he_t2 = rdaam_forward_diffusion(alpha,c0,c1,c2,c3,0.79,eta_q,L_dist,psi,omega,Etrap,Rjoules,E_L,log10D0L_a2_rdaam,n_iter,
                   U238_mol,U238_V,U235_mol,U235_V,Th232_mol,Th232_V,L,times...,T...)
    pre_he_t2 = (8*(U238_mol*exp(43.6*1e6*sec_in_yrs/τ38)-U238_mol)+
                 7*(U235_mol*exp(43.6*1e6*sec_in_yrs/τ35)-U235_mol))# Dodson value
    print("mass of he is: $(mass_he_t2) \n")
    print("predicted mass of he is: $(pre_he_t2) \n")
    # non-redistributed value
    @test isapprox(mass_he_t2/pre_he_t2,1.0; rtol = 1e-2)

    """
    wolf,98 fig. 5 - history 3
    """
    Ea = 36.2
    logD0_a2 = log10(10^(7.7))

    T1 = collect(LinRange(60.0,60.0,ceil(Int,n_T_pts/2)).+273.15)
    T2 = collect(LinRange(15.0,15.0,floor(Int,n_T_pts/2)).+273.15)
    T = vcat(T1,T2)
    times1 = collect(LinRange(100.0,20.0,ceil(Int,n_T_pts/2)+1).*3.1558e7*1e6)
    times2 = collect(LinRange(19.9,0.01,floor(Int,n_T_pts/2)).*3.1558e7*1e6)
    times = vcat(times1,times2)
    mass_he_t3 = rdaam_forward_diffusion(alpha,c0,c1,c2,c3,0.83,eta_q,L_dist,
                                         0.0,0.0,0.0,R,Ea,logD0_a2,n_iter,
                                         U238_mol,U238_V,U235_mol,U235_V,
                                         Th232_mol,Th232_V,L,times...,T...)
    pre_he_t3 = (8*(U238_mol*exp(41.4*1e6*sec_in_yrs/τ38)-U238_mol)+ # HeFTy output
                 7*(U235_mol*exp(41.4*1e6*sec_in_yrs/τ35)-U235_mol)) # HeFTy output
    print("Wolf test measured is $(mass_he_t3)! \n")
    print("Wolf test predicted is $(pre_he_t3)! \n")
    @test isapprox(mass_he_t3/pre_he_t3,1.0; atol = 1e-2) # <1 Ma error vs
end

@testset "he_vars_constraints.jl" begin
     # basic setup
     model1 = initialize_JuMP_model("mumps")
     @variable(model1, T)
     @test num_variables(model1)==1

     # add bounds
     model2 = initialize_JuMP_model("mumps")
     @variable(model2, 273.0<=T<=400.0)
     @test num_variables(model2)==1
     @test has_lower_bound(T)
     @test has_upper_bound(T)
     print("Basic JuMP tests passed!\n")
end

# basic model with custom constraints
@testset "custom constraints" begin
    # general params
    U238 = 4*1e-6
    R = 1.9872*1e-3
    L = 60*1e-4 # m
    n_T_pts = 100
    n_iter = 60.0
     model3 = initialize_JuMP_model("mumps",print_level=1)
     time_segs = 10
     (T,set_val) = define_variables!(time_segs-1,1,model3,0.1*ones(time_segs-1))
     @test num_variables(model3)==time_segs-1
     @test register_forward_model!(time_segs,model3)==true
     register_objective_function!(time_segs-1,model3)

     ## input
     # rdaam params
     # eqs, fanning curvilinear, ketcham et al 07
     alpha = 0.04672
     c0 = 0.39528
     c1 = 0.01073
     c2 = -65.12969
     c3 =  -7.91715
     rmr0 = 0.79
     eta_q = 0.91
     L_dist = 8.1*1e-4 # cm (!)

     # general params
     U238 = 28*1e-6
     R = 1.9872*1e-3
     L = 60*1e-4 # m
     times = collect(LinRange(120.0,0.01,time_segs).*3.1558e7*1e6)

     # preprocess concentrations
     (U238_V,U235_V,Th232_V,U238_mol,U235_mol,Th232_mol) = conc_to_atoms_per_volume(U238,0.0,density=3.20)

     L1 = 90*1e-4
     logD0_a2 = log10(10^1.5/(L1^2)) # cm^2/s
     Ea = 32.9

     constr1=@NLconstraint(model3,[i = 1:1],rdaam_forward_diffusion(alpha,c0,c1,c2,c3,rmr0,eta_q,L_dist,0.0,0.0,0.0,R,Ea,logD0_a2,n_iter,
                   U238_mol,U238_V,U235_mol,U235_V,Th232_mol,Th232_V,L,times...,T...)==3e-9)
     @NLobjective(model3,Min,mod_constraints(T...))
     optimize!(model3)
     @test termination_status(model3) == MOI.LOCALLY_SOLVED
end

@testset "one specified constraint" begin
    model4 = initialize_JuMP_model("mumps",print_level=1)
    time_segs = 10
    (T,set_dev) = define_variables!(time_segs-1,1,model4,0.1*ones(time_segs-1))
    register_objective_function!(time_segs-1,model4)
    register_forward_model!(time_segs,model4)
    ## input
    # rdaam params
    # eqs, fanning curvilinear, ketcham et al 07
    alpha = 0.04672
    c0 = 0.39528
    c1 = 0.01073
    c2 = -65.12969
    c3 =  -7.91715
    rmr0 = 0.79
    eta_q = 0.91
    L_dist = 8.1*1e-4 # cm (!)

    # general params
    U238 = 28*1e-6
    R = 1.9872*1e-3
    L = 60*1e-4 # m
    times = collect(LinRange(120.0,0.01,time_segs).*3.1558e7*1e6)

    # preprocess concentrations
    (U238_V,U235_V,Th232_V,U238_mol,U235_mol,Th232_mol) = ppm_to_atoms_per_volume(U238,0.0,density=3.20)

    L1 = 90*1e-4
    logD0_a2 = log10(10^1.5/(L1^2)) # cm^2/s
    Ea = 32.9
    density = 3.2
    He_conc = [3e-9]
    rdaam_define_constraints(model4,density,set_dev,He_conc,L,times,T,U238_mol,U235_mol,Th232_mol,U238_V,U235_V,Th232_V)
    optimize!(model4)
    @test termination_status(model4) == MOI.LOCALLY_SOLVED
end

@testset "FT_setup.jl" begin
     # Ketcham 07, 5_5M, Dpar case
    (l0m_K07_5_5Dpar,l0cm_K07_5_5Dpar)= init_track_len(calc_type="Dpar",
                                                Dpar=1.65,Ketcham07_5_5m=true)
    @test isapprox(l0m_K07_5_5Dpar,16.1,rtol=1e-3)
    @test isapprox(l0cm_K07_5_5Dpar,16.3,rtol=1e-3)

    (l0m_K07_5_5rmr0,l0cm_K07_5_5rmr0)= init_track_len(calc_type="rmr0",
                                                rmr0=1.65,Ketcham07_5_5m=true)
    @test isapprox(l0m_K07_5_5rmr0,16.14,rtol=1e-3)
    @test isapprox(l0cm_K07_5_5rmr0,16.35,rtol=1e-3)

    (l0m_K07_5_5cl_a,l0cm_K07_5_5cl_a)= init_track_len(calc_type="Cl_apfu",
                                                Cl=1.65,Ketcham07_5_5m=true)
    @test isapprox(l0m_K07_5_5cl_a,17.80,rtol=1e-3)
    @test isapprox(l0cm_K07_5_5cl_a,18.33,rtol=1e-3)

    (l0m_K07_5_5a_vol,l0cm_K07_5_5a_vol)= init_track_len(calc_type="a_len",
                                                a_vol=10.65,Ketcham07_5_5m=true)
    @test isapprox(l0m_K07_5_5a_vol,23.18,rtol=1e-3)
    @test isapprox(l0cm_K07_5_5a_vol,24.53,rtol=1e-3)

    # Donelick99 case
   (l0m_D99Dpar,l0cm_D99Dpar)= init_track_len(calc_type="Dpar",
                                               Dpar=1.65,Ketcham07_5_5m=false,
                                               Donelick99=true)
   @test isapprox(l0m_D99Dpar,16.1,rtol=1e-3)
   @test isapprox(l0cm_D99Dpar,16.3,rtol=1e-3)

   (l0m_D99rmr0,l0cm_D99rmr0)= init_track_len(calc_type="rmr0",
                                               rmr0=1.65,Donelick99=true)
   @test isapprox(l0m_D99rmr0,16.14,rtol=1e-3)
   @test isapprox(l0cm_D99rmr0,16.35,rtol=1e-3)

   (l0m_D99cl_a,l0cm_D99cl_a)= init_track_len(calc_type="Cl_apfu",
                                               Cl=1.65,Donelick99=true)
   @test isapprox(l0m_D99cl_a,17.80,rtol=1e-3)
   @test isapprox(l0cm_D99cl_a,18.33,rtol=1e-3)

   # Ketcham 07, 5_0M, Dpar case
  (l0m_K07_5_0Dpar,l0cm_K07_5_0Dpar)= init_track_len(calc_type="Dpar",
                                              Dpar=1.65,Ketcham07_5_5m=false,
                                              Ketcham07_5_0m=true)
  @test isapprox(l0m_K07_5_0Dpar,15.82,rtol=1e-3)
  @test isapprox(l0cm_K07_5_0Dpar,16.06,rtol=1e-3)

  (l0m_K07_5_0rmr0,l0cm_K07_5_0rmr0)= init_track_len(calc_type="rmr0",
                                              rmr0=1.65,Ketcham07_5_0m=true)
  @test isapprox(l0m_K07_5_0rmr0,15.94,rtol=1e-3)
  @test isapprox(l0cm_K07_5_0rmr0,16.19,rtol=1e-3)

  (l0m_K07_5_0cl_a,l0cm_K07_5_0cl_a)= init_track_len(calc_type="Cl_apfu",
                                              Cl=1.65,Ketcham07_5_0m=true)
  @test isapprox(l0m_K07_5_0cl_a,16.82,rtol=1e-3)
  @test isapprox(l0cm_K07_5_0cl_a,17.18,rtol=1e-3)

  (l0m_K07_5_0a_vol,l0cm_K07_5_0a_vol)= init_track_len(calc_type="a_len",
                                              a_vol=10.65,Ketcham07_5_0m=true)
  @test isapprox(l0m_K07_5_0a_vol,24.85,rtol=1e-3)
  @test isapprox(l0cm_K07_5_0a_vol,25.22,rtol=1e-3)

  # test custom functionality
  (l0m_cust,l0cm_cust)= init_track_len(custom=true,custom_param=1.65,
  custom_intercept=3.0,custom_slope=0.5,custom_c_intercept=2.0,
                                       custom_c_slope=1.0)
   @test isapprox(l0m_cust,(1.65*0.5+3.0),rtol=1e-3)
   @test isapprox(l0cm_cust,(1.65*1.0+2.0),rtol=1e-3)
end

# @testset "FT_forward.jl" begin
#     """
#     RDAAM type cooling
#     """
#     T = collect(LinRange(120.0,0.01,ceil(Int,n_T_pts)).+273.15)
#     times = collect(LinRange(120.0,0.01,ceil(Int,n_T_pts)+1).*3.1558e7*1e6)
#     alpha = 0.04672
#     beta = 0.0
#     c0 = 0.39528
#     c1 = 0.01073
#     c2 = -65.12969
#     c3 =  -7.91715

#     rcb2_test = FT_forward(alpha,beta,c0,c1,c2,c3,T,times)
# end


"""
    rte_lw_noscat_source!(
        src_lw::SourceLWNoScat{FT},
        op::OneScalar{FT},
        gcol,
        nlay,
    ) where {FT<:AbstractFloat}

Compute LW source function for upward and downward emission at levels using linear-in-tau assumption
See Clough et al., 1992, doi: 10.1029/92JD01419, Eq 13
"""
function rte_lw_noscat_source!(
    src_lw::SourceLWNoScat{FT},
    op::OneScalar{FT},
    gcol,
    nlay,
) where {FT<:AbstractFloat}
    # setting references
    @unpack src_up, src_dn, lev_source_inc, lev_source_dec, lay_source = src_lw
    #    @unpack nlay = slv.as
    Ds = op.angle_disc.gauss_Ds
    τ = op.τ

    τ_thresh = sqrt(eps(FT)) # or abs(eps(FT))?

    @inbounds for glay = 1:nlay
        τ_loc = τ[glay, gcol] * Ds[1]      # Optical path and transmission,

        trans = exp(-τ_loc)    # used in source function and transport calculations
        # Weighting factor. Use 2nd order series expansion when rounding error (~tau^2)
        # is of order epsilon (smallest difference from 1. in working precision)
        # Thanks to Peter Blossey

        fact = (τ_loc > τ_thresh) ? ((FT(1) - trans) / τ_loc - trans) :
            τ_loc * (FT(1 / 2) - FT(1 / 3) * τ_loc)

        # Equations below are developed in Clough et al., 1992, doi:10.1029/92JD01419, Eq 13

        src_up[glay, gcol] =
            (FT(1) - trans) * lev_source_inc[glay, gcol] +
            FT(2) * fact * (lay_source[glay, gcol] - lev_source_inc[glay, gcol])

        src_dn[glay, gcol] =
            (FT(1) - trans) * lev_source_dec[glay, gcol] +
            FT(2) * fact * (lay_source[glay, gcol] - lev_source_dec[glay, gcol])
    end
end

"""
    rte_lw_noscat_transport!(
        src_lw::SourceLWNoScat{FT},
        bcs_lw::LwBCs{FT},
        op::OneScalar{FT},
        gcol,
        flux::FluxLW{FT},
        igpt,
        ibnd,
        nlay,
        nlev,
    ) where {FT<:AbstractFloat}

Transport for no-scattering longwave problem.
"""
function rte_lw_noscat_transport!(
    src_lw::SourceLWNoScat{FT},
    bcs_lw::LwBCs{FT},
    op::OneScalar{FT},
    gcol,
    flux::FluxLW{FT},
    igpt,
    ibnd,
    nlay,
    nlev,
) where {FT<:AbstractFloat}
    # setting references
    @unpack src_up, src_dn, sfc_source = src_lw
    @unpack sfc_emis, inc_flux = bcs_lw
    @unpack flux_up, flux_dn, flux_net = flux

    Ds = op.angle_disc.gauss_Ds
    w_μ = op.angle_disc.gauss_wts
    n_μ = op.angle_disc.n_gauss_angles
    τ = op.τ

    @inbounds for ilev = 1:nlev
        flux_dn[ilev, gcol] = FT(0)
        flux_up[ilev, gcol] = FT(0)
    end

    if inc_flux ≠ nothing
        @inbounds flux_dn[nlev, gcol] = inc_flux[gcol, igpt]
    end
    # Transport is for intensity
    #   convert flux at top of domain to intensity assuming azimuthal isotropy
    @inbounds flux_dn[nlev, gcol] /= (FT(2) * FT(π) * w_μ[1])

    # Top of domain is index nlev
    # Downward propagation
    @inbounds for ilev = nlay:-1:1
        trans = exp(-τ[ilev, gcol] * Ds[1])
        flux_dn[ilev, gcol] = trans * flux_dn[ilev+1, gcol] + src_dn[ilev, gcol]
    end

    # Surface reflection and emission
    @inbounds flux_up[1, gcol] =
        flux_dn[1, gcol] * (FT(1) - sfc_emis[ibnd, gcol]) +
        sfc_emis[ibnd, gcol] * sfc_source[gcol]

    # Upward propagation
    @inbounds for ilev = 2:nlay+1
        trans = exp(-τ[ilev-1, gcol] * Ds[1])
        flux_up[ilev, gcol] =
            trans * flux_up[ilev-1, gcol] + src_up[ilev-1, gcol]
    end

    @inbounds for ilev = 1:nlev
        flux_up[ilev, gcol] *= (FT(2) * FT(π) * w_μ[1])
        flux_dn[ilev, gcol] *= (FT(2) * FT(π) * w_μ[1])
        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
    end
end

"""
    rte_lw_2stream_combine_sources!(src::SourceLW2Str, gcol, nlev, ncol)

RRTMGP provides source functions at each level using the spectral mapping
of each adjacent layer. This function combines these for two-stream calculations.
"""
function rte_lw_2stream_combine_sources!(src::SourceLW2Str, gcol, nlev, ncol)
    @inbounds src.lev_source[1, gcol] = src.lev_source_dec[1, gcol] # glev == 1
    @inbounds src.lev_source[nlev, gcol] = src.lev_source_inc[nlev-1, gcol] # glev == nlev
    @inbounds for glev = 2:nlev-1
        src.lev_source[glev, gcol] = sqrt(
            src.lev_source_dec[glev, gcol] * src.lev_source_inc[glev-1, gcol],
        )
    end
end

"""
    rte_lw_2stream_source!(
        op::TwoStream{FT},
        src_lw::SourceLW2Str{FT},
        gcol,
        nlay,
        ncol,
    ) where {FT<:AbstractFloat}

This function combines RRTMGP-specific sources at levels,
computes layer reflectance, transmittance, and
total source function at levels using linear-in-tau approximation.
"""
function rte_lw_2stream_source!(
    op::TwoStream{FT},
    src_lw::SourceLW2Str{FT},
    gcol,
    nlay,
    ncol,
) where {FT<:AbstractFloat}
    # Cell properties: reflection, transmission for diffuse radiation
    # Coupling coefficients needed for source function
    # -------------------------------------------------------------------------------------------------
    #
    # Longwave two-stream solutions to diffuse reflectance and transmittance for a layer
    #    with optical depth tau, single scattering albedo w0, and asymmetery parameter g.
    #
    # Equations are developed in Meador and Weaver, 1980,
    #    doi:10.1175/1520-0469(1980)037<0630:TSATRT>2.0.CO;2
    #
    # -------------------------------------------------------------------------------------------------
    # setting references
    @unpack τ, ssa, g = op
    @unpack Rdif, Tdif, lev_source, src_up, src_dn = src_lw

    lw_diff_sec = FT(1.66)

    @inbounds for glay = 1:nlay
        γ1 = lw_diff_sec * (1 - FT(0.5) * ssa[glay, gcol] * (1 + g[glay, gcol]))
        γ2 = lw_diff_sec * FT(0.5) * ssa[glay, gcol] * (1 - g[glay, gcol])
        k = sqrt(max((γ1 + γ2) * (γ1 - γ2), FT(1e-12)))

        # Refactored to avoid rounding errors when k, gamma1 are of very different magnitudes
        RT_term =
            1 / (
                k * (1 + exp(-2 * τ[glay, gcol] * k)) +
                γ1 * (1 - exp(-2 * τ[glay, gcol] * k))
            )

        @inbounds Rdif[glay, gcol] =
            RT_term * γ2 * (1 - exp(-2 * τ[glay, gcol] * k)) # Equation 25
        @inbounds Tdif[glay, gcol] = RT_term * 2 * k * exp(-τ[glay, gcol] * k) # Equation 26

        # Source function for diffuse radiation
        # Compute LW source function for upward and downward emission at levels using linear-in-tau assumption
        # This version straight from ECRAD
        # Source is provided as W/m2-str; factor of pi converts to flux units
        # lw_source_2str

        top = glay + 1
        bot = glay

        if τ[glay, gcol] > 1.0e-8
            # Toon et al. (JGR 1989) Eqs 26-27
            Z =
                (lev_source[bot, gcol] - lev_source[top, gcol]) /
                (τ[glay, gcol] * (γ1 + γ2))
            Zup_top = Z + lev_source[top, gcol]
            Zup_bottom = Z + lev_source[bot, gcol]
            Zdn_top = -Z + lev_source[top, gcol]
            Zdn_bottom = -Z + lev_source[bot, gcol]
            @inbounds src_up[glay, gcol] =
                pi * (
                    Zup_top - Rdif[glay, gcol] * Zdn_top -
                    Tdif[glay, gcol] * Zup_bottom
                )
            @inbounds src_dn[glay, gcol] =
                pi * (
                    Zdn_bottom - Rdif[glay, gcol] * Zup_bottom -
                    Tdif[glay, gcol] * Zdn_top
                )
        else
            @inbounds src_up[glay, gcol] = FT(0)
            @inbounds src_dn[glay, gcol] = FT(0)
        end
    end
end

"""
    adding_lw!(
        flux::FluxLW{FT},
        src_lw::SL,
        bcs_lw::BCL,
        gcol,
        igpt,
        ibnd,
        nlev,
        ncol,
    ) where {FT<:AbstractFloat,SL<:SourceLW2Str{FT},BCL<:LwBCs{FT}}

Transport of diffuse radiation through a vertically layered atmosphere, for the longwave problem.
Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08)
"""
function adding_lw!(
    flux::FluxLW{FT},
    src_lw::SL,
    bcs_lw::BCL,
    gcol,
    igpt,
    ibnd,
    nlev,
    ncol,
) where {FT<:AbstractFloat,SL<:SourceLW2Str{FT},BCL<:LwBCs{FT}}
    nlay = nlev - 1
    # setting references
    @unpack flux_up, flux_dn, flux_net = flux

    @unpack albedo, sfc_source, Rdif, Tdif, src_up, src_dn, src = src_lw
    @unpack inc_flux, sfc_emis = bcs_lw

    @inbounds for ilev = 1:nlev
        flux_dn[ilev, gcol] = FT(0)
        flux_up[ilev, gcol] = FT(0)
    end

    if inc_flux ≠ nothing
        @inbounds flux_dn[nlev, gcol] = inc_flux[gcol, igpt]
    end
    # Albedo of lowest level is the surface albedo...
    @inbounds albedo[1, gcol] = FT(1) - sfc_emis[ibnd, gcol]
    # ... and source of diffuse radiation is surface emission
    @inbounds src[1, gcol] = FT(π) * sfc_emis[ibnd, gcol] * sfc_source[gcol]

    # From bottom to top of atmosphere --
    #   compute albedo and source of upward radiation
    @inbounds for ilev = 1:nlay
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        albedo[ilev+1, gcol] =
            Rdif[ilev, gcol] +
            Tdif[ilev, gcol] * Tdif[ilev, gcol] * albedo[ilev, gcol] * denom # Equation 9
        #
        # Equation 11 -- source is emitted upward radiation at top of layer plus
        # radiation emitted at bottom of layer,
        # transmitted through the layer and reflected from layers below (Tdiff*src*albedo)
        src[ilev+1, gcol] =
            src_up[ilev, gcol] +
            Tdif[ilev, gcol] *
            denom *
            (src[ilev, gcol] + albedo[ilev, gcol] * src_dn[ilev, gcol])
    end

    # Eq 12, at the top of the domain upwelling diffuse is due to ...
    @inbounds flux_up[nlev, gcol] =
        flux_dn[nlev, gcol] * albedo[nlev, gcol] + # ... reflection of incident diffuse and
        src[nlev, gcol]                          # scattering by the direct beam below

    # From the top of the atmosphere downward -- compute fluxes
    @inbounds for ilev = nlay:-1:1
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        flux_dn[ilev, gcol] =
            (
                Tdif[ilev, gcol] * flux_dn[ilev+1, gcol] + # Equation 13
                Rdif[ilev, gcol] * src[ilev, gcol] +
                src_dn[ilev, gcol]
            ) * denom
        flux_up[ilev, gcol] =
            flux_dn[ilev, gcol] * albedo[ilev, gcol] + # Equation 12
            src[ilev, gcol]
    end

    @inbounds for ilev = 1:nlev
        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
    end
end
# ---------------------------------------------------------------
"""
    rte_sw_noscat_solve_kernel!(
        flux::FluxSW{FT},
        op::OneScalar{FT},
        bcs_sw::SwBCs{FT},
        solar_frac::FT,
        gcol,
        nlev,
    ) where {FT<:AbstractFloat}

No-scattering solver for the shortwave problem.
(Extinction-only i.e. solar direct beam)
"""
function rte_sw_noscat_solve_kernel!(
    flux::FluxSW{FT},
    op::OneScalar{FT},
    bcs_sw::SwBCs{FT},
    solar_frac::FT,
    gcol,
    nlev,
) where {FT<:AbstractFloat}
    @unpack toa_flux, zenith = bcs_sw
    τ = op.τ
    flux_dn_dir = flux.flux_dn_dir
    # downward propagation
    flux_dn_dir[nlev, gcol] = toa_flux[gcol] * solar_frac * cos(zenith[gcol])
    for ilev = nlev-1:-1:1
        flux_dn_dir[ilev, gcol] =
            flux_dn_dir[ilev+1, gcol] * exp(-τ[ilev, gcol] / cos(zenith[gcol]))
    end
end

"""
    sw_two_stream!(
        op::TwoStream{FT},
        src_sw::SourceSW2Str{FT},
        bcs_sw::SwBCs{FT},
        gcol,
        nlay,
    ) where {FT<:AbstractFloat}

Computes cell properties (transmittance and reflectance) for direct and diffuse radiation

Two-stream solutions to direct and diffuse reflectance and transmittance for a layer
with optical depth tau, single scattering albedo w0, and asymmetery parameter g.

Equations are developed in Meador and Weaver, 1980,
doi:10.1175/1520-0469(1980)037<0630:TSATRT>2.0.CO;2
"""
function sw_two_stream!(
    op::TwoStream{FT},
    src_sw::SourceSW2Str{FT},
    bcs_sw::SwBCs{FT},
    gcol,
    nlay,
) where {FT<:AbstractFloat}
    zenith = bcs_sw.zenith
    @unpack τ, ssa, g = op
    @unpack Rdif, Tdif, Rdir, Tdir, Tnoscat = src_sw

    @inbounds for glay = 1:nlay
        # Zdunkowski Practical Improved Flux Method "PIFM"
        #  (Zdunkowski et al., 1980;  Contributions to Atmospheric Physics 53, 147-66)
        γ1 =
            (FT(8) - ssa[glay, gcol] * (FT(5) + FT(3) * g[glay, gcol])) *
            FT(0.25)
        γ2 = FT(3) * (ssa[glay, gcol] * (FT(1) - g[glay, gcol])) * FT(0.25)
        γ3 = (FT(2) - (FT(3) * cos(zenith[gcol])) * g[glay, gcol]) * FT(0.25)
        γ4 = FT(1) - γ3
        α1 = γ1 * γ4 + γ2 * γ3                          # Eq. 16
        α2 = γ1 * γ3 + γ2 * γ4                          # Eq. 17
        k = sqrt(max((γ1 - γ2) * (γ1 + γ2), FT(1e-12)))

        exp_minusktau = exp(-τ[glay, gcol] * k)
        exp_minus2ktau = exp_minusktau * exp_minusktau

        # Refactored to avoid rounding errors when k, gamma1 are of very different magnitudes
        RT_term =
            FT(1) /
            (k * (FT(1) + exp_minus2ktau) + γ1 * (FT(1) - exp_minus2ktau))

        Rdif[glay, gcol] = RT_term * γ2 * (FT(1) - exp_minus2ktau) # Eqn. 25
        Tdif[glay, gcol] = RT_term * FT(2) * k * exp_minusktau     # Eqn. 26

        # Transmittance of direct, unscattered beam. Also used below
        Tnoscat[glay, gcol] = exp(-τ[glay, gcol] / cos(zenith[gcol]))

        # Direct reflect and transmission
        k_μ = k * cos(zenith[gcol])
        k_γ3 = k * γ3
        k_γ4 = k * γ4

        # Equation 14, multiplying top and bottom by exp(-k*tau)
        #   and rearranging to avoid div by 0.
        if abs(FT(1) - k_μ * k_μ) ≥ eps(FT)
            RT_term = ssa[glay, gcol] * RT_term / (FT(1) - k_μ * k_μ)
        else
            RT_term = ssa[glay, gcol] * RT_term / eps(FT)
        end

        Rdir[glay, gcol] =
            RT_term * (
                (FT(1) - k_μ) * (α2 + k_γ3) -
                (FT(1) + k_μ) * (α2 - k_γ3) * exp_minus2ktau -
                FT(2) * (k_γ3 - α2 * k_μ) * exp_minusktau * Tnoscat[glay, gcol]
            )
        #
        # Equation 15, multiplying top and bottom by exp(-k*tau),
        #   multiplying through by exp(-tau/mu0) to
        #   prefer underflow to overflow
        # Omitting direct transmittance
        #
        Tdir[glay, gcol] =
            -RT_term * (
                (FT(1) + k_μ) * (α1 + k_γ4) * Tnoscat[glay, gcol] -
                (FT(1) - k_μ) *
                (α1 - k_γ4) *
                exp_minus2ktau *
                Tnoscat[glay, gcol] - FT(2) * (k_γ4 + α1 * k_μ) * exp_minusktau
            )
    end
end

"""
    sw_source_2str!(
        src_sw::SourceSW2Str{FT},
        bcs_sw::SwBCs{FT},
        gcol,
        flux::FluxSW{FT},
        solar_frac::FT,
        ibnd,
        nlay,
    ) where {FT<:AbstractFloat}

Direct-beam and source for diffuse radiation
"""
function sw_source_2str!(
    src_sw::SourceSW2Str{FT},
    bcs_sw::SwBCs{FT},
    gcol,
    flux::FluxSW{FT},
    solar_frac::FT,
    ibnd,
    nlay,
) where {FT<:AbstractFloat}
    @unpack toa_flux, zenith, sfc_alb_direct = bcs_sw
    @unpack Rdir, Tdir, Tnoscat, src_up, src_dn, sfc_source = src_sw
    flux_dn_dir = flux.flux_dn_dir

    # layer index = level index
    # previous level is up (+1)
    flux_dn_dir[nlay+1, gcol] = toa_flux[gcol] * solar_frac * cos(zenith[gcol])
    for ilev = nlay:-1:1
        src_up[ilev, gcol] = Rdir[ilev, gcol] * flux_dn_dir[ilev+1, gcol]
        src_dn[ilev, gcol] = Tdir[ilev, gcol] * flux_dn_dir[ilev+1, gcol]
        flux_dn_dir[ilev, gcol] =
            Tnoscat[ilev, gcol] * flux_dn_dir[ilev+1, gcol]
    end
    sfc_source[gcol] = flux_dn_dir[1, gcol] * sfc_alb_direct[ibnd, gcol]
end

"""
    adding_sw!(
        src_sw::SourceSW2Str{FT},
        bcs_sw::SwBCs{FT},
        gcol,
        flux::FluxSW{FT},
        igpt,
        ibnd,
        nlev,
    ) where {FT<:AbstractFloat}

Transport for the two stream shortwave problem.

Transport of diffuse radiation through a vertically layered atmosphere.
Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08)
"""
function adding_sw!(
    src_sw::SourceSW2Str{FT},
    bcs_sw::SwBCs{FT},
    gcol,
    flux::FluxSW{FT},
    igpt,
    ibnd,
    nlev,
) where {FT<:AbstractFloat}
    nlay = nlev - 1
    # setting references
    @unpack flux_up, flux_dn, flux_net = flux

    @unpack albedo, sfc_source, Rdif, Tdif, src_up, src_dn, src = src_sw

    @inbounds for ilev = 1:nlev
        flux_dn[ilev, gcol] = FT(0)
        flux_up[ilev, gcol] = FT(0)
    end

    # Albedo of lowest level is the surface albedo...
    @inbounds albedo[1, gcol] = bcs_sw.sfc_alb_diffuse[ibnd, gcol]
    # ... and source of diffuse radiation is surface emission
    @inbounds src[1, gcol] = sfc_source[gcol]
    # From bottom to top of atmosphere --
    #   compute albedo and source of upward radiation
    @inbounds for ilev = 1:nlay
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        albedo[ilev+1, gcol] =
            Rdif[ilev, gcol] +
            Tdif[ilev, gcol] * Tdif[ilev, gcol] * albedo[ilev, gcol] * denom # Equation 9
        #
        # Equation 11 -- source is emitted upward radiation at top of layer plus
        # radiation emitted at bottom of layer,
        # transmitted through the layer and reflected from layers below (Tdiff*src*albedo)
        src[ilev+1, gcol] =
            src_up[ilev, gcol] +
            Tdif[ilev, gcol] *
            denom *
            (src[ilev, gcol] + albedo[ilev, gcol] * src_dn[ilev, gcol])
    end

    # Eq 12, at the top of the domain upwelling diffuse is due to ...
    flux_up[nlev, gcol] =
        flux_dn[nlev, gcol] * albedo[nlev, gcol] + # ... reflection of incident diffuse and
        src[nlev, gcol]                          # scattering by the direct beam below

    # From the top of the atmosphere downward -- compute fluxes
    @inbounds for ilev = nlay:-1:1
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        flux_dn[ilev, gcol] =
            (
                Tdif[ilev, gcol] * flux_dn[ilev+1, gcol] + # Equation 13
                Rdif[ilev, gcol] * src[ilev, gcol] +
                src_dn[ilev, gcol]
            ) * denom
        flux_up[ilev, gcol] =
            flux_dn[ilev, gcol] * albedo[ilev, gcol] + # Equation 12
            src[ilev, gcol]
    end

    @inbounds for ilev = 1:nlev
        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
    end
end
# ---------------------------------------------------------------


using LinearAlgebra
@testset "Gauss--Newton" begin
    coeff = [3; -1; 2.0; 0.1; 1.2]
    p =
        z ->
            coeff[1] * one(z) +
            coeff[2] * z +
            coeff[3] * z^2 +
            coeff[4] * z^3 +
            coeff[5] * z^4
    discr = collect(5 * exp.(1im * range(0, 2 * pi, length = 21))[1:end-1])

    coeff_pert = [3; -1; 2.0; 0.1; 1.2] + 0.9 * [-3; 1; -2.0; 0.1; -1.2]
    (graph, cref) = graph_monomial(coeff_pert)
    graph = complex(graph)
    opt_gauss_newton!(
        graph,
        p,
        discr,
        logger = 0,
        stoptol = 1e-14,
        γ0 = 0.85,
        cref = cref,
        errtype = :relerr,
        linlsqr = :svd,
        droptol = 1e-15,
    )

    @test all(
        abs.(eval_graph(graph, discr) .- p.(discr)) .< 1e-14 * abs.(p.(discr)),
    )
    A = [3 4; 5 6.6]
    PA = p(A)
    @test norm(eval_graph(graph, A) - PA) < 1e-14 * norm(PA)

    new_input = :B
    (graph, cref) = graph_monomial(coeff_pert, input = new_input)
    graph = complex(graph)
    opt_gauss_newton!(
        graph,
        p,
        discr,
        logger = 0,
        stoptol = 1e-14,
        input = new_input,
        γ0 = 0.85,
        cref = cref,
        errtype = :relerr,
        linlsqr = :svd,
        droptol = 1e-15,
    )

    @test all(
        abs.(eval_graph(graph, discr, input = :B) .- p.(discr)) .<
        1e-14 * abs.(p.(discr)),
    )
    A = [3 4; 5 6.6]
    PA = p(A)
    @test norm(eval_graph(graph, A, input = :B) - PA) < 1e-14 * norm(PA)

    coeff_pert = [3; -1; 2.0; 0.1; 1.2] + 0.9 * [-3; 1; -2.0; 0.1; -1.2]
    (graph, cref) = graph_monomial(coeff_pert)
    graph = complex(graph)
    opt_gauss_newton!(
        graph,
        p,
        discr,
        stoptol = 5e-13,
        cref = cref,
        errtype = :abserr,
    )
    @test all(abs.(eval_graph(graph, discr) .- p.(discr)) .< 5e-13)

    new_input = :Anew
    rename_node!(graph, :A, new_input, cref)
    graph = complex(graph)
    opt_gauss_newton!(
        graph,
        p,
        discr,
        stoptol = 5e-13,
        input = new_input,
        cref = cref,
        errtype = :abserr,
    )
    @test all(
        abs.(eval_graph(graph, discr, input = new_input) .- p.(discr)) .< 5e-13,
    )
end


using Plots

function step_func(x)
    return x > 0
end

function sigmoid(x)
    return 1 / (1 + exp(-x))
end

function relu(x)
    return max(0, x)
end

function identity_func(x)
    return x
end

function softmax(a)
    c = maximum(a)
    exp_a = exp.(a .- c)  # オーバーフロー対策
    sum_exp_a = sum(exp_a)
    y = exp_a ./ sum_exp_a

    return y
end

function init_network()
    network = Dict(
        "W1" => [0.1 0.3 0.5; 0.2 0.4 0.6],
        "b1" => [0.1 0.2 0.3],
        "W2" => [0.1 0.4; 0.2 0.5; 0.3 0.6],
        "b2" => [0.1 0.2],
        "W3" => [0.1 0.3; 0.2 0.4],
        "b3" => [0.1 0.2]
    )
end

function forward(network, x)
    W1, W2, W3 = network["W1"], network["W2"], network["W3"]
    b1, b2, b3 = network["b1"], network["b2"], network["b3"]

    a1 = x * W1 .+ b1
    z1 = sigmoid.(a1)
    a2 = z1 * W2 .+ b2
    z2 = sigmoid.(a2)
    a3 = z2 * W3 .+ b3
    y = identity_func.(a3)

    return y
end

function main()
    network = init_network()
    x = [1.0 0.5]
    y = forward(network, x)
    println(y)
end

function activationFuncTest()
    x = -5.0:0.1:5.0
    y1 = step_func.(x)
    plot(x, y1)
    y2 = sigmoid.(x)
    plot!(x, y2)

    y3 = relu.(x)
    plot(x, y3)

    a = [0.3 2.9 4.0]
    y4 = softmax(a)
    println(y4)
    println(sum(y4))

end

# activationFuncTest()
main()


import Distributions.MvNormalStats, Distributions.lpgamma, Distributions.suffstats, Distributions.mean

function suffstats{Tf <: FloatingPoint}(D::Type{MvNormal}, x::DenseArray{Tf, 2},
        w::DenseArray{Float64}, z::DenseArray{Tf, 2})
    d = size(x, 1)
    n = size(x, 2)

    tw = sum(w)
    s = zeros(d)
    for i=1:n
        @devec s[:] += x[:, i] .* w[i]
    end
    m = s * inv(tw)
    for i in 1:n
        z_i = view(z, :, i)
        x_i = view(x, :, i)
        @devec z_i[:] = (x_i .- m) .* sqrt(w[i])
    end
    s2 = A_mul_Bt(z, z)

    MvNormalStats(s, m, s2, tw)
end

function posterior_cool(prior::NormalWishart, ss::MvNormalStats)
    if ss.tw < eps() return prior end

    mu0 = prior.mu
    kappa0 = prior.kappa
    TC0 = prior.Tchol
    nu0 = prior.nu

    kappa = kappa0 + ss.tw
    nu = nu0 + ss.tw
    mu = (kappa0.*mu0 + ss.s) ./ kappa

    Lam0 = inv(TC0)
    z = prior.zeromean ? ss.m : ss.m - mu0
    Lam = inv(Lam0 + ss.s2 + kappa0*ss.tw/kappa*(z*z'))

    return NormalWishart(mu, kappa, cholfact(Lam), nu)
end

function expected_logdet(nw::NormalWishart)
    logd = 0

    for i=1:nw.dim
        logd += digamma(0.5 * (nw.nu + 1 - i))
    end

    logd += nw.dim * log(2)
    logd += logdet(nw.Tchol)

    return logd
end

function expected_truncation_level(nw::NormalWishart)
    return nw.Tchol[:U]' * nw.Tchol[:U] * nw.nu
end

function mean(nw::NormalWishart)
    T = expected_truncation_level(nw)
    return MvNormalCanon(T * nw.mu, T)
end

function lognorm(nw::NormalWishart)
    return (nw.dim / 2) * (log(2 * pi) - log(nw.kappa)) + (nw.nu / 2) * logdet(nw.Tchol) + (nw.dim * nw.nu / 2) * log(2.) + lpgamma(nw.dim, nw.nu / 2)
end

function entropy(nw::NormalWishart)
    en = 0.

    logd = expected_logdet(nw)

    en -= logd
    en -= nw.dim * (log(nw.kappa) - log(2*pi))
    en += nw.dim
    en /= 2

    en -= 0.5 * (nw.nu - nw.dim - 1) * logd
    en += nw.nu * nw.dim / 2
    en += 0.5 * nw.nu * logdet(nw.Tchol) + 0.5 * nw.nu * nw.dim * log(2.) + lpgamma(nw.dim, nw.nu/2)

    return en
end

function marginal_loglikelihood(prior::NormalWishart, posterior::NormalWishart, n::Float64)
    d = prior.dim
    return lognorm(posterior) - lognorm(prior) - (n*d/2) * log(2 * pi)
end

function gaussian_mixture{Tf <: FloatingPoint}(prior::NormalWishart, T::Int64, alpha::Float64, x::DenseArray{Tf, 2})
    dim, N = size(x)
    theta = Array(NormalWishart, T)

    buffer = similar(x)

    for k=1:T
        theta[k] = prior
    end

    function cluster_update(k::Int64, z::DenseArray{Float64})
        nw = posterior_cool(prior, suffstats(MvNormal, x, z, buffer))
        theta[k] = nw
    end

    function cluster_entropy(k::Int64)
        return entropy(theta[k])
    end

    function cluster_loglikelihood(k::Int64, z::DenseArray{Float64})
        post = posterior_cool(prior, suffstats(MvNormal, x, z, buffer))
        n = sum(z)

        ll = -lognorm(prior)

        nw = theta[k]
        ll -= nw.dim * n / 2 * log(2*pi)

        ll += (post.nu - post.dim)/2 * expected_logdet(nw)
        W = nw.Tchol[:U]' * nw.Tchol[:U]
        ll -= nw.nu * trace(inv(post.Tchol) * W) / 2

        dm = nw.Tchol[:U] * (nw.mu - post.mu)
        ll -= post.kappa * (nw.dim / nw.kappa + nw.nu * dot(dm, dm)) / 2

        return ll
    end

    function object_loglikelihood(k::Int64, i::Int64)
        nw = theta[k]
        ll = -nw.dim/2 * log(2*pi) + expected_logdet(nw)/2
        dx = nw.Tchol[:U] * (x[:, i] - nw.mu)
        ll -= (nw.dim / nw.kappa + nw.nu * dot(dx, dx)) / 2

        return ll
    end

    mm = TSBPMM(N, T, alpha, cluster_update,
            cluster_loglikelihood,
            object_loglikelihood,
            cluster_entropy; random_init=true)

    function predictive_loglikelihood{Tf <: FloatingPoint}(xt::DenseArray{Tf, 2})
        nt = size(xt, 2)
        logp = zeros(nt)

        π = zeros(T)
        pi!(π, mm)

        pred = map(predictive, theta)

        for i=1:nt
            for k=1:T
                logp[i] += π[k] * pdf(pred[k], x[:, i])
            end
        end
        log!(logp)

        return logp
    end

    return mm, theta, predictive_loglikelihood
end

export gaussian_mixture


"Z score a single cell data set

The data will be centered around +10 to make use of the sparse matrix data compression
Not expressed gene/cell combination will get the 0 value and genes lost in the normalization process will get a -1
"
function Zscore(data::SingleCellExpr)
  function scale(v::SparseVector)
    z = copy(v)
    z[z.nzind] = 10 + z[v.nzind] -mean(z[z.nzind]) / std(z[z.nzind])
    return(z)
  end
  function zScore(x::SparseVector)
    ##n = find(x .>= 0) ## there are no 0 vals ;-)
    if (length(find(x .== -1)) == 0 )
      x =   scale(x)
    else
      ok = find(x .>= 0)
      x[ok] = scale(x[ok])
    end
    return x
  end
  ## now I need to dig through the data in row direction (bad)
  nm = length( nonzeros(data.data) )
  pos = 0
  #define max result values
  v = zeros(nm)
  r = zeros(nm)
  c = zeros(nm)
  cells = size(data.data)[2]
  for i =1:size(data.data)[1]
    val = zScore( data.data[i,1:cells] );
    ## collect the Int64 > 0 values
    ok = (pos+1):(pos+length(val.nzind));
    v[ok] = collect(val.nzval);
    c[ok] = val.nzind;
    r[ok] = fill(i, length(val.nzind));
    pos = pos+length(val.nzind);
  end
  data.zscored = sparse( r[1:pos], c[1:pos] , v[1:pos] )

  return data
end


# Momentum-oriented technical indicator functions

"""
```
aroon(hl::Matrix{T}; n::Int64=25)::Array{Float64}
```

Aroon up/down/oscillator

*Output*

- Column 1: Aroon Up
- Column 2: Aroon Down
- Column 3: Aroon Oscillator
"""
function aroon(hl::AbstractMatrix{T}; n::Int64=25)::Matrix{Float64} where {T<:Real}
    @assert size(hl,2) == 2 "Argument `hl` must have exactly 2 columns."
    @assert n < size(hl,1) "Argument `n` must be less than the number of rows in argument `hl`."
    out = zeros(T, (size(hl,1),3))
    @inbounds for i in n:size(hl,1)
        out[i,1] = 100.0 * (n - findmax(hl[i-n+1:i,1])[2]) / 25.0
    end
    @inbounds for i in n:size(hl,1)
        out[i,2] = 100.0 * (n - findmin(hl[i-n+1:i,2])[2]) / 25.0
    end
    out[:,3] = out[:,1]-out[:,2]  # aroon oscillator
    out[1:n-1,:] .= NaN
    return out
end

"""
```
donch(hl::Matrix{T}; n::Int64=10, inclusive::Bool=true)::Array{Float64}
```

Donchian channel (if inclusive is set to true, will include current bar in calculations.)

*Output*

- Column 1: Lowest low of last `n` periods
- Column 2: Average of highest high and lowest low of last `n` periods
- Column 3: Highest high of last `n` periods
"""
function donch(hl::AbstractMatrix{T}; n::Int64=10, inclusive::Bool=true)::Matrix{Float64} where {T<:Real}
    @assert size(hl,2) == 2 "Argument `hl` must have exactly 2 columns."
    local lower::Array{T} = runmin(hl[:,2], n=n, cumulative=false, inclusive=inclusive)
    local upper::Array{T} = runmax(hl[:,1], n=n, cumulative=false, inclusive=inclusive)
    local middle::Array{T} = (lower .+ upper) ./ 2.0
    return [lower middle upper]
end

"""
```
momentum(x::Array{T}; n::Int64=1)::Array{Float64}
```

Momentum indicator (price now vs price `n` periods back)
"""
function momentum(x::AbstractArray{T}; n::Int64=1)::Array{Float64} where {T<:Real}
    @assert n>0 "Argument n must be positive."
    return diffn(x, n=n)
end

"""
```
roc(x::Array{T}; n::Int64=1)::Array{Float64}
```

Rate of change indicator (percent change between i'th observation and (i-n)'th observation)
"""
function roc(x::AbstractArray{T}; n::Int64=1)::Array{Float64} where {T<:Real}
    @assert n<size(x,1) && n>0 "Argument n out of bounds."
    out = zeros(size(x)) .* NaN
    @inbounds for i = (n+1):size(x,1)
        out[i] = x[i]/x[i-n] - 1.0
    end
    return out * 100.0
end

"""
```
macd(x::Array{T}; nfast::Int64=12, nslow::Int64=26, nsig::Int64=9)::Array{Float64}
```

Moving average convergence-divergence

*Output*

- Column 1: MACD
- Column 2: MACD Signal Line
- Column 3: MACD Histogram
"""
function macd(x::AbstractArray{T}; nfast::Int64=12, nslow::Int64=26, nsig::Int64=9,
    fastMA::Function=ema, slowMA::Function=ema, signalMA::Function=sma)::Matrix{Float64} where {T<:Real}
    out = zeros(T, (length(x),3))
    out[:,1] = fastMA(x, n=nfast) - slowMA(x, n=nslow)
    out[:,2] = signalMA(out[:,1], n=nsig)
    out[:,3] = out[:,1] - out[:,2]
    return out
end

"""
```
rsi(x::Array{T}; n::Int64=14, ma::Function=ema, args...)::Array{Float64}
```

Relative strength index
"""
function rsi(x::AbstractArray{T}; n::Int64=14, ma::Function=ema, args...)::Array{Float64} where {T<:Real}
    @assert n<size(x,1) && n>0 "Argument n is out of bounds."
    N = size(x,1)
    ups = zeros(N)
    dns = zeros(N)
    zro = 0.0
    dx = [NaN; ndims(x) > 1 ? diff(x, dims=1) : diff(x)]
    @inbounds for i=2:N
        if dx[i] > zro
            ups[i] = dx[i]
        elseif dx[i] < zro
            dns[i] = -dx[i]
        end
    end
    rs = [NaN; ma(ups[2:end], n=n; args...) ./ ma(dns[2:end], n=n; args...)]
    return 100.0 .- 100.0 ./ (1.0 .+ rs)
end

"""
```
adx(hlc::Array{T}; n::Int64=14, wilder=true)::Array{Float64}
```

Average directional index

*Output*

- Column 1: DI+
- Column 2: DI-
- Column 3: ADX
"""
function adx(hlc::AbstractArray{T}; n::Int64=14, ma::Function=ema, args...)::Matrix{Float64} where {T<:Real}
    @assert n<size(hlc,1) && n>0 "Argument n is out of bounds."
    if size(hlc,2) != 3
        error("HLC array must have three columns")
    end
    N = size(hlc,1)
    updm = zeros(N)
    dndm = zeros(N)
    updm[1] = dndm[1] = NaN
    @inbounds for i = 2:N
        upmove = hlc[i,1] - hlc[i-1,1]
        dnmove = hlc[i-1,2] - hlc[i,2]
        if upmove > dnmove && upmove > 0.0
            updm[i] = upmove
        elseif dnmove > upmove && dnmove > 0.0
            dndm[i] = dnmove
        end
    end
    dip = [NaN; ma(updm[2:N], n=n; args...)] ./ atr(hlc, n=n) * 100.0
    dim = [NaN; ma(dndm[2:N], n=n; args...)] ./ atr(hlc, n=n) * 100.0
    dmx = abs.(dip-dim) ./ (dip+dim)
    adx = [fill(NaN,n); ma(dmx[n+1:N], n=n; args...)] * 100.0
    return [dip dim adx]
end

"""
```
psar(hl::Array{T}; af_min::T=0.02, af_max::T=0.2, af_inc::T=af_min)::Array{Float64}
```

Parabolic stop and reverse (SAR)

*Arguments*
- `hl`: 2D array of high and low prices in first and second columns respectively
- `af_min`: starting/initial value for acceleration factor
- `af_max`: maximum acceleration factor (accel factor capped at this value)
- `af_inc`: increment to the acceleration factor (speed of increase in accel factor)
"""
function psar(hl::AbstractArray{T}; af_min::T=0.02, af_max::T=0.2, af_inc::T=af_min)::Array{Float64} where {T<:Real}
    @assert af_min<1.0 && af_min>0.0 "Argument af_min must be in [0,1]."
    @assert af_max<1.0 && af_max>0.0 "Argument af_max must be in [0,1]."
    @assert af_inc<1.0 && af_inc>0.0 "Argument af_inc must be in [0,1]."
    @assert size(hl,2) == 2 "Argument hl must have 2 columns."
    ls0 = 1
    ls = 0
    af0 = af_min
    af = 0.0
    ep0 = hl[1,1]
    ep = 0.0
    maxi = 0.0
    mini = 0.0
    sar = zeros(T,size(hl,1))
    sar[1] = hl[1,2] - std(hl[:,1]-hl[:,2])
    @inbounds for i = 2:size(hl,1)
        ls = ls0
        ep = ep0
        af = af0
        mini = min(hl[i-1,2], hl[i,2])
        maxi = max(hl[i-1,1], hl[i,1])
        # Long/short signals and local extrema
        if (ls == 1)
            ls0 = hl[i,2] > sar[i-1] ? 1 : -1
            ep0 = max(maxi, ep)
        else
            ls0 = hl[i,1] < sar[i-1] ? -1 : 1
            ep0 = min(mini, ep)
        end
        # Acceleration vector
        if ls0 == ls  # no signal change
            sar[i] = sar[i-1] + af*(ep-sar[i-1])
            af0 = (af == af_max) ? af_max : (af + af_inc)
            if ls0 == 1  # current long signal
                af0 = (ep0 > ep) ? af0 : af
                sar[i] = min(sar[i], mini)
            else  # current short signal
                af0 = (ep0 < ep) ? af0 : af
                sar[i] = max(sar[i], maxi)
            end
        else  # new signal
            af0 = af_min
            sar[i] = ep0
        end
    end
    return sar
end

"""
```
kst(x::Array{T};
    nroc::Array{Int64}=[10,15,20,30], navg::Array{Int64}=[10,10,10,15],
    wgts::Array{Int64}=collect(1:length(nroc)), ma::Function=sma)::Array{Float64}

```

KST (Know Sure Thing) -- smoothed and summed rates of change
"""
function kst(x::AbstractArray{T}; nroc::AbstractArray{Int64}=[10,15,20,30], navg::AbstractArray{Int64}=[10,10,10,15],
    wgts::AbstractArray{Int64}=collect(1:length(nroc)), ma::Function=sma)::Array{Float64} where {T<:Real}
    @assert length(nroc) == length(navg)
    @assert all(nroc.>0) && all(nroc.<size(x,1))
    @assert all(navg.>0) && all(navg.<size(x,1))
    N = length(x)
    k = length(nroc)
    out = zeros(size(x))
    @inbounds for j = 1:k
        out += ma(roc(x, n=nroc[j]), n=navg[j]) * wgts[j]
    end
    return out
end

"""
```
wpr(hlc::Matrix{T}, n::Int64=14)::Array{Float64}
```

Williams %R
"""
function wpr(hlc::AbstractMatrix{T}; n::Int64=14)::Array{Float64} where {T<:Real}
    hihi = runmax(hlc[:,1], n=n, cumulative=false)
    lolo = runmin(hlc[:,2], n=n, cumulative=false)
    return -100 * (hihi - hlc[:,3]) ./ (hihi - lolo)
end

"""
```
cci(hlc::Matrix{T}; n::Int64=20, c::T=0.015, ma::Function=sma)::Array{Float64}
```

Commodity channel index
"""
function cci(hlc::AbstractMatrix{T}; n::Int64=20, c::T=0.015, ma::Function=sma, args...)::Array{Float64} where {T<:Real}
    tp = (hlc[:,1] + hlc[:,2] + hlc[:,3]) ./ 3.0
    dev = runmad(tp, n=n, cumulative=false, fun=mean)
    avg = ma(tp, n=n; args...)
    return (tp - avg) ./ (c * dev)
end

"""
```
stoch(hlc::Matrix{T}; nK::Int64=14, nD::Int64=3, kind::Symbol=:fast, ma::Function=sma, args...)::Matrix{Float64}
```

Stochastic oscillator (fast or slow)
"""
function stoch(hlc::AbstractMatrix{T}; nK::Int64=14, nD::Int64=3,
    kind::Symbol=:fast, ma::Function=sma, args...)::Matrix{Float64} where {T<:Real}
    @assert kind == :fast || kind == :slow "Argument `kind` must be either :fast or :slow"
    @assert nK<size(hlc,1) && nK>0 "Argument `nK` out of bounds."
    @assert nD<size(hlc,1) && nD>0 "Argument `nD` out of bounds."
    hihi = runmax(hlc[:,1], n=nK, cumulative=false)
    lolo = runmin(hlc[:,2], n=nK, cumulative=false)
    out = zeros(T, (size(hlc,1),2))
    out[:,1] = (hlc[:,3]-lolo) ./ (hihi-lolo) * 100.0
    out[:,2] = ma(out[:,1], n=nD; args...)
    if kind == :slow
        out[:,1] = out[:,2]
        out[:,2] = ma(out[:,1], n=nD; args...)
    end
    return out
end

"""
```
smi(hlc::Matrix{T}; n::Int64=13, nFast::Int64=2, nSlow::Int64=25, nSig::Int64=9,
    maFast::Function=ema, maSlow::Function=ema, maSig::Function=sma)::Matrix{Float64}
```

SMI (stochastic momentum oscillator)
"""
function smi(hlc::AbstractMatrix{T}; n::Int64=13, nFast::Int64=2, nSlow::Int64=25, nSig::Int64=9,
    maFast::Function=ema, maSlow::Function=ema, maSig::Function=sma)::Matrix{Float64} where {T<:Real}
    hihi = runmax(hlc[:,1], n=n, cumulative=false)
    lolo = runmin(hlc[:,2], n=n, cumulative=false)
    hldif = hihi-lolo
    delta = hlc[:,3] - (hihi+lolo) / 2.0
    numer = maSlow(maFast(delta, n=nFast), n=nSlow)
    denom = maSlow(maFast(hldif, n=nFast), n=nSlow) / 2.0
    out = zeros(T, (size(hlc,1),2))
    out[:,1] = 100.0*(numer./denom)
    out[:,2] = maSig(out[:,1], n=nSig)
    return out
end


using ReachabilityAnalysis, SparseArrays, JLD2

LazySets.set_ztol(Float64, 1e-14)

const x25 = [zeros(24); 1.0; zeros(23)]
const x25e = vcat(x25, 0.0)

building_path = joinpath(@__DIR__, "building.jld2")

function building_BLDF01()
    @load building_path A B
    n = size(A, 1)
    U = Interval(0.8, 1.0)
    S = @system(x' = Ax + Bu, u ∈ U, x ∈ Universe(n))

    # 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)

    prob_BLDF01 = InitialValueProblem(S, X0)
end

using ReachabilityAnalysis: add_dimension

function building_BLDC01()
    @load building_path A B
    n = size(A, 1)
    U = Interval(0.8, 1.0)

    # 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)

    Ae = add_dimension(A)
    Ae[1:n, end] = B
    prob_BLDC01 = @ivp(x' = Ae * x, x(0) ∈ X0 × U)
end


##
cd(@__DIR__())
lines = readlines("input.txt")
tm = parse(Int, lines[1])
@show tm
ids = map(split(lines[2], ",")) do x
    x == "x" && return missing
    parse(Int128, x)
end |> skipmissing |> collect

##
ts = (tm .÷ ids .+ 1) .* ids
@show idn = ids[argmin(ts)]
@show wt = minimum(ts) - tm
@show idn * wt

## PARTE 2
lines = readlines("input.txt")
idsts = map(split(lines[2], ",")) do x
    x == "x" && return missing
    parse(Int128, x)
end |> enumerate
idsts = Iterators.filter(idsts) do (i,d)
    !ismissing(d)
end
ts = map(first, idsts) .- 1
ids = map(last, idsts)

##
n = 1068781
function find_tm(ids, ts; start = 1)
    n = start
    seqts = zeros(Int128, length(ids))
    ds = similar(ids)
    ds[1] = 0
    ds[2:end] .= 1
    @. seqts = ((n ÷ ids) + ds) * ids - ts
    while !all(==(seqts[1]), seqts[2:end])
        n += 1
        @. seqts = ((n ÷ ids) + ds) * ids - ts
    end
    n
end
function extended_gcd(a::T, b::T) where T
    old_r::T, r::T = a, b
    old_s::T, s::T = 1, 0
    old_t::T, t::T = 0, 1
    while r != 0
        q = old_r ÷ r
        old_r, r = r, old_r - q * r
        old_s, s = s, old_s - q * s
        old_t, t = t, old_t - q * t
    end
    (old_s, old_t), old_r, (t, s)
end
function cong_x(y::Union{Int128,BigInt}, p::Union{Int128,BigInt}, N::Integer = 1_000_000_000)
    c = zero(typeof(y))
    k = y ÷ p |> abs
    if y < 0
        y += k * p
    end
    while y < 0 && c <= N
        y += p
        c += 1
    end
    c = zero(typeof(y))
    k = y ÷ p |> abs
    if y > p
        y -= k * p
    end
    while y > p && y > 0 && c <= N
        y -= p
        c += 1
    end
    y
end
function find_tm_2(a::T, b, t1, t2, N::Integer = 1_000_000_000) where T
    (s, t), _, _ = extended_gcd(a, b)
    @show (s, t)
    x::T = t1 * t * b + t2 * s * a
    cong_x(x, a * b, N)
end
function find_tm_opt(ids, ts, N::Integer = 1_000_000_000)
    foldl(zip(ids[2:end], ts[2:end]); init=(ids[1],zero(Int128))) do (a, x), (b, t)
        @show (a, x), (b, t)
        a * b, find_tm_2(a, b, x, -t, N)
    end |> last
end

##
k = 9
@time x = find_tm_opt(ids[1:k], ts[1:k], 100_000_000_000)
@show x
@show (x .+ ts[1:k]) .% ids[1:k];


using FastGaussQuadrature

"""
ScaleGaussLegendre(N, a, b)

Uses `gausslegendre` from the `FastGaussQuadrature` package, and scales it to
calculate an integral from `a` to `b` instead of -1 to 1.
"""
function ScaleGaussLegendre(N, a, b)
    x,w = gausslegendre(N)
    mid = (b+a)/2
    L = (b-a)/2

    x = L*x .+ mid
    w = L * w

    return x,w
end

"""
ScaleGaussLaguerre(N, a, scale=1.)

Uses `gausslaguerre` from the `FastGaussQuadrature` package, and scales it to
calculate an integral from `a` to Inf instead of 0 to Inf. The `scale` parameter
allows for a typical length scale to be set.
"""
function ScaleGaussLaguerre(N, a, scale=1.)
    x,w = gausslaguerre(N)
    w = scale * exp.(log.(w) .+ x)

    # This is for the generalised Gauss Laguerre - but it should really be
    # done in the original func... but I don't see how that can be done
    # cleanly in matrix form.
    # x,w = gausslaguerre(N,2.0)
    # w = @. scale * exp(log(w) + x - 2*log(x))

    x = x*scale .+ a

    return x,w
end

"""
    Quadrature(en_target, ...)

This is the workhorse of the code. It figures out how to approximate the integral in the Lippmann-Schwinger equations and evaluates the Greens function too.

The standard grid is to put a Legendre region around the divergence at
k=sqrt(2m*en) of size `sym_size`. This includes a special point at the
divergence if on-shell. If this region does not go down to k=0, another Legendre
grid covers that inner range. Finally, a Laguerre grid covers the larger k
values.

The number of points in each region can be configured with `N` and if different
values are required `N_inner`, `N_mid`, and `N_outer`. The outer grid is
configured with `outer_style`, in which a Legendre grid can be chosen instead.

An alternative style is not yet implemented but includes an additional region
just around the divergence.
"""
function Quadrature(en_target ; N=10,
                    N_inner=N,
                    N_mid=N,
                    N_outer=N,
                    sym_size=0.1, #sym_inner_frac=0.01)
                    sym_inner_frac=nothing,
                    outer_style=(:laguerre, 1.0),
                    mass = 1.,
                    ħ = 1.
                    )
    κ = sqrt(2mass * en_target) / ħ

    if real(κ) < sym_size
        sym_size = real(κ)
    end

    Gfunc(x) = 1/(en_target - ħ^2*x^2/2mass)

    regions = []

    if sym_size == 0
        # This should only happen when the energy is negative.
    else
        if sym_inner_frac == nothing
            isreal(en_target) && @assert iseven(N_mid)

            # a) create one region spanning (κ - sym_size) : (κ + sym_szie)
            x,w = ScaleGaussLegendre(N_mid, real(κ) - sym_size, real(κ) + sym_size)

            w = w .* Gfunc.(x)
            push!(regions, (x,w))
        else
            # b) Create two regions from (κ - sym_size) : (κ - sym_inner_size) and
            # similar above. Then apply the special rule above.
            error("Not implemented")
            sym_inner_size = sym_size * sym_inner_frac
        end

        if isreal(κ)
            # Got to have the divergence always
            push!(regions, ([κ], [DivergenceQuad(κ)]))
        end
    end

    # Then create a Gauss-Laguerre region above and a Gauss-Legendre region below (if sym_size != k_target).
    if sym_size != real(κ)
        a = 0
        b = real(κ) - sym_size
        x,w = ScaleGaussLegendre(N_inner, a, b)
        w = w .* Gfunc.(x)
        lower = (x,w)

        pushfirst!(regions, lower)
    end

    if outer_style[1] == :laguerre
        laguerre_scale = outer_style[2]
        x,w = ScaleGaussLaguerre(N_outer, real(κ) + sym_size, laguerre_scale)
        w = w .* Gfunc.(x)
        outer = (x,w)
        push!(regions, outer)
    elseif outer_style[1] == :legendre
        legendre_max = outer_style[2]
        x,w = ScaleGaussLegendre(N_outer, real(κ) + sym_size, real(κ) + sym_size + legendre_max)
        w = w .* Gfunc.(x)
        outer = (x,w)
        push!(regions, outer)
    else
        error("Unknown outer_style $(outer_style[1])")
    end

    # x = vcat(first.(regions)...)
    # u = vcat(last.(regions)...)
    x,u = reduce(regions) do r1,r2
        x1,u1 = r1
        x2,u2 = r2
        x = union(x1,x2)
        u = zeros(ComplexF64, size(x))

        u[indexin(x1,x)] .+= u1
        u[indexin(x2,x)] .+= u2

        (x,u)
    end

    x = Vector{Float64}(x)
    inds = sortperm(x)
    x = x[inds]
    u = u[inds]
    return x,u,regions
end

function DivergenceQuad(qi)
    mass = ħ = 1.
    u = im*pi * 1/(ħ^2/mass * qi)
end

"""
    TaylorApprox(k)

This evaluates the divergent part of the integral, using a Taylor series expansion around that point.
"""
function TaylorApprox(k3)
    @assert length(k3) == 3
    @assert diff(k3) |> x -> (x[1] == x[2])
    qi = k3[2]
    dk = diff(k3)[1]

    @info "Including approximation for the divergent part"
    mass = ħ = 1.
    prefac = 2*mass/ħ^2 * dk/qi
    # f = k_list'.^2 .* Vmat

    u = zeros(3)
    u[:,3] = -prefac/2dk
    u[:,1] = prefac/2dk
    u[:,2] = prefac/2qi

    return u
end

# This was for a test calculating the K matrix instead of the T matrix.
function KMatQuadrature(en_target ; N=10,
                        N_inner=N,
                        N_mid=N,
                        N_outer=N,
                        sym_size=0.1, #sym_inner_frac=0.01)
                        sym_inner_frac=nothing,
                        laguerre_scale=1.0)
    if isreal(en_target)
        # Need to make sure there is a symmetric region in here.

        mass = ħ = 1.
        k_target = sqrt(2mass * en_target) / ħ

        if k_target < sym_size
            sym_size = k_target
        end

        function Gfunc(x)
            denom = en_target - ħ^2*x/2mass
            denom == 0. ? 0. : 1/denom
        end

        regions = []

        # Allow for two options
        if sym_inner_frac == nothing
            @assert isodd(N_mid)
            # a) create one region spanning (k_target - sym_size) : (k_target + sym_szie)
            x,w = ScaleGaussLegendre(N_mid, k_target - sym_size, k_target + sym_size)

            w .*= Gfunc.(x)
            push!(regions, (x,w))
        else
            # b) Create two regions from (k_target - sym_size) : (k_target - sym_inner_size) and
            # similar above. Then apply the special rule above.
            error("Not implemented")
            sym_inner_size = sym_size * sym_inner_frac
        end

        # Then create a Gauss-Laguerre region above and a Gauss-Legendre region below (if sym_size != k_target).
        if sym_size != k_target
            a = 0
            b = k_target - sym_size
            x,w = ScaleGaussLegendre(N_inner, a, b)
            w .*= Gfunc.(x)
            lower = (x,w)

            pushfirst!(regions, lower)
        end

        x,w = ScaleGaussLaguerre(N_outer, k_target + sym_size, laguerre_scale)
        w .*= Gfunc.(x)
        outer = (x,w)
        push!(regions, outer)

        # x = vcat(first.(regions)...)
        # u = vcat(last.(regions)...)
        x,u = reduce(regions) do r1,r2
            x1,u1 = r1
            x2,u2 = r2
            x = union(x1,x2)
            u = zeros(size(x))

            u[indexin(x1,x)] .+= u1
            u[indexin(x2,x)] .+= u2

            (x,u)
        end

        inds = sortperm(x)
        x = x[inds]
        u = u[inds]
        return x,u,regions
    else
        error("Not implemented")

        # Maybe do a GaussLegendre region around the real value of k_target, but
        # down to zero. Then throw in a Laguerre region for the outer.

    end
end


s = open("input.txt") do file
    read(file, String);
end

function ParseInput(s)
    startStateFlat = map(x->x=='#', replace(s,"\n"=>"") |> collect);
    dim = length(split(s,"\n")[1])
    startState = reshape(startStateFlat, (dim,dim));

    simSpace = falses(21, 21, 21, 21);

    halfDim = dim ÷ 2;

    simSpace[11-halfDim:11+halfDim - ifelse(dim%2==0,1,0),
        11-halfDim:11+halfDim - ifelse(dim%2==0,1,0),
        11, 11] .= startState;

    return simSpace;
end

function numNeighbors(space, loc)
    count = 0;
    for i = -1:1
        for j = -1:1
            for k = -1:1
                for l = -1:1
                    tempLoc = loc .+ [i,j,k,l];
                    if all(tempLoc .== loc)
                        continue;
                    end

                    if all(tempLoc.>0) && all(tempLoc.<22)
                        if space[tempLoc...]
                            count += 1;
                        end
                    end
                end
            end
        end
    end

    return count;
end

function BootUp(simSpace)
    nextSpace = falses(21,21,21,21);
    for step = 1:6
        for k = 1:21
            for i = 1:21
                for j = 1:21
                    for l = 1:21
                        curLoc = [i, j, k, l];
                        nCount = numNeighbors(simSpace, curLoc);

                        active = simSpace[curLoc...];
                        if active
                            # println(nCount);
                            if 1 < nCount < 4
                                nextSpace[curLoc...] = true;
                            else
                                nextSpace[curLoc...] = false;
                            end
                        else
                            if nCount == 3
                                nextSpace[curLoc...] = true;
                            else
                                nextSpace[curLoc...] = false;
                            end
                        end
                    end
                end
            end

            # println(k);
            # println(simSpace[:,:,k])
        end
        println("--------");
        simSpace .= nextSpace;
    end

    totalActive = sum(simSpace);
    println(totalActive);

    return totalActive;
end

simSpace = ParseInput(s);
# simSpace = ParseInput(".#.\n..#\n###");
answer = BootUp(simSpace)


using Plots;
pyplot();
using LinearAlgebra

function angular_position(points)
    cpoints = points[1] + im * points[2]
    return rad2deg(angle(cpoints))
end

function perturbed_circle(point, meanradius, amplitude, frequency)
    angle = angular_position(point)
    interfaceposition = meanradius + amplitude * sind(angle * frequency)
    distance = interfaceposition - norm(point)
    return distance
end

xrange = -1:1e-1:1
meanradius = 0.5
amplitude = 0.1
frequency = 5

plot(
    xrange,
    xrange,
    (x,y) -> -perturbed_circle([x,y], meanradius, 0., frequency),
    st = :surface,
    #camera = (-30, 30),
)
plot!(xrange,xrange,(x,y)->0.2,st=:surface)


using AbstractGPs
using ConjugateComputationVI
using Distributions
using Optim
using ParameterHandling
using Plots
using Random
using StatsBase
using StatsFuns
using TemporalGPs
using Zygote

using ConjugateComputationVI: GaussHermiteQuadrature, UnivariateFactorisedLikelihood

# Adjoint for the Poisson logpdf.
# Someone should implement this in the correct place.
# log(λ^x exp(-λ) / x!) =
# x log(λ) - λ - log(x!)
# dλ = x / λ - 1
Zygote.@adjoint function StatsFuns.poislogpdf(λ::Float64, x::Union{Float64, Int})
    function poislogpdf_pullback(Δ::Real)
        return Δ * (x / λ - 1), nothing
    end
    return StatsFuns.poislogpdf(λ, x), poislogpdf_pullback
end

function Zygote._pullback(ctx::Zygote.AContext, ::typeof(logpdf), d::Poisson, x)
    out, pb = Zygote._pullback(ctx, poislogpdf, d.λ, x)
    function logpdf_pullback(Δ::Real)
        _, Δλ, Δx = pb(Δ)
        return nothing, (λ=Δλ, ), Δx
    end
    return out, logpdf_pullback
end

# Specify a model.
θ_init = (scale=positive(1.0), stretch=positive(1.0));
θ_init_flat, unflatten = ParameterHandling.flatten(θ_init);

function build_latent_gp(θ::AbstractVector{<:Real})
    return build_latent_gp(ParameterHandling.value(unflatten(θ)))
end
function build_latent_gp(θ::NamedTuple)
    gp_naive = GP(θ.scale * AbstractGPs.transform(Matern52Kernel(), θ.stretch))
    gp = to_sde(gp_naive, SArrayStorage(Float64))
    lik = UnivariateFactorisedLikelihood(f -> Poisson{Float64}(exp(f)))
    return LatentGP(gp, lik, 1e-9)
end

# Generate synthetic data.
x = RegularSpacing(0.0, 0.001, 10_000);
y = rand(build_latent_gp(θ_init_flat)(x)).y;

f_approx_post, results_summary = ConjugateComputationVI.optimize_elbo(
    build_latent_gp,
    GaussHermiteQuadrature(10),
    x,
    y,
    θ_init_flat,
    BFGS(
        alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
        linesearch = Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(
        show_trace = true,
        iterations=25,
        f_calls_limit=50,
    ),
);

function approx_post_marginal_samples(x::AbstractVector, N::Int)
    return map(marginals(f_approx_post(x))) do latent_marginal
        f = rand(latent_marginal, N)
        return exp.(f)
    end
end

function approx_post_95_CI(x::AbstractVector, N::Int)
    samples = approx_post_marginal_samples(x, N)
    return quantile.(samples, Ref((0.025, 0.5, 0.975)))
end

# Plot the predictions.
qs = approx_post_95_CI(x, 100);
p1 = plot(
    x, getindex.(qs, 1);
    linealpha=0,
    fillrange=getindex.(qs, 3),
    label="95% CI",
    fillalpha=0.3,
    color=:blue,
);
scatter!(p1, x, y; markersize=1, label="Observations");

p2 = plot(f_approx_post(x, 1e-6); ribbon_scale=3, color=:blue, label="approx post latent");

plot(p1, p2; layout=(2, 1))


using MPI, Knet, ParticleSwarmOptimization
include("Args.jl")

function main()
    MPI.Init()
    args = parse_commandline()
    gpu(args["gpu"])
    comm = MPI.COMM_WORLD
    rank = MPI.Comm_rank(comm)
    world = MPI.Comm_size(comm)

    if rank == 0
        println("Parsed args:")

        for (arg,val) in args
            println("  $arg  =>  $val")
        end
    end

    x = [0.0 0.0 1.0 1.1; 0.0 1.0 0.0 1.1]
    y = [0.0 1.0 1.0 0.0]

    @knet function mlp(x; winit=Gaussian(0,.1))
        h1 = wbf(x; out=100, f=:sigm, winit=winit)
        h2 = wbf(h1; out=100, f=:sigm, winit=winit)
        return wbf(h2; out=1, f=:sigm, winit=winit)
    end

    problem = NNProblem(:mlp, args["l"], args["h"], Single(x, y), quadloss)

    #last parameter for the parallelization
    gbest = pso(problem, args["num"], args["epoch"], args["phi1"], args["phi2"], args["w"], args["vmin"], args["vmax"], true)

    if rank == 0
        println("\nBest: $(gbest.pbest_fitness)\n")
        println("Gold: $(y)")
        println("Pred: $(forw(gbest.x, x))")
    end

    MPI.Finalize()
end

main()


# operators imported and overridden from Base
import Base: +, -, *, /, ^, |, %, ==, !=, <, >, <<, >>, ~

# functions imported and overridden from Base
# convention: prefix `Base.` should always be present when overriding
# so comment the following will still work
import Base: show,convert, getproperty, propertynames, setproperty!
if isdefined(Base, :hasproperty) # Julia 1.2
    import Base: hasproperty
end
import Base: abs, inv, adjoint, exp, getindex

export Mv

export norm, rev, dual, involute, proj, refl, rot, exp_with_hint, scalar, even, odd
# export \cdot, \wedge, \intprod, \intprodr, \odot, \boxtimes, \circledast, \times
export +, -, *, /, ^, |, %, ==, !=, <, >, <<, >>, ~, ⋅, ∧, ⨼, ⨽, ⊙, ⊠, ⊛, ×
# Operator precedence: they have the same precedence, unlike in math
# julia> for op ∈ [:* :⋅ :∧ :⨼ :⨽ :⊙ :⊠ :⊛ :×]; println(String(op), "  ", Base.operator_precedence(op)) end
# *  13
# ⋅  13
# ∧  13
# ⨼  13
# ⨽  13
# ⊙  13
# ⊠  13
# ⊛  13
# ×  13

# experimental export \bar\times
# export ×̄

# \^-\^1
@define_postfix_symbol(⁻¹)
# # \^T
# @define_postfix_symbol(ᵀ)
# \^x
@define_postfix_symbol(ˣ)
# \doublepipe
@define_postfix_symbol(ǂ)
# \_+
@define_postfix_symbol(₊)
# \_-
@define_postfix_symbol(₋)
# \bot
# @define_postfix_symbol(⊥)

@doc raw"""
A wrapper class for `galgebra.mv.Mv`:

- all methods of `galgebra.mv.Mv` are delegated and can be called like in Python.
- enhanced with operator overriding and extra methods.
- automatically supports pretty printing like in Python.
"""
mutable struct Mv
    o::PyCall.PyObject
end

@define_show(Mv)
@delegate_properties(Mv, :o)
@delegate_doc(Mv)

"Addition."
@define_op(Mv, +, __add__)

"Subtraction."
@define_op(Mv, -, __sub__)

@doc raw"""
Geometric product.

``A * B \equiv A B``.
"""
@define_op(Mv, *, __mul__)

@doc raw"""
Division.

``A / B \equiv A B^{-1}``. Only valid when ``B`` has inverse.
"""
@define_op(Mv, /, __truediv__)

@doc raw"""
Comparisons of equality.

``A = B \equiv \mathrm{simplify}(A - B) = 0``
"""
@define_op(Mv, ==, __eq__)

@doc raw"""
Comparisons of inequality.

Hint: Type ≠ with `\neq`. Alternatively, use `!=`.

``A \neq B \equiv \mathrm{simplify}(A - B) \neq 0``
"""
@define_op(Mv, ≠, __ne__)
@define_op(Mv, !=, __ne__)

@doc raw"""
Wedge product.

Hint: type ∧ with `\wedge`.
"""
@define_op(Mv, ∧, __xor__)

@doc raw"""
Hestenes' inner product.

Hint: type ⋅ with `\cdot`. Alternatively, use `|`.

``\boldsymbol{A}_{r} \cdot \boldsymbol{B}_{s} \equiv \left\{\begin{array}{lr}{r \text { and } s \neq 0 :} & {\left\langle\boldsymbol{A}_{r} \boldsymbol{B}_{s}\right\rangle_{|r-s|}} \\ {r \text { or } s=0 :} & {0}\end{array}\right.``
"""
@define_op(Mv, ⋅, __or__)
@define_op(Mv, |, __or__)

@doc raw"""Left contraction, i.e. "contraction onto".

Hint: type ⨼ with `\intprod`. Alternatively, use `<`.

``A \rfloor B \equiv \sum\limits_{r, s}\left\langle\langle A\rangle_{r}\langle B\rangle_{s}\right\rangle_{s-r}``

In literature the notation is usually ``A \rfloor B``, but `\rfloor` is reserved by Julia.
"""
@define_op(Mv, ⨼, __lt__)
@define_op(Mv, <, __lt__)

@doc raw"""Right contraction, i.e. "contraction by".

Hint: type ⨽ with `\intprodr`. Alternatively, use `>`.

``A \lfloor B \equiv \sum\limits_{r, s}\left\langle\langle A\rangle_{r}\langle B\rangle_{s}\right\rangle_{r-s}``

In literature the notation is usually ``A \lfloor B``, but `\lfloor` is reserved by Julia.
"""
@define_op(Mv, ⨽, __gt__)
@define_op(Mv, >, __gt__)

@doc raw"""Commutator product.

Hint: type ⊠ with `\boxtimes`. Alternatively, use `>>`.

``A \underline{\times} B \equiv \dfrac{1}{2}(AB-BA)``.
"""
@define_op(Mv, ⊠, __rshift__)
@define_op(Mv, >>, __rshift__)

@doc raw"""Anti-commutator product.

Hint: type ⊙ with `\odot`. Alternatively, use `<<`.

``A \bar{\times} B \equiv \dfrac{1}{2}(AB+BA)``.
"""
@define_op(Mv, ⊙, __lshift__)
@define_op(Mv, <<, __lshift__)

# # experimental symbol for anti-comutator product: \bar\times
# A×̄B = (AB+BA)/2
# @define_op(Mv, ×̄, __lshift__)

@doc raw"""
Cross product for vectors in 3D.
"""
@define_op_with_impl(Mv, ×, mv.cross(x, y))

@doc raw"""Scalar product.

Hint: type ⊛ with `\circledast`. Alternatively, use `%`.

``A \circledast B \equiv \langle A B^{\dagger} \rangle``.

In literature the notation is usually ``\ast`` , but it's visually indistinguishable from `*`.
"""
@define_op_with_impl(Mv, ⊛, (x * ~y).scalar())
@define_op_with_impl(Mv, %, x ⊛ y)

"Unary negation."
@define_unary_op(Mv, -, __neg__)

@doc raw"""
Norm.

`norm(A)` = `A.norm()` ``\equiv \left\lVert A \right\rVert \equiv \sqrt{A \tilde{A}}``

Alternatively:

- `A.norm(hint="+")` ``\equiv \sqrt{A \tilde{A}}``
- `A.norm(hint="-")` ``\equiv \sqrt{- A \tilde{A}}``
- `A.norm(hint="0")` ``\equiv \sqrt{\left| A \tilde{A} \right|}``

Only valid when the result is a scalar.
"""
@define_unary_op(Mv, norm, norm)
@define_unary_op(Mv, Base.abs, norm)

@doc raw"""
Inverse.

`(A)⁻¹ = A^-1 = inv(A) = A.inv()` ``\equiv A^{-1}``

Hint: type ⁻¹ with `\^-\^1`.
"""
@define_unary_op(Mv, Base.inv, inv)
@define_postfix_op(Mv, ⁻¹, Base.inv)

@doc raw"""
Reversion.

`~A = A[:~] = rev(A) = A.rev()` ``\equiv \tilde{A} \equiv A^{\dagger}``

In literature the notation is usually ``\tilde{A}`` or ``A^{\dagger}``, the former is illegal syntax and `\dagger` in the latter is is reserved by Julia.
"""
@define_unary_op(Mv, ~, rev)
@define_unary_op(Mv, rev, rev)

# @deprecated
# @define_postfix_op(Mv, ᵀ, rev)

@doc raw"""
Dual, i.e. orthogonal complement, ``\Lambda^p \to \Lambda^{n-p}``.

`A'` ``\equiv A^{\bot} \equiv A I``

Note: call `Ga.dual_mode(mode)` to globally specify a different dual mode (`I+` is the default):

| dual_mode | ``A^{\bot}`` |
|-----------|--------------|
|   `+I`    |     ``IA``   |
|   `-I`    |    ``-IA``   |
|   `I+`    |     ``AI``   |
|   `I-`    |    ``-AI``   |
|  `+Iinv`  |  ``I^{-1}A`` |
|  `-Iinv`  | ``-I^{-1}A`` |
|  `Iinv+`  |  ``AI^{-1}`` |
|  `Iinv-`  | ``-AI^{-1}`` |
"""
@define_unary_op(Mv, Base.adjoint, dual)
@define_unary_op(Mv, dual, dual)
# @define_postfix_op(Mv, ⊥, dual)

@doc raw"""
Grade involution.

`(A)ˣ = A[:*] = involute(A)` ``\equiv A_+ - A_- \equiv`` `A.even() - A.odd()`

Hint: type ˣ with `\^x`.

In literature the notation is usually ``A^{*}``.
"""
@define_unary_op_with_impl(Mv, involute, x.even() - x.odd())
@define_postfix_op(Mv, ˣ, involute)

@doc raw"""
Clifford conjugate.

`(A)ǂ = A[:ǂ]` ``\equiv A^{*\dagger}``

Hint: type ǂ with `\doublepipe`.

In literature the notation is usually ``A^{\ddagger}``, but `\ddagger` is reserved by Julia.
"""
@define_unary_op_with_impl(Mv, Base.conj, involute(x).rev())
@define_postfix_op(Mv, ǂ, Base.conj)

@doc raw"""
Projection.

`proj(B, A)` `` \equiv P_{B}(A) \equiv`` `A.project_in_blade(B)`

Only valid if B is a blade.
"""
@define_op_with_impl(Mv, proj, mv.proj(x, y))

@doc raw"""
Reflection.

`refl(B, A)` `` \equiv \mathrm{Refl}_{B}(A) \equiv`` `A.reflect_in_blade(B)`

Only valid if B is a blade.
"""
@define_op_with_impl(Mv, refl, mv.refl(x, y))

@doc raw"""
Rotation.

Rotate the multivector `A` by the 2-blade `itheta`.

`rot(itheta, A)` ``\equiv A e^{I \theta} \equiv`` `A.rotate_multivector(itheta)`
"""
@define_op_with_impl(Mv, rot, rot_with_hint(x, y))
@pure rot_with_hint(itheta::Mv, A::Mv, hint::AbstractString="-") = mv.rot(itheta, A, hint)

@doc raw"Natural base exponential of X: ``e^X``"
@define_unary_op_with_impl(Mv, Base.exp, exp_with_hint(x))
@pure exp_with_hint(x::Mv, hint::AbstractString="-") = mv.exp(x, hint)

@doc raw"""
The `i`-th grade part.

`A[i] = A.grade(i)` ``\equiv \langle A B^{\dagger} \rangle_i``
"""
@pure Base.getindex(x::Mv, i::Integer) = x.grade(i)
@pure function Base.getindex(x::Mv, sym::Symbol)
    if sym == :+
        even(x)
    elseif sym == :-
        odd(x)
    elseif sym == :~
        rev(x)
    elseif sym == :⁻¹
        inv(x)
    elseif sym == :*
        involute(x)
    elseif sym == :ǂ
        conj(x)
    else
        throw(DomainError(x, "argument can only be one of :+, :-, :~, :⁻¹, :*, :ǂ"))
    end
end

@doc raw"""
Scalar (grade-0) part.

`scalar(A) = A.scalar()` ``\equiv \langle A B^{\dagger} \rangle \equiv \langle A B^{\dagger} \rangle_0``

Note: it returns a SymPy expression unlike A[0] which returns a Mv object
"""
@define_unary_op(Mv, scalar, scalar)

@doc raw"""
Even-grade part.

`A[:+] = (A)₊ = even(A) = A.even()` ``\equiv A_+``
"""
@define_unary_op(Mv, even, even)
@define_postfix_op(Mv, ₊, even)

@doc raw"""
Odd-grade part.

`A[:-] = (A)₋ = odd(A) = A.odd()` ``\equiv A_-``
"""
@define_unary_op(Mv, odd, odd)
@define_postfix_op(Mv, ₋, odd)

@define_lop(Mv, Number, +, __add__)
@define_rop(Mv, Number, +, __radd__)
@define_lop(Mv, Number, -, __sub__)
@define_rop(Mv, Number, -, __rsub__)
@define_lop(Mv, Number, *, __mul__)
@define_rop(Mv, Number, *, __rmul__)
@define_lop(Mv, Number, /, __truediv__)
@define_rop(Mv, Number, /, __rdiv__)
@define_lop(Mv, Number, ==, __eq__)
@define_lop(Mv, Number, !=, __ne__)

@pure function ^(x::Mv, y::Integer)
    if y < 0
        x.__pow__(abs(y)).inv()
    elseif y == 0
        1
    else
        x.__pow__(y)
    end
end


using Base.Test
using RandomFerns

features = reshape([0:0.01:10],1001,1);
labels = vec(sin(features) + (randn(size(features))/2));

println("\n ### k fold CV - Random Fern Regression ###")
r2 = kfoldCV_fern(labels, features, 10, 5)
@test mean(r2) > 0.3

println("\n ### k fold CV - Random Fern Ensemble Regression ###")
r2 = kfoldCV_fernEnsemble(labels, features, 10, 5, 100)
@test mean(r2) > 0.5


function createdgnodes(p,t,f,curvedboundary,curvedboundaryexpr,porder)

nve,ne = size(t);
nd=size(p,1);

elemtype = 0;
if (nd==2) && (nve==4)
    elemtype=1;
end
if (nd==3) && (nve==8)
    elemtype=1;
end

plocal,~,~,~,perm = masternodes(porder,nd,elemtype);
npl = size(plocal,1);

philocal = zeros(Float64, npl, nve);

if nd==1
    xi  = plocal[:,1];
    philocal[:,1] = 1.0 .- xi;
    philocal[:,2] = xi;
elseif nd==2 && nve==3 # tri
    xi  = plocal[:,1];
    eta = plocal[:,2];
    philocal[:,1] = 1.0 .- xi .- eta;
    philocal[:,2] = xi;
    philocal[:,3] = eta;
elseif nd==2 && nve==4 # quad
    xi  = plocal[:,1];
    eta = plocal[:,2];
    philocal[:,1] = (1.0 .- xi).*(1.0 .- eta);
    philocal[:,2] = xi.*(1.0 .- eta);
    philocal[:,3] = xi.*eta;
    philocal[:,4] = (1.0 .- xi).*eta;
elseif nd==3 && nve==4 # tet
    xi   = plocal[:,1];
    eta  = plocal[:,2];
    zeta = plocal[:,3];
    philocal[:,1] = 1.0 .- xi .- eta .- zeta;
    philocal[:,2] = xi;
    philocal[:,3] = eta;
    philocal[:,4] = zeta;
elseif nd==3 && nve==8 # hex
    xi   = plocal[:,1];
    eta  = plocal[:,2];
    zeta = plocal[:,3];
    philocal[:,1] = (1.0 .-xi).*(1.0 .-eta).*(1.0 .-zeta);
    philocal[:,2] = xi.*(1.0 .-eta).*(1.0 .-zeta);
    philocal[:,3] = xi.*eta.*(1.0 .-zeta);
    philocal[:,4] = (1.0 .-xi).*eta.*(1.0 .-zeta);
    philocal[:,5] = (1.0 .-xi).*(1.0 .-eta).*(zeta);
    philocal[:,6] = xi.*(1.0 .-eta).*(zeta);
    philocal[:,7] = xi.*eta.*(zeta);
    philocal[:,8] = (1.0 .-xi).*eta.*(zeta);
end

# Allocate nodes
dgnodes = zeros(Float64,npl,ne,nd);
for dim=1:nd
  for node=1:nve
    dp = reshape(philocal[:,node],npl,1)*reshape(p[dim,t[node,:]],1,ne);
    dgnodes[:,:,dim] = dgnodes[:,:,dim] + reshape(dp,npl,ne);
  end
end

dgnodes = permutedims(dgnodes,[1, 3, 2]);

if length(curvedboundaryexpr)>0 && porder>1 && maximum(curvedboundary[:])>0
    print("Project dgnodes onto the curved boundaries...\n");
    fd = curvedboundaryexpr;
    nd = size(p,1);
    nfe,ne = size(f);
    if nd == 2
        for i = 1:ne
            for j = 1:nfe
                if f[j,i]!=0 # boundary element
                    k = abs(f[j,i]); # get boundary index
                    if curvedboundary[k]==1 # if this boundary is curved
                        p = dgnodes[perm[:,j],:,i];
                        deps = 1e-8*maximum(maximum(p,dims=1)-minimum(p,dims=1));
                        d = fd[k](p');
                        dgradx = (fd[k]([p[:,1].+deps p[:,2]]') - d)/deps;
                        dgrady = (fd[k]([p[:,1] p[:,2].+deps]') - d)/deps;
                        dgrad2 = dgradx.*dgradx + dgrady.*dgrady;
                        dgrad2[dgrad2.==0] .= 1;
                        p = p - [d.*dgradx./dgrad2 d.*dgrady./dgrad2];
                        dgnodes[perm[:,j],:,i] = p;
                    end
                end
            end
        end
    elseif nd==3
        for i = 1:ne
            for j = 1:nfe
                if f[j,i]!=0 # boundary element
                    k = abs(f[j,i]); # get boundary index
                    if curvedboundary[k]==1 # if this boundary is curved
                        p = dgnodes[perm[:,j],:,i];
                        deps = 1e-8*maximum(maximum(p,dim=1)-minimum(p,dim=1));
                        d = fd[k](p');
                        dgradx = (fd[k]([p[:,1].+deps p[:,2] p[:,3]]')-d)/deps;
                        dgrady = (fd[k]([p[:,1] p[:,2].+deps p[:,3]]')-d)/deps;
                        dgradz = (fd[k]([p[:,1] p[:,2] p[:,3].+deps]')-d)/deps;
                        dgrad2 = dgradx.*dgradx+dgrady.*dgrady+dgradz.*dgradz;
                        dgrad2[dgrad2.==0] .= 1;
                        p = p-[d.*dgradx./dgrad2 d.*dgrady./dgrad2 d.*dgradz./dgrad2];
                        dgnodes[perm[:,j],:,i] = p;
                    end
                end
            end
        end
    end
end

return dgnodes

end


module TestGraphicalLinearAlgebra
using Test

using Random
using IterativeSolvers

using Catlab, Catlab.Theories, Catlab.WiringDiagrams, Catlab.Programs
using Catlab.LinearAlgebra.GraphicalLinearAlgebra
import LinearAlgebra: norm, svd

import LinearMaps: BlockDiagonalMap, UniformScaling
using LinearOperators

# Theories
###########

A, B = Ob(FreeLinearFunctions, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, A, B)

# Domains and codomains
@test dom(plus(f,g)) == A
@test codom(plus(f,g)) == B

# Unicode syntax
@test f+g == plus(f,g)

# Evaluation
############

# LinearMaps instance
#--------------------

f, g = LinearMap([0 1; 1 0]), LinearMap([1 2 3; 4 5 6; 7 8 9])
h, k = LinearMap([1 -1; -1 1; 0 1]), LinearMap([1 1; 2 2; 3 3])

x, y = [2, 1], [7, 3, 5]
@test (f⋅f)*x == f*(f*x)
@test (f⋅h)*x == h*(f*x)
@test (f⊕g)*[x;y] == [f*x; g*y]
@test swap(dom(f),dom(g)) * [x;y] == [y;x]

@test mcopy(dom(f))*x == [x;x]
@test delete(dom(f))*x == []
@test plus(dom(f))*[x;x] == 2x
@test zero(dom(f))*Float64[] == zero(x)

@test (h+k)*x == h*x + k*x
@test scalar(dom(f),3)*x == 3*x
@test antipode(dom(f))*x == -1*x
@test adjoint(g)*y == g'*y

# GMRES with LinearMaps
#----------------------

V, W = Ob(FreeLinearFunctions, :V, :W)
f, g, h = Hom(:f, V, W), Hom(:g, V, W), Hom(:h, W, W)

M = plus(f,g)⋅h
# M = plus(f,g)

d = Dict(f=>LinearMap([1 1 1 1; 2 2 2 3; 3 3 3 3.0]),
         g=>LinearMap([3 1 1 1; -1 -1 -1 -1; 0 2 0 0.0]),
         h=>LinearMap([1 4 -1; -1 1 0; 1 1 -1.0]),
         V=>LinearMapDom(4),
         W=>LinearMapDom(3))
F(ex) = functor((LinearMapDom, LinearMap), ex, generators=d)

Random.seed!(1) # Set seed to avoid sporadic test failures.
n = 4
m = 3
x = ones(Float64, n)
A = F(M)'*F(M)
b = A*x
x̂ᴰ = Matrix(A)\b

@test IterativeSolvers.zerox(F(M), b) == IterativeSolvers.zerox(Matrix(F(M)), b)
x̂ = gmres(A, b; maxiter=20)

b̂ = A*x̂
@test norm(b̂-b) < 1e-8

n = 32
M = plus(f,g)⋅h
M₄ = foldl(⊕, [M, M, M, M, M, M, M, M])
x = ones(Float64, n)
# A needs to be square and nonsingular
A = ( F(M₄)'*F(M₄) ) + UniformScaling(n)
b = A*x
x̂ᴰ = Matrix(A)\b

x̂, h = gmres(A, b; tol=1e-12, maxiter=160, log=true, verbose=false)

b̂ = A*x̂
@test norm(b̂-b) < 1e-8
@test norm(x̂-x)/norm(x) < 1e-8

Σ = svdl(A; k=10)[1]

# Eigensystem tests

λ, Q = powm(F(M)'*F(M))
@test norm((F(M)'*F(M))*Q - λ*Q) < 1e-3
λ, Q = powm(F(M⋅adjoint(M)))
@test norm(F(M⋅adjoint(M))*Q - λ*Q) < 1e-3

Σ, L = svdl(F(M))
@test all(Σ .>= 0)

λ, Q = powm(A)
@test norm(A*Q - λ*Q) < 1e-8

Σ, L = svdl(A)
@test all(Σ .>= 0)

M′ = ( M⊕id(V) ) + ( id(V)⊕M )
M₂ = ( M′⋅adjoint(M′) ) + id(V⊕V)
@test all(svd(Matrix(F(M₂))).S .> 1e-10)
b = ones(size(F(M₂), 1))
x̂ = gmres(F(M₂), b)
b̂ = F(M₂)*x̂
@test norm(b̂ .- 1) < 1e-12

#x₂ = vcat(x,-x)
#b₂= F(M₂)*x₂
#x̂₂ = gmres(F(M₂), b₂)
#b̂₂ = F(M₂)*x̂₂
#@test norm(b̂₂-b₂) < 1e-4
#@test norm(x̂₂-x₂) < 1e-4

# LinearOps instance
#-------------------

f, g = LinearOperator([0 1; 1 0]), LinearOperator([1 2 3; 4 5 6; 7 8 9])
h, k = LinearOperator([1 -1; -1 1; 0 1]), LinearOperator([1 1; 2 2; 3 3])

x, y = [2, 1], [7, 3, 5]
@test (f⋅f)*x == f*(f*x)
@test (f⋅h)*x == h*(f*x)
@test (f⊕g)*[x;y] == [f*x; g*y]
@test swap(dom(f),dom(g)) * [x;y] == [y;x]

@test mcopy(dom(f))*x == [x;x]
@test delete(dom(f))*x == []
@test plus(dom(f))*[x;x] == 2x
@test zero(dom(f))*Float64[] == zero(x)

@test (h+k)*x == h*x + k*x
@test scalar(dom(f),3)*x == 3*x
@test antipode(dom(f))*x == -1*x
@test adjoint(g)*y == g'*y

# LinearOps functor
#------------------

A, B = Ob(FreeLinearFunctions, :A, :B)

f, g  = Hom(:f, A, B), Hom(:g, A, B)
h, k  = Hom(:f, B, A), Hom(:g, B, A)

fop, gop = LinearOperator([0 1 -1; -1 1 0]), LinearOperator([1 2 3; 4 5 6])
hop, kop = LinearOperator([1 -1; -1 1; 0 1]), LinearOperator([1 1; 2 2; 3 3])

d = Dict(f=>fop, h=>hop, g => gop, k => kop)

F(ex)=functor((LinearOpDom, LinearOperator), ex, generators=d)

@test Matrix(F(compose(f,h))*[1 1; 2 3; 4 5]) == Matrix(hop*fop*[1 1; 2 3; 4 5])
@test Matrix(F(compose(g,k))*[1 1; 2 3; 4 5]) == Matrix(kop*gop*[1 1; 2 3; 4 5])

# Catlab evaluate
#----------------

A, B = Ob(FreeLinearFunctions, :A, :B)
f, g, h, k = Hom(:f, A, A), Hom(:g, B, B), Hom(:h, A, B), Hom(:k, A, B)
generators = Dict(
  A => 2, B => 3,
  f => [0 1; 1 0],
  g => [1 2 3; 4 5 6; 7 8 9],
  h => [1 -1; -1 1; 0 1],
  k => [1 1; 2 2; 3 3],
)
ev = (args...) -> evaluate(args...; generators=generators)

val = arg -> generators[arg]
x, y = [2, 1], [7, 3, 5]
@test ev(f, x) == val(f)*x

@test ev(f⋅f, x) == val(f)*(val(f)*x)
@test ev(f⋅h, x) == val(h)*(val(f)*x)
@test ev(f⊕g, x, y) == (val(f)*x, val(g)*y)

@test ev(plus(A), [1,3], [2,4]) == [3,7]
@test ev(zero(A)) == [0,0]
@test ev(h+k, x) == val(h)*x + val(k)*x
@test ev(scalar(A,3),x) == 3*x
@test ev(antipode(A),x) == -1*x

# Wiring diagrams
#################

A, B = Ob(FreeLinearFunctions, :A, :B)
f, g, h, k = Hom(:f, A, A), Hom(:g, B, B), Hom(:h, A, B), Hom(:k, A, B)

d = to_wiring_diagram(f⋅h⋅g)
@test nboxes(d) == 3
@test dom(d) == Ports([:A])
@test codom(d) == Ports([:B])

@present Mat(FreeLinearFunctions) begin
    X::Ob
    Y::Ob
    Z::Ob
    V::Ob
    W::Ob

    f::Hom(X,Y)
    g::Hom(Y,Z)
    h::Hom(Y⊕Z, V)
    fˣ::Hom(X,X)
    fʸ::Hom(Y,Y)
    fᶻ::Hom(Z,Z)
    fᵛ::Hom(V,V)
    fʷ::Hom(W,W)
end

@test_skip A = @program Mat (x::X,y::Y) begin
    v = h(f(x),g(y))
    return v
end

@test_skip codom(A) == generator(Mat, :V)

@test_skip A = @program Mat (x::X,y::Y) begin
    v = h(f(x),g(y))
    return fᵛ(v)
end

@test_skip codom(A) == generator(Mat, :V)

end


function fibonacci(n::Int)
    n==0 && return 0
    n==1 && return 1
    fibonacci(n-1) + fibonacci(n-2)
    end

print(fibonacci(40))