"""
    tiles = Tiler(areawidth, areaheight, nrows, ncols, margin=20)

A Tiler is an iterator that, for each iteration, returns a tuple of:

- the `x`/`y` point of the center of each tile in a set of tiles that divide up a
  rectangular space such as a page into rows and columns (relative to current 0/0)

- the number of the tile

`areawidth` and `areaheight` are the dimensions of the area to be tiled, `nrows`/`ncols`
are the number of rows and columns required, and `margin` is applied to all four
edges of the area before the function calculates the tile sizes required.

Tiler and Partition are similar:

- Partition lets you specify the width and height of a cell

- Tiler lets you specify how many rows and columns of cells you want, and a margin:

```
tiles = Tiler(1000, 800, 4, 5, margin=20)
for (pos, n) in tiles
    # the point pos is the center of the tile
end
```

You can access the calculated tile width and height like this:

```
tiles = Tiler(1000, 800, 4, 5, margin=20)
for (pos, n) in tiles
    ellipse(pos.x, pos.y, tiles.tilewidth, tiles.tileheight, :fill)
end
```

It's sometimes useful to know which row and column you're currently on.
`tiles.currentrow` and `tiles.currentcol` should have that information for you.

To use a Tiler to make grid points:

```
first.(collect(Tiler(800, 800, 4, 4)))
```

which returns an array of points that are the center points of the grid.
"""
mutable struct Tiler
    areawidth::Real
    areaheight::Real
    tilewidth::Real
    tileheight::Real
    nrows::Int
    ncols::Int
    currentrow::Int
    currentcol::Int
    margin::Real
    function Tiler(areawidth, areaheight, nrows::Int, ncols::Int; margin=10)
        tilewidth  = (areawidth - 2margin)/ncols
        tileheight = (areaheight - 2margin)/nrows
        currentrow=1
        currentcol=1
        new(areawidth, areaheight, tilewidth, tileheight, nrows, ncols, currentrow, currentcol, margin)
    end
end

function Base.iterate(pt::Tiler)
    x = -(pt.areawidth/2)  + pt.margin + (pt.tilewidth/2)
    y = -(pt.areaheight/2) + pt.margin + (pt.tileheight/2)
    tilenumber = 1
    x1 = x + pt.tilewidth
    y1 = y
    if x1 > (pt.areawidth/2) - pt.margin
        y1 += pt.tileheight
        x1 = -(pt.areawidth/2) + pt.margin + (pt.tilewidth/2)
    end
    pt.currentrow, pt.currentcol = (div(tilenumber-1, pt.ncols)+1, mod1(tilenumber, pt.ncols))
    return ((Point(x, y), tilenumber), (Point(x1, y1), tilenumber + 1))
end

function Base.iterate(pt::Tiler, state)
    if state[2] > (pt.nrows * pt.ncols)
        return
    end
    x = state[1].x
    y = state[1].y
    tilenumber = state[2]
    x1 = x + pt.tilewidth
    y1 = y
    if x1 > (pt.areawidth/2) - pt.margin
        y1 += pt.tileheight
        x1 = -(pt.areawidth/2) + pt.margin + (pt.tilewidth/2)
    end
    pt.currentrow, pt.currentcol = (div(tilenumber-1, pt.ncols)+1, mod1(tilenumber, pt.ncols))
    return ((Point(x, y), tilenumber), (Point(x1, y1), tilenumber + 1))
end

function Base.length(pt::Tiler)
    pt.nrows * pt.ncols
end

function Base.getindex(pt::Tiler, i::Int)
    1 <= i <= pt.ncols *  pt.nrows || throw(BoundsError(pt, i))
    xcoord = -pt.areawidth/2  + pt.margin + mod1(i, pt.ncols) * pt.tilewidth  - pt.tilewidth/2
    ycoord = -pt.areaheight/2 + pt.margin + (div(i - 1,  pt.ncols) * pt.tileheight) + pt.tileheight/2
    return (Point(xcoord, ycoord), i)
end

function Base.size(pt::Tiler)
    return (pt.nrows, pt.ncols)
end

Base.lastindex(pt::Tiler) = length(pt)
Base.firstindex(pt::Tiler) = 1

Base.eltype(::Type{Tiler}) = Tuple

Base.getindex(pt::Tiler, I) = [pt[i] for i in I]

"""
    GridRect(startpoint, xspacing, yspacing, width, height)

Define a rectangular grid, to start at `startpoint` and proceed along the x-axis in
steps of `xspacing`, then along the y-axis in steps of `yspacing`.

    GridRect(startpoint, xspacing=100.0, yspacing=100.0, width=1200.0, height=1200.0)

For a column, set the `xspacing` to 0:

    grid = GridRect(O, 0, 40)

To get points from the grid, use `nextgridpoint(g::Grid)`.

    julia> grid = GridRect(O, 0, 40);
    julia> nextgridpoint(grid)
    Luxor.Point(0.0, 0.0)

    julia> nextgridpoint(grid)
    Luxor.Point(0.0, 40.0)

When you run out of grid points, you'll wrap round and start again.
"""
mutable struct GridRect
    startpoint::Point
    currentpoint::Point
    xspacing::Float64
    yspacing::Float64
    width::Float64
    height::Float64
    rownumber::Int
    colnumber::Int
    function GridRect(startpoint, xspacing, yspacing, width=1200.0, height=1200.0)
        rownumber = 1
        colnumber = 0
        # find the "previous" point, so that the first call gets the real first point
        xspacing == 0 ? sp = startpoint.y - yspacing : sp = startpoint.y
        currentpoint = Point(startpoint.x - xspacing, sp)
        new(startpoint, currentpoint, xspacing, yspacing, width, height, rownumber, colnumber)
    end
end

"""
    GridHex(startpoint, radius, width=1200.0, height=1200.0)

Define a hexagonal grid, to start at `startpoint` and proceed along the x-axis and
then along the y-axis, `radius` is the radius of a circle that encloses each hexagon.
The distance in `x` between the centers of successive hexagons is:

``\\frac{\\sqrt{(3)} radius}{2}``

To get the next point from the grid, use `nextgridpoint(g::Grid)`.

When you run out of grid points, you'll wrap round and start again.
"""
mutable struct GridHex
    startpoint::Point
    radius::Float64
    currentpoint::Point
    width::Float64
    height::Float64
    rownumber::Int
    colnumber::Int
    function GridHex(startpoint, radius, width=1200.0, height=1200.0)
        rownumber = 1
        colnumber = 0
        # find the "previous" point, so that the first call gets the real first point
        currentpoint = Point(startpoint.x - (sqrt(3) * radius), startpoint.y)
        new(startpoint, radius, currentpoint, width, height, rownumber, colnumber)
    end
end

"""
    nextgridpoint(g::GridRect)

Returns the next available (or even the first) grid point of a grid.
"""
function nextgridpoint(g::GridRect)
    temp = Point(g.currentpoint.x + g.xspacing, g.currentpoint.y)
    if g.xspacing == 0
        g.rownumber += 1
        g.colnumber = 1
        g.currentpoint = Point(g.startpoint.x, temp.y + g.yspacing)
    else
        g.colnumber += 1
        g.currentpoint = temp
    end
    if abs(g.startpoint.x - temp.x) >= g.width # next row?
        g.rownumber += 1
        g.colnumber = 1
        g.currentpoint = Point(g.startpoint.x, temp.y + g.yspacing)
    end
    if g.currentpoint.y >= g.height   # finished, start again?
        g.rownumber = 1
        g.colnumber = 1
        g.currentpoint = Point(g.startpoint.x, g.startpoint.y)
    end
    return g.currentpoint
end

"""
    nextgridpoint(g::GridHex)

Returns the next available grid point of a hexagonal grid.
"""
function nextgridpoint(g::GridHex)
    g.currentpoint = Point(g.currentpoint.x + (g.radius * sqrt(3)), g.currentpoint.y)
    g.colnumber += 1
    if g.currentpoint.x >= g.width                 # next row
        g.rownumber += 1
        g.colnumber = 1
        if g.rownumber % 2 == 0
            g.currentpoint = Point(g.startpoint.x + (sqrt(3) * g.radius)/2, g.currentpoint.y + (3/2) * g.radius)
        else
            g.currentpoint = Point(g.startpoint.x, g.currentpoint.y + (3/2) * g.radius)
        end
    end
    if g.currentpoint.y >= g.height     # finished?
        g.rownumber = 1                 # start again...
        g.colnumber = 1
        g.currentpoint = Point(g.startpoint.x, g.startpoint.y)
    end
    return g.currentpoint
end

"""
    p = Partition(areawidth, areaheight, tilewidth, tileheight)

A Partition is an iterator that, for each iteration, returns a tuple of:

- the `x`/`y` point of the center of each tile in a set of tiles that divide up a
rectangular space such as a page into rows and columns (relative to current 0/0)

- the number of the tile

`areawidth` and `areaheight` are the dimensions of the area to be tiled,
`tilewidth`/`tileheight` are the dimensions of the tiles.

Tiler and Partition are similar:

- Partition lets you specify the width and height of a cell

- Tiler lets you specify how many rows and columns of cells you want, and a margin

    tiles = Partition(1200, 1200, 30, 30)
    for (pos, n) in tiles
        # the point pos is the center of the tile
    end

You can access the calculated tile width and height like this:

    tiles = Partition(1200, 1200, 30, 30)
    for (pos, n) in tiles
        ellipse(pos.x, pos.y, tiles.tilewidth, tiles.tileheight, :fill)
    end

It's sometimes useful to know which row and column you're currently on:

    tiles.currentrow
    tiles.currentcol

should have that information for you.

Unless the tilewidth and tileheight are exact multiples of the area width and height, you'll
see a border at the right and bottom where the tiles won't fit.
"""
mutable struct Partition
    areawidth::Real
    areaheight::Real
    tilewidth::Real
    tileheight::Real
    nrows::Int
    ncols::Int
    currentrow::Int
    currentcol::Int
    function Partition(areawidth, areaheight, tilewidth, tileheight)
        ncols = areawidth  / tilewidth  |> floor |> Integer
        nrows = areaheight / tileheight |> floor |> Integer
        if ncols < 1 || nrows < 1
            throw(error("partition(): not enough space for tiles that size"))
        end
        currentrow = 1
        currentcol = 1
        new(areawidth, areaheight, tilewidth, tileheight, nrows, ncols, currentrow, currentcol)
    end
end

function Base.iterate(pt::Partition)
    x = -(pt.areawidth/2)  + (pt.tilewidth/2)
    y = -(pt.areaheight/2) + (pt.tileheight/2)
    tilenumber = 1
    x1 = x + pt.tilewidth
    y1 = y
    if (x1 + pt.tilewidth/2) > (pt.areawidth/2)
        y1 += pt.tileheight
        x1 = -(pt.areawidth/2) + (pt.tilewidth/2)
    end
    pt.currentrow, pt.currentcol = (div(tilenumber-1, pt.ncols)+1, mod1(tilenumber, pt.ncols))
    return ((Point(x, y), tilenumber), (Point(x1, y1), tilenumber + 1))
end

function Base.iterate(pt::Partition, state)
    if state[2] > (pt.nrows * pt.ncols)
        return
    end
    x = state[1].x
    y = state[1].y
    tilenumber = state[2]
    x1 = x + pt.tilewidth
    y1 = y
    if (x1 + pt.tilewidth/2) > (pt.areawidth/2)
        y1 += pt.tileheight
        x1 = -(pt.areawidth/2) + (pt.tilewidth/2)
    end
    pt.currentrow, pt.currentcol = (div(tilenumber-1, pt.ncols)+1, mod1(tilenumber, pt.ncols))
    return ((Point(x, y), tilenumber), (Point(x1, y1), tilenumber + 1))
end

function Base.length(pt::Partition)
    pt.nrows * pt.ncols
end

function Base.getindex(pt::Partition, i::Int)
    1 <= i <= pt.ncols *  pt.nrows || throw(BoundsError(pt, i))
    xcoord = -pt.areawidth/2  + mod1(i, pt.ncols) * pt.tilewidth  - pt.tilewidth/2
    ycoord = -pt.areaheight/2 + (div(i - 1,  pt.ncols) * pt.tileheight) + pt.tileheight/2
    return (Point(xcoord, ycoord), i)
end

"""
    box(tiles::Tiler, n::Integer; action=:none, vertices=false, reversepath=false)
    box(tiles::Tiler, n::Integer, action::Symbol=:none; vertices=false, reversepath=false)

Draw a box in tile `n` of tiles `tiles`.
"""
function box(tiles::Tiler, n::Integer;
        action=:none,
        vertices=false,
        reversepath=false)
    tilew, tileh = tiles.tilewidth, tiles.tileheight
    if vertices == true || action == :none
        box(first(tiles[n]), tilew, tileh, action=:none, vertices=true, reversepath=reversepath)
    else
        box(first(tiles[n]), tilew, tileh, action=action, vertices=false, reversepath=reversepath)
    end
end

box(tiles::Tiler, n::Integer, action::Symbol; vertices=false, reversepath=false) =
        box(tiles, n, action=action, vertices=vertices, reversepath=reversepath)


using OrdinaryDiffEq, Test, RecursiveArrayTools, DiffEqDevTools, Statistics

u0 = fill(0., 2)
v0 = ones(2)
function f1_harmonic(dv,v,u,p,t)
  dv .= -u
end
function f2_harmonic(du,v,u,p,t)
  du .= v
end
function harmonic_analytic(y0, p, x)
  v0, u0 = y0
  ArrayPartition(-u0*sin(x) + v0*cos(x),u0*cos(x) + v0*sin(x))
end
ff_harmonic = DynamicalODEFunction(f1_harmonic,f2_harmonic;analytic=harmonic_analytic)
prob = DynamicalODEProblem(ff_harmonic,v0,u0,(0.0,5.0))

sol = solve(prob,SymplecticEuler(),dt=1/2)
sol_verlet = solve(prob,VelocityVerlet(),dt=1/100)
sol_ruth3 = solve(prob,Ruth3(),dt=1/100)

interp_time = 0:0.001:5
interp = sol(0.5)
interps = sol(interp_time)

sol_tsit5 = solve(prob,Tsit5())

prob = SecondOrderODEProblem(f1_harmonic,v0,u0,(0.0,5.0))

sol2 = solve(prob,SymplecticEuler(),dt=1/2)
sol2_verlet = solve(prob,VelocityVerlet(),dt=1/100)
sol2_ruth3 = solve(prob,Ruth3(),dt=1/100)

sol2_verlet(0.1)

@test sol[end][1] == sol2[end][1]
@test sol_verlet[end][1] == sol2_verlet[end][1]
@test sol_ruth3[end][1] == sol2_ruth3[end][1]
@test sol[end][3] == sol2[end][3]
@test sol_verlet[end][3] == sol2_verlet[end][3]
@test sol_ruth3[end][3] == sol2_ruth3[end][3]

prob = DynamicalODEProblem(ff_harmonic,v0,u0,(0.0,5.0))
println("Convergence tests")

dts = 1 .//2 .^(6:-1:3)
# Symplectic Euler
sim = test_convergence(dts,prob,SymplecticEuler(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 1 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 1 rtol = 1e-1
# Verlet
sim = test_convergence(dts,prob,VelocityVerlet(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1
# Test that position converges faster for Verlet
position_error = :final => [mean(sim[i].u[2].x[1] - sim[i].u_analytic[2].x[1]) for i in 1:length(sim)]
@test first(DiffEqDevTools.calc𝒪estimates(position_error).second) ≈ 4.0 rtol=1e-1

# 2nd Order Tableaus
sim = test_convergence(dts,prob,VerletLeapfrog(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1
sim = test_convergence(dts,prob,PseudoVerletLeapfrog(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1
sim = test_convergence(dts,prob,McAte2(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1

# Ruth
sim = test_convergence(dts,prob,Ruth3(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 3 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 3 rtol = 1e-1
sim = test_convergence(dts,prob,McAte3(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 3 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 3 rtol = 1e-1
sim = test_convergence(dts,prob,CandyRoz4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,McAte4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,McAte42(),dense_errors=true)
@test_broken sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test_broken sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,CalvoSanz4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

dts = 1 .//2 .^(4:-1:0)
sim = test_convergence(dts,prob,McAte5(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sim = test_convergence(dts,prob,Yoshida6(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 6 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4.69 rtol = 1e-1
sim = test_convergence(dts,prob,KahanLi6(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 6 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sim = test_convergence(dts,prob,McAte8(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 8 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sim = test_convergence(dts,prob,KahanLi8(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 8 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sol = solve(prob,Nystrom4(),dt=1/1000)

# Nyström method
dts = 1 .//2 .^(9:-1:6)
sim = test_convergence(dts,prob,Nystrom4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,Nystrom4VelocityIndependent(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,IRKN3(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 3 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 3 rtol = 1e-1
sim = test_convergence(dts,prob,IRKN4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
dts = 1.0 ./2.0 .^(5:-1:0)
sim = test_convergence(dts,prob,Nystrom5VelocityIndependent(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

dts = 1.0 ./2.0 .^(2:-1:-2)
sim = test_convergence(dts,prob,SofSpa10(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 10 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

# Methods need BigFloat to test convergence rate
dts = big"1.0"./big"2.0".^(5:-1:1)
prob_big = DynamicalODEProblem(ff_harmonic,[big"1.0",big"1.0"],
                               [big"0.0", big"0.0"],(big"0.",big"70."))
sim = test_convergence(dts,prob_big,DPRKN6(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 6 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 6 rtol = 1e-1
sim = test_convergence(dts,prob_big,DPRKN8(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 8 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob_big,DPRKN12(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 12 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob_big,ERKN4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob_big,ERKN5(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

# Adaptive methods regression test
sol = solve(prob, DPRKN6())
@test length(sol.u) < 20
sol = solve(prob, DPRKN8())
@test length(sol.u) < 13
sol = solve(prob, DPRKN12())
@test length(sol.u) < 10
sol = solve(prob, ERKN4(),reltol=1e-8)
@test length(sol.u) < 38
sol = solve(prob, ERKN5(),reltol=1e-8)
@test length(sol.u) < 34

# Test array partition outside of symplectic

function f_ap(du,u,p,t)
  du.x[1] .= -2u.x[2]
  du.x[2] .= u.x[1]
end

u = ArrayPartition((u0,v0))

prob = ODEProblem(f_ap,u,(0.0,5.0))
sol = solve(prob,Euler(),dt=1/100)

################# Out of place symplectic

println("Out of Place")

u0 = 0.0
v0 = 1.0
function f1_harmonic_nip(v,u,p,t)
  -u
end
function f2_harmonic_nip(v,u,p,t)
  v
end

ff_harmonic_nip = DynamicalODEFunction(f1_harmonic_nip,f2_harmonic_nip;
                                       analytic=harmonic_analytic)
prob = DynamicalODEProblem(ff_harmonic_nip,v0,u0,(0.0,5.0))

sol = solve(prob,SymplecticEuler(),dt=1/10)

dts = 1 .//2 .^(6:-1:3)
# Symplectic Euler
sim = test_convergence(dts,prob,SymplecticEuler(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 1 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 1 rtol = 1e-1
# Verlet
sim = test_convergence(dts,prob,VelocityVerlet(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1
# Test that position converges faster for Verlet
position_error = :final => [mean(sim[i].u[2].x[1] - sim[i].u_analytic[2].x[1]) for i in 1:length(sim)]
@test first(DiffEqDevTools.calc𝒪estimates(position_error).second) ≈ 4.0 rtol=1e-1

# 2nd Order Tableaus
sim = test_convergence(dts,prob,VerletLeapfrog(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1
sim = test_convergence(dts,prob,PseudoVerletLeapfrog(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1
sim = test_convergence(dts,prob,McAte2(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 2 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 2 rtol = 1e-1

# Ruth
sim = test_convergence(dts,prob,Ruth3(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 3 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 3 rtol = 1e-1
sim = test_convergence(dts,prob,McAte3(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 3 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 3 rtol = 1e-1
sim = test_convergence(dts,prob,CandyRoz4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,McAte4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,McAte42(),dense_errors=true)
@test_broken sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test_broken sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,CalvoSanz4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

dts = 1 .//2 .^(4:-1:0)
sim = test_convergence(dts,prob,McAte5(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sim = test_convergence(dts,prob,Yoshida6(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 6 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4.69 rtol = 1e-1
sim = test_convergence(dts,prob,KahanLi6(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 6 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sim = test_convergence(dts,prob,McAte8(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 8 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

sim = test_convergence(dts,prob,KahanLi8(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 8 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

# Nyström method
dts = 1 .//2 .^(9:-1:6)
sim = test_convergence(dts,prob,Nystrom4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob,Nystrom4VelocityIndependent(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
@test_broken sim = test_convergence(dts,prob,IRKN3(),dense_errors=true)
@test_broken sim.𝒪est[:l2] ≈ 3 rtol = 1e-1
@test_broken sim.𝒪est[:L2] ≈ 3 rtol = 1e-1
@test_broken sim = test_convergence(dts,prob,IRKN4(),dense_errors=true)
#@test_broken sim.𝒪est[:l2] ≈ 4 rtol = 1e-1
#@test_broken sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
dts = 1.0 ./2.0 .^(5:-1:0)
sim = test_convergence(dts,prob,Nystrom5VelocityIndependent(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

dts = 1.0 ./2.0 .^(2:-1:-2)
sim = test_convergence(dts,prob,SofSpa10(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 10 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

# Methods need BigFloat to test convergence rate
dts = big"1.0"./big"2.0".^(5:-1:1)
prob_big = DynamicalODEProblem(ff_harmonic_nip,big"1.0",big"0.0",
                                 (big"0.",big"70."))
sim = test_convergence(dts,prob_big,DPRKN6(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 6 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 6 rtol = 1e-1
sim = test_convergence(dts,prob_big,DPRKN8(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 8 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob_big,DPRKN12(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 12 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob_big,ERKN4(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1
sim = test_convergence(dts,prob_big,ERKN5(),dense_errors=true)
@test sim.𝒪est[:l2] ≈ 5 rtol = 1e-1
@test sim.𝒪est[:L2] ≈ 4 rtol = 1e-1

# Adaptive methods regression test
sol = solve(prob, DPRKN6())
@test length(sol.u) < 20
sol = solve(prob, DPRKN8())
@test length(sol.u) < 13
sol = solve(prob, DPRKN12())
@test length(sol.u) < 10
sol = solve(prob, ERKN4(),reltol=1e-8)
@test length(sol.u) < 38
sol = solve(prob, ERKN5(),reltol=1e-8)
@test length(sol.u) < 34


using JuMP
using MathOptInterface

optimizer_solver = optimizer_with_attributes(
    ActiveSetMethods.Optimizer,
    "external_optimizer" => GLPK.Optimizer,
    "OutputFlag" => 0,
)

model = Model(optimizer_solver)

@variable(model, X);
@variable(model, Y);
@objective(model, Min, X^2 + X);
@NLconstraint(model, X^2 - X == 2);
@NLconstraint(model, X * Y == 1);
@NLconstraint(model, X * Y >= 0);
@constraint(model, X >= -2);

JuMP.optimize!(model);

xsol = JuMP.value.(X)
ysol = JuMP.value.(Y)
status = termination_status(model)


function newton_raphson(f_,h_!,x; ρ = 0.5, c = 1e-4, tol = 1e-4, max_iter = 1000)
    grad= zeros(length(x))
    hes= zeros(length(x),length(x))
    y = 0.
    for i = 1:max_iter
        y = h_!(x, grad, hes)
        search_dir = -(cholfact(Positive, hes)\grad)
        (norm(search_dir) > tol) || break
        step_size = 1.
        while f_(x+search_dir*step_size) > y+c*step_size*(grad' * search_dir)
            step_size *= ρ
            step_size > 1e-10 || error("Linesearch failed! Problematic Hessian or Gradient?")
        end
        x = x + search_dir*step_size
    end
    return x, y, grad, hes
end


using SparseIR
using Test

import PyCall: pyimport, PyNULL, PyVector
import LegendrePolynomials: Pl

@testset "test_augmented_bosonic_basis" begin
    """Augmented bosonic basis"""
    wmax = 2
    beta = 1000
    basis = FiniteTempBasis(boson, beta, wmax, 1e-6)
    basis_legg = LegendreBasis(boson, beta, 2)
    basis_comp = CompositeBasis([basis_legg, basis])

    # G(tau) = c - e^{-tau*pole}/(1 - e^{-beta*pole})
    pole = 1.0
    c = 1e-2
    tau_smpl = TauSampling(basis_comp)
    gtau = c .- exp.(-tau_smpl.sampling_points .* pole) ./ (1 - exp(-beta * pole))
    gl_from_tau = fit(tau_smpl, gtau)

    gtau_reconst = evaluate(tau_smpl, gl_from_tau)
    @test isapprox(gtau, gtau_reconst; atol=1e-14 * maximum(abs.(gtau)), rtol=0)
end


ξ(Ma,a,b) = sqrt((1-Ma*a*b)^2-a^2)

function f!(F,x,Δxm,Δz2,Δz3,Ma,Δy)
    if (-0.8<=x[1]<=0.8) #&& (0.2<=abs(x[2])<=1)
        Δx20 = Δz2 * x[1]*x[2]/ξ(Ma,x[1],x[2])
        Δx21 = Δz2 * Ma*(1-Ma*x[1]*x[2])/ξ(Ma,x[1],x[2])
        Δx3  = Δz3 * x[1]*x[2]/sqrt(1-x[1]^2)
        F[1] = Δx20 + Δx21 + Δx3 - Δxm
        F[2] = x[2] - (Δxm - Δx21)/sqrt((Δxm-Δx21)^2+Δy^2)
        return nothing
    else
        return nothing
    end
end

"""
    r,pc_pm = refraction_correction(E::Environment[,h1=1.5,h2=z0-h1])

Compute propagation time correction and amplitude correction for 2D planar jet.

References:
-	Glegg, S., & Devenport, W. (2017). Aeroacoustics of low mach number flows: Fundamentals, analysis, and measurement. Chap. 10.
"""
function refraction_correction(E,h1=1.5,h2=E.z0-h1)
    r = zeros(E.N,E.M)
    pc_pm = similar(r)
    for n = 1:E.N
        for m = 1:E.M
            Δx,Δy,Δz = E.Rxy[:,n]-E.micgeom[:,m] # vector from mic to grid point
            r0 = sqrt(Δx^2+Δy^2+Δz^2)
            a0 = sqrt(Δx^2+Δy^2)/r0
            b0 = Δx/sqrt(Δx^2+Δy^2)
            Δz2,Δz3 = h1,h2

            res = nlsolve((F,x)->f!(F,x,Δx,Δz2,Δz3,E.Ma,Δy),[a0,b0],autoscale=false)
            a,b = res.zero

            Δx20 = Δz2*a*b/ξ(E.Ma,a,b)
            Δx3 = Δz3*a*b/sqrt(1-a^2)

            r[n,m] = Δz3/sqrt(1-a^2) + Δz2*(1-E.Ma*a*b)/ξ(E.Ma,a,b)

            # Amplitude correction:
            pc_pi = Δz2^2/(Δz2+Δz3)^2
            pi_pt = (sqrt(1-a^2)*(1-E.Ma*a*b)^2+ξ(E.Ma,a,b))^2/(4*ξ(E.Ma,a,b)^2)
            K(R,A,D) = sqrt(a^2*(D*sqrt(1-a^2)*(1-b^2)/ξ(E.Ma,a,b) + A*(b-D/R) )^2 + a^2*(1-b^2)*(A-D*sqrt(1-a^2)*b/ξ(E.Ma,a,b))^2 + (1-a^2)*(R-D*b)^2)

            d = 1-E.Ma*a*b
            r2 = (1-E.Ma*a*b)*Δz2/ξ(E.Ma,a,b)
            r3 = Δz3/sqrt(1-a^2)
            ρ2 = r2/d^2
            R2 = ρ2
            R3 = r3+ρ2
            A2 = sqrt(1-a^2)*ρ2/ξ(E.Ma,a,b)
            A3 = r3+sqrt(1-a^2)*ρ2/ξ(E.Ma,a,b)
            D2 = ρ2*E.Ma*a
            D3 = ρ2*E.Ma*a
            pt_pm = R3*K(R3,A3,D3)/(R2*K(R2,A2,D2))
            pc_pm[n,m] = pc_pi * pi_pt * pt_pm
        end
    end
    return r,pc_pm
end

"""
    steeringvectors(E::Environment)

Pre-compute steeringvectors for beamforming using an `Environment` with the needed parameters.
"""
function steeringvectors(E::Environment)
    @unpack fn,c,M,N,Nf,D,D0 = E
    kw = 2pi*fn/c
    vi = Array{ComplexF64,3}(undef,N,M,Nf)
    for j in 1:Nf
        vi[:,:,j] .= 1 ./(D.*D0.*sum(1 ./D.^2,dims=2)).*exp.(-im.*kw[j].*(D.-D0))
    end
    steeringvec = FreqArray(permutedims(vi,(2,1,3)),fn)
end

function steeringvectors!(E::Environment)
    @unpack fn,c,M,N,Nf,D,D0,Ma,h = E
    vi = Array{ComplexF64,3}(undef,N,M,Nf)

    if E.shear
        w = 2pi*fn
        dx,pc_pm = refraction_correction(E,h)
        ta = dx./c
        for j in 1:Nf
            vi[:,:,j] .= 1 ./(ta.*c.*D0.*sum(1 ./(ta.*c).^2,dims=2)).*exp.(-im.*w[j].*ta)
        end
        if E.ampcorr
            for j in 1:Nf
                vi[:,:,j] .*= pc_pm
            end
        end
    else
        kw = 2pi*fn/c
        for j in 1:Nf
            vi[:,:,j] .= 1 ./(D.*D0.*sum(1 ./D.^2,dims=2)).*exp.(-im.*kw[j].*(D.-D0))
        end
    end

    E.steeringvec = FreqArray(permutedims(vi,(2,1,3)),fn)
    return nothing
end


# constrained to one -liners for this one:
println("The answer to Euler Problem 6 is $(sum(1:100)^2 - sum(map(x->x^2, 1:100)))")


module RulesProbitOutTest

using Test
using ReactiveMP
using Random
using Distributions
using StatsFuns: normcdf, normccdf, normlogcdf, normlogccdf, normlogpdf, normpdf

import ReactiveMP: @test_rules

@testset "rules:Probit:out" begin

    @testset "Belief Propagation: (m_in::UnivariateNormalDistribution, )" begin

        @test_rules [ with_float_conversions = true ] Probit(:out, Marginalisation) [
            (input = (m_in = NormalMeanVariance(1, 0.5), ), output = Bernoulli(normcdf(1/sqrt(1+0.5)))),
            (input = (m_in = NormalMeanPrecision(1, 2), ), output = Bernoulli(normcdf(1/sqrt(1+0.5)))),
            (input = (m_in = NormalWeightedMeanPrecision(2, 2), ), output = Bernoulli(normcdf(1/sqrt(1+0.5)))),
            (input = (m_in = NormalMeanVariance(2, 0.25), ), output = Bernoulli(normcdf(2/sqrt(1+0.25)))),
            (input = (m_in = NormalMeanPrecision(2, 4), ), output = Bernoulli(normcdf(2/sqrt(1+0.25)))),
            (input = (m_in = NormalWeightedMeanPrecision(8, 4), ), output = Bernoulli(normcdf(2/sqrt(1+0.25)))),
        ]

    end

    @testset "Belief Propagation: (m_in::PointMass, )" begin

        @test_rules [ with_float_conversions = true ] Probit(:out, Marginalisation) [
            (input = (m_in = PointMass(1), ), output = Bernoulli(normcdf(1))),
            (input = (m_in = PointMass(2), ), output = Bernoulli(normcdf(2))),
            (input = (m_in = PointMass(3), ), output = Bernoulli(normcdf(3)))
        ]

    end

end

end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ b36fb690-3447-11eb-3629-edbddae4ebd1
arrivalStr,shuttleIds = readlines("./input.txt");

# ╔═╡ 47b5c4a0-3db9-11eb-3576-97091f1bf011
begin
	arrival = parse(Int,arrivalStr)
	shuttles = NamedTuple{(:shuttle, :delay),Tuple{Int64,Int64}}[]
	for (index,value) ∈ enumerate(split(shuttleIds,","))
		if value != "x"
			push!(shuttles, (shuttle = parse(Int,value), delay = mod(-1*(index-1),parse(Int,value))))
		end
	end
end

# ╔═╡ e1b52942-3db9-11eb-0e49-45d055bb09f0
shuttles

# ╔═╡ 3c75c812-3dbc-11eb-23a6-c53244ba0530
# Part 1
function part1(shuttles::Array{NamedTuple{(:shuttle, :delay),Tuple{Int64,Int64}}}, arrival::Int)
	waitTime, shuttleId = minimum(map( b -> (b.shuttle-arrival %b.shuttle, b.shuttle),shuttles))
	waitTime * shuttleId
end

# ╔═╡ d4e0a3d2-3dbd-11eb-19b9-5538a2737649
@show part1(shuttles,arrival)

# ╔═╡ 5edd004e-3dbe-11eb-1eba-477618416d09
# Chinese Reminder
# https://rosettacode.org/wiki/Chinese_remainder_theorem

function part2(shuttles::Array{NamedTuple{(:shuttle, :delay),Tuple{Int64,Int64}}})
	Π = prod(s -> s.shuttle, shuttles)
	sol = sum(s -> s.delay * (Π ÷ s.shuttle) * invmod(Π ÷ s.shuttle, s.shuttle),shuttles)
	mod(sol,Π)
end

# ╔═╡ 3e018d12-3dbf-11eb-1b42-e98c778fbb50
@show part2(shuttles)

# ╔═╡ Cell order:
# ╠═b36fb690-3447-11eb-3629-edbddae4ebd1
# ╠═47b5c4a0-3db9-11eb-3576-97091f1bf011
# ╠═e1b52942-3db9-11eb-0e49-45d055bb09f0
# ╠═3c75c812-3dbc-11eb-23a6-c53244ba0530
# ╠═d4e0a3d2-3dbd-11eb-19b9-5538a2737649
# ╠═5edd004e-3dbe-11eb-1eba-477618416d09
# ╠═3e018d12-3dbf-11eb-1b42-e98c778fbb50


include("./params.jl")
include("./fitgamma.jl")

function intensity_gamma(data)
    # Example usage:
    # x = 0.5-0.5*log(rand(1,1e3));
    # rs = intensity_gamma(x);
    #
    # Input argument
    # data: Spike-rate.
    #
    # Output argument
    # rs:   Estimated intensity function

    isi = data[2:end] - data[1:end-1]
    lambda, k = fitgamma(isi)

    Nspike = length(data)

    P = 0;
    prev_flg = 0;
    rs = zeros(L)
    for i=1:L

        if prev_flg == 0
            tau = i*dt
        else
            tau = i*dt - data[prev_flg]
        end

        p = 1.0 / gamma(k) * (k * lambda).^k * tau.^(k-1) .* exp(-k*lambda*tau); # ISI
        P = P + p*dt;

        rs[i] = p / (1.0 - P); # Intensity

        if prev_flg+1 < Nspike && data[prev_flg+1] <= i*dt

            P = 0;
            prev_flg += 1;

        end

    end

    rs
end


module DisposableResourceAllocation

using JuMP
using Gurobi

using LinearAlgebra
using MathOptInterface

export disposable_resource_allocation

function disposable_resource_allocation(
		initial_supply::Array{<:Real,1},
		supply::Array{<:Real,2},
		demand::Array{<:Real,2},
		adj_matrix::BitArray{2};
		objective::Symbol=:shortage,
		send_new_only::Bool=true,
		send_wait_period::Int=10,
		m::Real=1e-5,
		verbose::Bool=false,
)
	N, T = size(supply)
	@assert(size(initial_supply, 1) == N)
	@assert(size(demand, 1) == N)
	@assert(size(demand, 2) == T)
	@assert(size(adj_matrix, 1) == N)
	@assert(size(adj_matrix, 2) == N)
	@assert(objective in [:shortage, :overflow])

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	@variable(model, sent[1:N,1:N,1:T] >= 0)
	@variable(model, used[1:N,1:T] >= 0)
	@variable(model, obj_dummy[1:N,1:T] >= 0)

	@objective(model, Min, sum(obj_dummy) + m*sum(sent))

	if send_new_only
		@constraint(model, [t=1:T],
			sum(sent[:,:,t], dims=2) .<= max.(0, supply[:,t])
		)
	end

	@constraint(model, [t=1:T],
		sum(sent[:,:,t], dims=2) .+ used[:,t] .<=
			initial_supply
			.+ sum(supply[:,1:t], dims=2)
			.- sum(used[:,1:t-1], dims=2)
			.- sum(sent[:,:,1:t-1], dims=[2,3])
			.+ sum(sent[:,:,1:t-1], dims=[1,3])
	)

	for i = 1:N
		for j = 1:N
			if ~adj_matrix[i,j]
				@constraint(model, sum(sent[i,j,:]) .== 0)
			end
		end
	end

	if send_wait_period > 0
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,i,t]), sum(sent[i,:,t:min(t+send_wait_period,T)])] in MOI.SOS1([1.0, 1.0])
		)
	end

	flip_sign = (objective == :shortage) ? 1 : -1
	@constraint(model, [i=1:N,t=1:T],
		obj_dummy[i,t] >= flip_sign * (demand[i,t] - used[i,t])
	)

	optimize!(model)
	return model
end

end;


function solve()
    N, M, T = [parse(Int, x) for x in split(readline())]
    t = 0
    n = N
    for i=1:M
        A, B = [parse(Int, x) for x in split(readline())]
        n -= (A - t)
        if n ≤ 0
            return "No"
        end
        n += (B - A)
        n = min(n, N)
        t = B
    end
    n -= (T - t)
    n ≤ 0 ? "No" : "Yes"
end

println(solve())


using IterativeSolvers
using FactCheck

facts("Randomized linear algebra") do
    srand(1)
    m = n = 100
    B = randn(m, n)
    nB = norm(B)
    p = 1e-5 #Probability of failure

    for j=1:n
        @fact rnorm(B, 2j, p) --> greater_than_or_equal(nB)
        @fact rnorms(B, j, p) --> greater_than_or_equal(nB)
    end

    A = B*B'
    k = 1
    l, u = reigmin(A, k, p)
    @fact l <= eigmin(A) <= u --> true

    l, u = reigmax(A, k, p)
    @fact l <= eigmax(A) <= u --> true

    l, u = rcond(A, k, p)
    @fact l <= cond(A) <= u --> true
end


function printcircle(lo::Integer, hi::Integer, ndots::Integer; pad::Integer = 2)
    canvas = falses(2hi + 1, 2hi + 1)
    i = 0
    while i < ndots
        x, y = rand(-hi:hi, 2)
        if lo ^ 2 - 1 < x ^ 2 + y ^ 2 < hi ^ 2 + 1
            canvas[x + hi + 1, y + hi + 1] = true
            i += 1
        end
    end
    # print
    for i in 1:(2hi + 1)
        row = map(j -> j ? "\u25cf " : "  ", canvas[i, :])
        println(" " ^ pad, join(row))
    end
    return canvas
end

printcircle(10, 15, 100)


@add_problem constant function constant_Issue_166(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    # Issue #166
    α = Variable(5)
    fix!(α, ones(5,1))

    # has const vexity, but not at the head
    c = (rand(5,5) * α) * ones(1,5)

    β = Variable(5)
    β.value = ones(5)

    problem = minimize(sum(c * β), [β >= 0])
    handle_problem!(problem)
    if test
        @test problem.optval ≈ evaluate(sum(c * β)) atol=atol rtol=rtol
        @test problem.optval ≈ 0.0 atol=atol rtol=rtol
        @test β.value ≈ zeros(5) atol=atol rtol=rtol
    end
end

@add_problem constant function constant_Issue_228(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2)
    y = Variable(2)
    fix!(x, [1 1]')
    prob = minimize(y'*(x+[2 2]'), [y>=0])
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 0.0 atol=atol rtol=rtol
    end

    prob = minimize(x'*y, [y>=0])
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 0.0 atol=atol rtol=rtol
    end
end

@add_problem constant function constant_Test_double_fix!(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable()
    y = Variable()
    fix!(x, 1.0)
    prob = minimize(y*x, [y >= x, x >= 0.5])
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 1.0 atol=atol rtol=rtol
    end

    fix!(x, 2.0)
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 4.0 atol=atol rtol=rtol
    end

    free!(x)
    fix!(y, 1.0)
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 0.5 atol=atol rtol=rtol
    end
end

@add_problem constant function constant_fix!_and_multiply(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    p = Variable()
    fix!(p, 1.0)
    x = Variable(2,2)
    prob = minimize( tr(p*x), [ x >= 1 ])
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 2.0 atol=atol rtol=rtol
        @test evaluate( tr(p*x) ) ≈ 2.0 atol=atol rtol=rtol
    end
end

@add_problem constant function constant_fix!_with_complex_numbers(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = ComplexVariable()
    fix!(x, 1.0 + im*1.0)
    y = Variable()
    prob = minimize( real(x*y), [ y >= .5, real(x) >= .5, imag(x) >= 0])
    handle_problem!(prob)
    if test
        @test prob.optval ≈ .5 atol=atol rtol=rtol
        @test evaluate(real(x*y)) ≈ .5 atol=atol rtol=rtol
        @test evaluate(y) ≈ 0.5 atol=atol rtol=rtol
    end

    free!(x)
    if test
        fix!(y)
    else # if we haven't solved the problem, we need to fix! it to a number
        fix!(y, 0.5)
    end
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 0.25 atol=atol rtol=rtol
        @test evaluate(real(x*y)) ≈ 0.25 atol=atol rtol=rtol
        @test real(evaluate(x)) ≈ 0.5 atol=atol rtol=rtol
        @test evaluate(y) ≈ 0.5 atol=atol rtol=rtol
    end

    if test
        @test_throws DimensionMismatch fix!(x, rand(2))
        @test_throws DimensionMismatch fix!(x, rand(2,2))
    end
end

@add_problem constant function constant_fix!_with_vectors(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = ComplexVariable(5)
    fix!(x, ones(5) + im*ones(5))
    y = Variable()
    prob = minimize( real(y*sum(x)), [ y >= .5, real(x) >= .5, imag(x) >= 0])
    handle_problem!(prob)
    if test
        @test prob.optval ≈ 2.5 atol=atol rtol=rtol
        @test evaluate(real(y*sum(x))) ≈ 2.5 atol=atol rtol=rtol
        @test evaluate(y) ≈ 0.5 atol=atol rtol=rtol
    end

    free!(x)
    if test
        fix!(y)
    else # if we haven't solved the problem, we need to fix! it to a number
        fix!(y, 0.5)
    end

    handle_problem!(prob)
    if test
        @test prob.optval ≈ 1.25 atol=atol rtol=rtol
        @test evaluate(real(y*sum(x))) ≈ 1.25 atol=atol rtol=rtol
        @test real(evaluate(x)) ≈ 0.5*ones(5) atol=atol rtol=rtol
        @test evaluate(y) ≈ 0.5 atol=atol rtol=rtol
    end

    if test
        @test_throws DimensionMismatch fix!(x, rand(5,5))
        @test_throws DimensionMismatch fix!(x, rand(4))
    end

end


# Script that plots the time series of the model versus observations
# - observations are automatically extracted from tides and currents
# - only time series where both model and observations are sufficiently complete
#   are plotted
# - Plots daily maximum water levels from the timeseries in addition
#   to full timeseries and lowpass filtered timeseries
using NetCDF
using Plots
using Dates
using HTTP
using CSV
using Statistics
using TimeSeries

using FFTW
include("fourfilt.jl")

# using GR backend, behaves like matlab
gr()

nodatalength = 0.0 # decimal percentage of period (if non-zero then the lowpass filter does not work...)
# low pass filter length
          #days
lpwindow = 28*3600.0*24.0 # seconds
dpiset = 300

# filenames
filename = "ATT/Mv13.61.nc"
sn = ncread(filename, "station_name")
# read parts of the netcdf file, convert to datetimes
time = ncread(filename, "time")
zeta = ncread(filename, "zeta")
zeta[zeta .< -999] .= NaN
lon = ncread(filename, "x")
lat = ncread(filename, "y")
dt = DateTime(2004,6,15) + Dates.Second.(time)
# Make sure only from July 1st
zeta = zeta[:,dt .>= DateTime(2004,7,1)]
deleteat!(dt,dt .< DateTime(2004,7,1))
tmstep = (Dates.value(dt[2]-dt[1]))/1000 #seconds

# This is for reading observations through tidesandcurrents API
base = "https://tidesandcurrents.noaa.gov/api/datagetter?"
suf = "product=hourly_height&datum=MSL&units=metric&time_zone=gmt&application=ports_screen&format=csv"
sd  = join(["begin_date=",Dates.format(dt[1],"yyyymmdd"),"&"])
ed  = join(["end_date=",Dates.format(dt[end],"yyyymmdd"),"&"])

# legend names
linenames = [:Observed, :Model]

# initializing the scatter arrays
xp = Float64[]; yp = Float64[]; mv = Float64[]; ov = Float64[]

# loop over all stations
for ii = 1:size(sn,2)

    # only get stations in our AOI
    #if (lon[ii] > -75 || lon[ii] < -83 ||
    #lat[ii] > 36.5 || lat[ii] < 24); continue; end
    if (lon[ii] > -67 || lon[ii] < -100 ||
        lat[ii] > 46 || lat[ii] < 24); continue; end

    # if modelled doesn't have enough values
    if (sum(isnan.(zeta[ii,:])) > length(dt)*nodatalength); continue; end

    # get and plot the station data
    sta = join(["station=",rstrip(join(sn[:,ii])),"&"])
    res = HTTP.get(join([base,sd,ed,sta,suf]))
    mycsv = CSV.read(res.body)
    # if got error or missing time
    if (ismissing(mycsv[1,1]) || mycsv[1,1][1:5] == "Error"); continue; end
    dtn = DateTime.(mycsv[:,1],"yyyy-mm-dd HH:MM")
    # if observed doesn't have enough non-missing values
    if (sum(ismissing.(mycsv[:,2])) > length(dtn)*nodatalength); continue; end

    # Display the station IDs and indices
    display(ii)
    display(join(sn[:,ii]))

    # Process observed results
    zo = mycsv[:,2]
    zo[ismissing.(mycsv[:,2])] .= NaN
    # remove low-pass filtered timeseries
    tostep = (Dates.value(dtn[2]-dtn[1]))/1000 #seconds
    zofilt = fourfilt(zo,tostep,Inf,lpwindow)
    zoaf = zo .- zofilt
    # convert observed data into the timeseries array
    to = TimeArray(dtn,zoaf)
    # calculate daily maximum
    to = collapse(to,day,first,maximum)
    rename!(to, linenames[1])

    # Process modeled results
    zm = zeta[ii,:]
    # remove low-pass filtered timeseries
    zmfilt = fourfilt(zm,tmstep,Inf,lpwindow)
    zmaf = zm .- zmfilt
    # convert model data into the timeseries array
    tm = TimeArray(dt,zmaf)
    # calculate daily maximum
    tm = collapse(tm,day,first,maximum)
    rename!(tm, linenames[2])

    # plot the observations and model
    plt1 = plot(to,
        title = join(["Daily Maximum Water Levels at Station ID: ",
            rstrip(join(sn[:, ii]))]
        ),
        legend = :topright,
        ylabel = "Water Level (LMSL) [m]",
        xlabel = "DateTime [UTC]"
    )
    plot!(tm)

    plt2 = plot(dtn,zo)
    plot!(dt,zm)
    plot!(dtn,zofilt)
    plot!(dt,zmfilt)
    plt = plot(plt1,plt2,dpi=dpiset,layout = (2,1))

    display(plt) # display in plot pane

    # calculate max WL error
    vtm = values(tm)
    vto = values(to)
    vtm = maximum(vtm[.!isnan.(vtm)])
    vto = maximum(vto[.!isnan.(vto)])

    # add the location plot to show error at stations
    push!(xp,lon[ii])
    push!(yp,lat[ii])
    push!(mv,vtm)
    push!(ov,vto)

    savefig(join(["ATT/DailyWLPlots_Mv13/" rstrip(join(sn[:, ii]))]))

end

# write out the maximum water levels over this time period
using DataFrames
df = DataFrame(Lon = xp, Lat = yp, MaxWL_Obs = ov, MaxWL_Mod = mv)
CSV.write("ATT/DailyWLPlots_Mv13/MaxWLs.csv",df)

# I'm having trouble using scatter here because of clash between Plots and GR package
#plt = scatter(xp,yp,200*ones(length(xp)),error) #,legend=:none)
#display(plt) # display in plot pane
#savefig(join(["ATT/DailyWLPlots_Mv13/MaxWLError"))


function simulate(	model::String,
					m::Array{Float64, 1},
					c::Array{Float64, 1},
					ax::Float64,
					ay::Float64,
					az::Array{Float64, 1},
					Lx::Float64,
					Ly::Float64,
					Lz::Float64,
					number_of_frames::Array{Int64, 1},
					deltat::Float64,
					kmin::Int64,
					number_of_de_bins::Int64,
					ub_de::Float64)

	# Number of components in distribution.
	number_of_components::Int64 = length(m)
	fractions::Array{Float64, 1} = c / sum(c)
	cum_fractions::Array{Float64, 1} = cumsum(fractions)

	# Intensity of Poisson distribution of the number of particles. The factor
	# 1e12 takes into account that concentration is specified in particles/ml.
	lambda::Float64 = sum(c) * Lx * Ly * Lz / 1e12
	kmax::Int64 = maximum(number_of_frames)

	# Number of videos.
	number_of_videos::Int64 = length(number_of_frames)

	# Lower and upper bounds for the detection region.
	lbx::Float64 = 0.5 * (Lx - ax)
	ubx::Float64 = 0.5 * (Lx + ax)
	lby::Float64 = 0.5 * (Ly - ay)
	uby::Float64 = 0.5 * (Ly + ay)
	lbz::Array{Float64, 1} = 0.5 * (Lz - az)
	ubz::Array{Float64, 1} = 0.5 * (Lz + az)

	# Variables for storing positions, displacements, number of particles, etc
	x::Float64 = 0.0
	y::Float64 = 0.0
	z::Float64 = 0.0
	deltax::Float64 = 0.0
	deltay::Float64 = 0.0
	deltaz::Float64 = 0.0
	number_of_particles::Int64 = 0
	k::Int64 = 0
	de::Float64 = 0.0
	std_dev_random_walk::Float64 = 0.0
	index::Int64 = 0

	# Histogram matrix for number of positions and estimated diffusion coefficients of all
	# recorded trajectories.
	H::Array{Int64, 2} = zeros(kmax, number_of_de_bins)
	d_de::Float64 = ub_de / convert(Float64, number_of_de_bins)
	idx::Int64 = 0

	for current_video = 1:number_of_videos
		number_of_particles = rand_poisson(lambda)

		for current_particle = 1:number_of_particles
			# Generate random diffusion coefficent from distribution, or more precisely,
			# a random standard deviation for the displacements.
			index = rand_weighted_index(cum_fractions)
			std_dev_random_walk = sqrt(2.0 * m[index] * deltat)

			# Random initial position.
			x = Lx * rand()
			y = Ly * rand()
			z = Lz * rand()

			# Reset empirical diffusion coefficient.
			de = 0.0

			# In detection region or not? Check z limits first because they are more restrictive.
			#if (lbz[index] <= z <= ubz[index]) & (lbx <= x <= ubx) & (lby <= y <= uby)
			if (abs(z - 0.5 * Lz) <= 0.5 * az[index]) & (abs(x - 0.5 * Lx) <= 0.5 * ax) & (abs(y - 0.5 * Ly) <= 0.5 * ay)
				k = 1
			else
				k = 0
			end

			# Let particle diffuse through remainder of video.
			for current_frame = 2:number_of_frames[current_video]
				deltax = std_dev_random_walk * randn()
				deltay = std_dev_random_walk * randn()
				deltaz = std_dev_random_walk * randn()

				x = x + deltax
				y = y + deltay
				z = z + deltaz

				x = position_periodic(x, Lx)
				y = position_periodic(y, Ly)
				z = position_periodic(z, Lz)

				#if (lbz[index] <= z <= ubz[index]) & (lbx <= x <= ubx) & (lby <= y <= uby)
				if (abs(z - 0.5 * Lz) <= 0.5 * az[index]) & (abs(x - 0.5 * Lx) <= 0.5 * ax) & (abs(y - 0.5 * Ly) <= 0.5 * ay)
					k = k + 1
					if k >= 2 # The first point inside the detection region does not contributed toward the estimated diffusion coefficient.
						de = de + deltax^2 + deltay^2 # Only in x-y plane.
					end
				elseif k > 0
					if k >= kmin
						de = de / (convert(Float64, k - 1) * 4.0 * deltat) # The '4' comes from the 2D observations.
						idx = convert(Int64, ceil(de / d_de))
						if 1 <= idx <= number_of_de_bins
							H[k, idx] += 1
						end
					end
					k = 0
					de = 0.0
				end
			end

			# If particle is inside detection region in last frame, we have now missed that trajectory and need to add it now.
			if k >= kmin # Assuming kmin >= 2 otherwise the division fails.
				de = de / (convert(Float64, k - 1) * 4.0 * deltat) # The '4' comes from the 2D observations.
				idx = convert(Int64, ceil(de / d_de))
				if 1 <= idx <= number_of_de_bins
					H[k, idx] += 1
				end
			end

		end
	end

	return H
end


#=
ellipse.jl
=#

#using ImagePhantoms #: Object, Object2d

export Ellipse
export Circle
export phantom, radon, spectrum

"""
    Ellipse <: AbstractShape2
"""
struct Ellipse <: AbstractShape2 end

# constructors

"""
    Ellipse(cx, cy, rx=1, ry=rx, ϕ=0, value::Number=1)
    Ellipse(center::NTuple{2,RealU}, radii::NTuple{2,RealU}=(1,1), ϕ::RealU=0, v=1)
    Ellipse([6-vector])
Construct `Ellipse` object from parameters.
"""
function Ellipse(
    cx::RealU,
    cy::RealU,
    rx::RealU = oneunit(cx),
    ry::RealU = rx,
    ϕ::RealU = 0,
    value::Number = 1,
)
    (cx, cy, rx, ry) = promote(cx, cy, rx, ry)
    Object(Ellipse(), (cx,cy), (rx,ry), ϕ, value)
end

function Ellipse(
    center::NTuple{2,RealU},
    radii::NTuple{2,RealU} = (1,1) .* oneunit(center[1]),
    ϕ::RealU = 0,
    value::Number = 1,
)
    Ellipse(center..., radii..., ϕ, value)
end

function Ellipse(v::AbstractVector{<:Number})
    length(v) == 6 || throw(ArgumentError("$v wrong length"))
    Ellipse(v...)
end

# circles as a special case

"""
    Circle(x,y,r,v=1) (circle of radius `r` centered at `(x,y)`)
    Circle((x,y), r=1, v=1) ditto
    Circle([4-vector]) ditto
    Circle(r, v=1) centered at origin
Construct circle objects as special cases of `Ellipse` objects.
"""
Circle(cx::RealU, cy::RealU, r::RealU, v::Number = 1) =
    Ellipse(cx, cy,  r, r, 0, v)
Circle(center::NTuple{2,RealU}, r::RealU = oneunit(center[1]), v::Number = 1) =
    Circle(center..., r, v)
Circle(r::RealU, v::Number = 1) = Circle((zero(r), zero(r)), r, v)

function Circle(v::AbstractVector{<:Number})
    length(v) == 4 || throw(ArgumentError("$v wrong length"))
    Circle(v...)
end

# helper

# methods

"""
    phantom(ob::Object2d{Ellipse})
Returns function of `(x,y)` for making image.
"""
phantom(ob::Object2d{Ellipse}) =
    (x,y) -> (sum(abs2, coords(ob, x, y)) ≤ 1) * ob.value

"""
    radon_ellipse(r, ϕ, cx, cy, rx, ry, θ)
Radon transform of ellipse at point `(r,ϕ)`.
"""
function radon_ellipse(r, ϕ, cx, cy, rx, ry, θ)
    (sinϕ, cosϕ) = sincos(ϕ)
    r -= cx * cosϕ + cy * sinϕ # Radon translation property
    (sinϕ, cosϕ) = sincos(ϕ - θ) # Radon rotation property
    rp2 = abs2(rx * cosϕ) + abs2(ry * sinϕ) # square of projected radius
    return 2rx*ry / rp2 * sqrt(max(rp2 - abs2(r), 0*oneunit(r)^2))
end

"""
    radon(ob::Object2d{Ellipse})
Returns function of `(r,ϕ)` for making a sinogram.
"""
radon(ob::Object2d{Ellipse}) = (r,ϕ) -> ob.value *
    radon_ellipse(r, ϕ, ob.center..., ob.width..., ob.angle[1])

function spectrum_ellipse(fx, fy, cx, cy, rx, ry, θ)
    (kx, ky) = rotate2d(fx, fy, θ) # rect is rotated first, then translated
    return 2rx * exp(-2im*π*fx*cx) * # width=diameter=2*radius
           2ry * exp(-2im*π*fy*cy) * jinc(2sqrt(abs2(kx * rx) + abs2(ky * ry)))
end

"""
    spectrum(ob::Object2d{Ellipse})
Returns function of ``(f_x,f_y)`` for the spectrum (2D Fourier transform).
"""
spectrum(ob::Object2d{Ellipse}) = (fx,fy) -> ob.value *
    spectrum_ellipse(fx, fy, ob.center..., ob.width..., ob.angle[1])


rng = ParisiRapuano()
n=1000
@test all(0 .<= rand(rng, n) .< 1.)
# @test all(-5 .<= randn(rng, 1000) .< 5.)
n=100000
# FAILING. Some proble with Close1Open2?
#@test abs(mean(randn(rng, n))) < 5 / √n
for n=[1_000, 10_000]
    @test abs(mean(rand(rng, n))-0.5) < 20abs(mean(rand(n))-0.5)
    @test abs(var(rand(rng, n)) - 1/12) < 20abs(var(rand(n)) - 1/12)
end


using TypeVars
using Test

struct Zm{M, T<:Integer}
    val::T
end

import Base: +
+(a::Z, b::Z) where Z<:Zm = Z(mod(a.val + b.val, typevar(Z)))

@testset "TypeVars" begin

    m = big"2"^512 -1
    M = Symbol(hash(m))
    Z = Zm{M,BigInt}
    settypevar(Z, m)

    a = Z(m - 5)
    b = Z(m - 10)
    c = a + b
    @test c.val == m - 15
end


################################################################################
## ReverseKullbackLeibler
################################################################################

#=---------------
Evaluate
---------------=#
function evaluate{T<:AbstractFloat}(div::RKL, a::AbstractVector{T}, b::AbstractVector{T})
    if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end

    r = zero(T)

    for i in eachindex(a, b)
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        ui = ai/bi
        if ui > 0
            r += -bi*log(ui) + ai - bi
        else
            r = convert(T, Inf)
            break
        end
    end
    return r
end

function evaluate{T<:AbstractFloat}(div::RKL, a::AbstractVector{T})
    r = zero(T)
    for i in eachindex(a)
        @inbounds ai = a[i]
        if ai > 0
            r += - logmxp1(ai)
        else
            r = convert(T, Inf)
            break
        end
    end
    r
end

#=---------------
Gradient
---------------=#
function gradient{T<:AbstractFloat}(div::RKL, a::T, b::T)
    if a > 0 && b > 0
        u = - b/a + 1.0
    else
        u = convert(T, Inf)
    end
    return u
end

function gradient{T<:AbstractFloat}(div::RKL, a::T)
    if a > 0
        u = - 1.0 / a + 1.0
    else
        u = convert(T, Inf)
    end
    return u
end

#=---------------
hessian
---------------=#
function hessian{T<:AbstractFloat}(div::RKL, a::T, b::T)
    if a > 0 && b > 0
        u = b/a^2
    else
        u = convert(T, Inf)
    end
    return u
end

function hessian{T<:AbstractFloat}(div::RKL, a::T)
    if a > 0
        u = 1.0/a^2
    else
        u = convert(T, Inf)
    end
    u
end

################################################################################
## Modified Reverse Kullback-Leibler - MRKL
################################################################################

#=---------------
Evaluate
---------------=#
function evaluate{T<:AbstractFloat}(div::MRKL, a::AbstractVector{T}, b::AbstractVector{T})
    if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    f0, f1, f2, uϑ = div.m
    r = zero(T)
    for i = eachindex(a, b)
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        ui = ai/bi
        if ui >= uϑ
            r += (f0 + f1*(ui-uϑ) + .5*f2*(ui-uϑ)^2)*bi
        elseif ui > 0 && ui <u₀
            r += -bi*log(ui) + ai - bi
        else
            r = convert(T, Inf)
            break
        end
    end
    return r
end

function evaluate{T<:AbstractFloat}(div::MRKL, a::AbstractVector{T})
    f0, f1, f2, uϑ = div.m
    r = zero(T)
    for i in eachindex(a)
        @inbounds ai = a[i]
        if ai >= uϑ
            r += f0 + f1*(ai-uϑ) + .5*f2*(ai-uϑ)^2
        elseif ai > 0 && ai < uϑ
            r += -log(ai) + ai - 1.0
        else
            r = convert(T, Inf)
            break
        end
    end
    return r
end

#=---------------
gradient
---------------=#
function gradient{T<:AbstractFloat}(div::MRKL, a::T)
    f0, f1, f2, uϑ = div.m
    ui = a
    if ui > u₀
        u = (f1 + f2*(ui-uϑ))
    elseif ui <= uϑ
        u = gradient(div.d, a)
    end
    return u
end

function gradient{T<:AbstractFloat}(div::MRKL, a::T, b::T)
    f0, f1, f2, uϑ = div.m
    ui = a/b
    if ui > u₀
        u = f1 + f2*(ui-uϑ)
    elseif ui <= uϑ
        u = gradient(div.d, a, b)
    end
    return u
end

#=---------------
hessian
---------------=#
function hessian{T<:AbstractFloat}(div::MRKL, a::T)
    f0, f1, f2, uϑ = div.m
    if a >= uϑ
       u  = f2
    else
       u = hessian(div.d, a)
    end
    return u
end

function hessian{T<:AbstractFloat}(div::MRKL, a::T, b::T)
    f0, f1, f2, uϑ = div.m
    ui = a/b
    if ui >= u₀
        u  = f2/b
    else
        u = hessian(rkl, a, b)
    end
    return u
end

################################################################################
## Fully Modified Reverse Kullback-Leibler - FMRKL
##
################################################################################

#=---------------
Evaluate
---------------=#
function evaluate{T<:AbstractFloat}(div::FMRKL, a::AbstractVector{T}, b::AbstractVector{T})
    if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    r = zero(T)
    f0, f1, f2, uϑ, g0, g1, g2, uφ = div.m
    @inbounds for i = eachindex(a, b)
        ai = a[i]
        bi = a[i]
        ui = ai/bi
        if ui >= uϑ
            r += (f0 + f1*(ui-uϑ) + .5*f2*(ui-uϑ)^2)*bi
        elseif ui <= uφ
            r += (g0 + g1*(ui-uφ) + .5*g2*(ui-uφ)^2)*bi
        else
            r += -bi*log(ui) + ai - bi
        end
    end
    return r
end

function evaluate{T<:AbstractFloat}(div::FMRKL, a::AbstractVector{T})
    r = zero(T)
    f0, f1, f2, uϑ, g0, g1, g2, uφ = div.m
    @inbounds for i = eachindex(a)
        ui = a[i]
        if ui >= uϑ
            r += (f0 + f1*(ui-uϑ) + .5*f2*(ui-uϑ)^2)
        elseif ui <= uφ
            r += (g0 + g1*(ui-uφ) + .5*g2*(ui-uφ)^2)
        else
            r += -log(ui) + ui - 1.0
        end
    end
    return r
end

#=---------------
gradient
---------------=#
function gradient{T<:AbstractFloat}(div::FMRKL, a::T, b::T)
    f0, f1, f2, uϑ, g0, g1, g2, uφ = div.m
    @inbounds for i = eachindex(a, b)
        ai = a[i]
        bi = a[i]
        ui = ai/bi
        if ui >= uϑ
            u = f1 + f2*(ui-uϑ)
        elseif ui <= uφ
            u = g1 + g2*(ui-uφ)
        else
            u = 1-bi/ai
        end
    end
    return u
end

function gradient{T<:AbstractFloat}(div::FMRKL, a::T)
    f0, f1, f2, uϑ, g0, g1, g2, uφ = div.m
    if a >= uϑ
        u = (f1 + f2*(a-uϑ))
    elseif a <= uφ
        u = (g1 + g2*(a-uφ))
    else
        u = 1.0 - 1.0/a
    end
    return u
end

#=---------------
hessian
---------------=#
function hessian{T<:AbstractFloat}(div::FMRKL, a::T)
    f0, f1, f2, uϑ, g0, g1, g2, uφ = div.m
    if a >= uϑ
        u  = f2
    elseif a <= uφ
        u  = g2
    else
        u = 1/(a*a)
    end
    return u
end

function hessian{T<:AbstractFloat}(div::FMRKL, a::T, b::T)
    f0, f1, f2, uϑ, g0, g1, g2, uφ = div.m
    ui = a/b
    if ui >= uϑ
        u  = f2/b
    elseif ui <= uφ
        u  = g2/b
    else
        u = b/(a*a)
    end
    return u
end


function calc_perimeter_ngon(a,b)
if length(a) == length(b) && length(a) >= 2
    L = 0;
    L = L + abs(a[1]+im*b[1]);
    for i = 2:length(a)
        L = L + abs((a[i]-a[i-1])+im*(b[i]-b[i-1]));
    end
    L = L + abs(a[length(a)]+im*b[length(a)]);
    return L
end
end


@registerN DiscreteChebyshev AbstractCDOP2 α β
export DiscreteChebyshev

"""
    DiscreteChebyshev

This uses p22 of [Koepf and Schmersau](https://arxiv.org/pdf/math/9703217.pdf)  to define a two-parameter  family of *non* orthogonal polynomials.
See  the example  in [`DiscreteWeightFunction`](@ref) for implementing  the [DiscreteChebyshev](https://en.wikipedia.org/wiki/Discrete_Chebyshev_polynomials)  polynomials  from Wikipedia.

```jldoctest
julia> using Polynomials, SpecialPolynomials

julia> import SpecialPolynomials: Δₓ, ∇ₓ

julia> α,β = 1/2, 1
(0.5, 1)

julia> P  = DiscreteChebyshev{α,β}
DiscreteChebyshev{0.5, 1}

julia> i = 5
5

julia> yᵢ = basis(P, i)
DiscreteChebyshev(1.0⋅K⁽ᵅᵝ⁾₅(x))

julia> x = variable(P)
DiscreteChebyshev(- 2.0⋅K⁽ᵅᵝ⁾₀(x) + 2.0⋅K⁽ᵅᵝ⁾₁(x))

julia> a,b,c,d,e = SpecialPolynomials.abcde(P)
(a = 0, b = 0, c = 1, d = 0.5, e = 1)

julia> λᵢ  = -(a*i*(i-1)  + d*i)
-2.5

julia> Δₓ(∇ₓ(yᵢ)) +  (α*x + β) * Δₓ(yᵢ) ≈ -λᵢ*yᵢ # p22: "are not orthogonal, but satisfy the difference equation..."
true
```

"""
DiscreteChebyshev

basis_symbol(::Type{<:DiscreteChebyshev{α,β}}) where {α,β} = "K⁽ᵅᵝ⁾"
Polynomials.domain(::Type{<:DiscreteChebyshev{α,β}}) where {α,β} =
    Polynomials.Interval(-Inf, Inf)
abcde(::Type{<:DiscreteChebyshev{α,β}}) where {α,β} =
    NamedTuple{(:a, :b, :c, :d, :e)}((0, 0, 1, α, β))

function k1k0(P::Type{<:DiscreteChebyshev{α,β}}, n::Int) where {α,β}
    α
end


# Frank-Wolfe Methods
# CFW and BFW in Mitradijieva and Lindberg (2013)

# required packages: Graphs, Optim

include("misc.jl")

function BPR(x::Vector{Float64}, td::TA_Data)
    bpr = similar(x)
    for i in 1:length(bpr)
        bpr[i] = td.free_flow_time[i] * ( 1.0 + td.b[i] * (x[i]/td.capacity[i])^td.power[i] )
        bpr[i] += td.toll_factor * td.toll[i] + td.distance_factor * td.link_length[i]
    end
    return bpr
end

gradient = BPR

function objective(x::Vector{Float64}, td::TA_Data)
    # value = free_flow_time .* ( x + b.* ( x.^(power+1)) ./ (capacity.^power) ./ (power+1))
    # return sum(value)

    sum = 0.0
    for i in 1:length(x)
        sum += td.free_flow_time[i] * ( x[i] + td.b[i]* ( x[i]^(td.power[i]+1)) / (td.capacity[i]^td.power[i]) / (td.power[i]+1))
        sum += td.toll_factor * td.toll[i] + td.distance_factor * td.link_length[i]
    end
    return sum
end

function hessian(x::Vector{Float64}, td::TA_Data)
    no_arc = length(td.init_node)

    h = zeros(no_arc, no_arc)
    h_diag = hessian_diag(x, td)

    for i in 1:no_arc
        h[i,i] = h_diag[i]
    end

    return h

    #Link travel time = free flow time * ( 1 + b * (flow/capacity)^Power ).
end

function hessian_diag(x::Vector{Float64}, td::TA_Data)
    h_diag = Array{Float64}(undef, length(x))
    for i in 1:length(x)
        if td.power[i] >= 1.0
            h_diag[i] = td.free_flow_time[i] * td.b[i] * td.power[i] * (x[i]^(td.power[i]-1)) / (td.capacity[i]^td.power[i])
        else
            h_diag[i] = 0 # Some cases, power is zero.
        end
    end
    # h_diag = free_flow_time .* b .* power .* (x.^(power-1)) ./ (capacity.^power)

    return h_diag
    #Link travel time = free flow time * ( 1 + b * (flow/capacity)^Power ).
end

function all_or_nothing_single(travel_time::Vector{Float64}, td::TA_Data, graph, link_dic)
    local state::LightGraphs.DijkstraState{Float64, Int}
    x = zeros(size(td.init_node))

    for r in 1:size(td.travel_demand)[1]
        # for each origin node r, find shortest paths to all destination nodes
        state = TA_dijkstra_shortest_paths(graph, travel_time, r, td.init_node, td.term_node, td.first_thru_node)

        for s in 1:size(td.travel_demand)[2]
            # for each destination node s, find the shortest-path vector
            # load travel demand
            # x = x + travel_demand[r,s] * get_vector(state, r, s, link_dic)
            add_demand_vector!(x, td.travel_demand[r,s], state, r, s, link_dic)
        end
    end
    return x
end

# parallel computing version
# function all_or_nothing_parallel(travel_time::Vector{Float64}, td::TA_Data, graph, link_dic)
#     local state::LightGraphs.DijkstraState{Float64, Int}
#     vv = zeros(size(td.init_node))
#     x = zeros(size(td.init_node))

#     x = x .+ @distributed (+) for r in 1:size(td.travel_demand)[1]
#         # for each origin node r, find shortest paths to all destination nodes
#         # if there is any travel demand starting from node r.
#         vv = zeros(size(td.init_node))

#         if sum(td.travel_demand, dims=2)[r] > 0.0
#             state = TA_dijkstra_shortest_paths(graph, travel_time, r, td.init_node, td.term_node, td.first_thru_node)

#             for s=1:size(td.travel_demand)[2]
#                 # for each destination node s, find the shortest-path vector
#                 # v = get_vector(state, r, s, init_node, term_node)

#                 if td.travel_demand[r,s] > 0.0
#                     # load travel demand
#                     # vv = vv + travel_demand[r,s] * get_vector(state, r, s, link_dic)
#                     add_demand_vector!(vv, td.travel_demand[r,s], state, r, s, link_dic)
#                 end
#             end

#         end

#         vv
#     end

#     return x
# end

function all_or_nothing(travel_time::Vector{Float64}, td::TA_Data, graph, link_dic)
    # if nprocs() > 1 # if multiple CPU processes are available
    #     return all_or_nothing_parallel(travel_time, td, graph, link_dic)
    # else
    #     return all_or_nothing_single(travel_time, td, graph, link_dic)
    #     # when nprocs()==1, using @distributed just adds unnecessary setup time. I guess.
    # end

    return all_or_nothing_single(travel_time, td, graph, link_dic)
end

function ta_frank_wolfe(td::TA_Data; method=:bfw, max_iter_no=2000, step=:exact, log=:off, tol=1e-3)

    setup_time = time()

    if log==:on
        println("-------------------------------------")
        println("Network Name: $(ta_data.network_name)")
        println("Method: $method")
        println("Line Search Step: $step")
        println("Maximum Interation Number: $max_iter_no")
        println("Tolerance for AEC: $tol")
        println("Number of processors: ", nprocs())
    end

    # preparing a graph
    graph = create_graph(td.init_node, td.term_node)
    link_dic = sparse(td.init_node, td.term_node, collect(1:td.number_of_links))

    setup_time = time() - setup_time

    if log==:on
        println("Setup time = $setup_time seconds")
    end

    n_links = td.number_of_links

    iteration_time = time()

    # Finding a starting feasible solution
    travel_time = BPR(zeros(n_links), td)
    x0 = all_or_nothing(travel_time, td, graph, link_dic)

    # Initializing variables
    xk = x0
    tauk = 0.0
    yk_FW = x0
    sk_CFW = yk_FW
    Hk_diag = Vector{Float64}(undef, n_links)

    dk_FW = Vector{Float64}(undef, n_links)
    dk_bar = Vector{Float64}(undef, n_links)
    dk_CFW = Vector{Float64}(undef, n_links)
    dk = Vector{Float64}(undef, n_links)

    alphak = 0.0
    Nk = 0.0
    Dk = 0.0

    tauk = 0.0
    is_first_iteration = false
    is_second_iteration = false

    sk_BFW = yk_FW
    sk_BFW_old = yk_FW

    dk_bbar = Vector{Float64}(undef, n_links)
    muk = Vector{Float64}(undef, n_links)
    nuk = Vector{Float64}(undef, n_links)
    beta0 = 0.0
    beta1 = 0.0
    beta2 = 0.0

    for k in 1:max_iter_no
        # Finding yk
        travel_time = BPR(xk, td)
        yk_FW = all_or_nothing(travel_time, td, graph, link_dic)

        # Basic Frank-Wolfe Direction
        dk_FW = yk_FW - xk
        Hk_diag = hessian_diag(xk, td) # Hk_diag is a diagonal vector of matrix Hk

        # Finding a feasible direction
        if method == :fw # Original Frank-Wolfe
            dk = dk_FW
        elseif method == :cfw # Conjugate Direction F-W
            if k==1 || tauk > 0.999999 # If tauk=1, then start the process all over again.
                sk_CFW = yk_FW
                dk_CFW = sk_CFW - xk
            else
                dk_bar = sk_CFW - xk  # sk_CFW from the previous iteration k-1

                Nk = dot( dk_bar, Hk_diag .* dk_FW )
                Dk = dot( dk_bar, Hk_diag .* (dk_FW - dk_bar) )

                delta = 0.0001
                if Dk != 0.0 && 0.0 <= Nk/Dk <= 1.0 - delta
                    alphak = Nk/Dk
                elseif Dk != 0.0 && Nk/Dk > 1.0 - delta
                    alphak = 1.0 - delta
                else
                    alphak = 0.0
                end

                # Generating new sk_CFW and dk_CFW
                sk_CFW = alphak .* sk_CFW .+ (1.0 - alphak) .* yk_FW
                dk_CFW = sk_CFW .- xk
            end

            # Feasible Direction to Use for CFW
            dk = dk_CFW
        elseif method == :bfw # Bi-Conjugate Direction F-W

            if tauk > 0.999999
                is_first_iteration = true
                is_second_iteration = true
            end

            if k==1 || is_first_iteration       # First Iteration is like FW
                sk_BFW_old = yk_FW
                dk_BFW = dk_FW
                is_first_iteration = false
            elseif k==2 || is_second_iteration  # Second Iteration is like CFW
                dk_bar = sk_BFW_old - xk # sk_BFW_old from the previous iteration 1

                Nk = dot( dk_bar, Hk_diag .* dk_FW )
                Dk = dot( dk_bar, Hk_diag .* (dk_FW - dk_bar) )

                delta = 0.0001
                if Dk != 0.0 && 0.0 <= Nk/Dk <= 1.0 - delta
                    alphak = Nk/Dk
                elseif Dk != 0.0 && Nk/Dk > 1.0 - delta
                    alphak = 1.0 - delta
                else
                    alphak = 0.0
                end

                # Generating new sk_BFW and dk_BFW
                sk_BFW = alphak .* sk_BFW_old .+ (1-alphak) .* yk_FW
                dk_BFW = sk_BFW .- xk

                is_second_iteration = false
            else
                # println("over there $tauk")
                # sk_BFW, tauk is from iteration k-1
                # sk_BFW_old is from iteration k-2

                dk_bar  = sk_BFW - xk
                dk_bbar = tauk * sk_BFW - xk + (1.0 - tauk) * sk_BFW_old

                muk = - dot( dk_bbar, Hk_diag .* dk_FW ) / dot( dk_bbar, Hk_diag .* (sk_BFW_old - sk_BFW) )
                nuk = - dot( dk_bar, Hk_diag .* dk_FW ) / dot( dk_bar, Hk_diag .* dk_bar) + muk*tauk/(1-tauk)

                muk = max(0.0, muk)
                nuk = max(0.0, nuk)

                beta0 = 1.0 / ( 1.0 + muk + nuk )
                beta1 = nuk * beta0
                beta2 = muk * beta0

                sk_BFW_new = beta0 * yk_FW + beta1 * sk_BFW + beta2 * sk_BFW_old
                dk_BFW = sk_BFW_new - xk

                sk_BFW_old = sk_BFW
                sk_BFW = sk_BFW_new
            end

            # Feasible Direction to Use for BFW
            dk = dk_BFW
        else
            error("The type of Frank-Wolfe method is specified incorrectly. Use :fw, :cfw, or :bfw.")
        end
        # dk is now identified.

        if step == :exact
            # Line Search from xk in the direction dk
            optk = optimize(tau -> objective(xk+tau*dk, td), 0.0, 1.0, GoldenSection())
            tauk = optk.minimizer
        elseif step == :newton
            # Newton step
            tauk = - dot( gradient(xk, td), dk ) / dot( dk, Hk_diag.*dk )
            tauk = max(0.0, min(1.0, tauk))
        end

        # Average Excess Cost
        average_excess_cost = ( xk' * travel_time - yk_FW' * travel_time ) / sum(td.travel_demand)
        if log==:on
            @printf("k=%4d, tauk=%15.10f, objective=%15f, aec=%15.10f\n", k, tauk, objective(xk, td), average_excess_cost)
        end

        # rel_gap = ( objective(xk) - best_objective ) / best_objective

        # Convergence Test
        if average_excess_cost < tol
        # if rel_gap < tol
            break
        end

        # Update x
        new_x = xk + tauk * dk
        xk = new_x

        @assert minimum(xk) >= 0

    end

    iteration_time = time() - iteration_time

    if log==:on
        println("Iteration time = $iteration_time seconds")
    end

    return xk, travel_time, objective(xk, td)

end

#


module TestPredatorPrey
include("preamble_plots.jl")

using Bifurcations: resolved_points
using Bifurcations.Codim1
using Bifurcations.Codim2: NormalizingAS, BackReferencingAS
using Bifurcations.Examples.PredatorPrey

# Bifurcation points calculated in some version of Bifurcations.jl
CODIM1_KNOWN_POINTS = Dict(
    # PredatorPrey parameter (tuple) => points
    (0.0, 3, 5, 3) => [
        (Codim1.PointTypes.saddle_node,
         [0.0, 0.0, 0.6]),
        (Codim1.PointTypes.saddle_node,
         [1/3, 0, 0.821905296656178]),
        (Codim1.PointTypes.saddle_node,
         [0.41116155950979133, 0, 0.8329293250960671]),
        (Codim1.PointTypes.hopf,
         [1/3, 0.36576153597718997, 0.6715938474788321]),
    ],
)

@testset "smoke PredatorPrey" begin

    codim1_solver = init(PredatorPrey.prob, nominal_angle_rad=0.01)
    solve!(codim1_solver)

    @testset "known codim1 points" begin
        resolved = resolved_points(codim1_solver)
        known = CODIM1_KNOWN_POINTS[codim1_solver.prob.p.de_prob.p]
        @test length(resolved) == length(known)
        for (actual, desired) in zip(resolved, known)
            type_desired, u_desired = desired
            @test actual.point_type == type_desired
            @test actual.u ≈ u_desired  rtol=1e-4
        end
    end

    @testset "codim-2 ($(nameof(ASType))) from $(point.u0[[1, end]])" for
            ASType in [NormalizingAS, BackReferencingAS],
            point in special_intervals(codim1_solver)

        codim2_prob = BifurcationProblem(
            point,
            codim1_solver,
            (@lens _[3]),
            (0.0, 10.0),
            augmented_system = ASType(),
        )
        # TODO: Implement detection of Bogdanov-Takens bifurcation and
        # stop manually setting max_branches=0.
        codim2_solver = init(codim2_prob; max_branches=0)
        @test_nothrow solve!(codim2_solver)

        @testset "show" begin
            smoke_test_solver_show(codim2_solver)
        end

        @testset "plot" begin
            smoke_test_solver_plot(codim2_solver)
        end
    end
end

end  # module


# load packages
using CSV
using Random

"""
    set_up(;μ_θ_1::Real = -0.7, τ_θ_1::Real = 1.5, μ_θ_2::Real = 1.5, τ_θ_2::Real = 1.2, μ_θ_3::Real = -0.9, τ_θ_3::Real = 1.5, σ_ϵ::Real = 0.3, dt::Real=0.1, M::Int=5, N::Int=100)

Returns ground-truth parameter values and simulated data from the OU SDEMEM model.
"""
function set_up(;μ_θ_1::Real = -0.7, τ_θ_1::Real = 4,
                 μ_θ_2::Real = 2.3, τ_θ_2::Real = 10,
                 μ_θ_3::Real = -0.9, τ_θ_3::Real = 4,
                 σ_ϵ::Real = 0.3, dt::Real=0.05, M::Int=5, N::Int=100,seed::Int)

    # sample random effects
    Random.seed!(seed)

    ϕ = zeros(M,3)

    for i = 1:M
        ϕ[i,1] = μ_θ_1 + sqrt(1/τ_θ_1)*randn()
        ϕ[i,2] = μ_θ_2 + sqrt(1/τ_θ_2)*randn()
        ϕ[i,3] = μ_θ_3 + sqrt(1/τ_θ_3)*randn()
    end

    # model paramters
    η = [μ_θ_1; μ_θ_2; μ_θ_3; τ_θ_1; τ_θ_2; τ_θ_3]

    # define priors
    μ_0_1 = 0
    M_0_1 = 1
    α_1 = 2
    β_1 = 1

    μ_0_2 = 1
    M_0_2 = 1
    α_2 = 2
    β_2 = 1/2

    μ_0_3 = 0
    M_0_3 = 1
    α_3 = 2
    β_3 = 1

    prior_parameters_η = [μ_0_1 M_0_1 α_1 β_1;μ_0_2 M_0_2 α_2 β_2;μ_0_3 M_0_3 α_3 β_3]

    prior_parameters_σ_ϵ = [1; 1/0.4] # Gamma with scale parameter α = 1 and rate parameter β = 1/0.4

    # generate data

    Random.seed!(seed)
    y,x,t_vec = generate_data(N, M, η, ϕ, σ_ϵ,dt)

    # return: data, model parameteres, and parameteres for priors
    return y,x,t_vec,dt,η,σ_ϵ,ϕ,prior_parameters_η,prior_parameters_σ_ϵ

end

"""
    generate_data(N::Int, M::Int, η::Array, ϕ::Array, σ_ϵ::Real,dt::Real)

Generate data from the OU SDEMEM model.
"""
function generate_data(N::Int, M::Int, η::Array, ϕ::Array, σ_ϵ::Real,dt::Real)

    x_0 = zeros(M) # pre-allocate matriceis
    x = zeros(M,N)
    y = zeros(M,N)
    t_vec = zeros(N)

    for m = 1:M

        x[m,1] = x_0[m] # set start value
        y[m,1] = x[m,1] + σ_ϵ*randn()

        θ_1 = exp(ϕ[m,1]) # set parameters for subject m
        θ_2 = exp(ϕ[m,2])
        θ_3 = exp(ϕ[m,3])

        σ_cond = sqrt((θ_3^2/(2*θ_1))*(1-exp(-2*θ_1*dt))) # set latent process std for subject m

        # simulate process for subject m
        for t = 2:N
            t_vec[t] = t_vec[t-1] + dt
            μ_cond = θ_2 + (x[m,t-1]-θ_2)*exp(-θ_1*dt)
            x[m,t] = μ_cond + σ_cond*randn()
            y[m,t] = x[m,t] + σ_ϵ*randn()
        end
    end

    # return data
    return y,x,t_vec

end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

@testset "Basis Orthogonality"  begin

##

using Test
using ACE, JuLIP, JuLIP.Testing, QuadGK, LinearAlgebra
using ACE: TransformedJacobi, transform, transform_d, alloc_B, alloc_temp
using JuLIP: evaluate!

##

@info("Testing ortho-normality of Jacobi Polynomials")

α, β = 1 + rand(), 1 + rand()
N = 30
J = Jacobi(α, β, N)
Bj = alloc_B(J)
integrandJ = let J=J, B=Bj, α=α, β=β
   x ->  begin
            evaluate!(B, nothing, J, x)
            return B*B' * (1 - x)^α * (1+x)^β
         end
   end

Q = quadgk(integrandJ, -1, 1)[1]
println(@test round.(Q, digits=8) == Matrix(I, N+1, N+1))

##

@info("Testing ortho-normality of r-basis")

N = 20
rl, ru = 0.5, 3.0
fcut =  PolyCutoff2s(2, rl, ru)
trans = PolyTransform(2, 1.0)
P = TransformedJacobi(N, trans, fcut)
B = alloc_B(P)
tmp = alloc_temp(P)

integrand = let P = P, B = B, tmp = tmp
   r ->  begin
            evaluate!(B, tmp, P, r)
            return B * B' * abs(transform_d(P.trans, r))
         end
   end

Q = quadgk(integrand, rl, ru)[1]
println(@test round.(Q, digits=8) == Matrix(I, N+1, N+1))

##

@info("Testing (near-) orthonormality of r-basis via sampling")

# P = TransformedJacobi(...) from previous cell
Nsamples = 100_000

G = let
   G = zeros(length(P), length(P))
   for n = 1:Nsamples
      evaluate!(B, tmp, P, ACE.Utils.rand_radial(P))
      G += B * B'
   end
   G
end

println(@test cond(G) < 1.1)

##

@info("Testing (near-)orthonormality of Ylm-basis via sampling")

SH = ACE.SphericalHarmonics.SHBasis(5)

function gramian(SH::ACE.SphericalHarmonics.SHBasis, Nsamples=100_000)
   lenY = length(SH)
   G = zeros(ComplexF64, lenY, lenY)
   Y = alloc_B(SH)
   tmp = alloc_temp(SH)
   for n = 1:Nsamples
      evaluate!(Y, tmp, SH, ACE.Utils.rand_sphere())
      for i = 1:lenY, j = 1:lenY
         G[i,j] += Y[i] * Y[j]'
      end
   end
   return G / Nsamples
end

G = gramian(SH)
println(@test cond(G) < 1.1)

##

@info("Testing (near-)orthonormality of A-basis via sampling")

aceB = SHIPBasis( SparseSHIP(3, 5), trans, fcut )
function evalA(aceB, tmp, Rs)
   Zs = zeros(Int16, length(Rs))
   ACE.precompute_A!(tmp, aceB, Rs, Zs, 1)
   return tmp.A[1]
end

function A_gramian(aceB, Nsamples = 100_000)
   tmp = alloc_temp(aceB)
   lenA = length(tmp.A[1])
   G = zeros(ComplexF64, lenA, lenA)
   for n = 1:Nsamples
      R = ACE.rand_vec(aceB.J)
      A = evalA(aceB, tmp, [R])
      for i = 1:lenA, j = 1:lenA
         G[i,j] +=  A[i] * A[j]'
      end
   end
   return G
end

G = A_gramian(aceB)
println(@test cond(G) < 1.2)

end


using Base.Test
using QuantumOptics

b = SpinBasis(1//2)

psi1 = spinup(b)
psi2 = spindown(b)

rho = tensor(psi1, dagger(psi1))
sigma = tensor(psi2, dagger(psi2))

@test tracedistance(rho, sigma) == 1.
@test tracedistance(rho, rho) == 0.
@test tracedistance(sigma, sigma) == 0.


# SPDX-License-Identifier: MIT

using CUDA, CUDA.CUSPARSE, SparseArrays
import LogExpFunctions: logsumexp
using MarkovModels
import MarkovModels: logaddexp
using Test

const D = 5 # vector/matrix dimension
const B = 2 # batch size
const S = 3 # number of states
const N = 5 # number of frames

const T = Float64
const SF = LogSemifield{T}

function makefsm(SF, S)
    fsm = FSM{SF}()

    prev = addstate!(fsm, pdfindex = 1)
    setinit!(prev)
    link!(fsm, prev, prev)
    for s in 2:S
        state = addstate!(fsm, pdfindex = s)
        link!(fsm, prev, state)
        link!(fsm, state, state)
        prev = state
    end
    setfinal!(prev)
    renormalize!(fsm)
    fsm
end

function forward(A, init, lhs::AbstractMatrix{T}) where T
    N = size(lhs, 2)
    log_α = similar(lhs)
    log_α[:, 1] = lhs[:, 1] + init
    for n in 2:N
        log_α[:, n] = lhs[:,n] + dropdims(logsumexp(A .+ log_α[:,n-1], dims = 1),
                                          dims = 1)
    end
    log_α
end

function backward(Aᵀ, final, lhs::AbstractMatrix{T}) where T
    N = size(lhs, 2)
    log_β = similar(lhs)
    log_β[:, end] = final
    for n in N-1:-1:1
        log_β[:, n] = dropdims(logsumexp(Aᵀ .+ log_β[:,n+1] .+ lhs[:,n+1], dims = 1),
                               dims = 1)
    end
    log_β
end

function forward_backward(A, Aᵀ, init, final, lhs)
    log_α = forward(A, init, lhs)
    log_β = backward(Aᵀ, final, lhs)
    log_γ = log_α .+ log_β
    sums = logsumexp(log_γ, dims = 1)
    exp.(log_γ .- sums), minimum(sums)
end

@testset "logaddexp" begin
    @test (@inferred logaddexp(2.0, 3.0)) ≈ log(exp(2.0) + exp(3.0))
    @test (@inferred logaddexp(10002.0, 10003.0)) ≈ 10000 + logaddexp(2.0, 3.0)
end

if CUDA.functional()
    @testset "linalg" begin
        D = 5

        x = ones(T, D)
        y = spzeros(T, D)
        y[2] = T(3)
        xg = CuArray(x)
        yg = CuSparseVector(y)
        out = similar(xg)
        MarkovModels.elmul_svdv!(out, yg, xg)
        @test all(convert(Vector{T}, out) .≈ x .* y)

        x = ones(T, D)
        y = spzeros(T, D)
        xg = CuArray(x)
        yg = CuSparseVector(y)
        out = similar(xg)
        MarkovModels.elmul_svdv!(out, yg, xg)
        @test all(convert(Vector{T}, out) .≈ x .* y)

        for Yᵀ in [spdiagm(y), spzeros(T, length(y))]
            Yᵀ = spdiagm(y)
            Yᵀg = CuSparseMatrixCSC(Yᵀ)
            Yg = CuSparseMatrixCSR( Yᵀg.colPtr, Yᵀg.rowVal, Yᵀg.nzVal, (D,D))
            out = similar(xg)
            MarkovModels.mul_smdv!(out, Yg, xg)
            @test all(convert(Vector{T}, out) .≈ transpose(Yᵀ) * x)
        end

        x = ones(T, D)
        y = spzeros(T, D)
        X = ones(T, D, D)
        Y = spzeros(T, D, D)
        xg = CuArray(x)
        yg = CuSparseVector(y)
        Xg = CuArray(X)
        Yg = CuSparseMatrixCSC(Y)
        Yg = CuSparseMatrixCSR(Yg.colPtr, Yg.rowVal, Yg.nzVal, Yg.dims)
        out = similar(Xg)
        MarkovModels.elmul_svdm!(out, yg, Xg)
        @test all(convert(Matrix{T}, out) .≈ X .* reshape(y, :, 1))

        x = ones(T, D)
        y = spzeros(T, D)
        X = ones(T, D, D)
        Y = spzeros(T, D, D)
        Y[1,1] = T(1)
        Y[1,2] = T(2)
        Y[2,1] = T(3)
        xg = CuArray(x)
        yg = CuSparseVector(y)
        Xg = CuArray(X)
        Yg = CuSparseMatrixCSC(Y)
        Yg = CuSparseMatrixCSR(Yg.colPtr, Yg.rowVal, Yg.nzVal, Yg.dims)
        out = similar(Xg)
        MarkovModels.elmul_svdm!(out, yg, Xg)
        @test all(convert(Matrix{T}, out) .≈ X .* reshape(y, :, 1))

        A = spzeros(3, 3)
        Ag = CuSparseMatrixCSC(A)
        Ag = CuSparseMatrixCSR(Ag.colPtr, Ag.rowVal, Ag.nzVal, Ag.dims)
        Y = ones(3, 3)
        Yg = CuArray(Y)
        out = similar(Yg)
        MarkovModels.mul_smdm!(out, Ag, Yg)
        @test all(convert(Matrix{T}, out) .≈ A' * Y)

        A = spzeros(3, 3)
        A[1, 1] = 1/2
        A[1, 2] = 1/2
        A[2, 2] = 1/2
        A[2, 3] = 1/2
        A[3, 3] = 1/2
        Ag = CuSparseMatrixCSC(A)
        Ag = CuSparseMatrixCSR(Ag.colPtr, Ag.rowVal, Ag.nzVal, Ag.dims)
        Y = ones(3, 3)
        Yg = CuArray(Y)
        out = similar(Yg)
        MarkovModels.mul_smdm!(out, Ag, Yg)
        @test all(convert(Matrix{T}, out) .≈ A' * Y)

        A = ones(T, 3, 2)
        Ag = CuArray(A)
        B = ones(T, 2, 3)
        Bg = CuArray(B)
        out = similar(Ag, 3, 3)
        MarkovModels.mul_dmdm!(out, Ag, Bg)
        @test all(convert(Matrix{T}, out) .≈ A * B)
    end
end

@testset "forward_backward" begin
    lhs = ones(T, S, N)

    fsm = makefsm(SF, S)
    cfsm = compile(fsm, allocator = zeros)

    γ_ref, ttl_ref = forward_backward(
        convert(Matrix{T}, cfsm.A),
        convert(Matrix{T}, cfsm.Aᵀ),
        convert(Vector{T}, cfsm.π),
        convert(Vector{T}, cfsm.ω),
        convert(Matrix{T}, lhs)
    )

    γ_dcpu, ttl_dcpu = @inferred stateposteriors(cfsm, lhs)
    @test all(γ_ref .≈ γ_dcpu)
    @test ttl_ref ≈ ttl_dcpu

    cfsm = compile(fsm, allocator = spzeros)
    γ_scpu, ttl_scpu = @inferred stateposteriors(cfsm, lhs)
    @test all(γ_ref .≈ γ_scpu)
    @test ttl_ref ≈ ttl_scpu

    if CUDA.functional()
        cfsm = compile(fsm, allocator = zeros) |> gpu
        γ_dgpu, ttl_dgpu = @inferred stateposteriors(cfsm, CuArray(lhs))
        @test all(γ_ref .≈ convert(Matrix{T}, γ_dgpu))
        @test ttl_ref ≈ ttl_dgpu

        cfsm = compile(fsm, allocator = spzeros) |> gpu
        γ_sgpu, ttl_sgpu = @inferred stateposteriors(cfsm, CuArray(lhs))
        @test all(γ_ref .≈ convert(Matrix{T}, γ_sgpu))
        @test ttl_ref ≈ ttl_sgpu
    end
end

@testset "batch forward_backward" begin
    lhs = ones(T, S, N, B)

    fsm = makefsm(SF, S)
    cfsm = compile(fsm, allocator = zeros)

    γ_ref, ttl_ref = forward_backward(
        convert(Matrix{T}, cfsm.A),
        convert(Matrix{T}, cfsm.Aᵀ),
        convert(Vector{T}, cfsm.π),
        convert(Vector{T}, cfsm.ω),
        convert(Matrix{T}, lhs[:, :, 1])
    )

    γ_dcpu, ttl_dcpu = @inferred stateposteriors(cfsm, lhs)
    for b in 1:B
        @test all(γ_ref .≈ γ_dcpu[:,:,b])
        @test ttl_ref ≈ ttl_dcpu[b]
    end

    cfsm = compile(fsm, allocator = spzeros)
    γ_scpu, ttl_scpu = @inferred stateposteriors(cfsm, lhs)
    for b in 1:B
        @test all(γ_ref .≈ γ_scpu[:,:,b])
        @test ttl_ref ≈ ttl_scpu[b]
    end

    if CUDA.functional()
        cfsm = compile(fsm, allocator = zeros) |> gpu
        γ_dgpu, ttl_dgpu = @inferred stateposteriors(cfsm, CuArray(lhs))
        for b in 1:B
            @test all(γ_ref .≈ convert(Array{T,3}, γ_dgpu)[:,:,b])
            @test ttl_ref ≈ convert(Vector{T}, ttl_dgpu)[b]
        end

        cfsm = compile(fsm, allocator = spzeros) |> gpu
        γ_sgpu, ttl_sgpu = @inferred stateposteriors(cfsm, CuArray(lhs))
        for b in 1:B
            @test all(γ_ref .≈ convert(Array{T,3}, γ_sgpu)[:,:,b])
            @test ttl_ref ≈ convert(Vector{T}, ttl_sgpu)[b]
        end
    end
end


module QuadEig

using LinearAlgebra, SparseArrays, SuiteSparse

export linearize, deflate, eigbasis

### PencilScaling ##########################################################################
# Compute the optimal scaling
############################################################################################
struct PencilScaling{T}
    γ::T
    δ::T
end

function pencilscaling(A0::M, A1::M, A2::M; scaling = nothing, kw...) where {T,M<:AbstractMatrix{T}}
    scaling === nothing && return PencilScaling(T(1), T(1))
    n0, n1, n2 = opnorm.(Matrix.((A0, A1, A2)))
    τ = n1 / sqrt(n0 * n2)
    if τ <= 1
        γ = sqrt(n0 / n2)
        δ = 2 / (n0 + γ * n1)
    elseif scaling === -
        # distinct tropical roots, choose γ+ to favor large eigenvalues
        γ = n1 / n2
        δ = 1 / max(n2 * γ^2, n1 * γ, n0)
    else
        # distinct tropical roots, choose γ- to favor small eigenvalues
        γ = n0 / n1
        δ = 1 / max(n2 * γ^2, n1 * γ, n0)
    end
    return PencilScaling(T(γ), T(δ))
end

function scale!(A0, A1, A2, s::PencilScaling)
    γ, δ = scalingfactors(s)
    if γ != 1 || δ != 1
        A0 .*= γ^2 * δ
        A1 .*= γ * δ
        A2 .*= δ
    end
    return A0, A1, A2
end

scalingfactors(s::PencilScaling) = s.γ, s.δ

### QuadPencil #############################################################################
# Encodes the quadratic pencil problem
############################################################################################
struct QuadPencil{T<:Complex,M<:AbstractMatrix{T}}
    A0::M
    A1::M
    A2::M
    scaling::PencilScaling{T}
end

function quadpencil(A0, A1, A2; kw...)
    size(A0) == size(A1) == size(A2) && size(A0, 1) == size(A0, 2) ||
        throw(DimensionMismatch("Expected square matrices of the same size"))
    ptype = complex(promote_type(eltype(A0), eltype(A1), eltype(A2)))
    A0´ = copy_to_eltype(A0, ptype)
    A1´ = copy_to_eltype(A1, ptype)
    A2´ = copy_to_eltype(A2, ptype)
    s = pencilscaling(A0´, A1´, A2´; kw...)
    scale!(A0´, A1´, A2´, s)
    return QuadPencil(A0´, A1´, A2´, s)
end

# This also ensures any Adjoint/Transpose wrapper is removed
copy_to_eltype(A, ::Type{T}) where {T} = copy!(similar(A, T, size(A)), A)

Base.size(q::QuadPencil, n...) = size(q.A0, n...)

### QuadPencilQR ###########################################################################
# Pivoted QR-factorization
############################################################################################
struct QuadPencilPQR{T,M}
    pencil::QuadPencil{T,M}
    Q0::M
    RP0::M
    Q2´::M
    RP2::M
end

function pqr(pencil::QuadPencil)
    qr0 = pqr(pencil.A0)
    qr2 = pqr(pencil.A2)
    Q0, RP0 = getQ(qr0), getRP´(qr0)
    Q2´, RP2 = getQ´(qr2), getRP´(qr2)
    return QuadPencilPQR(pencil, Q0, RP0, Q2´, RP2)
end

# Sparse QR from SparseSuite is also pivoted
pqr(a::SparseMatrixCSC) = qr(a)
pqr(a) = qr!(copy(a), Val(true))

### Linearized Pencils #####################################################################
# Build second companion linearization, or its Q*C2*V rotation
############################################################################################
struct Linearization{T,M<:AbstractMatrix{T},R}
    pencilpqr::QuadPencilPQR{T,M}
    A::M
    B::M
    Q::M
    V::M
    deflate_tol::R
end

Linearization(q::QuadPencilPQR{T,M}, A::M, B::M, Q::M, V::M) where {T,M} =
    Linearization(q, A, B, Q, V, zero(real(T)))

linearize(A0, A1, A2; kw...) = linearize(pqr(quadpencil(A0, A1, A2; kw...)))

function linearize(q::QuadPencilPQR)
    A0, A1, A2 = q.pencil.A0, q.pencil.A1, q.pencil.A2
    o, z = one(A1), zero(A1)
    Q = [q.Q2´ z; z q.Q0']
    V = [o z; z q.Q0]
    A = [q.Q2´*A1 -q.Q2´*q.Q0; q.RP0 z]
    B = [q.RP2 z; z o]
    B .= .- B
    return Linearization(q, A, B, Q, V)
end

function Base.show(io::IO, l::Linearization{T,M}) where {T,M}
    print(io, summary(l), "\n",
"  Matrix size    : $(size(l.A, 1)) × $(size(l.A, 2))
  Matrix type    : $M
  Scalings γ, δ  : $(real.(scalingfactors(l)))
  Deflated       : $(deflationstring(l))")
end

Base.summary(l::Linearization{T}) where {T} =
    "Linearization{T}: second companion linearization of quadratic pencil"

deflationstring(l::Linearization) =
    isdeflated(l) ? "true ($(size(l.V, 1)) -> $(size(l.A, 1)))" : "false"

isdeflated(l) = !iszero(l.deflate_tol)

Base.size(l::Linearization, n...) = size(l.A, n...)

dimpencil(l::Linearization) = size(l.V, 1) ÷ 2

scalingfactors(l::Linearization) = scalingfactors(l.pencilpqr.pencil.scaling)

### deflate ################################################################################
# Compute deflated pencil
############################################################################################
deflate(A0, A1, A2; kw...) = deflate(linearize(A0, A1, A2; kw...); kw...)

function deflate(l::Linearization{T}; atol = sqrt(eps(real(T))), kw...) where {T}
    isdeflated(l) && return l
    n = size(l, 1) ÷ 2
    r0, r2 = nonzero_rows(l, atol)
    r0 == n || r2 == n && return l
    s = n - r2
    X = view(l.A, 1+r2:n, 1:r2+s+r0) # [X21 X22 X23]
    ZX´ = fod_z´(X, 1+s:s+r0+r2)
    A, B = deflatedAB(l.A, l.B, ZX´, r0, r2, s)
    V = view(l.V, :, 1:n+r0) * ZX´
    Q = l.Q[[1:r2; n+1:n+r0], :]
    q = l.pencilpqr
    return Linearization(q, A, B, Q, V, atol)
end

# get some columns of Z' in a full orthogonal decomposition of X
fod_z´(X::SparseMatrixCSC, cols) = _fod_z´(sparse(X'), cols)
fod_z´(X::SubArray{<:Any,2,<:SparseMatrixCSC}, cols) = _fod_z´(sparse(X'), cols)
fod_z´(X, cols = :) = _fod_z´(X', cols)
_fod_z´(X´, cols) = getQ(qr(X´), cols)

# carry out quadeig deflation
function deflatedAB(lA::SparseMatrixCSC, lB, ZX´, r0, r2, s)
    selected_rows = Isparse(size(lA, 1), [1:r2; 1+s+r2:s+r0+r2], :)
    A = view(selected_rows * lA, :, 1:s+r0+r2) * ZX´
    B = view(selected_rows * lB, :, 1:s+r0+r2) * ZX´
    return A, B
end

function deflatedAB(lA::AbstractMatrix, lB, ZX´, r0, r2, s)
    A = similar(lA, r0+r2, r0+r2)
    mul!(view(A, 1:r2, :), view(lA, 1:r2, 1:r2+s+r0), ZX´)
    mul!(view(A, 1+r2:r0+r2, :), view(lA, 1+s+r2:r2+s+r0, 1:r2+s+r0), ZX´)
    B = similar(A)
    mul!(view(B, 1:r2, :), view(lB, 1:r2, 1:r2+s+r0), ZX´)
    mul!(view(B, 1+r2:r0+r2, :), view(lB, 1+s+r2:r2+s+r0, 1:r2+s+r0), ZX´)
    return A, B
end

### quadeig ################################################################################
# Generalized Schur factorization (aka QZ factorization)
############################################################################################
function retarded_projector(A0, A1, A2; kw...)
    d´ = deflate(A0', A1', A2'; kw...)
    n = dimpencil(d´)
    atol = d´.deflate_tol
    γ, _ = scalingfactors(d´)

    # e = eigen(Matrix(d´.A'), Matrix(d´.B'))
    # b = d´.Q' * e.vectors
    # display(e.values)
    # @show real.(b[1:n, 18])
    # @show e.values[18] .* real.(b[1+n:2n, 18])

    λ, Z = retarded_deflated_eigbasis(d´.A', d´.B', γ, atol)
    basis = d´.Q' * Z

    nullbasis = nullspace_basis(A2, atol)

    # display(sparse(d´.Q'))
    A, B = AB_C2(A0',A1',A2'; kw...)
    # display(sparse(A'))
    display(d´.V' * A' * d´.Q' - d´.A')
    display(sparse(basis))

    Z11 = hcat(view(basis, 1:n, :), nullbasis)
    Z21 = hcat(view(basis, n+1:2n, :), zero(nullbasis))
    # display(sparse(Z11))
    # display(sparse(Z21))
    # projector = rdiv!(Z21, lu!(Z11))  # equivalent to Z21*Z11
    # return projector
end

function retarded_deflated_eigbasis(A, B, γ, atol)
    s = schur!(Matrix(A), Matrix(B))
    λs = chop!(γ .*  s.α ./ s.β, atol)
    which_retarded = abs2.(λs) .<= 1
    howmany_retarded = count(which_retarded)
    ordschur!(s, which_retarded)
    defbasis = view(s.Z, :, 1:howmany_retarded)
    λ = chop!(γ .*  s.α ./ s.β, atol)
    return λ, defbasis
end

function nullspace_basis(A, atol)
    q = pqr(A)
    RP´ = getRP´(q)
    r = nonzero_rows(RP´, atol)
    cols = r+1:size(A, 2)
    basis = getQ(q, cols)
    return basis
end

### Tools ##################################################################################

getQ(qr::Factorization, cols = :) = qr.Q * Idense(size(qr, 1), cols)
getQ(qr::SuiteSparse.SPQR.QRSparse, cols = :) =  Isparse(size(qr, 1), :, qr.prow) * sparse(qr.Q * Idense(size(qr, 1), cols))

getQ´(qr::Factorization, cols = :) = qr.Q' * Idense(size(qr, 1), cols)
getQ´(qr::SuiteSparse.SPQR.QRSparse, cols = :) = sparse((qr.Q * Idense(size(qr, 1), cols))') * Isparse(size(qr,1), qr.prow, :)

getRP´(qr::Factorization) = qr.R * qr.P'
getRP´(qr::SuiteSparse.SPQR.QRSparse) = qr.R * Isparse(size(qr, 2), qr.pcol, :)

getPR´(qr::Factorization) = qr.P * qr.R'
getPR´(qr::SuiteSparse.SPQR.QRSparse) = Isparse(size(qr, 2), :, qr.pcol) * qr.R'

Idense(n, ::Colon) = Matrix(I, n, n)

function Idense(n, cols)
    m = zeros(Bool, n, length(cols))
    for (j, col) in enumerate(cols)
        m[col, j] = true
    end
    return m
end

# Equivalent to I(n)[rows, cols], but faster
Isparse(n, rows, cols) = Isparse(inds(rows, n), inds(cols, n))

function Isparse(rows, cols)
    rowval = Int[]
    nzval = Bool[]
    colptr = Vector{Int}(undef, length(cols) + 1)
    colptr[1] = 1
    for (j, col) in enumerate(cols)
        push_rows!(rowval, nzval, rows, col)
        colptr[j+1] = length(rowval) + 1
    end
    return SparseMatrixCSC(length(rows), length(cols), colptr, rowval, nzval)
end

inds(::Colon, n) = 1:n
inds(is, n) = is

function push_rows!(rowval, nzval, rows::AbstractUnitRange, col)
    if col in rows
        push!(rowval, 1 + rows[1 + col - first(rows)] - first(rows))
        push!(nzval, true)
    end
    return nothing
end

function push_rows!(rowval, nzval, rows, col)
    for (j, row) in enumerate(rows)
        if row == col
            push!(rowval, j)
            push!(nzval, true)
        end
    end
    return nothing
end

function nonzero_rows(l::Linearization{T}, atol = sqrt(eps(real(T)))) where {T}
    n = size(l, 2) ÷ 2
    rank0, rank2 = nonzero_rows(view(l.A, n+1:2n, 1:n), atol), nonzero_rows(view(l.B, 1:n, 1:n), atol)
    return rank0, rank2
end

function nonzero_rows(p::QuadPencilPQR{T}, atol = sqrt(eps(real(T)))) where {T}
    rank0, rank2 = nonzero_rows(p.qr0.R, atol), nonzero_rows(p.qr2.R, atol)
    return rank0, rank2
end

function nonzero_rows(m::AbstractMatrix{T}, atol = sqrt(eps(real(T)))) where {T}
    n = 0
    for row in eachrow(m)
        all(z -> abs(z) < atol, row) && break
        n += 1
    end
    return n
end

function chop!(A::AbstractArray{T}, atol = sqrt(eps(real(T)))) where {T}
    for (i, a) in enumerate(A)
        if abs(a) < atol
            A[i] = zero(T)
        elseif abs(a) > 1/atol || isnan(a)
            A[i] = T(Inf)
        end
    end
    return A
end

function AB_C2(A0, A1, A2; kw...)
    p = quadpencil(A0, A1, A2; kw...)
    n = size(p, 1)
    o, z = one(p.A1), zero(p.A1)
    A = [p.A1 -o; p.A0 z]
    B = [-p.A2 z; z -o]
    return A, B
end

end # Module


include("TypesTest.jl")
include("BasisTest.jl")
include("NuclearRepulsionTest.jl")
include("OverlapTest.jl")
include("KineticEnergyTest.jl")
include("NuclearAttractionTest.jl")
include("ElecElecRepulsionTest.jl")

run_typesparticles_test()
run_typesbasis_test()
run_readbasis_test()
run_retreadbasis_test()
run_nuclearrepulsion_test()
run_buildoverlap_test()
run_buildkineticenergy_test()
run_buildnuclrattract_test()
run_buildelecelecrepul_test()