# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Script with functions to compute statistics for estimating sectoral trade elasticities (method of Caliendo and Parro, 2015)
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

"""
    CP_statistic(X::Matrix, τ_X::Matrix, N::Integer, S::Integer)

The function computes the LHS and RHS statistics of equation (23) in Caliendo and Parro (2015).

# Arguments
- `X::Matrix`: NS×N, origin country-industry destination country intermediate/final demand matrix X.
- `τ_X::Matrix`: NS×N, origin country-industry destination country intermediate/final demand tariff matrix τ_X.
- `N::Integer`: number of origin/destination countries.
- `S::Integer`: number of origin/destination industries.

# Output
- `lhs_temp::Vectpr{Float64}`: (Σᴺ⁻²n(n+1)/2)×1, LHS of eq. (23) in logs for intermediate/final demand per industry.
- `rhs_temp::Vectpr{Float64}`: (Σᴺ⁻²n(n+1)/2)×1, RHS of eq. (23) in logs for intermediate/final demand per industry.

# Examples
```julia-repl
julia> lhs_Z_temp, rhs_Z_temp, lhs_F_temp, rhs_F_temp = create_reg_data(X_temp, F_temp, τ_X_temp, τ_F_temp)
```
"""

function CP_statistic(X::Matrix, τ_X::Matrix, N::Integer, S::Integer)

    # initialize loop
    LHS = Float64[]
    RHS = Float64[]

    for i in 1:N-2 # i -> j
        for j in i+1:N-1 # j -> n
            for n in j+1:N # n -> i
                push!( LHS, ( X[i,j]*X[j,n]*X[n,i] ) / ( X[j,i]*X[n,j]*X[i,n] ) )
                push!( RHS, ( τ_X[i,j]*τ_X[j,n]*τ_X[n,i] ) / ( τ_X[j,i]*τ_X[n,j]*τ_X[i,n] ) )
            end
        end
    end

    # create series in logs
    lhs_temp = log.(LHS)
    rhs_temp = log.(RHS)

    return lhs_temp, rhs_temp
end

# ------------

"""
    elasticity_data(M::Matrix, τ::Matrix, demand::String, N::Integer, S::Integer)

A summary function to collect all industry LHS and RHS statistics of equation (23) in Caliendo and Parro (2015).
    Notice that intermediate and final demand are treated seperately.

# Arguments
- `M::Matrix`: NS×NS, origin country-industry destination country-industry intermediate demand matrix Z.
- `τ::Matrix`: NS×NS, origin country-industry destination country-industry intermediate input tariff matrix τ_Z.
- `N::Integer`: number of origin/destination countries.
- `S::Integer`: number of origin/destination industries.

# Output
- `lhs::Vectpr{Float64}`: S*(Σᴺ⁻²n(n+1)/2)×1, LHS of eq. (23) in logs for intermediate inputs and all industries.
- `rhs::Vectpr{Float64}`: S*(Σᴺ⁻²n(n+1)/2)×1, RHS of eq. (23) in logs for intermediate inputs and all industries.

# Examples
```julia-repl
julia> lhs, rhs = elasticity_data(Z, F, τ_Z, τ_F, N, S)
```
"""

function elasticity_data(M::Matrix, τ::Matrix, demand::String, N::Integer, S::Integer)

    if demand == "intermediate"
        # reduce intermediate demand to origin country-industry destination country
        X = [sum(M[i,j:j+S-1]) for i in 1:N*S, j in 1:S:N*S] # NS×N

    elseif demand == "total" # wrong
        X = [sum(M[i,j:j+S-1]) for i in 1:N*S, j in 1:S:N*S] # NS×N
        X = X .+ F
    elseif demand == "final" # wrong, dont understand how to reshape final demand matrix
        X = M
    else
        println(" × Failure! Specify demand schedule properly! \n")
    end

    τ_X = τ

    # initialize loop
    lhs = Float64[]
    rhs = Float64[]

    for j in 1:S

        # same sector in rows
        X_temp = X[j:S:(N-1)*S+j, :] # N×N
        τ_X_temp = τ_X[j:S:(N-1)*S+j, :] # N×N

        lhs_temp, rhs_temp = CP_statistic(X_temp, τ_X_temp, N, S)

        lhs = [lhs; lhs_temp]
        rhs = [rhs; rhs_temp]

    end

    # zero trade causes inf => create NaN instead which we let regression filter out later
    lhs = ifelse.(isinf.(lhs), NaN, lhs)
    rhs = ifelse.(isinf.(rhs), NaN, rhs)

    # check if computed length equals number of (theoretical) combinations
    combinations = Int64(sum([n*(n+1)/2 for n in 1:N-2])) # number of possible combinations from formula
    check = ([length(lhs) length(rhs)] .== S*combinations) == [1 1]
    if check == true
        println(" ✓ The length of the vector coincides with the theoretical length of: $combinations \n")
    else
        println(" × Failure! The length of the vector does not! coincide with the theoretical length of: $combinations \n")
    end

    return lhs, rhs

end


#=
  From srl-pp-tutorial-wasp-stockholm.pdf
  "Statistical Relational learning and Probabilistic Programming"
  by Luc De Raedt, Anton Dries, and Angelika Kimmig
  https://dtai.cs.kuleuven.be/problog/wasp17-tutorial.html

  slide 394f:
  """
  - Flipping a coin with unknown weight
  - Prior: uniform distribution on [0,1]
  - Observation: 5x heads in a row
  """

  The ProbLog model return the following which corresponds to the
  density of the probabilist
  """
   weight(c1,0.1): 3.3994697e-13
   weight(c1,0.3): 2.1679411e-06
   weight(c1,0.5): 0.0041497433
   weight(c1,0.7): 0.1317485
   weight(c1,0.9): 0.86409959         <----

   weight(c2,0.1): 3.2276726e-06
   weight(c2,0.3): 0.024109997
   weight(c2,0.5): 0.66724754         <----
   weight(c2,0.7): 0.30628626
   weight(c2,0.9): 0.0023529733
  """

  Here we observe 13 tosses and detecting the probability of throwing a head.

  In this version we simulate the ProbLog model more closely
  than in coins_learning.jl

  Result:
  Distributions of variable param (num:0)
  0.90000 =>    8646  (0.864600)
  0.70000 =>    1313  (0.131300)
  0.50000 =>      41  (0.004100)

  Distributions of variable ix (num:0)
  5.00000 =>    8646  (0.864600)
  4.00000 =>    1313  (0.131300)
  3.00000 =>      41  (0.004100)

  Distributions of variable ix1 (num:0)
  0.00000 =>   10000  (1.000000)

  Distributions of variable ix2 (num:0)
  0.00000 =>   10000  (1.000000)

  Distributions of variable ix3 (num:0)
  0.00000 =>    9959  (0.995900)
  1.00000 =>      41  (0.004100)

  Distributions of variable ix4 (num:0)
  0.00000 =>    8687  (0.868700)
  1.00000 =>    1313  (0.131300)

  Distributions of variable ix5 (num:0)
  1.00000 =>    8646  (0.864600)
  0.00000 =>    1354  (0.135400)

  Cf ~/webppl/coins_learning2.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function coins_learning(obs)
    params = [0.1,0.3,0.5,0.7,0.9];
    params_prior = [0.05,0.2,0.5,0.2,0.05];
    len = length(params)

    #  Generate the index for param/weight interval
    #  according to the priors in params_prior
    ix ~ Categorical(simplex(params_prior)) # [0,1,2,3,4]

    #  generate the param value to use for this data point
    obs ~ filldist(flip(params[ix]),length(obs))

    param ~ Dirac(params[ix])
    ix1 ~ Dirac(ix == 1)
    ix2 ~ Dirac(ix == 2)
    ix3 ~ Dirac(ix == 3)
    ix4 ~ Dirac(ix == 4)
    ix5 ~ Dirac(ix == 5)
end

# p: 0.9344
obs = [true for _ in 1:13]
# p: 0.5984
# obs = [true,false,true,true,true,true,true,false,false,true,false,false,true]
model = coins_learning(obs)
num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:param)
show_var_dist_pct(chns,:ix)
show_var_dist_pct(chns,:ix1)
show_var_dist_pct(chns,:ix2)
show_var_dist_pct(chns,:ix3)
show_var_dist_pct(chns,:ix4)
show_var_dist_pct(chns,:ix5)


export backtracking_linesearch

function backtracking_linesearch(model_orig, q, dobs, f_prev, g, proj; alpha=1f0, tau=.1f0, c1=1f-4, maxiter=10, verbose=false)

    # evaluate FWI objective as function of step size alpha
    function objective(alpha,p)
        model.m = proj(model_orig.m + alpha*reshape(p,model.n))

        # Set up linear operator and calculate data residual
        info = JUDI.TimeModeling.Info(prod(model.n), dobs.nsrc, JUDI.TimeModeling.get_computational_nt(q.geometry,dobs.geometry,model))
        F = JUDI.TimeModeling.judiModeling(info,model,q.geometry,dobs.geometry)
        dpred = F*q
        return .5f0*norm(dpred - dobs)^2
    end

    model = deepcopy(model_orig)    # don't modify original model
    p = -g/norm(g,Inf)  # normalized descent direction
    f_new = objective(alpha,p)
    iter = 1
    verbose == true && println("    Iter LS: ", iter, "; ", f_new, " <= ", f_prev + c1*alpha*dot(g,p), "; alpha: ", alpha)

    # sufficient decrease (Armijo) condition
    while f_new > f_prev + c1*alpha*dot(g,p) && iter < maxiter
        alpha *= tau
        f_new = objective(alpha,p)
        iter += 1
        verbose == true && println("    Iter LS: ", iter, "; ", f_new, " <= ", f_prev + c1*alpha*dot(g,p), "; alpha: ", alpha)
    end
    return alpha*p
end

function backtracking_linesearch(model_orig, q, dobs, f_prev, g, proj, objective; alpha=1f0, tau=.1f0, c1=1f-4, maxiter=10, verbose=false)

    model = deepcopy(model_orig)    # don't modify original model
    p = -g/norm(g,Inf)  # normalized descent direction
    model.m = proj(model_orig.m + alpha*reshape(p,model_orig.n))
    f_new = objective(model, q, dobs; compute_gradient=false)
    println("    LS: ",f_new)
    iter = 1
    verbose == true && println("    Iter LS: ", iter, "; ", f_new, " <= ", f_prev + c1*alpha*dot(g,p), "; alpha: ", alpha)

    # sufficient decrease (Armijo) condition
    while f_new > f_prev + c1*alpha*dot(g,p) && iter < maxiter
        alpha *= tau
        model.m = proj(model_orig.m + alpha*reshape(p,model_orig.n))
        f_new = objective(model, q, dobs; compute_gradient=false)
        println("    LS: ",f_new)
        iter += 1
        verbose == true && println("    Iter LS: ", iter, "; ", f_new, " <= ", f_prev + c1*alpha*dot(g,p), "; alpha: ", alpha)
    end
    return alpha*p
end


module PhaseSpaceTools

using Reexport
@reexport using Distributions
@reexport using Statistics
@reexport using LinearAlgebra
@reexport using Parameters

import GSL:sf_laguerre_n

export State, Coherent, Squeezed, Fock, Thermal, Crescent, Bogoliubov
export glauberP, husimiQ, wigner, positiveP, positiveW
export randnc, realnoise, realbridge
export reject, plaguerre, plaguerre_asymptotic, laguerren

include("types.jl")
include("helpers.jl")
include("methods.jl")

end


abstract type Grid end

function _is_monotonic_increasing{T}(x::Array{T, 1})
    if length(x) < 2
        return true
    end

    all(elem > 0 for elem in x[2:end] .- x[1:end-1])
end

struct RegularGrid <: Grid
    left :: Float64
    right :: Float64
    points :: Int
    vertices :: Array{Float64, 1}

    function RegularGrid(left=0, right=1, points=50)
        return new(left, right, points, linspace(left, right, points))
    end
end

struct ArbitraryGrid <: Grid
    vertices :: Array{Float64, 1}
    monotonic :: Bool

    ArbitraryGrid(data) = new(data, _is_monotonic_increasing(data))
end

grid_step(g::RegularGrid) = (g.right - g.left) / g.points

vertices(g::RegularGrid) = g.vertices
vertices(g::ArbitraryGrid) = g.vertices

is_monotonic(g::RegularGrid) = true
is_monotonic(g::ArbitraryGrid) = g.monotonic


export isean

function getweigth(len, index)
    if len === 14 || len === 8
        return index % 2 === 0 ? 1 : 3
    end

    return index % 2 === 0 ? 3 : 1
end

function isean(str::AbstractString)::Bool
    eanReg = r"^(\d{8}|\d{13}|\d{14})$"

    if !contains(str, eanReg)
        return false
    end

    checkDigit = parse(Int, str[end])
    checksum = sum(map(((index, value), ) -> value * getweigth(length(str), index) , enumerate(parse.(Int8, split(str[1:end-1], "")))))
    remainder = 10 - (checksum % 10)
    expectedCheckDigit = remainder < 10 ? remainder : 0

    return checkDigit === expectedCheckDigit
end


function gen_inv_quote(T, N, L)
    q, qa = initial_quote(L::Int, :U)

    for i ∈ 1:N
        lti = small_triangle(i)
        bti = lti+i
        Ui_i = Symbol(:Ui_, bti)
        U_i = Symbol(:U_, bti)
        push!(qa, :($Ui_i = 1 / $U_i))
        for j ∈ i+1:N
            ltj = small_triangle(j)
            Ui_i = Symbol(:Ui_, bti)
            U_i = Symbol(:U_, ltj+i)
            push!(qa, :(Utemp = $U_i * $Ui_i))
            for k ∈ i+1:j-1
                Ui_i = Symbol(:Ui_, small_triangle(k) + i)
                U_i = Symbol(:U_, ltj+k)
                push!(qa, :(Utemp += $U_i * $Ui_i))
            end
            Ui_i = Symbol(:Ui_, ltj+i)
            U_i = Symbol(:U_, ltj+j)
            push!(qa, :($Ui_i = -Utemp / $U_i))
        end
    end

    push!(q.args, :($T( ( @ntuple $L Ui ) )))
    q
end
function gen_ip_inv_quote!(qa, N, L, Ui = :Ui, U = :U, extract = extract_symbol)

    temp_num = 0

    for i ∈ 1:N
        lti = small_triangle(i)
        bti = lti+i
        Ui_i = extract(Ui, bti)
        U_i = extract(U, bti)
        push!(qa, :($Ui_i = 1 / $U_i))
        for j ∈ i+1:N
            ltj = small_triangle(j)
            Ui_i = extract(Ui, bti)
            U_i = extract(U, ltj+i)
            Utemp = Symbol(:t_, temp_num)
            temp_num += 1 #Comment this line on / off to change.
            push!(qa, :($Utemp = $U_i * $Ui_i)) # I'm curious. Woud renaming Utemps help optimization, or just be wasteful?
            for k ∈ i+1:j-1
                Ui_i = extract(Ui, small_triangle(k) + i)
                U_i = extract(U, ltj+k)
                push!(qa, :($Utemp += $U_i * $Ui_i))
            end
            Ui_i = extract(Ui, ltj+i)
            U_i = extract(U, ltj+j)
            push!(qa, :($Ui_i = -$Utemp / $U_i))
        end
    end

    # push!(q.args, :($T( ( @ntuple $L Ui ) )))
    qa
end

@generated function LinearAlgebra.inv(U::AbstractUpperTriangular{T,N,L}) where {T,N,L}
    gen_inv_quote(UpperTriangularMatrix{T,N,L}, N, L)
end
@generated function LinearAlgebra.inv(U::AbstractLowerTriangular{T,N,L}) where {T,N,L}
    gen_inv_quote(LowerTriangularMatrix{T,N,L}, N, L)
end

function gen_extracted_ip_inv_quote!(qa, N, L, Ui = :Ui, U = :U)
    for i ∈ 1:L
        push!(qa, :($(id_symbol(U, i)) = $(U)[$i]))
    end
    gen_ip_inv_quote!(qa, N, L, Ui, U, id_symbol)
    for i ∈ 1:L
        push!(qa, :($(Ui)[$i] = $(id_symbol(Ui, i))))
    end
end

@generated function LinearAlgebra.inv!(U::Union{UpperTriangularMatrix{T,N,L},LowerTriangularMatrix{T,N,L}}) where {T,N,L}
    q, qa = create_quote()
    gen_extracted_ip_inv_quote!(qa, N, L, :U, :U)
    push!(q.args, :U)
    q
end
@generated function LinearAlgebra.inv!(Ui::UpperTriangularMatrix{T,N,L}, U::AbstractUpperTriangular{T,N,L}) where {T,N,L}
    q, qa = create_quote()
    gen_extracted_ip_inv_quote!(qa, N, L, :Ui, :U)
    push!(q.args, :Ui)
    q
end
@generated function LinearAlgebra.inv!(Ui::LowerTriangularMatrix{T,N,L},U::AbstractLowerTriangular{T,N,L}) where {T,N,L}
    q, qa = create_quote()
    gen_extracted_ip_inv_quote!(qa, N, L, :Ui, :U)
    push!(q.args, :Ui)
    q
end


"""
    equispaced_line_mesh(
    x1min::FT,
    x1max::FT,
    n1,
    per = false,
) where {FT<:AbstractFloat}

This function generates an equispaced line mesh with user provided coordinates for left and right
domain vertices, number of mesh elements and periodicity information.
"""
function equispaced_line_mesh(
    x1min::FT,
    x1max::FT,
    n1,
    per = false,
) where {FT<:AbstractFloat}
    x1c = range(x1min, x1max; length = n1 + 1)
    return Mesh1D(x1c, per)
end

function Mesh1D(x1c, per = false)
    I = Int
    FT = eltype(x1c)
    nverts = length(x1c)
    nelems = nverts - 1
    nbndry = per ? 1 : 3
    coordinates = Array{FT}(undef, nverts, 1)
    coordinates[:, 1] .= x1c
    vertex_neighbors = zeros(I, nverts, 4)
    tag_names = (:interior, :left, :right)
    if per
        unique_verts = Array{I}(1:nverts-1)
        boundary_tag_names = tuple(tag_names[1])
        boundary_tags = [0]
        vertex_neighbors[1, :] .= [nelems, 2, 1, 1]
        vertex_neighbors[nverts, :] .= [nelems, 2, 1, 1]
        vertex_boundary = Vector(1:nverts-1)
        vertex_boundary_offset = [1, nverts]
    else
        unique_verts = Array{I}(1:nverts)
        vertex_neighbors[1, :] .= [0, 0, 1, 1]
        vertex_neighbors[nverts, :] .= [nelems, 2, 0, 0]
        if nelems > 1
            boundary_tag_names = tuple(tag_names...)
            boundary_tags = [0, 1, 2]
            vertex_boundary = vcat(Vector(2:nverts-1), 1, nverts)
            vertex_boundary_offset = [1, nverts - 1, nverts, nverts + 1]
        else
            boundary_tag_names = tuple(tag_names[2:end]...)
            boundary_tags = [1, 2]
            vertex_boundary = vcat(1, nverts)
            vertex_boundary_offset = [1, nverts, nverts + 1]
        end
    end
    for vt = 2:nverts-1
        vertex_neighbors[vt, :] .= [vt - 1, 2, vt, 1]
    end
    elem_verts = Array{I}(undef, nelems, 2)
    for el = 1:nelems
        elem_verts[el, :] .= [el, el + 1]
    end
    return Mesh1D(
        nverts,
        nelems,
        nbndry,
        coordinates,
        unique_verts,
        vertex_neighbors,
        vertex_boundary,
        boundary_tags,
        boundary_tag_names,
        vertex_boundary_offset,
        elem_verts,
    )
end


using BitTwiddlingConvenienceFunctions: intlog2, nextpow2, prevpow2
using Test

@testset "BitTwiddlingConvenienceFunctions.jl" begin

  @test all(i ->  intlog2(1 << i) == i, 0:(Int == Int64 ? 53 : 30))
  @test all(i -> nextpow2(i) == i, 0:2)
  for j in 1:10
    l, u = (1<<j)+1, 1<<(j+1)
    @test all(i -> nextpow2(i) == u, l:u)
  end

end


### A Pluto.jl notebook ###
# v0.12.8

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Statistics, DataFrames, CSV
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing4.07"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
data = CSV.read(datadir("purchaseData.csv"), copycols=true)

# ╔═╡ bda159e6-2440-11eb-01a6-f5f4d118214e
mean(data.Price)

# ╔═╡ bda18d2e-2440-11eb-101d-f5c5084a72a4
mean(skipmissing(data.Price))

# ╔═╡ bda23ce4-2440-11eb-0b01-f1e78f176026
coalesce.(data.Grade, "QQ")[1:4]

# ╔═╡ bdac24fc-2440-11eb-0d17-71f2e8eeffc9
first(dropmissing(data,:Price), 4)

# ╔═╡ bdaca91a-2440-11eb-3e00-75b670798eea
sum(ismissing.(data.Name))

# ╔═╡ bdb6f90e-2440-11eb-1e8f-1d6fedde42e0
findall(completecases(data))[1:4]

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing4.07"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═bda159e6-2440-11eb-01a6-f5f4d118214e
# ╠═bda18d2e-2440-11eb-101d-f5c5084a72a4
# ╠═bda23ce4-2440-11eb-0b01-f1e78f176026
# ╠═bdac24fc-2440-11eb-0d17-71f2e8eeffc9
# ╠═bdaca91a-2440-11eb-3e00-75b670798eea
# ╠═bdb6f90e-2440-11eb-1e8f-1d6fedde42e0
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"Return a string of the value as a percentage, rounded to one digit."
perc(x::Real) = @sprintf("%.1f%%", x*100)

"Return a string of `count` and percentage of `total` (in parentheses)."
count_and_perc(count, total) = "$(count)  ($(perc(count/total)))"


using WiltonInts84

using Base.Test
using FixedSizeArrays

include("num_quad.jl")

if !isdefined(:record)
    record = false
end

p1 = Vec(1.0, 0.0, 0.0)
p2 = Vec(0.0, 1.0, 0.0)
p3 = Vec(0.0, 0.0, 0.0)
z  = Vec(0.0,0.0,1.0)

N = 2
T = eltype(Float64)
P = Vec{3,T}
J = Vector{Vec{N+3,T}}()
L = Vector{Vec{N+3,P}}()
I = similar(J)
K = similar(L)

h = 1.2

# code 01 and 00
c = p3
r, R = 0.1, 1.1
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 3
@test length(q.arcs) == 1
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end
#error("stop")

# code 02
c = (p1+p2+p3)/3
r, R = 0.35, 200.0
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 6
@test length(q.arcs) == 3
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# code 10
c = (p1+p2+p3)/3
r, R = 0.1, 0.49
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 3
@test length(q.arcs) == 2
@test length(q.circles) == 1
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# code 11
c = p3
r, R = 0.1, 0.3
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 2
@test length(q.arcs) == 2
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# code 12
c = Vec(0.25,0.1,0.0)
r, R = 0.2, 0.3
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 3
@test length(q.arcs) == 2
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

c = Vec(0.1,0.25,0.0)
r, R = 0.2, 0.3
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 3
@test length(q.arcs) == 2
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

c = Vec(0.15,0.15,0.0)
r, R = 0.2, 0.3
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 4
@test length(q.arcs) == 3
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# code 20
c = p3
r, R = 0.1, 0.9
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 3
@test length(q.arcs) == 3
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

c = (p1+p2+p3)/3
r, R = 0.1, 0.45
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 3
@test length(q.arcs) == 3
@test length(q.circles) == 1
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# code 21 is impossible

# code 22
c = (p1+p2+p3)/3
r, R = 0.35, 0.45
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 6
@test length(q.arcs) == 6
@test length(q.circles) == 0
ctr = contour(p1,p2,p3,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# cover the remaining arc configurations
c = (p1+p2+p3)/3
r, R = 0.5, 0.6
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p3,p1,p2,c,r,R)
@test length(q.segments) == 4
@test length(q.arcs) == 4
@test length(q.circles) == 0
ctr = contour(p3,p1,p2,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p3,p1,p2,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

c = (p1+p2+p3)/3
r, R = 0.5, 0.6
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p2,p3,p1,c,r,R)
@test length(q.segments) == 4
@test length(q.arcs) == 4
@test length(q.circles) == 0
ctr = contour(p2,p3,p1,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p2,p3,p1,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

c = (p1+p2+p3)/3
r, R = 0.5, 0.6
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p3,p1,p2,c,r,R)
@test length(q.segments) == 4
@test length(q.arcs) == 4
@test length(q.circles) == 0
ctr = contour(p3,p1,p2,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p3,p1,p2,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

c = (p1+p2+p3)/3
r, R = 0.5, 0.6
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p3,p2,c,r,R)
@test length(q.segments) == 4
@test length(q.arcs) == 4
@test length(q.circles) == 0
ctr = contour(p1,p3,p2,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p3,p2,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# remaining circle case
c = (p1+p2+p3)/3
r, R = 0.1, 0.2
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 0
@test length(q.arcs) == 0
@test length(q.circles) == 2
ctr = contour(p1,p3,p2,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p3,p2,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# singularity outside, h = 0
c = Vec(0.5,-0.1,0.0)
r, R = 0.2, 0.25
c += Vec(0,0,h)
r, R = sqrt(r^2+h^2), sqrt(R^2+h^2)
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 2
@test length(q.arcs) == 2
@test length(q.circles) == 0
ctr = contour(p1,p3,p2,c,r,R)
A,B = wiltonints(ctr,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p3,p2,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# sphere tangent to triangle at vertex
c = Vec(0.0,1.0,0.1)
r, R = 0.0, 0.1
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 0
@test length(q.arcs) == 0
@test length(q.circles) == 0
A,B = wiltonints(q,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# circle through triangle vertex
c = Vec(-0.11,0.98,0.1)
r, R = 0.0, 0.15
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 1
@test length(q.arcs) == 1
@test length(q.circles) == 0
A,B = wiltonints(q,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

# circle through triangle vertex
c = Vec(-0.11,-0.02,0.1)
r, R = 0.0, 0.15
q = contour(p1,p2,p3,c,r,R)
@test length(q.segments) == 0
@test length(q.arcs) == 0
@test length(q.circles) == 0
A,B = wiltonints(q,c, Val{N})
push!(I,A); push!(K,B)
if record
    P,Q = dblquadints1(p1,p2,p3,c,Val{N},r,R)
    push!(J,P); push!(L,Q)
end

using JLD
fn = joinpath(dirname(@__FILE__),"dblquad1.jld")

# convert data to Julia data type
I = T[I[m][n] for m in eachindex(I), n in 1:length(eltype(I))]
K = T[K[m][n][p] for m in eachindex(K), n in 1:length(eltype(K)), p in 1:3]

if record == true
    J = T[J[m][n] for m in eachindex(J), n in 1:length(eltype(J))]
    L = T[L[m][n][p] for m in eachindex(L), n in 1:length(eltype(L)), p in 1:3]
    jldopen(fn,"w") do file
        write(file, "J", J)
        write(file, "L", L)
    end
else
    J,L = jldopen(fn,"r") do file
        read(file, "J"), read(file, "L")
    end
end

@test maximum(I-J) < 1.0e-5
@test maximum(K-L) < 1.0e-5


using IcanProj
using NetCDF
using DataFrames
using PyPlot
using Statistics
using CSV

function compute_errors(path_simulations, cfg, target_variable)

  file_fine = joinpath(path_simulations, "fsmres_forest", "results_$(cfg)", "$(target_variable)_1km.nc")

  file_coarse = joinpath(path_simulations, "fsmres_forest", "results_$(cfg)", "$(target_variable)_50km.nc")

  df_links = link_results(file_fine, file_coarse)

  data_coarse, data_aggregated, ngrids  = unify_results(file_fine, file_coarse, df_links, target_variable)

  bias = mean(data_coarse .- data_aggregated, dims = 1)

  bias = dropdims(bias, dims = 1)

  rmse = sqrt.(mean((data_coarse .- data_aggregated).^2 , dims = 1))

  rmse = dropdims(rmse, dims = 1)

  return rmse, bias

end

function get_metadata(path_simulations)

  file_fine = joinpath(path_simulations, "fsmres_forest", "results_1", "swe_1km.nc")

  file_coarse = joinpath(path_simulations, "fsmres_forest", "results_1", "swe_50km.nc")

  df_links = link_results(file_fine, file_coarse)

  df_agg = aggregate(df_links, :id_coarse, [mean, std])

  return df_agg

end

function compute_correlations(rmse, bias, df_links, target_variable)

  df = DataFrame()

  ikeep = .!isnan.(df_links[:elev_std])

  df[Symbol("r_" * target_variable * "_elev_for_bias")] = cor(bias[ikeep], df_links[:elev_std][ikeep])

  df[Symbol("r_" * target_variable * "_elev_for_rmse")] = cor(rmse[ikeep], df_links[:elev_std][ikeep])

  return df

end

function run_all()

  path_results = joinpath(dirname(pathof(IcanProj)), "..", "data")

  path_simulations = "/data04/jmg/fsm_simulations/netcdf"

  target_variables = ["swe", "latmo", "hatmo", "rnet"]

  df_links = get_metadata(path_simulations)

  for target_variable in target_variables

    @show target_variable

    df_all = []

    for cfg in 1:32

      @show cfg

      rmse, bias = compute_errors(path_simulations, cfg, target_variable)

      if cfg == 1
        df_all = compute_correlations(rmse, bias, df_links, target_variable)
      else
        df_tmp = compute_correlations(rmse, bias, df_links, target_variable)
        append!(df_all, df_tmp)
      end

    end

    file_final = joinpath(path_results, "correlation_error_$(target_variable).csv")

    df_all |> CSV.write(file_final, delim = ";")

  end

end

run_all()


struct ℝ{n} <: Manifold end

dim(::Type{ℝ{n}}) where n = n

# Creates aliases R1, TR1, etc.
for n = 1:12
    @eval const $(Symbol("R$n")) = ℝ{$n}
    @eval const $(Symbol("TR$n")) = T{ℝ{$n}}
    @eval export $(Symbol("R$n"))
    @eval export $(Symbol("TR$n"))
end
TR{n} = T{ℝ{n}}

# Assume all charts are global for now
isglobal(::Chart{k,ℝ{n}}) where {k,n} = true

emb_to_chart(p, C::Chart{1,ℝ{n}}) where n = p
chart_to_emb(p, C::Chart{1,ℝ{n}}) where n = p

# Not a diffeomorphism for all of R^n, but useful for manual checks
emb_to_chart(p, C::Chart{2,ℝ{n}}) where n = p.^2
chart_to_emb(p, C::Chart{2,ℝ{n}}) where n = p.^(1/2)

# Not a diffeomorphism for all of R^n, but useful for manual checks
emb_to_chart(p, C::Chart{3,ℝ{n}}) where n = p.^3
chart_to_emb(p, C::Chart{3,ℝ{n}}) where n = p.^(1/3)

emb_to_chart(p, C::Chart{4,ℝ{n}}) where n = p.^3 + p
function chart_to_emb(p, C::Chart{4,ℝ{n}}) where n
    a = sqrt(3)*sqrt.(27*p.^2 .+ 4) .+ 9*p
    (a/18).^(1/3) .- ((2/3)./a).^(1/3)
end

global_frame_coeffs(p, C::Chart{1,ℝ{n}}) where n = SMatrix{n,n,eltype(p)}(I)
global_coframe_coeffs(p, C::Chart{1,ℝ{n}}) where n = SMatrix{n,n,eltype(p)}(I)

function global_frame_coeffs(p, C::Chart{I,ℝ{n}}) where {n,I}
    C1 = Chart{1,ℝ{n}}()
    p1 = chart_transition(p, C, C1)
    B1 = global_frame_coeffs(p, C1)
    global_frame_coeffs_chart_transition(p1, B1, C1, C)
end

function global_coframe_coeffs(p, C::Chart{I,ℝ{n}}) where {n,I}
    C1 = Chart{1,ℝ{n}}()
    p1 = chart_transition(p, C, C1)
    β1 = global_frame_coeffs(p, C1)
    global_coframe_coeffs_chart_transition(p1, β1, C1, C)
end

inv_global_frame_coeffs(p, C) = inv(global_coframe_coeffs(p, C))
inv_global_coframe_coeffs(p, C) = inv(global_frame_coeffs(p, C))


mutable struct TimeStepWizard{T}
              cfl :: T
    diffusive_cfl :: T
       max_change :: T
       min_change :: T
           max_Δt :: T
           min_Δt :: T
               Δt :: T
end

"""
    TimeStepWizard(cfl=0.1, max_change=2.0, min_change=0.5, max_Δt=Inf, min_Δt=0.0, Δt=0.01)

A type for calculating adaptive time steps based on capping the CFL number at `cfl`.

On calling `update_Δt!(wizard, model)`, the `TimeStepWizard` computes a time-step such that
``cfl = max(u/Δx, v/Δy, w/Δz) Δt``, where ``max(u/Δx, v/Δy, w/Δz)`` is the maximum ratio
between model velocity and along-velocity grid spacing anywhere on the model grid. The new
`Δt` is constrained to change by a multiplicative factor no more than `max_change` or no
less than `min_change` from the previous `Δt`, and to be no greater in absolute magnitude
than `max_Δt` and no less than `min_Δt`.
"""
TimeStepWizard(; cfl=0.1, diffusive_cfl=Inf, max_change=2.0, min_change=0.5, max_Δt=Inf, min_Δt=0.0, Δt=0.01) =
        TimeStepWizard{typeof(Δt)}(cfl, diffusive_cfl, max_change, min_change, max_Δt, min_Δt, Δt)

"""
    update_Δt!(wizard, model)

Compute `wizard.Δt` given the velocities and diffusivities of `model`, and the parameters
of `wizard`.
"""
function update_Δt!(wizard, model)

    Δt = min(
             wizard.cfl * cell_advection_timescale(model),          # advective
             wizard.diffusive_cfl * cell_diffusion_timescale(model) # diffusive
            )

    # Put the kibosh on if needed
    Δt = min(wizard.max_change * wizard.Δt, Δt)
    Δt = max(wizard.min_change * wizard.Δt, Δt)
    Δt = clamp(Δt, wizard.min_Δt, wizard.max_Δt)

    wizard.Δt = Δt

    return nothing
end

(c::CFL{<:TimeStepWizard})(model) = c.Δt.Δt / c.timescale(model)


using LinearAlgebra: real
using PyPlot, SpecialFunctions, LinearAlgebra
using Peacock
polarisation = TM
eye(M) = Matrix{ComplexF64}(I, M, M)
diag_R(mat) = mat[diagind(mat)]

function mask(M)
    mat = zeros(M,M)
    for i = -(M-1):M-1
        mat = mat + diagm(i=>(-1)^(i)*ones(M-abs(i)))
    end
    return mat
end

r = 0.3
epf(x, y) = x^2 + y^2 < r^2 ? 9 : 1
muf(x,y) = 1
a1 = [1, 0]; a2 = [0, 1]
d1 = 0.001; d2 = 0.001
#geometry = Geometry(epf,muf,a1,a2,d1,d2)
#solver = Solver(geometry,11,11) #Solver里的Kx，Ky和下面的反过来了，相同的时候还是有问题，应该是a1，a2把方向给固定了

m = 5
M = (2*m+1)^2

kx = [i for i in -m:m]
kx = repeat(kx,1, 2*m+1)
ky = kx'
K = hcat(kx[:],ky[:])#为何没有*2pi
Kp = K
#K = hcat(diag_R(solver.Ky),diag_R(solver.Kx))/2/pi
#Kp = K

Kxp = [i for i in Kp[:,1]]
Kxp = repeat(Kxp,1,length(Kp[:,1]))
Kx = copy(Kxp)
Kxp = Kxp'

Kyp = [i for i in Kp[:,2]]
Kyp = repeat(Kyp,1,length(Kp[:,1]))
Ky = copy(Kyp)
Kyp = Kyp'

function geo(kx,ky)
    k=sqrt.(kx.^2 .+ ky.^2)
    eb = 1
    eab=8
    r=0.2
    ep=eab*r*besselj.(1,2*pi*k*r)./k.*cos.(pi*kx)
    ep[k.==0] .= (eb+eab*pi*r^2)
    return ep
end

#ep0 = solver.epc #不一样的原因是Peacock里Kx-Kxp对应的是qs，Ky-Kyp对应的是ps
#ep2 = geo(Ky-Kyp,Kx-Kxp)
epc = geo(Kx-Kxp,Ky-Kyp)
#epc = ep0.*mask(size(ep0)[1])
muc = 1

ii = 1
kym = (-1:0.01:1)*0.5
for omega = 0.4
kTD = []
kTDi = []
kTA = []
for ky in kym

    LHS11 = zeros(ComplexF64, M, M)
    LHS12 = eye(M)
    Kp = K + ones(M) * [0 ky]

    if polarisation == TE
        LHS21 = epc * (omega^2*muc*eye(M) .- inv(epc).*(Kp*transpose(Kp)))
        LHS22 = -epc*(inv(epc).*(K[:,1]*ones(1,M)+ones(M)*K[:,1]'))
        LHS = [LHS11 LHS12; LHS21 LHS22]
        RHS = epc
        label = L"H_z"

    elseif polarisation == TM
        LHS21 = (omega^2*epc - diagm(K[:,1].^2 + (ky .+ K[:,2]).^2))
        LHS22 = diagm(-2*K[:,1])
        LHS = [LHS11 LHS12; LHS21 LHS22]
        RHS = muc

        label = L"E_z"
    end
    kT, VT = eigen(LHS,sortby=x -> -abs(x)) #TODO注意这里的排序可能不同, 排序对结果不影响
    #kT, VT = eigen(LHS)
    VT = VT[1:M, :]
    #归一化，不用归一化也行
    VT = VT * diagm(vec(1 ./sqrt.(sum(abs.(VT).^2,dims=1))))
    maxVT,idVT=findmax(abs.(VT), dims=1)
    VT=VT*diagm(vec(maxVT./VT[idVT]))

    dT=1/omega*diag_R(VT'*inv(RHS)*(diagm(K[:,1])*VT+VT*diagm(kT)))

    i0= (abs.(real.(dT)) .> abs.(imag.(dT))).*((real.(dT)+imag.(dT)) .< 0) + (abs.(real.(dT)) .< abs.(imag.(dT))).*(imag.(kT).< 0)
    i0 = Bool.(i0)

    i3 = []
    for i4 = 0:1
        i0 = .!i0
        i1= 1:2*M
        i1 = i1[i0]
        i2 = sortperm(abs.(real.(kT[i0])))
        buf = sort(abs.(real.(kT[i0])))
        i1=i1[i2]
        i2 = (abs.([imag.(kT[i1]); Inf] - [Inf; imag.(kT[i1])] ) .> 1e-10) .| (abs.([imag.(kT[i1]); Inf]) .<= 1e-10) .| (abs.(abs.([real.(kT[i1]); Inf]) .- abs.([Inf; real.(kT[i1])])) .> 1e-10) .| (abs.(abs.([real.(kT[i1]); Inf]) .+ abs.([Inf; real.(kT[i1])]) .- 1) .>= 0.1)
        i1=i1[i2[1:(length(i2)-1)]]
        i3=[i3; i1[1:2m+1]]
    end
    kT=kT[i3]
    dT=dT[i3]
    VT=VT[:,i3]

    buffer=copy(kT)
    i1 = abs.(imag.(buffer)) .< 1e-6
    buffer1=real.(buffer)
    buffer1[.!i1] .= NaN #剔除所有虚部绝对值大于1e-5的值
    if length(buffer1) < length(kT)
        buffer1[length(buffer1)+1:lenm] .= NaN
    end
#    push!(kTD, buffer1')
	kTD = [kTD; buffer1[:]] #所有纯实数解
    buffer[i1] .= NaN
    if length(buffer) < length(kT)
        buffer[length(buffer)+1:lenm] .= NaN
    end
#    push!(kTDi, buffer')
	kTDi = [kTDi; buffer[:]] #所有复数解
    kTA = [kTA; kT[:]]
end
kTD = reshape(kTD,Int(length(kTD)/length(kym)),length(kym))
kTDi = reshape(kTDi,Int(length(kTDi)/length(kym)),length(kym))
kTA = reshape(kTA,Int(length(kTA)/length(kym)),length(kym))
kTD = transpose(kTD)
kTDi = transpose(kTDi)
kTA = transpose(kTA)
figure()
plot(real(kTA[:]),(kym*ones(1,2*(2m+1)))[:],".")
plot(imag(kTA[:]),(kym*ones(1,2*(2m+1)))[:],".")
xlabel(L"k_x(2\pi/a)")
ylabel(L"k_y(2\pi/a)")
title("ω=$omega(2πc/a)")
legend([L"Re(k_{x})",L"Im(k_x)"])
grid()
xlim([-1,1])
savefig("r_0p3_eps_9_omega_$omega.svg")
close()
end
#kTA
#plot(kTD,".")


"""
This is an example of how to extend CloudScat with new geometrical shapes.
Here we will build a cloud geometry consisting in concentric rings.  This
illustrates both the required extensions to define new geometrical shapes
and also how to include shapes defined implicitly by a mathematical function
(this is S = {x: f(x) < 0}).  This involves repeatedly solving the equation
f(x) = 0 along observation paths and therefore one should expect longer
simulation times.

To use this code run julia and type

julia> include("rings.jl"); run()

The output file will be based on this file's name so to run different simulations
you can just copy this file, change the configuration and run it as above with
the new file name.  You can also run from the command line as e.g.

> julia rings.jl

- - -
2020 Alejandro Luque IAA-CSIC (aluque@iaa.es)

"""

using CloudScat
using Parameters, Roots
const co = CloudScat.constants

# We need to import these names besides those exported by CloudScat
import CloudScat: interfaces!, inside, accum!, Accumulator

"""
Geometrical shape that encapsulates concentric rings of period `L`, amplitude
`A` and base altitude `z0`.  The top of the cloud is z = z0 + A * cos(2πρ/L),
with ρ = √(x²+y²).
"""
struct ConcentricRings <: Shape
    "Period (m)"
    L::Float64

    "Amplitude (m)"
    A::Float64

    "Base altitude (m)"
    z0::Float64
end

# This function provides a way to find intersections between the shape and
# straight lines between two points a and b.  We must find values of s such that
# (1 - s)a + sb is an interface point (0 < s < 1).  For each of these points
# we must call accum!(v, s).  Here we use the find_zeros function from the julia
# Roots package.  Note that in some cases some roots may not be found so use this
# with care.
function interfaces!(v::Accumulator, rings::ConcentricRings, a::Point, b::Point)
    function f(s)
        @unpack A, z0, L = rings
        p = @. (1 - s) * a + s * b
        return p[3] - (A * cos(2π * sqrt(p[1]^2 + p[2]^2) / L) + z0)
    end

    for s in find_zeros(f, 0.0, 1.0)
        accum!(v, s)
    end
end

# This function just checks whether a point is inside the shape.
function inside(rings::ConcentricRings, p::Point)
    @unpack A, z0, L = rings
    p[3] < (A * cos(2π * sqrt(p[1]^2 + p[2]^2) / L) + z0)
end

#
# Once a new shape is defined (here, ConcentricRings) one can use it and combine
# with other shapes to define a cloud geometry.
#

function run()
    params = Params(
        # Number of simulated photons
        N = 400000,

        # The wavelength of the scattered light
        λ = 337 * co.nano,

        # The source is defined as a straight line running from source_a
        # to source_b.  Within this line photons are distributed uniformly with
        # initially random, isotropic directions.  To use a point source set
        # source_b = source_a
        source_a = [0, 0, 10 * co.kilo],
        source_b = [0, 0, 10 * co.kilo],

        # Minimum altitude for observations
        zmin = 7 * co.kilo)

    # A cylinder is defined as Cylinder(zbottom, ztop, xcenter, ycenter, radius)
    c1 = Cylinder(7 * co.kilo,  15 * co.kilo, 0, 0, 20 * co.kilo)

    rings = ConcentricRings(2 * co.kilo, 2 * co.kilo, 10 * co.kilo)

    # We limit the extension of the concentric rings by intersecting the shape
    # with a 20-km-radius cylinder.
    cloud = intersect(c1, rings)

    # Geometry of the full domain.  Photons outside the domain are
    # immediately discarded.
    domain = Cylinder(7 * co.kilo, 60 * co.kilo, 0, 0, 200 * co.kilo)

    # Set a homogeneous cloud droplet density and radius.
    # See variable_droplet_radius.jl for how to deal with inhomogeneous clouds.
    # Radius of the scattering particles
    radius = 10e-6

    # Density of the scattering centers.
    nscat = 100 * co.centi^-3

    composition = Homogeneous(params.λ, nscat, radius)

    # Define the full simulation world
    world = World(cloud, domain, composition)

    # Observers record a time curve and an image of received light.  You can
    # set up as many observers as you wish but they have a significant impact
    # on performance.

    observers = [Observer(
        # Location of the observer
        position = [0, 0, 400 * co.kilo],

        # Sampling interval
        tsample = 1e-5,

        # Collected samples
        nsamples = 10000,

        # fov of the camera (half of the diagonal fov).
        fov = 40,

        # camera pixels
        pixels = 1024)]

    basename = splitext(splitdir(@__FILE__)[2])[1]

    # The saveto option specifies the output file.  Leave it as it is to base
    # the filename on the name of this file specifying the parameters.
    CloudScat.main(params, world, observers,
                   saveto="$(basename).h5")
end

# Check if the code has been 'included' or run from the shell.  In the latter
# case, run the simulation.
isinteractive() || run()


using Statistics

function find_maxima(w, n)
    maxima = []
    for i in (1+n):(length(w)-n)
        if issorted(w[i-n:i], lt=isless) && issorted(w[i:i+n], lt=isless, rev = true)
            append!(maxima, i)
        end
        end
    return maxima
end

function get_regression_αβ(t, W, n)
    α = zeros(axes(W,2))
    β = zeros(axes(W,2))
    for i in axes(W,2)
        m = find_maxima(W[:,i], n)
        #!! bump on tail
        m = m[2:5]
        β[i] = cov(t[m], log.(W[m,i]), corrected=false) / var(t[m], corrected=false)
        α[i] = mean(log.(W[m,i])) - β[i]*mean(t[m])
    end
    return α, β
end


#-------------------------
# importing Libraries
#-------------------------

using DataFrames
using Flux
using Flux: crossentropy
using Flux: @epochs
using Flux: argmax
using CSV
using LinearAlgebra
println("importing Libraries Completed")
#--------------------------
# Loading Data
#--------------------------

MNIST_Train = CSV.read("digit-recognizer/train.csv")  ### TRAIN DATA
MNIST_Test  = CSV.read("digit-recognizer/test.csv"); ### TEST DATA
println("importing Data Completed")

#--------------------------
# Extracting Feature and Label
#--------------------------

Feature = MNIST_Train[:,2:785]
Label   = MNIST_Train[:,1]
Test    = MNIST_Test
println("feature label array made")

#--------------------------
# Distributing Data into TRAIN DEV and TEST set.
#--------------------------

size_of_feature, feature_count = size(Feature)
distribution_perc = 0.8 ## 80%
size_of_train_data = convert(Int64, size_of_feature * distribution_perc)

x_train = Feature[1:size_of_train_data,:]
y_train = Label[1:size_of_train_data]
X_dev   = Feature[size_of_train_data:size_of_feature,:]
y_dev   = Label[size_of_train_data:size_of_feature]
println("distribution completed")

#--------------------------
# One Hot Encoding Of Label
#--------------------------

function one_hot_encoding(label)
	n_size = size(label)[1]
	label_arr = fill(0,( n_size, 10))
	for y in 1:n_size
		label_arr[y,label[y]+1] = 1
	end
	return label_arr
end

#--------------------------
# Function for Reshaping Feature to Image ((,784) => ( ,28, 28, 1 ,1))
#--------------------------

function reshaping(features)
	n_row, n_col = size(features)
	image_arr = fill(0.0,( n_row,28, 28, 1)) ### EMPTY ARRAY OF SIZE (33600, 28 ,28 1) AND (8401, 28, 28, 1)

    for index in 1:n_row
       image_arr[index,:,:,:] = reshape(convert(Array, features[index,:]), 28, 28, 1)/255## RESHAPING AND NORMALIZING BY DIVIDING BY 255\
    end

    return reshape(image_arr, (28,28,1,n_row))
end
#--------------------------
# OHE of Labels
#--------------------------

train_label = one_hot_encoding(y_train)
dev_label   = one_hot_encoding(y_dev)
println("label Encoded")
println(size(train_label))
#--------------------------
# Reshaphing The Data.
#--------------------------

train_image = reshaping(x_train)
dev_image   = reshaping(X_dev)
test_image  = reshaping(Test)
println("reshaping Done")

traindata = []
push!(traindata,(train_image,train_label))
println(size(traindata))
#--------------------------
# MODEL
#--------------------------

model = Chain(
    Conv((2,2), 1=>16, relu),
    x -> maxpool(x, (2,2)),
    Conv((2,2), 16=>8, relu),
    x -> maxpool(x, (2,2)),
    x -> reshape(x, :, size(x, 4)),
    Dense(288, 10),
    softmax)

loss(x, y) = crossentropy(model(x), y)
opt = ADAM(params(model))

@epochs 13 Flux.train!(loss, traindata, opt)

predicted = model(reshape(Float64.(dev_image), (28, 28, 1, size(dev_image)[1])))

accuracy = mean(argmax(predicted)-1 .== test_y)
println(accuracy)


"""
长方形有关的功能
"""

struct AbstractRectangle{T} <: AbstractPolygon{T}
    vertex :: Array{Point2D, 1}
    edges :: Array{Segment, 1}
    center :: Point2D
    width :: Float64
    height :: Float64
end

"""
构造一个平行坐标轴的长方形
"""
function Rectangle(center::Point2D, width, height)
    verts = [
        Point2D(center.x - width/2, center.y + height/2),
        Point2D(center.x - width/2, center.y - height/2),
        Point2D(center.x + width/2, center.y - height/2),
        Point2D(center.x + width/2, center.y + height/2)
    ]
    edges = [
        Segment(verts[1], verts[2]),
        Segment(verts[2], verts[3]),
        Segment(verts[3], verts[4]),
        Segment(verts[4], verts[1])
    ]
    return AbstractRectangle{:AXISRECT}(verts, edges, center, width, height)
end

function Square(center::Point2D, width)
    verts = [
        Point2D(center.x - width/2, center.y + width/2),
        Point2D(center.x - width/2, center.y - width/2),
        Point2D(center.x + width/2, center.y - width/2),
        Point2D(center.x + width/2, center.y + width/2)
    ]
    edges = [
        Segment(verts[1], verts[2]),
        Segment(verts[2], verts[3]),
        Segment(verts[3], verts[4]),
        Segment(verts[4], verts[1])
    ]
    return AbstractRectangle{:AXISSQUARE}(verts, edges, center, width, width)
end

"""
长方型的面积
"""
function area(rect::AbstractRectangle{N}) where N
    return rect.height * rect.width
end

#display函数在中间加::MIME"text/plain"
function Base.show(io::IO, tri::AbstractRectangle{N}) where N
    name = Dict(
        :AXISRECT => "Rectangle",
        :AXISSQUARE => "Square"
    )
    @printf(io, "%s\n center: %s\n width: %.6f\n height: %.6f",
    name[N], tri.center, tri.width, tri.height)
end


function timestep()
  dt = dt_max
  for i in active_particle
    @assert sph[i].P > 0

    s = index_list[i]
    vmax = 0
    for t = (s - Nngb) : (s + Nngb)
      t = order_periodic(t)
      k = index_order[t]
      vsig = cs(i) + cs(k) - 3 * eij(i, k) * vij(i, k)
      if vsig > vmax
        vmax = vsig
      end
    end
    dt_i = cfl * sph[i].hsml / vmax
    @assert dt_i > 0
    dt = min(dt, dt_i)
  end
  return dt
end

function update_time(t, dt)
  global t += dt
end


import Base: rand,
             ∈

export Ballp

"""
    Ballp{N<:AbstractFloat, VN<:AbstractVector{N}} <: AbstractCentrallySymmetric{N}

Type that represents a ball in the p-norm, for ``1 ≤ p ≤ ∞``.

It is defined as the set

```math
\\mathcal{B}_p^n(c, r) = \\{ x ∈ \\mathbb{R}^n : ‖ x - c ‖_p ≤ r \\},
```
where ``c ∈ \\mathbb{R}^n`` is its center and ``r ∈ \\mathbb{R}_+`` its radius.
Here ``‖ ⋅ ‖_p`` for ``1 ≤ p ≤ ∞`` denotes the vector ``p``-norm, defined as
``‖ x ‖_p = \\left( \\sum\\limits_{i=1}^n |x_i|^p \\right)^{1/p}`` for any
``x ∈ \\mathbb{R}^n``.

### Fields

- `p`      -- norm as a real scalar
- `center` -- center of the ball as a real vector
- `radius` -- radius of the ball as a scalar (``≥ 0``)

## Notes

The special cases ``p=1``, ``p=2`` and ``p=∞`` fall back to the specialized
types `Ball1`, `Ball2` and `BallInf`, respectively.

### Examples

A five-dimensional ball in the ``p=3/2`` norm centered at the origin of radius
0.5:

```jldoctest ballp_constructor
julia> B = Ballp(3/2, zeros(5), 0.5)
Ballp{Float64,Array{Float64,1}}(1.5, [0.0, 0.0, 0.0, 0.0, 0.0], 0.5)

julia> dim(B)
5
```

We evaluate the support vector in direction ``[1,2,…,5]``:

```jldoctest ballp_constructor
julia> σ([1., 2, 3, 4, 5], B)
5-element Array{Float64,1}:
 0.013516004434607558
 0.05406401773843023
 0.12164403991146802
 0.21625607095372093
 0.33790011086518895
```
"""
struct Ballp{N<:AbstractFloat, VN<:AbstractVector{N}} <: AbstractCentrallySymmetric{N}
    p::N
    center::VN
    radius::N

    # default constructor with domain constraint for radius and p
    function Ballp(p::N, center::VN, radius::N) where {N<:Real, VN<:AbstractVector{N}}
        @assert radius >= zero(N) "radius must not be negative"
        @assert p >= 1 "p must not be less than 1"
        if p == Inf
            return BallInf(center, radius)
        elseif p == 2
            return Ball2(center, radius)
        elseif p == 1
            return Ball1(center, radius)
        else
            return new{N, VN}(p, center, radius)
        end
    end
end

isoperationtype(::Type{<:Ballp}) = false
isconvextype(::Type{<:Ballp}) = true

# --- AbstractCentrallySymmetric interface functions ---

"""
    center(B::Ballp{N}) where {N<:AbstractFloat}

Return the center of a ball in the p-norm.

### Input

- `B` -- ball in the p-norm

### Output

The center of the ball in the p-norm.
"""
function center(B::Ballp{N}) where {N<:AbstractFloat}
    return B.center
end

# --- LazySet interface functions ---

"""
    σ(d::AbstractVector{N}, B::Ballp{N}) where {N<:AbstractFloat}

Return the support vector of a `Ballp` in a given direction.

### Input

- `d` -- direction
- `B` -- ball in the p-norm

### Output

The support vector in the given direction.
If the direction has norm zero, the center of the ball is returned.

### Algorithm

The support vector of the unit ball in the ``p``-norm along direction ``d`` is:

```math
σ_{\\mathcal{B}_p^n(0, 1)}(d) = \\dfrac{\\tilde{v}}{‖\\tilde{v}‖_q},
```
where ``\\tilde{v}_i = \\frac{|d_i|^q}{d_i}`` if ``d_i ≠ 0`` and
``tilde{v}_i = 0`` otherwise, for all ``i=1,…,n``, and ``q`` is the conjugate
number of ``p``.
By the affine transformation ``x = r\\tilde{x} + c``, one obtains that
the support vector of ``\\mathcal{B}_p^n(c, r)`` is

```math
σ_{\\mathcal{B}_p^n(c, r)}(d) = \\dfrac{v}{‖v‖_q},
```
where ``v_i = c_i + r\\frac{|d_i|^q}{d_i}`` if ``d_i ≠ 0`` and ``v_i = 0``
otherwise, for all ``i = 1, …, n``.
"""
function σ(d::AbstractVector{N}, B::Ballp{N}) where {N<:AbstractFloat}
    p = B.p
    q = p/(p-1)
    v = similar(d)
    @inbounds for (i, di) in enumerate(d)
        v[i] = di == zero(N) ? di : abs.(di).^q / di
    end
    vnorm = norm(v, p)
    svec = vnorm != zero(N) ? @.(B.center + B.radius * (v/vnorm)) : B.center
    return svec
end

"""
    ∈(x::AbstractVector{N}, B::Ballp{N}) where {N<:AbstractFloat}

Check whether a given point is contained in a ball in the p-norm.

### Input

- `x` -- point/vector
- `B` -- ball in the p-norm

### Output

`true` iff ``x ∈ B``.

### Notes

This implementation is worst-case optimized, i.e., it is optimistic and first
computes (see below) the whole sum before comparing to the radius.
In applications where the point is typically far away from the ball, a fail-fast
implementation with interleaved comparisons could be more efficient.

### Algorithm

Let ``B`` be an ``n``-dimensional ball in the p-norm with radius ``r`` and let
``c_i`` and ``x_i`` be the ball's center and the vector ``x`` in dimension
``i``, respectively.
Then ``x ∈ B`` iff ``\\left( ∑_{i=1}^n |c_i - x_i|^p \\right)^{1/p} ≤ r``.

### Examples

```jldoctest
julia> B = Ballp(1.5, [1., 1.], 1.)
Ballp{Float64,Array{Float64,1}}(1.5, [1.0, 1.0], 1.0)

julia> [.5, -.5] ∈ B
false

julia> [.5, 1.5] ∈ B
true
```
"""
function ∈(x::AbstractVector{N}, B::Ballp{N}) where {N<:AbstractFloat}
    @assert length(x) == dim(B)
    sum = zero(N)
    for i in eachindex(x)
        sum += abs(B.center[i] - x[i])^B.p
    end
    return sum^(one(N)/B.p) <= B.radius
end

"""
    rand(::Type{Ballp}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing)

Create a random ball in the p-norm.

### Input

- `Ballp` -- type for dispatch
- `N`     -- (optional, default: `Float64`) numeric type
- `dim`   -- (optional, default: 2) dimension
- `rng`   -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`  -- (optional, default: `nothing`) seed for reseeding

### Output

A random ball in the p-norm.

### Algorithm

The center and radius are normally distributed with mean 0 and standard
deviation 1.
Additionally, the radius is nonnegative.
The p-norm is a normally distributed number ≥ 1 with mean 1 and standard
deviation 1.
"""
function rand(::Type{Ballp};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing)
    rng = reseed(rng, seed)
    p = one(N) + abs(randn(rng, N))
    center = randn(rng, N, dim)
    radius = abs(randn(rng, N))
    return Ballp(p, center, radius)
end

"""
    translate(B::Ballp{N}, v::AbstractVector{N}) where {N<:AbstractFloat}

Translate (i.e., shift) a ball in the p-norm by a given vector.

### Input

- `B` -- ball in the p-norm
- `v` -- translation vector

### Output

A translated ball in the p- norm.

### Algorithm

We add the vector to the center of the ball.
"""
function translate(B::Ballp{N}, v::AbstractVector{N}) where {N<:AbstractFloat}
    @assert length(v) == dim(B) "cannot translate a $(dim(B))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    return Ballp(B.p, center(B) + v, B.radius)
end


using Keldysh, Test

@testset "dos" begin
  dos_integrator = Keldysh.GaussKronrodDOSIntegrator()
  # Flat DOS
  let ν=5.0, D=2.0, dos = Keldysh.flat_dos(ν=ν, D=D)
    α = ν*D
    moments = [dos_integrator(ω -> ω^n, dos) for n = 0:3]
    moments_ref = (1/π)*[D*(1 + coth(α)),
                         0.,
                         D^3*(π^2 + α^2)*(1 + coth(α))/(3α^2),
                         0.]
    @test isapprox(moments, moments_ref, atol=1e-10, rtol=1e-10)
  end

  # Gaussian DOS
  let ν = 2.0, ϵ = 1.0, dos = Keldysh.gaussian_dos(ν=ν, ϵ=ϵ)
    moments = [dos_integrator(ω -> ω^n, dos) for n = 0:6]
    moments_ref = [1.0,
                   ϵ,
                   ϵ^2 + 2ν,
                   ϵ^3 + 6ϵ*ν,
                   ϵ^4 + 12ϵ^2*ν + 12*ν^2,
                   ϵ^5 + 20ϵ^3*ν + 60ϵ*ν^2,
                   ϵ^6 + 30ϵ^4*ν + 180ϵ^2*ν^2 + 120ν^3
                  ]
    @test isapprox(moments, moments_ref, atol=1e-10, rtol=1e-10)
  end

  # Bethe DOS
  let t=2.0, dos = Keldysh.bethe_dos(t=t)
    moments = [dos_integrator(ω -> ω^n, dos) for n = 0:6]
    moments_ref = [1.0, 0., t^2, 0., 2t^4, 0., 5t^6]
    @test isapprox(moments, moments_ref, atol=1e-10, rtol=1e-10)
  end

  # Linear chain DOS
  let t=2.0, dos = Keldysh.chain_dos(t=t)
    moments = [dos_integrator(ω -> ω^n, dos) for n = 0:6]
    moments_ref = [1.0, 0., 2t^2, 0., 6t^4, 0., 20t^6]
    @test isapprox(moments, moments_ref, atol=1e-10, rtol=1e-10)
  end

  # Square lattice DOS
  let t=2.0, dos = Keldysh.square_dos(t=t)
    moments = [dos_integrator(ω -> ω^n, dos) for n = 0:6]
    moments_ref = [1.0, 0., 4t^2, 0., 36t^4, 0., 400t^6]
    @test isapprox(moments, moments_ref, atol=1e-10, rtol=1e-10)
  end
end


using Base.Test
using QuantumOptics

@testset "basis" begin

shape1 = [5]
shape2 = [2, 3]
shape3 = [6]

b1 = GenericBasis(shape1)
b2 = GenericBasis(shape2)
b3 = GenericBasis(shape3)

@test b1.shape == shape1
@test b2.shape == shape2
@test b1 != b2
@test b1 != FockBasis(2)
@test b1 == b1

@test tensor(b1) == b1
comp_b1 = tensor(b1, b2)
comp_uni = b1 ⊗ b2
comp_b2 = tensor(b1, b1, b2)
@test comp_b1.shape == [prod(shape1), prod(shape2)]
@test comp_uni.shape == [prod(shape1), prod(shape2)]
@test comp_b2.shape == [prod(shape1), prod(shape1), prod(shape2)]

@test b1^3 == CompositeBasis(b1, b1, b1)
@test (b1⊗b2)^2 == CompositeBasis(b1, b2, b1, b2)
@test_throws ArgumentError b1^(0)

comp_b1_b2 = tensor(comp_b1, comp_b2)
@test comp_b1_b2.shape == [prod(shape1), prod(shape2), prod(shape1), prod(shape1), prod(shape2)]
@test comp_b1_b2 == CompositeBasis(b1, b2, b1, b1, b2)

@test_throws ArgumentError tensor()
@test comp_b2.shape == tensor(b1, comp_b1).shape
@test comp_b2 == tensor(b1, comp_b1)

@test_throws ArgumentError ptrace(comp_b1, [1, 2])
@test ptrace(comp_b2, [1]) == ptrace(comp_b2, [2]) == comp_b1 == ptrace(comp_b2, 1)
@test ptrace(comp_b2, [1, 2]) == ptrace(comp_b1, [1])
@test ptrace(comp_b2, [2, 3]) == ptrace(comp_b1, [2])

comp1 = tensor(b1, b2, b3)
comp2 = tensor(b2, b1, b3)
@test permutesystems(comp1, [2,1,3]) == comp2

@test !QuantumOptics.bases.equal_bases([b1, b2], [b1, b3])
@test !QuantumOptics.bases.multiplicable(comp1, b1 ⊗ b2 ⊗ NLevelBasis(prod(b3.shape)))

end # testset


@testset "EuclideanVector"  begin

#---

using ACE
using Random, Printf, Test, LinearAlgebra, ACE.Testing
using ACE: evaluate, evaluate_d, SymmetricBasis, NaiveTotalDegree, PIBasis
using ACE.Random: rand_rot, rand_refl

# construct the 1p-basis
D = NaiveTotalDegree()
maxdeg = 6
ord = 3

B1p = ACE.Utils.RnYlm_1pbasis(; maxdeg=maxdeg, D = D)

# generate a configuration
nX = 10
Xs = rand(EuclideanVectorState, B1p.bases[1], nX)
cfg = ACEConfig(Xs)

#---

@info("SymmetricBasis construction and evaluation: EuclideanVector")

φ = ACE.EuclideanVector(Complex{Float64})
pibasis = PIBasis(B1p, ord, maxdeg; property = φ, isreal=false)
basis = SymmetricBasis(pibasis, φ)

BB = evaluate(basis, cfg)

# a stupid but necessary test
BB1 = basis.A2Bmap * evaluate(basis.pibasis, cfg)
println(@test isapprox(BB, BB1, rtol=1e-10))

##

@info("Test equivariance properties")

tol = 1E-9

@info("check for rotation, permutation and inversion equivariance")
for ntest = 1:20
   Xs = rand(EuclideanVectorState, B1p.bases[1], nX)
   BB = evaluate(basis, ACEConfig(Xs))
   Q = rand([-1,1]) * ACE.Random.rand_rot()
   Xs_rot = Ref(Q) .* shuffle(Xs)
   BB_rot = evaluate(basis, ACEConfig(Xs_rot))
   print_tf(@test all([ norm(Q' * b1 - b2) < tol
                        for (b1, b2) in zip(BB_rot, BB)  ]))
end
println()

##

@info(" ... derivatives")
for ntest = 1:30
   Us = randn(SVector{3, Float64}, length(Xs))
   C = randn(typeof(φ.val), length(basis))
   F = t -> sum( sum(c .* b.val)
                 for (c, b) in zip(C, ACE.evaluate(basis, ACEConfig(Xs + t[1] * Us))) )
   dF = t -> [ sum( sum(c .* db)
                    for (c, db) in zip(C, ACE.evaluate_d(basis, ACEConfig(Xs + t[1] * Us)) * Us) ) ]
   print_tf(@test fdtest(F, dF, [0.0], verbose=false))
end
println()

##

end


######
##FEM Heat dx Convergence Tests
######
using FiniteElementDiffEq, DiffEqDevTools

#Travis CI Test Setting
#Not good plots, but quick for unit tests
cs = cs_femheat_moving_dx

alg=FEMDiffEqHeatEuler(); println(alg) #Unstable due to μ
sim = test_convergence(cs,alg)

alg=FEMDiffEqHeatImplicitEuler(); println(alg)
sim2 = test_convergence(cs,alg)

alg=FEMDiffEqHeatCrankNicholson(); println(alg) #Bound by spatial discretization error at low dt, decrease dx for full convergence
cs = cs_femheat_moving_faster_dt
sim3 = test_convergence(cs,alg)

#plot(plot(sim),plot(sim2),plot(sim3),layout=@layout([a b c]),size=(1200,400))

#Returns true if all converge approximately dx^2
minimum([sim.𝒪est[:L2],sim2.𝒪est[:L2],sim3.𝒪est[:L2]] - 2 .<.1)


using StatsBase
using Base.Test

fnecdf = ecdf(randn(10000000))
@test isapprox(fnecdf([-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96]),
    [0.025, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.975], atol=1e-3)
@test isapprox(fnecdf(1.96), 0.975, atol=1e-3)
@test fnecdf([-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96]) ≈
    map(fnecdf, [-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96])

fnecdf = ecdf([0.5])
@test fnecdf([zeros(5000); ones(5000)]) == [zeros(5000); ones(5000)]


#
#      Sn - The manifold of the n-dimensional sphere
#  Point is a Point on the n-dimensional sphere.
#
import LinearAlgebra: norm, dot, nullspace
import Base: exp, log, show, cat
export Sphere, SnPoint, SnTVector,show, getValue
export distance, dot, exp, log, manifoldDimension, norm
export randomMPoint, opposite, parallelTransport, zeroTVector
export validateMPoint, validateTVector
export zeroTVector, typeofMPoint, typeofTVector
#
# Type definitions
#
@doc doc"""
    Sphere <: Manifold

The manifold $\mathcal M = \mathbb S^n$ of unit vectors in $\mathbb R^{n+1}$.
This manifold is a matrix manifold (see [`IsMatrixM`](@ref)) and embedded (see
[`IsEmbeddedM`](@ref)).

# Abbreviation

`Sn`

# Constructor

    Sphere(n)

generate the sphere $\mathbb S^n$

`
Its abbreviation is `Sn`.
"""
struct Sphere <: Manifold
  name::String
  dimension::Int
  abbreviation::String
  Sphere(dimension::Int) = new("$dimension-Sphere",dimension,"Sn($(dimension-1))")
end
@doc doc"""
    SnPoint <: MPoint

A point $x$ on the manifold $\mathcal M = \mathbb S^n$ represented by a unit
vector from $\mathbb R^{n+1}$
"""
struct SnPoint{T<:AbstractFloat} <: MPoint
  value::Vector{T}
  SnPoint{T}(value::Vector{T}) where {T <: AbstractFloat} = new(value)
end
SnPoint(value::Vector{T}) where {T <: AbstractFloat} = SnPoint{T}(value)
getValue(x::SnPoint) = x.value;

@doc doc"""
    SnTVector <: TVector

A tangent vector $\xi \in T_x\mathcal M$ on the manifold
$\mathcal M = \mathbb S^n$. For the representation the tangent space can be
given as $T_x\mathbb S^n = \bigl\{\xi \in \mathbb R^{n+1}
\big| \langle x,\xi\rangle = 0\bigr\}$, where $\langle\cdot,\cdot\rangle$
denotes the Euclidean inner product on $\mathbb R^{n+1}$.
"""
struct SnTVector{T <: AbstractFloat} <: TVector
  value::Vector{T}
  SnTVector{T}(value::Vector{T}) where {T <: AbstractFloat} = new(value)
end
SnTVector(value::Vector{T})  where {T <: AbstractFloat}  = SnTVector{T}(value)
getValue(ξ::SnTVector) = ξ.value;
# Traits
# ---
# (a) Sn is a MatrixManifold
@traitimpl IsMatrixM{Sphere}
@traitimpl IsMatrixP{SnPoint}
@traitimpl IsMatrixTV{SnTVector}
# (b) Sn is Embedded
@traitimpl IsEmbeddedM{Sphere}
@traitimpl IsEmbeddedP{SnPoint}
@traitimpl IsEmbeddedV{SnTVector}

# Functions
# ---
@doc doc"""
    distance(M,x,y)

Compute the Riemannian distance on $\mathcal M=\mathbb S^n$ embedded in
$\mathbb R^{n+1}$, which is given by

$ d_{\mathbb S^n}(x,y) = \operatorname{acos} \bigl(\langle x,y\rangle\bigr), $

where $\langle\cdot,\cdot\rangle$ denotes the Euclidean inner product
on $\mathbb R^{n+1}$.
"""
distance(M::Sphere,x::SnPoint,y::SnPoint) = acos(
	min( max(dot(getValue(x), getValue(y) ), -1.), 1. )
)

@doc doc"""
    dot(M,x,ξ,ν)
Compute the Riemannian inner product for two [`SnTVector`](@ref)s `ξ` and `ν`
from $T_x\mathcal M$ of the [`Sphere`](@ref) `M` given by
$\langle \xi, \nu \rangle_x = \langle \xi,\nu \rangle$, i.e. the inner product
in the embedded space $\mathbb R^{n+1}$.
"""
dot(M::Sphere, x::SnPoint, ξ::SnTVector, ν::SnTVector) = dot( getValue(ξ), getValue(ν) )

@doc doc"""
    exp(M,x,ξ[, t=1.0])
Compute the exponential map on the [`Sphere`](@ref) `M`$=\mathbb S^n$ with
respect to the [`SnPoint`](@ref) `x` and the [`SnTVector`](@ref) `ξ`, which can
be shortened with `t` to `tξ`. The formula reads

$\exp_x\xi = \cos(\lVert\xi\rVert_2)x + \sin(\lVert\xi\rVert_2)\frac{\xi}{\lVert\xi\rVert_2}.$
"""
function exp(M::Sphere,x::SnPoint,ξ::SnTVector,t::Float64=1.0)
  len = norm( getValue(ξ) )
  if len < eps(Float64)
    return x
  else
	xV = cos(t*len) * getValue(x)  +  (sin(t*len)/len) * getValue(ξ)
    return SnPoint( xV ./ norm(xV)  )
  end
end
@doc doc"""
    log(M,x,y)
Compute the logarithmic map on the [`Sphere`](@ref)
$\mathcal M=\mathbb S^n$, i.e. the [`SnTVector`](@ref) whose corresponding
[`geodesic`](@ref) starting from [`SnPoint`](@ref) `x` reaches the
[`SnPoint`](@ref)` y` after time 1 on the [`Sphere`](@ref) `M`.
The formula reads for $x\neq -y$

$\log_x y = d_{\mathbb S^n}(x,y)\frac{y-\langle x,y\rangle x}{\lVert y-\langle x,y\rangle x \rVert_2}.$
"""
function log(M::Sphere,x::SnPoint,y::SnPoint)
  scp = dot( getValue(x), getValue(y) )
  ξvalue = getValue(y) - scp*getValue(x)
  ξvnorm = norm(ξvalue)
  if (ξvnorm > eps(Float64))
    return SnTVector( ξvalue * acos(min(max(scp,-1.),1.) )/ξvnorm );
  else
    return zeroTVector(M,x)
  end
end
@doc doc"""
    manifoldDimension(x)
returns the dimension of the [`Sphere`](@ref) `M`$=\mathbb S^n$, the
[`SnPoint`](@ref) `x`, itself embedded in $\mathbb R^{n+1}$, belongs to.
"""
manifoldDimension(x::SnPoint)::Integer = length( getValue(x) )-1
"""
    manifoldDimension(M)
returns the dimension of the [`Sphere`](@ref) `M`.
"""
manifoldDimension(M::Sphere)::Integer = M.dimension
@doc doc"""
    norm(M,x,ξ)
Computes the norm of the [`SnTVector`](@ref) `ξ` in the tangent space
$T_x\mathcal M$ at [`SnPoint`](@ref) `x` of the [`Sphere`](@ref) `M`.
"""
norm(M::Sphere, x::SnPoint, ξ::SnTVector) = norm( getValue(ξ) )
@doc doc"""
    opposite(M,x)
returns the antipodal point of x, i.e. $ y = -x $.
"""
opposite(M::Sphere, x::SnPoint) = SnPoint( -getValue(x) )
@doc doc"""
    parallelTransport(M,x,y,ξ)
Compute the paralllel transport of the [`SnTVector`](@ref) `ξ` from
the tangent space $T_x\mathcal M$ at [`SnPoint`](@ref) `x` to
$T_y\mathcal M$ at [`SnPoint`](@ref)` y` on the [`Sphere`](@ref) `M` provided
that the corresponding [`geodesic`](@ref) $g(\cdot;x,y)$ is unique.
The formula reads

$P_{x\to y}(\xi) = \xi - \frac{\langle \log_xy,\xi\rangle_x}{d^2_{\mathbb S^n}(x,y)}
\bigl(\log_xy + \log_yx \bigr).$
"""
function parallelTransport(M::Sphere, x::SnPoint, y::SnPoint, ξ::SnTVector)
  ν = log(M,x,y)
  νL = norm(M,x,ν)
	return ξ - ( νL > 0 ? dot(M,x,ν,ξ)*(ν + log(M,y,x))/νL^2 : zeroTVector(M,x) )
end
@doc doc"""
    randomMPoint(M [,:Gaussian, σ=1.0])

return a random point on the Sphere by projecting a normal distirbuted vector
from within the embedding to the sphere.
"""
function randomMPoint(M::Sphere, ::Val{:Gaussian}, σ::Float64=1.0)::SnPoint
	v = σ * randn(manifoldDimension(M)+1);
	return SnPoint(v./norm(v))
end
@doc doc"""
    randomTVector(M,x [,:Gaussian,σ=1.0])

return a random tangent vector in the tangent space of the [`SnPoint`](@ref)
`x` on the [`Sphere`](@ref) `M`.
"""
function randomTVector(M::Sphere, x::SnPoint, ::Val{:Gaussian}, σ::Float64=1.0)
    n = σ * randn( size( getValue(x)) ) # Gaussian in embedding
	  nP = n - dot(n,getValue(x))*getValue(x) #project to TpM (keeps Gaussianness)
	  return SnTVector( nP )
end
tangentONB(M::Sphere, x::SnPoint, y::SnPoint) = tangentONB(M,x,log(M,x,y))
function tangentONB(M::Sphere,x::SnPoint,ξ::SnTVector)
    d = manifoldDimension(M)
    A = zeros(d+1,d+1)
    A[1,:] = transpose(getValue(x))
    A[2,:] = transpose(getValue(ξ))
    V = nullspace(A)
    κ = ones(d)
    if ξ != zeroTVector(M,x)
        # if we have a nonzero direction for the geodesic, add it and it gets curvature zero from the tensor
		ξ = ξ/norm(M,x,ξ)
		V = cat(getValue(ξ),V,dims=2)
        κ[1] = 0.0 # no curvature along the geodesic direction, if x!=y
    end
    Ξ = [ SnTVector(V[:,i]) for i in 1:d ]
    return Ξ,κ
end
typeofTVector(::Type{SnPoint{T}}) where T = SnTVector{T}
typeofMPoint(::Type{SnTVector{T}}) where T = SnPoint{T}
"""
    typicalDistance(M)

returns the typical distance on the [`Sphere`](@ref)` Sn`: π.
"""
typicalDistance(M::Sphere) = π

@doc doc"""
    validateMPoint(M,x)

validate, whether the [`SnPoint`](@ref) `x` is on the [`Sphere`](@ref) `M`$=\mathbb S^n$,
i.e. that the vector is of the correct dimension $n$ and its norm is $\lVert x \rVert = 1$.
"""
function validateMPoint(M::Sphere, x::SnPoint)
  if length(getValue(x)) ≠ M.dimension+1
    throw( ErrorException(
      "The Point $x is not on the $(M.name), since the vector dimension ($(length(getValue(x)))) is not $(M.dimension+1)."
    ))
  end
  if abs(norm(getValue(x)) - 1) >= 10^(-15)
    throw( ErrorException(
      "The Point $x is not on the $(M.name) since its norm $(norm(getValue(x))) ≠ 1"
    ))
  end
  return true
end

@doc doc"""
    validateTVector(M,x,ξ)

validate, whether the tangent vector [`SnTVector`](@ref) `ξ` is in the tangent
space of [`SnPoint`](@ref) `x` is on the [`Sphere`](@ref) `M`$=\mathbb S^n$,
i.e. that all three lengths are correct and $x^\mathrm{T}\xi = 0$.
"""
function validateTVector(M::Sphere,x::SnPoint,ξ::SnTVector)
  if (length(getValue(x)) ≠ length(getValue(ξ))) || (length(getValue(x)) ≠ M.dimension+1)
    throw( ErrorException(
      "The three dimensions of the $(M.name), the point x ($(length(getValue(x)))), and the tangent vector ($(length(getValue(ξ)))) don't match."
    ))
  end
  if dot(getValue(x),getValue(ξ)) >= 10^(-15)
    throw( ErrorException(
      "The tangent vector $ξ is not a tangent vector to $x on $(M.name) since the inner product is $( dot(getValue(x),getValue(ξ)) ) and not zero."
    ))
  end
  return true
end

@doc doc"""
    ξ = zeroTVector(M,x)

returns a zero vector in the tangent space $T_x\mathcal M$ of the
[`SnPoint`](@ref) $x\in\mathbb S^n$ on the [`Sphere`](@ref)` Sn`.
"""
zeroTVector(M::Sphere, x::SnPoint) = SnTVector(  zero( getValue(x) )  );
# Display
# ---
show(io::IO, M::Sphere) = print(io,"The $(M.name)")
show(io::IO, p::SnPoint) = print(io, "Sn($( getValue(p) ))")
show(io::IO, ξ::SnTVector) = print(io, "SnT($( getValue(ξ) ))")


# Car Escape
T = Float64;

# model
model = Dynamics.car
model_d = rk3(model)
n = model.n; m = model.m
x0 = [2.5;2.5;0.]
xf = [7.5;2.5;0.]

# cost
Q = (1e-3)*Diagonal(I,n)
R = (1e-2)*Diagonal(I,m)
Qf = 100.0*Diagonal(I,n)

# constraints
r = 0.5
s1 = 30; s2 = 50; s3 = 15

circles_escape = []

for i in range(0,stop=5,length=s1)
    push!(circles_escape,(0.,i,r))
end
for i in range(0,stop=5,length=s1)
    push!(circles_escape,(5.,i,r))
end
for i in range(0,stop=5,length=s1)
    push!(circles_escape,(10.,i,r))
end
for i in range(0,stop=10,length=s2)
    push!(circles_escape,(i,0.,r))
end
for i in range(0,stop=3,length=s3)
    push!(circles_escape,(i,5.,r))
end
for i in range(5,stop=8,length=s3)
    push!(circles_escape,(i,5.,r))
end

n_circles_escape = 3*s1 + s2 + 2*s3

function cI_escape(c,x,u)
    for i = 1:n_circles_escape
        c[i] = circle_constraint(x,circles_escape[i][1],circles_escape[i][2],circles_escape[i][3])
    end
end

trap = Constraint{Inequality}(cI_escape,n,m,n_circles_escape,:trap)
bnd = BoundConstraint(n,m,u_min=-5.,u_max=5.)
goal = goal_constraint(xf)

N = 101
tf = 3.0
U = [ones(m) for k = 1:N-1]
obj = LQRObjective(Q,R,Qf,xf,N)

constraints = Constraints(N)
constraints[1] += bnd
for k = 2:N-1
    constraints[k] += trap + bnd
end
constraints[N] += goal

car_escape = Problem(model_d,obj,constraints=constraints,N=N,tf=tf,x0=x0,xf=xf)
initial_controls!(car_escape, U);

X_guess = [2.5 2.5 0.;4. 5. .785;5. 6.25 0.;7.5 6.25 -.261;9 5. -1.57;7.5 2.5 0.]
X0_escape = interp_rows(N,tf,Array(X_guess'))

copyto!(car_escape.X,X0_escape)

# plot escape
function plot_escape(X,x0=x0,xf=xf,X0=X0_escape)
    X_array = to_array(X)
    plot(labels="")
    plot_obstacles(circles_escape,:grey)
    plot!(X0[1,:],X0[2,:],width=2,color=:purple,label="X_guess",axis=:off)
    plot!((x0[1],x0[2]),marker=:circle,color=:red,label="")
    plot!((xf[1],xf[2]),marker=:circle,color=:green,label="",legend=:left)
    plot!(X_array[1,:],X_array[2,:],color=:blue,width=2,label="Solution")
end


using LoopVectorization
using Plots
using UnPack

using StartUpDG
using StartUpDG.ExplicitTimestepUtils

N = 3
K1D = 16
T = .5
CFL = 1.

# init ref element and mesh
rd = init_reference_tri(N)
VX,VY,EToV = uniform_tri_mesh(K1D)
md = init_DG_mesh(VX,VY,EToV,rd)

# Make domain periodic
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
build_periodic_boundary_maps!(md,rd,LX,LY)

#####
##### Define initial coefficients and time-stepping
#####

# "Define the initial conditions by interpolation"
@unpack x,y = md
p = @. exp(-100*(x^2+y^2))
u = @. 0*x
v = @. 0*x

# "Define function to evaluate the RHS: Q = (p,u,v)"
function rhs(Q, rd::RefElemData, md::MeshData)

    @unpack Dr,Ds,LIFT,Vf = rd
    @unpack rxJ,sxJ,ryJ,syJ,J = md
    @unpack nxJ,nyJ,sJ = md
    @unpack mapP,mapB = md

    Qf = (x->Vf*x).(Q)
    QP = (x->x[mapP]).(Qf)
    dp,du,dv = QP.-Qf

    tau = .5
    dun   = @. (du*nxJ + dv*nyJ)/sJ
    pflux = @. .5*(du*nxJ + dv*nyJ) - tau*dp*sJ
    uflux = @. .5*dp*nxJ - tau*dun*nxJ
    vflux = @. .5*dp*nyJ - tau*dun*nyJ

    pr,ur,vr = (x->Dr*x).(Q)
    ps,us,vs = (x->Ds*x).(Q)
    dpdx = @. rxJ*pr + sxJ*ps
    dpdy = @. ryJ*pr + syJ*ps
    dudx = @. rxJ*ur + sxJ*us
    dvdy = @. ryJ*vr + syJ*vs
    rhsp = dudx + dvdy + LIFT*pflux
    rhsu = dpdx        + LIFT*uflux
    rhsv = dpdy        + LIFT*vflux

    return map(x->-x./J,(rhsp,rhsu,rhsv))
end

function rhs_avx(Q, rd::RefElemData, md::MeshData)

    @unpack Dr,Ds,LIFT,Vf = rd
    @unpack rxJ,sxJ,ryJ,syJ,J = md
    @unpack nxJ,nyJ,sJ = md
    @unpack mapP,mapB = md

    Qf = (x->Vf*x).(Q)
    QP = (x->x[mapP]).(Qf)
    dp,du,dv = QP.-Qf

    tau = .5
    dun   = @avx @. (du*nxJ + dv*nyJ)/sJ
    pflux = @avx @. .5*(du*nxJ + dv*nyJ) - tau*dp*sJ
    uflux = @avx @. .5*dp*nxJ - tau*dun*nxJ
    vflux = @avx @. .5*dp*nyJ - tau*dun*nyJ

    pr,ur,vr = (x->Dr*x).(Q)
    ps,us,vs = (x->Ds*x).(Q)
    dpdx = @avx @. rxJ*pr + sxJ*ps
    dpdy = @avx @. ryJ*pr + syJ*ps
    dudx = @avx @. rxJ*ur + sxJ*us
    dvdy = @avx @. ryJ*vr + syJ*vs
    rhsp = dudx + dvdy + LIFT*pflux
    rhsu = dpdx        + LIFT*uflux
    rhsv = dpdy        + LIFT*vflux

    return (x->@avx -x./J).((rhsp,rhsu,rhsv))
end

# "Time integration coefficients"
rk4a,rk4b,rk4c = ck45()
CN = (N+1)*(N+2)/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

# "Perform time-stepping"
Q = (p,u,v) # use mutable arrays for Q, resQ for scoping
resQ = zero.(Q)
for i = 1:Nsteps
    for INTRK = 1:5
        rhsQ = rhs_avx(Q,rd,md)
        bcopy!.(resQ, @. rk4a[INTRK]*resQ + dt*rhsQ)
        bcopy!.(Q,@. Q + rk4b[INTRK]*resQ)
    end

    if i%10==0 || i==Nsteps
        println("Time step $i out of $Nsteps")
    end
end

# # dopri initialization
# rka,rkE,rkc = dp56()
# PIparams = init_PI_controller(5)
# Qtmp = similar.(Q)
# rhsQrk = ntuple(x->zero.(Q),length(rkE))
# prevErrEst = nothing
# rhsQ = rhs(Q,rd,md)
# bcopy!.(rhsQrk[1],rhsQ) # initialize DOPRI rhs (FSAL property)
#
# t = 0.0
# i = 0
# dthist = [dt]
#
# while t < T
#     for INTRK = 2:7
#         k = zero.(Qtmp)
#         for s = 1:INTRK-1
#             bcopy!.(k, @. k + rka[INTRK,s]*rhsQrk[s])
#         end
#         bcopy!.(Qtmp, @. Q + dt*k)
#         rhsQ = rhs_avx(Qtmp,rd,md)
#         bcopy!.(rhsQrk[INTRK],rhsQ)
#     end
#
#     global t,dt,i,prevErrEst # modify compute_adaptive_dt!(...prevErrEst)
#     accept_step, dtnew, prevErrEst = compute_adaptive_dt(
#                                         Q,rhsQrk,dt,rkE,PIparams,prevErrEst)
#     if accept_step
#         t += dt
#         bcopy!.(Q, Qtmp)
#         bcopy!.(rhsQrk[1], rhsQrk[7]) # use FSAL property
#         push!(dthist,dt) # store dt history
#     end
#     dt = min(T-t,dtnew)
#     i = i + 1  # number of total steps attempted
#
#     if i%10==0
#         println("i = $i, t = $t, dt = $dtnew, errEst = $prevErrEst")
#     end
# end

# plotting nodes
gr(aspect_ratio=1,legend=false)
using NodesAndModes.Tri
rp,sp = equi_nodes_2D(10)
Vp = vandermonde_2D(N,rp,sp)/rd.VDM
scatter((x->Vp*x).((x,y,Q[1])),zcolor=Vp*Q[1], msw=0,ms=2,cam=(0,90))


# f(x)=tanh(x)

sf=1.0

Ftanh(x::Array{Float64,2})=(sf*tanh(x),[]);
Ftanh_inplace(value,aux,x::Array{Float64,2})=copy!(value,sf*tanh(x))

Dtanh(derivativeIDX,f_c,faux_c,grad_c,grad_n,x)=axpy!(sf,grad_c.*(1.-(f_c/sf).^2),grad_n)

if PROC=="GPU" ## TODO
    function Ftanh(x::CudaArray)
        tmp=copy(x)
        tanh!(sf,x,tmp)
        return (tmp,[]) # problem here is that I don't know how to free tmp -- this will cause memory leak I think.
    end
    Ftanh_inplace(value::CudaArray,aux,x::CudaArray)=tanh!(sf,x,value)

    Dtanh(derivativeIDX,f_c,faux_c,grad_c::CudaArray,grad_n::CudaArray,x::CudaArray)=Dtanh!(sf,grad_c,f_c,grad_n)
end

Derivative[Ftanh]=Dtanh
Inplace[Ftanh]=Ftanh_inplace

import Base.tanh
tanh(A::ADnode)=ADnode(Ftanh,A)
export tanh


# Exports
export access_geometry, visible
export spacecraft_compute_access
export compute_access
export parallel_compute_access
export create_lookup_location_opportunity

###################
# Access Geometry #
###################

"""
Compute the view geometry from an observer to a specific request.

Arguments:
- `sat_ecef:Array{<:Real, 1}` Earth-fixed satellite position
- `request::Request` Request object

Returns:
- `look_angle::Float64` Look angle from the satellite to the request center. Equivalent to off-nadiar angle. [deg]
- `range::Float64` Range from the observing satellite to the location. [m]
"""
function access_geometry(sat_ecef::Array{<:Real, 1}, request::Request)
    # Satellite state
    r_sat = sat_ecef[1:3]

    # Geodetic sub-satellte point
    sat_geod 	 = sECEFtoGEOD(r_sat)
    sub_sat_geod = [sat_geod[1], sat_geod[2], 0.0]
    sub_sat_ecef = sGEODtoECEF(sub_sat_geod)

    # Look angle
    nadir_dir  = (sub_sat_ecef - r_sat)/norm(sub_sat_ecef - r_sat)
    location_dir = (request.ecef - r_sat)/norm(request.ecef - r_sat)
    look_angle = acos(dot(nadir_dir, location_dir))*180.0/pi

    # Distance from sub-satellite point to location along direct line of sight
    range = norm(r_sat - request.ecef) # range to location

    return look_angle, range
end

"""
Compute the view geometry from an observer to a specific station.

Arguments:
- `sat_ecef:Array{<:Real, 1}` Earth-fixed satellite position
- `station::GroundStation` Station object

Returns:
- `elevation::Float64` Elevation of satellite with respect to station [deg]
"""
function access_geometry(sat_ecef::Array{<:Real, 1}, station::GroundStation)
    # Satellite state
    r_sat     = sat_ecef[1:3]
    r_station = station.ecef

    azimuth, elevation, range = sSEZtoAZEL(sECEFtoSEZ(r_station, r_sat), use_degrees=true)

    return elevation, range
end

"""
Computes whether an request is visible from a given state and return boolean true/false
result.

Arguments:
- `sat_ecef::Array{<:Real, 1}` Satellite position in Earth fixed frame
- `request::Request` Request being viewed

Returns:
- `visible::Bool` Indication of whether request is visible or not
"""
function visible(sat_ecef::Array{<:Real, 1}, request::Request)
    # Compute access geometry
    look_angle, range = access_geometry(sat_ecef, request)

    r = norm(sat_ecef[1:3])
    # theta_max = request.look_angle_max*pi/180.0
    # max_range = -r*cos(theta_max) - sqrt(r^2*cos(theta_max)^2 - (r^2-R_EARTH^2)^2)

    # To be fixed, no idea why this is complex:
    max_range = R_EARTH
    # max_range = 1.5*sqrt(r^2 - R_EARTH^2)

    # @debug "max range: $max_range"
    if request.look_angle_min <= look_angle &&
       look_angle <= request.look_angle_max &&
       range <= max_range

        return true
    else
        return false
    end
end

"""
Computes whether an station is visible from a given state and return boolean true/false
result.

Arguments:
- `sat_ecef::Array{<:Real, 1}` Satellite position in Earth fixed frame
- `station::GroundStation` Station being communicatted with

Returns:
- `visible::Bool` Indication of whether station is visible or not
"""
function visible(sat_ecef::Array{<:Real, 1}, station::GroundStation)
    elevation, range = access_geometry(sat_ecef, station)

    r = norm(sat_ecef[1:3])

    if station.elevation_min <= elevation
        # dist = norm(sat_ecef[1:3] - station.ecef[1:3])
        return true
    else
        return false
    end
end

###################
# Location Access #
###################

"""
Find the time when a location transitions between being accessible/not-accessible.
"""
function find_access_boundary(tle::TLE, location::Location, epc0::Epoch, step::Real; tol::Real=1.0e-1)

    # println("Step size: $step")

    epc = deepcopy(epc0)

    if abs(step) < tol
        # If time step is below desired tolerance do nothing
        return epc0
    else
        x_ecef   = sECItoECEF(epc, state(tle, epc))
        initial_visibility = visible(x_ecef, location)

        # While visibility is the same as it was initially step in time and
        # recompute visibility
        while visible(x_ecef, location) == initial_visibility
            epc    = epc + step
            x_ecef = sECItoECEF(epc, state(tle, epc))
        end

        # Reverse search direction and half step size
        next_step = -sign(step)*max(abs(step)/2, tol/2)

        return find_access_boundary(tle, location, epc, next_step, tol=tol)
    end
end

function spacecraft_compute_access(problem::SatPlanningProblem, spacecraft::Spacecraft,
            locations::Array{<:Location, 1};
            macro_step::Real=60.0,
            tol::Real=0.01,
            orbit_fraction::Real=0.5,
            id_offset::Integer=0)

    # Access is computed over planning horizon of problem (t_start -> t_end)

    # Compute orbital period
    T = orbit_period(sCARTtoOSC(state(spacecraft.tle, problem.t_start), use_degrees=true)[1])

    opportunities = Array{Opportunity, 1}(undef, 0)

    next_step = macro_step

    for loc in locations
        # Perform macro adaptive stepsize search
        epc = problem.t_start
        while epc < problem.t_end

            # println("Current step: $epc")

            x_ecef = sECItoECEF(epc, state(spacecraft.tle, epc))

            if visible(x_ecef, loc)
                # println("Found instant of visibility - $(visible(x_ecef, loc)) - $epc")

                # Search for AOS (before initial guess epoch)
                window_open = find_access_boundary(spacecraft.tle, loc, epc, -macro_step, tol=tol)
                # println("Found window_open boundary: $window_open")

                # Search for LOS (after initial guess epoch)
                window_close = find_access_boundary(spacecraft.tle, loc, epc, macro_step, tol=tol)
                # println("Found window_close boundary: $window_close")

                # Convert window_open/window_close to elapsed time values
                # window_open = get_rel_time(problem, window_open)
                # window_close = get_rel_time(problem, window_close)

                # If zero-doppler collection is required updated pass-times and geometry profile
                collect_duration = 0
                if typeof(loc) == Request && loc.require_zero_doppler == true
                    epc_mid = window_open + (window_close - window_open)/2.0
                    window_open  = epc_mid - loc.collect_duration/2.0
                    window_close = epc_mid + loc.collect_duration/2.0
                end

                # Create Oppourtunity
                if (window_close - window_open) > 0.0
                    if typeof(loc) == Request
                        collect = Collect(window_open,
                            window_close,
                            spacecraft=spacecraft,
                            location=loc)

                        push!(opportunities, collect)
                    elseif typeof(loc) == GroundStation
                        contact = Contact(window_open,
                            window_close,
                            spacecraft=spacecraft,
                            location=loc)

                        push!(opportunities, contact)
                    else
                        throw(ErrorException("Unknown type of location: $loc"))
                    end
                end

                # Step a most of an orbit because there (likely) won't be another
                # Acces until at least 1 orbit later
                epc += T*orbit_fraction

            else
                # Take a macro step because there's a good chance there won't be an
                # opportunity until the next one
                epc += macro_step
            end
        end
    end

    # Sort opportunities in ascending order
    sort!(opportunities, by = x -> x.t_start)

    # Apply IDs for collects in ascending time order
    for (idx, opp) in enumerate(opportunities)
        # Override computed opportunity id
        opp.id = idx + id_offset
    end

    return opportunities
end

function compute_access(problem::SatPlanningProblem;
            macro_step::Real=60.0, tol::Real=0.01,
            orbit_fraction::Real=0.5)

    # All Opportunities
    all_opportunities = Array{Opportunity, 1}(undef, 0)

    # Compute all opportunities for each spacecraft
    for spacecraft in problem.spacecraft
        @debug "Computing access for spacecraft: $(spacecraft.id)"
        opportunities = spacecraft_compute_access(problem,
                            spacecraft,
                            problem.locations,
                            macro_step=macro_step, tol=tol,
                            orbit_fraction=orbit_fraction)

        # Add opportunities to list of all opportunities
        push!(all_opportunities, opportunities...)
    end

    # Sort opportunities in ascending order
    sort!(all_opportunities, by = x -> x.t_start)

    # Apply IDs for collects in ascending time order
    id_offset = length(problem.opportunities)
    for (idx, opp) in enumerate(all_opportunities)
        # Override computed opportunity id
        opp.id = idx
    end

    # Add opportunities to problem
    problem.opportunities = all_opportunities

    # Create separate arrays of contacts and requests
    problem.contacts = filter(x -> typeof(x) == Contact, problem.opportunities)
    problem.collects = filter(x -> typeof(x) == Collect, problem.opportunities)

    # Zero Location Opportunity Counts
    for loc in problem.locations
        problem.lt_loc_opps[loc.id] = Integer[]
    end

    # Update opportunity lookup table
    for opp in problem.opportunities
        problem.lt_opportunities[opp.id] = opp

        if typeof(opp) == Contact
            problem.lt_contacts[opp.id] = opp
        elseif typeof(opp) == Collect
            problem.lt_collects[opp.id] = opp
        end

        # Location -> Opportunity Lookup
        push!(problem.lt_loc_opps[opp.location.id], opp.id)
    end

    # Update opportunities with number of collects remaining
    for (img_id, opp_ids) in problem.lt_loc_opps
        # Update opportunities with number of remaining collect opps
        for (i, opp_id) in enumerate(reverse(opp_ids))
            problem.lt_opportunities[opp_id].nr = i-1
        end
    end

    return
end

function parallel_compute_access(problem::SatPlanningProblem;
            macro_step::Real=60.0, tol::Real=0.01,
            orbit_fraction::Real=0.5)

    # Create anonymous function to map array inputs to
    # Must be declared at start of function for julia reasons...
    # x is a tuple (spacecraft, locations)
    fn(p, a, b, c) = x -> spacecraft_compute_access(p, x[1], x[2], macro_step=a, tol=b, orbit_fraction=c)

    # Create work assignments
    nw = nworkers()        # Number of workers
    lw = length(problem.locations) # Length of work

    # @debug "Have $lw items of work and $nw workers")

    wa = floor(Int, lw/nw) # Average work per worker
    wr = lw - nw*wa        # Remaining worker

    # Construct assignments (function inputs) for Workers
    assignments = Tuple{Spacecraft, Array{<:Location, 1}}[]

    for sc in problem.spacecraft
         # @debug "Worker 1 convering $(1):$(wa+wr)")
        push!(assignments, (sc, problem.locations[1:(wa+wr)]))
        for i in 2:nw
            # @debug "Worker $i convering $(1+wr+(i-1)*wa):$(i*wa+wr)")
            push!(assignments, (sc, problem.locations[(1+wr+(i-1)*wa):(i*wa+wr)]))
        end
    end

    # Execute Work in paralle
    results = pmap(fn(problem, macro_step, tol, orbit_fraction), assignments)

    # Aggregate and process results

    # Add opportunities to problem
    problem.opportunities = vcat(results...)

    # Sort opportunities in ascending order
    sort!(problem.opportunities, by = x -> x.t_start)

    # Create separate arrays of contacts and requests
    problem.contacts = filter(x -> typeof(x) == Contact, problem.opportunities)
    problem.collects = filter(x -> typeof(x) == Collect, problem.opportunities)

    # Zero Location Opportunity Counts
    for loc in problem.locations
        problem.lt_loc_opps[loc.id] = Integer[]
    end

    # Update opportunity lookup table
    for (id, opp) in enumerate(problem.opportunities)
        # Assign global opportunity ID
        opp.id = id

        # Create lookup entry
        problem.lt_opportunities[opp.id] = opp

        if typeof(opp) == Contact
            problem.lt_contacts[opp.id] = opp
        elseif typeof(opp) == Collect
            problem.lt_collects[opp.id] = opp
        end

        # Location -> Opportunity Lookup
        push!(problem.lt_loc_opps[opp.location.id], opp.id)
    end

    return
end

###############################
# Action Space Precomputation #
###############################

export precompute_action_space
function precompute_action_space(problem::SatPlanningProblem; enable_resources::Bool=false)

    # First compute feasible transition lookup
    lt_feasible_transitions = Dict{Tuple{Integer, Integer}, Bool}()

    println("Computing Feasibility Lookup Table:")
    for start_opp_idx in 1:length(problem.opportunities)
        start_opp_opp = problem.lt_opportunities[start_opp_idx]

        if start_opp_idx % 100 == 0
            println("Current Index: $start_opp_idx")
        end

        for end_opp_idx in (start_opp_idx+1):length(problem.opportunities)
            end_opp_opp = problem.lt_opportunities[end_opp_idx]

            # Check feasibility constraints
            valid = true

            for constraint in problem.constraints
                if valid == false
                    break
                end

                valid = valid && constraint(problem.lt_opportunities[start_opp_idx], problem.lt_opportunities[end_opp_idx])
            end

            # Push
            lt_feasible_transitions[(start_opp_idx, end_opp_idx)] = valid
        end
    end

    println("Current Index: $(length(problem.opportunities))")

    # Dict mapping (last collect/downlink action) -> {all possible actions}
    sunpoint_actions = Tuple{Opportunity, Integer}[]
    lt_feasible_actions = Dict{Tuple{Integer, Integer}, AbstractVector{Opportunity}}()

    # Below might work but it doesn't solve the problem of having sunpointed actions
    # be part of the
    println("Computing Action Space Lookup:")
    for ca_idx in 1:length(problem.opportunities)
        ca_opp = problem.lt_opportunities[ca_idx]

        if ca_idx % 100 == 0
            println("Current Index: $ca_idx")
        end

        # Create sunpoint action once for time
        if enable_resources == true && ca_idx < length(problem.opportunities)
            sp_action = Sunpoint(id=length(problem.opportunities)+length(sunpoint_actions)+1, t_start=problem.opportunities[ca_idx+1].t_start)
            push!(sunpoint_actions, (sp_action, ca_idx+1)) # Record equivalent index of sunpoint action
        end

        for lcdo_idx in 1:ca_idx
            lcdo_opp = problem.lt_opportunities[lcdo_idx]

            # Create Array for (current action, last action)
            lt_feasible_actions[(lcdo_opp.id, ca_opp.id)] = Opportunity[]

            for fa_idx in (ca_idx+1):length(problem.opportunities)
                if lt_feasible_transitions[(lcdo_idx, fa_idx)] == true
                    push!(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)], problem.opportunities[fa_idx])
                end

                if problem.solve_breadth > 0 && length(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)]) > problem.solve_breadth
                    break
                end
            end

            # Add sunpoint action as feasible action
            if enable_resources == true && ca_idx < length(problem.opportunities)
                push!(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)], sp_action)
            end
        end
    end

    println("Current Index: $(length(problem.opportunities))")

    if enable_resources == true
        println("Computing Sunpoint Indices:")
        for (sp_action, ca_idx) in sunpoint_actions
            ca_opp = problem.lt_opportunities[ca_idx]

            if ca_idx % 100 == 0
                println("Current Sunpoint Index: $ca_idx")
            end

            for lcdo_idx in 1:ca_idx
                lcdo_opp = problem.lt_opportunities[lcdo_idx]

                # Copy Exisiting feasible actions
                lt_feasible_actions[(lcdo_opp.id, sp_action.id)] = copy(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)])
            end
        end

        println("Current Sunpoint Index: $(length(problem.opportunities))")
    end

    # Update Problem Values
    problem.lt_feasible_actions = lt_feasible_actions
    if enable_resources == true
        problem.actions = vcat(problem.opportunities, [x[1] for x in sunpoint_actions])
    else
        problem.actions = problem.opportunities
    end
    sort!(problem.actions, by = x -> x.t_start)

    # Create Action Lookup Table
    lt_actions = Dict{Integer, Opportunity}()
    for action in problem.actions
        lt_actions[action.id] = action
    end
    problem.lt_actions = lt_actions
end

export pas_fa_work
function pas_fa_work(problem::SatPlanningProblem, start_indices::AbstractVector{Integer})
    # Array to store feasible transitoins
    feasible_transitions = Tuple{Integer, Integer, Bool}[]

    # println("Checking in! - $(length(problem.opportunities)) - Work: $start_indices")

    # For Opportunity compute all feasible transitions
    for start_opp_idx in start_indices
        start_opp_opp = problem.lt_opportunities[start_opp_idx]

        for end_opp_idx in (start_opp_idx+1):length(problem.opportunities)
            end_opp_opp = problem.lt_opportunities[end_opp_idx]

            # Check feasibility constraints
            valid = true

            for constraint in problem.constraints
                if valid == false
                    break
                end

                valid = valid && constraint(problem.lt_opportunities[start_opp_idx], problem.lt_opportunities[end_opp_idx])
            end

            # Push
            push!(feasible_transitions, (start_opp_idx, end_opp_idx, valid))
        end
    end

    return feasible_transitions
end

export pas_cas_work
function pas_cas_work(problem::SatPlanningProblem, lt_transitions::Dict{Tuple{Integer, Integer}, Bool}, current_index::AbstractVector{Integer}; enable_resources::Bool=false)
    println("Checking in!")
    flush(stdout)
    sunpoint_actions = Tuple{Opportunity, Integer}[]
    lt_feasible_actions = Dict{Tuple{Integer, Integer}, AbstractVector{Opportunity}}()

    for ca_idx in current_index
        ca_opp = problem.lt_opportunities[ca_idx]

        # Create sunpoint action once for time
        if enable_resources == true && ca_idx < length(problem.opportunities)
            sp_action = Sunpoint(id=length(problem.opportunities)+length(sunpoint_actions)+1, t_start=problem.opportunities[ca_idx+1].t_start)
            push!(sunpoint_actions, (sp_action, ca_idx+1)) # Record equivalent index of sunpoint action
        end

        for lcdo_idx in 1:ca_idx
            lcdo_opp = problem.lt_opportunities[lcdo_idx]

            # Create Array for (current action, last action)
            lt_feasible_actions[(lcdo_opp.id, ca_opp.id)] = Opportunity[]

            for fa_idx in (ca_idx+1):length(problem.opportunities)
                if lt_transitions[(lcdo_idx, fa_idx)] == true
                    push!(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)], problem.opportunities[fa_idx])
                end

                if problem.solve_breadth > 0 && length(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)]) > problem.solve_breadth
                    break
                end
            end

            # Add sunpoint action as feasible action
            if enable_resources == true && ca_idx < length(problem.opportunities)
                push!(lt_feasible_actions[(lcdo_opp.id, ca_opp.id)], sp_action)
            end
        end
    end

    return lt_feasible_actions, sunpoint_actions
end

export parallel_precompute_action_space
function parallel_precompute_action_space(problem::SatPlanningProblem; enable_resources::Bool=false)

    ## Compute Feasible Transitions
    fn(p) = w -> pas_fa_work(p, w)

    # Create work assignments to minimize message passing
    lw = length(problem.opportunities)  # Length of work
    nw = nworkers()                     # Number of workers
    wa = floor(Int, lw/nw)              # Average work per worker
    wr = lw - nw*wa                     # Remaining worker

    # Construct Work Assignment
    work = AbstractVector{Integer}[]

    push!(work, collect(1:(wa+wr)))
    for i in 2:nw
        push!(work, collect((1+wr+(i-1)*wa):(i*wa+wr)))
    end

    # Execute work in parallel, aggregate
    feasibility_results = vcat(pmap(fn(problem), work)...)

    # Aggrecate results
    lt_feasible_transitions = Dict{Tuple{Integer, Integer}, Bool}()

    for item in feasibility_results
        lt_feasible_transitions[(item[1], item[2])] = item[3]
    end

    ## Compute Action Space

    fn(p, lt_f, er) = w -> pas_cas_work(p, lt_f, w, enable_resources=er)

    # Create work assignments to minimize message passing
    lw = length(problem.opportunities)  # Length of work
    nw = nworkers()                     # Number of workers
    wa = floor(Int, lw/nw)              # Average work per worker
    wr = lw - nw*wa                     # Remaining worker

    # Construct Work Assignment
    work = AbstractVector{Integer}[]

    cids = collect(shuffle(1:length(problem.opportunities)))
    push!(work, collect(cids[1:(wa+wr)]))
    for i in 2:nw
        push!(work, collect(cids[(1+wr+(i-1)*wa):(i*wa+wr)]))
    end

    # Execute work in parallel, aggregate
    println("Starting Action Space Compute")
    action_space_results = pmap(fn(problem, lt_feasible_transitions, enable_resources), work)

    # Aggrecate results
    sunpoint_actions = Tuple{Opportunity, Integer}[]
    lt_feasible_actions = Dict{Tuple{Integer, Integer}, AbstractVector{Opportunity}}()

    for (lt_feas, sp_actions) in action_space_results

        # Copy Feasibility Lookup
        for (k, v) in lt_feas
            lt_feasible_actions[k] = v
        end

        # Copy Sunpoint Actions and update ID
        for spa in sp_actions
            # Update ID
            spa[1].id = length(problem.opportunities)+length(sunpoint_actions)+1
            push!(sunpoint_actions, spa)
        end

    end

    return lt_feasible_actions

    ## Compute Sunpoint Actions

    # Update Problem Values
    # problem.lt_feasible_actions = lt_feasible_actions
    # if enable_resources == true
    #     problem.actions = vcat(problem.opportunities, [x[1] for x in sunpoint_actions])
    # else
    #     problem.actions = problem.opportunities
    # end
    # sort!(problem.actions, by = x -> x.t_start)

    # # Create Action Lookup Table
    # lt_actions = Dict{Integer, Opportunity}()
    # for action in problem.actions
    #     lt_actions[action.id] = action
    # end
    # problem.lt_actions = lt_actions
end