using ValueOrientedRiskManagementInsurance
using DataFrames
using Base.Test

println("start ECModel2.jl...")

lob_strings = ["Fire", "Liability", "Theft"]
df_claims = Vector(length(lob_strings))
claims = Vector(length(lob_strings))
mack = Vector(length(lob_strings))
β = Vector(length(lob_strings))
for i = 1:length(df_claims)
  df_claims[i] =
    readtable("test/ECModel_Input_" * lob_strings[i] * ".csv",
              header = false)
  mack[i] = Mack(df_claims[i])
  β[i] = mack[i].futureclaims / sum(mack[i].futureclaims)
end

function logpar2statpar(m_s::Vector)
  Real[exp(m_s[1] + m_s[2]^2 / 2),
       √(exp(m_s[2]^2-1)) * exp(m_s[1] + m_s[2]^2 / 2)]
end

function statpar2logpar(μ_σ::Vector)
  vc = μ_σ[2]/μ_σ[1]
  Real[log(μ_σ[1]) - 0.5 * (1+log(vc^2)), √(1+log(vc^2))]
end

logpar2statpar([0,1])
x = logpar2statpar([0,1])

statpar2logpar(x)

statpar2logpar(x[1],x[2])
#
# mean, sde -> logmean, logsigma
# R[i] = LogNormal(log_mean[i], log_sigma[i])
# v⋅mack[i].β * R[i]
#
# v: stochastic discont
# β
# R[i] = stochastic total undiscounted future paiments,
#

println("...end ECModel2.jl")


# ---
# title: 1356. Sort Integers by The Number of 1 Bits
# id: problem1356
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Sort, Bit Manipulation
# link: <https://leetcode.com/problems/sort-integers-by-the-number-of-1-bits/description/>
# hidden: true
# ---
#
# Given an integer array `arr`. You have to sort the integers in the array in
# ascending order by the number of **1 's** in their binary representation and
# in case of two or more integers have the same number of **1 's** you have to
# sort them in ascending order.
#
# Return _the sorted array_.
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [0,1,2,3,4,5,6,7,8]
#     Output: [0,1,2,4,8,3,5,6,7]
#     **Explantion:** [0] is the only integer with 0 bits.
#     [1,2,4,8] all have 1 bit.
#     [3,5,6] have 2 bits.
#     [7] has 3 bits.
#     The sorted array by bits is [0,1,2,4,8,3,5,6,7]
#
#
# **Example 2:**
#
#
#
#     Input: arr = [1024,512,256,128,64,32,16,8,4,2,1]
#     Output: [1,2,4,8,16,32,64,128,256,512,1024]
#     **Explantion:** All integers have 1 bit in the binary representation, you should just sort them in ascending order.
#
#
# **Example 3:**
#
#
#
#     Input: arr = [10000,10000]
#     Output: [10000,10000]
#
#
# **Example 4:**
#
#
#
#     Input: arr = [2,3,5,7,11,13,17,19]
#     Output: [2,3,5,17,7,11,13,19]
#
#
# **Example 5:**
#
#
#
#     Input: arr = [10,100,1000,10000]
#     Output: [10,100,10000,1000]
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 500`
#   * `0 <= arr[i] <= 10^4`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Distributions

# In the literature Cauchy distributions have been used for sampling the
# f and cr constants used in DE.
function sample_bimodal_cauchy_once(cauchyDistrs, cutoffProb = 0.5)
  index = (rand() < cutoffProb) ? 1 : 2
  rand(cauchyDistrs[index])
end

# When sampling it is common to truncate in either or both ends.
function sample_bimodal_cauchy(cauchyDistrs; cutoffProb = 0.5,
  truncateAbove1 = true, truncateBelow0 = true)
  value = sample_bimodal_cauchy_once(cauchyDistrs, cutoffProb)
  if value > 1.0
    if truncateAbove1
      return 1.0
    else
      return sample_bimodal_cauchy(cauchyDistrs; cutoffProb = cutoffProb, truncateAbove1 = truncateAbove1, truncateBelow0 = truncateBelow0)
    end
  elseif value < 0.0
    if truncateBelow0
      return 0.0
    else
      return sample_bimodal_cauchy(cauchyDistrs; cutoffProb = cutoffProb, truncateAbove1 = truncateAbove1, truncateBelow0 = truncateBelow0)
    end
  else
    return value
  end
end

# For sampling f in DE, bimodal_cauchy(0.65, 0.1, 1.0, 0.1) has been proposed.
# For sampling cr in DE, bimodal_cauchy(0.1, 0.1, 0.95, 0.1) has been proposed.
function bimodal_cauchy(location1, scale1, location2, scale2)
  (Cauchy(location1, scale1), Cauchy(location2, scale2))
end


### ============== ### ============== ###
##    3D Simple Vicksek Model          ##
##    Grid version (particle search)   ##
##    Martin Zumaya Hernandez          ##
##    30 / 06 / 2017                   ##
### ============== ### ============== ###

### ============== ### ============== ###

using CollectiveDynamics.SVM3D

### =============== ### =============== ###
###   DEFINITION OF INITIAL PARAMETERS  ###
### =============== ### =============== ###

N   = parse(Int, ARGS[1]) # number of particles
ρ   = parse(Float64, ARGS[2]) # density
η   = parse(Float64, ARGS[3]) # noise intensity
T   = parse(Int, ARGS[4]) # integration time steps
rep = parse(Int, ARGS[5]) # ensemble index

### =============== ### =============== ###

v0 = 1.0 # particle's speed
dt = 1.0 # integration time step

### =============== ### =============== ###

M = 10 # number of boxes per dimension

l = 0.5 # interaction range is double the distance each particle moves in one time step
# l = 0.1 # interaction range is ten times the distance each particle moves in one time step

r0 = (v0 * dt) / l # local interaction range

L  = cbrt(N / ρ) # size of box

### =============== ### =============== ###

cell_size = step(linspace(0., L , M))

### ============== ### ============== ###

box = Box(L, M)
flock = Flock(N, L, dt, v0, r0, η)

### ============== ### ============== ###

output_path = set_output_data_structure_vsk("SVM_3D", N, ρ)

pos_file = open(output_path * "/pos_$(rep).dat", "w+")
vel_file = open(output_path * "/vel_$(rep).dat", "w+")

times = [convert(Int, exp10(i)) for i in 0:T]

for i in 1:(length(times) - 1)

    if i > 1

        for t in (times[i]+1):times[i+1]

            evolve_system(flock, box, cell_size)

            if t % times[i] == 0 || t % times[i-1] == 0
                println("//////// ", t)
                write(pos_file, vcat(flock.pos...))
                write(vel_file, vcat(flock.vel...))
            end
        end

    else

        for t in (times[i]+1):times[i+1]

            evolve_system(flock, box, cell_size)

            if t % times[i] == 0
                println("//////// ", t)
                write(pos_file, vcat(flock.pos...))
                write(vel_file, vcat(flock.vel...))
            end
        end

    end

end

close(pos_file)
close(vel_file)

println("Done all")


# A port of IDONE to Julia
# Mostly to get an idea of the performance overhead
# of using Python instead.

using LinearAlgebra
using Distributions

struct IDONEModel
    # Dimensionality
    d :: Int64
    # Bounds
    lb :: Vector{Float64}
    ub :: Vector{Float64}
    # Regularization parameter
    reg :: Float64
    # Weights of ReLU basis functions
    # As IDONE is for integers, these weights as well as biases
    # are integers as well. Store as float anyways.
    W :: Array{Float64, 2}
    # Biases of ReLU basis functions
    b :: Array{Float64, 1}
    # Basis function weights.
    # Updated using Recursive Least Squares
    c :: Vector{Float64}
    Δc :: Vector{Float64}
    # Keep track of the last x update.
    last_x_upd :: Vector{Float64}
    # Recursive Least Squares covariance matrix
    P :: Array{Float64, 2}
end

function generateWb(d :: Int64, lb :: Vector{Float64}, ub :: Vector{Float64})
    @assert d == length(lb)
    @assert d == length(ub)

    # Model offset
    n_basis = 1
    # One variable dependent (incl correction for d = 1)
    n_basis += (2 * round(Int64, sum(ub .- lb)))
    # Subsequent variable dependent
    n_basis += (2 * round(Int64, sum(1 .+ (ub[1:end-1] .- lb[2:end]) .- (lb[2:end] .- ub[1:end-1]))))

    # Actually construct the matrix
    W = zeros(Float64, (n_basis, d))
    b = zeros(Float64, n_basis)

    # Model offset
    fni = 1
    b[fni] = 1
    fni += 1

    # One variable dependent
    for k in 1:d
        for i in lb[k]:ub[k]
            if i == lb[k]
                W[fni, k] = 1
                b[fni] = -i
                fni += 1
            elseif i == ub[k]
                W[fni, k] = -1
                b[fni] = i
                fni += 1
            else
                W[fni, k] = -1
                b[fni] = i
                W[fni+1, k] = 1
                b[fni+1] = -i
                fni += 2
            end
        end
    end

    # Two subsequent variables
    for k in 2:d
        for i in (lb[k] - ub[k - 1]):(ub[k] - lb[k - 1] + 1)
            if i == lb[k] - ub[k-1]
                W[fni, k] = 1
                W[fni, k - 1] = -1
                b[fni] = -i
                fni += 1
            elseif i == ub[k] - lb[k - 1] + 1
                W[fni, k] = -1
                W[fni, k - 1] = 1
                b[fni] = i
                fni += 1
            else
                W[fni, k] = -1
                W[fni, k - 1] = 1
                b[fni] = i
                W[fni+1, k] = 1
                W[fni+1, k - 1] = -1
                b[fni+1] = -i
                fni += 2
            end
        end
    end

    return W, b
end

function generate_model(d :: Int64, lb :: Vector{Float64}, ub :: Vector{Float64})
    W, b = generateWb(d, lb, ub)
    D = length(b)
    c = ones(D)
    Δc = zeros(D)
    # No initial model offset
    c[1] = 0
    # reg = 1e-3
    reg = 1e-8
    P = Matrix(I / reg, D, D)
    return IDONEModel(d, lb, ub, reg, W, b, c, Δc, zeros(d), P)
end

function update_model!(model :: IDONEModel, x :: Vector{Float64}, y :: Float64)
    Zx = Z(model, x)
    model.last_x_upd .= x
    PZx = model.P * Zx
    g = PZx ./ (1 + Zx' * PZx)
    model.P .= model.P .- g * PZx'
    model.Δc .= (y - Zx' * model.c) .* g
    # model.c .= model.c .+ (y .- Zx' * model.c) .* g
    model.c .+= model.Δc
end

# Model evaluation
function id_relu(x)
    return max(0, x)
end

function ∇id_relu(x)
    return (x > 0) + 0.5 * (x == 0)
end

function Z(model :: IDONEModel, x :: Vector{Float64})
    return id_relu.(model.W * x .+ model.b)
end

function ∇Z(model :: IDONEModel, x :: Vector{Float64})
    return ∇id_relu.(model.W * x .+ model.b)
end

function ∇Z!(g :: Vector{Float64}, model :: IDONEModel, x :: Vector{Int64})
    resize!(g, length(model.c))
    mul!(g, model.W, x)
    g .+= model.b
    g .= ∇id_relu.(g)
end

function predict(model :: IDONEModel, x :: Vector{Float64})
    y_pred = dot(model.c, Z(model, x))
    return y_pred
end

function predict_v_and_diff(model :: IDONEModel, x :: Vector{Float64})
    Zx = Z(model, x)
    y_pred = dot(model.c, Zx)
    y_pred_diff = dot(model.Δc, Zx)
    return y_pred, y_pred_diff
end

function ∇predict(model :: IDONEModel, x :: Vector{Float64})
    return model.c' * Diagonal(∇Z(model, x)) * model.W
end

function ∇predict!(g :: Vector{Float64}, model :: IDONEModel, x :: Vector{Float64})
    gz = ∇Z(model, x)
    gz .*= model.c
    # gz' * W = W' * gz
    # Because apparently the other way around is not recognized.
    # Maybe a bug in Julia?
    mul!(g, model.W', gz)
end

function predict_dist_m2(model :: IDONEModel, x :: Vector{Float64})
    y_pred, y_pred_diff = predict_v_and_diff(model, x)
    return Normal(y_pred, abs(y_pred_diff) / sqrt(2))
end

function predict_EI_fg!(g :: Union{Vector{Float64}, Nothing}, model :: IDONEModel, x :: Vector{Float64}, f_best :: Float64)
    # Mixed Zx, gradient.
    vs = model.W * x .+ model.b
    Zx = id_relu.(vs)
    gZx = ∇id_relu.(vs)

    # Predict y value
    y_pred = dot(model.c, Zx)
    # Predict y variance.
    y_pred_diff = dot(model.Δc, Zx)
    # Predict gradient for y value.
    if !isnothing(g)
        gZx .*= model.c
        g_y_pred = model.W' * gZx
        # And for its' corresponding variance.
        gZx .*= model.Δc
        g_y_pred_diff = model.W' * gZx
    end
    # Estimated distribution
    μ = y_pred

    if x == model.last_x_upd
        #  Special case. Define σ = 0.
        if !isnothing(g)
            g .= g_y_pred
        end
        return y_pred - f_best
    end

    σ = abs(y_pred_diff) / sqrt(2)
    # Update gradient for σ
    if !isnothing(g)
        g_y_pred_diff .*= sign(y_pred_diff)
        g_y_pred_diff .*= sqrt(2)
    end

    # Calculate Expected Improvement and its corresponding gradient.
    n01 = Normal()

    # Calculate expected improvement
    # Calculate difference between computed mu and f_best
    # update gradient g_y_pred to be the gradient of diff_f_best_μ
    diff_f_best_μ = f_best - μ

    if !isnothing(g)
        g_y_pred .*= -1
    end
    #
    diff_f_best_μ_div_sigma = diff_f_best_μ / σ

    # cdf_diff_f_best_μ = cdf(n01, diff_f_best_μ_div_sigma)
    # pdf_diff_f_best_μ = pdf(n01, diff_f_best_μ_div_sigma)
    # 1 - to flip (we want to minimize the value, not maximize.)
    # integral goes from (f_best − μ)/σ to ∞ instead of -∞ to (f_best − μ)/σ.
    cdf_diff_f_best_μ = 1 - cdf(n01, diff_f_best_μ_div_sigma)
    pdf_diff_f_best_μ = - pdf(n01, diff_f_best_μ_div_sigma)
    ei = diff_f_best_μ * cdf_diff_f_best_μ + σ * pdf_diff_f_best_μ

    if !isnothing(g)
        g .=
            # Grad of diff_f_best_μ * cdf_diff_f_best_μ
            g_y_pred .* cdf_diff_f_best_μ .+
            (diff_f_best_μ * pdf_diff_f_best_μ) .* g_y_pred ./ g_y_pred_diff .+
            # Grad of σ * pdf_diff_f_best_μ
            g_y_pred_diff .* pdf_diff_f_best_μ .+
            (σ * -diff_f_best_μ_div_sigma * pdf_diff_f_best_μ) .* g_y_pred ./ g_y_pred_diff
    end
    return ei#, g
end

##

function max_EI(μ, σ, f_best)
    if σ < 1e-8
        # Small sigma -> full certainty.
        # Which means μ will be sampled with 100% certainty:
        # eg. it is either a full on improvement (if μ - f_best > 0)
        # or no improvement at all.
        return max(μ - f_best, 0.)
    end
    u = Normal()
    f_best_min_μ = f_best - μ
    f_best_min_μ_div_σ = f_best_min_μ / σ
    return f_best_min_μ * cdf(u, f_best_min_μ_div_σ) + σ * pdf(u, f_best_min_μ_div_σ)
end

function max_EI(n :: Normal, f_best)
    return max_EI(n.μ, n.σ, f_best)
end

function min_EI(μ, σ, f_best)
    if σ < 1e-8
        # Small sigma -> full certainty.
        # Which means μ will be sampled with 100% certainty:
        # eg. it is either a full on improvement (if μ - f_best > 0)
        # or no improvement at all.
        return min(μ - f_best, 0.)
    end
    u = Normal()
    f_best_min_μ = f_best - μ
    f_best_min_μ_div_σ = f_best_min_μ / σ
    return f_best_min_μ * (1. - cdf(u, f_best_min_μ_div_σ)) - σ * pdf(u, f_best_min_μ_div_σ)
end

function min_EI(n :: Normal, f_best)
    return min_EI(n.μ, n.σ, f_best)
end

##
using NLSolversBase
using LineSearches
using Optim

struct IDONE
    model :: IDONEModel
    f :: Function
    best_x :: Vector{Float64}
    best_y :: Ref{Float64}
    y0 :: Float64
    max_iter :: Int64
    current_iter :: Ref{Int64}
end

function scale_f(y0 :: Float64, y :: Float64)
    if abs(y0) > 1e-8
        return (y - y0) / abs(y0)
    else
        return (y - y0)
    end
end

function inv_scale_f(y0 :: Float64, y :: Float64)
    if abs(y0) > 1e-8
        return y * abs(y0) + y0
    else
        return y + y0
    end
end

function optimize_surrogate(model :: IDONEModel, initial_x :: Vector{Float64}, f_best :: Float64)

    df = OnceDifferentiable(
        (x) -> predict(model, x),
        (g, x) -> ∇predict!(g, model, x),
        ones(model.d)
        )

    # function fg!(F, G, x)
    #     return predict_EI_fg!(G, model, x, f_best)
    # end

    # df = OnceDifferentiable(only_fg!(fg!), ones(model.d))

    # -- Optimize using Optim.jl
    # This seems to occasonally run into errors, likely related
    # to their implementation, as manual inspection indicates the
    # model is well behaved at the point of error.
    # Backtracking Linesearch has issues with NaN values as well.
    # Which seems to be caused by being used together with Fminbox.
    optimizer = Fminbox()
    # optimizer = Fminbox(ParticleSwarm())
    # optimizer = Fminbox(LBFGS(linesearch=LineSearches.BackTracking()))
    options = Optim.Options(iterations=50)

    surrogate_solution = nothing
    try
        # surrogate_solution = optimize(df, model.lb, model.ub, initial_x, optimizer, options)
        surrogate_solution = optimize(df, model.lb .- 0.2, model.ub .+ 0.2, initial_x, optimizer, options)
    catch e
        # println("Optimizer failed: $e")
    end
    retries = 0
    while (isnothing(surrogate_solution) || any(isnan.(surrogate_solution.minimizer)) || isnan(surrogate_solution.minimum)) && retries < 10
        try
            surrogate_solution = optimize(df, model.lb, model.ub, initial_x, optimizer, options)
        catch e
            # println("Optimizer failed: $e")
        end
        retries += 1
    end

    if retries >= 10
        return
    end
    # println("Iteration [$(idone.current_iter[])/$(idone.max_iter)] Model indicates $(inv_scale(idone.y0, surrogate_solution.minimum)) at $(surrogate_solution.minimizer).")
    next_x = round.(surrogate_solution.minimizer)
    pred_y = surrogate_solution.minimum

    return next_x, pred_y
end

function variation!(next_x :: Vector{Float64}, idone :: IDONE)
    # Explore if not the last evaluation.
    if idone.current_iter != idone.max_iter
        # Random fuzzy exploration.
        for j in 1:idone.model.d
            # 1/d probability of exploring.
            rand() > 1/idone.model.d && continue
            lbounded = next_x[j] == idone.model.lb[j]
            ubounded = next_x[j] == idone.model.lb[j]
            # Bounds only allow one value.
            lbounded && ubounded && continue
            # Exploration direction is fixed.
            lbounded && (next_x[j] += 1; continue)
            ubounded && (next_x[j] -= 1; continue)
            # Pick randomly otherwise
            next_x[j] += rand(Bool) * 2 - 1
        end
    end
end

function step!(idone :: IDONE)

    # Obtain next point according to model.
    # next_x, pred_y = optimize_surrogate(idone.model, idone.best_x)
    result = optimize_surrogate(idone.model, idone.best_x, idone.best_y[])
    if isnothing(result)
        return
    end

    next_x, pred_y = result

    # Clamp it to be certain.
    next_x .= clamp.(next_x, idone.model.lb, idone.model.ub)
    # Apply variation
    variation!(next_x, idone)

    # @assert all(idone.model.lb .<= next_x)
    # @assert all(next_x .<= idone.model.ub)
    next_x .= clamp.(next_x, idone.model.lb, idone.model.ub)

    # Evaluate function at suggested point next_x
    next_y = idone.f(next_x)
    scaled_next_y = scale_f(idone.y0[], next_y)

    # println("Iteration [$(idone.current_iter[])/$(idone.max_iter)] Scaled $(scaled_next_y) found $(pred_y) predicted")

    # println("Predicted before update: $(predict(idone.model, next_x))")
    update_model!(idone.model, next_x, scaled_next_y)
    # println("Predicted after update: $(predict(idone.model, next_x)), should be $(scaled_next_y)")

    idone.current_iter[] % 1 == 0 &&
        println("Iteration [$(idone.current_iter[])/$(idone.max_iter)] Found value $(next_y) at $(next_x)")

    # Update best
    if idone.best_y[] > next_y
        idone.best_x .= next_x
        idone.best_y[] = next_y
    end

    # Next iteration
    idone.current_iter[] += 1
end

function optimize_IDONE(f :: Function, d :: Int64, lb :: Vector{Float64}, ub :: Vector{Float64}, n_eval=20)
    model = generate_model(d, lb, ub)
    # Start in the middle
    best_x = round.(lb .+ (ub .- lb) ./ 2)
    # println("$(best_x)")
    y0 = f(best_x)
    best_y = y0
    # Note, scaled_y = (y - y0) / abs(y0)
    update_model!(model, best_x, 0.0)

    #
    id = IDONE(model, f, best_x, Ref(best_y), y0, n_eval-1, Ref(1))

    for _ in 1:n_eval-1
        step!(id)
    end

    return id.best_y[], id.best_x, id.model
end

rosen_wrong(x) = sum(100.0 .* (x[1:end] .- x[end:-1:1] .^ 2.0) .^ 2.0 .+ (1 .- x[end:-1:1]) .^ 2.0)
rosen(x) = sum(100.0 .* (x[2:end] .- x[1:end-1] .^ 2.0) .^ 2.0 .+ (1 .- x[1:end-1]) .^ 2.0)

function nqueens(x)
    n = length(x)
    if !(all(x .> 0) && all(x .<= n))
        error("Incorrect indices in $x")
    end
    queens_in_column = zeros(length(x))
    queens_in_primary_diagonal = zeros(2 * length(x) - 1)
    queens_in_perp_diagonal = zeros(2 * length(x) - 1)
    score = 0.
    for (i, p) in enumerate(round.(Int64, x))
        score += queens_in_column[p]
        queens_in_column[p] += 1
        Pdiag = p - i + n
        score += queens_in_primary_diagonal[Pdiag]
        queens_in_primary_diagonal[Pdiag] += 1
        Pdiag = i + p - 1
        score += queens_in_perp_diagonal[Pdiag]
        queens_in_perp_diagonal[Pdiag] += 1
    end
    return score
end


module MMKP

using JuMP

export load_folder

struct Item
    profit::Int64
    resources::Vector{Int64}
end
function Item(v::Vector{Int64})
    Item(first(v), v[2:end])
end

const Group = Vector{Item}

struct MMKP_Data
    resource_bounds::Vector{Int}
    groups::Vector{Group}
end

mutable struct MMKP_Prob
    id::Dict
    model::Union{Model,Nothing}
end

include("loaders.jl")
include("moi_model.jl")
include("load_mps.jl")
problems = load_problems()

function get_id(path)
    class, instance = splitpath(path)[end-1:end]
    Dict("class"=>class, "instance"=>first(split(instance, ".")))
end

function load_file(path, loader; n_threads=4)
    id = get_id(path)
    problem = loader(path)
    model = jump_model(problem, n_threads)
    MMKP_Prob(id, model)
end

_get_loader(path) = first(filter(i-> !isnothing(i),
        map(fl->occursin(fl[1], path) ? fl[2] : nothing, folder_keys)))
function load_subfolder(path)
    l = _get_loader(path)
    map(f->load_file(f, l), readdir(path, join=true))
end

function load_folder(path=joinpath(@__DIR__, "benchmark_problems"))
    vcat(load_subfolder.(readdir(path, join=true))...)
end

end


"""
$(TYPEDEF)

An expression representing the application of a function to symbolic arguments.

# Fields
$(FIELDS)

# Examples

Operations can be built by application of most built-in mathematical functions
to other [`Expression`](@ref) instances:

```jldoctest
julia> using ModelingToolkit

julia> @variables x y;

julia> op1 = sin(x)
sin(x())

julia> typeof(op1.op)
typeof(sin)

julia> op1.args
1-element Array{Expression,1}:
 x()

julia> op2 = x + y
x() + y()

julia> typeof(op2.op)
typeof(+)

julia> op2.args
2-element Array{Expression,1}:
 x()
 y()
```
"""
struct Operation <: Expression
    """The function to be applied."""
    op::Function
    """The arguments the function is applied to."""
    args::Vector{Expression}
end

Base.isequal(x::Operation,y::Operation) =
    x.op == y.op && length(x.args) == length(y.args) && all(isequal.(x.args,y.args))
Base.isequal(::Operation, ::Number   ) = false
Base.isequal(::Number   , ::Operation) = false
Base.isequal(::Operation, ::Variable ) = false
Base.isequal(::Variable , ::Operation) = false
Base.isequal(::Operation, ::Constant ) = false
Base.isequal(::Constant , ::Operation) = false

Base.show(io::IO, O::Operation) = print(io, convert(Expr, O))

# For inv
Base.convert(::Type{Operation}, x::Number) = Operation(identity, Expression[Constant(x)])
Base.convert(::Type{Operation}, x::Operation) = x
Base.convert(::Type{Operation}, x::Expression) = Operation(identity, Expression[x])
Operation(x) = convert(Operation, x)

#convert to Expr
Base.Expr(op::Operation) = simplified_expr(op)
Base.convert(::Type{Expr},x::Operation) = Expr(x)


"""
    DifferentialOps

Module defining types describing differential operators and compositions
thereof.
"""
module DifferentialOps

export
    AbstractDifferentialOp,
    Derivative,
    ScaledDerivative,
    DifferentialOpSum,
    LeftNormal,
    RightNormal,
    max_order

import LinearAlgebra: dot

"""
    AbstractDifferentialOp

Represents a general differential operator.
"""
abstract type AbstractDifferentialOp end

"""
    AbstractNormalDirection

Represents the normal direction on a given domain boundary.
"""
abstract type AbstractNormalDirection end

"""
    LeftNormal <: AbstractNormalDirection

Specifies the normal direction on the left boundary of a 1D domain.

The left normal direction goes opposite to the coordinate axis.
"""
struct LeftNormal <: AbstractNormalDirection end

"""
    RightNormal <: AbstractNormalDirection

Specifies the normal direction on the right boundary of a 1D domain.

The right normal direction is equal to that of the coordinate axis.
"""
struct RightNormal <: AbstractNormalDirection end

Broadcast.broadcastable(op::AbstractDifferentialOp) = Ref(op)
Broadcast.broadcastable(dir::AbstractNormalDirection) = Ref(dir)

"""
    max_order(op::AbstractDifferentialOp)
    max_order(ops...)

Get maximum derivative order of one or more differential operators.
"""
function max_order end

"""
    dot(op::AbstractDifferentialOp, dir::AbstractNormalDirection) -> AbstractDifferentialOp

Project derivative along a normal direction.

This should be used to convert from a [normal
derivative](https://en.wikipedia.org/wiki/Directional_derivative#Normal_derivative)
at the boundaries, to a derivative along the coordinate axes of the domain.

In practice, this returns `op` for [`RightNormal`](@ref).
For [`LeftNormal`](@ref), it multiplies the odd-order derivatives by -1.
"""
dot(op::AbstractDifferentialOp, dir::AbstractNormalDirection) = dot(dir, op)

dot(::RightNormal, op) = op

max_order(ops::Vararg{AbstractDifferentialOp}) = max(max_order.(ops)...)

"""
    Derivative{n} <: AbstractDifferentialOp

Specifies the `n`-th derivative of a function.
"""
struct Derivative{n} <: AbstractDifferentialOp end

Derivative(n::Integer) = Derivative{n}()
Derivative() = Derivative(1)

function Base.literal_pow(::typeof(^), op::Derivative{n}, ::Val{p}) where {n,p}
    Derivative(n * p)
end

max_order(::Derivative{n}) where {n} = n

# We return a ScaledDerivative for type stability.
dot(dir::LeftNormal, D::Derivative) = dot(dir, 1 * D) :: ScaledDerivative

Base.show(io::IO, D::Derivative{n}) where {n} = print(io, "D{", n, "}")

"""
    ScaledDerivative{n} <: AbstractDifferentialOp

`n`-th derivative of a function scaled by a constant coefficient.
"""
struct ScaledDerivative{n,T<:Number} <: AbstractDifferentialOp
    D :: Derivative{n}
    α :: T
    function ScaledDerivative(D::Derivative{n}, α::T) where {n,T}
        # If we allow α = 0, then the output of functions such as max_order
        # is not a compile-time constant.
        iszero(α) && throw(ArgumentError("scale factor cannot be zero"))
        new{n,T}(D, α)
    end
end

Base.show(io::IO, S::ScaledDerivative) = print(io, S.α, " * ", S.D)
max_order(S::ScaledDerivative) = max_order(S.D)
Base.:-(S::ScaledDerivative) = ScaledDerivative(S.D, -S.α)

Base.:(==)(x::ScaledDerivative, y::ScaledDerivative) = x.D == y.D && x.α == y.α
Base.:(==)(x::ScaledDerivative, y::Derivative) = x.D == y && x.α == 1
Base.:(==)(x::Derivative, y::ScaledDerivative) = y == x

function dot(::LeftNormal, S::ScaledDerivative{n}) where {n}
    r = isodd(n) ? -1 : 1
    ScaledDerivative(S.D, r * S.α)
end

Base.:*(α::Number, D::Derivative) = ScaledDerivative(D, α)
Base.:*(D::Derivative, α) = α * D
Base.:-(D::Derivative) = -1 * D

"""
    DifferentialOpSum <: AbstractDifferentialOp

Sum of two differential operators.
"""
struct DifferentialOpSum{
        OpA <: AbstractDifferentialOp,
        OpB <: AbstractDifferentialOp,
    } <: AbstractDifferentialOp
    a :: OpA
    b :: OpB
end

Base.:+(a::AbstractDifferentialOp, b::AbstractDifferentialOp) = DifferentialOpSum(a, b)
Base.:-(a::AbstractDifferentialOp, b::AbstractDifferentialOp) = DifferentialOpSum(a, -b)

Base.:(==)(x::DifferentialOpSum, y::DifferentialOpSum) =
    (x.a == y.a && x.b == y.b) || (x.a == y.b && x.b == y.a)

Base.show(io::IO, D::DifferentialOpSum) = print(io, D.a, " + ", D.b)

max_order(D::DifferentialOpSum) = max_order(D.a, D.b)

dot(dir::LeftNormal, D::DifferentialOpSum) = DifferentialOpSum(dot.(dir, (D.a, D.b))...)

end


# function plotmesh(m :: TriMesh;
#                         linewidth :: Real = 1,
#                         marker :: String = "None",
#                         markersize :: Real = 10,
#                         color = :black)
#     p = plot()
#     for edge in eachcol(m.edge)
#         xy = m.point[:, edge]
#         plot!(p, xy[1, :], xy[2, :], linewidth=linewidth,
#             markersize=markersize, color=color, legend=false)
#     end
#     return p
# end
#

"""
    make_border_points(points, n, expansion, k = size(points, 2)) -> Matrix

Create a set of border points around the hull of a domain.

# Arguments
- `points`: points to form a hull around, as a matrix with a point in each column
- `n`: number of edge points
- `expansion`: factor to extend the hull beyond the provided points
- `k`: number of neighbors to use when determining hull. Default is all points, giving a convex hull; can be decreased to make the hull concave.

# Returns
- `Matrix`: coordinates of domain border
"""
function make_border_points(points, n, expansion, k=size(points, 2))
    hull = concave_hull(collect(eachcol(points)), k)
    cm = mean(points, dims=2)
    points2 = cm .+ 1.01 * (points .- cm)
    edge_points = reduce(hcat, [p for p in eachcol(points2) if ! in_hull(p, hull)])

    θ = [atan(y, x) for (x, y) in eachcol(edge_points .- cm)]
    x = edge_points[1, sortperm(θ)]
    y = edge_points[2, sortperm(θ)]
    θ = [θ; θ[1]]
    x = [x; x[1]]
    y = [y; y[1]]

    ds = vec(sqrt.(sum(diff([x y]', dims=2).^2, dims=1)))
    s = cumsum([0; ds])
    xinterp = LinearInterpolation(s, x)
    yinterp = LinearInterpolation(s, y)

    Δs = (s[end] - s[1]) / n
    ss = s[1]:Δs:s[end]
    border_points = [xinterp(ss) yinterp(ss)]'
    border_points = (border_points .- cm) * expansion .+ cm
end

"""
    inverse_distance_weights(dists::AbstractVector) -> Vector

Calculate the relative weights via inverse distance.

# Arguments
- `dists::AbstractVector`: vector of distances

# Returns
- `Vector`: relative weight of points at each distance
"""
function inverse_distance_weights(dists::AbstractVector)
    w = [d > 0 ? 1/d : 1/(d+eps()) for d in dists]
    return w / sum(w)
end

"""
    observation_matrix(mesh::TriMesh, points::AbstractMatrix) -> SparseMatricCSC

Create matrix of weights to linearly interpolate from nodes to observation locations.

# Arguments
- `mesh::TriMesh`: a mesh from TriangleMesh.jl
- `points::AbstractMatrix`: observation locations in columns

# Returns
- `SparseMatrixCSC`: sparse matrix with weights
"""
function observation_matrix(mesh::TriMesh, points::AbstractMatrix)
    npoint = size(points, 2)
    nmesh = size(mesh.point, 2)
    tree = KDTree(mesh.point)
    idx, dists = knn(tree, points, 3)

    ii = repeat(1:npoint, inner=3)
    jj = reduce(vcat, idx)
    ww = mapreduce(inverse_distance_weights, vcat, dists)
    return sparse(ii, jj, ww, npoint, nmesh)
end


const YEAR  = 365.242
const GNEWT = 39.4845/YEAR^2  # Units of MSUN*AU^3/YEAR^2
const NDIM  = 3
const third = 1.0/3.0

include("../src/init_nbody.jl")

#@testset "init_nbody" begin

elements = readdlm("elements.txt",',')

#elements[2:3,1] /= 100.0
#n_body = 4
n_body = 3
t0 = 7257.93115525-7300.0
elements[:,3] .-= 7300.0

jac_init     = zeros(Float64,7*n_body,7*n_body)
jac_init_big = zeros(BigFloat,7*n_body,7*n_body)
jac_init_num = zeros(BigFloat,7*n_body,7*n_body)
x,v = init_nbody(elements,t0,n_body,jac_init)
elements_big = big.(elements)
t0big = big(t0)
xbig,vbig = init_nbody(elements_big,t0big,n_body,jac_init_big)
elements0 = copy(elements)
#dq = big.([1e-10,1e-5,1e-6,1e-6,1e-6,1e-5,1e-5])
#dq = big.([1e-10,1e-8,1e-8,1e-8,1e-8,1e-8,1e-8])
dq = big.([1e-12,1e-10,1e-10,1e-10,1e-10,1e-10,1e-10])

#dq = big.([1e-15,1e-15,1e-15,1e-15,1e-15,1e-15,1e-15])
t0big = big(t0)
# Now, compute derivatives numerically:
for j=1:n_body
  for k=1:7
    elementsbig = big.(elements0)
    dq0 = dq[k]; if j==1 && k==1 ; dq0 = big(1e-15); end
    elementsbig[j,k] += dq0
    xp,vp = init_nbody(elementsbig,t0big,n_body)
    elementsbig[j,k] -= 2dq0
    xm,vm = init_nbody(elementsbig,t0big,n_body)
    for l=1:n_body, p=1:3
      i1 = (l-1)*7+p
      if k == 1
        j1 = j*7
      else
        j1 = (j-1)*7+k-1
      end
      jac_init_num[i1,  j1] = (xp[p,l]-xm[p,l])/dq0*.5
      jac1 = jac_init[i1,j1]; jac2 = jac_init_num[i1,j1]
      if abs(jac1-jac2) > 1e-4*abs(jac1+jac2) && abs(jac1+jac2) > 1e-14
        println(l," ",p," ",j," ",k," ",jac_init_num[i1,j1]," ",jac_init[i1,j1]," ",jac_init_num[i1,j1]/jac_init[i1,j1])
      end
      jac_init_num[i1+3,j1] = (vp[p,l]-vm[p,l])/dq0*.5
      jac1 = jac_init[i1+3,j1]; jac2 = jac_init_num[i1+3,j1]
      if abs(jac1-jac2) > 1e-4*abs(jac1+jac2) && abs(jac1+jac2) > 1e-14
        println(l," ",p+3," ",j," ",k," ",jac_init_num[i1+3,j1]," ",jac_init[i1+3,j1]," ",jac_init_num[i1+3,j1]/jac_init[i1+3,j1])
      end
    end
  end
  jac_init_num[j*7,j*7]=1.0
end

#println("Maximum jac_init-jac_init_num: ",maximum(abs.(jac_init-jac_init_num)))
println("Maximum jac_init-jac_init_num: ",maximum(abs.(asinh.(jac_init)-asinh.(jac_init_num))))
println("Maximum jac_init-jac_init_big: ",maximum(abs.(asinh.(jac_init)-asinh.(jac_init_big))))
println("Maximum jac_init/jac_init_big-1: ",maximum(abs.(jac_init[jac_init_big .!= 0.0]./jac_init_big[jac_init_big .!= 0.0].-1)))
println("Maximum jac_init_big-jac_init_num: ",maximum(abs.(asinh.(jac_init_num)-asinh.(jac_init_big))))
println(convert(Array{Float64,2},jac_init_big-jac_init_num))
#@test isapprox(jac_init_num,jac_init)
@test isapprox(jac_init_num,jac_init;norm=maxabs)
@test isapprox(jac_init,convert(Array{Float64,2},jac_init_big);norm=maxabs)
#end

for i=1:21, j=1:21
  if jac_init_big[i,j] != 0.0
    dij = abs(jac_init[i,j]/jac_init_big[i,j]-1)
    if dij > 1e-12
      jac_max = dij
      println(i," ",j," ",convert(Float64,jac_max)," ",jac_init[i,j]," ",convert(Float64,jac_init_big[i,j])," ",convert(Float64,jac_init[i,j]-jac_init_big[i,j]))
    end
  end
end


export FixedParameterHomotopy, fix_parameters

"""
    FixedParameterHomotopy(H:AbstractHomotopy, parameters)

Construct a homotopy from the given [`AbstractHomotopy`](@ref) `H` with the given `parameters`
fixed.
"""
struct FixedParameterHomotopy{S<:AbstractHomotopy,T} <: AbstractHomotopy
    homotopy::S
    parameters::Vector{T}
end
FixedParameterHomotopy(H::Homotopy, p; compile::Union{Bool,Symbol} = COMPILE_DEFAULT[]) =
    FixedParameterHomotopy(fixed(H; compile = compile), p)
FixedParameterHomotopy(
    H::AbstractHomotopy,
    p;
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
) = FixedParameterHomotopy(H, p)
Base.size(H::FixedParameterHomotopy) = size(H.homotopy)

ModelKit.variables(H::FixedParameterHomotopy) = variables(H.homotopy)
ModelKit.parameters(H::FixedParameterHomotopy) = Variable[]
ModelKit.variable_groups(H::FixedParameterHomotopy) = variable_groups(H.homotopy)

(H::FixedParameterHomotopy)(x, t, p = nothing) = H.homotopy(x, t, H.parameters)

ModelKit.evaluate!(u, H::FixedParameterHomotopy, x, t) =
    evaluate!(u, H.homotopy, x, t, H.parameters)
ModelKit.evaluate_and_jacobian!(u, U, H::FixedParameterHomotopy, x, t) =
    evaluate_and_jacobian!(u, U, H.homotopy, x, t, H.parameters)
ModelKit.taylor!(u, v::Val, H::FixedParameterHomotopy, tx, t) =
    taylor!(u, v, H.homotopy, tx, t, H.parameters)

"""
    fix_parameters(H::Union{Homotopy,AbstractHomotopy}, p; compile::Union{Bool,Symbol} = $(COMPILE_DEFAULT[]))

Fix the parameters of the given homotopy `H`. Returns a [`FixedParameterHomotopy`](@ref).
"""
fix_parameters(
    H::Union{Homotopy,AbstractHomotopy},
    p;
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
) = FixedParameterHomotopy(H, p; compile = compile)


# indicator of a PSD

"""
  IndPSD()

Returns the function `g = ind{A : A ⪰ 0}`, i.e. the indicator of the positive semidefinite cone.
The argument to the function can be either a Symmetric or Hermitian object.
From Julia 0.5, the argument can also be an AbstractVector{Float64} holding a symmetric matrix in (lower triangular) packed storage.
"""

immutable IndPSD <: IndicatorConvexCone end

function (f::IndPSD){T <: RealOrComplex}(X::HermOrSym{T})
  F = eigfact(X);
  for i in eachindex(F.values)
    #Do we allow for some tolerance here?
    if F.values[i] <= -1e-14
      return +Inf
    end
  end
  return 0.0
end

function prox!{T <: RealOrComplex}(f::IndPSD, X::HermOrSym{T}, Y::HermOrSym{T}, gamma::Real=1.0)
  n = size(X, 1);
  F = eigfact(X);
  for i in eachindex(F.values)
    F.values[i] = max.(0.0, F.values[i]);
  end
  for i = 1:n
    for j = 1:n
      Y.data[i,j] = 0.0
      for k = 1:n
        Y.data[i,j] += F.vectors[i,k]*F.values[k]*F.vectors[j,k]
      end
    end
  end
  return 0.0
end

fun_name(f::IndPSD) = "indicator of positive semidefinite cone"
fun_dom(f::IndPSD) = "Symmetric, Hermitian, AbstractArray{Float64}"
fun_expr(f::IndPSD) = "x ↦ 0 if A ⪰ 0, +∞ otherwise"
fun_params(f::IndPSD) = "none"

function prox_naive{T <: RealOrComplex}(f::IndPSD, X::HermOrSym{T}, gamma::Real=1.0)
  F = eigfact(X);
  return F.vectors * diagm(max.(0.0, F.values)) * F.vectors', 0.0;
end

### Below: with AbstractVector argument

function (f::IndPSD){T <: Float64}(x::AbstractVector{T})
  Z = dspev!('N', 'L', copy(x))
  for i in 1:length(Z)
    #Do we allow for some tolerance here?
    if Z[i] <= -1e-14
      return +Inf
    end
  end
  return 0.0
end

function prox!(f::IndPSD, x::AbstractVector{Float64}, y::AbstractVector{Float64}, gamma::Real=1.0)
  y[:] = x              # Copy x since dspev! corrupts input
  (W, Z) = dspevV!('L', y)
  W = max.(W, 0.0)         # NonNeg eigenvalues
  M = Z.*W'             # Equivalent to Z*diagm(W) without constructing W matrix
  M = M*Z'              # Now let M = Z*diagm(W)*Z'
  n = length(W)
  k = 1
  for j in 1:n, i in j:n  # Store lower diagonal of M in y
    y[k] = M[i,j]
    k = k+1
  end
  return 0.0
end

function prox_naive{T<:Float64}(f::IndPSD, x::AbstractVector{T}, gamma::Real=1.0)
  n = Int(sqrt(1/4+2*length(x))-1/2)  # Formula for size of matrix
  X = Array{T,2}(n,n)
  k = 1
  for j = 1:n, i = j:n                # Store y in M
    X[i,j] = x[k]                   # Lower half
    if i != j
      X[j,i] = x[k]               # Strictly upper half
    end
    k = k+1
  end
  X, v = prox_naive(f, Symmetric(X), gamma)
  y = similar(x)
  k = 1
  for j = 1:n, i = j:n                # Store Lower half of X in y
    y[k] = X[i,j]
    k = k+1
  end
  return y, 0.0
end


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

using Test, Tensors, LinearAlgebra

# Kronecker delta
@test delta(1, 1) == 1
@test delta(1, 2) == 0
@test_throws MethodError delta(1.0, 2.0)  # indices must be integers
@test_throws MethodError delta(1, BigInt(2))  # both must have the same type
@test delta(1, 2) isa Int  # the output type matches the input
@test delta(BigInt(1), BigInt(2)) isa BigInt

# Various tensors
let Z3 = zeros(3, 3),
    O3 = ones(3, 3)
    @test isapprox(tovoigt(II()), I)
    @test isapprox(tovoigt(IT()), [I   Z3;
                                   Z3  Z3])

    @test isapprox(tovoigt(IS()), [I   Z3;
                                   Z3  1//2*I])
    @test isapprox(tovoigt(IA()), [Z3  Z3;
                                   Z3  1//2*I])

    @test isapprox(tovoigt(IV()), [1//3*O3  Z3;
                                   Z3       Z3])
    @test isapprox(tovoigt(ID()), [(I - 1//3*O3)  Z3;
                                   Z3             1//2*I])

    @test let lambda = 10.0,
              mu = 1.0
        isapprox(tovoigt(isotropic_elasticity_tensor(lambda, mu)), [(lambda*O3 + 2*mu*I)  Z3;
                                                                    Z3                    mu*I])
    end
end

# Lamé parameters for isotropic solids
@test all(isapprox(result, expected)
        for (result, expected) in zip(lame(1e11, 0.3), (5.769230769230769e10, 3.846153846153846e10)))
@test all(isapprox(result, expected)
        for (result, expected) in zip(delame(lame(1e11, 0.3)...), (1e11, 0.3)))

# Mutating function to non-mutating function conversion
let  # introduce a local scope so the name `f!` is only defined locally for this test.
    function f!(out, x)
        out[:] = [sin(elt) for elt in x]
        return nothing
    end

    let
        out = [0.0]
        @test all([f!(out, [pi/4]) == nothing,
                 isapprox(out, [1/sqrt(2)])])
    end

    let
        out = [0.0]
        f = debang(f!)
        @test f isa Function
        @test all([isapprox(f([pi/4]), [1/sqrt(2)]),
                 out == [0.0]])
    end
end

# Newton root finder
let g(x) = [(1 - x[1]^2) + x[2]],
    x0 = [0.8, 0.2]
    @test !isapprox(g(x0), [0.0], atol=1e-15)  # find_root should have to actually do something
    @test isapprox(g(find_root(g, x0)), [0.0], atol=1e-15)
end


module TestingMayOptimizeLinearAlgebraMethods

export ldiv_tests, cholesky_tests

using Test
using LinearAlgebra
using LinearAlgebra: AbstractTriangular
using MayOptimize
using MayOptimize:
    CodeChoice,
    Standard,
    AVX
using MayOptimize.LinearAlgebraMethods:
    AbstractAlgorithm,
    CholeskyBanachiewiczLower,
    CholeskyBanachiewiczLowerI,
    CholeskyBanachiewiczLowerII,
    CholeskyBanachiewiczUpper,
    CholeskyCroutLower,
    CholeskyCroutUpper,
    CholeskyCroutUpperI,
    CholeskyCroutUpperII,
    CholeskyFactorization,
    CholeskyLower,
    CholeskyUpper,
    Floats,
    exec!

# Create an array initiallized with NaN's.
nans(A::AbstractArray) = nans(eltype(A), size(A))
nans(::Type{T}, dims::Integer...) where {T} = nans(T, dims)
nans(::Type{T}, dims::Tuple{Vararg{Integer}}) where {T<:AbstractFloat} =
    fill!(Array{T}(undef, dims), T(NaN))
nans(::Type{Complex{T}}, dims::Tuple{Vararg{Integer}}) where {T<:AbstractFloat} =
    fill!(Array{Complex{T}}(undef, dims), Complex{T}(T(NaN),T(NaN)))

function check(pred::Function, A::AbstractMatrix)
    I, J = axes(A)
    for j ∈ J
        for i ∈ I
            pred(A[i,j], i, j) || return false
        end
    end
    return true
end

function check(pred::Function, A::AbstractMatrix, B::AbstractMatrix)
    @assert axes(A) == axes(B)
    I, J = axes(A)
    for j ∈ J
        for i ∈ I
            pred(A[i,j], B[i,j], i, j) || return false
        end
    end
    return true
end

# Generate random values uniformly distributed in [-1,1].
generate_array(T::Type{<:Real}, dims...) = 2*rand(T, dims...) .- 1
generate_array(T::Type{<:Complex{<:Real}}, dims...) =
    2*rand(T, dims...) .- Complex(1,1)

# Build a random Hermitian positive definite matrix.
function generate_symmetric_definite_matrix(T::Type, n::Integer)
    H = generate_array(T, 10n, n)
    A = H'*H
    return (A + A')/2
end

function ldiv_tests(;
                    types = (Float32,
                             Float64,
                             Complex{Float64},),
                    sizes = (30,),
                    opts  = (Standard,
                             Debug,
                             InBounds,
                             Vectorize,))
    @testset "ldiv!(opt, [y,] A, b) " begin
        for T in types, n in sizes
            A = generate_symmetric_definite_matrix(T, n)
            C = cholesky(A)
            L = C.L
            @test L*L' ≈ A
            R = C.U
            @test R'*R ≈ A
            w = Array{T}(undef, n)
            x = 2*rand(T, n) .- 1
            y = A*x
            z = Array{T}(undef, n)
            for opt in opts
                # ldiv! on lower triangular matrix L.
                @test ldiv!(opt, L, copyto!(z, y)) ≈ L\y
                @test ldiv!(opt, z, L, y) ≈ L\y
                @test ldiv!(opt, L', copyto!(w, z)) ≈ L'\z
                @test ldiv!(opt, w, L', z) ≈ L'\z
                @test ldiv!(opt, L', ldiv!(opt, L, copyto!(z, y))) ≈ x
                @test ldiv!(opt, w, L', ldiv!(opt, z, L, y)) ≈ x
                #@test lmul!(opt, L, lmul!(opt, L', copyto!(z, x))) ≈ y
                #@test lmul!(opt, L, lmul!(opt, z, L', x)) ≈ y
                # ldiv! on upper triangular matrix R.
                @test ldiv!(opt, R', copyto!(z, y)) ≈ R'\y
                @test ldiv!(opt, z, R', y) ≈ R'\y
                @test ldiv!(opt, R, copyto!(w, z)) ≈ R\z
                @test ldiv!(opt, w, R, z) ≈ R\z
                @test ldiv!(opt, R, ldiv!(opt, R', copyto!(z, y))) ≈ x
                @test ldiv!(opt, w, R, ldiv!(opt, z, R', y)) ≈ x
                #@test lmul!(opt, R', lmul!(opt, R, copyto!(z, x))) ≈ y
                #@test lmul!(opt, R', lmul!(opt, z, R, x)) ≈ y
            end
        end
    end
end

function cholesky_tests(;
                        algs  = (CholeskyCroutLower,
                                 CholeskyBanachiewiczLowerI,
                                 CholeskyBanachiewiczLowerII,
                                 CholeskyCroutUpperI,
                                 CholeskyCroutUpperII,
                                 CholeskyBanachiewiczUpper,),
                        types = (Float64,
                                 Complex{Float64},),
                        sizes = (30,),
                        opts  = (Debug,
                                 InBounds,
                                 Vectorize,))
    @testset "Cholesky factorization" begin
        for T in types, n in sizes
            # Build a random Hermitian positive definite matrix.
            A = generate_symmetric_definite_matrix(T, n)
            C = cholesky(A)
            w = Array{T}(undef, n)
            x = 2*rand(T, n) .- 1
            y = A*x
            z = Array{T}(undef, n)
            Cs = cholesky(Standard, A)
            @test Cs.L == C.L
            @test Cs.U == C.U
            for opt in opts, algtype in algs
                if algtype <: CholeskyLower
                    alg = algtype(opt)
                    # Test out-of-place Cholesky factorization.  Checking the validity
                    # of the decomposition and that the strict upper part of L has been
                    # untouched (it is filled with NaN's).
                    L = LowerTriangular(exec!(alg, nans(A), A))
                    @test L*L' ≈ A
                    @test L ≈ C.L
                    @test check((a,i,j) -> (i ≥ j || isnan(a)), parent(L))
                    # Idem for the in-place operation version.
                    X = copy(A)
                    L = LowerTriangular(exec!(alg, X))
                    @test L*L' ≈ A
                    @test L ≈ C.L
                    @test check((a,b,i,j) -> (i ≥ j || a === b), A, X)
                elseif algtype <: CholeskyUpper
                    alg = algtype(opt)
                    # Test out-of-place Cholesky factorization.  Checking the validity
                    # of the decomposition and that the strict lower part of R has been
                    # untouched (it is filled with NaN's).
                    R = UpperTriangular(exec!(alg, nans(A), A))
                    @test R'*R ≈ A
                    @test R ≈ C.U
                    @test check((a,i,j) -> (i ≤ j || isnan(a)), parent(R))
                    # Idem for the in-place operation version.
                    X = copy(A)
                    R = UpperTriangular(exec!(alg, X))
                    @test R'*R ≈ A
                    @test R ≈ C.U
                    @test check((a,b,i,j) -> (i ≤ j || a === b), A, X)
                end
            end
        end
    end
end

end # module


export randEllipic

"""
-  Look for a `phat`, such that when used in NORTA, it will return correlations approximately `ρ`
- `d` : distribution
- `ρ` : target correlation
- NORTA method for generating Random Vectors with Arbitrary Marginal Distribution
    fand CorrelationMatrices first appears in the following paper
    https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.48.281&rep=rep1&type=pdf
    The generated vector entries usually have a smaller correlation than what is aimed for.
"""
struct preNORTA
    d::D where D<:S
    ρ::Float64
    ρhat::Float64
    function preNORTA(d,ρ,ρhat=ρ)
        L , H = minmax(sign(ρ),0) # Lower bound, Higher bound
        ϵ = min(0.01,(1-abs(ρ))/10)
        while true
            N = [rand(MvNormal([1 ρhat;ρhat 1])) for _=1:10000]
            X = []
            Y = []

            for i in N
               append!(X,quantile(d,cdf(Normal(),i[1])))
               append!(Y,quantile(d,cdf(Normal(),i[2])))
            end

            temp = cor(X,Y)
            if abs(temp-ρ) < ϵ || H - L < ϵ
                break
            end
            if temp < ρ
                ρhat, L = (ρhat+H)/2 ,  (L+ρhat*4)/5
            else
                ρhat, H = (ρhat+L)/2 , (H+ρhat*4)/5
            end
        end
        new(d,ρ,ρhat)
    end
end
"""
```julia
randEllipic(d, n; r , diag, norm)

randEllipic(n; r, norm)
```
- `d` : default `Normal()`, entry distribution
- `n`  : dimensions
- `r` : default `0.5`, the correlation of ``H_{ij},H_{ji}`` pairs
- `norm` : default `false`, if `norm` set to `true`, then the matrix will be normalized with ``n^{-1/2}``.
- `diag` : default `diag=d`, the distribution for diagonal entries.

# Examples
Generate a random elliptic matrix, with entries from ``\\mathscr{N}(0,1)`` and ``\\rho(H_{ij},H_{ji})=0.5``
```julia
randEllipic(500)

500×500 Matrix{Float64}:
  2.03417    -0.424289    1.28267   …  -0.114754  -1.96059
  0.44479    -1.45563     1.32828       1.00149    0.45786
  1.56525     0.0832211  -0.186738     -1.3914     1.04151
 -0.11633    -0.483301   -1.81348      -1.57536    0.514818
  ⋮                                 ⋱
  1.24274    -0.411623   -1.04984      -0.812778  -1.84479
 -0.0817287  -0.254886    0.674914      0.756269  -0.0296209
 -1.48281     0.51675    -1.58041       0.156923   0.244599
  0.852339    1.04593    -0.119082      1.43634    0.114493
```

Generate a normalized random elliptic matrix, with entries `Poisson(10)` and ``\\rho(H_{ij},H_{ji})=0.1``
```julia
using Distributions
randEllipic(Poisson(10),500, r=0.1 , norm=true)

500×500 Matrix{Float64}:
  0.268328    -0.0413153  -0.0175096   …   0.0190835   0.0201304
 -0.00599949   0.447214   -0.0175878       0.0112805   0.100704
 -0.0258879    0.0219927   0.402492       -0.04749    -0.050853
  0.0219071    0.0119609  -0.00448502      0.0043233   0.0404757
  ⋮                                    ⋱
 -0.0145467    0.0800297   0.00247891      0.0189267   0.071565
  0.016412     0.0334019   0.0663348      -0.0180889   0.023773
 -0.0485914   -0.0575288  -0.0409827       0.491935   -0.0969691
  0.0405447    0.0503843   0.00624668      0.0558304   0.402492

```
"""
function randEllipic(d::D, n::Int; r = 0.5::Float64, diag=d::D, norm = false::Bool) where  D<:S
    if r==1
        return randHermitian(d,n,diag=diag,norm=norm)
    end
    # TODO: take care of the r=-1 case.

    r = preNORTA(d,r).ρhat
    M = zeros(n,n)
    for i in 1:n
        M[i,i] = rand(diag)
        for j in i+1:n
            M[i,j],M[j,i] = rand(MvNormal([1 r;r 1]))
        end
    end

    if norm
        M /= sqrt(n)
    end
    return M
end

function randEllipic(n::Int; r=0.5::Float64, norm = false::Bool)
    return randEllipic(Normal(),n, r=r, norm=norm)
end

"""
```julia
randStochastic(n; type, norm)
```
- `n`: dimension
- `type` : default `type = 3`, `3` for doubly randStochastic, `1` for row, `2` for column stochastic
- `norm` : default `false`, if set to `true`, the matrix will be normalized by ``\\sqrt{n}`` (*not a typo*)
# Examples

Generates a 3 by 3 random doubly stochastic  matrix
```julia
randStochastic(3)

3×3 Matrix{Float64}:
 0.132593  0.216041  0.651367
 0.484097  0.320777  0.195126
 0.261495  0.537825  0.20068
```

Generates a 3 by 3 normalized random row stochastic  matrix
```julia
randStochastic(3,type = 1)

3×3 Matrix{Float64}:
 0.220849   0.146942  0.632209
 0.188052   0.26294   0.549008
 0.0170714  0.524574  0.458355
```

Generates a 3 by 3 normalized random column stochastic  matrix
```julia
randStochastic(3,type=2,norm=true)

3×3 Matrix{Float64}:
 0.583396  0.608739  0.732921
 0.672821  0.078786  0.302657
 0.475834  1.04453   0.696473
```
"""
function randStochastic(n::Int; type = 3 ::Int, norm = false::Bool)
    if type == 3
       M = randUnitary(n)
       M = M.*conj(M)
    else
        M = rand(n,n)
        for i in 1:n
            M[i,:] /= sum(M[i,:])
        end
        if type == 2
            M = M'
        end
    end

    if norm
        M*=sqrt(n)
    end

    return M
end


##############################   Coulomb  ###################################
mutable struct Coulomb <: EmpiricalPotential
    q_1
    q_2
    ϵ0
    rcutoff
    species :: AbstractVector
end

get_parameters(c::Coulomb) = Parameter{}(())
set_parameters(p::Parameter{}, c::Coulomb) = c

deserialize_parameters(p::Parameter{()}, c::Coulomb) = []
serialize_parameters(p::Vector, c::Coulomb) = Parameter{()}( () )

get_hyperparameters(c::Coulomb) = Parameter{:rcutoff}((c.rcutoff, ))
set_hyperparameters(p::Parameter{(:rcutoff,)}, c::Coulomb) = Coulomb(c.q_1, c.q_2, c.ϵ0, p.rcutoff, c.species)

deserialize_hyperparameters(p::Parameter{(:rcutoff,)}, c::Coulomb) = [p.rcutoff]
serialize_hyperparameters(p::Vector, c::Coulomb) = Parameter{(:rcutoff, )}( (p[1],) )

# ############################# Energies ##########################################

function potential_energy(R::AbstractFloat, c::Coulomb)
    return c.q_1 * c.q_2 / (4.0 * π * c.ϵ0 * R)
end

# ############################### Forces ##########################################

function force(R::AbstractFloat, r::SVector{3,<:AbstractFloat}, c::Coulomb)
    SVector(c.q_1 * c.q_2 / (4.0 * π * c.ϵ0 * R^2) .* r ./ R)
end

# ##############################   Gradients  ###################################

function grad_potential_energy(r::Vector{<:Real}, p::Coulomb)
    error("The Coulomb potential has no trainable parameters")
end
function grad_force(r:: Vector{<:Real}, p::Coulomb)
    error("The Coulomb potential has no trainable parameters")
end

function grad_virial(r::Vector{<:Real}, p::Coulomb)
    error("The Coulomb potential has no trainable parameters")
end


# Forest-Ruth and Position Extended Forest-Ruth Like algorithm
# Leapfrog is 2nd order, FR and PEFRL are 4th order.

function fr_step!(r::AbstractVecOrMat{V3{T}}, v::AbstractVecOrMat{V3{T}},
         m::AbstractVector; G, dt) where T
    nplanets = length(m)
    a = zeros(V3{T}, nplanets)
    θ = 1/(2-2^(1/3))
    @inbounds begin
        for j=1:nplanets
            r[j] += θ/2*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += θ*dt*a[j]
            r[j] += (1-θ)/2*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += (1-2θ)*dt*a[j]
            r[j] += (1-θ)/2*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += θ*dt*a[j]
            r[j] += θ/2*dt*v[j]
        end
    end
    return r, v
end

function pefrl_step!(r::AbstractVecOrMat{V3{T}}, v::AbstractVecOrMat{V3{T}},
         m::AbstractVector; G, dt) where T
    nplanets = length(m)
    a = zeros(V3{T}, nplanets)
    ξ = +0.1786178958448091E+00
    λ = -0.2123418310626054E+00
    χ = -0.6626458266981849E-01
    @inbounds begin
        for j=1:nplanets
            r[j] = r[j]
            v[j] = v[j]
            r[j] += ξ*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += (1-2λ)/2*dt*a[j]
            r[j] += χ*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += λ*dt*a[j]
            r[j] += (1-2(ξ+χ))*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += λ*dt*a[j]
            r[j] += χ*dt*v[j]
        end
        update_acceleration!(a, r, m, G)
        for j=1:nplanets
            v[j] += (1-2λ)/2*dt*a[j]
            r[j] += ξ*dt*v[j]
        end
    end
    return r, v
end

function fr!(r::AbstractVecOrMat{V3{T}}, v::AbstractVecOrMat{V3{T}},
         m::AbstractVector; G=G_year_AU, n, dt) where T
    leapfrog_looper!(FR,r, v, m; G=G, n=n, dt=dt)
end

function pefrl!(r::AbstractVecOrMat{V3{T}}, v::AbstractVecOrMat{V3{T}},
         m::AbstractVector; G=G_year_AU, n, dt) where T
    leapfrog_looper!(PEFRL,r, v, m; G=G, n=n, dt=dt)
end

function fr(planets::AbstractVector{Body{T}}; G=G_year_AU, n, dt, keep_history=false) where {T}
    leapfrog(FR, planets; G=G, n=n, dt=dt, keep_history=keep_history)
end

function pefrl(planets::AbstractVector{Body{T}}; G=G_year_AU, n, dt, keep_history=false) where {T}
    leapfrog(PEFRL, planets; G=G, n=n, dt=dt, keep_history=keep_history)
end

steper(::Type{FR}) = fr_step!
steper(::Type{PEFRL}) = pefrl_step!


export Dist2Set, LossTotal, IoU, g_CQ, MultipleSplitFeasCQ_fun_grad

function Dist2Set(input,P,TrOpts)
  #active_z_slice is only used for hyperspectral imaging where we map a 3D or 4D
  #data volume to a 2D map of the earth. this 2D map is located at slice number
  #active_z_slice in a x-y-z-nchan-n_ex tensor

  #initialize some stuff
  size_input = size(input)
  pen_grad   = zeros(Float32,size(input))
  pen_value  = 0f0

  #use softmax of final network state for constraining
  #the output for segmentation problems (don't use softmax for non-linear regression)
  #input      = softmax(Array(input)[:,:,TrainOpts.channels,:],dims=3)

  input      = softmax(input[:,:,TrOpts.active_channels,:],dims=3)

  for j in TrOpts.active_channels #loop over channels, each channel has a (different) corresponding projector
    input_slice        = input[:,:,j,1]
    pen_value_slice    = 0f0 #
    P_input            = P[j](input_slice) #project onto constraint set (can be an intersection)
    pen_value_slice    = 0.5f0*norm(vec(P_input)-vec(input_slice))^2 #squared point-to-set distance functino
    pen_grad[:,:,j,1] .= input_slice .- P_input # gradient of point-to-set distance function

    pen_value = pen_value + pen_value_slice #accumulate penalty value (loss) over channels
  end #end channels loop

  #at this point,the softmax of the network output has been projected onto the intersection. Now we need to 'complete' the gradient
  #via the last part of the chain-rule for distance-function-squared and softmax
  pen_grad[:,:,TrOpts.active_channels,1] .= ∇softmax(pen_grad[:,:,TrOpts.active_channels,1], input[:,:,TrOpts.active_channels,1]; dims=3)

  return pen_value, pen_grad
end

function g_CQ(input,TD_OP,P_sub,alpha_CQ,TrOpts)

  #initialize some stuff
  size_input = size(input)
  pen_grad   = zeros(Float32,size(input))
  pen_value  = 0f0

  #use softmax of final network state for constraining
  #the output for segmentation problems (don't use softmax for non-linear regression)
  input      = softmax(input[:,:,TrOpts.active_channels,:],dims=3)

  for j in TrOpts.active_channels #loop over channels, each channel has a (different) corresponding projector
    input_slice        = input[:,:,j,1]
    pen_value_slice    = 0f0 #
    #P_input            = P[j](input_slice) #project onto constraint set (can be an intersection)
    #pen_value_slice    = 0.5f0*norm(vec(P_input)-vec(input_slice))^2 #squared point-to-set distance functino
    f_temp,g_temp      = MultipleSplitFeasCQ_fun_grad(vec(deepcopy(input_slice)),TD_OP[j],P_sub[j],alpha_CQ[j])
    pen_grad[:,:,j,1] .= reshape(g_temp,size(input_slice)) # gradient of point-to-set distance function

    pen_value = pen_value + f_temp #accumulate penalty value (loss) over channels
  end #end channels loop

  #at this point,the softmax of the network output has been projected onto the intersection. Now we need to 'complete' the gradient
  #via the last part of the chain-rule for distance-function-squared and softmax
  pen_grad[:,:,TrOpts.active_channels,1] .= ∇softmax(pen_grad[:,:,TrOpts.active_channels,1], input[:,:,TrOpts.active_channels,1]; dims=3)

  return pen_value, pen_grad
end

function MultipleSplitFeasCQ_fun_grad(x::AbstractVector{T},A,P_sub,alpha::AbstractVector) where T

# returns the gradient of the sum of squared distance functions for the split-feasibility
#problem: A_i x = y_i , y_i \in C_i
# f = \sum_{i=1}^p \alpha_i/2 \| P_C_i (A_i x) - A_i x \|_2^2
# g = \sum_{i=1}^p \alpha_i/2 A_i^t (A_i x - P_C_i (A_i x) )

#serial implementation
#(each projection or mat-vec can still use multiple threads)

f    = zeros(T,length(alpha))#allocate function value per constraint/linear operator
ng   = zeros(T,length(alpha))#allocate norm of gradient per constraint/linear operator
grad = zeros(T,length(x))

for i = 1:length(A)
  y     = A[i]*x
  PcAx  = P_sub[i](deepcopy(y))
  f[i]  = norm(PcAx - y,2).^2
  grad .= grad .+ alpha[i].*A[i]'*(y.-PcAx)
end

return sum(f),grad
end

function LossTotal(HN,TrOpts,X0::AbstractArray{T, N},label,P,image_weights,active_z_slice::Int) where {T, N}
  #loss for hyperspectral imaging with 5d input
  alpha=TrOpts.alpha; use_gpu=TrOpts.use_gpu; lossf=TrOpts.lossf; lossg=TrOpts.lossg; flip_dims=TrOpts.flip_dims; permute_dims=TrOpts.permute_dims; active_channels=TrOpts.active_channels;

    if (isempty(flip_dims) && isempty(permute_dims)) == false
      X0, label, image_weights = AugmentDataLabel(X0, label, image_weights,TrOpts)#optional: augment data
    end

    Y_curr, Y_new, lgdet = HN.forward(X0,X0)
    if use_gpu == true
      Y_new = Y_new|>cpu
    end

    if isempty(label)==false
        lval         = lossf(Y_new[:,:,active_z_slice,active_channels,1],label[:,:,active_z_slice,active_channels,1],image_weights)
        (grad,dummy) = lossg(Y_new[:,:,active_z_slice,active_channels,1],label[:,:,active_z_slice,active_channels,1],image_weights)
      else
        lval = 0.0
    end

  if alpha>0
    dc2,dc2_grad = Dist2Set(Y_new[:,:,active_z_slice,:,:],P,TrOpts)

    #grad  = grad + alpha*dc2_grad[:,:,active_channels,1]
  else
    dc2 = 0f0
  end

  if isempty(label)==false && alpha>0
      grad  = grad + alpha*dc2_grad[:,:,active_z_slice,active_channels,1]
  elseif isempty(label)==true
      n = size(dc2_grad)
      #dc2_grad = reshape(dc2_grad,n[1],n[2],1,n[3],1)
      grad  = alpha*dc2_grad[:,:,active_channels,1]
  end

   ΔY_curr= zeros(Float32,size(Y_new))
   ΔY_curr[:,:,active_z_slice,active_channels,1] .= grad
   ΔY_new = zeros(Float32,size(Y_new))

   if use_gpu == true
     Y_new   = Y_new|>gpu
     ΔY_curr = ΔY_curr|>gpu
     ΔY_new  = ΔY_new|>gpu
   end
   ΔY_curr, ΔY_new, Y_curr, Y_new = HN.backward(ΔY_curr, ΔY_new, Y_curr, Y_new)

   return lval, dc2
end

function LossTotal(HN,TrOpts,X0::AbstractArray{T, N},label,P,image_weights,active_z_slice::Array{Any,1},P_mode::String,TD_OP,P_sub,alpha_CQ) where {T, N}
    alpha=TrOpts.alpha; use_gpu=TrOpts.use_gpu; lossf=TrOpts.lossf; lossg=TrOpts.lossg; active_channels=TrOpts.active_channels;

    Y_curr, Y_new, lgdet = HN.forward(X0,X0)
    if use_gpu == true
      Y_new = Y_new|>cpu
    end

    #initialize gradient
    if N==4
      grad = zeros(Float32,size(Y_new[:,:,active_channels,1]))
    elseif N==5
      grad = zeros(Float32,size(Y_new[:,:,:,active_channels,1]))
    end

    if isempty(label)==false
      if N==4
        lval         = TrOpts.lossf(Y_new[:,:,active_channels,1],label,image_weights)
        (grad,dummy) = TrOpts.lossg(Y_new[:,:,active_channels,1],label,image_weights)
      elseif N==5
        lval         = TrOpts.lossf(Y_new[:,:,:,active_channels,1],label,image_weights)
        (grad,dummy) = TrOpts.lossg(Y_new[:,:,:,active_channels,1],label,image_weights)
      end
    else
      lval = 0.0
    end

  if (alpha[1][1]>0 || alpha_CQ[1][1]>0)==true
    if P_mode == "Proj_intersection"
      dc2,dc2_grad = Dist2Set(Y_new,P,active_channels)
      dc2_grad .= alpha.*dc2_grad
    elseif P_mode == "g_CQ"
      dc2,dc2_grad = g_CQ(Y_new,TD_OP,P_sub,alpha_CQ,TrOpts)
    end

    if N==4
      grad  = grad + dc2_grad[:,:,active_channels,1]
    elseif N==5
      grad  = grad + dc2_grad[:,:,:,active_channels,1]
    end
  else
    dc2 = 0f0
  end

  #optional: randomly set parts of the gradient to zero
  if TrOpts.rand_grad_perc_zero!=0.0
    rand_mask = CartesianIndices(size(grad)[1:end-2])
    rand_mask = shuffle(vec(rand_mask))[1:round(Int,TrOpts.rand_grad_perc_zero*prod(size(rand_mask)))]
    if N==4;grad[rand_mask,:,:].=0f0;elseif N==5;grad[rand_mask,:,:,:].=0f0;end;
  end

   ΔY_curr= zeros(Float32,size(Y_new))
   if N==4
     ΔY_curr[:,:,active_channels,1] .= grad
   elseif N==5
     ΔY_curr[:,:,:,active_channels,1] .= grad
   end
   ΔY_new = zeros(Float32,size(Y_new))

   if use_gpu == true
     Y_new   = Y_new|>gpu
     ΔY_curr = ΔY_curr|>gpu
     ΔY_new  = ΔY_new|>gpu
   end
   ΔY_curr, ΔY_new, Y_curr, Y_new = HN.backward(ΔY_curr, ΔY_new, Y_curr, Y_new)

   return lval, dc2
end

function IoU(HN,data,labels)

if length(size(data[1]))==5
 return 1.0, 1.0
end

threshold = 0.65
IoU_pos = zeros(length(data))
IoU_neg = zeros(length(data))

  for i=1:length(data)
    ~, prediction, ~ = HN.forward(data[i],data[i])
    prediction = prediction |> cpu
    prediction[:,:,1:2,1].=softmax(prediction[:,:,1:2,1],dims=3);

    pred_thres = zeros(Int,size(prediction)[1:2])
    pos_inds   = findall(prediction[:,:,1] .> threshold)
    pred_thres[pos_inds] .= 1

    prediction_for_acc = pred_thres[5:end-4,5:end-4]
    mask_for_acc       = labels[i][5:end-4,5:end-4,1]

    pos_pred_inds = findall(prediction_for_acc.==1)
    neg_pred_inds = findall(prediction_for_acc.==0)

    true_pos_inds = findall(mask_for_acc .== 1)
    true_neg_inds = findall(mask_for_acc .== 0)

    #IoU
    IoU_pos[i] = length(intersect(pos_pred_inds,true_pos_inds))/length(union(pos_pred_inds,true_pos_inds))
    IoU_neg[i] = length(intersect(neg_pred_inds,true_neg_inds))/length(union(neg_pred_inds,true_neg_inds))
  end

return IoU_pos, IoU_neg
end

#Loss for applications other than Hyperspectral. Uses projection onto intersection.
function LossTotal(HN,TrOpts,X0::AbstractArray{T, N},label,P,image_weights,active_z_slice::Array{Any,1}) where {T, N}
    alpha=TrOpts.alpha; use_gpu=TrOpts.use_gpu; lossf=TrOpts.lossf; lossg=TrOpts.lossg; active_channels=TrOpts.active_channels;

    Y_curr, Y_new, lgdet = HN.forward(X0,X0)
    if use_gpu == true
      Y_new = Y_new|>cpu
    end

    #initialize gradient
    if N==4
      grad = zeros(Float32,size(Y_new[:,:,active_channels,1]))
    elseif N==5
      grad = zeros(Float32,size(Y_new[:,:,:,active_channels,1]))
    end

    if isempty(label)==false
      if N==4
        lval         = lossf(Y_new[:,:,active_channels,1],label,image_weights)
        (grad,dummy) = lossg(Y_new[:,:,active_channels,1],label,image_weights)
      elseif N==5
        lval         = lossf(Y_new[:,:,:,active_channels,1],label,image_weights)
        (grad,dummy) = lossg(Y_new[:,:,:,active_channels,1],label,image_weights)
      end
    else
      lval = 0.0
    end
    # # #take equal number of random pixels from each class for the gradient
    # pos_inds = findall(label[:,:,1,1].==1)
    # neg_inds = findall(label[:,:,1,1].==0)
    # npix = min(length(pos_inds),length(neg_inds))
    # npix = round(Int,npix/2)
    # pos_inds_select = shuffle(pos_inds)[1:npix]
    # neg_inds_select = shuffle(neg_inds)[1:npix]
    # random_mask = zeros(Float32,size(grad))
    # random_mask[pos_inds_select,1:2,1].=1
    # random_mask[neg_inds_select,1:2,1].=1

    #grad .= grad.*random_mask

  if alpha>0
    #if P_mode == "Proj_intersection"
      dc2,dc2_grad = Dist2Set(Y_new,P,TrOpts)
    #elseif P_mode == "g_CQ"
    #  dc2,dc2_grad = g_CQ(Y_new,TD_OP,P_sub,alpha_CQ,active_channels)
    #end
    # if (norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) > 10f0
    #   @warn "(norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) > 10f0"
    # elseif (norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) < 0.1f0
    #   @warn "(norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) < 0.1f0"
    # end
    if N==4
      grad  = grad + alpha*dc2_grad[:,:,active_channels,1]
    elseif N==5
      grad  = grad + alpha*dc2_grad[:,:,:,active_channels,1]
    end
  else
    dc2 = 0f0
  end

  #optional: randomly set parts of the gradient to zero
  if TrOpts.rand_grad_perc_zero!=0.0
    rand_mask = CartesianIndices(size(grad)[1:end-2])
    rand_mask = shuffle(vec(rand_mask))[1:round(Int,TrOpts.rand_grad_perc_zero*prod(size(rand_mask)))]
    if N==4;grad[rand_mask,:,:].=0f0;elseif N==5;grad[rand_mask,:,:,:].=0f0;end;
  end

   ΔY_curr= zeros(Float32,size(Y_new))
   if N==4
     ΔY_curr[:,:,active_channels,1] .= grad
   elseif N==5
     ΔY_curr[:,:,:,active_channels,1] .= grad
   end
   ΔY_new = zeros(Float32,size(Y_new))

   if use_gpu == true
     Y_new   = Y_new|>gpu
     ΔY_curr = ΔY_curr|>gpu
     ΔY_new  = ΔY_new|>gpu
   end
   ΔY_curr, ΔY_new, Y_curr, Y_new = HN.backward(ΔY_curr, ΔY_new, Y_curr, Y_new)

   return lval, dc2
end


using JuMP
using GLPK
using Plots  #; plotlyjs()
using Measures
using DataFrames
using NPZ
using CSV
using Dates
using DataFrames
using Statistics
using CartesianGeneticProgramming
using Cambrian
import Cambrian.mutate
using Plots, Measures
using StatsPlots
include("../utils.jl")
using Dierckx
using BSplineKit
using Distances

struct Month
    topo0
    topo1

    tp
    hs
    direction

    shoreline_change
    shoreline_change_class  # (landward, seaward)
end

function fit_poly_inf_norm(x,y,n)
	m = Model(GLPK.Optimizer)

    @variable(m, alpha[1:n+1])
	@variable(m, t >= 0.)
	@constraint(m, [i=1:length(y)], sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i] <= t);
	@constraint(m, [i=1:length(y)], -t <= sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i]);
	@objective(m, Min, t)

	optimize!(m)

    return JuMP.value.(alpha)
end

f2(x, args) = (args[2] * x) .+ args[1]

dataset_path = get_dataset_path(;dataset_name="series_dataset_path")
n_series = 6
series = []
count_months = 0
xs_csv_path = get_dataset_path(;dataset_name="xs_csv_path")
xs_csv = CSV.read(xs_csv_path, DataFrame)
xs = xs_csv.x

changes = []

for s in 1:n_series
    push!(series, [])
    series_dir = joinpath(dataset_path, string(s))

    files = readdir(series_dir)
    topo_files = filter(f-> !occursin("_waves.csv", f), files)
    topo_files = map(f->replace(f, "_topo.npy" => ""), topo_files)

    month_numbers = sort(map(f->parse(Int64, f), topo_files))
    n_months = length(month_numbers)
    global count_months += n_months

    println(s, " - ", n_months)

    for mi in 1:n_months
        i = month_numbers[mi]
        topo_path = string(dataset_path, '/', s, '/', i, "_topo.npy")
        if i == 312
            println("TEST")
            exit()
            continue
        end
        topo = npzread(topo_path)

        topo0 = topo[1, :]
        topo1 = topo[2, :]

        t0_index_last_positive = findlast(topo0.>0)
        t0_index_first_negative = findfirst(topo0.<0)

        t0_x_last_positive = xs[t0_index_last_positive]
        t0_x_first_negative = xs[t0_index_first_negative]
        t0_value_last_positive = topo0[t0_index_last_positive]
        t0_value_first_negative = topo0[t0_index_first_negative]

        t0_coeffs = fit_poly_inf_norm([t0_value_last_positive, t0_value_first_negative], [t0_x_last_positive, t0_x_first_negative], 1)
        t0_x_at_zero = f2([0], t0_coeffs)[1]

        t1_index_last_positive = findlast(topo1.>0)
        t1_index_first_negative = findfirst(topo1.<0)

        t1_x_last_positive = xs[t1_index_last_positive]
        t1_x_first_negative = xs[t1_index_first_negative]
        t1_value_last_positive = topo1[t1_index_last_positive]
        t1_value_first_negative = topo1[t1_index_first_negative]

        t1_coeffs = fit_poly_inf_norm([t1_value_last_positive, t1_value_first_negative], [t1_x_last_positive, t1_x_first_negative], 1)
        t1_x_at_zero = f2([0], t1_coeffs)[1]

        shoreline_change = t1_x_at_zero - t0_x_at_zero
        shoreline_change = norm_to_range(shoreline_change, -15.0, 15.0, -1.0, 1.0)
        if shoreline_change < 0
            class = (0, 1)  # (landward, seaward)
        else
            class = (1, 0)  # (landward, seaward)
        end

        waves_path = string(dataset_path, '/', s, '/', i, "_waves.csv")
        waves_csv = CSV.read(waves_path, DataFrame)

        tp_month = norm_to_range(waves_csv.tp, 0.0, 25.0, -1.0, 1.0)
        hs_month = norm_to_range(waves_csv.hs, 0.0, 10.0, -1.0, 1.0)
        direction_month = norm_to_range(waves_csv.direction, 0.0, 360.0, -1.0, 1.0)

        month = Month(topo0, topo1, tp_month, hs_month, direction_month, shoreline_change, class)
        push!(series[end], month)
    end
end

real_class = []
real_change = []
tp_min = []
tp_max = []
tp_avg = []
hs_min = []
hs_max = []
hs_avg = []
dir_min = []
dir_max = []
dir_avg = []

for si in 1:length(series)
    global correct
    global wrong
    months = series[si]
    for mi in 1:length(months)
        month = months[mi]
        push!(real_class, month.shoreline_change_class)
        push!(real_change, norm_to_range(month.shoreline_change, -1.0, 1.0, -15.0, 15.0))
        push!(tp_min, minimum(month.tp))
        push!(tp_max, maximum(month.tp))
        push!(tp_avg, mean(month.tp))
        push!(hs_min, minimum(month.hs))
        push!(hs_max, maximum(month.hs))
        push!(hs_avg, mean(month.hs))
        push!(dir_min, minimum(month.direction))
        push!(dir_max, maximum(month.direction))
        push!(dir_avg, mean(month.direction))
    end
end

println("Max change: ", maximum(real_change))
println("Min change: ", minimum(real_change))

# p_tp_min = scatter(tp_min, preds, xlabel="min(Tp)", ylabel="Pred (1=True, 0=False)")
# p_tp_max = scatter(tp_max, preds, xlabel="max(Tp)", ylabel="Pred (1=True, 0=False)")
# p_tp_avg = scatter(tp_avg, preds, xlabel="mean(Tp)", ylabel="Pred (1=True, 0=False)")
# tp_plot = plot(p_tp_min, p_tp_max, p_tp_avg, layout=(1, 3))

# p_hs_min = scatter(hs_min, preds, xlabel="min(Hs)", ylabel="Pred (1=True, 0=False)")
# p_hs_max = scatter(hs_max, preds, xlabel="max(Hs)", ylabel="Pred (1=True, 0=False)")
# p_hs_avg = scatter(hs_avg, preds, xlabel="mean(Hs)", ylabel="Pred (1=True, 0=False)")
# hs_plot = plot(p_hs_min, p_hs_max, p_hs_avg, layout=(1, 3))

# p_dir_min = scatter(dir_min, preds, xlabel="min(Dir)", ylabel="Pred (1=True, 0=False)")
# p_dir_max = scatter(dir_max, preds, xlabel="max(Dir)", ylabel="Pred (1=True, 0=False)")
# p_dir_avg = scatter(dir_avg, preds, xlabel="mean(Dir)", ylabel="Pred (1=True, 0=False)")
# dir_plot = plot(p_dir_min, p_dir_max, p_dir_avg, layout=(1, 3))

# display(plot(tp_plot, hs_plot, dir_plot, layout=(3, 1), size = (1600, 800)))

# display(plot(real_change, legend=false, ylabel="Shoreline change [m]", xlabel="Months", size=(1280, 300)))
# readline()
# exit()

################## dx correlations ##################
p_tp_min = scatter(tp_min, real_change, xlabel="min(Tp)", ylabel="Real change", title=string("Correlation: ", round(cor(tp_min, real_change), digits=2)))
p_tp_max = scatter(tp_max, real_change, xlabel="max(Tp)", ylabel="Real change", title=string("Correlation: ", round(cor(tp_max, real_change), digits=2)))
p_tp_avg = scatter(tp_avg, real_change, xlabel="mean(Tp)", ylabel="Real change", title=string("Correlation: ", round(cor(tp_avg, real_change), digits=2)))
tp_plot = plot(p_tp_min, p_tp_max, p_tp_avg, layout=(1, 3), bottom_margin=20mm)

p_hs_min = scatter(hs_min, real_change, xlabel="min(Hs)", ylabel="Real change", title=string("Correlation: ", round(cor(hs_min, real_change), digits=2)))
p_hs_max = scatter(hs_max, real_change, xlabel="max(Hs)", ylabel="Real change", title=string("Correlation: ", round(cor(hs_max, real_change), digits=2)))
p_hs_avg = scatter(hs_avg, real_change, xlabel="mean(Hs)", ylabel="Real change", title=string("Correlation: ", round(cor(hs_avg, real_change), digits=2)))
hs_plot = plot(p_hs_min, p_hs_max, p_hs_avg, layout=(1, 3), bottom_margin=20mm)

p_dir_min = scatter(dir_min, real_change, xlabel="min(dir)", ylabel="Real change", title=string("Correlation: ", round(cor(dir_min, real_change), digits=2)))
p_dir_max = scatter(dir_max, real_change, xlabel="max(dir)", ylabel="Real change", title=string("Correlation: ", round(cor(dir_max, real_change), digits=2)))
p_dir_avg = scatter(dir_avg, real_change, xlabel="mean(dir)", ylabel="Real change", title=string("Correlation: ", round(cor(dir_avg, real_change), digits=2)))
dir_plot = plot(p_dir_min, p_dir_max, p_dir_avg, layout=(1, 3))

dx_correlations_plot = plot(tp_plot, hs_plot, dir_plot, layout=(3, 1), size = (1600, 800), margin=10mm)

################## tp correlations ##################
p_hs_min = scatter(hs_min, tp_min, xlabel="min(Hs)", ylabel="min(Tp)", title=string("Correlation: ", round(cor(hs_min, tp_min), digits=2)))
p_hs_max = scatter(hs_max, tp_max, xlabel="max(Hs)", ylabel="max(Tp)", title=string("Correlation: ", round(cor(hs_max, tp_max), digits=2)))
p_hs_avg = scatter(hs_avg, tp_avg, xlabel="mean(Hs)", ylabel="mean(Tp)", title=string("Correlation: ", round(cor(hs_avg, tp_avg), digits=2)))
hs_plot = plot(p_hs_min, p_hs_max, p_hs_avg, layout=(1, 3), bottom_margin=20mm)

p_dir_min = scatter(dir_min, tp_min, xlabel="min(dir)", ylabel="min(Tp)", title=string("Correlation: ", round(cor(dir_min, tp_min), digits=2)))
p_dir_max = scatter(dir_max, tp_max, xlabel="max(dir)", ylabel="max(Tp)", title=string("Correlation: ", round(cor(dir_max, tp_max), digits=2)))
p_dir_avg = scatter(dir_avg, tp_avg, xlabel="mean(dir)", ylabel="mean(Tp)", title=string("Correlation: ", round(cor(dir_avg, tp_avg), digits=2)))
dir_plot = plot(p_dir_min, p_dir_max, p_dir_avg, layout=(1, 3))

tp_correlations_plot = plot(hs_plot, dir_plot, layout=(2, 1), size = (1600, 600), margin=10mm)
display(plot(tp_correlations_plot))
readline()


module InterpolationWrapper

using JavaCall

const AkimaSplineInterpolator = @jimport org.hipparchus.analysis.interpolation.AkimaSplineInterpolator
const BicubicFunction = @jimport org.hipparchus.analysis.interpolation.BicubicFunction
const BicubicInterpolatingFunction = @jimport org.hipparchus.analysis.interpolation.BicubicInterpolatingFunction
const BicubicInterpolator = @jimport org.hipparchus.analysis.interpolation.BicubicInterpolator
const BilinearInterpolatingFunction = @jimport org.hipparchus.analysis.interpolation.BilinearInterpolatingFunction
const BilinearInterpolator = @jimport org.hipparchus.analysis.interpolation.BilinearInterpolator
const BivariateFunction = @jimport org.hipparchus.analysis.BivariateFunction
const BivariateGridInterpolator = @jimport org.hipparchus.analysis.interpolation.BivariateGridInterpolator
const Class = @jimport java.lang.Class
const Derivative = @jimport org.hipparchus.analysis.differentiation.Derivative
const DerivativeStructure = @jimport org.hipparchus.analysis.differentiation.DerivativeStructure
const DividedDifferenceInterpolator = @jimport org.hipparchus.analysis.interpolation.DividedDifferenceInterpolator
const Field = @jimport org.hipparchus.Field
const FieldBivariateFunction = @jimport org.hipparchus.analysis.FieldBivariateFunction
const FieldElement = @jimport org.hipparchus.FieldElement
const FieldHermiteInterpolator = @jimport org.hipparchus.analysis.interpolation.FieldHermiteInterpolator
const FieldPolynomialSplineFunction = @jimport org.hipparchus.analysis.polynomials.FieldPolynomialSplineFunction
const FieldUnivariateInterpolator = @jimport org.hipparchus.analysis.interpolation.FieldUnivariateInterpolator
const GridAxis = @jimport org.hipparchus.analysis.interpolation.GridAxis
const HermiteInterpolator = @jimport org.hipparchus.analysis.interpolation.HermiteInterpolator
const InterpolatingMicrosphere = @jimport org.hipparchus.analysis.interpolation.InterpolatingMicrosphere
const InterpolatingMicrosphere2D = @jimport org.hipparchus.analysis.interpolation.InterpolatingMicrosphere2D
const JString = @jimport java.lang.String
const LinearInterpolator = @jimport org.hipparchus.analysis.interpolation.LinearInterpolator
const LoessInterpolator = @jimport org.hipparchus.analysis.interpolation.LoessInterpolator
const MicrosphereProjectionInterpolator = @jimport org.hipparchus.analysis.interpolation.MicrosphereProjectionInterpolator
const MultivariateFunction = @jimport org.hipparchus.analysis.MultivariateFunction
const MultivariateInterpolator = @jimport org.hipparchus.analysis.interpolation.MultivariateInterpolator
const NevilleInterpolator = @jimport org.hipparchus.analysis.interpolation.NevilleInterpolator
const Object = @jimport java.lang.Object
const PiecewiseBicubicSplineInterpolatingFunction = @jimport org.hipparchus.analysis.interpolation.PiecewiseBicubicSplineInterpolatingFunction
const PiecewiseBicubicSplineInterpolator = @jimport org.hipparchus.analysis.interpolation.PiecewiseBicubicSplineInterpolator
const PolynomialFunction = @jimport org.hipparchus.analysis.polynomials.PolynomialFunction
const PolynomialFunctionLagrangeForm = @jimport org.hipparchus.analysis.polynomials.PolynomialFunctionLagrangeForm
const PolynomialFunctionNewtonForm = @jimport org.hipparchus.analysis.polynomials.PolynomialFunctionNewtonForm
const PolynomialSplineFunction = @jimport org.hipparchus.analysis.polynomials.PolynomialSplineFunction
const RealFieldBivariateFunction = @jimport org.hipparchus.analysis.RealFieldBivariateFunction
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const RealFieldUnivariateFunction = @jimport org.hipparchus.analysis.RealFieldUnivariateFunction
const SplineInterpolator = @jimport org.hipparchus.analysis.interpolation.SplineInterpolator
const TricubicFunction = @jimport org.hipparchus.analysis.interpolation.TricubicFunction
const TricubicInterpolatingFunction = @jimport org.hipparchus.analysis.interpolation.TricubicInterpolatingFunction
const TricubicInterpolator = @jimport org.hipparchus.analysis.interpolation.TricubicInterpolator
const TrivariateFunction = @jimport org.hipparchus.analysis.TrivariateFunction
const TrivariateGridInterpolator = @jimport org.hipparchus.analysis.interpolation.TrivariateGridInterpolator
const UnitSphereRandomVectorGenerator = @jimport org.hipparchus.random.UnitSphereRandomVectorGenerator
const UnivariateFunction = @jimport org.hipparchus.analysis.UnivariateFunction
const UnivariateInterpolator = @jimport org.hipparchus.analysis.interpolation.UnivariateInterpolator
const UnivariatePeriodicInterpolator = @jimport org.hipparchus.analysis.interpolation.UnivariatePeriodicInterpolator

export AkimaSplineInterpolator
export BicubicFunction
export BicubicInterpolatingFunction
export BicubicInterpolator
export BilinearInterpolatingFunction
export BilinearInterpolator
export BivariateGridInterpolator
export DividedDifferenceInterpolator
export FieldHermiteInterpolator
export FieldUnivariateInterpolator
export GridAxis
export HermiteInterpolator
export InterpolatingMicrosphere
export InterpolatingMicrosphere2D
export LinearInterpolator
export LoessInterpolator
export MicrosphereProjectionInterpolator
export MultivariateInterpolator
export NevilleInterpolator
export PiecewiseBicubicSplineInterpolatingFunction
export PiecewiseBicubicSplineInterpolator
export SplineInterpolator
export TricubicFunction
export TricubicInterpolatingFunction
export TricubicInterpolator
export TrivariateGridInterpolator
export UnivariateInterpolator
export UnivariatePeriodicInterpolator
export add_sample_point
export copy
export derivatives
export equals
export get_class
export get_dimension
export get_n
export get_polynomials
export get_size
export get_x_inf
export get_x_sup
export get_y_inf
export get_y_sup
export hash_code
export interpolate
export interpolation_index
export is_valid_point
export node
export notify
export notify_all
export size
export smooth
export to_real_field_bivariate_function
export to_string
export value
export wait

include("akima_spline_interpolator.jl")
include("bicubic_function.jl")
include("bicubic_interpolating_function.jl")
include("bicubic_interpolator.jl")
include("bilinear_interpolating_function.jl")
include("bilinear_interpolator.jl")
include("bivariate_grid_interpolator.jl")
include("divided_difference_interpolator.jl")
include("field_hermite_interpolator.jl")
include("field_univariate_interpolator.jl")
include("grid_axis.jl")
include("hermite_interpolator.jl")
include("interpolating_microsphere.jl")
include("interpolating_microsphere2_d.jl")
include("linear_interpolator.jl")
include("loess_interpolator.jl")
include("microsphere_projection_interpolator.jl")
include("multivariate_interpolator.jl")
include("neville_interpolator.jl")
include("piecewise_bicubic_spline_interpolating_function.jl")
include("piecewise_bicubic_spline_interpolator.jl")
include("spline_interpolator.jl")
include("tricubic_function.jl")
include("tricubic_interpolating_function.jl")
include("tricubic_interpolator.jl")
include("trivariate_grid_interpolator.jl")
include("univariate_interpolator.jl")
include("univariate_periodic_interpolator.jl")

end


function MP.differentiate(m::Monomial{C}, x::PolyVar{C}) where C
    i = findfirst(isequal(x), _vars(m))
    if (i == nothing || i == 0) || m.z[i] == 0
        zeroterm(m)
    else
        z = copy(m.z)
        z[i] -= 1
        m.z[i] * Monomial(_vars(m), z)
    end
end

function MP.differentiate(p::Polynomial{C, T}, x::PolyVar{C}) where {C, T}
    # grlex order preserved
    i = something(findfirst(isequal(x), _vars(p)), 0)
    S = typeof(zero(T) * 0)
    if iszero(i)
        zero(Polynomial{C, S})
    else
        keep = findall(z -> z[i] > 0, p.x.Z)
        Z = copy.(p.x.Z[keep])
        a = Vector{S}(undef, length(keep))
        for j in 1:length(Z)
            a[j] = p.a[keep[j]] * Z[j][i]
            Z[j][i] -= 1
        end
        Polynomial(a, MonomialVector(_vars(p), Z))
    end
end


type BlackScholesLattice{T <: AbstractBinomialTree} <: TreeLattice
  tree::T
  riskFreeRate::Float64
  dt::Float64
  discountFactor::Float64
  pd::Float64
  pu::Float64
  treeLattice::TreeLattice1D

  function BlackScholesLattice(tree::T, riskFreeRate::Float64, endTime::Float64, steps::Int)
    dt = endTime / steps
    discountFactor = exp(-riskFreeRate * dt)
    pd = probability(tree, 0, 0, 0)
    pu = probability(tree, 0, 0, 1)
    bsm_lattice = new{T}(tree, riskFreeRate, dt, discountFactor, pd, pu)

    tg = TimeGrid(endTime, steps)
    lat = TreeLattice1D(tg, 2, bsm_lattice)
    bsm_lattice.treeLattice = lat

    return bsm_lattice
  end
end

BlackScholesLattice{T <: AbstractBinomialTree}(tree::T, riskFreeRate::Float64, endTime::Float64, steps::Int) = BlackScholesLattice{T}(tree, riskFreeRate, endTime, steps)

get_underlying(bsm::BlackScholesLattice, i::Int, idx::Int) = get_underlying(bsm.tree, i, idx)

descendant(bsm::BlackScholesLattice, i::Int, idx::Int, branch::Int) = descendant(bsm.tree, i, idx, branch)

probability(bsm::BlackScholesLattice, i::Int, idx::Int, branch::Int) = probability(bsm.tree, i, idx, branch)

get_size(bsm::BlackScholesLattice, i::Int) = get_size(bsm.tree, i)

discount(bsm::BlackScholesLattice, ::Int, ::Int) = bsm.discountFactor

function step_back!(bsm::BlackScholesLattice, i::Int, values::Vector{Float64}, newValues::Vector{Float64})
  for j = 1:i
    newValues[j] = (bsm.pd * values[j] + psm.pu * values[j+1]) * bsm.discountFactor
  end

  return newValues
end


# This file was generated, do not modify it.

using MLJ

@load LinearRegressor pkg=MLJLinearModels

using RDatasets, DataFrames
boston = dataset("MASS", "Boston")
first(boston, 3)

describe(boston, :mean, :std, :eltype)

using ScientificTypes
data = coerce(boston, :Tax=>Continuous, :Rad=>Continuous);

y = data.MedV
X = select(data, Not(:MedV));

mdl = LinearRegressor()

mach = machine(mdl, X, y)
fit!(mach)

fp = fitted_params(mach)
@show round.(fp.coefs[1:3], sigdigits=3)
@show round(fp.intercept, sigdigits=3)

ŷ = predict(mach, X)
round(rms(ŷ, y), sigdigits=4)

X2 = hcat(X, X.LStat .* X.Age);

rename!(X2, :x1 => :interaction);

mach = machine(mdl, X2, y)
fit!(mach)
ŷ = predict(mach, X2)
round(rms(ŷ, y), sigdigits=4)

X3 = hcat(X.LStat, X.LStat.^2)
machine(mdl, X3, y)
fit!(mach)
ŷ = predict(mach, X3)
round(rms(ŷ, y), sigdigits=4)


using Distributions
using Base.Test

@testset "non-central T-Distributions" begin
    ν = 50
    λ = -5

    d = NoncentralT(ν, λ)
    @test_throws MethodError NoncentralT(ν, complex(λ))
    @test d == typeof(d)(params(d)...)
    @test mean(d) ≈ -5.0766 atol = 1e-5

    λ = 5
    d = NoncentralT(ν, λ)
    @test mean(d) ≈ 5.0766 atol=1e-5

    λ = 0
    d = NoncentralT(ν, λ)
    @test var(d) == var(TDist(ν))
end


using Unitful
using DataStructures

t = range(0.0, stop=10.0, length=100)

result = OrderedDict{String,Any}()

result["time"] = t*u"s"
result["r"]    = [[0.4 * cos(t[i]),
                   0.5 * sin(t[i]),
                   0.3 * cos(t[i])] for i in eachindex(t)]*u"m"

plot(result, ["r", "r[2]", "r[2:3]"], heading="Plot vector signals")

println("\n... result info:")
printResultInfo(result)


using ITensors

import ITensors: op

function op(::OpName"expτSS", ::SiteType"S=1/2", s1::Index, s2::Index; τ)
  h =
    1 / 2 * op("S+", s1) * op("S-", s2) +
    1 / 2 * op("S-", s1) * op("S+", s2) +
    op("Sz", s1) * op("Sz", s2)
  return exp(τ * h)
end

function main(; N=10, cutoff=1E-8, δt=0.1, ttotal=5.0)
  # Compute the number of steps to do
  Nsteps = Int(ttotal / δt)

  # Make an array of 'site' indices
  s = siteinds("S=1/2", N; conserve_qns=true)

  # Make gates (1,2),(2,3),(3,4),...
  gates = ops([("expτSS", (n, n + 1), (τ=-δt * im / 2,)) for n in 1:(N - 1)], s)

  # Include gates in reverse order too
  # (N,N-1),(N-1,N-2),...
  append!(gates, reverse(gates))

  # Function that measures <Sz> on site n
  function measure_Sz(psi::MPS, n)
    psi = orthogonalize(psi, n)
    sn = siteind(psi, n)
    Sz = scalar(dag(prime(psi[n], "Site")) * op("Sz", sn) * psi[n])
    return real(Sz)
  end

  # Initialize psi to be a product state (alternating up and down)
  psi0 = productMPS(s, n -> isodd(n) ? "Up" : "Dn")

  c = div(N, 2)

  # Compute and print initial <Sz> value
  t = 0.0
  Sz = measure_Sz(psi0, c)
  println("$t $Sz")

  # Do the time evolution by applying the gates
  # for Nsteps steps
  psi = psi0
  for step in 1:Nsteps
    psi = apply(gates, psi; cutoff=cutoff)
    t += δt
    Sz = measure_Sz(psi, c)
    println("$t $Sz")
  end

  # Now do the same evolution with an MPO
  rho0 = MPO(psi0)
  rho = rho0
  for step in 1:Nsteps
    rho = apply(gates, rho; cutoff=cutoff, apply_dag=true)
    t += δt
  end
  @show inner(psi, rho, psi)
  @show inner(psi, psi)
  @show tr(rho)

  return nothing
end


"""
    erig(b::Basis, i::Int, j::Int; alpha=1.0, scale=1.0)
    erig(b::Basis, i::Int, j::Int, n::Int, m::Int; alpha=1.0, scale=1.0)

Electron repulsion integral with Gaussian repulsion scale*exp(-α*r^2).

For erig4 `i` and `j` are indices for particale 1 and `n` and `m` for particale 2.
"""
function erig(b::AbstractBasis, i::Int, j::Int; Ve=x->exp(-x^2))
    return Ve(b[i]-b[j])
end

function erig(b::AbstractBasis, i::Int, j::Int, n::Int, m::Int; Ve=x->exp(-x^2))::Float64
    # TODO this function might not be type stable
    T = Float64   # This need to change for AD to work eg. Base.return_type()
    # Basis functions are orthogonal
    i != j && return zero(T)
    n != m && return zero(T)
    return Ve(b[i]-b[n])
end

function fock_matrix(b::AbstractBasis; Ve=x->exp(-x^2), Vn=x->-2exp(-0.5x^2))
    orbitals = initial_orbitals(b)
    return fock_matrix(b, orbitals; Ve=Ve, Vn=Vn)
end

function fock_matrix(b::AbstractBasis, orbitals::AbstractMatrix; Ve=x->exp(-x^2), Vn=x->-2exp(-0.5x^2))
    @argcheck size(orbitals) == (length(b), length(b))
    h₁ = one_electron_operator(b, Vn)
    J = coulomb_matrix(b, orbitals)
    #K = exchange_matrix(b, orbitals) # For more than 2 electrons
    return h₁ + J  # h₁ + 2*J - K
end

function fock_matrix!(f::AbstractMatrix, b::AbstractBasis, orbitals::AbstractMatrix; Ve=x->exp(-x^2), Vn=x->-2exp(-0.5x^2))
    @argcheck size(f) == size(orbitals) == (length(b), length(b))
    h₁ = one_electron_operator(b, Vn)
    J = coulomb_matrix(b, orbitals)
    #K = exchange_matrix(b, orbitals) # For more than 2 electrons
    return f .= h₁ + J  # h₁ + 2*J - K
end

function coulomb_matrix(b::AbstractBasis, orbitals::AbstractMatrix)
    C = zeros(size(orbitals))
    return coulomb_matrix!(C, b, orbitals)
end

function coulomb_matrix!(C::AbstractMatrix, b::AbstractBasis, orbitals::AbstractMatrix)
    @argcheck size(C) == size(orbitals) == (length(b), length(b))
    l = length(b)
    w = get_weight(b)

    # Two electrons in total
    ρ = orbitals[:,1] * orbitals[:,1]'
    for i in 1:l
        # C[i,j]=0 for i!=j
        # also erig(b,i,i,n,m) = 0 for n!=m
        C[i,i] = sum( n -> erig(b,i,i,n,n) * ρ[n,n] * w[n], 1:l)
        C[i,i] *= w[i]
    end
    return C
end

function exchange_matrix(b::AbstractBasis, orbitals::AbstractMatrix)
    @argcheck size(orbitals) == (length(b), length(b))
    l = length(b)
    K = zeros(l,l)
    w = get_weight(b)

    # Two electrons in total
    ρ = orbitals[:,1] * orbitals[:,1]'
    Threads.@threads for i in 1:l
        for j in 1:l
            for n in 1:l
                K[i,j] += sum( k-> erig(b, i,n,j,k) * ρ[n,k], 1:l)
            end
            K[i,j] *= w[i] * w[j]
        end
    end
    return K
end

function bracket(
        b::AbstractBasis,
        psi1::AbstractVector,
        op::AbstractMatrix,
        psi2::AbstractVector
    )
    w = get_weight(b)
    return (conj.(psi1).*w)' * op * psi2
end

function bracket(b::AbstractBasis, psi1::AbstractVector, psi2::AbstractVector)
    w = get_weight(b)
    return sum( conj.(psi1) .* w .* psi2 )
end

function kinetic_energy(b::AbstractBasis; mass=1.0)
    ∇ = derivative_matrix(b)
    g = metric_tensor(b)
    return (0.5/mass) * ∇' * g * ∇ # = 0.5 * (g∇)' * ∇
end

function one_electron_operator(b, Vn; mass=1.0)
    g = metric_tensor(b)
    V = g * diagm( Vn.(b) )
    return kinetic_energy(b; mass=mass) + V
end


# Steps to install Python in Atom
# https://python.plainenglish.io/6-simple-steps-for-a-python-setup-in-atom-ca3100711f62
using PyCall
math = pyimport("math")
math.sin(math.pi / 4)

#Format: YYY-MM-DD
run(`python UC_BAU_postprocess.py '2019-01-01' '2019-01-02'`)


using Dates

export days, yearfraction

#export AbstractDayCountConvention

#export DayCount30_360, DayCountAct,  DayCountAct_Act

export DayCount30A_360, DayCount30E_360, DayCount30E_360ISDA
export DayCountAct_360, DayCountAct_364, DayCountAct_365Fixed, DayCountAct_365Actual, DayCountAct_365Long, DayCountAct_365_25
export DayCountAct_ActISDA

abstract type AbstractDayCountConvention end

abstract type DayCount30_360 <: AbstractDayCountConvention end
abstract type DayCountAct <: AbstractDayCountConvention end
abstract type DayCountAct_Act <: AbstractDayCountConvention end

abstract type DayCount30A_360 <: DayCount30_360 end
abstract type DayCount30E_360 <: DayCount30_360 end
abstract type DayCount30E_360ISDA <: DayCount30_360 end

abstract type DayCountAct_360 <: DayCountAct end
abstract type DayCountAct_364 <: DayCountAct end
abstract type DayCountAct_365Fixed <: DayCountAct end
abstract type DayCountAct_365_25 <: DayCountAct end
abstract type DayCountAct_365Actual <: DayCountAct end
abstract type DayCountAct_365Long <: DayCountAct end

abstract type DayCountAct_ActISDA <: DayCountAct_Act end

@inline daysFormula(y1::Int, y2::Int, m1::Int, m2::Int, d1::Int, d2::Int)::Int = 360(y2-y1) + 30(m2-m1) + (d2-d1)
"""
    function days(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCount30A_360

Calculates days between startDate and endDate using 2006 ISDA definitions 4.16f.
also known as 30/360 U.S. Municipal, 30/360 Bond Basis, 30/360 ISDA
"""
days(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCount30A_360 = begin
    d1 = min(day(startDate), 30)
    de = day(endDate)
    d2 = if d1 == 30
        min(de, 30)
    else
        de
    end
    daysFormula(year(startDate), year(endDate), month(startDate), month(endDate), d1, d2)
end
days(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCount30E_360 = begin
    daysFormula(
        year(startDate), year(endDate),
        month(startDate), month(endDate),
        min(day(startDate), 30), min(day(endDate), 30))
end
days(startDate::Date, endDate::Date, terminationDate::Date, ::Type{T}) where T <: DayCount30E_360ISDA = begin
    d1 = if startDate == Dates.lastdayofmonth(startDate)
        30
    else
        day(startDate)
    end

    d2 = if endDate == Dates.lastdayofmonth(endDate) && !(month(endDate) == 2 && endDate == terminationDate)
        30
    else
        day(endDate)
    end
    daysFormula(year(startDate), year(endDate), month(startDate), month(endDate), d1, d2)
end
yearfraction(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCount30_360 =
    days(startDate, endDate, T) / 360.0
yearfraction(startDate::Date, endDate::Date, terminationDate::Date, ::Type{T}) where T <: DayCount30_360 = days(startDate, endDate, terminationDate, T) / 360.0

days(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCountAct =
    (endDate - startDate).value
yearfraction(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCountAct_360 =
    days(startDate, endDate, T) / 360.0
yearfraction(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCountAct_364 =
    days(startDate, endDate, T) / 364.0
yearfraction(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCountAct_365Fixed =
    days(startDate, endDate, T) / 365.0
yearfraction(startDate::Date, endDate::Date, ::Type{T}) where T <: DayCountAct_365Actual = begin
    d = days(startDate, endDate, T)
    for y ∈ year(startDate):year(endDate)
        if Dates.isleapyear(y) && startDate < Date(y,2,29) ≤ endDate
            return d / 366.0
        end
    end
    d / 365.0
end