using Revise
	using DiffEqOperators, ForwardDiff, DifferentialEquations
	using BifurcationKit, LinearAlgebra, Plots, SparseArrays, Parameters, Setfield
	const BK = BifurcationKit

norminf(x) = norm(x, Inf)

function Laplacian2D(Nx, Ny, lx, ly, bc = :Dirichlet)
	hx = 2lx/Nx
	hy = 2ly/Ny
	D2x = CenteredDifference(2, 2, hx, Nx)
	D2y = CenteredDifference(2, 2, hy, Ny)
	if bc == :Neumann
		Qx = Neumann0BC(hx)
		Qy = Neumann0BC(hy)
	elseif  bc == :Dirichlet
		Qx = Dirichlet0BC(typeof(hx))
		Qy = Dirichlet0BC(typeof(hy))
	end
	D2xsp = sparse(D2x * Qx)[1]
	D2ysp = sparse(D2y * Qy)[1]

	A = kron(sparse(I, Ny, Ny), D2xsp) + kron(D2ysp, sparse(I, Nx, Nx))
	return A, D2x
end

function NL!(f, u, p, t = 0.)
	@unpack r, μ, ν, c3, c5 = p
	n = div(length(u), 2)
	u1 = @view u[1:n]
	u2 = @view u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1 = @view f[1:n]
	f2 = @view f[n+1:2n]

	@. f1 .= r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
	@. f2 .= r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2

	return f
end

NL(u, p) = NL!(similar(u), u, p)

function Fcgl!(f, u, p, t = 0.)
	mul!(f, p.Δ, u)
	f .= f .+ NL(u, p)
end

Fcgl(u, p, t = 0.) = Fcgl!(similar(u), u, p, t)

# computation of the first derivative
# d1Fcgl(x, p, dx) = ForwardDiff.derivative(t -> Fcgl(x .+ t .* dx, p), 0.)

d1NL(x, p, dx) = ForwardDiff.derivative(t -> NL(x .+ t .* dx, p), 0.)

function dFcgl(x, p, dx)
	f = similar(dx)
	mul!(f, p.Δ, dx)

	nl = d1NL(x, p, dx)
	f .= f .+ nl
end

function Jcgl(u, p, t = 0.)
	@unpack r, μ, ν, c3, c5, Δ = p

	n = div(length(u), 2)
	u1 = @view u[1:n]
	u2 = @view u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1u = zero(u1)
	f2u = zero(u1)
	f1v = zero(u1)
	f2v = zero(u1)

	@. f1u =  r - 2 * u1 * (c3 * u1 - μ * u2) - c3 * ua - 4 * c5 * ua * u1^2 - c5 * ua^2
	@. f1v = -ν - 2 * u2 * (c3 * u1 - μ * u2)  + μ * ua - 4 * c5 * ua * u1 * u2
	@. f2u =  ν - 2 * u1 * (c3 * u2 + μ * u1)  - μ * ua - 4 * c5 * ua * u1 * u2
	@. f2v =  r - 2 * u2 * (c3 * u2 + μ * u1) - c3 * ua - 4 * c5 * ua * u2 ^2 - c5 * ua^2

	jacdiag = vcat(f1u, f2v)

	Δ + spdiagm(0 => jacdiag, n => f1v, -n => f2u)
end

####################################################################################################
Nx = 41*1
	Ny = 21*1
	n = Nx*Ny
	lx = pi
	ly = pi/2

	Δ = Laplacian2D(Nx, Ny, lx, ly)[1]
	par_cgl = (r = 0.5, μ = 0.1, ν = 1.0, c3 = -1.0, c5 = 1.0, Δ = blockdiag(Δ, Δ))
	sol0 = 0.1rand(2Nx, Ny)
	sol0_f = vec(sol0)

####################################################################################################
eigls = EigArpack(1.0, :LM)
# eigls = eig_MF_KrylovKit(tol = 1e-8, dim = 60, x₀ = rand(ComplexF64, Nx*Ny), verbose = 1)
opt_newton = NewtonPar(tol = 1e-9, verbose = true, eigsolver = eigls, maxIter = 20)
opts_br = ContinuationPar(dsmax = 0.02, ds = 0.01, pMax = 2., detectBifurcation = 3, nev = 15, newtonOptions = (@set opt_newton.verbose = false), nInversion = 4)

	br, = @time continuation(Fcgl, Jcgl, vec(sol0), par_cgl, (@lens _.r), opts_br, verbosity = 0)

plot(br)
####################################################################################################
# Look for periodic orbits
f1 = DiffEqArrayOperator(par_cgl.Δ)
f2 = NL!
prob_sp = SplitODEProblem(f1, f2, sol0_f, (0.0, 120.0), @set par_cgl.r = 1.2; reltol = 1e-8, dt = 0.1)
prob = ODEProblem(Fcgl, sol0_f, (0.0, 120.0), (@set par_cgl.r = 1.2))#, jac = Jcgl, jac_prototype = Jcgl(sol0_f, par_cgl))
####################################################################################################
# sol = @time solve(prob, Vern9(); abstol=1e-14, reltol=1e-14)
sol = @time solve(prob_sp, ETDRK2(krylov=true); abstol=1e-14, reltol=1e-14, dt = 0.1) #1.78s
# sol = @time solve(prob, LawsonEuler(krylov=true, m=50); abstol=1e-14, reltol=1e-14, dt = 0.1)
# sol = @time solve(prob_sp, CNAB2(linsolve=LinSolveGMRES()); abstol=1e-14, reltol=1e-14, dt = 0.03)

plot(sol.t, [norm(v[1:Nx*Ny], Inf) for v in sol.u], xlims=(115, 120))

# plotting the solution as a movie
for ii = 1:20:length(sol.t)
	# heatmap(reshape(sol[1:Nx*Ny,ii],Nx,Ny),title="$(sol.t[ii])") |> display
end

####################################################################################################
# this encodes the functional for the Shooting problem
probSh = ShootingProblem(
	# we pass the ODEProblem encoding the flow and the time stepper
	prob_sp, ETDRK2(krylov = true),
	[sol[:, end]], abstol = 1e-10, reltol = 1e-8)

initpo = vcat(sol(116.), 4.9) |> vec
	probSh(initpo, @set par_cgl.r = 1.2) |> norminf

ls = GMRESIterativeSolvers(reltol = 1e-4, N = 2n + 1, maxiter = 50, verbose = false)
	optn = NewtonPar(verbose = true, tol = 1e-9,  maxIter = 25, linsolver = ls)
outpo, = @time newton(probSh, initpo, (@set par_cgl.r = 1.2), optn; normN = norminf)

heatmap(reshape(outpo[1:Nx*Ny], Nx, Ny), color = :viridis)

eig = EigKrylovKit(tol = 1e-7, x₀ = rand(2Nx*Ny), verbose = 2, dim = 40)
	opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.03, ds= -0.01, pMax = 2.5, maxSteps = 32, newtonOptions = (@set optn.eigsolver = eig), nev = 15, precisionStability = 1e-3, detectBifurcation = 0, plotEveryStep = 1)
br_po, upo, = @time continuation(probSh, outpo, (@set par_cgl.r = 1.2), (@lens _.r),
		opts_po_cont;
		verbosity = 3,
		plot = true,
		# tangentAlgo = BorderedPred(),
		linearAlgo = MatrixFreeBLS(@set ls.N = probSh.M*2n+2),
		# callbackN = cb_ss,
		plotSolution = (x, p; kwargs...) -> heatmap!(reshape(x[1:Nx*Ny], Nx, Ny); color=:viridis, kwargs...),
		recordFromSolution = (u, p; k...) -> BK.getAmplitude(probSh, u, (@set par_cgl.r = p); ratio = 2), normC = norminf)

####################################################################################################
# automatic branch switching
jet  = BK.getJet(Fcgl, Jcgl)

ls = GMRESIterativeSolvers(reltol = 1e-4, maxiter = 50, verbose = false)
	optn = NewtonPar(verbose = true, tol = 1e-9,  maxIter = 25, linsolver = ls)
eig = EigKrylovKit(tol = 1e-7, x₀ = rand(2Nx*Ny), verbose = 2, dim = 40)
	opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.02, ds= 0.01, pMax = 2.5, maxSteps = 32, newtonOptions = (@set optn.eigsolver = eig), nev = 15, precisionStability = 1e-3, detectBifurcation = 0, plotEveryStep = 1)

Mt = 1 # number of time sections
br_po, = continuation(
	jet...,	br, 2,
	# arguments for continuation
	opts_po_cont,
	# probSh;
	ShootingProblem(Mt, prob_sp, ETDRK2(krylov = true); abstol = 1e-10, reltol = 1e-8) ;
	verbosity = 3, plot = true, ampfactor = 1.5, δp = 0.01,
	# callbackN = (state; kwargs...) -> (println("--> amplitude = ", BK.amplitude(state.x, n, M; ratio = 2));true),
	linearAlgo = MatrixFreeBLS(@set ls.N = Mt*2n+2),
	finaliseSolution = (z, tau, step, contResult; k...) ->begin
		BK.haseigenvalues(contResult) && Base.display(contResult.eig[end].eigenvals)
		return true
	end,
	recordFromSolution = (u, p; k...) -> BK.getAmplitude(p.prob, u, (@set par_cgl.r = p.p); ratio = 2),
	normC = norminf)

#ShootingProblem(1, par_cgl, prob_sp, ETDRK2(krylov = true)) ;


# Runs all unit tests
using LinearLeastSquares
using Base.Test

info("JULIA VERSION: $(VERSION)") # important as Julia versions are volatile
include("helpers.jl")
include("add_subtract.jl")
include("multiply_divide.jl")
include("stack.jl")
include("diag.jl")
include("repmat.jl")
include("index.jl")
include("reshape.jl")
include("transpose.jl")
include("sum_squares.jl")
include("compound.jl")
include("constraints.jl")
include("problems.jl")


using MultivariateOrthogonalPolynomials, DifferentialEquations

Z = Zernike(1)
W = Weighted(Z)
xy = axes(W,1)
Δ = Z \ Laplacian(xy) * W


function mandel(z)
    c = z
    maxiter = 80
    for n = 1:maxiter
        if abs(z) > 2
            return n-1
        end
        z = z^2 + c
    end
    return maxiter
end

function randmatstat(t)
    n = 5
    v = zeros(t)
    w = zeros(t)
    for i = 1:t
        a = randn(n,n)
        b = randn(n,n)
        c = randn(n,n)
        d = randn(n,n)
        P = [a b c d]
        Q = [a b; c d]
        v[i] = trace((P.'*P)^4)
        w[i] = trace((Q.'*Q)^4)
    end
    std(v)/mean(v), std(w)/mean(w)
end


module AeroMDAO

#----------------------IMPORTS--------------------------------#
using LinearAlgebra
using StaticArrays
using CoordinateTransformations, Rotations
using ForwardDiff, DiffResults
using PrettyTables

## Math tools
#==========================================================================================#

include("Tools/MathTools.jl")
import .MathTools: tupvector, fwdsum, fwddiff, fwddiv, weighted_vector, vectarray, slope, splitat, adj3, columns, extend_yz, reflect_mapper, cosine_interp, structtolist, inverse_rotation, rotation, affine_2D, Point2D, Point3D, x, y, z, reshape_array, midpair_map, partition, uniform_spacing, linear_spacing, cosine_spacing, sine_spacing

export Point2D, Point3D, tupvector, tuparray, midpair_map, cosine_interp, sine_spacing, rotation, inverse_rotation, affine_2D, fwdsum, fwddiff, fwddiv, partition, weighted_vector, uniform_spacing, linear_spacing, cosine_spacing, sine_spacing, structtolist, vectarray

## Non-dimensionalization
#==========================================================================================#

include("Tools/NonDimensional.jl")
import .NonDimensional: dynamic_pressure, force_coefficient, moment_coefficient, rate_coefficient, pressure_coefficient, aerodynamic_coefficients, print_coefficients, reynolds_number, print_derivatives, force, moment

export dynamic_pressure, force_coefficient, moment_coefficient, rate_coefficient, pressure_coefficient, aerodynamic_coefficients, print_coefficients, reynolds_number, print_derivatives, force, moment

## Panels
#===========================================================================#

include("Geometry/PanelGeometry/PanelGeometry.jl")
import .PanelGeometry: AbstractPanel, AbstractPanel2D, Panel2D, WakePanel2D, AbstractPanel3D, Panel3D, p1, p2, p3, p4, zero, collocation_point, panel_length, transform_panel, transform_panel_points, panel_angle, panel_tangent, panel_normal, panel_location, panel_area, panel_coords, transform, midpoint, panel_points, wake_panel, wake_panels, reverse_panel, panel_velocity, panel_scalar, trailing_edge_panel, get_surface_values, panel_vector, panel_dist, average_chord, average_width, wetted_area

export AbstractPanel, AbstractPanel2D, Panel2D, WakePanel2D, AbstractPanel3D, Panel3D, Point2D, collocation_point, p1, p2, p3, p4, transform, panel_normal, midpoint, panel_location, panel_tangent, panel_points, panel_dist, wake_panel, wake_panels, panel_area, reverse_panel, panel_length, transform_panel, panel_angle, panel_vector, panel_velocity, panel_scalar, trailing_edge_panel, get_surface_values, average_chord, average_width, wetted_area

## Wing geometry
#==========================================================================================#

include("Geometry/AircraftGeometry/AircraftGeometry.jl")
import .AircraftGeometry: Aircraft, Foil, kulfan_CST, naca4, camber_CST, paneller, read_foil, split_foil, foil_camthick, camthick_foil, camber_thickness, cosine_foil, camthick_to_CST, coords_to_CST, max_thickness_to_chord_ratio_location, Fuselage, projected_area, length, cosine_spacing, HalfWing, HalfWingSection, Wing, WingSection, mean_aerodynamic_chord, span, aspect_ratio, projected_area, taper_ratio, info, leading_edge, trailing_edge, chop_leading_edge, chop_trailing_edge, chop_wing, wing_bounds, paneller, mesh_horseshoes, mesh_wing, mesh_cambers, make_panels, vlmesh_wing, max_tbyc_sweeps, mean_aerodynamic_center, panel_wing, wetted_area, number_of_spanwise_panels, spanwise_spacing, coordinates, camber_coordinates, surface_coordinates,foils, chords, twists, spans, dihedrals, sweeps, left, right

export Aircraft, Foil, kulfan_CST, naca4, camber_CST, paneller, read_foil, split_foil, foil_camthick, camthick_foil, camber_thickness, cosine_foil, camthick_to_CST, coords_to_CST, max_thickness_to_chord_ratio_location, # 2D setups
Fuselage, projected_area, length, cosine_spacing, # Fuselage
HalfWing, HalfWingSection, Wing, WingSection, mean_aerodynamic_chord, span, aspect_ratio, projected_area, taper_ratio, info, leading_edge, trailing_edge, chop_leading_edge, chop_trailing_edge, chop_wing, wing_bounds, paneller, mesh_horseshoes, mesh_wing, mesh_cambers, make_panels, vlmesh_wing, max_tbyc_sweeps, mean_aerodynamic_center, panel_wing, wetted_area, number_of_spanwise_panels, spanwise_spacing, coordinates, camber_coordinates, surface_coordinates, foils, chords, twists, spans, dihedrals, sweeps, left, right

## Laplace
#==========================================================================================#

include("Tools/Laplace.jl")
import .Laplace: Uniform2D, Freestream, velocity, potential, stream, aircraft_velocity, cartesian_to_freestream, freestream_to_cartesian

export Uniform2D, velocity, Freestream, aircraft_velocity, stream, vortex_stream_1, vortex_stream_2, source_stream, cartesian_to_freestream, freestream_to_cartesian

## Aerodynamic analyses
#==========================================================================================#

## Doublet-source panel method

include("Aerodynamics/DoubletSource/DoubletSource.jl")
import .DoubletSource: solve_problem, doublet_matrix, source_matrix, boundary_vector, wake_panels, source_strengths, tangential_velocities, solve_strengths, eval_coefficients, lift_coefficient

export doublet_matrix, source_matrix, boundary_vector, wake_panels, source_strengths, tangential_velocities, solve_strengths

## Linear-strength source and vorticity panel method

include("Aerodynamics/LinearVortexSource/LinearVortexSource.jl")
import .LinearVortexSource: total_velocity, source_velocity, vortex_velocity, vortex_influence_matrix, source_influence_matrix, neumann_boundary_condition, kutta, two_point_matrix, linear_source_matrix, linear_vortex_matrix, constant_source_matrix, constant_source_boundary_condition

export total_velocity, source_velocity, vortex_velocity, vortex_influence_matrix, source_influence_matrix, neumann_boundary_condition, kutta, two_point_matrix, linear_source_matrix, linear_vortex_matrix, constant_source_matrix, constant_source_boundary_condition

## Vortex lattice

include("Aerodynamics/VortexLattice/VortexLattice.jl")
import .VortexLattice: Horseshoe, streamlines, solve_system, transform, bound_leg, bound_leg_center, bound_leg_vector, r1, r2, points, horseshoe_line, horseshoe_point, make_horseshoes, nearfield_drag, evaluate_coefficients, body_to_wind_axes, body_to_stability_axes, stability_to_body_axes, wind_to_body_axes, case_dynamics, VLMState, VLMSystem, compute_influence_matrix!, compute_boundary_condition!, compute_collocation_points!, compute_horseshoes!, solve_system!, compute_farfield_forces!, solve_case!, compute_wake_properties!, normals, horseshoes, collocation_points, AIC, RHS, circulations, surface_force_coefficients, surface_moment_coefficients, farfield_force_coefficients

export Horseshoe, streamlines, solve_system, transform, bound_leg, bound_leg_center, bound_leg_vector, r1, r2, points, horseshoe_line, horseshoe_point, make_horseshoes, nearfield_drag, evaluate_coefficients, body_to_wind_axes, body_to_stability_axes, stability_to_body_axes, wind_to_body_axes, case_dynamics,  # Pure
VLMState, VLMSystem, compute_influence_matrix!, compute_boundary_condition!, compute_collocation_points!, compute_horseshoes!, solve_system!, compute_farfield_forces!, solve_case!, compute_wake_properties!, normals, horseshoes, collocation_points, AIC, RHS, circulations, surface_force_coefficients, surface_moment_coefficients, farfield_force_coefficients # Impure

## Profile drag estimation

include("Aerodynamics/profile_drag.jl")

export wetted_area_drag, profile_drag_coefficient

## Cases

include("Aerodynamics/Cases/cases.jl")
include("Aerodynamics/Cases/stability_cases.jl")
include("Aerodynamics/Cases/foil_cases.jl")

export solve_case, solve_stability_case, streamlines, print_case, print_info

## Post-processing
#==========================================================================================#

include("Tools/plot_tools.jl")

export plot_panels, plot_streams, plot_wing, plot_surface

end


using SVM
using Base.Test
srand(1)
iris = readcsv(joinpath(dirname(@__FILE__), "iris.csv"))
X = map(Float64, iris[:, 1:4]')
p, n = size(X)
train = rand(Bool, n)
Y = [species == "setosa" ? 1.0 : -1.0 for species in iris[:, 5]]
model = svm(X[:,train], Y[train])
@test predict(model, X[:, ~train]) == Y[~train]


"""
Computes Z4 topological invariants given an irvsp.out file, range of bands [nstart,nend] for systems with TRS and IS
"""
function compute_Z4(myFile::String, nstart::Int64, nend::Int64)
    parity = parse_parity_from_irvsp(myFile, nstart, nend)
    # Compute the Strong Index
    nₖ⁺ = sum(x->x>0, parity, dims=1) # Compute number of +1 parity eigenvalues
    nₖ⁻ = sum(x->x<0, parity, dims=1) # Compute number of -1 parity eigenvalues
    κ = mod(1/4*sum( nₖ⁺ .- nₖ⁻), 4) # Compute κ, the Z₄ invariant
    ν₀ = mod(κ, 2)

    # Compute the Weak Indices
    nₖ₁⁺ = nₖ⁺[[2,5,6,8]]
    nₖ₁⁻ = nₖ⁻[[2,5,6,8]]
    ν₁ = mod(1/4*sum(nₖ₁⁺.- nₖ₁⁻), 2)
    nₖ₂⁺ = nₖ⁺[[3,5,7,8]]
    nₖ₂⁻ = nₖ⁻[[3,5,7,8]]
    ν₂ = mod(1/4*sum(nₖ₂⁺ .- nₖ₂⁻), 2)
    nₖ₃⁺ = nₖ⁺[[4,6,7,8]]
    nₖ₃⁻ = nₖ⁻[[4,6,7,8]]
    ν₃ = mod(1/4*sum(nₖ₃⁺.- nₖ₃⁻), 2)

    Z₄ = "Z₄ = (κ;ν₁ν₂ν₃) = ($(trunc(Int,κ));$(trunc(Int,ν₁))$(trunc(Int,ν₂))$(trunc(Int,ν₃)))"
    Z₂ = "Z₂ = (ν₀;ν₁ν₂ν₃) = ($(trunc(Int,ν₀));$(trunc(Int,ν₁))$(trunc(Int,ν₂))$(trunc(Int,ν₃))), where ν₀ = κ mod 2"
    if κ ∈ (1, 3)
        result = "strong topological insulator"
    elseif κ == 0
        result = "trivial insulator"
        if (ν₁ == 1 || ν₂ == 1 || ν₃ == 1)
            result = "weak topological insulator"
        end
    elseif κ == 2
        result = "higher-order topological insulator"
    end

    println(Z₄)
    println(Z₂)
    println("System is a $result")
end

"""
Computes Z2 topological invarilnts given an irvsp.out file, range of bands [nstart,nend] for systems with TRS and IS
"""
function compute_Z2(myFile::String, nstart::Int64, nend::Int64)
    parity = parse_parity_from_irvsp(myFile, nstart, nend)
    δ = ones(8)
    for n = 1:2:(nend-nstart+1)
        δ .*= parity[n, :]
    end
    minus1_ν₀ = prod(δ)
    minus1_ν₁ = prod(δ[[2,5,6,8]])
    minus1_ν₂ = prod(δ[[3,5,7,8]])
    minus1_ν₃ = prod(δ[[4,6,7,8]])
    ν₀ = minus1_ν₀ == 1 ? 0 : 1
    ν₁ = minus1_ν₁ == 1 ? 0 : 1
    ν₂ = minus1_ν₂ == 1 ? 0 : 1
    ν₃ = minus1_ν₃ == 1 ? 0 : 1
    if ν₀ == 1
        result = "strong topological insulator"
    elseif ν₀ == 0
        result = "trivial insulator"
        if (ν₁ == 1 || ν₂ == 1 || ν₃ == 1)
            result = "weak topological insulator"
        end
    end
    Z₂ = "Z₂ = (ν₀;ν₁ν₂ν₃) = ($(trunc(Int,ν₀));$(trunc(Int,ν₁))$(trunc(Int,ν₂))$(trunc(Int,ν₃)))"
    println(Z₂)
    println("System is a $result")
end


# Copyright (c) 2019 Iain Dunning, and contributors                             #src
#                                                                               #src
# This work is licensed under a Creative Commons Attribution-ShareAlike 4.0     #src
# International License. URL: http://creativecommons.org/licenses/by-sa/4.0.    #src

# # Rocket Control

# **Originally Contributed by**: Iain Dunning

# This tutorial shows how to solve a nonlinear rocketry control problem.
# The problem was drawn from the [COPS3](http://www.mcs.anl.gov/~more/cops/cops3.pdf)
# benchmark.

# Our goal is to maximize the final altitude of a vertically launched rocket.

# We can control the thrust of the rocket, and must take account of the rocket
# mass, fuel consumption rate, gravity, and aerodynamic drag.

# Let us consider the basic description of the model (for the full description,
# including parameters for the rocket, see the COPS3 PDF)

# ### Overview

# We will use a discretized model of time, with a fixed number of time steps,
# $n$.

# We will make the time step size $\Delta t$, and thus the final time
# $t_f = n \cdot \Delta t$, a variable in the problem. To approximate the
# derivatives in the problem we will use the [trapezoidal rule](http://en.wikipedia.org/wiki/Trapezoidal_rule).

# ### State and Control

# We will have three state variables:
#
# * Velocity, $v$
# * Altitude, $h$
# * Mass of rocket and remaining fuel, $m$
#
# and a single control variable, thrust $T$.

# Our goal is thus to maximize $h(t_f)$.

# Each of these corresponds to a JuMP variable indexed by the time step.

# ### Dynamics

# We have three equations that control the dynamics of the rocket:
#
# Rate of ascent: $$h^\prime = v$$
# Acceleration: $$v^\prime = \frac{T - D(h,v)}{m} - g(h)$$
# Rate of mass loss: $$m^\prime = -\frac{T}{c}$$
#
# where drag $D(h,v)$ is a function of altitude and velocity, and gravity
# $g(h)$ is a function of altitude.

# These forces are defined as
#
# $$D(h,v) = D_c v^2 exp\left( -h_c \left( \frac{h-h(0)}{h(0)} \right) \right)$$
# and
# $$g(h) = g_0 \left( \frac{h(0)}{h} \right)^2$$
#
# The three rate equations correspond to JuMP constraints, and for convenience
# we will represent the forces with nonlinear expressions.

using JuMP
import Ipopt
import Plots

# Create JuMP model, using Ipopt as the solver

rocket = Model(Ipopt.Optimizer)
set_silent(rocket)

# ## Constants

# Note that all parameters in the model have been normalized
# to be dimensionless. See the COPS3 paper for more info.

h_0 = 1    # Initial height
v_0 = 0    # Initial velocity
m_0 = 1    # Initial mass
g_0 = 1    # Gravity at the surface

T_c = 3.5  # Used for thrust
h_c = 500  # Used for drag
v_c = 620  # Used for drag
m_c = 0.6  # Fraction of initial mass left at end

c     = 0.5 * sqrt(g_0 * h_0)  # Thrust-to-fuel mass
m_f   = m_c * m_0              # Final mass
D_c   = 0.5 * v_c * m_0 / g_0  # Drag scaling
T_max = T_c * g_0 * m_0        # Maximum thrust

n = 800    # Time steps

# ## Decision variables

@variables(rocket, begin
    Δt ≥ 0, (start = 1/n) # Time step
    ## State variables
    v[1:n] ≥ 0            # Velocity
    h[1:n] ≥ h_0          # Height
    m_f ≤ m[1:n] ≤ m_0    # Mass
    ## Control variables
    0 ≤ T[1:n] ≤ T_max    # Thrust
end)

# ## Objective

# The objective is to maximize altitude at end of time of flight.

@objective(rocket, Max, h[n])

# ## Initial conditions

fix(v[1], v_0; force = true)
fix(h[1], h_0; force = true)
fix(m[1], m_0; force = true)
fix(m[n], m_f; force = true)

# ## Forces

@NLexpressions(rocket, begin
    ## Drag(h,v) = Dc v^2 exp( -hc * (h - h0) / h0 )
    drag[j = 1:n], D_c * (v[j]^2) * exp(-h_c * (h[j] - h_0) / h_0)
    ## Grav(h)   = go * (h0 / h)^2
    grav[j = 1:n], g_0 * (h_0 / h[j])^2
    ## Time of flight
    t_f, Δt * n
end)

# ## Dynamics

for j in 2:n
    ## h' = v
    ## Rectangular integration
    ## @NLconstraint(rocket, h[j] == h[j - 1] + Δt * v[j - 1])
    ## Trapezoidal integration
    @NLconstraint(rocket, h[j] == h[j - 1] + 0.5 * Δt * (v[j] + v[j - 1]))
    ## v' = (T-D(h,v))/m - g(h)
    ## Rectangular integration
    ## @NLconstraint(
    ##     rocket,
    ##     v[j] == v[j - 1] + Δt *((T[j - 1] - drag[j - 1]) / m[j - 1] - grav[j - 1])
    ## )
    ## Trapezoidal integration
    @NLconstraint(
        rocket,
        v[j] == v[j-1] +
            0.5 * Δt * (
                (T[j] - drag[j] - m[j] * grav[j]) / m[j] +
                (T[j - 1] - drag[j - 1] - m[j - 1] * grav[j - 1]) / m[j - 1]
            )
    )
    ## m' = -T/c
    ## Rectangular integration
    ## @NLconstraint(rocket, m[j] == m[j - 1] - Δt * T[j - 1] / c)
    ## Trapezoidal integration
    @NLconstraint(rocket, m[j] == m[j - 1] - 0.5 * Δt * (T[j] + T[j-1]) / c)
end

# Solve for the control and state
println("Solving...")
status = optimize!(rocket)

# ## Display results

println("Max height: ", objective_value(rocket))

#-

function my_plot(y, ylabel)
    return Plots.plot(
        (1:n) * value.(Δt),
        value.(y)[:];
        xlabel = "Time (s)",
        ylabel = ylabel,
    )
end

Plots.plot(
    my_plot(h, "Altitude"),
    my_plot(m, "Mass"),
    my_plot(v, "Velocity"),
    my_plot(T, "Thrust");
    layout = (2, 2),
    legend = false,
    margin = 1Plots.cm,
)


VERSION >= v"0.4.0" && __precompile__()

module Jackknife

"""
    Jackknife.leaveoneout(estimator::Function, x)

Estimate the parameter `estimator(x)` for each subsample of `x`, systematically
omitting each index one at a time. The result is a vector of length `length(x)-1`
of parameter estimates.
"""
function leaveoneout{T<:Real}(estimator::Function, x::AbstractVector{T})
    n = length(x)
    n > 1 || throw(ArgumentError("The sample must have size > 1"))

    S = typeof(estimator(x[1]))
    θ = Array{S}(n)

    @inbounds for i = 1:n
        θ[i] = estimator(vcat(x[1:i-1], x[i+1:end]))
    end

    return θ
end

"""
    Jackknife.variance(estimator::Function, x)

Compute the jackknife estimate of the variance of `estimator`, which is given as a
function that computes a point estimate when passed a real-valued vector `x`.
"""
function variance{T<:Real}(estimator::Function, x::AbstractVector{T})
    θ = leaveoneout(estimator, x)
    n = length(x)
    return Base.var(θ) * (n - 1)^2 / n
end

"""
    Jackknife.bias(estimator::Function, x)

Compute the jackknife estimate of the bias of `estimator`, which is given as a
function that computes a point estimate when passed a real-valued vector `x`.
"""
function bias{T<:Real}(estimator::Function, x::AbstractVector{T})
    θ = leaveoneout(estimator, x)
    return (length(x) - 1) * (mean(θ) - estimator(x))
end

"""
    Jackknife.estimate(estimator::Function, x)

Compute the bias-corrected jackknife estimate of the parameter `estimator(x)`.
"""
function estimate{T<:Real}(estimator::Function, x::AbstractVector{T})
    θ = leaveoneout(estimator, x)
    n = length(x)
    return n * estimator(x) - (n - 1) * mean(θ)
end

end # module


# using Revise
# using Pkg
# cd("..")
# Pkg.activate(".")

using Test
using WeightedOnlineStats
import OnlineStatsBase: eachcol, eachrow
using StatsBase
import OnlineStats: Extrema
using Statistics
using Random
using MultivariateStats

Random.seed!(123)
l = 1000
x = rand(l);
xmis = convert(Array{Union{Float64, Missing}}, x);
x2 = rand(l, 5)
x2mis = convert(Array{Union{Float64, Missing}}, x2)
x2mis[end, 1] = missing
x2mis[end - 1, 1] = missing
w = rand(l);
wmis = convert(Array{Union{Float64, Missing}}, w);
wmis[end] = missing
wmis[end - 2] = missing

Base.isapprox(x::Tuple, y::Tuple) =
    reduce(&, map((a, b) -> a ≈ b, x, y))

missing_to_nan(x::Array{Union{Missing, T}}) where T = map(y -> y === missing ? T(NaN) : y, x)

function test_fit(T::Type{<: WeightedOnlineStats.WeightedOnlineStat{S}},
                  data, weights, unpack_fun, jfun) where S
    l = length(data)
    @assert l == length(weights)

    o = T()
    for (xi, wi) in zip(data, weights)
        fit!(o, xi, wi)
    end

    @test unpack_fun(o) ≈ jfun(data, weights)
    @test eltype(unpack_fun(o)) == eltype(o)
end

include("test_hist.jl")
include("test_sum.jl")
include("test_mean.jl")
include("test_var.jl")
include("test_cov.jl")
include("test_pca.jl")


# ---
# title: 1009. Complement of Base 10 Integer
# id: problem1009
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Math
# link: <https://leetcode.com/problems/complement-of-base-10-integer/description/>
# hidden: true
# ---
#
# Every non-negative integer `N` has a binary representation.  For example, `5`
# can be represented as `"101"` in binary, `11` as `"1011"` in binary, and so
# on.  Note that except for `N = 0`, there are no leading zeroes in any binary
# representation.
#
# The _complement_  of a binary representation is the number in binary you get
# when changing every `1` to a `0` and `0` to a `1`.  For example, the
# complement of `"101"` in binary is `"010"` in binary.
#
# For a given number `N` in base-10, return the complement of it's binary
# representation as a base-10 integer.
#
#
#
# **Example 1:**
#
#
#
#     Input: 5
#     Output: 2
#     Explanation: 5 is "101" in binary, with complement "010" in binary, which is 2 in base-10.
#
#
# **Example 2:**
#
#
#
#     Input: 7
#     Output: 0
#     Explanation: 7 is "111" in binary, with complement "000" in binary, which is 0 in base-10.
#
#
# **Example 3:**
#
#
#
#     Input: 10
#     Output: 5
#     Explanation: 10 is "1010" in binary, with complement "0101" in binary, which is 5 in base-10.
#
#
#
#
# **Note:**
#
#   1. `0 <= N < 10^9`
#   2. This question is the same as 476: <https://leetcode.comhttps://leetcode.com/problems/number-complement/>
#
#
## @lc code=start
using LeetCode

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


export Payment, CashFlow, npv, irr

struct Payment{T, S<:Real}
    date::T
    amount::S
end

struct CashFlow{T, S<:Real, R<:DayCountConvention}
    flow::Array{Payment{T,S},1}
    daycount::R
    basedate::Date
end

Base.length(cf::CashFlow) = length(cf.flows)

Base.convert(::Type{Payment{T}}, p::Payment{Date}, daycount, base = today()) where {T<:Real} =
    Payment(yearfraction(daycount, base, p.date, T), p.amount)

Base.convert(::Type{CashFlow{T}}, cf::CashFlow{Date}) where {T<:Real} =
    CashFlow([convert(Payment{T}, p, cf.daycount, cf.basedate) for p in cf.flow], cf.daycount, cf.basedate)

Base.show(io::IO, cf::CashFlow) = begin
    println(io, rpad("Day count convention:", 23, ' ') * "$(cf.daycount)")
    println(io, rpad("Base date:", 23, ' ') * "$(cf.basedate)")
    println(io, "Cash flow:")
    for p in cf.flow
        println(io, rpad("  $(p.date)", 23, ' ') * "$(p.amount)")
    end
end

npv(flow, rate) = sum(p.amount / (1 + rate)^p.date for p in flow)
npv(cf::CashFlow{T}, rate::T) where {T<:Real} = npv(cf.flow, rate)
npv(cf::CashFlow, rate::T) where {T<:Real} = npv(convert(CashFlow{T}, cf), rate)

irr(flow, r0) = fzero(r -> npv(flow, r), r0)
irr(cf::CashFlow{T}, r0::T = 0.1) where {T<:Real} = irr(cf.flow, r0)
irr(cf::CashFlow, r0::T = 0.1) where {T<:Real} = irr(convert(CashFlow{T}, cf), r0)


function update_mean(μ::Real, n::Integer, x::Real; remove::Bool=false)
    mask = 2 * Int(!remove) - 1
    new_n = n + mask
    if new_n == 0
        return 0.0, 0
    elseif new_n < 0
        error("Cannot remove an element from an empty collection")
    else
        sum_x = μ * n
        new_sum_x = sum_x + mask * x
        new_μ = new_sum_x / new_n
        return new_μ, new_n
    end
end

function update_mean_std(μ::Real, σ::Real, n::Integer, x::Real; remove::Bool=false)
    new_μ, new_n = update_mean(μ, n, x; remove=remove)
    if new_n == 0
        return 0.0, 0.0, 0
    end
    σ² = σ^2
    μ_x² = σ² + μ^2
    new_μ_x², new_n = update_mean(μ_x², n, x^2; remove=remove)
    new_σ² = new_μ_x² - new_μ^2
    # if new_σ² < -0.1
    #     @warn "precision error?" new_σ²
    # end
    new_σ² = max(0, new_σ²)
    new_σ = √(new_σ²)
    return new_μ, new_σ, new_n
end

function update_mean_std_corr(μ_x::Real, μ_y::Real, σ_x::Real, σ_y::Real, ρ::Real, n::Integer, x::Real, y::Real; remove::Bool=false)
    new_μ_x, new_σ_x, new_n = update_mean_std(μ_x, σ_x, n, x; remove=remove)
    if new_n == 0
        return 0.0, 0.0, 0.0, 0.0, 0.0, 0
    end
    new_μ_y, new_σ_y, new_n = update_mean_std(μ_y, σ_y, n, y; remove=remove)
    cov = ρ * σ_x * σ_y
    μ_xy = cov + μ_x * μ_y
    new_μ_xy, new_n = update_mean(μ_xy, n, x * y; remove=remove)
    new_cov = new_μ_xy - new_μ_x * new_μ_y
    if abs(new_cov) < 1e-9
        new_ρ = 0.0
    else
        new_ρ = new_cov / (new_σ_x * new_σ_y)
    end
    # if abs(new_ρ) > 1 + 1e-9
    #     @warn "precision error?" new_ρ
    # end
    new_ρ = min(max(-1, new_ρ), 1)
    return new_μ_x, new_μ_y, new_σ_x, new_σ_y, new_ρ, new_n
end

# stats = (0.0, 0.0, 0.0, 0.0, 0.0, 0)

# stats = update_mean_std_corr(stats..., 1, 1)
# stats = update_mean_std_corr(stats..., 2, 1.4)
# stats = update_mean_std_corr(stats..., 1.4, 1.5)

function get_trace_distribution(w_trace, possible_weights)
    dist = Dict{Float64,Float64}()
    for w in possible_weights
        dist[round(w, digits=2)] = 0
    end
    for w in w_trace
        dist[round(w, digits=2)] = dist[round(w, digits=2)] + 1 / length(w_trace)
    end
    return dist
end


#workspace()

using IntervalArithmetic
using EAGODomainReduction

single_expr = :(x[1]+x[2]^2-exp(x[1]))
tape_out = Generate_Tape(single_expr,2,-2,4)
dag_out = getDAG(single_expr,[1,2])

#=
generates the tape list for the constraints -1 <= (x[1]+x[2]^2-exp(x[1]) <= 2
and 2 <= :(x[1]*x[2]-x[5]) <= 4 and runs the contractor 6 times in place on the
variable box X = [1,20]^5
=#
mult_expr = [:(x[1]+x[2]^2)
             :(x[1]+x[5])
            ]
tapelist_out = Generate_TapeList(mult_expr,5,[-1.0,2.0],[2.0,3.0])
X = [Interval(0.5,20.0) for i=1:5]
DAGContractor!(X,tapelist_out,6)
println("contracted X value: ", X)

#=
generates the tape list for the constraints -1 <= (x[1]+x[2]^2-exp(x[1]) <= 2
and 2 <= :(x[1]*x[2]-x[5]) <= 4 with x[5]==5 and runs the contractor 6 times in place on the
variable box X = [1,20]^5
=#

X = [Interval(-10.0,20.0) for i=1:5]
ftapelist_out = Generate_Fixed_TapeList(mult_expr,4,[-8.0,0.5],[2.0,4.5],[[5.0]])
DAGContractor!(X,ftapelist_out,6)
println("contracted X value: ", X)


# Spherical harmonic I/O, storage, and conversions

@testset "SHCilmToCindex and SHCindexToCilm" begin
    Random.seed!(0)
    # Choose random values
    lmax = 2
    cilm = randn(2, lmax + 1, lmax + 1)

    cindex = SHCilmToCindex(cilm)
    @test size(cindex) == (2, (lmax + 1) * (lmax + 2) ÷ 2)
    # These loops test all values in cindex
    for l in 0:lmax, m in 0:l, i in 1:(1 + (m > 0))
        @test cindex[i, l * (l + 1) ÷ 2 + m + 1] == cilm[i, l + 1, m + 1]
    end

    cilm′ = SHCindexToCilm(cindex)
    for l in 0:lmax, m in 0:l, i in 1:(1 + (m > 0))
        @test isapprox(cilm′[i, l + 1, m + 1], cilm[i, l + 1, m + 1])
    end

    cindex′ = SHCilmToCindex(cilm′)
    @test isapprox(cindex′, cindex)
end

@testset "SHCilmToVector SHVectorToCilm, and YlmIndexVector" begin
    Random.seed!(0)
    # Choose random values
    lmax = 2
    cilm = randn(2, lmax + 1, lmax + 1)

    vector = SHCilmToVector(cilm, lmax)
    @test length(vector) == (lmax + 1)^2
    coverage = falses(length(vector))
    for l in 0:lmax, m in 0:l, i in 1:(1 + (m > 0))
        index = YlmIndexVector(i, l, m)
        @test !coverage[index]
        coverage[index] = true
        @test vector[index] == cilm[i, l + 1, m + 1]
    end
    @test all(coverage)

    cilm′ = SHVectorToCilm(vector, lmax)
    for l in 0:lmax, m in 0:l, i in 1:(1 + (m > 0))
        @test isapprox(cilm′[i, l + 1, m + 1], cilm[i, l + 1, m + 1])
    end

    vector′ = SHCilmToVector(cilm′, lmax)
    @test isapprox(vector′, vector)
end


## initialize module
using Pkg;Pkg.activate(".")
include("../src/RecursiveClusterAlgorithm.jl")
using .RecursiveClusterAlgorithm

using VortexDistributions
using Test, Plots

## set up mock data
N = 10
xi = rand(N)
yi = rand(N)
tree = spanning_tree(xi,yi)

## make cluster using vortices and tree
vort = PointVortex.(xi,yi,one.(xi))
c1 = Cluster(vort,tree)
@test c1.vortices |> length == N

## plot cluster
p1 = plot()
plot_cluster!(p1,c1) # check that the plot method still works

## do some random tests of robustness
N = 30
xi,yi = randn(N),randn(N)
tree = spanning_tree(xi,yi)
c1 = Cluster(PointVortex.(xi,yi,one.(xi)),tree)
p1 = plot(grid=false,axis=false)
plot_cluster!(p1,c1)

xi,yi = randn(N),randn(N)
xi .+= 5
tree = spanning_tree(xi,yi)
c1 = Cluster(PointVortex.(xi,yi,-one.(xi)),tree)
plot_cluster!(p1,c1)


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using ViewerGL; GL = ViewerGL

# input of primitive shapes
W,(WW,EW,FW) = Lar.cuboid([1,1], true)
square = (W,EW)
# hierarchical assembly
model = Lar.Struct([
            Lar.Struct([
                square,
                Lar.t(.175,.175), Lar.s(.25,.25),
                square
            ]),
            Lar.t(.825,.825), Lar.s(-.25,-.25),
            square
        ])
# visualization of generated wire-frame model
V,EV = Lar.struct2lar(model)
VV = [[k] for k=1:size(V,2)]
GL.VIEW(GL.numbering(.4)((V,[VV, EV])));

triangles = Lar.triangulate2d(V, EV)
GL.VIEW([ GL.GLGrid(V,triangles) ]);

# generate edges
ev = map(sort,cat([[[u,v], [v,w], [w,u]] for (u,v,w) in triangles]))
# remove duplicated edges from triangulation
nodupev = collect(Set(ev))
GL.VIEW(GL.numbering(.25)((V,[VV, nodupev, triangles]),GL.COLORS[1],0.2));


module ShaleDrillingLikelihood_SumFunctions_Test

DOBTIME = false

using ShaleDrillingLikelihood

using Test
using StatsFuns
using Random
using BenchmarkTools
using InteractiveUtils

using StatsFuns: softmax!

using ShaleDrillingLikelihood: logsumexp!,
    sumprod3,
    softmax3test!,
    logsumexp3test!

using Base: OneTo

@testset "sumprod3" begin
    for n in (1,2,3,4,10,)
        x = randn(n)
        y = randn(n)
        z = randn(n)
        @test sum(x .* y .* z) ≈ sumprod3(x,y,z)
    end
end

@testset "Custom logsumexp functions" begin

    for n in (1,2,3,4,10,100,1000,10_000)
        # println("test logsumexp! for $n elements")
        y = randn(n)
        x = view(y, 1:n)
        z1 = similar(y)
        z2 = similar(y)

        bmark = logsumexp(x)
        @test bmark ≈ logsumexp!(z1,x)
        # @show @code_warntype logsumexp!(z1,x)

        fill!(z1,0)
        softmax!(z1,y)
        logsumexp!(z2,y)
        @test sum(z2) ≈ 1
        @test z1 ≈ z2

        if DOBTIME
            @btime logsumexp($x)
            @btime softmax!($z1,$y)
            @btime logsumexp!($z1,$y)
        end
    end

end

println("printme")

@testset "softmax3! functionw" begin

    ni, nj, nk = 100, 100, 10
    x0 = rand(ni, nj, nk)
    q0 = similar(x0)

    # x = view(x0, :, :, OneTo(nk))
    # q = view(q0, :, :, OneTo(nk))
    x = x0
    q = q0

    tmpmax = zeros(ni,nj)
    lse = similar(tmpmax)

    qcopy = similar(q)
    softmax3test!(qcopy, x)
    @test sum(qcopy; dims=3) ≈ ones(ni,nj)

    lsetest = logsumexp3test!(x)

    # softmax3 w/ matrices
    # @show @code_warntype softmax3!(q, lse, tmpmax, x, 1)
    fill!(q,0)
    fill!(lse,0)
    softmax3!(q, lse, tmpmax, x, 1)
    @test all(sum(q; dims=3) .<= 1)

    fill!(q,0)
    fill!(lse,0)
    softmax3!(q, lse, tmpmax, x)
    @test sum(q; dims=3) ≈ ones(ni,nj)
    @test q ≈ qcopy

    # test w/ view
    nkm1 = nk-1
    @views qvw = q[:,:,1:nkm1]
    @views xvw = x[:,:,1:nkm1]
    # @show @code_warntype softmax3!(qvw, lse, tmpmax, xvw, 1)
    fill!(q,0)
    fill!(lse,0)
    softmax3!(qvw, lse, tmpmax, xvw, 1)
    @test all(sum(qvw; dims=3) .<= 1)

    fill!(q,0)
    fill!(lse,0)
    softmax3!(qvw, lse, tmpmax, xvw)
    @test sum(qvw; dims=3) ≈ ones(ni,nj)

    # test logsumexp
    fill!(q,0)
    fill!(lse,0)
    logsumexp_and_softmax!(lse, q, tmpmax, x, 1)
    @test lse ≈ lsetest
    @test all(q[:,:,1] .< 1)
    @test q[:,:,1] ≈ qcopy[:,:,1]
    @test !(q ≈ qcopy)

    if DOBTIME
        @btime softmax3!($q, $lse, $tmpmax, $x)
        @btime softmax3test!($q, $x)
    end

end

end


using Plots,LaTeXStrings
using LinearAlgebra,FundamentalsNumericalComputation

gr()
default(legend=nothing,linewidth=2)
plt = plot(aspect_ratio=1,showaxis=false,grid=false,xaxis=:none,yaxis=false,layout=(1,2))

u = .3 - .2im;
v = exp(0.1im)*u;

plot!([0,real(u)],0.1.+[0,imag(u)],l=:darkred,subplot=1)
plot!([0, real(v)],0.1.+[0, imag(v)],l=:darkblue,subplot=1)
plot!(real([u, v]),0.1.+imag([u, v]),l=:black,subplot=1)
xlims!(0 ,.34)
ylims!(-.25, .3)

u = .3 - .2im;
v = 1im*u;
plot!([0, real(u)],[0, imag(u)],l=:darkred,subplot=2)
plot!([0 ,real(v)],[0, imag(v)],l=:darkblue,subplot=2)
plot!(real([u, v]),imag([u, v]),l=:black,subplot=2)
xlims!(0 ,.34)
ylims!(-.25, .3)

annotate!(.15,-0.02,text(L"\mathbf{u}",:right),subplot=1)
annotate!(.18,0.02,text(L"\mathbf{v}",:left),subplot=1)
annotate!(.325,-0.09,text(L"\mathbf{u}-\mathbf{v}",:left),subplot=1)
annotate!(.15,-0.12,text(L"\mathbf{u}",:right),subplot=2)
annotate!(.06,0.15,text(L"\mathbf{v}",:left),subplot=2)
annotate!(.27,0.12,text(L"\mathbf{u}-\mathbf{v}",:left),subplot=2)

savefig("nonorthogonal.svg")


# ## Axis
# ### axis

#md # ![FILE_NAME.png](images/FILE_NAME.png)

using Measures
gr()

plot(sin, 0, 2π,
    xticks=0:0.5:2π,
    xrotation=60,
    xtickfont=font(30),
    bottom_margin=15mm)


using LinearAlgebra
using Printf

struct RankInfo
    label::AbstractString
    count::Int
    score_values::Array{Int}
    reward::Int
    penalty::Int
end

mutable struct ScoreTransitionMatrix
    transition_matrix::Array{Float64, 2}
    rank_info::RankInfo
end

mutable struct ScoreSteadyState
    probabilities::Array{Float64}
    rank_info::RankInfo
end

function steady_state(transition_matrix::Array{Float64, 2})
    values = eigvals(transition_matrix')
    vectors = eigvecs(transition_matrix')
    for i=1:length(values)
        if abs(values[i] - 1.0) < 1e-9
            vec = real(vectors[:,i])
            return vec/sum(vec)
        end
    end
    return zeros(length(values))
end

function score_steady_state(score_transition_matrix::ScoreTransitionMatrix)
    steady_state_vector = steady_state(score_transition_matrix.transition_matrix)
    return ScoreSteadyState(steady_state_vector, score_transition_matrix.rank_info)
end

function rank_matrix_from_steady_states(steady_states::Array{ScoreSteadyState})
    count = sum([steady_states[i].rank_info.count for i=1:length(steady_states)])

    rank_matrix = Matrix{Float64}(I, count, count)

    i = 1
    for idx=1:length(steady_states)
        rank_info = steady_states[idx].rank_info
        probabilities = steady_states[idx].probabilities
        for offset=1:rank_info.count
            if i > 1
                for j=1:length(rank_info.score_values)
                    if rank_info.score_values[j] + rank_info.penalty < rank_info.score_values[1]
                        rank_matrix[i,i-1] += 0.5 * probabilities[j]
                    end
                end
                rank_matrix[i,i] -= rank_matrix[i,i-1]
            end
            if i < count
                for j=1:length(rank_info.score_values)
                    if rank_info.score_values[j] + rank_info.reward >= 100
                        rank_matrix[i,i+1] += 0.5 * probabilities[j]
                    end
                end
                rank_matrix[i,i] -= rank_matrix[i,i+1]
            end
            i = i + 1
        end
    end

    return rank_matrix
end

function report_score_steady_states(score_steady_states::Array{ScoreSteadyState})
    for i=1:length(score_steady_states)
        ss = score_steady_states[i]
        println(ss.rank_info.label, ": ")
        for j=1:length(ss.probabilities)
            println("  ", ss.rank_info.score_values[j], " => ", ss.probabilities[j])
        end
    end
end

function report_rank_steady_state(rank_steady_state)
    println("Rank steady state: ", rank_steady_state)

    @printf("C: %.2f%%\n", 100 * sum(rank_steady_state[1:3]))
    @printf("B: %.2f%%\n", 100 * sum(rank_steady_state[4:6]))
    @printf("A: %.2f%%\n", 100 * sum(rank_steady_state[7:9]))
    @printf("S: %.2f%%\n", 100 * rank_steady_state[10])
    @printf("S+: %.2f%%\n", 100 * rank_steady_state[11])
end

function default_score_transitions()
                            # 0  10  20  30  40  50  60  70  80  90
    c_minus_matrix = Float64[0.5  0  0.5  0   0   0   0   0   0   0; #  0
                             0.5  0   0  0.5  0   0   0   0   0   0; # 10
                              0  0.5  0   0  0.5  0   0   0   0   0; # 20
                              0   0  0.5  0   0  0.5  0   0   0   0; # 30
                              0   0   0  0.5  0   0  0.5  0   0   0; # 40
                              0   0   0   0  0.5  0   0  0.5  0   0; # 50
                              0   0   0   0   0  0.5  0   0  0.5  0; # 60
                              0   0   0   0   0   0  0.5  0   0  0.5;# 70
                              0   0   0   0   0   0   0  1.0  0   0; # 80
                              0   0   0   0   0   0   0  0.5 0.5  0 ]# 90

    c_minus = ScoreTransitionMatrix(c_minus_matrix, RankInfo("C-", 1, [10*i for i=0:9], 20, -10))

                      # 0   5  10   15  20  25  30  35  40  45  50  55  60  65  70  75  80  85  90  95
    c_matrix = Float64[ 0   0   0   0.5  0   0  0.5  0   0   0   0   0   0   0   0   0   0   0   0   0; #  0
                        0   0   0    0  0.5  0  0.5  0   0   0   0   0   0   0   0   0   0   0   0   0; #  5
                       0.5  0   0    0   0  0.5  0   0   0   0   0   0   0   0   0   0   0   0   0   0; # 10
                        0  0.5  0    0   0   0  0.5  0   0   0   0   0   0   0   0   0   0   0   0   0; # 15
                        0   0  0.5   0   0   0   0  0.5  0   0   0   0   0   0   0   0   0   0   0   0; # 20
                        0   0   0   0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0   0   0   0; # 25
                        0   0   0    0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0   0   0; # 30
                        0   0   0    0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0   0; # 35
                        0   0   0    0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0; # 40
                        0   0   0    0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0; # 45
                        0   0   0    0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0; # 50
                        0   0   0    0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0; # 55
                        0   0   0    0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0; # 60
                        0   0   0    0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0; # 65
                        0   0   0    0   0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0; # 70
                        0   0   0    0   0   0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0; # 75
                        0   0   0    0   0   0   0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5; # 80
                        0   0   0    0   0   0   0   0   0   0   0   0   0   0  0.5 0.5  0   0   0   0; # 85
                        0   0   0    0   0   0   0   0   0   0   0   0   0   0  0.5  0  0.5  0   0   0; # 90
                        0   0   0    0   0   0   0   0   0   0   0   0   0   0  0.5  0   0  0.5  0   0] # 95

    c = ScoreTransitionMatrix(c_matrix, RankInfo("C", 1, [5*i for i=0:19], 15, -10))

    c_plus_matrix = zeros(50, 50)

    for i=1:50
        if i+6 <= 50
            c_plus_matrix[i,i+6] = 0.5 # +12
        else
            c_plus_matrix[i,36] = 0.5 # 70
        end

        if i-5 >= 1
            c_plus_matrix[i,i-5] = 0.5 # -10
        else
            c_plus_matrix[i,16] = 0.5 # 30
        end
    end

    c_plus_and_b_minus = ScoreTransitionMatrix(c_plus_matrix, RankInfo("C+/B-", 2, [2*i for i=0:49], 12, -10))

                      # 0  10  20  30  40  50  60  70  80  90
    b_matrix = Float64[ 0  0.5  0  0.5  0   0   0   0   0   0; #  0
                       0.5  0  0.5  0   0   0   0   0   0   0; # 10
                        0  0.5  0  0.5  0   0   0   0   0   0; # 20
                        0   0  0.5  0  0.5  0   0   0   0   0; # 30
                        0   0   0  0.5  0  0.5  0   0   0   0; # 40
                        0   0   0   0  0.5  0  0.5  0   0   0; # 50
                        0   0   0   0   0  0.5  0  0.5  0   0; # 60
                        0   0   0   0   0   0  0.5  0  0.5  0; # 70
                        0   0   0   0   0   0   0  0.5  0  0.5;# 80
                        0   0   0   0   0   0   0  0.5 0.5  0];# 90

    b_thru_a_plus = ScoreTransitionMatrix(b_matrix, RankInfo("B/B+/A-/A/A+", 5, [10*i for i=0:9], 10, -10))

    s_matrix = zeros(100, 100)

    for i=1:100
        if i <= 40
            reward = 5
        else
            reward = 4
        end
        if i <= 80
            penalty = 5
        else
            penalty = 6
        end

        if i+reward <= 100
            s_matrix[i,i+reward] = 0.5
        else
            s_matrix[i,71] = 0.5
        end

        if i-penalty >= 1
            s_matrix[i,i-penalty] = 0.5
        else
            s_matrix[i,31] = 0.5
        end
    end

    s = ScoreTransitionMatrix(s_matrix, RankInfo("S", 1, [i for i=0:99], 4, -5))

    s_plus_matrix = zeros(100, 100)

    for i=1:100
        if i <= 40
            reward = 4
        elseif i <= 80
            reward = 3
        else
            reward = 2
        end
        if i <= 40
            penalty = 4
        else
            penalty = 5
        end

        if i+reward <= 100
            s_plus_matrix[i,i+reward] = 0.5
        else
            s_plus_matrix[i,100] = 0.5
        end

        if i-penalty >= 1
            s_plus_matrix[i,i-penalty] = 0.5
        else
            s_plus_matrix[i,31] = 0.5
        end
    end

    s_plus = ScoreTransitionMatrix(s_plus_matrix, RankInfo("S+", 1, [i for i=0:99], 2, -4))

#                          # 0  10  20  30  40  50  60  70  80  90 100
#   a_plus_matrix = Float64[ 0  0.5  0  0.5  0   0   0   0   0   0   0; #  0
#                           0.5  0  0.5  0   0   0   0   0   0   0   0; # 10
#                            0  0.5  0  0.5  0   0   0   0   0   0   0; # 20
#                            0   0  0.5  0  0.5  0   0   0   0   0   0; # 30
#                            0   0   0  0.5  0  0.5  0   0   0   0   0; # 40
#                            0   0   0   0  0.5  0  0.5  0   0   0   0; # 50
#                            0   0   0   0   0  0.5  0  0.5  0   0   0; # 60
#                            0   0   0   0   0   0  0.5  0  0.5  0   0; # 70
#                            0   0   0   0   0   0   0  0.5  0  0.5  0;# 80
#                            0   0   0   0   0   0   0   0  0.5  0  0.5;# 90
#                            0   0   0   0   0   0   0   0   0  0.5 0.5];# 100

#   a_plus = ScoreTransitionMatrix(a_plus_matrix, RankInfo("A+", 1, [10*i for i in 0:10], 10, -10))

    return ScoreTransitionMatrix[ c_minus, c, c_plus_and_b_minus, b_thru_a_plus, s, s_plus ]
end

function proposed_score_transitions()
                            # 0  10  20  30  40  50  60  70  80  90
    c_minus_matrix = Float64[0.5  0  0.5  0   0   0   0   0   0   0; #  0
                             0.5  0   0  0.5  0   0   0   0   0   0; # 10
                              0  0.5  0   0  0.5  0   0   0   0   0; # 20
                              0   0  0.5  0   0  0.5  0   0   0   0; # 30
                              0   0   0  0.5  0   0  0.5  0   0   0; # 40
                              0   0   0   0  0.5  0   0  0.5  0   0; # 50
                              0   0   0   0   0  0.5  0   0  0.5  0; # 60
                              0   0   0   0   0   0  0.5  0   0  0.5;# 70
                              0   0   0   0   0   0   0  1.0  0   0; # 80
                              0   0   0   0   0   0   0  0.5 0.5  0 ]# 90

    c_minus = ScoreTransitionMatrix(c_minus_matrix, RankInfo("C-", 1, [10*i for i=0:9], 20, -10))

                      #  5  10   15  20  25  30  35  40  45  50  55  60  65  70  75  80  85  90  95
    c_matrix0= Float64[  0   0    0  0.5  0  0.5  0   0   0   0   0   0   0   0   0   0   0   0   0; #  5
                         0   0    0   0  0.5 0.5  0   0   0   0   0   0   0   0   0   0   0   0   0; # 10
                        0.5  0    0   0   0  0.5  0   0   0   0   0   0   0   0   0   0   0   0   0; # 15
                         0  0.5   0   0   0   0  0.5  0   0   0   0   0   0   0   0   0   0   0   0; # 20
                         0   0   0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0   0   0   0; # 25
                         0   0    0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0   0   0; # 30
                         0   0    0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0   0; # 35
                         0   0    0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0   0; # 40
                         0   0    0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0   0; # 45
                         0   0    0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0   0; # 50
                         0   0    0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0   0; # 55
                         0   0    0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0   0; # 60
                         0   0    0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0   0; # 65
                         0   0    0   0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0   0; # 70
                         0   0    0   0   0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5  0; # 75
                         0   0    0   0   0   0   0   0   0   0   0   0   0  0.5  0   0   0   0  0.5; # 80
                         0   0    0   0   0   0   0   0   0   0   0   0   0  0.5 0.5  0   0   0   0; # 85
                         0   0    0   0   0   0   0   0   0   0   0   0   0  0.5  0  0.5  0   0   0; # 90
                         0   0    0   0   0   0   0   0   0   0   0   0   0  0.5  0   0  0.5  0   0] # 95

    c = ScoreTransitionMatrix(c_matrix0, RankInfo("C", 1, [5*i for i=1:19], 15, -10))

    c_plus_matrix0 = zeros(49, 49)

    for i=1:49
        if i+6 <= 49
            c_plus_matrix0[i,i+6] = 0.5 # +12
        else
            c_plus_matrix0[i,35] = 0.5 # 70
        end

        if i-5 >= 1
            c_plus_matrix0[i,i-5] = 0.5 # -10
        else
            c_plus_matrix0[i,15] = 0.5 # 30
        end
    end

    c_plus_and_b_minus = ScoreTransitionMatrix(c_plus_matrix0, RankInfo("C+/B-", 2, [2*i for i=1:49], 12, -10))

                      # 10  20  30  40  50  60  70  80  90
    b_matrix0= Float64[  0  0.5 0.5  0   0   0   0   0   0; # 10
                        0.5  0  0.5  0   0   0   0   0   0; # 20
                         0  0.5  0  0.5  0   0   0   0   0; # 30
                         0   0  0.5  0  0.5  0   0   0   0; # 40
                         0   0   0  0.5  0  0.5  0   0   0; # 50
                         0   0   0   0  0.5  0  0.5  0   0; # 60
                         0   0   0   0   0  0.5  0  0.5  0; # 70
                         0   0   0   0   0   0  0.5  0  0.5;# 80
                         0   0   0   0   0   0  0.5 0.5  0];# 90

    b_thru_a_plus = ScoreTransitionMatrix(b_matrix0, RankInfo("B/B+/A-/A/A+", 5, [10*i for i=1:9], 10, -10))

    s_matrix0 = zeros(99, 99)

    for i=1:99
        if i <= 39
            reward = 5
        else
            reward = 4
        end
        if i <= 79
            penalty = 5
        else
            penalty = 6
        end

        if i+reward <= 99
            s_matrix0[i,i+reward] = 0.5
        else
            s_matrix0[i,70] = 0.5
        end

        if i-penalty >= 1
            s_matrix0[i,i-penalty] = 0.5
        else
            s_matrix0[i,30] = 0.5
        end
    end

    s = ScoreTransitionMatrix(s_matrix0, RankInfo("S", 1, [i for i=1:99], 4, -5))

    s_plus_matrix0 = zeros(99, 99)

    for i=1:99
        if i <= 39
            reward = 4
        elseif i <= 79
            reward = 3
        else
            reward = 2
        end
        if i <= 39
            penalty = 4
        else
            penalty = 5
        end

        if i+reward <= 99
            s_plus_matrix0[i,i+reward] = 0.5
        else
            s_plus_matrix0[i,99] = 0.5
        end

        if i-penalty >= 1
            s_plus_matrix0[i,i-penalty] = 0.5
        else
            s_plus_matrix0[i,30] = 0.5
        end
    end

    s_plus = ScoreTransitionMatrix(s_plus_matrix0, RankInfo("S+", 1, [i for i=1:99], 2, -4))
#                          # 10  20  30  40  50  60  70  80  90 100
#   a_plus_matrix0= Float64[  0  0.5 0.5  0   0   0   0   0   0   0; # 10
#                            0.5  0  0.5  0   0   0   0   0   0   0; # 20
#                             0  0.5  0  0.5  0   0   0   0   0   0; # 30
#                             0   0  0.5  0  0.5  0   0   0   0   0; # 40
#                             0   0   0  0.5  0  0.5  0   0   0   0; # 50
#                             0   0   0   0  0.5  0  0.5  0   0   0; # 60
#                             0   0   0   0   0  0.5  0  0.5  0   0; # 70
#                             0   0   0   0   0   0  0.5  0  0.5  0;# 80
#                             0   0   0   0   0   0   0  0.5  0  0.5;# 90
#                             0   0   0   0   0   0   0   0  0.5 0.5];# 100

#   a_plus = ScoreTransitionMatrix(a_plus_matrix0, RankInfo("A+", 1, [10*i for i in 1:10], 10, -10))

    return ScoreTransitionMatrix[ c_minus, c, c_plus_and_b_minus, b_thru_a_plus, s, s_plus ]
end

function solve_and_report(score_transitions::Array{ScoreTransitionMatrix})
    score_steady_states = [score_steady_state(score_transitions[i]) for i in 1:length(score_transitions)]
    report_score_steady_states(score_steady_states)

    rank_matrix = rank_matrix_from_steady_states(score_steady_states)

    println("Rank transition matrix: ", rank_matrix)

    rank_steady_state = steady_state(rank_matrix)

    report_rank_steady_state(rank_steady_state)
end

println("\n *** Current Splatoon algorithm (demote when score < 0) *** ")
solve_and_report(default_score_transitions())

println("\n *** Proposed fix (demote when score = 0) *** ")
solve_and_report(proposed_score_transitions())


@defcomp LGforcing begin

    c_LGconcentration = Parameter(index=[time], unit="ppbv")
    c0_LGconcbaseyr = Parameter(unit="ppbv", default=0.11)
    f_LGforcing = Variable(index=[time], unit="W/m2")
    f0_LGforcingbase = Parameter(unit="W/m2", default=0.022)
    fslope_LGforcingslope = Parameter(unit="W/m2", default=0.2)

    function run_timestep(p, v, d, t)

        # eq.13 in Hope 2006
        v.f_LGforcing[t] = p.f0_LGforcingbase + p.fslope_LGforcingslope * (p.c_LGconcentration[t] - p.c0_LGconcbaseyr)

    end
end


function mccormick_envelope(model::Model, w::Variable, u::Variable, v::Variable, urange::ClosedInterval, vrange::ClosedInterval)
    # convex approximation of w == u * v on u ∈ [umin, umax], v ∈ [vmin, vmax]
    umin, umax = minimum(urange), maximum(urange)
    vmin, vmax = minimum(vrange), maximum(vrange)
    vertices = [umin vmin umin * vmin; umax vmax umax * vmax; umax vmin umax * vmin; umin vmax umin * vmax]
    SimpleHRepresentation(hrep(polyhedron(SimpleVRepresentation(vertices), CDDLibrary())))
end

function piecewise_mccormick_envelope_constraints(model::Model, u::Variable, v::Variable, w::Variable,
    urange::Range, vrange::ClosedInterval)
    vars = [u; v; w]
    if length(urange) == 2
        envelope = mccormick_envelope(model, u, v, w, first(urange)..last(urange), vrange)
        polyhedron_constraints(model, envelope, vars)
        nothing
    else
        N = length(urange) - 1
        envelopes = Vector{SimpleHRepresentation}(N)
        for i = 1 : N
            envelopes[i] = mccormick_envelope(model, u, v, w, urange[i]..urange[i + 1], vrange)
        end
        hull_reformulation(model, envelopes, vars)
        nothing
    end
end


using HarmonicOrthogonalPolynomials, Plot
plotly()

G = HarmonicOrthogonalPolynomials.grid(SphericalHarmonic()[:,Block.(Base.OneTo(10))])
n,m = size(G)
g = append!(vec(G), [SphericalCoordinate(0,0), SphericalCoordinate(π,0)])
x,y,z = ntuple(k ->getindex.(g,k), 3);
N = length(g)

i = [range(0; length=m, step=n); ]
j = [range(n; length=m-1, step=n); 0]
k = fill(N-2, m)

for ν = range(0; length=m-1, step=n)
    append!(i, range(ν; length=n-1))
    append!(i, range(ν+n; length=n-1))
    append!(j, range(ν+1; length=n-1))
    append!(j, range(ν+1; length=n-1))
    append!(k, range(ν+n; length=n-1))
    append!(k, range(ν+n+1; length=n-1))
end

append!(i, range(n*(m-1); length=n-1))
append!(i, range(0; length=n-1))
append!(j, range(n*(m-1)+1; length=n-1))
append!(j, range(n*(m-1)+1; length=n-1))
append!(k, range(0; length=n-1))
append!(k, range(1; length=n-1))

append!(i, [range(2n-1; length=m-1, step=n); n-1])
append!(j, range(n-1; length=m, step=n))
append!(k, fill(N-1, m))

mesh3d(x,y,z; connections = (i,j,k))


"""
    Code for a super-Gaussian mask shape for use with CCF
Author: Eric Ford
Created: September 2020
"""

using QuadGK

"""   SuperGaussianCCFMask
A truncated Gaussian mask with two parameters, it's stdandard deviation and where to truncate it, both as a velocity in m/s.
Mask weights are stored separately in a line list.

TODO: Repalce Gaussian with super-Gaussian
Warning:  Not implemented/tested yet.
"""
struct SuperGaussianCCFMask <: AbstractCCFMaskShape
    σ_sqrt2::Float64
    power::Float64
    half_width_truncation::Float64
    normalization::Float64

    """ `SuperGaussianCCFMask( σ ; half_truncation_width_in_σ=2 )` """
    function SuperGaussianCCFMask(σ::Real, p::Real, w::Real=2 )
        @assert 0 < σ <= 300000   # 300km/s is arbitrary choice for an upper limit
        @assert 1 <= p <= 2
        @assert 0 < w <= 4
        function integrand(x)
            exp(-(0.5*(x/σ)^2)^p)
        end
        integral = quadgk(integrand, -w*σ, w*σ)[1]
        norm = 1.0/integral
        new(σ*sqrt(2.0),p,σ*w,norm)
    end
end

""" `SuperGaussianCCFMask( inst ; scale_factor )` """
function SuperGaussianCCFMask(inst::InstT; power::Real = default_supergaussian_ccf_exponent, fwhm::Real = default_supergaussian_ccf_fwhm,
        σ_scale_factor::Real = 1, truncation_scale_factor::Real = default_supergaussian_ccf_truncation_scale_factor ) where { InstT<:AbstractInstrument }
    # TODO: Update default values
    σ = σ_scale_factor * fwhm/sqrt(8 * log(2)^(1/power))  # From Ryan Petersburg email 9/11/2020
    w = truncation_scale_factor
    SuperGaussianCCFMask(σ,power,w/2)
end

λ_min(m::SuperGaussianCCFMask,λ::Real) = λ/calc_doppler_factor(m.half_width_truncation)
λ_max(m::SuperGaussianCCFMask,λ::Real) = λ*calc_doppler_factor(m.half_width_truncation)

function integrate(m::SuperGaussianCCFMask, v_lo::Real,v_hi::Real)
    quadgk(m, v_lo, v_hi, atol=1e-4)[1]
end

""" Functor for returning PSF for Δv <= half_width.  """
function (m::SuperGaussianCCFMask)(Δv::Real)
    if abs2(Δv) > abs2(m.half_width_truncation)
        return zero(Δv)
    else
        return m.normalization*exp(-((Δv/m.σ_sqrt2)^2)^m.power)
    end
end

function mask_with_increased_fwhm(m::SuperGaussianCCFMask, Δfwhm::Real )
    # WARNING: Haven't checked how good of an approximation this is
    σ = m.σ_sqrt2/sqrt(2)
    fwhm_orig = σ * sqrt(8 * log(2)^(1/m.power))
    fwhm_new = sqrt(fwhm_orig^2 + Δfwhm^2)
    σ_new = fwhm_new/sqrt(8 * log(2)^(1/m.power))
    return SuperGaussianCCFMask(σ_new, m.power, m.half_width_truncation * (σ_new * sqrt(2) / σ_sqrt2) )
end


using KernelFunctions
using MLDataUtils
using Zygote
using Flux
using Distributions, LinearAlgebra
using Plots

Flux.@functor SqExponentialKernel
Flux.@functor ScaleTransform
Flux.@functor KernelSum
Flux.@functor Matern32Kernel

xmin = -3; xmax = 3
x = range(xmin,xmax,length=100)
x_test = range(xmin,xmax,length=300)
x,y = noisy_function(sinc,x;noise=0.1)
X = reshape(x,:,1)
X_test = reshape(x_test,:,1)
k = SqExponentialKernel(1.0)#+Matern32Kernel(2.0)
λ = [-1.0]
f(x,k,λ) = kernelmatrix(k,x,X,obsdim=1)*inv(kernelmatrix(k,X,obsdim=1)+exp(λ[1])*I)*y
f(X,k,1.0)
loss(k,λ) = f(X,k,λ) |> ŷ ->sum(y-ŷ)/length(y)+exp(λ[1])*norm(ŷ)
loss(k,λ)
ps = Flux.params(k)
push!(ps,λ)
opt = Flux.Momentum(0.1)
##
for i in 1:10
    grads = Zygote.gradient(()->loss(k,λ),ps)
    Flux.Optimise.update!(opt,ps,grads)
    p = Plots.scatter(x,y,lab="data",title="Loss = $(loss(k,λ))")
    Plots.plot!(x_test,f(X_test,k,λ),lab="Prediction",lw=3.0)
    display(p)
end


using ForwardDiff
using Roots
using StatsBase
using Plots
using Measures

include("helper_functions.jl")

begin
    λ1 = 1
    λ2 = 1.5
    α = 1
    β = 1
    μ(α) = 1 / (1.5 * (α))
    ν(α) = 1 / (0.7 * (α))
    parameters = (λ1, λ2, α, μ, ν)
    he_parameters = (λ1, λ2, α, μ, ν, β)
    opponent_parameters = (λ2, λ1, α, μ, ν)
    opponent_he_parameters = (λ2, λ1, α, μ, ν, β)
    qs = 0:0.001:1
end

plot(qs, he_pstar.(qs, he_parameters...))
plot!(qs, pstar.(qs, parameters...))

total_cost(he_equilibrium_probability(he_parameters...)..., parameters...)
total_cost(equilibrium_probability(parameters...)..., parameters...)

plot(qs, he_cost.(he_pstar.(qs, he_parameters...), qs, he_parameters...), ylims = (0, 50))
plot!(qs, cost.(pstar.(qs, parameters...), qs, parameters...))
plot!(qs, he_cost.(pstar.(qs, parameters...), qs, parameters..., 1))

iter = 0.001:0.001:2

begin
	probability_tuples = equilibrium_probability.(λ1, λ2, iter, μ, ν)
    he_probability_tuples = he_equilibrium_probability.(λ1, λ2, iter, μ, ν, β)
	response_probs = first.(probability_tuples)
	response_qrobs = last.(probability_tuples)
    he_response_probs = first.(he_probability_tuples)
	he_response_qrobs = last.(he_probability_tuples)

    plot(title = "Traffic at each server for varying α, λ = 1, and λ2 = $(dp_string(λ2, 3S))",
        legend = :outerright,
        size = (800, 600))
    plot!(iter, response_probs, label = "p⋆", lc = :blue, lw = 1.2)
    plot!(iter, response_qrobs, label = "q⋆", lc = :red, lw = 1.2)
    plot!(iter, he_response_probs, label = "he_p⋆", ls = :dash, lc = :blue, lw = 1.2)
    plot!(iter, he_response_qrobs, label = "he_q⋆", ls = :dash, lc = :red, lw = 1.2)
    plot!(ylims = (0,1), xlims = (0,1.76))
    plot!(xlabel = "α", ylabel = "probability")
    plot!(margin = 2mm)
end


module Multroot

export multroot

using ..Polynomials
using LinearAlgebra

"""
    multroot(p; verbose=false, kwargs...)

Use `multroot` algorithm of
[Zeng](https://www.ams.org/journals/mcom/2005-74-250/S0025-5718-04-01692-8/S0025-5718-04-01692-8.pdf)
to identify roots of polynomials with suspected multiplicities over
`Float64` values, typically.

Example:

```
julia> using Polynomials

julia> p = fromroots([sqrt(2), sqrt(2), sqrt(2), 1, 1])
Polynomial(-2.8284271247461907 + 11.656854249492383*x - 19.07106781186548*x^2 + 15.485281374238573*x^3 - 6.242640687119286*x^4 + 1.0*x^5)

julia> roots(p)
5-element Array{Complex{Float64},1}:
  0.999999677417768 + 0.0im
 1.0000003225831504 + 0.0im
 1.4141705716005981 + 0.0im
 1.4142350577588885 - 3.72273772278647e-5im
 1.4142350577588885 + 3.72273772278647e-5im

julia> Polynomials.Multroot.multroot(p)
(values = [0.9999999999999993, 1.4142135623730958], multiplicities = [2, 3], κ = 5.218455674370636, ϵ = 2.8736226244218195e-16)
```

The algorithm has two stages. First it uses `pejorative_manifold` to
identify the number of distinct roots and their multiplicities. This
uses the fact if `p=Π(x-zᵢ)ˡⁱ`, `u=gcd(p, p′)`, and `u⋅v=p` then
`v=Π(x-zi)` is square free and contains the roots of `p` and `u` holds
the multiplicity details, which are identified by recursive usage of
`ncgd`, which identifies `u` and `v` above even if numeric
uncertainties are present.

Second it uses `pejorative_root` to improve a set of initial guesses
for the roots under the assumption the multiplicity structure is
correct using a Newton iteration scheme.

The following tolerances, passed through to `pejorative_manifold` by
`kwargs...`, are all used in the first stage, to identify the
multiplicity structure:

* `θ`: the singular value threshold, set to `1e-8`. This is used by
  `Polynomials.ngcd` to assess if a matrix is rank deficient by
  comparing the smallest singular value to `θ ⋅ ||p||₂`.

* `ρ`: the initial residual tolerance, set to `1e-10`. This is passed
  to `Polynomials.ngcd`, the GCD finding algorithm as a relative tolerance.

* `ϕ`: A scale factor, set to `100`. As the `ngcd` algorithm is called
  recursively, this allows the residual tolerance to scale up to match
  increasing numeric errors.

Returns a named tuple with

* `values`: the identified roots
* `multiplicities`: the corresponding multiplicities
* `κ`: the estimated condition number
* `ϵ`: the backward error, `||p̃ - p||_w`.

If the wrong multiplicity structure is identified in step 1, then
typically either `κ` or `ϵ` will be large. The estimated forward
error, `||z̃ - z||₂`, is bounded up to higher order terms by `κ ⋅ ϵ`.
This will be displayed if `verbose=true` is specified.

For polynomials of degree 20 or higher, it is often the case the `l`
is misidentified.

"""
function multroot(p::Polynomials.StandardBasisPolynomial{T}; verbose=false,
                  kwargs...) where {T}

    # degenerate case, constant
    degree(p) == 0 && return (values=T[], multiplicities=Int[], κ=NaN, ϵ=NaN)

    # degenerate case, all zeros
    if (nz = findfirst(!iszero, coeffs(p))) == length(coeffs(p))
        return (values=zeros(T,1), multiplicities=[nz-1], κ=NaN, ϵ=NaN)
    end

    z, l = pejorative_manifold(p; kwargs...)
    z̃ = pejorative_root(p, z, l)
    κ, ϵ = stats(p, z̃, l)

    verbose && show_stats(κ, ϵ)
    (values = z̃, multiplicities = l, κ = κ, ϵ = ϵ)

end

# The multiplicity structure, `l`, gives rise to a pejorative manifold `Πₗ = {Gₗ(z) | z∈ Cᵐ}`.
function pejorative_manifold(p::Polynomials.StandardBasisPolynomial{T};
                        ρ = 1e-10, # initial residual tolerance
                        θ = 1e-8,  # zero singular-value threshold
                        ϕ = 100)  where {T}

    zT = zero(float(real(T)))
    u, v, w, θ′, κ = Polynomials.ngcd(p, derivative(p),
                                       atol=ρ*norm(p), satol = θ*norm(p),
                                       rtol = zT, srtol = zT)
    zs = roots(v)
    nrts = length(zs)
    ls = ones(Int, nrts)

    while !Polynomials.isconstant(u)

        normp = 1 + norm(u, 2)
        ρ′ = max(ρ*normp, ϕ * θ′)  # paper uses just latter
        u, v, w, θ′, κ = Polynomials.ngcd(u, derivative(u),
                                          atol=ρ′, satol = θ*normp,
                                          rtol = zT, srtol = zT,
                                          minⱼ = degree(u) - nrts # truncate search
                                          )

        rts = roots(v)
        nrts = length(rts)

        for z in rts
            tmp, ind = findmin(abs.(zs .- z))
            ls[ind] = ls[ind] + 1
        end

    end

    zs, ls # estimate for roots, multiplicities

end

"""
    pejorative_root(p, zs, ls; kwargs...)

Find a *pejorative* *root* for `p` given multiplicity structure `ls` and initial guess `zs`.

The pejorative manifold for a multplicity structure `l` is denoted `{Gₗ(z) | z ∈ Cᵐ}`. A pejorative
root is a least squares minimizer of `F(z) = W ⋅ [Gₗ(z) - a]`. Here `a ~ (p_{n-1}, p_{n-2}, …, p_1, p_0) / p_n` and `W` are weights given by `min(1, 1/|aᵢ|)`. When `l` is the mathematically correct structure, then `F` will be `0` for a pejorative root. If `l` is not correct, then the backward error `||p̃ - p||_w` is typically large, where `p̃ = Π (x-z̃ᵢ)ˡⁱ`.

This follows Algorithm 1 of [Zeng](https://www.ams.org/journals/mcom/2005-74-250/S0025-5718-04-01692-8/S0025-5718-04-01692-8.pdf)
"""
function pejorative_root(p::Polynomials.StandardBasisPolynomial,
                         zs::Vector{S}, ls::Vector{Int}; kwargs...) where {S}
    ps = (reverse ∘ coeffs)(p)
    pejorative_root(ps, zs, ls; kwargs...)
end

# p is [1, a_{n-1}, a_{n-2}, ..., a_1, a_0]
function pejorative_root(p, zs::Vector{S}, ls::Vector{Int};
                 τ = eps(float(real(S)))^(2/3),
                 maxsteps = 100, # linear or quadratic
                 verbose=false
                 ) where {S}

    ## Solve WJ Δz = W(Gl(z) - a)
    ## using weights min(1/|aᵢ|), i ≠ 1

    m,n = sum(ls), length(zs)

    # storage
    a = p[2:end]./p[1]     # a ~ (p[n-1], p[n-2], ..., p[0])/p[n]
    W = Diagonal([min(1, 1/abs(aᵢ)) for aᵢ in a])
    J = zeros(S, m, n)
    G = zeros(S,  1 + m)
    Δₖ = zeros(S, n)
    zₖs = copy(zs)  # updated

    cvg = false
    δₖ₀ = -Inf
    for ctr in 1:maxsteps

        evalJ!(J, zₖs, ls)
        evalG!(G, zₖs, ls)
        Δₖ .= (W*J) \ (W*(view(G, 2:1+m) .- a)) # weighted least squares

        δₖ₁ = norm(Δₖ, 2)
        Δ = δₖ₀ - δₖ₁

        if ctr > 10 && δₖ₁ >= δₖ₀
            δₖ₀ < δₖ₁ && @warn "Increasing Δ, terminating search"
            cvg = true
            break
        end

        zₖs .-= Δₖ

        if δₖ₁^2 <= (δₖ₀ - δₖ₁) * τ
            cvg = true
            break
        end

        δₖ₀ = δₖ₁
    end
    verbose && show_stats(stats(p, zₖs, ls)...)

    if cvg
        return zₖs
    else
        @info ("""
The multiplicity count may be in error: the initial guess for the roots failed
to converge to a pejorative root.
""")
        return(zₖs)
    end

end

# Helper: compute Π (x-zᵢ)^lᵢ directly
function evalG!(G, zs::Vector{T}, ls::Vector) where {T}

    G .= zero(T)
    G[1] = one(T)

    ctr = 1
    for (z,l) in zip(zs, ls)
        for _ in 1:l
            for j in length(G)-1:-1:1#ctr
                G[1 + j] -= z * G[j]
            end
            ctr += 1
        end
    end

    nothing
end

# For eachcolumn, computes (3.8)
# -lⱼ * (x - xⱼ)^(lⱼ-1) * Π_{k ≠ j} (x - zₖ)^lₖ
function evalJ!(J, zs::Vector{T}, ls::Vector) where {T}

    J .= 0
    n, m = sum(ls), length(zs)

    evalG!(view(J, 1:1+n-m, 1), zs, ls .- 1)
    for (j′, lⱼ) in enumerate(reverse(ls))
        for i in 1+n-m:-1:1
            J[i, m - j′ + 1] = -lⱼ * J[i, 1]
        end
    end

    for  (j, lⱼ) in enumerate(ls)
        for (k, zₖ) in enumerate(zs)
            k == j && continue
            for i in n-1:-1:1
                J[1+i,j] -= zₖ * J[i,j]
            end
        end
    end
    nothing
end

# Defn (3.5) condition number of z with respect to the multiplicity l and weight W
cond_zl(p::AbstractPolynomial, zs::Vector{S}, ls) where {S}  = cond_zl(reverse(coeffs(p)), zs, ls)
function cond_zl(p, zs::Vector{S}, ls) where {S}
    J = zeros(S, sum(ls), length(zs))
    W = diagm(0 => [min(1, 1/abs(aᵢ)) for aᵢ in p[2:end]])
    evalJ!(J, zs, ls)
    F = qr(W*J)
    _, σ, _ = Polynomials.NGCD.smallest_singular_value(F.R)
    1 / σ
end

backward_error(p::AbstractPolynomial, zs::Vector{S}, ls) where {S}  =
    backward_error(reverse(coeffs(p)), zs, ls)
function backward_error(p, z̃s::Vector{S}, ls) where {S}
    # || ã - a ||w
    G = zeros(S,  1 + sum(ls))
    evalG!(G, z̃s, ls)
    a = p[2:end]./p[1]
    W = Diagonal([min(1, 1/abs(aᵢ)) for aᵢ in a])
    u = G[2:end] .- a
    norm(W*u,2)
end

function stats(p, zs, ls)
    cond_zl(p, zs, ls), backward_error(p, zs, ls)
end

function show_stats(κ, ϵ)
    println("")
    println("Condition number κ(z; l) = ", κ)
    println("Backward error ||ã - a||w = ", ϵ)
    println("Estimated forward root error ||z̃ - z||₂ = ", κ * ϵ)
    println("")
end

end


# This file is a part of AstroLib.jl. License is MIT "Expat".

function precess_cd(cd::AbstractMatrix{T}, epoch1::T, epoch2::T, crval_old::AbstractVector{T},
                    crval_new::AbstractVector{T}, FK4::Bool) where {T<:AbstractFloat}
    t = (epoch2 - epoch1) / 1000

    if FK4
        st = (epoch1 - 1900) / 1000
        c = sec2rad(t*(20046.85 - st*(85.33 + st*0.37) + t*(-42.67 - st*0.37 - t*41.8)))
    else
        st = (epoch1 - 2000) / 1000
        c = sec2rad(t*(20043.109 - st*(85.33 + st*0.217) + t*(-42.665 - st*0.217 - t*41.8)))
    end
    pole_ra = zero(T)
    pole_dec = T(90)

    if (epoch1 == 2000 && epoch2 == 1950) || (epoch1 == 1950 && epoch2 == 2000)
        pole_ra, pole_dec = bprecess(pole_ra, pole_dec)
    else
        pole_ra, pole_dec = precess(pole_ra, pole_dec, epoch1, epoch2, FK4=FK4)
    end
    sind1, cosd1 = sincos(deg2rad(crval_old[2]))
    sind2, cosd2 = sincos(deg2rad(crval_new[2]))
    sinra = sind(crval_new[1] - pole_ra)
    sin_c, cos_c = sincos(c)
    cosfi = (cos_c - sind1 * sind2) / (cosd1 * cosd2)
    sinfi = (abs(sin_c) * sinra) / cosd1
    return cd * SMatrix{2,2}(cosfi, -sinfi, sinfi, cosfi)
end

"""
    precess_cd(cd, epoch1, epoch2, crval_old, crval_new[, FK4=true]) -> cd

### Purpose ###

Precess the coordinate description matrix.

### Explanation ###

The coordinate matrix is precessed from epoch1 to epoch2.

### Arguments ###

* `cd`: 2 x 2 coordinate description matrix in degrees
* `epoch1`: original equinox of coordinates, scalar
* `epoch2`: equinox of precessed coordinates, scalar
* `crval_old`: 2 element vector containing right ascension and declination
  in degrees of the reference pixel in the original equinox
* `crval_new`: 2 element vector giving crval in the new equinox
* `FK4` (optional boolean keyword): if this keyword is set to `true`,
  then the precession constants are taken in the FK4 reference frame. When it
  is `false`, the default is the FK5 frame

### Output ###

* `cd`: coordinate description containing precessed values

### Example ###

```jldoctest
julia> using AstroLib

julia> precess_cd([20 60; 45 45], 1950, 2000, [34, 58], [12, 83])
2×2 Array{Float64,2}:
  48.8944  147.075
 110.188   110.365
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
This function should not be used for values more than 2.5 centuries from the year 1900.
This function calls [`sec2rad`](@ref), [`precess`](@ref) and [`bprecess`](@ref).
"""
precess_cd(cd::AbstractMatrix{<:Real}, epoch1::Real, epoch2::Real, crval_old::AbstractVector{<:Real},
           crval_new::AbstractVector{<:Real}, FK4::Bool=false) =
               precess_cd(float(cd), promote(float(epoch1), float(epoch2))...,
                          float(crval_old), float(crval_new), FK4)


using Test
using MeshArrays

@testset "MeshArrays tests:" begin
for nTopo=1:3
    if nTopo==1; grTopo="cs"; nFaces=6; N=200;
    elseif nTopo==2; grTopo="llc"; nFaces=5; N=200;
    elseif nTopo==3; grTopo="ll"; nFaces=1; N=500;
    end;
    Npt=nFaces*N*N
    GridVariables=GridOfOnes(grTopo,nFaces,N)
    mygrid=GridVariables["XC"].grid
    @test mygrid.class == grTopo
    Rini= 0.; Rend= 0.;
    (Rini,Rend,DXCsm,DYCsm)=MeshArrays.demo2(GridVariables);
    @test isa(Rend,MeshArray)
    @test sum(isfinite.(Rend)) == Npt
    Sini=sqrt(sum(Rini*Rini)/(Npt-1.0))
    Send=sqrt(sum(Rend*Rend)/(Npt-1.0))
    #println([Sini Send])
    @test isapprox(Sini,1.000; atol=1e-2)
    @test isapprox(Send,0.093; atol=1e-2)
end;
end;


module TestSphere
    using Evolutionary
    using Base.Test

    function test_result(result::Vector, fitness::Float64, N::Int, tol::Float64)
        @test length(result) == N
        @test_approx_eq_eps fitness 0.0 tol
    end

    # Objective function
    sphere{T <: Vector}(x::T) = sum(x.*x)

    # Parameters
    N = 30
    P = 25
    terminate(strategy) = strategy[:σ] < 1e-10
    initial = ones(N)

    # Testing: (μ/μ_I, λ)-σ-Self-Adaptation-ES
    # with isotropic mutation operator y' := y + σ(N_1(0, 1), ..., N_N(0, 1))
    result, fitness, cnt = es(sphere, N;
        initStrategy = strategy(σ = 1.0, τ = 1/sqrt(2*N)),
        recombination = average, srecombination = averageSigma1,
        mutation = isotropic, smutation = isotropicSigma,
        termination = terminate, selection=:comma,
        μ = 3, λ = P, iterations = 1000)
    println("(3/3,$(P))-σ-SA-ES => F: $(fitness), C: $(cnt)")
    test_result(result, fitness, N, 1e-5)

    # Testing: GA
    result, fitness, cnt =
        ga( sphere, N;
            populationSize = 4P,
            mutationRate = 0.05,
            ɛ = 0.1,
            selection = sus,
            crossover = intermediate(0.25),
            mutation = domainrange(fill(0.5,N)),
            tol = 1e-5, tolIter = 15)
    println("GA(pop=$(P),xover=0.8,μ=0.1) => F: $(fitness), C: $(cnt)")
    test_result(result, fitness, N, 1e-2)

end


# Observer tests

# Test constructors
@test size(LinearObserver(ones(2,2)).H) == (2,2)
@test size(LinearObserver(1.0).H) == (1,1)
@test size(LinearObserver(1.0, 1.0).R) == (1,1)
@test size(LinearObserver(eye(3), eye(3)).H) == (3,3)
@test size(LinearObserver(ones(3)).R) == (3,3)
@test size(LinearObserver(ones(3), eye(3)).H) == (3,1)
@test size(LinearObserver(ones(3)').R) == (1,1)
@test size(LinearObserver(ones(3)', 1.0).R) == (1,1)
@test size(LinearObserver(ones(3)', ones(1,1)').R) == (1,1)
@test size(LinearObserver(ones(3)', 0.1).H) == (1,3)

# Test compatibility methods
lin_obs = LinearObserver(ones(2,2), eye(2))
StateEstimation.assert_compatibility(lin_obs, DiscreteState(ones(2)))
StateEstimation.assert_compatibility(DiscreteState(ones(2)), lin_obs)
@test_throws DimensionMismatch StateEstimation.assert_compatibility(lin_obs,
    DiscreteState(ones(3)))
lin_obs = LinearObserver(ones(2)', 2.0)
StateEstimation.assert_compatibility(DiscreteState([1.0]), lin_obs)
@test_throws DimensionMismatch StateEstimation.assert_compatibility(
    DiscreteState(ones(2)), lin_obs)

# Test discrete prediction methods
lin_obs = LinearObserver(ones(2,2), eye(2))
@test predict(DiscreteState(ones(2)), lin_obs) == DiscreteState(2*ones(2), 0)
@test predict(UncertainDiscreteState(ones(2), eye(2)), lin_obs) ==
    UncertainDiscreteState(2*ones(2), lin_obs.H*lin_obs.H'+lin_obs.R, 0)

# Test continuous prediction methods
lin_obs = LinearObserver(ones(2)', 1.0)
@test predict(ContinuousState(ones(2)), lin_obs) == ContinuousState(2.0, 0.0)
@test predict(UncertainContinuousState(ones(2), eye(2)), lin_obs) ==
    UncertainContinuousState(2.0*ones(1), lin_obs.H*lin_obs.H'+lin_obs.R, 0.0)

# Test observability methods
linear_sys = LinearSystem(Float64[[1, -2] [-3, -4]]', 0.1*eye(2))
linear_obs = LinearObserver(reshape(Float64[1, 2], 1, 2), reshape([0.1], 1, 1))
@test observable(linear_sys, linear_obs) == false
@test observable(eye(2), eye(2))

# Test simulation methods
lin_sys = LinearSystem([[0.0, 1.0] [-1.0, 1.0]], 2*eye(2))
lin_obs = LinearObserver(ones(2)', 1.0)
continuous_state = ContinuousState(ones(2))
uncertain_discrete_state = UncertainDiscreteState(ones(2), eye(2))
@test typeof(simulate(continuous_state, lin_obs)) <: ContinuousState
@test typeof(simulate(uncertain_discrete_state, lin_obs)) <: DiscreteState
sim_state, sim_meas = simulate(continuous_state, lin_sys, lin_obs, 1.0)
@test typeof(sim_state) <: ContinuousState
@test typeof(sim_meas) <: ContinuousState

# Test nonlinear constructors
discrete_nl_fcn(t, x::Vector) = x.^2;
discrete_nl_jac(t, x::Vector) = diagm(2*x)
NonlinearObserver(discrete_nl_fcn, discrete_nl_jac, eye(2))
NonlinearObserver(discrete_nl_fcn, discrete_nl_jac, 1.0)
NonlinearObserver(LinearObserver(ones(2,2), eye(2)))

# Test nonlinear absolute state prediction methods
nonlin_obs = NonlinearObserver(discrete_nl_fcn, discrete_nl_jac, eye(3))
@test predict(DiscreteState([0.0, 1.0, 2.0]), nonlin_obs) ==
    DiscreteState([0.0, 1.0, 4.0], 0)
@test predict(ContinuousState([0.0, 2.0, 4.0], 1.0), nonlin_obs) ==
    ContinuousState([0.0, 4.0, 16.0], 1.0)

# Test nonlinear uncertain state prediction methods
@test predict(UncertainDiscreteState([0.0, 1.0, 2.0], eye(3)), nonlin_obs) ==
    UncertainDiscreteState([0.0, 1.0, 4.0], diagm([1.0, 5.0, 17.0]), 0)
@test predict(UncertainContinuousState([0.0, 1.0, 2.0], eye(3)), nonlin_obs) ==
    UncertainContinuousState([0.0, 1.0, 4.0], diagm([1.0, 5.0, 17.0]), 0.0)


"""
"""
forwardinterp

function forwardinterp(bc, n)
    opn = _forwardinterp.(bc, n)
    eye = Eye{Bool}.(n)
    _forwardinterp(opn, eye)
end

_forwardinterp(::Dirichlet, n::Int) =
    Bidiagonal(ones(n) / 2, ones(n-1) / 2, :U)

_forwardinterp(::Periodic, n::Int) =
    spdiagm(1-n => ones(1) / 2,
            0 => ones(n) / 2,
            1 => ones(n-1) / 2)

_forwardinterp(opn::NTuple{1}, eye::NTuple{1}) = opn

_forwardinterp(opn::NTuple{2,Any}, eye::NTuple{2,Any}) =
    ApplyArray(kron, eye[2], opn[1]),
    ApplyArray(kron, opn[2], eye[1])

_forwardinterp(opn::NTuple{3,Any}, eye::NTuple{3,Any}) =
    ApplyArray(kron, eye[3], eye[2], opn[1]),
    ApplyArray(kron, eye[3], opn[2], eye[1]),
    ApplyArray(kron, opn[3], eye[2], eye[1])

"""
"""
backwardinterp

function backwardinterp(bc, n)
    opn = _backwardinterp.(bc, n)
    eye = Eye{Bool}.(n)
    _backwardinterp(opn, eye)
end

_backwardinterp(::Dirichlet, n::Int) =
    Bidiagonal(ones(n) / 2, ones(n-1) / 2, :L)

_backwardinterp(::Periodic, n::Int) =
    spdiagm(-1 => ones(n-1) / 2,
            0 => ones(n) / 2,
            n - 1 => ones(1) / 2)

_backwardinterp(opn::NTuple{1}, eye::NTuple{1}) = opn

_backwardinterp(opn::NTuple{2,Any}, eye::NTuple{2,Any}) =
    ApplyArray(kron, eye[2], opn[1]),
    ApplyArray(kron, opn[2], eye[1])

_backwardinterp(opn::NTuple{3,Any}, eye::NTuple{3,Any}) =
    ApplyArray(kron, eye[3], eye[2], opn[1]),
    ApplyArray(kron, eye[3], opn[2], eye[1]),
    ApplyArray(kron, opn[3], eye[2], eye[1])


"""
    ProjectedPointDistribution(M::Manifold, d, proj!, p)

Generates a random point in ambient space of `M` and projects it to `M`
using function `proj!`. Generated arrays are of type `TResult`, which can be
specified by providing the `p` argument.
"""
struct ProjectedPointDistribution{TResult,TM<:Manifold,TD<:Distribution,TProj} <:
       MPointDistribution{TM}
    manifold::TM
    distribution::TD
    proj!::TProj
end

function ProjectedPointDistribution(
    M::Manifold,
    d::Distribution,
    proj!,
    ::TResult,
) where {TResult}
    return ProjectedPointDistribution{TResult,typeof(M),typeof(d),typeof(proj!)}(
        M,
        d,
        proj!,
    )
end

"""
    projected_distribution(M::Manifold, d, [p=rand(d)])

Wrap the standard distribution `d` into a manifold-valued distribution. Generated
points will be of similar type to `p`. By default, the type is not changed.
"""
function projected_distribution(M::Manifold, d, p=rand(d))
    return ProjectedPointDistribution(M, d, project!, p)
end

function Random.rand(
    rng::AbstractRNG,
    d::ProjectedPointDistribution{TResult},
) where {TResult}
    p = convert(TResult, rand(rng, d.distribution))
    return d.proj!(d.manifold, p, p)
end

function Distributions._rand!(
    rng::AbstractRNG,
    d::ProjectedPointDistribution,
    p::AbstractArray{<:Number},
)
    Distributions._rand!(rng, d.distribution, p)
    return d.proj!(d.manifold, p, p)
end

Distributions.support(d::ProjectedPointDistribution) = MPointSupport(d.manifold)

"""
    ProjectedFVectorDistribution(type::VectorBundleFibers, p, d, project!)

Generates a random vector from ambient space of manifold `type.manifold`
at point `p` and projects it to vector space of type `type` using function
`project!`, see [`project`](@ref) for documentation.
Generated arrays are of type `TResult`.
"""
struct ProjectedFVectorDistribution{
    TResult,
    TSpace<:VectorBundleFibers,
    ManifoldPoint,
    TD<:Distribution,
    TProj,
} <: FVectorDistribution{TSpace,ManifoldPoint}
    type::TSpace
    point::ManifoldPoint
    distribution::TD
    project!::TProj
end

function ProjectedFVectorDistribution(
    type::VectorBundleFibers,
    p,
    d::Distribution,
    project!,
    ::TResult,
) where {TResult}
    return ProjectedFVectorDistribution{
        TResult,
        typeof(type),
        typeof(p),
        typeof(d),
        typeof(project!),
    }(
        type,
        p,
        d,
        project!,
    )
end

function Random.rand(
    rng::AbstractRNG,
    d::ProjectedFVectorDistribution{TResult},
) where {TResult}
    X = convert(TResult, reshape(rand(rng, d.distribution), size(d.point)))
    return d.project!(d.type, X, d.point, X)
end

function Distributions._rand!(
    rng::AbstractRNG,
    d::ProjectedFVectorDistribution,
    X::AbstractArray{<:Number},
)
    # calling _rand!(rng, d.d, v) doesn't work for all arrays types
    return copyto!(X, rand(rng, d))
end

"""
    normal_tvector_distribution(M::Euclidean, p, σ)

Normal distribution in ambient space with standard deviation `σ`
projected to tangent space at `p`.
"""
function normal_tvector_distribution(M::Manifold, p, σ)
    d = Distributions.MvNormal(zero(vec(p)), σ)
    return ProjectedFVectorDistribution(TangentBundleFibers(M), p, d, project!, p)
end

function Distributions.support(tvd::ProjectedFVectorDistribution)
    return FVectorSupport(tvd.type, tvd.point)
end


# See https://en.wikipedia.org/wiki/Polynomial

using Polynomials

expressions = [
    :(roots(Polynomial([24 + 7im, -50, 35, -10, 1])))
    :(roots(Polynomial([-26, 24 - 44im, 14 + 33im, -9 - 5im, 1])))
]

for expr in expressions
    println("[] :julia-answer \"\"\"", escape_string(string(expr)), " = ", eval(expr), "\"\"\".")
end


# Sigma point transform functions

"""
unscented_transform(b::GaussianBelief, λ::Int=2, α::Number=1,
    β::Number=0; decomp_method::String = "cholesky")

Convert a single GaussianBelief (mean and covariance) to set of 2n+1 sigma
points, with n being the dimensionality of the state space. Return an array of
the points, and arrays for weights used in mean and covariance calculations.
Uses formulation from ProbRob for α/β parameters on a separate covariance
weighting, although this is not necessary.

"""
function unscented_transform(b::GaussianBelief, λ::Number=2, α::Number=1,
    β::Number=0; decomp_method::String = "cholesky")

    # compute state space dimensionality
    n = length(b.μ)

    # compute weights
    w_μ = 1/(2.0*(n+λ))*ones(2*n+1)
    w_μ[1] = λ/(λ+n)
    w_Σ = copy(w_μ)
    w_Σ[1] += (1 - α^2 + β) # Per ProbRob formulation

    # compute Σ^0.5 using different methods
    if decomp_method == "cholesky" # default
        s = cholesky(b.Σ).L
    elseif decomp_method == "unit-axes"
        # TODO: SVD decomp method
        error("Unit axes method not yet implemented.")
    elseif decomp_method == "tilt-axes"
        # TODO: SVD decomp method with tilt transformation
        error("Tilt axes method not yet implemented.")
    else
        error("Undefined unscented transform Σ sqrt.")
    end

    # compute sigma points
    points = [b.μ]
    for i in 1:n
        push!(points, b.μ + sqrt(n+λ)*s[:,i])
        push!(points, b.μ - sqrt(n+λ)*s[:,i])
    end

    return points, w_μ, w_Σ
end

"""
    unscented_transform_inverse(points::Vector{AbstractVector}, w_μ::Vector,
        w_Σ::Vector)

Convert a 2n+1 sigma points and weights back to a single measure for
mean and covariance (GaussianBelief).

Uses formulation from ProbRob for α/β parameters on a separate covariance
weighting, although this is not necessary.

"""
function unscented_transform_inverse(points::Vector{A}, w_μ::Vector{B},
    w_Σ::Vector{C}) where {A<:AbstractVector, B<:Number, C<:Number}

    catpoints = reduce(hcat, points)
    # calculate weighted mean
    μ = catpoints * w_μ

    # calculated weighted covariance
    diff = catpoints .- μ
    scaled = diff .* sqrt.(w_Σ)'
    Σ = scaled * scaled'

    return GaussianBelief(μ, Σ)
end

# Unscented Kalman Filter functions

"""
    predict(filter::UnscentedKalmanFilter, b0::GaussianBelief, u::AbstractVector)

Uses Unscented Kalman filter to run prediction step on gaussian belief b0,
given control vector u.
"""
function predict(filter::UnscentedKalmanFilter, b0::GaussianBelief, u::AbstractVector{<:Number})

    # Motion update

    n = length(b0.μ)

    # approximate Gaussian belief with sigma points
    points, w_μ, w_Σ = unscented_transform(b0, filter.λ, filter.α, filter.β)

    # iterate over each sigma point and propagate it through motion function
    pointsp = [predict(filter.d, point, u) for point in points]

    # apply inverse unscented transform to approximate new Gaussian
    bp = unscented_transform_inverse(pointsp, w_μ, w_Σ)

    # add process noise
    Σp = bp.Σ + filter.d.W

    return GaussianBelief(bp.μ, Σp)
end

"""
    measure(filter::UnscentedKalmanFilter, bp::GaussianBelief, y::AbstractVector;
        u::AbstractVector = [false])

Uses Unscented Kalman filter to run measurement update on predicted gaussian
belief bp, given measurement vector y. If u is specified and filter.o.D has
been declared, then matrix D will be factored into the y predictions.
"""
function measure(filter::UnscentedKalmanFilter, bp::GaussianBelief, y::AbstractVector{<:Number};
                u::AbstractVector{<:Number} = [false])

    # Measurement update

    # approximate Gaussian belief with sigma points
    points, w_μ, w_Σ = unscented_transform(bp, filter.λ, filter.α, filter.β)

    # iterate over sigma points and computed expected measurement
    ysp = [measure(filter.o, point, u) for point in points]

    catpoints = reduce(hcat, ysp)
    # compute expected measurement
    yp = catpoints * w_μ

    # compute marginal covariance components
    ydiff = catpoints .- yp
    scaled_ydiff = ydiff .* sqrt.(w_Σ)'
    Σ_Y = scaled_ydiff * scaled_ydiff' + filter.o.V

    xdiff = reduce(hcat, points) .- bp.μ
    scaled_xdiff = xdiff .* sqrt.(w_Σ)'
    Σ_XY = scaled_xdiff * scaled_ydiff'

    # measurement update
    K = Σ_XY / Σ_Y
    μn = bp.μ + K * (y - yp)
    Σn = bp.Σ - K * Σ_XY'
    return GaussianBelief(μn, Σn)
end


using Zygote
using Zygote:@adjoint
using Random

export gradient_check

@adjoint function qr(A)
		Q, R = qr(A)
    (Q, R), dy -> (qr_back(A, Q, R, dy...),)
end

@adjoint function qr(A::AbstractMatrix, pivot::Val{true})
		Q, R, P = qr(A, pivot)
    (Q, R, P), dy -> (qr_back(Q*R, Q, R, dy[1], dy[2])*P',nothing)
end

@adjoint function lq(A)
    L, Q = lq(A)
    (L, Q), dy -> (lq_back(A, L, Q, dy...),)
end

@adjoint function svd(A)
    U, S, V = svd(A)
    (U, S, V), dy -> (svd_back(U, S, V, dy...),)
end

@adjoint function rsvd(A, args...; kwargs...)
    U, S, V = rsvd(A, args...; kwargs...)
    (U, S, V), dy -> (svd_back(U, S, V, dy...),)
end

@adjoint function symeigen(A)
    E, U = symeigen(A)
    (E, U), adjy -> (symeigen_back(E, U, adjy...),)
end

function gradient_check(f, args...; η = 1e-5)
    g = gradient(f, args...)
    dy_expect = η*sum(abs2.(g[1]))
    dy = f(args...)-f([gi == nothing ? arg : arg.-η.*gi for (arg, gi) in zip(args, g)]...)
    @show dy
    @show dy_expect
    isapprox(dy, dy_expect, rtol=1e-2, atol=1e-8)
end

@adjoint Random.seed!(n) = Random.seed!(n), _ -> nothing


using Distributions
using Test
using DataFrames
using CSVFiles
using DataStructures

import GlobalSensitivityAnalysis: borehole, split_output, sample

################################################################################
## JULIA
################################################################################

# define the (uncertain) parameters of the problem and their distributions
data = SobolData(
    params = OrderedDict(
        :rw => Normal(0.10, 0.0161812),
        :r => LogNormal(7.71, 1.0056),
        :Tu => Uniform(63070, 115600),
        :Hu => Uniform(990, 1110),
        :Tl => Uniform(63.1, 116),
        :Hl => Uniform(700, 820),
        :L => Uniform(1120, 1680),
        :Kw => Uniform(9855, 12045)
    ),
    N = 1000
)

N = data.N
D = length(data.params)

# sampling
julia_samples = sample(data)
julia_borehole = borehole(julia_samples)

# analysis
julia_A, julia_B, julia_AB, julia_BA = split_output(julia_borehole, N, D, data.calc_second_order)
julia_results = analyze(data, julia_borehole, num_resamples = 10_000)

################################################################################
## Python
################################################################################

# sampling
py_samples = load("data/sobol/py_borehole/py_samples.csv", header_exists=false, colnames = ["rw", "r", "Tu", "Hu", "Tl", "Hl", "L", "Kw"]) |> DataFrame
py_borehole = load("data/sobol/py_borehole/py_borehole.csv", header_exists=false) |> DataFrame

# analysis
py_A = load("data/sobol/py_borehole/py_A.csv", header_exists=false) |> DataFrame
py_B = load("data/sobol/py_borehole/py_B.csv", header_exists=false) |> DataFrame
py_AB = load("data/sobol/py_borehole/py_AB.csv", header_exists=false) |> DataFrame
py_BA = load("data/sobol/py_borehole/py_BA.csv", header_exists=false) |> DataFrame

py_firstorder = load("data/sobol/py_borehole/py_firstorder.csv", header_exists=false) |> DataFrame
py_secondorder = load("data/sobol/py_borehole/py_secondorder.csv", header_exists=false) |> DataFrame
py_totalorder = load("data/sobol/py_borehole/py_totalorder.csv", header_exists=false) |> DataFrame

py_firstorder_conf = load("data/sobol/py_borehole/py_firstorder_conf.csv", header_exists=false) |> DataFrame
py_secondorder_conf = load("data/sobol/py_borehole/py_secondorder_conf.csv", header_exists=false) |> DataFrame
py_totalorder_conf = load("data/sobol/py_borehole/py_totalorder_conf.csv", header_exists=false) |> DataFrame

################################################################################
## Testing
################################################################################

@testset "Borehole Function - Sampling" begin
    @test julia_samples ≈ Matrix(py_samples) atol = ATOL_sample
end

@testset "Borehole Function - Analysis" begin

    @test julia_borehole ≈ Matrix(py_borehole) atol = ATOL_sobol
    @test julia_A ≈ Matrix(py_A) atol = ATOL_sobol
    @test julia_B ≈ Matrix(py_B) atol = ATOL_sobol
    @test julia_AB ≈ Matrix(py_AB) atol = ATOL_sobol
    @test julia_BA ≈ Matrix(py_BA) atol = ATOL_sobol

    @test julia_results[:firstorder] ≈ Matrix(py_firstorder) atol = ATOL_sobol
    @test julia_results[:totalorder]≈ Matrix(py_totalorder) atol = ATOL_sobol

    for i = 1:D
        for j = i+1:D
            @test julia_results[:secondorder][i,j] ≈ Matrix(py_secondorder)[i,j] atol = ATOL_sobol
        end
    end

    # test confidence intervals
    @test julia_results[:firstorder_conf] ≈ Matrix(py_firstorder_conf) atol = ATOL_CI
    @test julia_results[:totalorder_conf] ≈ Matrix(py_totalorder_conf) atol = ATOL_CI

    for i = 1:D
        for j = i+1:D
            @test julia_results[:secondorder_conf][i,j] ≈ Matrix(py_secondorder_conf)[i,j] atol = ATOL_CI
        end
    end

end


# Function for calculating ZV-MCMC estimators using linear polynomial
function linearZv(mcmc::Array{Float64, 2}, z::Array{Float64, 2})
  nPars = size(mcmc, 2)

  covAll = Array(Float64, nPars+1, nPars+1, nPars)
  precision = Array(Float64, nPars, nPars, nPars)
  sigma = Array(Float64, nPars, nPars)
  a = Array(Float64, nPars, nPars)

  for i = 1:nPars
     covAll[:, :, i] = cov([z mcmc[:, i]])
     precision[:, :, i] = inv(covAll[1:nPars, 1:nPars, i])
     sigma[:, i] = covAll[1:nPars, nPars+1, i]
     a[:, i] = -precision[:, :, i]*sigma[:, i]
  end

  zvMcmc = mcmc+z*a

  return zvMcmc, a
end

# Function for calculating ZV-MCMC estimators using quadratic polynomial
function quadraticZv(mcmc::Array{Float64, 2}, z::Array{Float64, 2})
  nData, nPars = size(mcmc)
  k = convert(Int, nPars*(nPars+3)/2)
  l = 2*nPars+1

  zQuadratic = Array(Float64, nData, k)
  covAll = Array(Float64, k+1, k+1, nPars)
  precision = Array(Float64, k, k, nPars)
  sigma = Array(Float64, k, nPars)
  a = Array(Float64, k, nPars)

  zQuadratic[:, 1:nPars] = z
  zQuadratic[:, (nPars+1):(2*nPars)] = 2*z.*mcmc-1
  for i = 1:(nPars-1)
    for j = (i+1):nPars
      zQuadratic[:, l] = mcmc[:, i].*z[:, j]+mcmc[:, j].*z[:, i]
      l += 1
    end
  end

  for i = 1:nPars
    covAll[:, :, i] = cov([zQuadratic mcmc[:, i]]);
    precision[:, :, i] =
      inv(covAll[1:k, 1:k, i])
    sigma[:, i] = covAll[1:k, k+1, i]
    a[:, i] = -precision[:, :, i]*sigma[:, i]
  end

  zvMcmc = mcmc+zQuadratic*a;

  return zvMcmc, a
end