"""
    box dynamics
    - 3D particle with orientation and 8 corners subject to contact forces

    - configuration: q = (px, py, pz, rx, ry, rz) ∈ R⁶
        - orientation : modified Rodrigues angles

    - contacts (8x)
        - impact force (magnitude): n ∈ R₊
        - friction force: b ∈ R²
            - contact force: λ = (b, n) ∈ R² × R₊
            - friction coefficient: μ ∈ R₊

    Discrete Mechanics and Variational Integrators
        pg. 363
"""

struct Box
	n::Int
	m::Int
	d::Int

    mass # mass
    J # inertia
    μ # friction coefficient
    g # gravity

    r             # corner length
	n_corners     # number of corners
    corner_offset # precomputed corner offsets

    nq # configurations dimension
end

# Methods
function mass_matrix(model::Box)
	Diagonal(@SVector [model.mass, model.mass, model.mass,
		model.J, model.J, model.J])
end

function gravity(model::Box)
	@SVector [0., 0., model.mass * model.g, 0., 0., 0.]
end

function kinematics(model::Box, q)
    p = view(q, 1:3)
    r = view(q, 4:6)

    R = MRP(r...)

    SVector{24}([(p + R * model.corner_offset[1])...,
                 (p + R * model.corner_offset[2])...,
                 (p + R * model.corner_offset[3])...,
                 (p + R * model.corner_offset[4])...,
                 (p + R * model.corner_offset[5])...,
                 (p + R * model.corner_offset[6])...,
                 (p + R * model.corner_offset[7])...,
                 (p + R * model.corner_offset[8])...])
end

function jacobian(model::Box, q)
	k(z) = kinematics(model, z)
	ForwardDiff.jacobian(k, q)
end

function signed_distance(model::Box, q)
	idx = collect([3, 6, 9, 12, 15, 18, 21, 24])
	kinematics(model, q)[idx]
end

function P_func(model::Box, q)
	idx = collect([1, 2, 4, 5, 7, 8, 10, 11, 13, 14, 16, 17, 19, 20, 22, 23])
	k(z) = kinematics(model, z)[idx]
	ForwardDiff.jacobian(k, q)
end

function B_func(model::Box, q)
	# p = view(q, 1:3)
    r = view(q, 4:6)

    R = MRP(r...)
    # @SMatrix [0. 0. 0. R[1,1] R[2,1] R[3,1];
    #           0. 0. 0. R[1,2] R[2,2] R[3,2];
    #           0. 0. 0. R[1,3] R[2,3] R[3,3]]
	@SMatrix [1. 0. 0. 0.0 0.0 0.0;
			  0. 1. 0. 0.0 0.0 0.0;
			  0. 0. 1. 0.0 0.0 0.0;
			  0. 0. 0. R[1,1] R[2,1] R[3,1];
		      0. 0. 0. R[1,2] R[2,2] R[3,2];
		      0. 0. 0. R[1,3] R[2,3] R[3,3]]
end

# dynamics
function dynamics(model, q1, q2, q3, u1, λ, h)
      nq = model.nq
      SVector{nq}(mass_matrix(model) * (2.0 * q2 - q1 - q3) / h
            - h * gravity(model)
            + h * jacobian(model, q3)' * λ
			+ h * transpose(B_func(model, q3)) * u1)
end

# Kinematics
num_contacts = 8
d = 0.5
c8 = @SVector [d, d, d]
c2 = @SVector [d, d, -d]
c3 = @SVector [d, -d, d]
c4 = @SVector [d, -d, -d]
c5 = @SVector [-d, d, d]
c6 = @SVector [-d, d, -d]
c7 = @SVector [-d, -d, d]
c1 = @SVector [-d, -d, -d]

corner_offset = @SVector [c1, c2, c3, c4, c5, c6, c7, c8]

# Model
nq = 6
nu = 6
model = Box(2 * nq, nu, 0,
			1.0,
			1.0 / 12.0 * 1.0 * ((2.0 * d)^2 + (2.0 * d)^2),
 			0.5,
			9.81,
            d, num_contacts, corner_offset,
            nq)

# qq = rand(nq)
# kinematics(model, qq)
# ϕ_func(model, qq)
# jacobian(model, qq)
# P_func(model, qq)
num_contacts
# var
num_var = nq + num_contacts + num_contacts + 3 * num_contacts + 3 * num_contacts
function unpack(z)
	q = view(z, 1:nq)
	n = view(z, nq .+ (1:num_contacts))
	sϕ = view(z, nq + num_contacts .+ (1:num_contacts))
	b = view(z, nq + 2 * num_contacts .+ (1:3 * num_contacts))
	sb = view(z, nq + 2 * num_contacts + 3 * num_contacts .+ (1:3 * num_contacts))

	b_traj = [view(b, (i - 1) * 3 .+ (1:3)) for i = 1:num_contacts]
	sb_traj = [view(sb, (i - 1) * 3 .+ (1:3)) for i = 1:num_contacts]

	return q, n, sϕ, b_traj, sb_traj
end

function initialize(q2, num_var; z_init = 1.0)
	z = z_init * ones(num_var)

	z[1:nq] = copy(q2)
	z[nq + num_contacts .+ (1:num_contacts)] = signed_distance(model, q2)
	z[nq + 2 * num_contacts .+ (1:3 * num_contacts)] = vcat([[model.μ * view(z, nq .+ (1:num_contacts))[i]; 0.1 * model.μ * view(z, nq .+ (1:num_contacts))[i]; 0.1 * model.μ * view(z, nq .+ (1:num_contacts))[i]] for i = 1:num_contacts]...)
	z[nq + 2 * num_contacts + 3 * num_contacts .+ (1:3 * num_contacts)] = vcat([[10.0 * z_init; 1.0 * z_init; 1.0 * z_init] for i = 1:num_contacts]...)
	return z
end

function visualize!(vis, model::Box, q;
        Δt = 0.1)

	default_background!(vis)

	setobject!(vis["box"], GeometryBasics.Rect(Vec(-1.0 * model.r,
		-1.0 * model.r,
		-1.0 * model.r),
		Vec(2.0 * model.r, 2.0 * model.r, 2.0 * model.r)),
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

    for i = 1:model.n_corners
        setobject!(vis["corner$i"], GeometryBasics.Sphere(Point3f0(0),
            convert(Float32, 0.05)),
            MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))
    end

    anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

    for t = 1:length(q)
        MeshCat.atframe(anim, t) do

            settransform!(vis["box"],
				compose(Translation(q[t][1:3]...), LinearMap(MRP(q[t][4:6]...))))

            for i = 1:model.n_corners
                settransform!(vis["corner$i"],
                    Translation((q[t][1:3] + MRP(q[t][4:6]...) * (corner_offset[i]))...))
            end
        end
    end
    # settransform!(vis["/Cameras/default"], compose(Translation(-1, -1, 0),LinearMap(RotZ(pi/2))))
    MeshCat.setanimation!(vis, anim)
end


"""
OGIP(args...)

Calcute the Original Gas In Place either Deterministic or Probabilistic

``OGIP=\\frac { 4.356\\times { 10 }^{ -5 }\\quad A\\quad h\\quad Phi\\quad (1-Sw) }{ Bg } `` in Bscf or MMMscfd

 Bg can be supplied or be calculated with a Pressure, Temperature in °F and Z

``Bg=\\frac { 0.02827\\quad z\\quad T }{ P } `` in ft3/scf. From Tarek Ahmed Reservoir Handbook 4ed pg 66
 ## *Deterministic Estimation*
 ### Example
 ```julia
     OGIP(Area=600,Height=30,Phi=0.3,Sw=0.3, Temp=180, Pres=2800, z=0.99)
     Original Gas In Place Deterministic Results
     Area = 600 Acre
     Height= 30 ft
     Phi= 0.3
     Sw= 0.3
     Bg= 0.006393798639642858 ft3|scf
    ---------
    OGIP=25.752578284072662 Bscf
```
## Probabilistic estimations
### Example
```julia
    OGIP(Area=600,Height=30,Phi=0.3,Sw=0.3, Temp=180, Pres=2800, z=0.99,PhiDist=Normal(0.18,0.03), SwDist=Normal(0.3,0.05))

    Original Gas In Place Probabilistic Results
    Area = 600 Acre
    Height= 30 ft
    Bg= 0.006393798639642858 ft3|scf
    Percentiles [0.1, 0.5, 0.9]
    Phi       [0.141553, 0.18, 0.218447]
    Sw       [0.235922, 0.3, 0.364078]
    ---------
    OGIP=      [12.1409, 15.4248, 19.1437] Bscf
```
Number of samples n can be set. Default n=1000

Percentiles can be set. Default Perc=[0.1,0.5,0.9]

 A plot can be display by setting DistHist=true. Default DisHist=false
"""
function OGIP(;Area=640,Height=1,Phi=0, Sw=0, Bg=false,
                Pres=0, Temp=0, z=0,
                PhiDist=Normal(), SwDist=Normal(), n=1000,
                DisHist=false,
                Perc=[0.1,0.5,0.9],
                Phicolor= :darkred, Phialpha= 0.2,
                Swcolor= :darkblue, Swalpha= 0.2)

          #When is only a well for the analysis the area is ussually difined as 640 acres
     #https://www.spec2000.net/16-ooip.htm

    if Bg==false
        TempR=459.67+Temp
        Bg=0.02827*z*TempR/Pres
    end

    if PhiDist==Normal() && SwDist==Normal()
        Ogip=0.000043560*Area*Height*Phi*(1-Sw)/Bg
        println("Original Gas In Place Deterministic Results \n Area = $Area Acre \n Height= $Height ft \n Phi= $Phi \n Sw= $Sw \n Bg= $Bg ft3|scf \n---------\n OGIP=$Ogip Bscf")
    elseif PhiDist!=Normal() && SwDist!=Normal()
        PhiSample=rand(Truncated(PhiDist,0,1),n)
        SwSample=rand(Truncated(SwDist,0,1),n)
        Ogip=map((x,y)->0.000043560*Area*Height*x*(1-y)/Bg,PhiSample,SwSample)
        if DisHist!=false

        p1=prophist(Ogip,Normal,title="Original Gas In Place [Bscf]")
        p2=prophist(PhiSample,Normal,seriescolor= Phicolor, seriesalpha= Phialpha, title="Porosity")
        p3=prophist(SwSample,LogNormal,seriescolor= Swcolor, seriesalpha= Swalpha,title="Water Saturation")
        p=Plots.plot(p1,p2,p3, layout=(3,1))
        display(p)
        end

            PhiQ=quantile.(Truncated(PhiDist,0,1),Perc)
            SwQ=quantile.(Truncated(SwDist,0,1),Perc)
            OgipQ=quantile!(Ogip,Perc)

        println("Original Gas In Place Probabilistic Results \n Area = $Area Acre \n Height= $Height ft \n Bg= $Bg ft3|scf \n Percentiles $Perc \n Phi       $PhiQ \n Sw       $SwQ \n --------- \n OGIP=      $OgipQ Bscf")

    elseif PhiDist!=Normal()
        PhiSample=rand(Truncated(PhiDist,0,1),n)
        Ogip=map(x->0.000043560*Area*Height*x*(1-Sw)/Bg,PhiSample)

        if DisHist!=false

        p1=prophist(Ogip,Normal,title="Original Gas In Place [Bscf]")
        p2=prophist(PhiSample,Normal,seriescolor= Phicolor, seriesalpha= Phialpha,title="Porosity")
        p=Plots.plot(p1,p2, layout=(2,1))
            display(p)
            end
        PhiQ=quantile.(Truncated(PhiDist,0,1),Perc)
        OgipQ=quantile!(Ogip,Perc)

println("Original Gas In Place Probabilistic Results \n Area = $Area Acre \n Height= $Height ft \n Bg= $Bg ft3|scf \n Sw= $Sw \n Percentiles $Perc \n Phi       $PhiQ \n---------\n OGIP       $OgipQ Bscf")

    else
        SwSample=rand(Truncated(SwDist,0,1),n)
        Ogip=map(x->0.000043560*Area*Height*Phi*(1-x)/Bg,SwSample)

         if DisHist!=false

        p1=prophist(Ogip,Normal,title="Original Gas In Place [Bscf]")
        p2=prophist(SwSample,LogNormal,seriescolor= Swcolor, seriesalpha= Swalpha,title="Water Saturation")
        p=Plots.plot(p1,p2, layout=(2,1))
        display(p)
            end
        SwQ=quantile.(Truncated(SwDist,0,1),Perc)
        OgipQ=quantile!(Ogip,Perc)

println("Original Gas In Place Probabilistic Results \n Area = $Area Acre \n Height= $Height ft \n Bg= $Bg ft3|scf \n Phi= $Phi \n \n Percentiles $Perc \n Sw       $SwQ \n---------\n OGIP       $OgipQ Bscf")

    end

end


Nx=200
Lx=1.0
u=1.0
dt=Lx/u/Nx/10
t_end=0.5*Lx/u
m=createMesh1D(Nx, Lx)
u_face=createFaceVariable(m, [u])
c_init=createCellVariable(m, 0.0)
c_init_up=createCellVariable(m, 0.0)
c=createCellVariable(m, 0.0)
c_upwind=createCellVariable(m, 0.0)

BC=createBC(m)
BC.left.a[:]=0.0
BC.left.b[:]=1.0
BC.left.c[:]=1.0
Mbc, RHSbc=boundaryConditionTerm(BC)

Mconv_up=convectionUpwindTerm(u_face)

flux_limiter=fluxLimiter("SUPERBEE")
for t in dt:dt:t_end
  Mt, RHSt=transientTerm(c_init, dt)
  Mt_up, RHSt_up=transientTerm(c_init_up, dt)
  c_upwind=solveLinearPDE(m, Mt_up+Mconv_up+Mbc, RHSt_up+RHSbc)
  for j in 1:10
    RHSconv=convectionTvdRHS(u_face, c, flux_limiter)
    c=solveLinearPDE(m, Mt+Mconv_up+Mbc, RHSt+RHSbc+RHSconv)
  end
  c_init=copyCell(c)
  c_init_up=copyCell(c_upwind)
end
Plots.plot(c.value[2:end-1])
Plots.plot!(c_upwind.value[2:end-1])


using GeometricFlux
using Flux
using Flux: onehotbatch, onecold, logitcrossentropy, throttle
using Flux: @epochs
using Flux.Data: DataLoader
using JLD2
using Statistics: mean
using SparseArrays
using LightGraphs.SimpleGraphs
using LightGraphs: adjacency_matrix
using CUDA

@load "data/cora_features.jld2" features
@load "data/cora_labels.jld2" labels
@load "data/cora_graph.jld2" g

num_nodes = 2708
num_features = 1433
heads  = 8
hidden = 8
target_catg = 7
epochs = 10

## Preprocessing data
train_X = Matrix{Float32}(features) |> gpu  # dim: num_features * num_nodes
train_y = Matrix{Float32}(labels) |> gpu  # dim: target_catg * num_nodes
adj_mat = Matrix{Float32}(adjacency_matrix(g)) |> gpu

## Model
model = Chain(GATConv(g, num_features=>hidden, heads=heads),
              Dropout(0.6),
              GATConv(g, hidden*heads=>target_catg, heads=heads, concat=false)
              ) |> gpu
# test model
# @show model(train_X)

## Loss
loss(x, y) = logitcrossentropy(model(x), y)
accuracy(x, y) = mean(onecold(cpu(model(x))) .== onecold(cpu(y)))

# test loss
# @show loss(train_X, train_y)

# test gradient
# @show gradient(X -> loss(X, train_y), train_X)

## Training
ps = Flux.params(model)
train_data = DataLoader(train_X, train_y, batchsize=num_nodes)
opt = ADAM(0.01)
evalcb() = @show(accuracy(train_X, train_y))

@epochs epochs Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))


using RvSpectraKitLearn
using MultivariateStats

# Likely need to update this to point to directory containing the spectra you wish to analyze
include(joinpath(Pkg.dir("RvSpectraKitLearn"),"path_to_spectra.jl"))
cd(path_to_spectra)

(datafiles, phases) = get_filenames(); # num_max_files=4);
(lambda, obs_noisefree) = read_filelist(datafiles);

sampling = 3.0
planet_amplitude = 10.0
rvs_true = -planet_amplitude*cos.(2pi*phases)

snr = 100.0
obs = RvSpectraKitLearn.make_noisy_spectra(obs_noisefree,snr,sampling=sampling)

pca_out = fit_pca_default(obs);
pca_scores = transform(pca_out,obs)   # compute scores using Julia's PCA routine for comparison's sake
pca_eford_out = fit_pca_eford(obs);   # compute first few scores using Eric's itterative PCA routine

doppler_comp_simple = calc_doppler_component_simple(lambda,obs);
genpca_simple_out = fit_gen_pca_rv(obs,doppler_comp_simple)
rvs_out_simple = est_rvs_from_pc(obs,genpca_simple_out[1],genpca_simple_out[2][:,1])

doppler_comp_gp = calc_doppler_component_gp(lambda,obs);
genpca_gp_out = fit_gen_pca_rv(obs,doppler_comp_gp)
rvs_out_gp = est_rvs_from_pc(obs,genpca_gp_out[1],genpca_gp_out[2][:,1])
rms_simple = sqrt(sumabs2(rvs_out_simple.-rvs_true)/length(rvs_true))
rms_gp = sqrt(sumabs2(rvs_out_gp.-rvs_true)/length(rvs_true))
println("snr = ", snr, " rms = ", rms_simple, "  ", rms_gp)

using PyPlot
plot(phases,rvs_true,"b-")
 plot(phases,rvs_out_simple,"r.")
 plot(phases,rvs_out_gp,"g.")
 plot(phases,rvs_out_simple.-rvs_true,"r.")
 plot(phases,rvs_out_gp.-rvs_true,"g.")
 xlabel("Phase")
 ylabel("Velocity (m/s)")


###################################################
## geometry.jl
## Geometry tools
###################################################

"""
    Return bounding box of polygon/ list of points
"""
function boundingBox{T<:AbstractFloat}(poly::Vector{Tuple{T,T}})
    minX::T = Inf; maxX::T = -Inf; minY::T = Inf; maxY::T = -Inf
    for (x,y) in poly
        minX = min(minX,x)
        maxX = max(maxX,x)
        minY = min(minY,y)
        maxY = max(maxY,y)
    end
    return (minX,maxX,minY,maxY)
end

"""
    Scales a given box (i.e. 5% bigger if scale=1.05)
"""
function scaleBox(box, scale)
    minX, maxX, minY, maxY = box
    lenX = maxX - minX; midX = (minX + maxX)/2
    lenY = maxY - minY; midY = (minY + maxY)/2
    return (midX - lenX*scale/2, midX + lenX*scale/2, midY-lenY*scale/2, midY+lenY*scale/2)
end

"""
    Distance between the two points in meters (true geodesic distance from LLA coordinates)
"""
distanceGeo(pt1::Node, pt2::Node) = distanceGeo(pt1.lon, pt1.lat, pt2.lon, pt2.lat)

function distanceGeo(lon1, lat1, lon2, lat2)
    dLat = toradians(lat2 - lat1)
    dLon = toradians(lon2 - lon1)
    lat1 = toradians(lat1)
    lat2 = toradians(lat2)
    a = sin(dLat/2)^2 + sin(dLon/2)^2 * cos(lat1) * cos(lat2)
    2.0 * atan2(sqrt(a), sqrt(1-a)) * 6373.0 * 1000
end

"""
    Distance between the two points in meters (from coordinates)
"""
distanceCoord(pt1::Node, pt2::Node) = distanceCoord(pt1.x, pt1.y, pt2.x, pt2.y)

function distanceCoord{T<:AbstractFloat}(x1::T, y1::T, x2::T, y2::T)
    return sqrt((x2-x1)^2+(y2-y1)^2)
end

"""
    Check is a point is inside a polygon
"""
function pointInsidePolygon{T<:AbstractFloat}(x::AbstractFloat,y::AbstractFloat,poly::Vector{Tuple{T,T}})
    n = length(poly)
    inside =false

    p1x,p1y = poly[1]
    for i in 0:n
        p2x,p2y = poly[i % n + 1]
        if y > min(p1y,p2y) && y <= max(p1y,p2y) && x <= max(p1x,p2x)
            if p1y != p2y
                xinters = (y-p1y)*(p2x-p1x)/(p2y-p1y)+p1x
            end
            if p1x == p2x || x <= xinters
                inside = !inside
            end
        end
        p1x,p1y = p2x,p2y
    end
    return inside
end

"""
    `toENU`, projects latitude and longitude to ENU coordinate system
"""
function toENU(lon, lat, center)
    enu = ENU(LLA(lat,lon), LLA(center[2],center[1]))
    return enu.east, enu.north
end
toENU(lon, lat, n::Network) = toENU(lon, lat, n.projcenter)
"""
    `updateProjection`
    updates the (x,y) projection of the nodes of the network given their latitude/longitude
"""
function updateProjection!(n::Network)
    bounds = boundingBox([(n.lon,n.lat) for n in n.nodes])
    n.projcenter = ((bounds[2]+bounds[1])/2, (bounds[4]+bounds[3])/2)
    nodes = Array{Node}(length(n.nodes))
    for (i,no) in enumerate(n.nodes)
        x,y = toENU(no.lon,no.lat,n.projcenter)
        nodes[i] = Node(x,y,no.lon,no.lat)
    end
    n.nodes = nodes
    n
end


using GeometricIntegrators.BasisFunctions
using GeometricIntegrators.CommonFunctions
using GeometricIntegrators.Config
using GeometricIntegrators.Integrators
using GeometricIntegrators.Quadratures
using GeometricIntegrators.Solvers
using GeometricIntegrators.TestProblems.HarmonicOscillatorProblem
using GeometricIntegrators.Utils
using Test

using GeometricIntegrators.TestProblems.HarmonicOscillatorProblem: Δt, nt, refx, refq, refp

iode = harmonic_oscillator_iode()

QGau4 = GaussLegendreQuadrature(4)
BGau4 = LagrangeBasis(nodes(QGau4))

### CGVI Integrators ###

cgint = IntegratorCGVI(iode, BGau4, QGau4, Δt)
cgsol = integrate(cgint, nt)

@test rel_err(cgsol.q, refx) < 1E-7

### DGVI Integrators ###

dgint = IntegratorDGVI(iode, BGau4, QGau4, Δt)
dgsol = integrate(dgint, nt)

@test rel_err(dgsol.q, refx) < 1E-7


"""
    solve(problem::AnalyticalMatrixFormalism, gradient)

Compute the signal in a multilayered cylinder or sphere using an analytical
matrix formalism solution.

This function is based on the following articles and corresponding code:
    [1] D. S. Grebenkov, NMR Survey of Reflected Brownian Motion,
        Rev. Mod.Phys. 79, 1077 (2007)
    [2] D. S. Grebenkov, Pulsed-gradient spin-echo monitoring of restricted diffusion in
        multilayered structures, J. Magn. Reson. 205, 181-195 (2010).
"""
function solve(problem::AnalyticalMatrixFormalism, gradient::ScalarGradient)
    (; analytical_laplace, lap_mat, volumes) = problem
    (; r, ρ, γ) = analytical_laplace
    (; Λ, Br, B, U) = lap_mat

    gradient.profile isa PGSE || error("Only implemeted for PGSE")

    w = γ * gradient.amplitude * 1e12 * r[end]
    δ = gradient.profile.δ * 1e-6
    Δ = gradient.profile.Δ * 1e-6

    U₀ = complex(U)
    U = complex(U)

    expmv!(-δ, Λ + Br + im * w * B, U)
    expmv!(-(Δ - δ), Λ + Br, U)
    expmv!(-δ, Λ + Br - im * w * B, U)

    U₀'U * volumes'ρ
end


update{N, T, T2}(b::HyperRectangle{N, T}, v::Vec{N, T2}) =
    update(b, Vec{N, T}(v))
function update{N, T}(b::HyperRectangle{N, T}, v::Vec{N, T})
    m = min(minimum(b), v)
    mm = max(maximum(b), v)-m
    HyperRectangle{N, T}(m, mm)
end

# Min maximum distance functions between hrectangle and point for a given dimension
@inline function min_dist_dim{N, T}(rect::HyperRectangle{N, T}, p::Vec{N, T}, dim::Int)
    max(zero(T), max(minimum(rect)[dim] - p[dim], p[dim] - maximum(rect)[dim]))
end

@inline function max_dist_dim{N, T}(rect::HyperRectangle{N, T}, p::Vec{N, T}, dim::Int)
    max(maximum(rect)[dim] - p[dim], p[dim] - minimum(rect)[dim])
end

@inline function min_dist_dim{N, T}(rect1::HyperRectangle{N, T},
                                    rect2::HyperRectangle{N, T},
                                    dim::Int)
    max(zero(T), max(minimum(rect1)[dim] - maximum(rect2)[dim],
                     minimum(rect2)[dim] - maximum(rect1)[dim]))
end

@inline function max_dist_dim{N, T}(rect1::HyperRectangle{N, T},
                                    rect2::HyperRectangle{N, T},
                                    dim::Int)
    max(maximum(rect1)[dim] - minimum(rect2)[dim],
        maximum(rect2)[dim] - minimum(rect1)[dim])
end

# Total minimum maximum distance functions
@inline function min_euclideansq{N, T}(rect::HyperRectangle{N, T},
                                       p::Union{Vec{N, T},
                                       HyperRectangle{N, T}})
    minimum_dist = T(0.0)
    for dim in 1:length(p)
        d = min_dist_dim(rect, p, dim)
        minimum_dist += d*d
    end
    return minimum_dist
end

# could add an @symmetric macro, which defines min_euclidean(pt, s) from
# min_euclidean(s, pt) automatically etc.
@inline min_euclidean(pt1::Vec, pt2::Vec) = norm(pt1-pt2)
min_euclidean(pt::Vec, s::Simplex) = √(sqdist(pt, s))
min_euclidean(s::Simplex, pt::Vec) = min_euclidean(pt, s)
min_euclidean(r1::HyperRectangle, r2::Vec) = sqrt(min_euclideansq(r1, r2))
min_euclidean(r1::Vec, r2::HyperRectangle) = sqrt(min_euclideansq(r1, r2))
min_euclidean(r1::HyperRectangle, r2::HyperRectangle) = sqrt(min_euclideansq(r1, r2))
min_euclidean(c1::AbstractConvexHull, pt::Vec) = gjk(c1,pt)
min_euclidean(pt::Vec, c2::AbstractConvexHull) = gjk(pt,c2)
min_euclidean(c1::AbstractConvexHull, c2::AbstractConvexHull) = gjk(c1,c2)

@inline function max_euclideansq{N, T}(rect::HyperRectangle{N, T}, p::Union{Vec{N, T}, HyperRectangle{N, T}})
    maximum_dist = T(0.0)
    for dim in 1:length(p)
        d = max_dist_dim(rect, p, dim)
        maximum_dist += d*d
    end
    return maximum_dist
end
function max_euclidean{N, T}(rect::HyperRectangle{N, T}, p::Union{Vec{N, T}, HyperRectangle{N, T}})
    sqrt(max_euclideansq(rect, p))
end

# Functions that return both minimum and maximum for convenience
@inline function minmax_dist_dim{N, T}(rect::HyperRectangle{N, T}, p::Union{Vec{N, T}, HyperRectangle{N, T}}, dim::Int)
    minimum_d = min_dist_dim(rect, p, dim)
    maximum_d = max_dist_dim(rect, p, dim)
    return minimum_d, maximum_d
end

@inline function minmax_euclideansq{N, T}(rect::HyperRectangle{N, T}, p::Union{Vec{N, T}, HyperRectangle{N, T}})
    minimum_dist = min_euclideansq(rect, p)
    maximum_dist = max_euclideansq(rect, p)
    return minimum_dist, maximum_dist
end

function minmax_euclidean{N, T}(rect::HyperRectangle{N, T}, p::Union{Vec{N, T}, HyperRectangle{N, T}})
    minimumsq, maximumsq = minmax_euclideansq(rect, p)
    return sqrt(minimumsq), sqrt(maximumsq)
end


@info "Starting precompile exec"

ENV["PYTHON"] = "/srv/conda/envs/notebook/bin/python"

using LinearAlgebra, Statistics, DataFrames, Plots, SymPy, StatsPlots, Distributions

mean(randn(100))
randn(10, 10) \ randn(10)

df = DataFrame(a=[1,2,3], b=[4,5,6])
@df df scatter(:a, :b)

@vars x
expand((x+1)^10)

mean(Normal(0, 1))


using BenchmarkTools
using FiSUDiPPuS

options = joinpath(@__DIR__, "../data/default.jl")
include(options)

x,y = FiSUDiPPuS.get_data(settings) # settings is a dict in the options file

## Convert to correct types
start = settings[:datatype].(settings[:start])
bleach_weight = settings[:datatype](settings[:bleach_weight])

@btime FiSUDiPPuS.model($x[:,:], $start, $settings[:N], $bleach_weight);

x = 3000.0
A = [1.0, -1.0]
Γ = [8.0, 8.0]
ω = [2900.0, 2910.0]
@btime FiSUDiPPuS.sfspec($x, $A, $ω, $Γ)

# 496.331 μs (4505 allocations: 510.06 KiB)
# 336.573 μs (2259 allocations: 194.02 KiB)
# 296.531 μs (22 allocations: 89.84 KiB)
# 298.902 μs (19 allocations: 89.52 KiB)
# 196.338 μs (20 allocations: 89.63 KiB)

# 573.098 ns (11 allocations: 1.19 KiB)
# 447.626 ns (9 allocations: 928 bytes)


# Make sure PERK.krr computes Eq. 11 in Nataraj et al. for exact kernels
function test_krr_1()

    Random.seed!(0)
    T = 100
    xtrain = LinRange(0, 1, T)
    ytrain = xtrain .+ 0.1 .* randn(T)
    kernel = EuclideanKernel()
    ρ = 0.01
    trainData = PERK.krr_train(xtrain, ytrain, kernel, ρ)
    ytest = 0.6
    xhat = PERK.krr(ytest, trainData, kernel)

    K = ytrain * ytrain'
    M = I - ones(T, T) / T
    k = ytest * conj(ytrain) - K * ones(T) / T
    x_correct = xtrain' * (ones(T) / T + M * ((M * K * M + T * ρ * I) \ k))

    # For line coverage, access Q, L, and T in trainData as well
    return xhat ≈ x_correct &&
           trainData.Q == length(ytest) &&
           trainData.L == length(xhat) &&
           trainData.T == 100

end

function test_krr_2()

    Random.seed!(0)
    T = 100
    D = 3
    xtrain = LinRange(0, 1, T)
    ytrain = xtrain' .+ 0.1 .* randn(D, T)
    Λ = ones(D)
    kernel = GaussianKernel(Λ)
    ρ = 0.01
    trainData = PERK.krr_train(xtrain, ytrain, kernel, ρ)
    ytest = transpose([0.6 0.57 0.67])
    xhat = PERK.krr(ytest, trainData, kernel)

    gauss = (p, q) -> exp(-0.5 * norm((1 ./ Λ) .* (p - q))^2)
    K = [gauss(ytrain[:,m], ytrain[:,n]) for m = 1:T, n = 1:T]
    M = I - ones(T, T) / T
    k = [gauss(vec(ytest), ytrain[:,n]) for n = 1:T] - K * ones(T) / T
    x_correct = xtrain' * (ones(T) / T + M * ((M * K * M + T * ρ * I) \ k))

    return xhat[] ≈ x_correct

end

# Make sure PERK.krr with EuclideanKernel produces the same result as ridge
# regression
function test_krr_3()

    Random.seed!(0)
    T = 100
    xtrain = LinRange(0, 1, T)
    ytrain = xtrain .+ 0.1 .* randn(T)
    kernel = EuclideanKernel()
    ρ = 0.01
    trainData = PERK.krr_train(xtrain, ytrain, kernel, ρ)
    N = 10
    ytest = LinRange(0.1, 0.9, N)
    xhat = PERK.krr(ytest, trainData, kernel)

    xtilde = [xtrain .- mean(xtrain); 0]
    Ytilde = [ytrain .- mean(ytrain); sqrt(T * ρ)]
    w = Ytilde \ xtilde
    b = mean(xtrain) - w * mean(ytrain)
    x_correct = w .* ytest .+ b

    return xhat ≈ x_correct

end

function test_krr_4()

    Random.seed!(0)
    T = 100
    D = 3
    xtrain = LinRange(0, 1, T)
    ytrain = xtrain' .+ 0.1 .* randn(D, T)
    kernel = EuclideanKernel()
    ρ = 0.01
    trainData = PERK.krr_train(xtrain, ytrain, kernel, ρ)
    N = 10
    ytest = LinRange(0.1, 0.9, N)' .+ 0.1 .* randn(D, N)
    xhat = PERK.krr(ytest, trainData, kernel)

    xtilde = [xtrain .- mean(xtrain); zeros(D)]
    Ytilde = [(ytrain .- mean(ytrain, dims = 2))'; sqrt(T * ρ) * I]
    w = Ytilde \ xtilde
    b = mean(xtrain) - w' * dropdims(mean(ytrain, dims = 2), dims = 2)
    x_correct = [w' * ytest[:,n] + b for n = 1:N]

    return xhat ≈ x_correct

end

# Other tests
function test_krr_5()

    rng = StableRNG(0)
    f = x -> exp(-30 / x)
    xtrue = 100
    y = f(xtrue)
    T = 200
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    H = 100
    kernel = GaussianRFF(H, λ * mean(y))
    ρ = 2.0^-20
    (ytrain, xtrain) = generatenoisydata(rng, T, xDists, noiseDist, signalModels)
    trainData = PERK.krr_train(xtrain, ytrain, kernel, ρ, randn(rng, H),
        rand(rng, H))
    xhat = PERK.krr(y, trainData, kernel)

    # For code coverage, call
    # - krr_train without rng and freq and phase
    # - generatenoisydata without rng
    # - addnoise! without rng
    # - train without rng
    # (results not used)
    PERK.krr_train(xtrain, ytrain, kernel, ρ)
    generatenoisydata(T, xDists, noiseDist, signalModels)
    PERK.addnoise!([y], noiseDist)
    PERK.train(T, xDists, noiseDist, signalModels, kernel, ρ)

    error_rel = abs(xhat - xtrue) / xtrue
    # For line coverage, access Q, L, and H in trainData as well
    return isapprox(error_rel, 0.07054474887124201, atol = 1e-7) &&
           trainData.Q == length(y) &&
           trainData.L == length(xtrue) &&
           trainData.H == H

end

function test_krr_6()

    rng = StableRNG(0)
    f = x -> exp(-30 / x)
    xtrue = 100
    y = f(xtrue)
    T = 200
    xDists = Uniform(10, 500)
    noiseDist = Normal(0, 0.01)
    signalModels = f
    λ = 2.0^-1.5
    H = 100
    kernel = GaussianRFF(H, λ * mean(y))
    ρ = 2.0^-20
    (ytrain, xtrain) = generatenoisydata(rng, T, xDists, noiseDist, signalModels)
    trainData = PERK.krr_train(xtrain, ytrain, kernel, ρ, randn(rng, H, 1),
        rand(rng, H))
    xhat = PERK.krr(y, trainData, kernel)

    error_rel = abs(xhat - xtrue) / xtrue
    return isapprox(error_rel, 0.07054474887124201, atol = 1e-7)

end

@testset "Kernel Ridge Regression" begin

    @test test_krr_1()
    @test test_krr_2()
    @test test_krr_3()
    @test test_krr_4()
    @test test_krr_5()
    @test test_krr_6()

end


"""
    fixed_knockoffs(X::Matrix{T}; method=:sdp)

Creates fixed knockoffs based on equation (2.2)-(2.4) of
"Controlling the false discovery rate via Knockoffs" by Barber and Candes (2015)

# Inputs
+ `X`: A `n × p` numeric matrix, each row is a sample, and each column is normalized
to mean 0 variance 1 with unit norm.
+ `method`: :equi for equi-distant knockoffs (eq 2.3) or :sdp for SDP
knockoffs (eq 2.4)

# Output
+ `Knockoff`: A struct containing the original `X` and its knockoff `X̃`,
in addition to other variables (e.g. `s`)
"""
function fixed_knockoffs(X::Matrix{T}, method::Symbol) where T <: AbstractFloat
    n, p = size(X)
    n ≥ 2p || error("fixed_knockoffs: only works for n ≥ 2p case! sorry!")
    # compute gram matrix using full svd
    U, σ, V = svd(X, full=true)
    Σ = V * Diagonal(σ)^2 * V'
    Σinv = V * inv(Diagonal(σ)^2) * V'
    # compute s vector using the specified method
    if method == :equi
        λmin = typemax(T)
        for σi in σ
            σi^2 < λmin && (λmin = σi^2)
        end
        s = min(1, 2λmin) .* ones(size(Σ, 1))
    elseif method == :sdp
        svar = Variable(p)
        problem = maximize(sum(svar), svar ≥ 0, 1 ≥ svar, 2Σ - Diagonal(svar) in :SDP)
        solve!(problem, () -> SCS.Optimizer(verbose=false))
        s = clamp.(evaluate(svar), 0, 1) # for numeric stability
    else
        error("fixed_knockoffs: method can only be :equi or :sdp")
    end
    # compute Ũ such that Ũ'X = 0
    Ũ = @view(U[:, p+1:2p])
    # compute C such that C'C = 2D - D*inv(Σ)*D via eigendecomposition (cholesky not stable)
    D = Diagonal(s)
    γ, P = eigen(2D - D*Σinv*D)
    clamp!(γ, 0, typemax(T)) # numerical stability
    C = Diagonal(sqrt.(γ)) * P
    # compute knockoffs
    X̃ = X * (I - Σinv*D) + Ũ * C
    return Knockoff(X, X̃, s, Σ, Σinv)
end


using OrdinaryDiffEq
import OrdinaryDiffEq: ODEProblem
using Petri

export linreg, peakgap, paramsweep, fluxes, u₀, makeplots_seir, makeplots_seird, savedata

#this function was removed in Julia 0.7
"""    linreg(x,y)

perform a simple linear regression for scalar case. Coefficients are returned
as (intercept, scalar)
"""
linreg(x, y) = hcat(fill!(similar(x), 1), x) \ y

"""    peakgap(s1, s2)

identify the time delay between the peak in two time series. Time series should
be stored as (x(t), t) for consistency with the tuples(sol) from OrdinaryDiffEq.
"""
function peakgap(s1, s2)
    @show p1, i1 = findmax(first.(s1))
    @show p2, i2 = findmax(first.(s2))
    return s2[i2][end] - s1[i1][end]
end

"""    peakgap(sol, i::Int, j::Int)

identify the time delay between the peak in two dimensions of an ODE solution structure.
"""
function peakgap(sol, i::Int, j::Int)
    s1 = [(u[i],t) for (u,t) in tuples(sol)]
    s2 = [(u[j],t) for (u,t) in tuples(sol)]
    return peakgap(s1,s2)
end

"""    paramsweep(f::Function, m::Petri.Model, initials, tspan, params)

solve an ODEProblem for each parameter setting in params and apply the function f to the solution.

Usage:
    paramsweep(m,initials,tspan,[p1,p2,p3]) do sol
        return f(sol)
    end
"""
function paramsweep(f::Function, m::Petri.Model, initials, tspan, params)
    map(params) do p
        prob = ODEProblem(fluxes(m), initials, tspan, p)
        sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
        return f(sol)
    end
end

"""    fluxes(m::Petri.Model)

a PetriNet interpreter that computes the mass action kinetics of a petri net.

Usage: pass this to ODEProblem to set up an ODE for a given model.
"""
function fluxes(m::Petri.Model)
    S = m.S
    T = m.Δ
    nS = length(S)
    nT = length(T)
    ϕ = zeros(Float64, nT)
    f(du, u, p, t) = begin
        for (i, t) in enumerate(T)
            ins = t[1]
            # TODO: accomodate multiplicites here
            ϕ[i] = p[i]*prod(u[ins])
        end
        for i in S
            du[i] = 0
        end
        for (i, t) in enumerate(T)
            ins = t[1]
            out = t[2]
            for s in ins
                # TODO: accomodate multiplicites here
                du[s] -= ϕ[i]
            end
            for s in out
                # TODO: accomodate multiplicites here
                du[s] += ϕ[i]
            end
        end
        return du
    end
    return f
end

u₀(m::Petri.Model) = begin
    zeros(Float64, length(m.S))
end

u₀(m::Petri.Model, initialS, initialI=1) = begin
    u0=zeros(Float64, length(m.S))
    u0[1] = initialS
    u0[2] = initialI
    return u0
end

function savedata(sol::ODESolution, colnames, fname::String)
    open(fname, "w") do fp
        println(fp, "time, $colnames")
        map(tuples(sol)) do (u,t)
            print(fp, "$t")
            map(u) do x
                print(fp, ",$x")
            end
            println(fp, "")
        end
    end
end


const NDIMS = 2

type Body
	mass::Float64
	pos::Vector{Float64}
	vel::Vector{Float64}

	initial_energy

	step
	prev_accel::Array{Float64, 2}

	function Body(mass=0.0, pos=Vector{Float64}(NDIMS), vel=Vector{Float64}(NDIMS))
		accel_hist_len = 4
		accel_hist = zeros(Float64, NDIMS, accel_hist_len)

		new(mass, pos, vel, 0, 0, accel_hist)
	end
end

######## IO ###########

import Base.show
function show(io::IO, a::Body)
	print(io, "mass: ", a.mass, "\n")
	print(io, "pos: ", join(a.pos, ", "), "\n")
	print(io, "vel: ", join(a.vel, ", "), "\n")
end

import Base.print
function print(io::IO, a::Body)
	@printf io "%24.16e\n" a.mass

	map(Array[a.pos, a.vel]) do b
		map(x-> @printf(io,"%24.16e",x), b)
		print(io, "\n")
	end
end

function read_body()
	mass = parse(Float64, readline())
	pos = map(x -> parse(Float64, x), split(readline()))
	vel = map(x -> parse(Float64, x), split(readline()))
	return Body(mass, pos, vel)
end

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

function accel(body::Body)
	r2 = dot(body.pos, body.pos)
	r3 = r2 * sqrt(r2) # we want r^(2/3)
	acc = body.pos * (-body.mass / r3)
	return acc
end

function jerk(body::Body)
	r2 = dot(body.pos, body.pos)
	r3 = r2 * sqrt(r2) # we want r^(2/3)
	j = (body.vel+body.pos*(-3*(dot(body.pos,body.vel))/r2))*(-body.mass/r3)
	return j
end

function prev_accel(body::Body, n::Int)
	return body.prev_accel[:,n]
end

function set_prev_accel!(body::Body, n::Int, prev::Vector)
	body.prev_accel[:,n] = prev
end

####### Energy #########

function kin_energy(body::Body)
	return 0.5 * dot(body.vel, body.vel)
end

function pot_energy(body::Body)
	return -body.mass / sqrt(dot(body.pos, body.pos))
end

function total_energy(body::Body)
	return kin_energy(body) + pot_energy(body)
end

function init_energy!(body::Body)
	body.initial_energy = total_energy(body)
end


@testset "Verify vander for ProHermiteBasis: evaluation, first and second derivatives" begin
    Ne = 200
    x = randn(Ne)
    #Test all the basis
    B = ProHermiteBasis(6)
    ktab = [0; 1; 2]
    for k in ktab
        # Test evaluation
        dV = vander(B, k, x)
        dVtrue = zeros(Ne, size(B))

        for i=1:B.m
            dVtrue[:,i] .= derivative(FamilyScaledProHermite[i], k , x)
        end
        @test norm(dV - dVtrue)<1e-6
    end
end

@testset "Verify vander for PhyHermiteBasis: evaluation, first and second derivatives" begin
    Ne = 200
    x = randn(Ne)
    #Test all the basis
    B = PhyHermiteBasis(6)
    ktab = [0; 1; 2]
    for k in ktab
        # Test evaluation
        dV = vander(B, k, x)
        dVtrue = zeros(Ne, size(B))

        for i=1:B.m
            dVtrue[:,i] .= derivative(FamilyScaledPhyHermite[i], k , x)
        end
        @test norm(dV - dVtrue)<1e-6
    end
end

@testset "Verify vander for CstProHermiteBasis: evaluation, first and second derivatives" begin
    Ne = 200
    x = randn(Ne)
    #Test all the basis
    B = CstProHermiteBasis(6)
    ktab = [0; 1; 2]
    for k in ktab
        # Test evaluation
        dV = vander(B, k, x)
        dVtrue = zeros(Ne, size(B))

        for i=1:B.m
            if i==1
                Pik = derivative(FamilyProPolyHermite[1], k)
                dVtrue[:,i] .= Pik.(x)
            else
                dVtrue[:,i] .= derivative(FamilyScaledProHermite[i-1], k , x)
            end
        end
        @test norm(dV - dVtrue)<1e-6
    end
end

@testset "Verify vander for CstPhyHermiteBasis: evaluation, first and second derivatives" begin
    Ne = 200
    x = randn(Ne)
    #Test all the basis
    B = CstPhyHermiteBasis(6)
    ktab = [0; 1; 2]

    for k in ktab
        # Test evaluation
        dV = vander(B, k, x)
        dVtrue = zeros(Ne, size(B))

        for i=1:B.m
            if i==1
                Pik = derivative(FamilyProPolyHermite[1], k)
                dVtrue[:,i] .= Pik.(x)
            else
                dVtrue[:,i] .= derivative(FamilyScaledPhyHermite[i-1], k , x)
            end
        end
        @test norm(dV - dVtrue)<1e-6
    end
end

@testset "Verify vander for CstLinProHermiteBasis: evaluation, first and second derivatives" begin
    Ne = 200
    x = randn(Ne)
    #Test all the basis
    B = CstLinProHermiteBasis(6)
    ktab = [0; 1; 2]

    for k in ktab
        # Test evaluation
        dV = vander(B, k, x)
        dVtrue = zeros(Ne, size(B))

        for i=1:B.m
            if i==1
                Pik = derivative(FamilyProPolyHermite[1], k)
                dVtrue[:,i] .= Pik.(x)
            elseif i==2
                Pik = derivative(FamilyProPolyHermite[2], k)
                dVtrue[:,i] .= Pik.(x)
            else
                dVtrue[:,i] .= derivative(FamilyScaledProHermite[i-2], k , x)
            end
        end
        @test norm(dV - dVtrue)<1e-6
    end
end

@testset "Verify vander for CstLinPhyHermiteBasis: evaluation, first and second derivatives" begin
    Ne = 200
    x = randn(Ne)
    #Test all the basis
    B = CstLinPhyHermiteBasis(6)
    ktab = [0; 1; 2]

    for k in ktab
        # Test evaluation
        dV = vander(B, k, x)
        dVtrue = zeros(Ne, size(B))

        for i=1:B.m
            if i==1
                Pik = derivative(FamilyProPolyHermite[1], k)
                dVtrue[:,i] .= Pik.(x)
            elseif i==2
                Pik = derivative(FamilyProPolyHermite[2], k)
                dVtrue[:,i] .= Pik.(x)
            else
                dVtrue[:,i] .= derivative(FamilyScaledPhyHermite[i-2], k , x)
            end
        end
        @test norm(dV - dVtrue)<1e-6
    end
end
#
#
# @testset "Verify vander for basis: evaluation, first and second derivatives" begin
#     Ne = 200
#     x = randn(Ne)
#     #Test all the basis
#     B = Basis(5)#CstProHermiteBasis(5; scaled = true)
#     # Btab = [CstProHermiteBasis(5; scaled = false);
#     #         CstProHermiteBasis(5; scaled = true)];
#             # CstPhyHermiteBasis(5; scaled = false);
#             # CstPhyHermiteBasis(5; scaled = true);
#             # CstLinPhyHermiteBasis(5; scaled = false);
#             # CstLinPhyHermiteBasis(5; scaled = true);
#             # CstLinProHermiteBasis(5; scaled = false);
#             # CstLinProHermiteBasis(5; scaled = true)]
#             # Test evaluation, derivative and second derivative
#     ktab = [0; 1; 2]
#
#     for k in ktab
#         # Test evaluation
#         dV = vander(B, k, x)
#         dVtrue = zeros(Ne, size(B))
#
#         for i=1:B.m
#             if i==1
#                 Pik = derivative(FamilyProPolyHermite[1], k)
#                 dVtrue[:,i] .= Pik.(x)
#             else
#                 dVtrue[:,i] .= derivative(FamilyScaledProHermite[i-1], k , x)
#             end
#         end
#         @test norm(dV - dVtrue)<1e-6
#     end
# end

# @testset "Test defined families list" begin
#
#     B = CstPhyHermiteBasis(10; scaled =false)
#     @test size(B,1)==12
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#
#     for i=2:size(B,1)
#         @test B[i] == FamilyPhyHermite[i-1]
#     end
#
#
#     B = CstPhyHermiteBasis(10; scaled =true)
#     @test size(B,1)==12
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#
#     for i=2:size(B,1)
#         @test B[i] == FamilyScaledPhyHermite[i-1]
#     end
#
#
#     B = CstProHermiteBasis(10; scaled =false)
#     @test size(B,1)==12
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#
#     for i=2:size(B,1)
#         @test B[i] == FamilyProHermite[i-1]
#     end
#
#
#     B = CstProHermiteBasis(10; scaled =true)
#     @test size(B,1)==12
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#
#     for i=2:size(B,1)
#         @test B[i] == FamilyScaledProHermite[i-1]
#     end
#
#     B = CstLinPhyHermiteBasis(10; scaled =false)
#     @test size(B,1)==13
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#     @test B[2] ==B.f[2] == FamilyProPolyHermite[2]
#
#
#     for i=3:size(B,1)
#         @test B[i] == FamilyPhyHermite[i-2]
#     end
#
#     B = CstLinPhyHermiteBasis(10; scaled = true)
#     @test size(B,1)==13
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#     @test B[2] ==B.f[2] == FamilyProPolyHermite[2]
#
#
#     for i=3:size(B,1)
#         @test B[i] == FamilyScaledPhyHermite[i-2]
#     end
#
#     B = CstLinProHermiteBasis(10; scaled =false)
#     @test size(B,1)==13
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#     @test B[2] ==B.f[2] == FamilyProPolyHermite[2]
#
#
#     for i=3:size(B,1)
#         @test B[i] == FamilyProHermite[i-2]
#     end
#
#     B = CstLinProHermiteBasis(10; scaled = true)
#     @test size(B,1)==13
#
#     @test B[1] ==B.f[1] == FamilyProPolyHermite[1]
#     @test B[2] ==B.f[2] == FamilyProPolyHermite[2]
#
#
#     for i=3:size(B,1)
#         @test B[i] == FamilyScaledProHermite[i-2]
#     end
# end

# @testset "Test evaluation of the basis" begin
#     # @testset "Test evaluation of the basis" begin
#         Ne = 10
#         x = 0.5*randn(Ne)
#
#         B = CstPhyHermiteBasis(5; scaled =false)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         for i=0:5
#             @test abs(Bx[i+2] -  PhyHermite(i)(x[1]))< 1e-8
#         end
#
#
#         B = CstPhyHermiteBasis(5; scaled = true)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         for i=0:5
#             @test abs(Bx[i+2] -  PhyHermite(i; scaled = true)(x[1]))< 1e-8
#         end
#
#         B = CstProHermiteBasis(5; scaled =false)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         for i=0:5
#             @test abs(Bx[i+2] -  ProHermite(i)(x[1]))< 1e-8
#         end
#
#
#         B = CstProHermiteBasis(5; scaled = true)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         for i=0:5
#             @test abs(Bx[i+2] -  ProHermite(i; scaled = true)(x[1]))< 1e-8
#         end
#
#         B = CstLinPhyHermiteBasis(5; scaled =false)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         @test Bx[2] == x[1]
#
#         for i=0:5
#             @test abs(Bx[i+3] -  PhyHermite(i)(x[1]))< 1e-8
#         end
#
#         B = CstLinPhyHermiteBasis(5; scaled =true)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         @test Bx[2] == x[1]
#
#         for i=0:5
#             @test abs(Bx[i+3] -  PhyHermite(i; scaled = true)(x[1]))< 1e-8
#         end
#
#         B = CstLinProHermiteBasis(5; scaled =false)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         @test Bx[2] == x[1]
#
#         for i=0:5
#             @test abs(Bx[i+3] -  ProHermite(i)(x[1]))< 1e-8
#         end
#
#         B = CstLinProHermiteBasis(5; scaled =true)
#         m = size(B,1)
#
#         Bx = B(x[1])
#         @test Bx[1] == 1.0
#         @test Bx[2] == x[1]
#
#         for i=0:5
#             @test abs(Bx[i+3] -  ProHermite(i; scaled = true)(x[1]))< 1e-8
#         end
# end


#
# finds local minimum between st1 and st3 along contour
# works by taking minimum over uniform sampling of points and then iterating once.
# returns t value of minimum along contour: tmin
# also returns integrand value at minimum: fmin
#
function locate_minimum(st1::Complex{Float64}, st3::Complex{Float64}, a::Float64, b::Float64, x::Float64, y::Float64, derivs::Bool)
    rst1,rst3 = real(st1),real(st3)
    t1,t2 = rst1,rst3
    tmin,fmin = 0.0,0.0
    # number of loops
    for k=1:n_its
        linspace!(lm_ts,t1,t2,n_t)
        # Precalculations
        for i=1:n_t
            temp1c=2.0(lm_ts[i] - rst3) + W
            temp2c=-4.0(lm_ts[i] - rst1) + V
            f = M_F_1 * atan(temp1c) - M_F_2
            g = M_G_1 * atan(temp2c) - M_G_2
            lm_gam[i] = complex(lm_ts[i],f*g)
        end
        integrand!(lm_integ,lm_integx,lm_integy,lm_gam,a,b,x,y,derivs,n_t)
        # function to be minimized: square of two norm of u, ux, uy
        for i=1:n_t
            test_min[i] = abs2(lm_integ[i])+abs2(lm_integx[i])+abs2(lm_integy[i])
        end

        fmin,idx = findmin(test_min)
        if idx == 1 || idx == n_t || k == n_its
            tmin = lm_ts[idx]
        else
            t1,t2 = lm_ts[idx-1],lm_ts[idx+1]
        end
    end
    return tmin,fmin
end


module Allen2012

using JuMP

# CODE ABANDONED: REASON: Seems like the family of constraints (17) that
# should prevent overlapping boxes does not work. There are more recent
# works with formulations that does not seem to share the same mistake,
# so we will use one of those.

"""
    build(model, pqr, LWH; allow_rot = false)
    build(model, p, q, r, L, W, H; allow_rot = false)

Add the variables and constraints of Allen 2012 3D-packing (maximize profit
of packed items, one single container, weakly heterogenous boxes, free rotation
can be allowed or not) to the model.

For more information see: Allen, Sam D., Edmund K. Burke, and Jakub Mareček.
“A Space-Indexed Formulation of Packing Boxes into a Larger Box.” Operations
Research Letters 40, no. 1 (January 1, 2012): 20–24.
https://doi.org/10.1016/j.orl.2011.10.008.

# Arguments
- `model`: JuMP model to which the variables and contraints will be added to.
- `pqr`: three sequences of numbers (with the same length) denoting
  respectively the length/width/height of the boxes/cartons
- `LWH`: three numbers denoting respectively the length/width/height of the
  container
- `A`: a sequence of numbers (with the same lenght as p/q/r) denoting the
  amount of each item available for packing
- `w`: a sequence of numbers (with the same lenght as p/q/r) denoting the
  worth/profit of each item
- `allow_rot`: if the model should use orientation variables
  TODO: add the name of the orientation variables here
"""
function build(model, p, q, r, L, W, H, A, w; allow_rot = false)
  build(model, (p, q, r), (L, W, H), A, w; allow_rot = allow_rot)
end
function build(model, pqr, LWH, A, w; allow_rot = false)
  # TODO: allow the caller to ask for rotation
  @assert !allow_rot

  @assert length(pqr) == 3
  @assert (isone ∘ length ∘ unique ∘ map)(length, pqr)
  L1, L2, L3 = pqr
  n = length(L1)
  @assert length(A) == N
  @assert length(w) == N
  #@assert length(LWH) == 3
  #@assert (isone ∘ length ∘ unique ∘ map)(length, LWH)
  DX, DY, DZ = LWH
  #m = length(LWH[1])
  #container_volumes = map(prod, zip(LWH...))
  # highest length/width/height of a container
  #ML, MW, MH = map((xs)->max(xs...), LWH)
  #ML, MW, MH = (10000, 10000, 10000)

  # To avoid creating unnecessary constraints, and also to simplify the code,
  # two lists are created. One contains the invalid positions and another the
  # valid ones. The objective function and constraints iterate such lists.
  invalid = Vector{NTuple{4,Int64}}()
  valid = Vector{NTuple{4,Int64}}()
  for t = 1:n, x = 1:DX, y = 1:DY, z = 1:DZ
    if (x > DX - L1[t] + 1) || (y > DY - L2[t] + 1) || (z > DZ - L3[t] + 1)
      push!(invalid, (t, x, y, z))
    else
      push!(valid, (t, x, y, z))
    end
  end

  @variables model begin
    u[1:n, 1:DX, 1:DY, 1:DZ], Bin
  end

  # Objective funtion Eq. (13)
  @objective(model, Max, sum(u[t,x,y,z] * w[t] for (t, x, y, z) = valid))

  # Constraint (14) -- the lower-left-bottom corner of two or more boxes do
  # not share the same point.
  for x = 1:DX, y = 1:DY, z = 1:DZ
    @constraint(model, sum(u[t, x, y, z] for t = 1:n) <= 1)
  end

  # Constraint (15) -- u[t,x,y,z] == 0 for all (x + L_{1t} > D_X) OR
  # (y + L_{2t} > D_Y) OR (z + L_{3t} > D_Z); i.e., the boxes stay inside
  # the container.
  for (t, x, y, z) = invalid
    setlowerbound(u[t, x, y, z], 0.0)
    setupperbound(u[t, x, y, z], 0.0)
  end

  # Constraint (16) -- we do not pack more boxes of some type than the ones
  # available.
  for t = 1:n
    @constraint(model,
      sum(u[t, x, y, z] for x = 1:DX, y = 1:DY, z = 1:DZ) <= A[t]
    )
  end

  # Constraint (17) -- boxes do not overlap.
  for (t, x, y, z) = valid
    @constraint(model, sum(u[t, x′, y′, z′] for x′ = x:(x + L1[t]),
      y′ = y:(y + L2[t]), z′ = z:(z + L2[t])) <= 1)
  end

  model
end


include("alternatingMinimization.jl")
include("convexRelaxation.jl")

struct node
    zero_indices::Array{Any, 1}
    one_indices::Array{Any, 1}
    upper_bound::Float64
    lower_bound::Float64
    relaxed_binary_matrix::Array{Float64, 2}
end

function isTerminal(one_indices, zero_indices, k_sparse, matrix_dim)
    """
    This function evaluates whether the current node in the Branch and Bound
    tree is a terminal node.

    :param zero_indices: List of 2-tuples of Int64 where each entry consists of
                         an index (i, j) such that Y_ij is constrained to take
                         value 0.
    :param one_indices: List of 2-tuples of Int64 where each entry consists of
                        an index (i, j) such that S_ij is constrained to take
                        value 1 where S is the binary matrix denoting the
                        sparsity pattern of Y.
    :param k_sparse: The maximum sparsity of the sparse matrix (Int64).
    :param matrix_dim: The row and column dimension of the data (Int64).

    :return: True if this is a terminal node, false otherwise (Bool).
    """

    num_fixed_ones = size(one_indices)[1]
    num_fixed_zeros = size(zero_indices)[1]

    if num_fixed_ones == k_sparse
        return true
    end

    if num_fixed_zeros == matrix_dim ^ 2 - k_sparse
        return true
    end

    return false

end

function SLR_BnB(U, mu, lambda, k_sparse, k_rank; epsilon=0.1,
    iter_min_improv=1e-12, symmetric_tree=false, output_to_file=false,
    output_file_name="temp.txt")
    """
    This function computes a certifiably near-optimal solution to the problem
        min ||U - X - Y||_F^2 + lambda * ||X||_F^2 + mu * ||Y||_F^2
        subject to rank(X) <= k_rank, ||Y||_0 <= k_sparse

    by executing a custom branch and bound algorithm.

    :param U: An arbitrary n-by-n matrix.
    :param mu: The regularization parameter for the sparse matrix penalty
                   (Float64).
    :param lambda: The regularization parameter for the low rank matrix penalty
                   (Float64).
    :param k_sparse: The maximum sparsity of the sparse matrix (Int64).
    :param k_rank: The maximum rank of the low rank matrix (Int64).
    :param epsilon: The Branch and Bound termination criterion. This function
                    will terminate with an epsilon optimal solution to the
                    optimization problem (Float64).
    :param iter_min_improv: The minimal fractional decrease in the objective
                            value required for the alternating minimization
                            procedure to continue iterating.
    :param symmetric_tree: If true, Branch and Bound will only explore symmetric
                           sparsity patterns (Bool).
    :param output_to_file: If true, progress output will be printed to file. If
                           false, progress output will be printed to the console
                           (Bool).
    :param output_file_name: Name of the file progress output will be printed to
                             when output_to_file=true (String).

    :return: This function returns eight values:
             1) A tuple of two n-by-n arrays that correspond to the globally
                optimal solution to the optimization problem. The first element
                in the tuple is the optimal low rank matrix and the second
                element in the tuple is the optimal sparse matrix.
             2) The final global upper bound (Float64).
             3) The final global lower bound (Float64).
             4) The number of nodes explored during the optimization process
                (Int64).
             5) A list of values of type Float64 representing the evolution of
                the upper bound during the optimization procedure.
             6) A list of values of type Float64 representing the evolution of
                the lower bound during the optimization procedure.
             7) The total time elapsed during the optimiation process
                (milliseconds).
             8) The number of terminal nodes explored during the optimization
                process (Int64).

                best feasible solution found
             by this method (the first element in the tuple is the matrix X and
             the second element is the matrix Y). The second value is the
             objective value (Float64) achieved by the returned feasible
             solution.
    """

    start_time = now()

    # Solve the root node
    init_zero_indices = []
    init_one_indices = []

    upper_bound_sol, upper_bound_obj = SLR_AM(U, mu, lambda, k_sparse, k_rank,
        random_restarts=0, zero_indices=init_zero_indices,
        one_indices=init_one_indices, min_improvement=iter_min_improv)

    lower_bound_sol, lower_bound_obj = solve_perspective_relaxation(U, mu,
        lambda, k_sparse, k_rank, zero_indices=init_zero_indices,
        one_indices=init_one_indices)

    root_optimality_gap = (upper_bound_obj - lower_bound_obj) / upper_bound_obj
    if output_to_file
        open(output_file_name, "a") do io
            write(io, "Root Node has been solved.\n")
            write(io, "Root Node upper bound is: $upper_bound_obj\n")
            write(io, "Root Node lower bound is: $lower_bound_obj\n")
            write(io, "Root Node optimality gap is: $root_optimality_gap\n\n")
        end
    else
        println("Root Node has been solved.")
        println("Root Node upper bound is: $upper_bound_obj")
        println("Root Node lower bound is: $lower_bound_obj")
        println("Root Node optimality gap is: $root_optimality_gap")
        println()
    end

    # Initialize the node list and global bounds
    master_node_list = []
    root_node = node(init_zero_indices,
                     init_one_indices,
                     upper_bound_obj,
                     lower_bound_obj,
                     lower_bound_sol[3])
    push!(master_node_list, root_node)
    global_upper_bound = upper_bound_obj
    global_lower_bound = lower_bound_obj
    global_X_sol = upper_bound_sol[1]
    global_Y_sol = upper_bound_sol[2]

    num_explored_nodes = 1
    upper_bound_hist = [global_upper_bound]
    lower_bound_hist = [global_lower_bound]

    n = size(U)[1]

    terminal_nodes = 0
    # Main branch and bound loop
    while (global_upper_bound - global_lower_bound) /
                                            global_upper_bound > epsilon

        """
        # Select the most recently added node (DFS)
        current_node = pop!(master_node_list)
        """
        current_node_index = findall(this_node->
            this_node.lower_bound == global_lower_bound, master_node_list)[1]
        current_node = master_node_list[current_node_index]
        deleteat!(master_node_list, current_node_index)

        # If the node is a terminal node, place it back on the master list
        if isTerminal(current_node.one_indices, current_node.zero_indices,
                      k_sparse, n)
            prepend!(master_node_list, current_node)
            continue
        end
        # Select entry to branch on
        index = argmin(abs.(current_node.relaxed_binary_matrix .- 0.5))
        i = index[1]
        j = index[2]

        # Construct the two child sets of indices from this parent node
        one_condition = size(current_node.one_indices)[1] <= k_sparse - 2
        zero_condition = size(current_node.zero_indices)[1] <= n^2 - k_sparse - 2
        sym_branch_feasible = one_condition && zero_condition
        if symmetric_tree && sym_branch_feasible && (i != j)
            new_index_zero = (append!(copy(current_node.zero_indices),
                              [(i, j), (j, i)]), copy(current_node.one_indices))
            new_index_one = (copy(current_node.zero_indices),
                             append!(copy(current_node.one_indices),
                             [(i, j), (j, i)]))
            new_index_list = [new_index_zero, new_index_one]
        else
            new_index_zero = (append!(copy(current_node.zero_indices), [(i, j)]),
                              copy(current_node.one_indices))
            new_index_one = (copy(current_node.zero_indices),
                             append!(copy(current_node.one_indices), [(i, j)]))
            new_index_list = [new_index_zero, new_index_one]
        end

        for (zero_indices, one_indices) in new_index_list

            upper_bound_sol, upper_bound_obj = SLR_AM(U, mu, lambda, k_sparse,
                        k_rank, random_restarts=0, zero_indices=zero_indices,
                        one_indices=one_indices,
                        min_improvement=iter_min_improv)

            if isTerminal(one_indices, zero_indices, k_sparse, n)
                lower_bound_sol = upper_bound_sol
                lower_bound_obj = upper_bound_obj
                relaxed_binary_matrix = abs.(lower_bound_sol[2]) .> 1e-10
                terminal_nodes = terminal_nodes + 1
            else
                lower_bound_sol, lower_bound_obj = solve_perspective_relaxation(
                    U, mu, lambda, k_sparse, k_rank, zero_indices=zero_indices,
                    one_indices=one_indices)
                relaxed_binary_matrix = lower_bound_sol[3]
            end

            if upper_bound_obj < global_upper_bound
                global_upper_bound = upper_bound_obj
                global_X_sol = upper_bound_sol[1]
                global_Y_sol = upper_bound_sol[2]

                current_gap = global_upper_bound - global_lower_bound
                current_gap = current_gap / global_upper_bound

                println("The new upper bound is: $global_upper_bound")
                println("The current lower bound is: $global_lower_bound")
                println("The current optimality gap is: $current_gap")
                println()

                if output_to_file
                    open(output_file_name, "a") do io
                        write(io, "The new upper bound is: $global_upper_bound\n")
                        write(io, "The current lower bound is: $global_lower_bound\n")
                        write(io, "The current optimality gap is: $current_gap\n\n")
                    end
                else
                    println("The new upper bound is: $global_upper_bound")
                    println("The current lower bound is: $global_lower_bound")
                    println("The current optimality gap is: $current_gap")
                    println()
                end

                new_master_node_list = []
                for this_node in master_node_list
                    if this_node.lower_bound < global_upper_bound
                        push!(new_master_node_list, this_node)
                    end
                end

                pre_size = size(master_node_list)[1] + 1
                post_size = size(new_master_node_list)[1] + 1

                if output_to_file
                    open(output_file_name, "a") do io
                        write(io, "Tree size before pruning: $pre_size\n")
                        write(io, "Tree size after pruning: $post_size\n\n")
                    end
                else
                    println("Tree size before pruning: $pre_size")
                    println("Tree size after pruning: $post_size")
                    println()
                end

                master_node_list = new_master_node_list
            end

            if lower_bound_obj < global_upper_bound
                new_node = node(zero_indices,
                                one_indices,
                                upper_bound_obj,
                                lower_bound_obj,
                                relaxed_binary_matrix)
                push!(master_node_list, new_node)
            end

            num_explored_nodes = num_explored_nodes + 1
            append!(upper_bound_hist, global_upper_bound)
            append!(lower_bound_hist, global_lower_bound)

            if num_explored_nodes % 200 == 0
                current_time = now() - start_time

                if output_to_file
                    open(output_file_name, "a") do io
                        write(io, "$num_explored_nodes nodes have been explored.\n")
                        write(io, "Current elapsed time: $current_time\n")
                        write(io, "The current lower bound is: $global_lower_bound\n")
                    end
                else
                    println("$num_explored_nodes nodes have been explored.")
                    println("Current elapsed time: $current_time")
                    println("The current lower bound is: $global_lower_bound")
                end
            end

        end

        if size(master_node_list)[1] == 0
            if output_to_file
                open(output_file_name, "a") do io
                    write(io, "All nodes have been explored")
                end
            else
                println("All nodes have been explored")
            end
            break
        end

        global_lower_bound = master_node_list[1].lower_bound
        for current_node in master_node_list
            if current_node.lower_bound < global_lower_bound
                global_lower_bound = current_node.lower_bound
            end
        end
    end

    total_time = now() - start_time

    if output_to_file
        open(output_file_name, "a") do io
            write(io, "An optimal solution has been found!")
        end
    else
        println("An optimal solution has been found!")
    end

    return (global_X_sol, global_Y_sol),
           global_upper_bound,
           global_lower_bound,
           num_explored_nodes,
           upper_bound_hist,
           lower_bound_hist,
           total_time,
           terminal_nodes

end;


using PyPlot
# == Set seed and generate a_t sequence == #
srand(123)
n = 100
a_seq = sin.(linspace(0, 5*pi, n)) + 2 + 0.1 * randn(n)

# == Model parameters == #
γ = 0.8
m = 1
d = γ * [1, -1]
h = 1.0

# == Initial conditions == #
y_m = [2.0]

testlq = LQFilter(d, h, y_m)
y_hist, L, U, y = optimal_y(testlq, a_seq)
y = y[end:-1:1]  # reverse y

# == Plot simulation results == #
fig, ax = subplots(figsize=(10, 6.1))
ax[:set_xlabel]("Time")

# == Some fancy plotting stuff -- simplify if you prefer == #
bbox = (0., 1.01, 1., .101)

time = 1:length(y)
ax[:set_xlim](0, maximum(time))
ax[:plot](time, a_seq / h, "k-o", ms=4, lw=2, alpha=0.6, label=L"$a_t$")
ax[:plot](time, y, "b-o", ms=4, lw=2, alpha=0.6, label=L"$y_t$")
ax[:grid]()
ax[:legend](ncol=2, bbox_to_anchor= bbox, loc = 3, mode = "expand", fontsize= 16)


abstract type Wavelet end

function minphase(wav::Array{T}) where T<:AbstractFloat
    n = length(wav)

    # zero pad
    wav = [wav; zeros(T,n)]

    # fft
    WAV = fft(wav)

    # log(|WAV|)
    logabswav = log.(abs.(WAV) .+ Base.eps(T))

    # Hilbert transform of log(|WAV|)
    hlogabswav = imag.(hilbert(logabswav)) # DSP.hilbert gives the analytic signal

    # min-phase:
    PHI = -hlogabswav

    # min-phase wavelet:
    WAV = abs.(WAV) .* exp.(im * PHI)
    wav = real.(ifft(WAV))
    wav[1:n]
end

function phaseshift(wav::Array{T}, phi::T) where T<:AbstractFloat
    if isapprox(phi, 0.0)
        return wav
    end
    n = length(wav)

    # zero pad
    wav = [wav; zeros(n)]

    # fft
    WAV = rfft(wav)

    # phase shift
    for i = 1:length(WAV)
        WAV[i] *= exp.(-im*phi)
    end
    irfft(WAV, length(wav))[1:n]
end

function costap(w::Wavelet, t::Array{T}) where T<:AbstractFloat
    if w.tmax < 0
        return ones(T,length(t))
    end
    itstrt = clamp(round(Int, .9*w.tmax/t[end]),1,length(t))
    itends = clamp(round(Int, (w.tmax/t[end])*length(t)),1,length(t))
    tap = ones(T,length(t))
    for i = 1:length(t)
        if itstrt <= i <= itends
            tap[i] = .5*(1.0 + cos(pi*(i-itstrt)/(itends-itstrt)))
        elseif i > itends
            tap[i] = 0.0
        end
    end
    tap
end

mutable struct WaveletSine <: Wavelet
    a::Float64
    f::Float64
    integrate::Bool
end
WaveletSine(;a = 1.0, f = 10.0, integrate = false) = WaveletSine(a, f, integrate)
copy(w::WaveletSine) = WaveletSine(w.a,w.f,w.integrate)

mutable struct WaveletRicker <: Wavelet
    a::Float64
    f::Float64
    tmax::Float64
    phase::Float64
    integrate::Bool
end
WaveletRicker(;a = 1.0, f = 15.0, tmax = -1.0, phase = 0.0, integrate = false) = WaveletRicker(a, f, tmax, phase, integrate)
copy(w::WaveletRicker) = WaveletRicker(w.a,w.f,w.tmax,w.phase,w.integrate)

mutable struct WaveletMinPhaseRicker <: Wavelet
    a::Float64
    f::Float64
    tmax::Float64
    integrate::Bool
end
WaveletMinPhaseRicker(;a = 1.0, f = 15.0, tmax = -1.0, integrate = false) = WaveletMinPhaseRicker(a, f, tmax, integrate)
copy(w::WaveletMinPhaseRicker) = WaveletMinPhaseRicker(w.a,w.f,w.tmax,w.integrate)

mutable struct WaveletDerivRicker <: Wavelet
    a::Float64
    f::Float64
    tmax::Float64
end
WaveletDerivRicker(;a = 1.0, f = 15.0, tmax = -1.0) = WaveletDerivRicker(a, f, tmax)
copy(w::WaveletDerivRicker) = WaveletDerivRicker(w.a,w.f,w.tmax)

mutable struct WaveletCausalRicker <: Wavelet
    a::Float64
    f::Float64
    tol::Float64
    tmax::Float64
    integrate::Bool
end
WaveletCausalRicker(;a = 1.0, f = 15.0, tol = 1.e-4, tmax = -1.0, integrate = false) = WaveletCausalRicker(a, f, tol, tmax, integrate)
copy(w::WaveletCausalRicker) = WaveletCausalRicker(w.a,w.f,w.tol,w.tmax,w.integrate)

mutable struct WaveletOrmsby <: Wavelet
    a::Float64
    f1::Float64
    f2::Float64
    f3::Float64
    f4::Float64
    tmax::Float64
    integrate::Bool
end
WaveletOrmsby(;a = 1.0, f1 = 5.0, f2 = 10.0, f3 = 40.0, f4 = 45.0, tmax = -1.0, integrate = false) = WaveletOrmsby(a, f1, f2, f3, f4, tmax, integrate)
copy(w::WaveletOrmsby) = WaveletOrmsby(w.a,w.f1,w.f2,w.f3,w.f4,w.tmax,w.integrate)

mutable struct WaveletMinPhaseOrmsby <: Wavelet
    a::Float64
    f1::Float64
    f2::Float64
    f3::Float64
    f4::Float64
    tmax::Float64
    integrate::Bool
end
WaveletMinPhaseOrmsby(;a = 1.0, f1 = 5.0, f2 = 10.0, f3 = 40.0, f4 = 45.0, tmax = -1.0, integrate = false) = WaveletMinPhaseOrmsby(a, f1, f2, f3, f4, tmax, integrate)
copy(w::WaveletMinPhaseOrmsby) = WaveletMinPhaseOrmsby(w.a,w.f1,w.f2,w.f3,w.f4,w.tmax,w.integrate)

function get(w::WaveletSine, t::Array{T}) where T<:AbstractFloat
    wav = T(w.a) .* sin.(T(2 * pi * w.f) .* t)
    w.integrate ? cumsum(wav, dims=1) : wav
end

function get(w::WaveletRicker, t::Array{T}) where T<:AbstractFloat
    wav = convert(Array{T},w.a * (1.0 .- 2.0 * (pi * w.f * t).^2) .* exp.(-(pi * w.f * t).^2))
    wav = phaseshift(wav, T(w.phase))
    wav .*= costap(w,t)
    w.integrate ? cumsum(wav, dims=1) : wav
end

function get(w::WaveletMinPhaseRicker, t::Array{T}) where T<:AbstractFloat
    wav = minphase(get(WaveletRicker(w.a,w.f,-1.0,0.0,false),t .- mean(t)))
    wav .*= costap(w,t)
    w.integrate ? cumsum(wav, dims=1) : wav
end

function get(w::WaveletDerivRicker, t::Array{T}) where T<:AbstractFloat
    wav1 = get(WaveletRicker(w.a,w.f,-1.0,0.0,false),t)
    wav = similar(wav1)
    wav[1] = wav1[1]
    for i = 2:length(wav1)
        wav[i] = wav1[i] - wav1[i-1]
    end
    wav .* costap(w,t)
end

function get(w::WaveletCausalRicker, t::Array{T}) where T<:AbstractFloat
    wav = get(WaveletRicker(w.a,w.f,-1.0,0.0,false),t .- mean(t))
    idx = findall(ww->abs(ww)>w.tol*w.a,wav)[1]
    rng = idx:length(wav)
    wav[1:length(rng)] = wav[rng]
    wav[length(rng)+1:end] .= 0.0
    wav .*= costap(w,t)
    w.integrate ? cumsum(wav, dims=1) : wav
end

function get(w::WaveletOrmsby, t::Array{T}) where T<:AbstractFloat
    numerator(f,t) = sinc.(f*t).^2 .* (pi*f)^2
    pf43 = T(pi*(w.f4-w.f3))
    pf21 = T(pi*(w.f2-w.f1))
    wav = numerator(T(w.f4),t)/pf43 - numerator(T(w.f3),t)/pf43 - numerator(T(w.f2),t)/pf21 + numerator(T(w.f1),t)/pf21
    wav ./= maximum(abs, wav)
    wav .*= w.a
    wav .-= mean(wav)
    wav = wav .* costap(w,t)
    w.integrate ? cumsum(wav, dims=1) : wav
end

function get(w::WaveletMinPhaseOrmsby, t::Array{T}) where T<:AbstractFloat
    wav = minphase(get(WaveletOrmsby(w.a,w.f1,w.f2,w.f3,w.f4,w.tmax,false),t .- mean(t)))
    wav ./= maximum(abs, wav)
    wav .*= w.a
    wav .-= mean(wav)
    wav .*= costap(w,t)
    w.integrate ? cumsum(wav, dims=1) : wav
end


# Spatial calculations

export
    vector1D,
    vector,
    wrapcoords,
    wrapcoordsvec

"""
    vector1D(c1, c2, side_length)

Displacement between two 1D coordinate values from c1 to c2, accounting for
the bounding box.
The minimum image convention is used, so the displacement is to the closest
version of the coordinate accounting for the periodic boundaries.
"""
function vector1D(c1, c2, side_length)
    if c1 < c2
        return (c2 - c1) < (c1 - c2 + side_length) ? (c2 - c1) : (c2 - c1 - side_length)
    else
        return (c1 - c2) < (c2 - c1 + side_length) ? (c2 - c1) : (c2 - c1 + side_length)
    end
end

"""
    vector(c1, c2, box_size)

Displacement between two coordinate values, accounting for the bounding box.
The minimum image convention is used, so the displacement is to the closest
version of the coordinates accounting for the periodic boundaries.
"""
vector(c1, c2, box_size) = vector1D.(c1, c2, box_size)

@generated function vector(c1::SVector{N}, c2::SVector{N}, box_size) where N
    quote
        Base.Cartesian.@ncall $N SVector{$N} i->vector1D(c1[i], c2[i], box_size[i])
    end
end

sqdistance(i, j, coords, box_size) = sum(abs2, vector(coords[i], coords[j], box_size))

"""
    wrapcoords(c, side_length)

Ensure a 1D coordinate is within the simulation box and return the coordinate.
"""
wrapcoords(c, side_length) = c - floor(c / side_length) * side_length

"""
    wrapcoordsvec(c, box_size)

Ensure a coordinate is within the simulation box and return the coordinate.
"""
wrapcoordsvec(v, box_size) = wrapcoords.(v, box_size)


using Grid
using Cubature
using Distributions
using Iterators

iters = 0
function f(x)
    global iters
    iters += 1
    prod(pdf(Poisson(3.0), floor(x)))
end

function gen_grid(mins, maxs, N=30)
    D = length(mins)
    stepsizes = ceil(float(maxs .- mins) ./N)
    ranges = [mins[i]:stepsizes[i]:maxs[i] for i=1:D]
    grid = hcat(map(collect, collect(product(ranges...)))...)'
    return map(length,ranges), stepsizes, grid
end

const tol = 0.001
const upp = 100

println("####################")
@time val = hquadrature(f, [0], [upp], abstol=tol)
println(val)
println("1D used $iters iters")

println("####################")
iters = 0
@time val = hcubature(f, zeros(2), zeros(2).+upp, abstol=tol)
println(val)
println("2D used $iters iters")

println("####################")
iters = 0
@time val = hcubature(f, zeros(3), zeros(3).+upp, abstol=tol)
println(val)
println("3D used $iters iters")

println("####################")
println("Now using Grid...")
println("####################")

valsi = nothing
iiters = 0
function fi(x)
    global iters, iiters
    v,g,h = valgradhess(valsi,x...)
    if abs(prod(g)) > 1e-15 || abs(det(h)) > 1e-15
        f(x)
    else
        iiters += 1
        valsi[x...]
    end
end

N = 20

println("####################")
_,_,pts = gen_grid([0],[upp], N)
@time vals = [f(pt) for pt in pts]
valsi = InterpGrid(vals, BCnearest, InterpQuadratic)
iters = 0
iiters = 0
@time val = hquadrature(fi, [0], [upp], abstol=tol)
println(val)
println("1D used $iters actual iters, $iiters interp iters")

println("####################")
lens,_,pts = gen_grid(zeros(2), zeros(2).+upp, N)
@time vals = [f(pts[i,:]) for i=1:size(pts,1)]
vals = reshape(vals, lens...)
valsi = InterpGrid(vals, BCnearest, InterpQuadratic)
iters = 0
iiters = 0
@time val = hcubature(fi, zeros(2), zeros(2).+upp, abstol=tol)
println(val)
println("2D used $iters actual iters, $iiters interp iters")

#println("####################")
#lens,_,pts = gen_grid(zeros(3), zeros(3).+upp, N)
#@time vals = [f(pts[i,:]) for i=1:size(pts,1)]
#vals = reshape(vals, lens...)
#valsi = InterpGrid(vals, BCnearest, InterpQuadratic)
#iters = 0
#iiters = 0
#@time val = hcubature(fi, zeros(3), zeros(3).+upp, abstol=tol)
#println(val)
#println("3D used $iters actual iters, $iiters interp iters")

hess = reshape([det(valgradhess(valsi, pts[i,:]...)[3]) for i=1:size(pts,1)], lens...);
grad = reshape([prod(valgrad(valsi, pts[i,:]...)[2]) for i=1:size(pts,1)], lens...);


# A* shortest-path algorithm

#################################################
#
#  A* shortest-path algorithm
#
#################################################

module AStar

# The enqueue! and dequeue! methods defined in Base.Collections (needed for
# PriorityQueues) conflict with those used for queues. Hence we wrap the A*
# code in its own module.

using Graphs
using Base.Collections
using Compat

export shortest_path

function a_star_impl!{V,D}(
    graph::AbstractGraph{V},# the graph
    frontier,               # an initialized heap containing the active vertices
    colormap::Vector{Int},  # an (initialized) color-map to indicate status of vertices
    edge_dists::AbstractEdgePropertyInspector{D},  # cost of each edge
    heuristic::Function,    # heuristic fn (under)estimating distance to target
    t::V)  # the end vertex

    while !isempty(frontier)
        (cost_so_far, path, u) = dequeue!(frontier)
        if u == t
            return path
        end

        for edge in out_edges(u, graph)
            v = target(edge)
            if colormap[v] < 2
                colormap[v] = 1
                new_path = cat(1, path, edge)
                path_cost = cost_so_far + edge_property(edge_dists, edge, graph)
                enqueue!(frontier,
                        (path_cost, new_path, v),
                        path_cost + heuristic(v))
            end
        end
        colormap[u] = 2
    end
    nothing
end

function shortest_path{V,E,D}(
    graph::AbstractGraph{V,E},  # the graph
    edge_dists::AbstractEdgePropertyInspector{D},      # cost of each edge
    s::V,                       # the start vertex
    t::V,                       # the end vertex
    heuristic::Function = n -> 0)
            # heuristic (under)estimating distance to target
    frontier = VERSION < v"0.4-" ? PriorityQueue{@compat(Tuple{D,Array{E,1},V}),D}() : PriorityQueue(@compat(Tuple{D,Array{E,1},V}),D)
    frontier[(zero(D), E[], s)] = zero(D)
    colormap = zeros(Int, num_vertices(graph))
    colormap[s] = 1
    a_star_impl!(graph, frontier, colormap, edge_dists, heuristic, t)
end

function shortest_path{V,E,D}(
    graph::AbstractGraph{V,E},  # the graph
    edge_dists::Vector{D},      # cost of each edge
    s::V,                       # the start vertex
    t::V,                       # the end vertex
    heuristic::Function = n -> 0)
    edge_len::AbstractEdgePropertyInspector{D} = VectorEdgePropertyInspector(edge_dists)
    shortest_path(graph, edge_len, s, t, heuristic)
end

end

using .AStar


function crc32(crc::Int, str::String)
    table = zeros(UInt32, 256)

    for i in 0:255
        tmp = i
        for j in 0:7
            if tmp & 1 == 1
                tmp >>= 1
                tmp ⊻= 0xedb88320
            else
                tmp >>= 1
            end
        end

        table[i + 1] = tmp
    end

    crc ⊻= 0xffffffff

    for i in UInt32.(collect(str))
        crc = (crc >> 8) ⊻ table[(crc & 0xff) ⊻ i + 1]
    end

    crc ⊻ 0xffffffff
end

str = "The quick brown fox jumps over the lazy dog"
crc = crc32(0, str)
assert(crc == 0x414fa339)
println("Message: ", str)
println("Checksum: ", hex(crc))


using ApproxFun, SingularIntegralEquations, Test

@testset "Fractal" begin
    d = Segment(0.,1.)
    c = Circle()

    for set in (:cantor,:thincantor,:thinnercantor,:thinnestcantor)
        @eval begin
            d = Segment(0.,1.)
            c = Circle()
            @test $set(d,0) == d
            @test $set(d,0,4) == d
            @test $set(d,1) == Segment(0.0,1/3) ∪ Segment(2/3,1.0)
            @test $set(d,1,4) == Segment(0.0,1/4) ∪ Segment(3/4,1.0)
            @test $set(c,0) == c
            @test $set(c,1) == Arc(0.0,1.0,(-2.6179938779914944,-0.523598775598299)) ∪
                                Arc(0.0,1.0,(0.5235987755982987,2.6179938779914944))
        end
    end

    @test cantor(d,2) == Segment(0.0,1/9)∪Segment(2/9,1/3)∪Segment(2/3,7/9)∪Segment(8/9,1.0)
    @test cantor(d,2,4) == Segment(0.0,1/16)∪Segment(3/16,1/4)∪Segment(3/4,13/16)∪Segment(15/16,1.0)

    @test thincantor(d,2) == Segment(0.0,1/12)∪Segment(1/4,1/3)∪Segment(2/3,3/4)∪Segment(11/12,1.0)
    @test thincantor(d,2,4) == Segment(0.0,1/20)∪Segment(1/5,1/4)∪Segment(3/4,4/5)∪Segment(19/20,1.0)

    @test thinnercantor(d,2) == Segment(0.0,1/27)∪Segment(8/27,1/3)∪Segment(2/3,19/27)∪Segment(26/27,1.0)
    @test thinnercantor(d,2,4) == Segment(0.0,1/64)∪Segment(15/64,1/4)∪Segment(3/4,49/64)∪Segment(63/64,1.0)
    @test thinnercantor(d,3,4) == Segment(0.0,1/4096)∪Segment(63/4096,1/64)∪Segment(15/64,961/4096)∪Segment(1023/4096,1/4)∪Segment(3/4,3073/4096)∪Segment(3135/4096,49/64)∪Segment(63/64,4033/4096)∪Segment(4095/4096,1/1)

    @test thinnestcantor(d,2) == Segment(0.0,1/27)∪Segment(8/27,1/3)∪Segment(2/3,19/27)∪Segment(26/27,1.0)
    @test thinnestcantor(d,2,4) == Segment(0.0,1/64)∪Segment(15/64,1/4)∪Segment(3/4,49/64)∪Segment(63/64,1.0)
    @test thinnestcantor(d,3,4) == Segment(0.0,1/16384)∪Segment(255/16384,1/64)∪Segment(15/64,3841/16384)∪Segment(4095/16384,1/4)∪Segment(3/4,12289/16384)∪Segment(12543/16384,49/64)∪Segment(63/64,16129/16384)∪Segment(16383/16384,1/1)

    @test smithvolterracantor(d,0) == d
    @test smithvolterracantor(d,1) == Segment(0.0,3/8)∪Segment(5/8,1.0)
    @test smithvolterracantor(d,2) == Segment(0.0,5/32)∪Segment(7/32,3/8)∪Segment(5/8,25/32)∪Segment(27/32,1.0)
end


# SPDX-License-Identifier: X11
# 2020-08-31
# Tsundoku (300pt)

function main()
  N, M, K = parse.(Int, split(readline(), " "))
  A = parse.(Int, split(readline(), " "))
  B = parse.(Int, split(readline(), " "))
  As = runsum(A)
  Bs = runsum(B)

  j = M
  ans = 0
  for i ∈ 0:N
    As[i + 1] > K && break

    while j > 0 && As[i + 1] + Bs[j + 1] > K
      j -= 1
    end

    ans = max(ans, i + j)
  end

  println(ans)
end

function runsum(A)
  As = Vector{Int}(undef, 1 + lastindex(A))
  As[1] = 0
  for i ∈ 1:lastindex(A)
    As[i + 1] = As[i] + A[i]
  end
  return As
end

main()


"""
    sysinv = inv(sys; atol = 0, atol1 = atol, atol2 = atol, rtol, checkinv = true)

Compute for a descriptor system `sys = (A-λE,B,C,D)` with the transfer function matrix `G(λ)`, a descriptor realization of its inverse system
`sysinv = (Ai-λEi,Bi,Ci,Di)`, such that the transfer function matrix `Ginv(λ)` of `sysinv` is the inverse of `G(λ)` (i.e., `G(λ)*Ginv(λ) = I`).
The realization of `sysinv` is determined using inversion-free formulas and the invertibility condition is checked, unless `checkinv = false`.

The keyword arguments `atol1`, `atol2` and `rtol` specify, respectively,
the absolute tolerance for the nonzero elements of `A`, `B`, `C`, `D`, the absolute tolerance for the nonzero elements of `E`,
and the relative tolerance for the nonzero elements of `A`, `B`, `C`, `D` and `E`.
The default relative tolerance is `n*ϵ`, where `n` is the order of the square matrices `A` and `E`, and  `ϵ` is the working machine epsilon.
The keyword argument `atol` can be used to simultaneously set `atol1 = atol` and `atol2 = atol`.
"""
function inv(sys::DescriptorStateSpace{T}; checkinv::Bool = true,
             atol::Real = zero(real(T)), atol1::Real = atol, atol2::Real = atol,
             rtol::Real = (size(sys.A,1)+1)*eps(real(float(one(T))))*iszero(min(atol1,atol2))) where T
   p, m = sys.ny, sys.nu
   n = sys.nx
   m == p  || error("The system must have the same number of inputs and outputs")
   Ai = [sys.A sys.B; sys.C sys.D]
   Bi = [zeros(T,n,m); -Matrix{T}(I,m,m)]
   Ci = [zeros(T,m,n) Matrix{T}(I,m,m)]
   Di = zeros(T,m,m)
   Ei = (sys.E == I) ? blockdiag(Matrix{T}(I,n,n),zeros(T,m,m)) :
                       blockdiag(sys.E,zeros(T,m,m))
   checkinv && (MatrixPencils.isregular(Ai, Ei, atol1 = atol1, atol2 = atol2, rtol = rtol ) ||
                error("The system is not invertible"))
   return DescriptorStateSpace{T}(Ai, Ei, Bi, Ci, Di, sys.Ts)
end
"""
    sysldiv = ldiv(sys1, sys2; atol = 0, atol1 = atol, atol2 = atol, rtol = n*ϵ, checkinv = true)
    sysldiv = sys1 \\ sys2

Compute for the descriptor systems `sys1 = (A1-λE1,B1,C1,D1)` with the transfer function matrix `G1(λ)` and
`sys2 = (A2-λE2,B2,C2,D2)` with the transfer function matrix `G2(λ)`, a descriptor realization
`sysldiv = (Ai-λEi,Bi,Ci,Di)` of `sysldiv = inv(sys1)*sys2`,
whose transfer-function matrix `Gli(λ)` represents the result of the left division `Gli(λ) = inv(G1(λ))*G2(λ)`.
The realization of `sysldiv` is determined using inversion-free formulas and the invertibility condition for `sys1` is checked,
unless `checkinv = false`.

The keyword arguments `atol1`, `atol2` and `rtol` specify, respectively,
the absolute tolerance for the nonzero elements of `A1`, `B1`, `C1`, `D1`, `A2`, `B2`, `C2`, `D2`,
the absolute tolerance for the nonzero elements of `E1` and `E2`,
and the relative tolerance for the nonzero elements of `A1`, `B1`, `C1`, `D1`, `A2`, `B2`, `C2`, `D2`, `E1` and `E2`.
The default relative tolerance is `n*ϵ`, where `n` is the maximum of orders of the square matrices `A1` and `A2`, and  `ϵ` is the working machine epsilon.
The keyword argument `atol` can be used to simultaneously set `atol1 = atol` and `atol2 = atol`.
"""
function ldiv(sys1::DescriptorStateSpace{T1}, sys2::DescriptorStateSpace{T2};
              atol::Real = zero(real(T1)), atol1::Real = atol, atol2::Real = atol,
              rtol::Real = (max(size(sys1.A,1),size(sys2.A,1))*eps(real(float(one(T1)))))*iszero(min(atol1,atol2))) where {T1,T2}
   p1, m1 = sys1.ny, sys1.nu
   m1 == p1  || error("The system SYS1 must be square")
   p2, m2 = sys2.ny, sys2.nu
   p1 == p2  || error("The systems SYS1 and SYS2 must have the same number of outputs")
   sys1.Ts == sys2.Ts ||  error("The systems SYS1 and SYS2 must have same sampling time")
   T = promote_type(T1,T2)
   n1 = sys1.nx
   n2 = sys2.nx
   if n1 > 0
      A, E, B, C, D = dssdata(T,[sys1 sys2])
      Ai = [A B[:,1:m1]; C D[:,1:m1]]
      Ei = E == I ? [I zeros(T,n1+n2,m1); zeros(T,m1,n1+n2+m1)]  : blockdiag(E,zeros(T,m1,m1))
      MatrixPencils.isregular(Ai, Ei, atol1 = atol1, atol2 = atol2, rtol = rtol) ||
                 error("The system SYS1 is not invertible")
      Bi = [B[:,m1+1:m1+m2]; D[:,m1+1:m1+m2]]
      Ci = [zeros(T,p1,n1+n2) -I]
      Di = zeros(T,p1,m2)
   else
      D1 = copy_oftype(sys1.D,T)
      LUD = lu(D1)
      (norm(D1,Inf) <= atol1 || rcond(LUD.U) <= 10*m1*eps(real(float(one(T))))) &&
               error("The system SYS2 is not invertible")
      Ai, Ei, Bi, Ci, Di = dssdata(T,sys2)
      ldiv!(LUD,Ci); ldiv!(LUD,Di)
   end

   return DescriptorStateSpace{T}(Ai, Ei, Bi, Ci, Di, sys2.Ts)
end
function (\)(sys1::DescriptorStateSpace, sys2::DescriptorStateSpace; kwargs...)
   ldiv(sys1,sys2; kwargs...)
end
"""
    sysrdiv = rdiv(sys1, sys2; atol = 0, atol1 = atol, atol2 = atol, rtol = n*ϵ, checkinv = true)
    sysrdiv = sys1 / sys2

Compute for the descriptor systems `sys1 = (A1-λE1,B1,C1,D1)` with the transfer function matrix `G1(λ)` and
`sys2 = (A2-λE2,B2,C2,D2)` with the transfer function matrix `G2(λ)`, a descriptor realization
`sysrdiv = (Ai-λEi,Bi,Ci,Di)` of `sysrdiv = sys1*inv(sys2)`,
whose transfer-function matrix `Gri(λ)` represents the result of the right division `Gri(λ) = G1(λ)*inv(G2(λ))`.
The realization of `sysrdiv` is determined using inversion-free formulas and the invertibility condition for `sys2` is checked,
unless `checkinv = false`.

The keyword arguments `atol1`, `atol2` and `rtol` specify, respectively,
the absolute tolerance for the nonzero elements of `A1`, `B1`, `C1`, `D1`, `A2`, `B2`, `C2`, `D2`,
the absolute tolerance for the nonzero elements of `E1` and `E2`,
and the relative tolerance for the nonzero elements of `A1`, `B1`, `C1`, `D1`, `A2`, `B2`, `C2`, `D2`, `E1` and `E2`.
The default relative tolerance is `n*ϵ`, where `n` is the maximum of orders of the square matrices `A1` and `A2`,
and  `ϵ` is the working machine epsilon.
The keyword argument `atol` can be used to simultaneously set `atol1 = atol` and `atol2 = atol`.
"""
function rdiv(sys1::DescriptorStateSpace{T1}, sys2::DescriptorStateSpace{T2};
              atol::Real = zero(real(T1)), atol1::Real = atol, atol2::Real = atol,
              rtol::Real = max(size(sys1.A,1),size(sys2.A,1))*eps(real(float(one(T1))))*iszero(min(atol1,atol2))) where {T1,T2}
   p2, m2 = sys2.ny, sys2.nu
   m2 == p2  || error("The system SYS2 must be square")
   p1, m1 = sys1.ny, sys1.nu
   m1 == m2  || error("The systems SYS1 and SYS2 must have the same number of inputs")
   sys1.Ts == sys2.Ts ||  error("The systems SYS1 and SYS2 must have same sampling time")
   T = promote_type(T1,T2)
   n1 = sys1.nx
   n2 = sys2.nx
   if n2 > 0
      A, E, B, C, D = dssdata(T,[sys2; sys1])
      Ai = [A B; C[1:p2,:] D[1:p2,:]]
      Ei = E == I ? [I zeros(T,n1+n2,p2); zeros(T,p2,n1+n2+p2)] : blockdiag(E,zeros(T,p2,p2))
      MatrixPencils.isregular(Ai, Ei, atol1 = atol1, atol2 = atol2, rtol = rtol) ||
                 error("The system SYS2 is not invertible")
      Ci = [C[p2+1:p1+p2,:] D[p2+1:p1+p2,:]]
      Bi = [zeros(T,n1+n2,m1); -I]
      Di = zeros(T,p1,m1)
   else
      D2 = copy_oftype(sys2.D,T)
      LUD = lu(D2)
      (norm(D2,Inf) <= atol1 || rcond(LUD.U) <= 10*m1*eps(real(float(one(T))))) &&
                  error("The system SYS2 is not invertible")
      Ai, Ei, Bi, Ci, Di = dssdata(T,sys1)
      rdiv!(Bi,LUD); rdiv!(Di,LUD)
   end

   return DescriptorStateSpace{T}(Ai, Ei, Bi, Ci, Di, sys1.Ts)
end
function (/)(sys1::DescriptorStateSpace, sys2::DescriptorStateSpace; kwargs...)
    rdiv(sys1,sys2; kwargs...)
end
"""
    sysdual = gdual(sys, rev = false)
    sysdual = transpose(sys, rev = false)

Compute for a descriptor system `sys = (A-λE,B,C,D)` with the transfer function matrix `G(λ)`,
the descriptor system realization of its dual system
`sysdual = (Ad-λEd,Bd,Cd,Dd)`, where `Ad = transpose(A)`, `Ed = transpose(E)`, `Bd = transpose(C)`,
`Cd = transpose(B)` and `Dd = transpose(D)`,
such that the transfer function matrix `Gdual(λ)` of `sysdual` is the transpose of `G(λ)`
(i.e., `Gdual(λ) = transpose(G(λ))`).

If `rev = true`, the tranposition is combined with the reverse permutation of the state variables, such that
`sysdual = (P*Ad*P-λP*Ed*P,P*Bd,Cd*P,Dd)`, where `P` is the permutation matrix with ones down the second diagonal.
"""
function gdual(SYS::DescriptorStateSpace{T};kwargs...) where T
    transpose(SYS;kwargs...)
end
function transpose(SYS::DescriptorStateSpace{T}; rev = false) where T
   if rev
       return DescriptorStateSpace{T}(reverse(reverse(transpose(SYS.A),dims=1),dims=2),
                                      SYS.E == I ? I : reverse(reverse(transpose(SYS.E),dims=1),dims=2),
                                      reverse(transpose(SYS.C),dims=1), reverse(transpose(SYS.B),dims=2), copy(transpose(SYS.D)), SYS.Ts)
   else
       return DescriptorStateSpace{T}(copy(transpose(SYS.A)), SYS.E == I ? I : copy(transpose(SYS.E)), copy(transpose(SYS.C)), copy(transpose(SYS.B)),
                                      copy(transpose(SYS.D)), SYS.Ts)
   end
end

"""
    sysconj = ctranspose(sys)
    sysconj = sys'

Compute the conjugate transpose (or adjoint) of a descriptor system (see [`adjoint`](@ref)).
"""
function ctranspose(sys::DescriptorStateSpace{T}) where T
   if sys.Ts == 0
      # continuous-time case
      return DescriptorStateSpace{T}(copy(-sys.A'), sys.E == I ? I : copy(sys.E'), copy(-sys.C'), copy(sys.B'), copy(sys.D'), sys.Ts)
   else
      # discrete-time case
      (nx, ny, nu) = (sys.nx, sys.ny, sys.nu)
      n = nx+ny
      return DescriptorStateSpace{T}(sys.E == I ? Matrix{T}(I,n,n) : blockdiag(sys.E',Matrix{T}(I,ny,ny)),
                                     [sys.A' sys.C'; zeros(T,ny,n)],
                                     [zeros(T,nx,ny) ; -I],
                                     [sys.B' zeros(T,nu,ny)], copy(sys.D'), sys.Ts)
   end
end
"""
    sysconj = adjoint(sys)
    sysconj = sys'

Compute for a descriptor system `sys = (A-λE,B,C,D)` with the transfer function matrix `G(λ)`,
the descriptor system realization of its adjoint (also called _conjugate transpose_) system
`sysconj = (Ac-λEc,Bc,Cc,Dc)`, such that the transfer function matrix `Gconj(λ)` of `sysconj`
is the appropriate conjugate transpose of `G(λ)`, as follows:
for a continuous-time system with `λ = s`, `Gconj(s) := transpose(G(-s))`, while
for a discrete-time system with `λ = z`, `Gconj(z) := transpose(G(1/z))`.
"""
function adjoint(sys::DescriptorStateSpace)
    ctranspose(sys)
end


using PyPlot
using JLD
function plot_errorbar()
	X = load("../data/sens/lorenz/dJds_vs_params.jld")
	rho_arr = X["s"]
	dJds_arr = X["dJds"]
	err_arr = sqrt.(X["var_dJds"])/4
	@show err_arr
	sel = err_arr .<= 0.2
	rho_arr = rho_arr[sel]
	err_arr = err_arr[sel]
	dJds_arr = dJds_arr[sel]

	fig, ax = subplots(1,1)
	ax.plot(rho_arr, dJds_arr, "X", ms=10.0)
	ax.errorbar(rho_arr, dJds_arr, yerr=err_arr, linewidth=3.0, color="blue",fmt="none")
    ax.xaxis.set_tick_params(labelsize=28)
    ax.yaxis.set_tick_params(labelsize=28)
    ax.set_xlabel(L"$s$",fontsize=28)
	ax.set_ylabel(L"$d_{s}\langle J\rangle$",fontsize=28)
	ax.grid(true)

end
function plot_lines()
	X = load("../data/sens/lorenz/dvarz_ds2_800.jld")
	rho = X["rho"]
	dJds = X["dJds"]
	err = sqrt.(X["var_dJds"])/4

	@show err
	X = load("../data/obj_erg_avg/lorenz/zm28sq.jld")
    s2 = X["s2"]
    J = X["J"]
    fig, ax = subplots(1,1)
    ax.plot(s2, J, ".", ms=25.0)
    ax.xaxis.set_tick_params(labelsize=28)
    ax.yaxis.set_tick_params(labelsize=28)
    ax.set_xlabel(L"$s$",fontsize=28)
    ax.set_ylabel(L"$\langle J\rangle$",fontsize=28)
    ax.grid(true)

    eps = 2e-1

	X = load("../data/obj_erg_avg/lorenz/zm28sq_dJds.jld")
	J = X["J"]
	s2 = X["s2"]

	sel = rho .<= 30.5

	dJds = dJds[sel]
	err = err[sel]

	J = J[sel]
	s2 = s2[sel]
	rho = rho[sel]
	s2 = s2[err .<= 1.8]
	J = J[err .<= 1.8]
    dJds = dJds[err .<= 1.8]
	n = size(dJds)[1]
	J_pts = reshape([J .- eps*dJds J .+ eps*dJds], n,
                    2)'
    s_pts = reshape([s2 .- eps s2 .+ eps], n,
                    2)'

    ax.plot(s_pts[:,1:n-1], J_pts[:,1:n-1], "k",lw=3.0)
    ax.plot(s_pts[:,n], J_pts[:,n], "k",lw=4.0,label=L"$d_s\langle J\rangle$ from S3")

    ax.legend(fontsize=28)

	X = load("../data/sens/lorenz/dvarz_ds2_1500.jld")
	rho = X["rho"]
	dJds = X["dJds"]
	err = sqrt.(X["var_dJds"])/4
	@show err
	X = load("../data/obj_erg_avg/lorenz/zm28sq_dJds.jld")
	J = X["J"]
	s2 = X["s2"]

	sel = .!(30 .>= rho .>= 29)
	dJds = dJds[sel]
	err = err[sel]
	s2 = s2[sel]
	J = J[sel]

	s2 = s2[err .<= 1.8]
	J = J[err .<= 1.8]
    dJds = dJds[err .<= 1.8]
	n = size(dJds)[1]
	J_pts = reshape([J .- eps*dJds J .+ eps*dJds], n,
                    2)'
    s_pts = reshape([s2 .- eps s2 .+ eps], n,
                    2)'

    ax.plot(s_pts[:,1:n-1], J_pts[:,1:n-1], "k",lw=3.0)
    ax.plot(s_pts[:,n], J_pts[:,n], "k",lw=4.0,label=L"$d_s\langle J\rangle$ from S3")

end


n,m = 3,2

# Custom Equality Constraint
p1 = 3
c(v,x,u) = begin v[1]  = x[1]^2 + x[2]^2 - 5; v[2:3] =  u - ones(2,1) end
c(v,x) = c(v,x,zeros(2))
jacob_c(x,u) = [2x[1] 2x[2] 0 0 0;
                0     0     0 1 0;
                0     0     0 0 1];
jacob_c(x) = [2x[1] 2x[2] 0;
                0     0     0;
                0     0     0];

v = zeros(p1)
x = [1.,2,3]
u = [-5.,5]
c(v,x,u)
@test v == [0,-6,4]

c(v,x)
@test v == [0, -1, -1]

J = zeros(p1,n+m)
J_term = zeros(p1,n)
ForwardDiff.jacobian!(J_term,c,v,x)
@test J_term == jacob_c(x)

import TrajectoryOptimization: generate_jacobian
∇c, = generate_jacobian(c, n, m, p1)
methods(∇c)
∇c(J,x,u)
@test J == jacob_c(x,u)
@test ∇c(J_term,x) == jacob_c(x)

# Test constraint function
con = Constraint{Equality}(c,n,m,p1,:custom)
evaluate!(v, con, x, u)
@test v == [0,-6,4]
evaluate!(v, con, x)
@test v == [0, -1, -1]
@test length(con) == p1

# Test constraint jacobian
con.∇c(J,x,u);
jacobian!(J, con, x, u)
@test J == jacob_c(x,u)
jacobian!(J_term, con, x)
@test J_term == jacob_c(x)

# Custom inequality constraint
p2 = 2
c2(v,x,u) = begin v[1] = sin(x[1]); v[2] = sin(x[3]) end
∇c2(Z,x,u) = begin Z[1,1] = cos(x[1]); Z[2,3] = cos(x[3]); end
con2 = Constraint{Inequality}(c2,∇c2,n,m,p2,:ineq)

# Bound constraint
x_max = [5,5,Inf]
x_min = [-10,-5,0.]
u_max = 0.
u_min = -10.
#
# p3 = 2(n+m)
# bnd = BoundConstraint(n,m,x_max=x_max,x_min=x_min,u_min=u_min,u_max=u_max, trim=false)
# v = zeros(p3)
# evaluate!(v, bnd, x, u)
# @test v == [-4,-3,-Inf,-5,5,-11,-7,-3,-5,-15]
# C = PartedArray(zeros(p3,n+m),create_partition2((p3,),(n,m),Val((:xx,:xu))))
# jacobian!(C, bnd, x, u)
# @test C.xx == [Diagonal(1I,n); zeros(m,n); -Diagonal(1I,n); zeros(m,n)]
# @test C.xu == [zeros(n,m); Diagonal(1I,m); zeros(n,m); -Diagonal(1I,m)]
# @test length(bnd) == p3
#
# v = zeros(2n)
# evaluate!(v, bnd, x)
# @test v == [-4,-3,-Inf,-11,-7,-3]
# C = zeros(2n, n)
#
# C = zeros(2n-1,n)
# bnd.active.x_all
# bnd.jac[bnd.active.x_all,1:n]
# bnd.active.x_min
# jacobian!(C, bnd, x)
# @test jacobian!(C, bnd, x) == [Diagonal(1.0I,n)[1:(n-1),:]; -Diagonal(1.0I,n)]
# @test length(bnd,:terminal) == 2n

# Trimmed bound constraint
bnd = BoundConstraint(n,m,x_max=x_max,x_min=x_min,u_min=u_min,u_max=u_max)
p3 = 2(n+m)-1
p3_N = 2n-1
v = zeros(p3)
evaluate!(v, bnd, x, u)
@test v == [-4,-3,-5,5,-11,-7,-3,-5,-15]
C = zeros(p3, n+m)
jac_bnd = [Diagonal(1I,n+m)[[true;true;false;true;true],:]; -Diagonal(1I,n+m)]
@test jacobian!(C, bnd, x, u) == jac_bnd
@test length(bnd,:stage) == p3

v = zeros(p3_N)
evaluate!(v, bnd, x)
@test v == [-4,-3,-11,-7,-3]
C = zeros(p3_N, n)
jacobian!(C, bnd, x)
jac_bnd_term = [Diagonal(1I,n)[[true;true;false],:]; -Diagonal(1I,n)]
@test C == jac_bnd_term

# Create Constraint Set
C = [con,con2,bnd]
@test C isa ConstraintSet
C = con + con2 + bnd
@test C isa ConstraintSet

@test findall(C,Inequality) == [false,true,true]
@test split(C) == ([con2,bnd],[con,])
@test count_constraints(C) == (p2+p3,p1)
@test inequalities(C) == [con2,bnd]
@test equalities(C) == [con,]
@test bounds(C) == [bnd,]
@test labels(C) == [:custom,:ineq,:bound]
@test TrajectoryOptimization.is_terminal.(C) == [true,false,true]

# Terminal Constraint
cterm(v,x) = begin v[1] = x[1] - 5; v[2] = x[1]*x[2] end
∇cterm(x) = [1 0 0; x[2] x[1] 0]
∇cterm(A,x) = copyto!(A,∇cterm(x))
p_term = 2
v = zeros(p_term)
cterm(v,x)
con_term = Constraint{Equality}(cterm,∇cterm,n,p_term,:terminal)
v2 = zeros(p_term)
con_term.c(v2,x)
@test v == v2
A = zeros(p_term,n)
@test con_term.∇c(A,x) == ∇cterm(x)

C_term = [con,bnd,con_term]
C2 = [con,con2,bnd,con_term]
@test C2 isa ConstraintSet
@test terminal(C2) == C_term
@test terminal(C_term) == C_term
@test stage(C2) == C
count_constraints(C2)
@test split(C2) == ([con2,bnd],[con,con_term])
@test split(C2) == (inequalities(C2),equalities(C2))
@test bounds(C2) == [bnd]
@test labels(C2) == [:custom,:ineq,:bound,:terminal]
@test TrajectoryOptimization.num_stage_constraints(C2) == p1+p2+p3
@test TrajectoryOptimization.num_terminal_constraints(C2) == p3_N + p_term + p1

import TrajectoryOptimization: num_stage_constraints, num_terminal_constraints
num_stage_constraints(C2)
num_constraints(C2,:stage)

v_stage = PartedVector(C2)
v_term = PartedVector(C2,:terminal)
v_stage2 = PartedVector(C2,:stage)
v_term2 = PartedVector(C2,:terminal)
evaluate!(v_stage,C,x,u)
evaluate!(v_term,C_term,x)
evaluate!(v_stage2,C2,x,u)
evaluate!(v_term2,C2,x)
@test v_stage == v_stage2
@test v_term == v_term2

c_jac = PartedMatrix(C,n,m)
@test size(c_jac) == (p1+p2+p3,n+m)
@test size(c_jac.custom) == (p1,n+m)
@test size(c_jac.x) == (p1+p2+p3,n)
@test size(c_jac.u) == (p1+p2+p3,m)
jacobian!(c_jac, C, x, u)
@test c_jac.custom == jacob_c(x,u)
@test c_jac.bound == jac_bnd

c_jac = PartedMatrix(C2,n,m,:terminal)
@test size(c_jac) == (p1+p3_N+p_term,n)
@test size(c_jac.custom) == (p1,n)
@test size(c_jac.ineq) == (0,n)
@test size(c_jac.u) == (p1+p3_N+p_term,0)
@test size(c_jac.terminal) == (p_term,n)
jacobian!(c_jac, C2, x)
jacobian!(c_jac.custom, con, x)
@test c_jac.custom == jacob_c(x)
@test c_jac.bound == jac_bnd_term

# Augment State
m_inf = m+n
con_inf = TrajectoryOptimization.infeasible_constraints(n,m)
u_inf = [u; 5; -5; 10]
v = zeros(n)
con_inf.c(v,x,u_inf[con_inf.inds[2]])
@test v == [5,-5,10]

C_inf = [con,con2,bnd,con_inf]
v_stage = PartedVector(stage(C))
v_inf = PartedVector(stage(C_inf))
TrajectoryOptimization.evaluate!(v_stage,C,x,u_inf)
TrajectoryOptimization.evaluate!(v_inf,C_inf,x,u_inf)
@test v_inf == [v_stage;5;-5;10]

C2
C3 = copy(C2)
@test TrajectoryOptimization.remove_bounds!(C3) == [bnd]
@test C3 == [con,con2,con_term]
@test pop!(C3,:ineq) == con2
@test C3 == [con,con_term]

# Constraints
CS = [[con], con+con2+bnd, con+con2+bnd, con+bnd]
N = length(CS)
@test CS isa Vector{<:ConstraintSet}
PC = Constraints(CS)
PC[1] += bnd
@test PC[1] == [con,bnd]
@test num_constraints(PC) == [p1+p3, p1+p2+p3, p1+p2+p3, p1+p3_N]
PC[N] += con_term
@test num_constraints(PC) == [p1+p3, p1+p2+p3, p1+p2+p3, p1+p3_N+p_term]
PC2 = copy(PC)
pop!(PC2[2],:ineq)
@test num_constraints(PC2) == [p1+p3, p1+p3, p1+p2+p3, p1+p3_N+p_term]
@test num_constraints(PC) == [p1+p3, p1+p2+p3, p1+p2+p3, p1+p3_N+p_term]

Cval, ∇Cval = TrajectoryOptimization.init_constraint_trajectories(PC,n,m,N)
@test length.(Cval) == num_constraints(PC)
X = [float.(x) for k = 1:N]
U = [float.(u) for k = 1:N-1]
TrajectoryOptimization.update_constraints!(Cval, PC, X, U)
v = zeros(p1)
c(v,x,u)
@test Cval[1].custom == v
c(v,x)
@test Cval[N].custom == v

jacobian!(∇Cval, PC, X, U)
@test ∇Cval[N] == c_jac