immutable CoordinateSystem
    center::XYZ
    rotation_matrix::Mat3x3
end

function CoordinateSystem(pc::XYZ,pz::XYZ,pxz::XYZ)
    vecz = normalize(pz - pc)
    vecy = normalize(cross(vecz,pxz-pc))
    vecx = normalize(cross(vecy,vecz))
    CoordinateSystem(pc,Mat3x3([Array(vecx) Array(vecy) Array(vecz)]))
end

rotate_from_global(csys::CoordinateSystem,p::XYZ) = csys.rotation_matrix' * p
rotate_to_global(csys::CoordinateSystem,p::XYZ) = csys.rotation_matrix * p

rotate_from_global(csys::CoordinateSystem,p::Mat3x3) = csys.rotation_matrix' * p * csys.rotation_matrix
rotate_to_global(csys::CoordinateSystem,p::Mat3x3) = csys.rotation_matrix * p * csys.rotation_matrix'

from_global(csys::CoordinateSystem,p::XYZ) = rotate_from_global(csys,p-csys.center)
to_global(csys::CoordinateSystem,p::XYZ) = rotate_to_global(csys,p) + csys.center


using Test

using MultilayerPerceptron

#===== Helpers =====#
function cross_entropy(ŷ::AbstractMatrix{Float64}, y::AbstractMatrix{Float64})::Float64
    n_batches = size(ŷ, 2)
    -sum(@. y * log(ŷ)) / n_batches
end

function binary_cross_entropy(ŷ::Float64, y::Bool)::Float64
    ϵ = eps(ŷ)
    -log(ϵ + (y ? ŷ : 1.0 - ŷ))
end

# https://math.stackexchange.com/a/2503773
function ∇binary_cross_entropy(ŷ::Float64, y::Bool)::Float64
    let y = Float64(y)
        (ŷ - y) / (ŷ - ŷ * ŷ)
    end
end

@testset "Activation" begin
    sz = (5, 3)
    y = rand(Float64, sz) * 10
    ŷ = rand(Float64, sz) * 10
    (e, _) = logit_cross_entropy(ŷ, y)
    @test e ≈ cross_entropy(softmax(ŷ), y)

    len = 5
    ŷ = rand(Float64, 1, len)
    y = rand(Bool, 1, len)
    (e, ∇) = logit_binary_cross_entropy(ŷ, y)
    logistic_ŷ = logistic.(ŷ)
    @test e ≈ sum(binary_cross_entropy.(logistic_ŷ, y)) / len
    @test ∇ ≈ @. ∇binary_cross_entropy(logistic_ŷ, y) * ∇logistic(ŷ, logistic_ŷ)
end

@testset "Xor" begin
    N = 1000

    inputs = unique(Iterators.product(0:1, 0:1))
    train_x = hcat(map(ab -> [Float64.(ab)...], inputs)...)
    train_y = reshape(map(((a, b),) -> Float64(a ⊻ b), inputs), 1, :)

    hidden_layer = Layer((2, 2); activation_fn=tanh, grad_fn=∇tanh)
    output_layer = Layer((1, 2))

    mlp = MLP([hidden_layer, output_layer])
    fit!(mlp, train_x, train_y; learning_rate=0.2, batch_size=2, n_epochs=N)

    test_x, test_y = train_x, train_y
    @test round.(Int, predict(mlp, Float64.(test_x))) == test_y
end


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

### Gauss quadrature rules for prismatic elements (wedge)

""" Gauss-Legendre quadrature, 6 point rule on wedge. """
function get_quadrature_points(::Type{Val{:GLWED6}})
    w = 1.0/6.0
    a = sqrt(1.0/3.0)
    weights = (w, w, w, w, w, w)
    points = ((0.5, 0.0, -a), (0.0, 0.5, -a), (0.5, 0.5, -a),
              (0.5, 0.0,  a), (0.0, 0.5,  a), (0.5, 0.5,  a))
    return zip(weights, points)
end

function get_order(::Type{Val{:GLWED6}})
    return 1
end

""" Gauss-Legendre quadrature, 6 point rule on wedge. """
function get_quadrature_points(::Type{Val{:GLWED6B}})
    w = 1.0/6.0
    a = sqrt(1.0/3.0)
    weights = (w, w, w, w, w, w)
    points = ((2.0/3.0, 1.0/6.0, -a), (1.0/6.0, 2.0/3.0, -a), (1.0/6.0, 1.0/6.0, -a),
              (2.0/3.0, 1.0/6.0,  a), (1.0/6.0, 2.0/3.0,  a), (1.0/6.0, 1.0/6.0,  a))
    return zip(weights, points)
end

function get_order(::Type{Val{:GLWED6B}})
    return 1
end

""" Gauss-Legendre quadrature, 21 point rule on wedge. """
function get_quadrature_points(::Type{Val{:GLWED21}})
    alpha = sqrt(3/5)
    c1 = 5/9
    c2 = 8/9
    a = (6+sqrt(15))/21
    b = (6-sqrt(15))/21

    weights = (
        c1*9/80,
        c1*((155+sqrt(15))/2400),
        c1*((155+sqrt(15))/2400),
        c1*((155+sqrt(15))/2400),
        c1*((155-sqrt(15))/2400),
        c1*((155-sqrt(15))/2400),
        c1*((155-sqrt(15))/2400),
        c2*9/80,
        c2*((155+sqrt(15))/2400),
        c2*((155+sqrt(15))/2400),
        c2*((155+sqrt(15))/2400),
        c2*((155-sqrt(15))/2400),
        c2*((155-sqrt(15))/2400),
        c2*((155-sqrt(15))/2400),
        c1*9/80,
        c1*((155+sqrt(15))/2400),
        c1*((155+sqrt(15))/2400),
        c1*((155+sqrt(15))/2400),
        c1*((155-sqrt(15))/2400),
        c1*((155-sqrt(15))/2400),
        c1*((155-sqrt(15))/2400))

    points = (
          (1/3, 1/3, -alpha),
          (a, a, -alpha),
          (1-2a, a, -alpha),
          (a, 1-2a, -alpha),
          (b, b, -alpha),
          (1-2b, b, -alpha),
          (b, 1-2b, -alpha),
          (1/3, 1/3, 0),
          (a, a, 0),
          (1-2a, a, 0),
          (a, 1-2a, 0),
          (b, b, 0),
          (1-2b, b, 0),
          (b, 1-2b, 0),
          (1/3, 1/3, alpha),
          (a, a, alpha),
          (1-2a, a, alpha),
          (a, 1-2a, alpha),
          (b, b, alpha),
          (1-2b, b, alpha),
          (b, 1-2b, alpha))

    return zip(weights, points)
end

function get_order(::Type{Val{:GLWED21}})
    return 5
end


using Test
using LinearAlgebra
using LinearAlgebra: norm
using Ferrite
using StaticArrays

using TopOpt
using TrussTopOpt.TrussTopOptProblems
using TrussTopOpt.TrussTopOptProblems: buckling
using Crayons.Box
using GLMakie

problem_json = ["mgz_geom_stiff_ex9.1.json", "buckling_2d_debug.json"]
ins_dir = joinpath(@__DIR__, "instances", "fea_examples");

# @testset "Buckling problem solve - $(problem_json[i])" for i in 1:length(problem_json)
    i = 1
    file_name = problem_json[i]
    problem_file = joinpath(ins_dir, file_name)

    node_points, elements, mats, crosssecs, fixities, _ = parse_truss_json(problem_file);

    P_val = 400 # kN
    P = Dict(2 => SVector{2}(P_val * [0.05, -1.0]))
    nnode = length(node_points)

    # Newton iteration setup
    max_newton_itr = 50
    NEWTON_TOL = 1e-8

    function dict2vec(dict, veclen)
        v = zeros(veclen)
        for (nid, val) in dict
            v[2*nid-1:2*nid] = val
        end
        return v
    end

    # Δu = Dict(nid => SVector{2,Float64}(0,0) for nid in keys(node_points))

    updated_nodes = copy(node_points)
    problem = TrussProblem(Val{:Linear}, node_points, elements, P, fixities, mats, crosssecs);
    solver = FEASolver(Displacement, Direct, problem)
    solver()
    K, Kσ = buckling(problem, solver.globalinfo, solver.elementinfo; u=zeros(2*nnode))
    Kt = K + Kσ
    P_vec = solver.globalinfo.f
    # * initialize displacements as zero
    ut = zeros(nnode*2)

    # tracking iteration history for plotting
    ub_t = []
    p_t = []

    newton_itr = -1

    # Newton-Raphson Iterations
    # https://people.duke.edu/~hpgavin/cee421/truss-finite-def.pdf
    # https://kristofferc.github.io/JuAFEM.jl/dev/examples/hyperelasticity/
    while true
        global newton_itr, Kt, P_vec, ut, updated_nodes
        newton_itr += 1
        reaction = Kt * ut
        @show res_norm = norm(reaction - P_vec)
        push!(ub_t, ut[2*2-1])
        push!(p_t,-reaction[2*2])

        if res_norm < NEWTON_TOL
            break
        elseif newton_itr > max_newton_itr
            break
            # error("Reached maximum Newton iterations, aborting")
        end

        # construct new stiffness matrix
        ut = Kt \ P_vec
        updated_nodes = Dict(nid => pt + ut[2*nid-1:2*nid] for (nid, pt) in updated_nodes)
        problem = TrussProblem(Val{:Linear}, updated_nodes, elements, P, fixities, mats, crosssecs);
        solver = FEASolver(Displacement, Direct, problem)
        solver()
        K, Kσ = buckling(problem, solver.globalinfo, solver.elementinfo; u=ut)
        Kt = K + Kσ
    end
    scatter(ub_t, p_t)
    lines!(ub_t, p_t, color = :blue, linewidth = 3)
    current_figure()

    # scene, layout = draw_truss_problem(problem; u=v, default_exagg=1.0, exagg_range=10.0,
    #     default_load_scale=0.2, default_support_scale=0.2, default_arrow_size=0.03)

# end # end test set

# =#


@testset "Magnitude of completeness" begin
    b = 1
    M0 = 0.5
    mags = rand(GRclassic(b,M0), 100000)
    filter!(x->rand() < cdf(Normal(0.8,0.2),x), mags)
    ΔM = 0.1

    x, y = histmag(mags,ΔM)
    @test round(maxcurv(mags, ΔM), digits=2) == 0.95
    @test maxcurv(mags, ΔM) == maxcurv(x,y)
    @test_throws ErrorException("The two input arrays must have the same number of elements") maxcurv([0.1],[1.2,1.6])

    @test round(mbass(mags, ΔM), digits=2) == 1.25
    @test mbass(mags, ΔM) == mbass(x,y)
    @test_throws ErrorException("The two input arrays must have the same number of elements") mbass([0.1],[1.2,1.6])
    @test isnan( mbass(collect(0.1:0.1:1.1),fill(11,11)) )
    @test isnan( mbass([0.1,0.2],[1.2,1.6]) )

    @test round(mbs(mags,ΔM,winsize=5, minevents=50), digits=2) == 1.55
    @test isnan( mbs([0.1,0.2],0.1) )

end


#!/usr/bin/env julia
using Plots, Random, Printf

function initCondGen() #get random initial conditions for mass/radius, position, and velocity
    function getMass(nBodies) #generate random masses that better reflect actual stellar populations
        mList=zeros(nBodies)
        N=(0.5^(-1.3)-150^(-1.3))/1.3 #crude approximation of IMF integral assuming alpha = 2.3, stellar mass range of 0.5:150 solar masses
        rescale=1e6
        max=floor(Int,N*rescale)
        for i=1:nBodies
            intTarget=rand(0:max,1)[1]/rescale
            m=(0.5^(-1.3)-intTarget*1.3)^(-1/1.3) #just algebra from above
            mList[i]=round(m,digits=2)
        end
        return mList
    end
    m=rand(1:1500,3)./10 #3 random masses between 0.1 and 150 solar masses, uniform distribution
    #m=getMass(3) #get mass from IMF -- this way is kind of boring...so not using it, but left here in case I change my mind?
    rad=m.^0.8 #3 radii based on masses in solar units
    m=m.*2e30 #convert to SI kg
    rad=rad.*7e8 #convert to SI m
    pos1=rand(-10:10,3) #random initial coordinates x & y for first body, AU
    function genPos2(pos1)
        accept2=false
        while accept2==false
            pos2=rand(-10:10,3) #random initial coordinates for second body, AU
            dist21=sqrt((pos1[1]-pos2[1])^2+(pos1[2]-pos2[2])^2+(pos1[3]-pos2[3])^2)
            if (dist21*1.5e11)>(rad[1]+rad[2]) #they aren't touching
                accept2=true
                return pos2
            end
        end
    end
    pos2=genPos2(pos1)
    function genPos3(pos1,pos2)
        accept3=false
        while accept3==false
            pos3=rand(-10:10,3) #random initial coordinates for third body, AU
            dist31=sqrt((pos1[1]-pos3[1])^2+(pos1[2]-pos3[2])^2+(pos1[3]-pos3[3])^2)
            dist32=sqrt((pos2[1]-pos3[1])^2+(pos2[2]-pos3[2])^2+(pos2[3]-pos3[3])^2)
            if (dist31*1.5e11)>(rad[1]+rad[3]) && (dist32*1.5e11)>(rad[2]+rad[3]) #3rd isn't touching either
                accept3=true
                return pos3
            end
        end
    end
    pos3=genPos3(pos1,pos2)
    pos=[pos1[1],pos1[2],pos1[3],pos2[1],pos2[2],pos2[3],pos3[1],pos3[2],pos3[3]].*1.5e11 #convert accepted positions to SI, m
    v=rand(-7e3:7e3,9) #random xyz velocities with mag between -10 & 10 km/s, totally arbitrary...
    #r=[x1,y1,x2,y2,x3,y3,v1x,v1y,v2x,v2y,v3x,v3y]
    r=[pos[1],pos[2],pos[3],pos[4],pos[5],pos[6],pos[7],pos[8],pos[9],v[1],v[2],v[3],v[4],v[5],v[6],v[7],v[8],v[9]]
    open("initCond.txt","w") do f #save initial conditions to file in folder where script is run
        write(f,"m1=$(@sprintf("%.1f",(m[1]/2e30))) m2=$(@sprintf("%.1f",(m[2]/2e30))) m3=$(@sprintf("%.1f",(m[3]/2e30))) (solar masses)\nv1x=$(v[1]/1e3) v1y=$(v[2]/1e3) v1z=$(v[3]/1e3) v2x=$(v[4]/1e3) v2y=$(v[5]/1e3) v2z=$(v[6]/1e3) v3x=$(v[7]/1e3) v3y=$(v[8]/1e3) v3z=$(v[9]/1e3) (km/s)\nx1=$(pos1[1]) y1=$(pos1[2]) z1=$(pos1[3]) x2=$(pos2[1]) y2=$(pos2[2]) z2=$(pos2[3]) x3=$(pos3[1]) y3=$(pos3[2]) z3=$(pos3[3]) (AU from center)")
    end
    return r, rad, m
end

function dR(r,m) #function we will use RK4 on to approximate solution
    G=6.67408313131313e-11# Nm^2/kg^2
    M1,M2,M3=m[1],m[2],m[3] #kg
    x1,x2,x3=r[1],r[4],r[7] #m
    y1,y2,y3=r[2],r[5],r[8] #m
    z1,z2,z3=r[3],r[6],r[9]
    c1,c2,c3=G*M1,G*M2,G*M3 #Nm^2/kg
    r1_2=sqrt((x1-x2)^2+(y1-y2)^2+(z1-z2)^2) #distance from 1->2, m
    r1_3=sqrt((x1-x3)^2+(y1-y3)^2+(z1-z3)^2) #distance from 1->3, m
    r2_3=sqrt((x2-x3)^2+(y2-y3)^2+(z2-z3)^2) #distance from 2->3, m

    v1X,v2X,v3X=r[10],r[13],r[16] #these are our change in position after dt (dr/dt*dt=dr)
    v1Y,v2Y,v3Y=r[11],r[14],r[17] #m after * dt
    v1Z,v2Z,v3Z=r[12],r[15],r[18]

    #get change in velocity from accelerations (d^2r/dt^2*dt=dv/dt*dt=dv)
    dx1=-(c2*(x1-x2)/(r1_2^3))-(c3*(x1-x3)/(r1_3^3)) #d^2x/dt^2 for 1, m/s after * dt
    dx2=-(c1*(x2-x1)/(r1_2^3))-(c3*(x2-x3)/(r2_3^3)) #d^2x/dt^2 for 2, m/s
    dx3=-(c1*(x3-x1)/(r1_3^3))-(c2*(x3-x2)/(r2_3^3)) #d^2x/dt^2 for 3, m/s
    dy1=-(c2*(y1-y2)/(r1_2^3))-(c3*(y1-y3)/(r1_3^3)) #d^2y/dt^2 for 1, m/s
    dy2=-(c1*(y2-y1)/(r1_2^3))-(c3*(y2-y3)/(r2_3^3)) #d^2y/dt^2 for 2, m/s
    dy3=-(c1*(y3-y1)/(r1_3^3))-(c2*(y3-y2)/(r2_3^3)) #d^2y/dt^2 for 3, m/s
    dz1=-(c2*(z1-z2)/(r1_2^3))-(c3*(z1-z3)/(r1_3^3)) #d^2y/dt^2 for 1, m/s
    dz2=-(c1*(z2-z1)/(r1_2^3))-(c3*(z2-z3)/(r2_3^3)) #d^2y/dt^2 for 2, m/s
    dz3=-(c1*(z3-z1)/(r1_3^3))-(c2*(z3-z2)/(r2_3^3)) #d^2y/dt^2 for 3, m/s

    return [v1X,v1Y,v1Z,v2X,v2Y,v2Z,v3X,v3Y,v3Z,dx1,dy1,dz1,dx2,dy2,dz2,dx3,dy3,dz3]
end

function gen3Body(stopCond=[10,100],numSteps=10000) #default stop conditions of 10 yrs and 100 AU sep
    tStop=stopCond[1]*365*24*3600 #convert to SI s
    sepStop=stopCond[2]*1.5e11 #convert to SI m
    stop=false
    currentT=0
    t=range(0,stop=tStop,length=(numSteps+1)) #+1 because I don't want 0 to count
    stepSize=tStop/numSteps
    x1=zeros(length(t))
    y1=zeros(length(t))
    z1=zeros(length(t))
    x2=zeros(length(t))
    y2=zeros(length(t))
    z2=zeros(length(t))
    x3=zeros(length(t))
    y3=zeros(length(t))
    z3=zeros(length(t))
    r,rad,m=initCondGen()
    min12=rad[1]+rad[2]
    min13=rad[1]+rad[3]
    min23=rad[2]+rad[3]
    i=1
    stopT=maximum(t)

    #implement RK4 to model solutions to differential equations
    while stop==false
        if currentT==stopT || currentT>stopT #in case of rounding error or something
            stop=true
        elseif i>numSteps+1 #inf loop failsafe
            stop=true
            println("error: shouldn't have gotten here")
        else
            x1[i]=r[1]
            y1[i]=r[2]
            z1[i]=r[3]
            x2[i]=r[4]
            y2[i]=r[5]
            z2[i]=r[6]
            x3[i]=r[7]
            y3[i]=r[8]
            z3[i]=r[9]

            k1=stepSize*dR(r,m)
            k2=stepSize*dR(r.+0.5.*k1,m)
            k3=stepSize*dR(r.+0.5.*k2,m)
            k4=stepSize*dR(r.+k3,m)
            r+=(k1.+2.0*k2.+2.0.*k3.+k4)./6

            #check separation after each dt step
            sep12=sqrt((x1[i]-x2[i])^2+(y1[i]-y2[i])^2+(z1[i]-z2[i])^2)
            sep13=sqrt((x1[i]-x3[i])^2+(y1[i]-y3[i])^2+(z1[i]-z3[i])^2)
            sep23=sqrt((x3[i]-x2[i])^2+(y3[i]-y2[i])^2+(z3[i]-z2[i])^2)

            if sep12<min12 || sep13<min13 || sep23<min23 || sep12>sepStop || sep13>sepStop || sep23>sepStop
                stop=true #stop if collision happens or body is ejected
                t=range(0,stop=currentT,length=i) #t should match pos vectors
                x1=x1[1:i] #don't want trailing zeros
                y1=y1[1:i]
                z1=z1[1:i]
                x2=x2[1:i]
                y2=y2[1:i]
                z2=z2[1:i]
                x3=x3[1:i]
                y3=y3[1:i]
                z3=z3[1:i]
            end
            i+=1
            currentT+=stepSize #next step
        end
    end
    return [x1,y1,z1,x2,y2,z2,x3,y3,z3], t, m, rad
end

function getInteresting3Body(minTime=0) #in years, defaults to 0
    #sometimes random conditions result in a really short animation where things
    #just crash into each other/fly away, so this function throws away those
    yearSec=365*24*3600
    interesting=false
    i=1
    while interesting==false
        plotData,t,m,rad=gen3Body([60,150],600000)
        if (maximum(t)/yearSec)>minTime #only return if simulation runs for longer than minTime
            println(maximum(t)/yearSec) #tell me how many years we are simulating
            open("cron_log.txt","a") do f #for cron logging, a flag = append
                write(f,"$(maximum(t)/yearSec)\n")
            end
            return plotData,t,m,rad
            interesting=true
        elseif i>5000 #computationally expensive so don't want to go forever
            interesting=true #render it anyways I guess because sometimes it's fun?
            println("did not find interesting solution in number of tries allotted, running anyways")
            println(maximum(t)/yearSec) #how many years simulation runs for
            open("cron_log.txt","a") do f #for cron logging
                write(f,"$(maximum(t)/yearSec)\n")
            end
            return plotData,t,m,rad
        end
        i+=1
    end
end

function getLims(pos,padding) #determines plot limits at each frame, padding in units of pos
    x=[pos[1],pos[4],pos[7]]
    xMin=minimum(x)
    xMax=maximum(x)
    dx=xMax-xMin
    y=[pos[2],pos[5],pos[8]]
    yMin=minimum(y)
    yMax=maximum(y)
    dy=yMax-yMin
    z=[pos[3],pos[6],pos[9]]
    zMin=minimum(z)
    zMax=maximum(z)
    dz=zMax-zMin
    dList=[dx,dy,dz]
    if maximum(dList)==dx
        #use x for square
        xlims=[xMin-padding,xMax+padding]
        ylims=[yMin-padding,yMin+dx+padding]
        zlims=[zMin-padding,zMin+dx+padding]
    elseif maximum(dList)==dy
        #use y for square
        xlims=[xMin-padding,xMin+dy+padding]
        ylims=[yMin-padding,yMax+padding]
        zlims=[zMin-padding,zMin+dy+padding]
    else
        #use z for cube
        xlims=[xMin-padding,xMin+dz+padding]
        ylims=[yMin-padding,yMin+dz+padding]
        zlims=[zMin-padding,zMax+padding]
    end
    return xlims,ylims,zlims
end

function getColors(m,c) #places colors of objects according to mass/size
    #c=[:biggest,:medium,:smallest] (order of input colors)
    maxM=maximum(m)
    minM=minimum(m)
    colors=[:blue,:blue,:blue] #testing
    if m[1]==maxM
        colors[1]=c[1]
        if m[2]==minM
            colors[2]=c[3]
            colors[3]=c[2]
        else
            colors[3]=c[3]
            colors[2]=c[2]
        end
    elseif m[2]==maxM
        colors[2]=c[1]
        if m[1]==minM
            colors[1]=c[3]
            colors[3]=c[2]
        else
            colors[3]=c[3]
            colors[1]=c[2]
        end
    else
        colors[3]=c[1]
        if m[1]==minM
            colors[1]=c[3]
            colors[2]=c[2]
        else
            colors[2]=c[3]
            colors[1]=c[2]
        end
    end
    return colors
end

function makeCircleVals(r,center=[0,0,0])
    #(x,y,z)=(rcos(theta)sin(phi),rsin(theta)sin(phi),rcos(phi))
    xOffset=center[1]
    yOffset=center[2]
    zOffset=center[3]
    xVals=[]
    yVals=[]
    zVals=[]
    for i=0:pi/128:2*pi
        for j=0:pi/64:pi
            x=r*cos(i)*sin(j)+xOffset
            y=r*sin(i)*sin(j)+yOffset
            z=r*cos(j)+zOffset
            push!(xVals,x)
            push!(yVals,y)
            push!(zVals,z)
        end
    end
    return xVals,yVals,zVals
end

plotData,t,m,rad=getInteresting3Body(15)
c=[:DodgerBlue,:Gold,:Tomato] #most massive to least massive, also roughly corresponds to temp
colors=getColors(m,c)
#adding fake stars
numStars=2500
starsX=zeros(numStars)
starsY=zeros(numStars)
starsZ=zeros(numStars)
for i=1:numStars
    num=rand(-200:200,3) #box size is 70 AU but we need some extra padding for movement
    starsX[i]=num[1]
    starsY[i]=num[2]
    starsZ[i]=num[3]
end

#this new way runs significantly faster (~2x improvement over @anim)
#Downside is it spams folder with png images of every frame and must manually compile with ffmpeg
#Comment out and use older way (after this below) if performance/specific formatting is not an issue

plotLoadPath="/home/kirk/Documents/3Body/tmpPlots/"
threeBodyAnim=Animation(plotLoadPath,String[])
for i=1:333:length(t) #this makes animation scale ~1 sec/year with other conditions
    GR.inline("png") #added to eneable cron/jobber compatibility, also this makes frames generate WAY faster? Prior to adding this when run from cron/jobber frames would stop generating at 408 for some reason.
    gr(legendfontcolor = plot_color(:white)) #legendfontcolor=:white plot arg broken right now (at least in this backend)
    print("$(@sprintf("%.2f",i/length(t)*100)) % complete\r") #output percent tracker
    pos=[plotData[1][i],plotData[2][i],plotData[3][i],plotData[4][i],plotData[5][i],plotData[6][i],plotData[7][i],plotData[8][i],plotData[9][i]] #current pos
    limx,limy,limz=getLims(pos./1.5e11,10) #convert to AU, 10 AU padding
    p=plot3d(plotData[1][1:33:i]./1.5e11,plotData[2][1:33:i]./1.5e11,plotData[3][1:33:i]./1.5e11,label="",linecolor=colors[1],linealpha=max.((1:33:i) .+ 10000 .- i,2500)/10000) #plot orbits up to i
    p=plot3d!(plotData[4][1:33:i]./1.5e11,plotData[5][1:33:i]./1.5e11,plotData[6][1:33:i]./1.5e11,label="",linecolor=colors[2],linealpha=max.((1:33:i) .+ 10000 .- i,2500)/10000) #linealpha argument causes lines to decay
    p=plot3d!(plotData[7][1:33:i]./1.5e11,plotData[8][1:33:i]./1.5e11,plotData[9][1:33:i]./1.5e11,label="",linecolor=colors[3],linealpha=max.((1:33:i) .+ 10000 .- i,2500)/10000) #example: alpha=max.((1:i) .+ 100 .- i,0) causes only last 100 to be visible
    p=scatter3d!(starsX,starsY,starsZ,markercolor=:white,markersize=:1,label="") #fake background stars
    star1=makeCircleVals(rad[1],[plotData[1][i],plotData[2][i],plotData[3][i]]) #generate spheres with appropriate sizes for each star
    star2=makeCircleVals(rad[2],[plotData[4][i],plotData[5][i],plotData[6][i]]) #at current positions
    star3=makeCircleVals(rad[3],[plotData[7][i],plotData[8][i],plotData[9][i]])
    p=plot3d!(star1[1]./1.5e11,star1[2]./1.5e11,star1[3]./1.5e11,label="$(@sprintf("%.1f", m[1]./2e30))",color=colors[1],fill=true) #plot star circles with labels
    p=plot3d!(star2[1]./1.5e11,star2[2]./1.5e11,star2[3]./1.5e11,label="$(@sprintf("%.1f", m[2]./2e30))",color=colors[2],fill=true)
    p=plot3d!(star3[1]./1.5e11,star3[2]./1.5e11,star3[3]./1.5e11,label="$(@sprintf("%.1f", m[3]./2e30))",color=colors[3],fill=true)
    p=plot3d!(background_color=:black,background_color_legend=:transparent,foreground_color_legend=:transparent,
        background_color_outside=:white,aspect_ratio=:equal,legendtitlefontcolor=:white) #formatting for plot frame
    p=plot3d!(title="Random Three-Body Problem\nt: $(@sprintf("%0.2f",t[i]/365/24/3600)) years after start",
        legend=:best,xaxis=("x: AU",(limx[1],limx[2]),font(9,"Courier")),yaxis=("y: AU",(limy[1],limy[2]),font(9,"Courier")),zaxis=("z: AU",(limz[1],limz[2]),font(9,"Courier")),
        gridalpha=0.5,gridcolor=:white,titlefont=font(14,"Courier"),size=(720,721),legendfontsize=8,legendtitle="Mass (in solar masses)",legendtitlefontsize=8) #add in axes/title/legend with formatting
    frame(threeBodyAnim,p) #generate the frame
    closeall() #close plots
end

#threeBodyFile="3Body_fps30.mp4"

#crf is compression value (17 or 18 "visually lossless"), pix_fmt is for twitter specific vid req, -b:v specifies target bitrate, -vcodec specifies codec (h264 in this case) -y says overwrite existing file
#run( `ffmpeg -framerate 30 -i $plotLoadPath"%06d.png" -vcodec libx264 -pix_fmt yuv420p -profile:v high -b:v 2048K -y -vf "scale=720:720,setdar=1/1" $threeBodyFile` ) #-vf scale=720:72 -crf 25
#run( `ffmpeg -framerate 30 -i $plotLoadPath"%06d.png" -c:v libx264 -preset slow -coder 1 -movflags +faststart -g 15 -crf 18 -pix_fmt yuv420p -profile:v high -y -bf 2 -vf "scale=720:720,setdar=1/1" $threeBodyFile` ) #all this bullshit to hopefully satisfy twitter requirements

#NOTE: moved ffmpeg commands to shell script

#old (simpler) way of generating animation
#uncomment and use this way if you just want a simple animation saved and don't
#care about performance/specific formatting of video.

# threeBodyAnim=@animate for i=1:length(t)
#     gr(legendfontcolor = plot_color(:white)) #plot arg broken right now in Julia
#     print("$(@sprintf("%.2f",i/length(t)*100)) % complete\r") #output percent tracker
#     pos=[plotData[1][i],plotData[2][i],plotData[3][i],plotData[4][i],plotData[5][i],plotData[6][i]] #current pos
#     limx,limy=getLims(pos./1.5e11,5) #convert to AU, 5 AU padding
#     plot(plotData[1][1:i]./1.5e11,plotData[2][1:i]./1.5e11,label="",linecolor=colors[1])
#     plot!(plotData[3][1:i]./1.5e11,plotData[4][1:i]./1.5e11,label="",linecolor=colors[2])
#     plot!(plotData[5][1:i]./1.5e11,plotData[6][1:i]./1.5e11,label="",linecolor=colors[3])
#     scatter!(starsX,starsY,markercolor=:white,markersize=:1,label="") #fake background stars
#     star1=makeCircleVals(rad[1],[plotData[1][i],plotData[2][i]])
#     star2=makeCircleVals(rad[2],[plotData[3][i],plotData[4][i]])
#     star3=makeCircleVals(rad[3],[plotData[5][i],plotData[6][i]])
#     plot!(star1[1]./1.5e11,star1[2]./1.5e11,label="$(@sprintf("%.1f", m[1]./2e30))",color=colors[1],fill=true)
#     plot!(star2[1]./1.5e11,star2[2]./1.5e11,label="$(@sprintf("%.1f", m[2]./2e30))",color=colors[2],fill=true)
#     plot!(star3[1]./1.5e11,star3[2]./1.5e11,label="$(@sprintf("%.1f", m[3]./2e30))",color=colors[3],fill=true)
#     plot!(background_color=:black,background_color_legend=:transparent,background_color_outside=:white,aspect_ratio=:equal,legendtitlefontcolor=:white) #legendfontcolor=:white
#     plot!(xlabel="x: AU",ylabel="y: AU",title="Random Three Body Problem\nt: $(@sprintf("%0.2f",t[i]/365/24/3600)) yrs after start",
#         legend=:best,xaxis=("x: AU",(limx[1],limx[2]),font(12,"Courier")),yaxis=("y: AU",(limy[1],limy[2]),font(12,"Courier")),
#         grid=false,titlefont=font(24,"Courier"),size=(720,720),legendfontsize=12,legendtitle="Mass (in solar masses)",legendtitlefontsize=14)
#     end every 25

#mp4(threeBodyAnim,"3Body_fps30.mp4",fps=30)
#OR
#gif(threeBodyAnim,"3Body_fps30.gif",fps=30)


using FileIO,Optim
include("model.jl")
img = load("raw/509.tiff")
x0 = Float32[290 265 385; 310 360 315] / 512
ygold = convert(Array{Float32}, img.data)
#ygold[ygold.<0.05] = 0.0
x = similar(ygold,(2,3))
x6 = reshape(x, 6)
minit(x)
ypred = similar(ygold)
dx = similar(x)

f(x)=Float32(mlossxy(reshape(x,(2,3)),ygold))

function g(x,dx)
    x33=reshape(x,(2,3))
    dx33=reshape(dx,(2,3))
    mforw(x33,ypred)
    mback(x33,ypred,ygold,dx33)
    dx33
end
# optimize(f, g, x, LBFGS())


include("SortSparse.jl")
include("KoLesky.jl")
include("Cholesky.jl")
include("Utils.jl")
using LinearAlgebra
using RecursiveArrayTools
using Distributions
using Distances
using Random
using Profile
using JLD
using MAT
using BlockArrays
Random.seed!(123)

NTrain = 2000
NTest = 5
N = NTest + NTrain
d =  2
x = randn(d, N)

l = 1.0
σ = Diagonal(2.0 * ones(N))
cov = r -> exp(-r/l)

ρ = 4.0
λ = 1.3

# reordering points and forming skeletons
@time colptr, rowval, P, revP, distances = sortSparseRev(x, ρ, 1)
x = x[:, P]
skeletons = construct_skeletons(colptr, rowval, distances, λ)

# Creating
Θ = cov.(pairwise(Euclidean(), x, dims=2))
Θ = Θ + σ

L = assembleL(skeletons, x, cov)
U = sparse(triu(L * L') + min.(inv(σ), 1e30))
icholU_high_level!(U)

ΘApprox = inv(Matrix(L' * L)) * U' * U * σ

@show KL(Θ, ΘApprox)
@show KL(ΘApprox, Θ)


# BATs pdf, cdf, quantile, sampler defined below
# Parameter layout: (κ₀, τ₀, ϕ₀, κ₁, τ₁, ϕ₁, ν)
# (κ₀, τ₀, ϕ₀) are shape, scale, location for lower Tail
# (κ₁, τ₁, ϕ₁) are shape, scale, location for upper Tail

# Psi function, its derivative, and its inverse. Branched like this to stabilize
# numerics for large argument x.
Ψ(x)  = (x < 25.0) ? log1p(exp(x))           : x
dΨ(x) = (x < 25.0) ? exp(x)/(1.0+exp(x))     : 1.0
iΨ(x) = (x < 25.0) ? log(max(0.0, expm1(x))) : x # avoids an error with log of negative number

# One half of the H function, and the extension relevant for the density.
# Trying to be careful here about numerics, but it is kind of fishy as κ->0.
function H_part(x, κ, τ, ϕ)
  if isapprox(κ, 0.0, atol = 1.0e-8) # cutoff chosen by eye.
      a = Ψ((x - ϕ) / τ)
      return exp(a) - 0.5 * (a^2) * exp(a) * κ
  end
  (1.0 + κ * Ψ((x - ϕ) / τ))^(1.0 / κ)
end

H(x, κ₀, τ₀, ϕ₀, κ₁, τ₁, ϕ₁) = H_part(x, κ₁, τ₁, ϕ₁) - H_part(-x, κ₀, τ₀, -ϕ₀)
dH_part(x, κ, τ, ϕ) = ForwardDiff.derivative(z -> H_part(z, κ, τ, ϕ), x)

# Full PDF in Distributions.jl format. Note that the numerics
# here can be concerning: the return line out_1*(out_2+out_3) can commonly look
# like subnormal_tiny * (super_huge + super_huge), which is a floating point
# nightmare.
function Distributions.pdf(d::BulkAndTailsDist, x::Real)
  # Check support.
  ((x ≥ maximum(d)) || (x ≤ minimum(d))) && return 0.0
  # Unpack parameters for readability.
  (κ₀, τ₀, ϕ₀, κ₁, τ₁, ϕ₁, ν) = params(d)
  # Evaluate the density if we're inside the defined support.
  out_1 = Distributions.pdf(TDist(ν),H(x,κ₀,τ₀,ϕ₀,κ₁,τ₁,ϕ₁))
  out_2 = dH_part(x, κ₁, τ₁, ϕ₁)
  out_3 = dH_part(-x, κ₀, τ₀, -ϕ₀)
  max(0.0, out_1*(out_2+out_3)) # numerically dubious.
end

# Note: this t-cdf does not work with autodiff. Would need hand-coded tcdf function.
function Distributions.cdf(d::BulkAndTailsDist, x::Real)
  if (x ≤ minimum(d))
    return 0.0
  elseif (x ≥ maximum(d))
    return 1.0
  else
    (κ₀, τ₀, ϕ₀, κ₁, τ₁, ϕ₁, ν) = params(d)
    return Distributions.cdf(TDist(ν),H(x,κ₀,τ₀,ϕ₀,κ₁,τ₁,ϕ₁))
  end
end

function Distributions.quantile(d::BulkAndTailsDist, p::Real)
  Roots.find_zero(x -> Distributions.cdf(d,x) - p, (minimum(d),maximum(d)), Roots.Bisection())
end

function Distributions.logpdf(d::BulkAndTailsDist, x::Real)
  # Check support.
  ((x ≥ maximum(d)) || (x ≤ minimum(d))) && return -Inf
  # Unpack parameters for readability.
  (κ₀, τ₀, ϕ₀, κ₁, τ₁, ϕ₁, ν) = params(d)
  # Evaluate the density if we're inside the defined support.
  out_1 = Distributions.logpdf(TDist(ν),H(x,κ₀,τ₀,ϕ₀,κ₁,τ₁,ϕ₁))
  out_2 = dH_part(x, κ₁, τ₁, ϕ₁)
  out_3 = dH_part(-x, κ₀, τ₀, -ϕ₀)
  out_1 + log(out_2+out_3)
end

# Note: this t-cdf does not work with autodiff. Would need hand-coded logtcdf function.
function Distributions.logcdf(d::BulkAndTailsDist, x::Real)
  if (x ≤ minimum(d))
    return -Inf
  elseif (x ≥ maximum(d))
    return 0.0
  else
    (κ₀, τ₀, ϕ₀, κ₁, τ₁, ϕ₁, ν) = params(d)
    return Distributions.logcdf(TDist(ν),H(x,κ₀,τ₀,ϕ₀,κ₁,τ₁,ϕ₁))
  end
end

function Base.rand(d::BulkAndTailsDist; rng=Random.GLOBAL_RNG)
  quantile(d,rand(rng))
end

# Define R-friendly version.
batspdf(x, parms) = Distributions.pdf(BulkAndTailsDist(parms),x)
batscdf(x, parms) = Distributions.cdf(BulkAndTailsDist(parms),x)
batsquantile(x, parms) = Distributions.quantile(BulkAndTailsDist(parms),x)
batslogpdf(x, parms) = Distributions.logpdf(BulkAndTailsDist(parms),x)
batslogcdf(x, parms) = Distributions.logcdf(BulkAndTailsDist(parms),x)

# Overload for vector inputs, because R users like to not know when they are
# broadcasting functions.
batspdf(xv::Vector, parms) = [batspdf(x, parms) for x in xv]
batscdf(xv::Vector, parms) = [batscdf(x, parms) for x in xv]
batsquantile(xv::Vector, parms) = [batsquantile(x, parms) for x in xv]
batslogpdf(xv::Vector, parms) = [batslogpdf(x, parms) for x in xv]
batslogcdf(xv::Vector, parms) = [batslogcdf(x, parms) for x in xv]


module TestExteriorCalculus
using Test

using Catlab
using CombinatorialSpaces.ExteriorCalculus

@present Diffusion2DQuantities(FreeExtCalc2D) begin
  X::Space
  C::Hom(munit(), Form0(X))     # concentration
  ϕ::Hom(munit(), DualForm1(X)) # negative diffusion flux
  k::Hom(Form1(X), Form1(X))    # diffusivity (usually scalar multiplication)
end

@present Diffusion2D <: Diffusion2DQuantities begin
  # Fick's first law
  ϕ == C ⋅ d₀(X) ⋅ k ⋅ ⋆₁(X)
  # Diffusion equation
  C ⋅ ∂ₜ(Form0(X)) == ϕ ⋅ dual_d₁(X) ⋅ ⋆₀⁻¹(X)
end

@present Diffusion2D′ <: Diffusion2DQuantities begin
  # Diffusion equation
  C ⋅ ∂ₜ(Form0(X)) == C ⋅ d₀(X) ⋅ k ⋅ δ₁(X)
end

X, C, ϕ = Diffusion2D[:X], Diffusion2D[:C], Diffusion2D[:ϕ]
@test codom(C) == Form0(X)
@test codom(ϕ) == DualForm1(X)

end


using JuMP
using CPLEX
include("parser.jl")

function dual_relax(MyFileName::String)
  n, s, t, S, d1, d2, p, ph, d, D = read_instance(MyFileName)

  # Create the model
  m = Model(CPLEX.Optimizer)

  ## Variables
  # x relaxe, y relaxe
  @variable(m, x[1:n,1:n] >= 0) # certains arcs n'existent pas, ils ne sont pas utilises dans les contraintes
  @variable(m, y[1:n] >= 0)
  @variable(m, t1 >= 0)
  @variable(m, t2 >= 0)
  @variable(m, z[1:n,1:n] >= 0, Int) # meme remarque
  @variable(m, zprim[1:n] >= 0, Int)

  ## Constraints
  @constraint(m,[i in 1:n], y[i] <= 1) # relaxation
  @constraint(m,[i in 1:n, j in 1:n], x[i,j] <= 1) # relaxation
  @constraint(m,sum(x[s,j] for j in 1:n if d[s,j] != 0) == 1)
  @constraint(m,sum(x[i,t] for i in 1:n if d[i,t] != 0) == 1)
  @constraint(m,[v in 1:n; v != s && v != t], sum(x[i,v] for i in 1:n if d[i,v] != 0) == sum(x[v,j] for j in 1:n if d[v,j] != 0))
  @constraint(m,[v in 1:n; v != t], y[v] == sum(x[v,j] for j in 1:n if d[v,j] != 0))
  @constraint(m,[i in 1:n; i != t], y[t] == sum(x[i,t] for i in 1:n if d[i,t] != 0))
  @constraint(m, t2*d2 + sum(p[i]*y[i] + 2*zprim[i] for i in 1:n) <= S)
  @constraint(m, [i in 1:n, j in 1:n ; d[i,j] != 0], t1 + z[i,j] >= d[i,j]*x[i,j])
  @constraint(m, [i in 1:n], t2 + zprim[i] >= ph[i]*y[i])

  ## Objective
  @objective(m, Min, d1*t1 + sum(d[i,j]*x[i,j] + D[i,j]*z[i,j] for i in 1:n, j in 1:n if d[i,j] != 0))

  #resolution
  optimize!(m)

  return JuMP.objective_value.(m) # enlever . et JuMP
end

function Trouve_min(P,Q,d)
  min1 = -1
  min2 = -1
  dist = -1
  for p in P
    for q in Q
      if d[p,q] > 0 && (dist == -1 || d[p,q] < dist)
        min1 = p
        min2 = q
        dist = d[p,q]
      end
    end
  end
  return min1,min2,dist
end

function Dijkstra(n,s,t,d)
  M = Array{Int64,2}(zeros(n,3))
  # 1:Distance à s
  # 2:Prédecesseur
  # 3:Nombre de prédecesseurs
  P = [s]
  Q = [i for i in 1:n if i!=s]
  for i in 1:n
    if i != s
      M[i,1] = -1
      M[i,2] = -1
    end
  end
  for q in Q
    if d[s,q] > 0
      M[q,1] = d[q,s]
      M[q,2] = s
      M[q,3] = 1
    end
  end
  while size(Q)[1] > 0
    p,q,dist = Trouve_min(P,Q,d)
    push!(P,q)
    j = 1
    while Q[j] != q
      j = j+1
    end
    deleteat!(Q,j)
    for r in Q
      if (d[q,r] > 0) && (M[q,1] + d[q,r] < M[r,1] || M[r,1] == -1)
        M[r,1] = M[q,1] + d[q,r]
        M[r,2] = q
        M[r,3] = M[q,3]+1
      end
    end
  end
  return M
end

function sortdistance(n,d)
  Sorted_d = Array{Int64}(zeros(0))
  Sorted_i = Array{Int64}(zeros(0))
  Sorted_j = Array{Int64}(zeros(0))
  for i in 1:n
    for j in 1:n
      k = 1
      dist = d[i,j]
      done = false
      while k <= size(Sorted_d)[1] && done == false
        if dist < Sorted_d[k]
          k += 1
        else
          done = true
        end
      end
      insert!(Sorted_d,k,dist)
      insert!(Sorted_i,k,i)
      insert!(Sorted_j,k,j)
    end
  end
  return Sorted_i,Sorted_j
end

function evaluatedist(L,Sorted_i,Sorted_j,d,D,d1)
  distance = 0
  for l in 1:(size(L)[1]-1)
    distance += d[L[l],L[l+1]]
  end
  remaining = d1
  k = size(L)[1] - 1
  acc = 1
  while k > 0 && remaining > 0
    i,j = Sorted_i[acc],Sorted_j[acc]
    for l in 1:(size(L)[1]-1)
      if L[l] == i && L[l+1] == j
        k -= 1
        if D[i,j] < remaining
          distance += d[i,j]*D[i,j]
          remaining -= D[i,j]
        else
          distance += d[i,j] * remaining
          remaining = 0
        end
      end
    end
    acc = acc + 1
  end
  return distance
end

function sortweight(n,ph)
  Sorted = []
  for i in 1:n
    k = 1
    done = false
    while k <= size(Sorted)[1] && done == false
      if ph[i] < ph[Sorted[k]]
        k += 1
      else
        done =true
      end
    end
    insert!(Sorted,k,i)
  end
  return Sorted
end

function evaluateweight(L,Sorted,ph,d2)
  k = size(L)[1]
  acc = 1
  remaining = d2
  weight = 0
  while k > 0 && remaining > 0
    i = Sorted[acc]
    if i in L
      k -= 1
      if 2 < remaining
        weight += 2*ph[i]
        remaining -= 2
      else
        weight += remaining * ph[i]
        remaining = 0
      end
    end
    acc += 1
  end
  return weight
end

function TotalValue(L,Value,Weight,Relativeweight,Order,Limit)
  total = 0
  weight = Array{Int64}(zeros(0))
  relativeweight = Array{Int64}(zeros(0))
  for l in L
    total += value[l]
    push!(weight,Weight[l])
    push!(relativeweight,Relativeweight[l])
  end
  total += evaluate(L,weight,relativeweight,Order,Limit)
  return total
end

function coeff(i,j)
  return 1
end

function heuristic(MyFileName::String)
  n, s, t, S, d1, d2, p, ph, d, D = read_instance(MyFileName)
  global_tol = 5
  instant_tol = 2
  timelimit = 60
  relax_dual_value = dual_relax(MyFileName)
#function heuristic(instant_tol,global_tol,n, s, t, S, d1, d2, p, ph, d, D,timelimit)
  Ones = Array{Int64,2}(zeros(n,n))
  for i in 1:n
    for j in 1:n
    Ones[i,j] = 1
    end
  end
  H1 = Dijkstra(n,s,t,Ones)
  H2 = Dijkstra(n,t,s,Ones)
  Length1 = Array{Int64}(zeros(n))
  for i in [1:n]
    Length1[i] = H1[i,3] + H2[i,3]
  end

  New_d = Array{Int64,2}(zeros(n,n))
  for j in 1:n
    for i in 1:n
      if d[i,j] > 0
        length = (Length1[i]+Length1[j])/2
        extra = min(floor(d1/length),coeff(i,j)*D[i,j]*d[i,j])
        New_d[i,j] = d[i,j] + extra
      end
    end
  end
  H3 = Dijkstra(n,s,t,New_d)
  H4 = Dijkstra(n,t,s,New_d)
  Length2 = Array{Int64}(zeros(n))
  for i in [1:n]
    Length2[i] = H3[i,3] + H4[i,3]
  end

  Sortedweight = sortweight(n,p)
  println(Sortedweight)
  I,J = sortdistance(n,d)
  pmin = p[Sortedweight[n]]
  pmax = p[Sortedweight[1]]
  P_lim = [S - H2[i,3]*pmin for i in 1:n]
  Distance = n * d[I[1],J[1]]
  Path = [s]
  Weight = 0

  DecisionTree = [[s,p[s],global_tol,false]]
  DecisionMatrix = [[] for i in 1:n]
  Current = s
  t0 = time()
  Started = false
  iter = 0
  while (time() - t0 < timelimit) && (Started == false || size(DecisionTree)[1] > 1)
    iter+=1
    k = size(DecisionTree)[1]
    u,tol = DecisionTree[k][1],DecisionTree[k][3]
    for v in 1:n
      if New_d[u,v] > 0 && ((H4[v,3] - H4[u,3]) <= (-1 + min(tol,instant_tol)))
        i = 1
        done = false
        while i < size(DecisionMatrix[u])[1] && done == false
          if (H4[v,3],New_d[u,v]) > (H4[DecisionMatrix[u][i],3],New_d[u,DecisionMatrix[u][i]])
            i += 1
          else
            done = true
          end
        end
        insert!(DecisionMatrix[u],i,v)
      end
    end
    println(DecisionMatrix[s])
    forward = false
    while forward == false && (k > 1 || Started == false)
      Started = true
      while size(DecisionMatrix[u])[1] == 0 && k > 1
        pop!(DecisionTree)
        k -= 1
        if k > 0
          u,tol = DecisionTree[k][1],DecisionTree[k][3]
          if size(DecisionMatrix[u])[1] > 0
            popfirst!(DecisionMatrix[u])
          end
        else
          Current = s
        end
      end
      if size(DecisionMatrix[u])[1] > 0
        v = DecisionMatrix[u][1]
        w = DecisionTree[k][2] + p[v]
        checkedweight = DecisionTree[k][4]
        if checkedweight == false && S - w < 2*d2*pmax
          checkedweight = true
          L = [DecisionTree[i][1] for i in 1:k]
          push!(L,v)
          w += evaluateweight(L,Sortedweight,ph,d2)
        end
        path = [DecisionTree[i][1] for i in 1:size(DecisionTree)[1]]
        if w > P_lim[v] || v in path
         popfirst!(DecisionMatrix[u])
        elseif v == t
          push!(path,v)
          weight = evaluateweight(path,Sortedweight,ph,d2)
          for i in path
            weight += p[i]
          end
          distance = evaluatedist(path,I,J,d,D,d1)
          if weight <= S && distance < Distance
            Path = path
            Weight = weight
            Distance = distance
          end
          popfirst!(DecisionMatrix[u])
        else
          push!(DecisionTree,[v,w,tol-(H4[v,3]-H4[u,3]+1),checkedweight])
          forward = true
          Current = v
        end
      end
    end
  end
  println(Path," ",Weight," ",Distance," ",(Distance/relax_dual_value)-1," ",iter)
  return Path,Weight,Distance,(Distance/relax_dual_value)-1
end


using CompScienceMeshes, BEAST
Γ = readmesh(joinpath(@__DIR__,"sphere2.in"))
@show numcells(Γ)
X = raviartthomas(Γ)

Δt, Nt = 0.6, 200
T = timebasisshiftedlagrange(Δt, Nt, 3)
U = timebasisdelta(Δt, Nt)

V = X ⊗ T
W = X ⊗ U

duration = 20 * Δt
delay = 1.5 * duration
amplitude = 1.0
gaussian = creategaussian(duration, delay, amplitude)

direction, polarisation = ẑ, x̂
E = BEAST.planewave(polarisation, direction, derive(gaussian), 1.0)

SL = TDMaxwell3D.singlelayer(speedoflight=1.0, numdiffs=1)

@hilbertspace j
@hilbertspace j′
tdefie = @discretise SL[j′,j] == -1.0E[j′]   j∈V  j′∈W
xefie = solve(tdefie)

Xefie, Δω, ω0 = fouriertransform(xefie, Δt, 0.0, 2)
ω = collect(ω0 .+ (0:Nt-1)*Δω)
_, i1 = findmin(abs.(ω.-1.0))

ω1 = ω[i1]
ue = Xefie[:,i1] / fouriertransform(gaussian)(ω1)


function optimalAssignments(estimator::SBMEstimator, dataset::Dataset, sbm::SBM; time_limit::Float64=400.0)::Tuple{OptResults, Matrix{Int}}
    return optimalAssignments(dataset, sbm, time_limit=time_limit, verbose=estimator.verbose)
end

function optimalAssignments(dataset::Dataset, sbm::SBM; time_limit::Float64=400.0, verbose::Bool=false)::Tuple{OptResults, Matrix{Int}}
    """
    Solves the SBM optimization problem with its descriptive formulation,
    given the matrix of probabilities w.

    Parameters
    ----------
    dataset : Dataset
        Dataset representing an observed graph.
    sbm : SBM
        Stochastic Block Model

    Returns
    -------
    (obj_value, elapsed, x_opt) : type
        Value of objective, elapsed time, and optimal values for x
    """
    A = dataset.A
    q = dataset.n_communities
    n = dataset.n
    m = dataset.m
    k = dataset.k
    w = sbm.w

    if verbose
        println("q = $q groups\nn = $n nodes\nm = $m edges");
        println("w=")
        println(w)
        println("Time limit: $time_limit seconds")
    end

    # convert from seconds to milli-seconds
    time_limit = Int(round(time_limit*1000))

    w[w .== 0] .= 1e-30
    W = zeros(n,n,q,q);
    for i=1:n, j=1:n, g=1:q, h=1:q
        if A[i,j] != 0
            W[i,j,g,h] = 0.5*( -A[i,j]*log(w[g,h]) + ((k[i]*k[j])/(2*m))*w[g,h] )
        else
            W[i,j,g,h] = 0.5*((k[i]*k[j])/(2*m))*w[g,h]
        end
    end

    ############################
    # TODO: set time limit and verbose options
    model = Model(solver=GLPKSolverMIP(tm_lim=time_limit)) # version 0.18 of JuMP

    # Variables
    @variable(model, 0 <= y[1:n,1:n,1:q,1:q] <= 1)
    @variable(model, x[1:n,1:q], Bin);

    # Objective
    @objective(model, Min, sum(
            (W[i,j,g,h] * y[i,j,g,h])
            for g=1:q, h=1:q, i=1:n, j=1:n));

    # Constraints
    @constraint(model, con1[i=1:n, j=1:n, g=1:q, h=1:q], x[i,g] - y[i,j,g,h] >= 0)
    @constraint(model, con2[i=1:n, j=1:n, g=1:q, h=1:q], x[j,h] - y[i,j,g,h] >= 0)
    @constraint(model, con3[i=1:n, j=1:n, g=1:q, h=1:q], 1 - x[i,g] - x[j,h] + y[i,j,g,h] >= 0)

    # Assignment constraints
    @constraint(model, assign[i = 1:n], sum(x[i,g] for g=1:q) == 1); # assignment constraint

    # Symmetry breaking constraints
    @constraint(model, conObj1, x[1,1] == 1); # object 1 must be in cluster 1
    @constraint(model, conObj[j=2:n], (1 - x[j,2]) <= sum((1 - x[i,1]) for i=2:(j-1)) + x[j,1] ) # if objects 2,...,j-1 are in cluster 1 and object j is not, then object j must be in cluster 2

    # Solve
    start = time()
    status = solve(model) # version 0.18 of JuMP
    solvetime = time() - start
    obj_lb = getobjectivebound(model);
    if isnan(obj_lb)
        obj_lb = -Inf
    end
    obj_ub = getobjectivevalue(model);
    if isnan(obj_ub)
        obj_ub = Inf
    end
    iterations = nothing
    nodecount = nothing
    lazycount = nothing
    ############################
    # Recover variable values
    x = getvalue(x)
    x = round.(x[:,:])
    x[map(v -> isnan(v) , x)] .= 0 # replace NaN values with 0
    x = Int.(x)

    opt_results = OptResults(obj_lb, obj_ub, status, solvetime, iterations, nodecount, lazycount)
    if verbose
        displayResults(opt_results)
    end
    return opt_results, x
end

function MINLP(estimator::SBMEstimator, dataset::Dataset)
    throw("not implemented")
end


# coupling.jl

"""
Implementation of Frohlich's Hamiltonian and coupling parameter α applied to the Polaron model.

See Frohlich 1952:
https://www.jstor.org/stable/99166

"""

"""
frohlich_α(ϵ_optic::Float64, ϵ_static::Float64, freq::Float64, m_eff::Float64)

    Calculates the Frohlich alpha parameter, for a given dielectric constant, frequency (f) of phonon in Hertz, and effective mass (in units of the bare electron mass).

    See Equation 2.37.

"""
function frohlich_α(ϵ_optic, ϵ_static, freq, m_eff)

    ω = 2 * π * freq

    α = 0.5 / (4 * π * ϵ_0) *
       (1 / ϵ_optic - 1 / ϵ_static) *
       (eV^2 / (ħ * ω)) *
       sqrt(2 * m_e * m_eff * ω / ħ)

    return α
end


using BenchmarkTools, FinancialMonteCarlo, DualNumbers

S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 10000;
Nstep = 30;
sigma = dual(0.2, 1.0);
mc = MonteCarloConfiguration(Nsim, Nstep);
toll = 1e-3;

rfCurve = ZeroRate(r);

FwdData = Forward(T * 2.0)
EUData = EuropeanOption(T * 1.1, K)
AMData = AmericanOption(T * 4.0, K)
BarrierData = BarrierOptionDownOut(T * 0.5, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T * 3.3)
AsianFixedStrikeData = AsianFixedStrikeOption(T * 1.2, K)
Model = BlackScholesProcess(sigma, Underlying(S0, d));

optionDatas = [FwdData, EUData, AMData, BarrierData, AsianFloatingStrikeData, AsianFixedStrikeData]

@btime (FwdPrice, EuPrice, AMPrice, BarrierPrice, AsianPrice1, AsianPrice2) = pricer(Model, rfCurve, mc, optionDatas);


# ---
# title: Lines and markers
# cover: assets/lines_and_markers.png
# description: Basic linear and scatter plots.
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

# ## A simple scatter plot

using AlgebraOfGraphics, CairoMakie
set_aog_theme!() #src

df = (x=rand(100), y=rand(100))
xy = data(df) * mapping(:x, :y)
draw(xy)

# ## A simple lines plot

x = range(-π, π, length=100)
y = sin.(x)
df = (; x, y)
xy = data(df) * mapping(:x, :y)
layer = visual(Lines)
draw(layer * xy)

# ## Lines and scatter combined plot

x = range(-π, π, length=100)
y = sin.(x)
df = (; x, y)
xy = data(df) * mapping(:x, :y)
layers = visual(Scatter) + visual(Lines)
draw(layers * xy)

#

x = range(-π, π, length=100)
y = sin.(x)
df1 = (; x, y)
df2 = (x=rand(10), y=rand(10))
layers = data(df1) * visual(Lines) + data(df2) * visual(Scatter)
fg = draw(layers * mapping(:x, :y))

# save cover image #src
mkpath("assets") #src
save("assets/lines_and_markers.png", fg) #src


using Test
using JuliaEuler

@testset "Examples 25 - 50" begin

    @testset "Ex25" begin
        @test ex25(3) == 12
        @test ex25(1000) == 4782

    end

    @testset "Ex26" begin
        @test ex26(10) == 7
        @test ex26(1000) == 983
    end

    @testset "Ex27" begin
        @test ex27(1000,1000) == -59231
        @test ex27(80,1602) == -126479
    end

    @testset "Ex28" begin
        @test ex28(3) == 25
        @test ex28(5) == 101
        @test ex28(7) == 261
        @test ex28(1001) == 669171001
    end

    @testset "Ex29" begin
        @test ex29(5,5) == 15
        @test ex29(100,100) == 9183
    end
end


using GLMNet

function logrange(min, max, num)
  shift = 1.0 - min
  exp(linrange(log(min + shift), log(max + shift), num)) - shift
end

macro validate_and_init_for_stability_selection()
  esc(quote
    @assert 0.0 <= alpha <= 1.0
    X = float64(X)
    y = float64(y)
    N, P = size(X)
    halfsize = ifloor(N / 2)
    counts = zeros(Int64, P, nlambda)
    updatecounts = zeros(Int64, nlambda)
    Xn, xcolnorms = normalize_columns(X)
    Xnw = zeros(Float64, N, P)
    indices = collect(1:N)
  end)
end

function normalize_columns(A)
  acolnorms = map((j) -> norm(A[:,j]), 1:size(A, 2))
  An = broadcast(/, A, acolnorms')
  return (An, acolnorms)
end

macro update_counts(counts, updatecounts, r)
  quote
    for li in 1:length($(r).lambda)
      $(updatecounts)[li] += 1
      for bi in 1:P
        if $(r).betas[bi, li] != 0.0
          $(counts)[bi, li] += 1
        end
      end
    end
  end
end

# Stability selection in the spirit of Meinshausen&Buhlman but with the Complementary Pairs
# variant described in
#   Rajen Shah and Richard J. Samworth, "Variable selection with error control: Another
#   look at Stability Selection", arxiv.org, 2011, http://arxiv.org/pdf/1105.5578.pdf
#
#   selfreqs = stability_selection(X, y; options...)
#
# Inputs:
#  X is the N*P design matrix
#  y is the N*1 variable to predict
#
# Output:
#  selfreqs is the probability that a feature was selected during the first
#             maxit iterations of the GLMNet LASSO algorithm. The features are
#             randomly reweighted in [alpha, 1.0] and the probability is estimated
#             by numbootstraps samples.
#
# Options:
#  maxit is the maximum number of iterations of the LASSO
#  alpha is the minimum weight of a feature
#  numbootstraps is the number of bootstrap samples
#  parallel is true if it should run the
function stability_selection(X::AbstractMatrix, y::AbstractVector;
                             numbootstraps = 50,
                             maxit = 500,
                             alpha = 0.20,
                             nlambda = 25,
                             lambdas = [],
                             options...)

  @validate_and_init_for_stability_selection

  # Run one glmnet fit with only a few lambdas so we can extract their min and max.
  if length(lambdas) == 0
    res = glmnet(Xn, y; nlambda = 3, options...)
    lambdamax, lambdamin = res.lambda[1], res.lambda[end]
    lambdas = reverse(logrange(lambdamin, lambdamax, nlambda))
  end

  # Now take numbootstraps (parallel) samples
  map(1:numbootstraps) do i

    # Reweight each covariate in each round. We can do this via the penalty_factors
    # option instead?!
    weights = alpha + (1.0 - alpha) * rand(P)
    for j in 1:P
      for i in 1:N
        Xnw[i, j] = Xn[i, j] / weights[j]
      end
    end

    # Randomly split sample in two sets
    shuffle!(indices)
    idxs1 = indices[1:halfsize]
    idxs2 = indices[(halfsize+1):(2*halfsize)] # Can't take to end since they need to be same size

    # Run GLMNet on each of the samples and count which variables are selected
    r1 = glmnet(Xnw[idxs1,:], y[idxs1]; lambda = lambdas, standardize = false, intercept = true, maxit = maxit)
    r2 = glmnet(Xnw[idxs2,:], y[idxs2]; lambda = lambdas, standardize = false, intercept = true, maxit = maxit)

    @update_counts(counts, updatecounts, r1)
    @update_counts(counts, updatecounts, r2)
  end

  # Complementary Pairs SS just sums all counts per feature
  #freqs = sum(counts, 2) / sum(updatecounts)

  # Original muhleinbein paper used the max??
  freqs = broadcast(*, counts, (1 ./ updatecounts)')
  # Return the max frequency achieved per feature
  freqs = maximum(freqs, 2)

  freqs[:]
end

# This implements the treshold (tau) selection procedure as described
# at the end of page 10 in the Rajan&Shamworth paper.
#
# Input:
#  P is the number of original features
#  l is max number of low probability features selected, on average
#  B is number of bootstraps, i.e. 2*B is the number of subsets used
function select_treshold(P, l = 0.50, B = 50)
  q = sqrt(0.8*l*P)
  minD(q/P, B)
end

function minD(theta, B, rs = [-0.5, -0.25])
  minimum(vcat(ones(B), tailprobabilities(theta^2, B, rs[1]), tailprobabilities(theta, B, rs[2])))
end

# Calculate the tail probabilities (based on r-concave funcs) for each
# relevant tau. Based on Shah and Shamworth paper and their R code:
#   http://www.statslab.cam.ac.uk/~rds37/papers/r_concave_tail.R
function tailprobabilities(eta, B, r)
  MAXa = 100000
  MINa = 0.00001
  s = -1/r
  etaB = eta * B
  k_start = iceil(2 * etaB) + 1
  output = ones(B)
  if k_start > B
    return(output)
  end

  a_vec = MAXa * ones(B)

end

function indover(v, treshold = 0.70)
  collect(1:length(v))[v .> treshold]
end

N, P, k = (50, 200, 5)
coefs = shuffle(vcat(100*rand(k), zeros(P-k)))
X = randn(N, P)
errors = 1*randn(N)
y = X * coefs .+ errors

tic()
selfreqs = stability_selection(X, y; maxit = 500, alpha = 0.20, nlambda = 20)
t = toq()
reverse(sort(coefs))[1:k]
reverse(sort(selfreqs))[1:(2k)]
indover(coefs, 0.0)
indover(selfreqs, min(N/P, 1.0)) # Shah&Samworth recommends N/P as the treshold, I modified to min(N/P, 1.0) since frequencies <= 1.0


export closure, small_generating_set

################################################################################
#
#  Find identity with respect to operation
#
################################################################################

# It is assumed that the elements have finite order with respect to op.
function find_identity(S, op)
  @assert length(S) > 0
  g = S[1]
  h = g
  while true
    hh = op(h, g)
    if hh == g
      return h
    end
    h = hh
  end
end

################################################################################
#
#  Small generating set
#
################################################################################

function small_generating_set(G, op)
  i = find_identity(G, op)
  return small_generating_set(G, op, i)
end

function small_generating_set(G, op, id)
  orderG = length(G)

  if length(G) == 1
    return G
  end

  function non_trivial_randelem()
    x = rand(G)
    while (x == id)
      x = rand(G)
    end
    return x
  end

  firsttry = 10
  secondtry = 20
  # First try one element
  for i in 1:firsttry
    gen = non_trivial_randelem()
    if length(closure([gen], op, id)) == orderG
      return [gen]
    end
  end

  for i in 1:secondtry
    gens = [non_trivial_randelem(), non_trivial_randelem()]
    if length(closure(gens, op, id)) == orderG
      return unique(gens)
    end
  end

  # Now use that unconditionally log_2(|G|) elements generate G

  b = ceil(Int, log(2, orderG))
  @assert orderG <= 2^b

  j = 0
  while true
    if j > 2^20
      error("Something wrong with generator search")
    end
    j = j + 1
    gens = [non_trivial_randelem() for i in 1:b]
    if length(closure(gens, op, id)) == orderG
      return unique(gens)
    end
  end
end

################################################################################
#
#  Computing closure under group operation
#
################################################################################

# It is assumed that S is nonempty and that the group generated by S under op
# is finite.
function closure(S, op)
  i = find_identity(S, op)
  return closure(S, op, i)
end

function closure(S, op, id)
  if length(S) == 0
    return [id]
  elseif length(S) == 1
    return _closing_under_one_generator(S[1], op, id)
  else
    return _closing_under_generators_dimino(S, op, id)
  end
end

function _closing_under_generators_naive(S, op, id)
  list = push!(copy(S), id)
  stable = false
  while !stable
    stable = true
    for g in list
      for s in S
        m = op(g, s)
        if !(m in list)
          push!(list, m)
          stable = false
        end
      end
    end
  end
  return list
end

function _closing_under_one_generator(x, op, id)
  elements = [x]
  y = x
  while !(y == id)
    y = op(y, x)
    push!(elements, y)
  end
  return elements
end

function _closing_under_generators_dimino(S, op, id)

  t = length(S)
  order = 1
  elements = [id]
  g = S[1]

  while g != id
    order = order +1
    push!(elements, g)
    g = op(g, S[1])
  end

  for i in 2:t
    if !(S[i] in elements)
      previous_order = order
      order = order + 1
      push!(elements, S[i])
      for j in 2:previous_order
        order = order + 1
        push!(elements, op(elements[j], S[i]))
      end

      rep_pos = previous_order + 1
      while rep_pos <= order
        for k in 1:i
          s = S[k]
          elt = op(elements[rep_pos], s)
          if !(elt in elements)
            order = order + 1
            push!(elements, elt)
            for j in 2:previous_order
              order = order + 1
              push!(elements, op(elements[j], elt))
            end
          end
        end
        rep_pos = rep_pos + previous_order
      end
    end
  end
  return elements
end

################################################################################
#
#  Multiplication table
#
################################################################################

# Construct multiplication table of G under op
function _multiplication_table(G, op)
  l = length(G)
  z = Matrix{Int}(l, l)
  for i in 1:l
    for j in 1:l
      p = op(G[i], G[j])
      for k in 1:l
        if p == G[k]
          z[i, j] = k
          break
        end
      end
    end
  end
  return z
end

################################################################################
#
#  Generic group given by multiplication table
#
################################################################################

# The elements are just 1..n and i * j = G.mult_table[i, j]
mutable struct GrpGen
  identity::Int
  order::Int
  mult_table::Array{Int, 2}
  gens::Array{Int, 1}

  function GrpGen(M::Array{Int, 2})
    z = new()
    z.mult_table = M
    z.identity = find_identity(z)
    return z
  end
end

struct GrpGenElem
  group::GrpGen
  i::Int
end

################################################################################
#
#  Compute generic group from anything
#
################################################################################

function generic_group(G, op)
  return GrpGen(_multiplication_table(G, op))
end

################################################################################
#
#  Parent
#
################################################################################

function parent(g::GrpGenElem)
  return g.group
end

################################################################################
#
#  Construct the ith element
#
################################################################################

function getindex(G::GrpGen, i::Int)
  return GrpGenElem(G, i)
end

function ==(g::GrpGenElem, h::GrpGenElem)
  return parent(g) == parent(h) && g.i == h.i
end

################################################################################
#
#  Order
#
################################################################################

function order(G::GrpGen)
  return size(G.mult_table, 1)
end

################################################################################
#
#  Order of an element
#
################################################################################

function order(g::GrpGenElem)
  k = 2
  h = g * g
  while h != g
    h = g * h
    k = k + 1
  end
  return k - 1
end

################################################################################
#
#  Identity
#
################################################################################

function find_identity(G::GrpGen)
  return _find_identity(G.mult_table)
end

function _find_identity(m::Array{Int, 2})
  return find_identity([1], (i, j) -> m[i, j])
end

function id(G::GrpGen)
  return GrpGenElem(G, G.identity)
end

################################################################################
#
#  Multiplication
#
################################################################################

function *(g::GrpGenElem, h::GrpGenElem)
  G = parent(g)
  return GrpGenElem(G, G.mult_table[g.i, h.i])
end

################################################################################
#
#  Inverse
#
################################################################################

function inv(g::GrpGenElem)
  G = parent(g)
  m = G.mult_table
  i = g.i
  l = size(m, 1)
  ide = id(G)
  for j in 1:l
    if m[i, j] == G.identity
      return GrpGenElem(G, j)
    end
  end
  error("Not possible")
end

################################################################################
#
#  String I/O
#
################################################################################

function Base.show(io::IO, G::GrpGen)
  print(io, "Generic group with multiplication table\n")
  println(io, G.mult_table)
end

function Base.show(io::IO, g::GrpGenElem)
  print(io, "Element of generic group ($(g.i))\n")
end

################################################################################
#
#  Is abelian?
#
################################################################################

function isabelian(G::GrpGen)
  return defines_abelian_group(G.mult_table)
end

function defines_abelian_group(m)
  return issymmetric(m)
end

################################################################################
#
#  Isomorphism for specific groups
#
################################################################################

function isisomorphic_to_16T7(G::GrpGen)
  ordershape = [ (1, 1), (2, 3), (4, 12) ]
  # Use algorithm of Tarjahn (or something like that)
  # 16T7 has the following presentation:
  #
  #  $.1^4 = Id($)
  #  $.2^4 = Id($)
  #  $.3^2 = Id($)
  #  $.2^-1 * $.1^2 * $.2^-1 = Id($)
  #  $.1^-1 * $.2^-1 * $.1 * $.2^-1 = Id($)
  #  $.1^-1 * $.3 * $.1 * $.3 = Id($)
  #  $.2^-1 * $.3 * $.2 * $.3 = Id($)

  relations = [[(2, -1), (1, 2), (2, -1)],
               [(1, -1), (2, -1), (1, 1), (2, -1)],
               [(1, -1), (3, 1), (1, 1), (3, 1)],
               [(2, -1), (3, 1), (2, 1), (3, 1)]]

  l = order(G)

  if l != 16
    return false
  end

  if isabelian(G)
    return false
  end

  elements_by_orders = Dict{Int, Array{GrpGenElem, 1}}()

  for i in 1:l
    g = G[i]
    o = order(g)
    if haskey(elements_by_orders, o)
      push!(elements_by_orders[o], g)
    else
      elements_by_orders[o] = [g]
    end
  end

  for (o, no) in ordershape
    if !haskey(elements_by_orders, o)
      return false
    else
      elts = elements_by_orders[o]
      if length(elts) != no
        return false
      end
    end
  end

  ide = id(G)

  img_of_gens = Iterators.product(elements_by_orders[4],
                                  elements_by_orders[4],
                                  elements_by_orders[2])

  for (g1, g2, g3) in img_of_gens
    g2inv = inv(g2)
    g1inv = inv(g1)
    #  $.2^-1 * $.1^2 * $.2^-1 = Id($)
    z = g2inv
    z = z * g1
    z = z * g1
    z = z * g2inv
    if z != ide
      continue
    end
    #  $.1^-1 * $.2^-1 * $.1 * $.2^-1 = Id($)
    z = g1inv
    z = z * g2inv
    z = z * g1
    z = z * g2inv
    if z != ide
      continue
    end
    #  $.1^-1 * $.3 * $.1 * $.3 = Id($)
    z = g1inv
    z = z * g3
    z = z * g1
    z = z * g3
    if z != ide
      continue
    end
    #  $.2^-1 * $.3 * $.2 * $.3 = Id($)
    z = g2inv
    z = z * g3
    z = z * g2
    z = z * g3
    if z != ide
      continue
    end

    cl = closure([g1, g2, g3], *, ide)

    if length(cl) < 16
      continue
    end
    return true
  end
  return false
end

function isisomorphic_to_8T5(G::GrpGen)
  l = order(G)

  if l != 8
    return false
  end

  # First check if it is abelian or not.
  if isabelian(G)
    return false
  end

  # Now G is D_8 (aka D_4 or 8T4) or Q_8
  # But D_8 has 5 elements of order 2, where as Q_8 has only 1 element of order 2.

  z = 0
  for i in 1:l
    g = G[i]
    o = order(g)
    if o == 2
      z += 1
    end
    if z > 1
      return false
    end
  end
  @assert z == 1
  return true
end


#=
Tools for working with Markov Chains

@author : Spencer Lyon, Zac Cranko

@date: 07/10/2014

References
----------

http://quant-econ.net/jl/finite_markov.html
=#
import LightGraphs: DiGraph, period, attracting_components,
                    strongly_connected_components, is_strongly_connected

@inline check_stochastic_matrix(P) = maxabs(sum(P, 2) - 1) < 1e-15 ? true : false

"""
Finite-state discrete-time Markov chain.

It stores useful information such as the stationary distributions, and
communication, recurrent, and cyclic classes, and allows simulation of state
transitions.

##### Fields

- `p::Matrix` The transition matrix. Must be square, all elements must be
positive, and all rows must sum to unity
"""
type MarkovChain{T, TM<:AbstractMatrix, TV<:AbstractVector}
    p::TM # valid stochastic matrix
    state_values::TV

    function MarkovChain(p::AbstractMatrix, state_values)
        n, m = size(p)

        eltype(p) != T &&
            throw(ArgumentError("Types must be consistent with given matrix"))

        n != m &&
            throw(DimensionMismatch("stochastic matrix must be square"))

        any(p.<0) &&
            throw(ArgumentError("stochastic matrix must have nonnegative elements"))

        !check_stochastic_matrix(p) &&
            throw(ArgumentError("stochastic matrix rows must sum to 1"))

        length(state_values) != n &&
            throw(DimensionMismatch("state_values should have $n elements"))

        return new(p, state_values)
    end
end

# Provide constructor that infers T from eltype of matrix
MarkovChain(p::AbstractMatrix, state_values=1:size(p, 1)) =
    MarkovChain{eltype(p), typeof(p), typeof(state_values)}(p, state_values)

"Number of states in the markov chain `mc`"
n_states(mc::MarkovChain) = size(mc.p, 1)

function Base.show{T,TM}(io::IO, mc::MarkovChain{T,TM})
    println(io, "Discrete Markov Chain")
    println(io, "stochastic matrix of type $TM:")
    println(io, mc.p)
end

"""
solve x(P-I)=0 using an algorithm presented by Grassmann-Taksar-Heyman (GTH)

##### Arguments

- `p::Matrix` : valid stochastic matrix

##### Returns

- `x::Matrix`: A matrix whose columns contain stationary vectors of `p`

##### References

The following references were consulted for the GTH algorithm

- W. K. Grassmann, M. I. Taksar and D. P. Heyman, "Regenerative Analysis and
Steady State Distributions for Markov Chains, " Operations Research (1985),
1107-1116.
- W. J. Stewart, Probability, Markov Chains, Queues, and Simulation, Princeton
University Press, 2009.

"""
gth_solve{T<:Integer}(a::AbstractMatrix{T}) =
    gth_solve(convert(AbstractArray{Rational, 2}, a))

# solve x(P-I)=0 by the GTH method
function gth_solve{T<:Real}(original::AbstractMatrix{T})
    # NOTE: the full here will convert a sparse matrix to dense before the gth
    #       algorithm begins. Given the nature of the algorithm this will
    #       almost certainly be more efficient than operating on the sparse
    #       matrix itself.
    a = copy(full(original))
    n = size(a, 1)
    x = zeros(T, n)

    @inbounds for k in 1:n-1
        scale = sum(a[k, k+1:n])
        if scale <= zero(T)
            # There is one (and only one) recurrent class contained in
            # {1, ..., k};
            # compute the solution associated with that recurrent class.
            n = k
            break
        end
        a[k+1:n, k] /= scale

        for j in k+1:n, i in k+1:n
            a[i, j] += a[i, k] * a[k, j]
        end
    end

    # backsubstitution
    x[n] = 1
    @inbounds for k in n-1:-1:1, i in k+1:n
        x[k] += x[i] * a[i, k]
    end

    # normalisation
    x / sum(x)
end

"""
Find the recurrent classes of the `MarkovChain`

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix

##### Returns

- `x::Vector{Vector}`: A `Vector` containing `Vector{Int}`s that describe the
recurrent classes of the transition matrix for p

"""
recurrent_classes(mc::MarkovChain) = attracting_components(DiGraph(mc.p))

communication_classes(mc::MarkovChain) = strongly_connected_components(DiGraph(mc.p))

is_irreducible(mc::MarkovChain) =  is_strongly_connected(DiGraph(mc.p))

is_aperiodic(mc::MarkovChain) = period(mc) == 1

function period(mc::MarkovChain)
    g = DiGraph(mc.p)
    recurrent = attracting_components(g)
    periods   = Int[]

    for r in recurrent
        push!(periods, period(g[r]))
    end
    return lcm(periods)
end
"""
calculate the stationary distributions associated with a N-state markov chain

##### Arguments

- `mc::MarkovChain` : MarkovChain instance containing a valid stochastic matrix

##### Returns

- `dists::Matrix{Float64}`: N x M matrix where each column is a stationary
distribution of `mc.p`

"""
function mc_compute_stationary{T}(mc::MarkovChain{T})
    recurrent = recurrent_classes(mc)

    # unique recurrent class
    length(recurrent) == 1 && return gth_solve(mc.p)

    # more than one recurrent classes
    stationary_dists = Array(T, n_states(mc), length(recurrent))
    for (i, class) in enumerate(recurrent)
        dist        = zeros(T, n_states(mc))
        dist[class] = gth_solve(mc.p[class, class]) # recast to correct dimension
        stationary_dists[:, i] = dist
    end
    return stationary_dists
end

"""
Simulate time series of state transitions of the Markov chain `mc`.

The sample path from the `j`-th repetition of the simulation with initial state
`init[i]` is stored in the `(j-1)*num_reps+i`-th column of the matrix X.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of each simulation.
- `init::Vector{Int}` : Vector containing initial states.
- `;num_reps::Int(1)` : Number of repetitions of simulation for each element
of `init`

##### Returns

- `X::Matrix{Int}` : Array containing the sample paths as columns, of shape
(ts_length, k), where k = length(init)* num_reps

"""
function simulate(mc::MarkovChain, ts_length::Int, init::Vector{Int};
                  num_reps::Int=1)
    k = length(init)*num_reps
    X = Array(Int, ts_length, k)
    X[1, :] = repmat(init, num_reps)

    simulate!(mc, X)
    return X
end

"""
Simulate time series of state transitions of the Markov chain `mc`.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of each simulation.
- `init::Int` : Initial state.
- `;num_reps::Int(1)` : Number of repetitions of simulation

##### Returns

- `X::Matrix{Int}` : Array containing the sample paths as columns, of shape
(ts_length, k), where k = num_reps

"""
simulate(mc::MarkovChain, ts_length::Int, init::Int; num_reps::Int=1) =
    simulate(mc, ts_length, [init], num_reps=num_reps)

"""
Simulate time series of state transitions of the Markov chain `mc`.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of each simulation.
- `;num_reps::Union{Int, Void}(nothing)` : Number of repetitions of simulation.

##### Returns

- `X::Matrix{Int}` : Array containing the sample paths as columns, of shape
(ts_length, k), where k = num_reps

"""
function simulate(mc::MarkovChain, ts_length::Int; num_reps::Int=1)
    init = rand(1:size(mc.p, 1), num_reps)
    return simulate(mc, ts_length, init)
end

"""
Fill `X` with sample paths of the Markov chain `mc` as columns.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `X::Matrix{Int}` : Preallocated matrix of integers to be filled with sample
paths of the markov chain `mc`. The elements in `X[1, :]` will be used as the
initial states.

"""
function simulate!(mc::MarkovChain, X::Matrix{Int})
    n = size(mc.p, 1)

    # NOTE: ensure dense array and transpose before slicing the array. Then
    #       when passing to DiscreteRV use `sub` to avoid allocating again
    p = full(mc.p)'
    P_dist = [DiscreteRV(sub(p, :, i)) for i in 1:n]

    ts_length, k = size(X)

    for i in 1:k
        for t in 1:ts_length-1
            X[t+1, i] = draw(P_dist[X[t]])
        end
    end
    X
end

"""
Simulate time series of state transitions of the Markov chain `mc`.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of each simulation.
- `init_state::Int(rand(1:n_states(mc)))` : Initial state.

##### Returns

- `x::Vector`: A vector of transition indices for a single simulation
"""
function simulation(mc::MarkovChain, ts_length::Int,
                    init_state::Int=rand(1:n_states(mc)))
    simulate(mc, ts_length, init_state; num_reps=1)[:, 1]
end

# ----------------------- #
# simulate_values methods #
# ----------------------- #

function simulate_values(mc::MarkovChain,
                         ts_length::Int,
                         init_state::Vector{Int};
                         num_reps::Int=1)
    k = length(init_state)*num_reps
    X = Array(eltype(mc.state_values), ts_length, k)
    init_state = repmat(init_state, num_reps)
    X[1, :] = mc.state_values[init_state]

    simulate_values!(mc, init_state, X)
    return X
end

function simulate_values(mc::MarkovChain, ts_length::Int, init_state::Int;
                         num_reps::Int=1)
    simulate_values(mc, ts_length, [init_state], num_reps=num_reps)
end

function simulate_values(mc::MarkovChain, ts_length::Int; num_reps::Int=1)
    init_state = rand(1:size(mc.p, 1), num_reps)
    simulate_values(mc, ts_length, init_state)
end

function simulate_values!(mc::MarkovChain, init_state::Vector{Int}, X::Matrix)
    n = size(mc.p, 1)

    # NOTE: ensure dense array and transpose before slicing the array. Then
    #       when passing to DiscreteRV use `sub` to avoid allocating again
    p = full(mc.p)'
    P_dist = [DiscreteRV(sub(p, :, i)) for i in 1:n]

    ts_length, k = size(X)

    for i in 1:k
        i_state = init_state[i]
        for t in 1:ts_length-1
            i_state = draw(P_dist[i_state])
            X[t+1, i] = mc.state_values[i_state]
        end
    end
    X
end

@doc """ Like `simulate(::MarkovChain, args...; kwargs...)`, but instead of
returning integers specifying the state indices, this routine returns the
values of the `mc.state_values` at each of those indices. See docstring
for `simulate` for more information
""" simulate_values, simulate_values!

"""
Simulate time series of state transitions of the Markov chain `mc`.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of each simulation.
- `init_state::Int(rand(1:n_states(mc)))` : Initial state.

##### Returns

- `x::Vector`: A vector of state values along a simulated path
"""
function value_simulation(mc::MarkovChain, ts_length::Int,
                          init_state::Int=rand(1:n_states(mc)))
    simulate_values(mc, ts_length, init_state; num_reps=1)[:, 1]
end


# # SingleLayerQG decaying 2D turbulence with and without finite Rossby radius of deformation
#
#md # This example can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/singlelayerqg_decaying_barotropic_equivalentbarotropic.ipynb).
#
# We use here the `SingleLayerQG` module to simulate decaying two-dimensional turbulence and
# investigate how does a finite Rossby radius of deformation affects its evolution.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add GeophysicalFlows, Printf, Random, Plots"
# ```

# ## Let's begin
# Let's load `GeophysicalFlows.jl` and some other needed packages.
#
using GeophysicalFlows, Printf, Random, Plots

using GeophysicalFlows: peakedisotropicspectrum
using Random: seed!

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical, domain, and simulation parameters
#
# First, we pick some numerical and physical parameters for our model.

n, L  = 128, 2π             # grid resolution and domain length
deformation_radius = 0.35   # the deformation radius
nothing # hide

## Then we pick the time-stepper parameters
    dt = 1e-2  # timestep
nsteps = 4000  # total number of steps
 nsubs = 20    # number of steps between each plot
nothing # hide

# ## Problem setup
# We initialize two problems by providing a set of keyword arguments to the `Problem` constructor.
# The two problems are otherwise the same, except one has an infinite deformation radius, `prob_bqg`,
# and the other has finite deformation radius, `prob_eqbqg`.
prob_bqg = SingleLayerQG.Problem(dev; nx=n, Lx=L,
                                      dt=dt, stepper="FilteredRK4", aliased_fraction=0)
prob_eqbqg = SingleLayerQG.Problem(dev; nx=n, Lx=L, deformation_radius = deformation_radius,
                                      dt=dt, stepper="FilteredRK4", aliased_fraction=0)
nothing # hide

# ## Setting initial conditions

# For initial condition we construct a relative vorticity with energy most energy around total
# wavenumber ``k_0``.
seed!(1234)
k₀, E₀ = 6, 0.5
∇²ψ₀ = peakedisotropicspectrum(prob_bqg.grid, k₀, E₀, mask=prob_bqg.timestepper.filter)
nothing # hide

# `SingleLayerQG` allows us to set up the initial ``q`` for each problem via `set_q!()` function.
# To initialize both `prob_bqg` and `prob_eqbqg` with the same flow, we first use function
# `SingleLayerQG.streamfunctionfrompv!` to get the streamfunction that corresponds to the
# relative vorticity we computed above. This works in the purely barotropic problem, `prob_bqg`
# since in that case the QGPV is simply the relative vorticity.
∇²ψ₀h = rfft(∇²ψ₀)
ψ₀h = @. 0 * ∇²ψ₀h
SingleLayerQG.streamfunctionfrompv!(ψ₀h, ∇²ψ₀h, prob_bqg.params, prob_bqg.grid)
nothing # hide

# and then use the streamfunction to compute the corresponding ``q_0`` for each problem,
q₀_bqg   = irfft(-prob_bqg.grid.Krsq .* ψ₀h, prob_bqg.grid.nx)
q₀_eqbqg = irfft(-(prob_eqbqg.grid.Krsq .+ 1/prob_eqbqg.params.deformation_radius^2) .* ψ₀h, prob_bqg.grid.nx)
nothing # hide

# Now we can initialize our problems with the same flow.
SingleLayerQG.set_q!(prob_bqg, q₀_bqg)
SingleLayerQG.set_q!(prob_eqbqg, q₀_eqbqg)
nothing # hide

# Let's plot the initial vorticity field for each problem. A function that returns relative
# vorticity from each problem's state variable will prove useful. Note that when plotting, we
# decorate the variable to be plotted with `Array()` to make sure it is brought back on the
# CPU when the variable lives on the GPU.
relativevorticity(prob) = irfft(-prob.grid.Krsq .* prob.vars.ψh, prob.grid.nx)

x, y = prob_bqg.grid.x, prob_bqg.grid.y

p_bqg = heatmap(x, y, Array(relativevorticity(prob_bqg)'),
         aspectratio = 1,
                   c = :balance,
                clim = (-40, 40),
               xlims = (-L/2, L/2),
               ylims = (-L/2, L/2),
              xticks = -3:3,
              yticks = -3:3,
              xlabel = "x",
              ylabel = "y",
               title = "barotropic\n ∇²ψ, t=" * @sprintf("%.2f", prob_bqg.clock.t),
          framestyle = :box)

p_eqbqg = heatmap(x, y, Array(relativevorticity(prob_eqbqg)'),
         aspectratio = 1,
                   c = :balance,
                clim = (-40, 40),
               xlims = (-L/2, L/2),
               ylims = (-L/2, L/2),
              xticks = -3:3,
              yticks = -3:3,
              xlabel = "x",
              ylabel = "y",
               title = "equivalent barotropic; deformation radius: " * @sprintf("%.2f", prob_eqbqg.params.deformation_radius) * "\n ∇²ψ, t=" * @sprintf("%.2f", prob_eqbqg.clock.t),
          framestyle = :box)

l = @layout Plots.grid(1, 2)
p = plot(p_bqg, p_eqbqg, layout = l, size = (800, 380))

# ## Time-stepping the `Problem` forward

# Now we time-step both problems forward and animate the relative vorticity in each case.

startwalltime = time()

cfl(prob) = prob.clock.dt * maximum([maximum(prob.vars.u) / prob.grid.dx, maximum(prob.vars.v) / prob.grid.dy])

anim = @animate for j = 0:Int(nsteps/nsubs)
  if j % (1000 / nsubs) == 0
    log_bqg = @sprintf("barotropic; step: %04d, t: %d, cfl: %.2f, walltime: %.2f min",
        prob_bqg.clock.step, prob_bqg.clock.t, cfl(prob_bqg), (time()-startwalltime)/60)
    println(log_bqg)

    log_eqbqg = @sprintf("equivalent barotropic; step: %04d, t: %d, cfl: %.2f, walltime: %.2f min",
        prob_eqbqg.clock.step, prob_eqbqg.clock.t, cfl(prob_eqbqg), (time()-startwalltime)/60)
    println(log_eqbqg)
  end

  p[1][1][:z] = Array(relativevorticity(prob_bqg))
  p[1][:title] = "barotropic\n ∇²ψ, t=" * @sprintf("%.2f", prob_bqg.clock.t)
  p[2][1][:z] = Array(relativevorticity(prob_eqbqg))
  p[2][:title] = "equivalent barotropic; deformation radius: " * @sprintf("%.2f", prob_eqbqg.params.deformation_radius) * "\n ∇²ψ, t=" * @sprintf("%.2f", prob_eqbqg.clock.t)

  stepforward!(prob_bqg, nsubs)
  SingleLayerQG.updatevars!(prob_bqg)

  stepforward!(prob_eqbqg, nsubs)
  SingleLayerQG.updatevars!(prob_eqbqg)
end

mp4(anim, "singlelayerqg_barotropic_equivalentbarotropic.mp4", fps=18)


module Iteration
    include("readinputfile.jl")
    include("shoot.jl")
    include("solve_tf.jl")
    using Printf
    using .Readinputfile
    using .Shoot
    using .Solve_TF

    function construct(opt)
        rif_val = Readinputfile.construct(opt.inpname, opt.usethread)
        data = Readinputfile.readfile(rif_val)

        dx = data.xmax / float(data.grid_num)

        shoot_val = Shoot.construct(dx, data)
        xarray, yarray = Shoot.shootf!(shoot_val)
        solve_tf_param, solve_tf_val = Solve_TF.construct(data, xarray, yarray)
        return data, solve_tf_param, solve_tf_val, yarray
    end

    getNormRD(newarray, oldarray, x) = let
        tmp = newarray .- oldarray
        return sqrt(sum(tmp .* tmp))
    end

    function iteration!(data, solve_tf_param, solve_tf_val, yarray)
        convergence_flag = false

        for iter in 0:data.iteration_maxiter - 1
            newyarray = Solve_TF.solvetf!(solve_tf_param, solve_tf_val, yarray)
            normrd = getNormRD(newyarray, yarray, solve_tf_val.node_x_glo)

            @printf("Iteration # %d: NormRD = %.15f\n", iter + 1, normrd)
            if normrd <= data.iteration_criterion
                convergence_flag = true
                break
            end

            yarray = simple_mixing(data, newyarray, yarray)
        end

        if convergence_flag
            @printf("計算が収束しました．\n")
        else
            @printf("計算が収束しませんでした．プログラムを終了します．\n")
            exit(1)
        end

        return solve_tf_val.node_x_glo
    end

    simple_mixing(data, newarray, oldarray) = let
        return data.iteration_mixing_weight .* newarray .+ (1.0 - data.iteration_mixing_weight) .* oldarray
    end
end


using PushGP: BoolOr, Dup

@testset "Examples from PushGP docs" begin

@testset "Example 1 from Push3 docs" begin
    # Simple example from the page https://faculty.hampshire.edu/lspector/push3-description.html
    # ( 2 3 INTEGER.* 4.1 5.2 FLOAT.+ TRUE FALSE BOOLEAN.OR )
    interp = PushInterpreter()
    p = [2, 3, Mul{Int}(), 4.1, 5.2, Add{Float64}(), true, false, BoolOr()]
    execute!(interp, p)
    @test length(stack(interp, Bool)) == 1
    @test last(stack(interp, Bool)) == true
    @test length(stack(interp, Float64)) == 1
    @test last(stack(interp, Float64)) == 9.3
    @test length(stack(interp, Int)) == 1
    @test last(stack(interp, Int)) == 6
end # @testset "Example 1 from Push3 docs" begin

@testset "Example 2 from Push3 docs" begin
    # Example 2 from the page https://faculty.hampshire.edu/lspector/push3-description.html
    # ( 5 1.23 INTEGER.+ ( 4 ) INTEGER.- 5.67 FLOAT.* )
    interp = PushInterpreter()
    p = [5, 1.23, Add{Int}(), [4], Sub{Int}(), 5.67, Mul{Float64}()]
    execute!(interp, p)
    @test length(stack(interp, Int)) == 1
    @test last(stack(interp, Int)) == 1
    @test length(stack(interp, Float64)) == 1
    @test last(stack(interp, Float64)) == 6.9741
end # @testset "Example 2 from Push3 docs" begin

@testset "Example 3 from Push3 docs" begin
    # Example 3 from the page https://faculty.hampshire.edu/lspector/push3-description.html
    # ( 5 INTEGER.DUP INTEGER.+ )
    interp = PushInterpreter()
    p = [5, Dup{Int}(), Add{Int}()]
    execute!(interp, p)
    @test length(stack(interp, Int)) == 1
    @test last(stack(interp, Int)) == 10
end

end # @testset "Examples from PushGP docs" begin


using ModelingToolkitStandardLibrary.Thermal, ModelingToolkit, OrdinaryDiffEq, Test
@parameters t

# Test HeatCapacitor, TemperatureSensor, RelativeTemperatureSensor, FixedTemperature
@testset "Heat systems" begin
    T, C, G = 10.0, 10.0, 10.0
    @variables final_T(t)
    @named mass1         = HeatCapacitor(C=C)
    @named mass2         = HeatCapacitor(C=C)
    @named th_conductor  = ThermalConductor(G=G)
    @named reltem_sensor = RelativeTemperatureSensor()
    @named T_sensor1     = TemperatureSensor()
    @named T_sensor2     = TemperatureSensor()
    @named tem_src       = FixedTemperature(T=T)

    @info "Building a single-body system..."
    eqs = [
        connect(mass1.hp, th_conductor.hp1)
        connect(th_conductor.hp2, reltem_sensor.hp1)
        connect(reltem_sensor.hp2, tem_src.hp)
    ]
    @named h1 = ODESystem(eqs, t, systems=[mass1, reltem_sensor, tem_src, th_conductor])
    sys = structural_simplify(h1)

    u0 = [
        mass1.T        => 2.0
        th_conductor.T => 10.0
    ]
    prob = ODEProblem(sys, u0, (0, 2.0))
    sol = solve(prob, Rosenbrock23())

    temperatures = reduce(hcat, sol.u)
    # Check if Relative temperature sensor reads the temperature of heat capacitor
    # when connected to a thermal conductor and a fixed temperature source
    @test sol[reltem_sensor.T] == temperatures[1, :] - temperatures[2, :] - sol[tem_src.hp.T]

    @info "Building a two-body system..."
    eqs = [
        connect(T_sensor1.hp, mass1.hp, th_conductor.hp1)
        connect(th_conductor.hp2, mass2.hp, T_sensor2.hp)
        final_T ~ (mass1.C * mass1.T + mass2.C * mass2.T) /
        (mass1.C + mass2.C)
    ]
    @named h2 = ODESystem(eqs, t, [final_T], [],
                        systems=[mass1, mass2, T_sensor1, T_sensor2, th_conductor])
    sys = structural_simplify(h2)

    u0 = [
        mass1.T => 1.0
        mass2.T => 10.0
        final_T => 12
    ]
    prob = ODEProblem(sys, u0, (0, 3.0))
    sol = solve(prob, Rosenbrock23())

    m1, m2 = sol.u[end]
    @test m1 ≈ m2 atol=1e-1
    mass_T = reduce(hcat, sol.u)
    @test sol[T_sensor1.T] == mass_T[1, :]
    @test sol[T_sensor2.T] == mass_T[2, :]
end

# Test HeatFlowSensor, FixedHeatFlow, ThermalResistor, ThermalConductor, ThermalGround
@testset "Heat flow system" begin
    C, G, R = 10, 10, 10
    @named flow_src     = FixedHeatFlow(Q_flow=50, α=100)
    @named mass1        = HeatCapacitor(C=C)
    @named hf_sensor1   = HeatFlowSensor()
    @named hf_sensor2   = HeatFlowSensor()
    @named th_conductor = ThermalConductor(G=G)
    @named th_resistor  = ThermalResistor(R=R)
    @named th_ground    = ThermalGround()

    @info "Building a heat-flow system..."
    eqs = [
        connect(mass1.hp, th_resistor.hp1, th_conductor.hp1)
        connect(th_conductor.hp2, flow_src.hp, hf_sensor1.hp1, hf_sensor2.hp1)
        connect(th_resistor.hp2, hf_sensor1.hp2, hf_sensor2.hp2, th_ground.hp)
    ]
    @named h2 = ODESystem(eqs, t,
                          systems=[mass1, hf_sensor1, hf_sensor2,
                                   th_resistor, flow_src, th_ground, th_conductor])
    sys = structural_simplify(h2)

    u0 = [
        mass1.T            => 10.0
        th_resistor.Q_flow => 1.0
    ]
    prob = ODEProblem(sys, u0, (0, 3.0))
    sol  = solve(prob, Rosenbrock23())

    @test sol[th_conductor.T]*G == sol[th_conductor.Q_flow]
    @test sol[th_conductor.Q_flow] ≈ sol[hf_sensor1.Q_flow] + sol[flow_src.hp.Q_flow]

    @test sol[mass1.T] == sol[th_resistor.T]
    @test sol[th_resistor.T]./R ≈ sol[th_resistor.Q_flow]

end

# Tests ConvectiveResistor and includes FixedTemperature and ThermalResistor
@testset "Piston cylinder wall" begin
    Tᵧ, Tᵪ = 1000, 10 # ᵧ -> gas and ᵪ -> coolant
    Rᵧ, Rᵪ = 50e-4, 10e-4 # R = 1/h; h is convection co-efficient
    R_wall = 1.5e-4
    @named coolant     = ConvectiveResistor(R=Rᵪ)
    @named gas         = ConvectiveResistor(R=Rᵧ)
    @named wall        = ThermalResistor(R=R_wall)
    @named gas_tem     = FixedTemperature(T=Tᵧ)
    @named coolant_tem = FixedTemperature(T=Tᵪ)

    @info "Building a piston-cylinder..."
    eqs = [
        connect(gas_tem.hp, gas.solidport)
        connect(gas.fluidport, wall.hp1)
        connect(wall.hp2, coolant.fluidport)
        connect(coolant.solidport, coolant_tem.hp)
    ]
    @named piston = ODESystem(eqs, t, systems=[gas_tem, wall, gas, coolant, coolant_tem])
    sys = structural_simplify(piston)

    u0 = [
        coolant.dT  => 5.0
        wall.Q_flow => 10.0
    ]
    prob = ODEProblem(sys, u0, (0, 3.0))
    sol = solve(prob, Rosenbrock23())

    # Heat-flow-rate is equal in magnitude
    # and opposite in direction
    @test sol[gas.Q_flow] + sol[coolant.Q_flow] == zeros(length(sol))
end

# Test ConvectiveConductor, BodyRadiation
@testset "Radiator system" begin
    Tᵧ, Tᵪ = 1000, 10 # ᵧ -> gas and ᵪ -> coolant
    R_wall = 10
    G = 0.04
    σ   = 5.6703744191844294e-8 # Stefan-Boltzmann constant

    @named base        = ThermalResistor(R=R_wall)
    @named gas_tem     = FixedTemperature(T=Tᵧ)
    @named coolant_tem = FixedTemperature(T=Tᵪ)
    @named radiator    = BodyRadiation(G=G)
    @named ground      = ThermalGround()
    @named dissipator  = ConvectiveConductor(G=10)

    @info "Building a radiator..."
    eqs = [
        connect(gas_tem.hp, radiator.hp1, base.hp1, dissipator.solidport)
        connect(base.hp2, radiator.hp2, coolant_tem.hp, dissipator.fluidport)
    ]
    @named rad = ODESystem(eqs, t, systems=[base, gas_tem, radiator, dissipator, coolant_tem])
    sys = structural_simplify(rad)

    u0 = [
        base.Q_flow => 10
        dissipator.Q_flow => 10
    ]
    prob = ODEProblem(sys, u0, (0, 3.0))
    sol = solve(prob, Rosenbrock23())

    @test sol[dissipator.dT] == sol[radiator.hp1.T] - sol[radiator.hp2.T]
    rad_Q_flow = G*σ*(Tᵧ^4 - Tᵪ^4)
    @test sol[radiator.Q_flow] == fill(rad_Q_flow, length(sol[radiator.Q_flow]))
end

@testset "Thermal Collector" begin
    @named flow_src    = FixedHeatFlow(Q_flow=50, α=100)
    @named hf_sensor   = HeatFlowSensor()
    @named th_ground   = ThermalGround()
    @named collector   = ThermalCollector(N=2)
    @named th_resistor = ThermalResistor(R=10)
    @named tem_src     = FixedTemperature(T=10)

    @info "Building a heat collector..."
    eqs = [
        connect(flow_src.hp, collector.hp1, th_resistor.hp1)
        connect(tem_src.hp, collector.hp2)
        connect(hf_sensor.hp1, collector.collector_port)
        connect(hf_sensor.hp2, th_ground.hp, th_resistor.hp2)
        ]
    @named coll = ODESystem(eqs, t,
                            systems=[hf_sensor,flow_src, tem_src,
                            collector, th_resistor])
    sys = structural_simplify(coll)

    u0 = [
        th_resistor.Q_flow => 1.0
    ]
    prob = ODEProblem(sys, u0, (0, 3.0))
    sol  = solve(prob, Rosenbrock23())

    @test sol[collector.collector_port.Q_flow] + sol[collector.hp1.Q_flow] + sol[collector.hp2.Q_flow] ==
        zeros(length(sol[collector.collector_port.Q_flow]))
    @test sol[collector.collector_port.T] == sol[collector.hp1.T] == sol[collector.hp2.T]
end


@testset "Arithmetic" begin
    @testset "Sequence" begin
        a = Sequence(ParameterSpace() × (Taylor(1) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1))^1, collect(1.0:13.0))
        b = Sequence(ParameterSpace() × (Taylor(2) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1))^1, collect(1.0:19.0))
        c = Sequence(ParameterSpace() × (Taylor(0) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1))^1, collect(1.0:7.0))
        d = Sequence(ParameterSpace() × (Taylor(2) ⊗ Fourier(0, 1.0) ⊗ Chebyshev(1))^1, collect(1.0:7.0))

        @test +(4.0\(8.0*a*3.0))/3.0 == a + a == a - (-a) == a +̄ a == a -̄ (-a)

        @test a + b == a - (-b) == ladd!(a, copy(b)) == lsub!(a, -b)
        @test a +̄ b == a -̄ (-b) == radd!(copy(a), b) == rsub!(copy(a), -b)

        @test a + c == a - (-c) == radd!(copy(a), c) == rsub!(copy(a), -c)
        @test a +̄ c == a -̄ (-c) == ladd!(a, copy(c)) == lsub!(a, -c)

        @test a + d == a - (-d)
        @test a +̄ d == a -̄ (-d)
    end

    @testset "LinearOperator" begin
        𝒮₁ = ParameterSpace() × (Taylor(1) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1))^1
        𝒮₂ = ParameterSpace() × (Taylor(2) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1))^1
        𝒮₃ = ParameterSpace() × (Taylor(0) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1))^1
        𝒮₄ = ParameterSpace() × (Taylor(2) ⊗ Fourier(0, 1.0) ⊗ Chebyshev(1))^1

        A = LinearOperator(𝒮₁, 𝒮₁, [i+j for i ∈ indices(𝒮₁), j ∈ indices(𝒮₁)])
        B = LinearOperator(𝒮₂, 𝒮₁, [i+j for i ∈ indices(𝒮₁), j ∈ indices(𝒮₂)])
        C = LinearOperator(𝒮₃, 𝒮₁, [i+j for i ∈ indices(𝒮₁), j ∈ indices(𝒮₃)])
        D = LinearOperator(𝒮₄, 𝒮₁, [i+j for i ∈ indices(𝒮₁), j ∈ indices(𝒮₄)])

        @test +(4.0\(8.0*A*3.0))/3.0 == A + A == A - (-A) == A +̄ A == A -̄ (-A)
        @test A == (A - I) + I == -(I - (I + A)) == radd!(rsub!(copy(A), I), I) == -lsub!(I, ladd!(I, copy(A)))

        @test A + B == A - (-B) == ladd!(A, copy(B)) == lsub!(A, -B)
        @test A +̄ B == A -̄ (-B) == radd!(copy(A), B) == rsub!(copy(A), -B)

        @test A + C == A - (-C) == radd!(copy(A), C) == rsub!(copy(A), -C)
        @test A +̄ C == A -̄ (-C) == ladd!(A, copy(C)) == lsub!(A, -C)

        @test A + D == A - (-D)
        @test A +̄ D == A -̄ (-D)

        @test A ∘ A == A*A == A^2
        @test A ∘ B == mul!(similar(B), A, B, true, false)
        @test B ∘ A == mul!(similar(A), B, A, true, false)
        @test A ∘ C == mul!(similar(C), A, C, true, false)
        @test C ∘ A == mul!(similar(A), C, A, true, false)
        @test A ∘ D == mul!(similar(D), A, D, true, false)
        @test D ∘ A == mul!(similar(A), D, A, true, false)
    end

    @testset "Convolution" begin
        function conv(a, b)
            n = fft_size(space(a), space(b))
            space_c = image(*, space(a), space(b))
            c = Sequence(space_c, Vector{Float64}(undef, dimension(space_c)))
            rifft!(c, fft(a, n) .* fft(b, n))
            return c
        end

        a = Sequence(Taylor(1), [1.0, 2.0])
        b = Sequence(Taylor(2), [1.0, 2.0, 3.0])
        @test fft_size(space(a), space(a)) == fft_size(space(a), 2)
        @test conv(a, a) ≈ a * a == a ^ 2
        @test conv(a, b) ≈ a * b
        @test a *̄ a == a ^̄ 2

        a = Sequence(Fourier(1, 1.0), [1.0, 2.0, 3.0])
        b = Sequence(Fourier(2, 1.0), [1.0, 2.0, 3.0, 4.0, 5.0])
        @test fft_size(space(a), space(a)) == fft_size(space(a), 2)
        @test conv(a, a) ≈ a * a == a ^ 2
        @test conv(a, b) ≈ a * b
        @test a *̄ a == a ^̄ 2

        a = Sequence(Chebyshev(1), [1.0, 2.0])
        b = Sequence(Chebyshev(2), [1.0, 2.0, 3.0])
        @test fft_size(space(a), space(a)) == fft_size(space(a), 2)
        @test conv(a, a) ≈ a * a == a ^ 2
        @test conv(a, b) ≈ a * b
        @test a *̄ a == a ^̄ 2

        a = Sequence(Taylor(1) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(1), collect(1.0:12.0))
        b = Sequence(Taylor(2) ⊗ Fourier(0, 1.0) ⊗ Chebyshev(1), collect(1.0:6.0))
        @test fft_size(space(a), space(a)) == fft_size(space(a), 2)
        @test conv(a, a) ≈ a * a == a ^ 2
        @test conv(a, b) ≈ a * b
        @test a *̄ a == a ^̄ 2
    end
end