# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %%
"""https://github.com/moruten/julia-code/blob/99b50bf329195ba102a9496602a22138f07d186b/Test-2021-8-12-no6.ipynb"""
module Original

using Plots
using Random

function any_width_rand!(n_rand,width_l,width_r)
    rands=zeros(n_rand)
    L=width_r - width_l
    for i=1:n_rand
        rands[i] = width_l + L*rand()
    end
    return rands
end

function action_cal(x)
    exp(0.5*x^2)
end

function MP_test(x_ini,x_pro)
    S_ini = action_cal(x_ini)
    S_pro = action_cal(x_pro)

    if exp(S_ini - S_pro) > rand()
        return x_pro
    else
        return x_ini
    end
end

n_rand=1_000_000

MP_rands=zeros(n_rand)

x_propose  = 0.0
x_ini      = 0.0
width_rand = 2.0

@time for i=2:n_rand
    x_propose   = x_ini + any_width_rand!(1,-width_rand,width_rand)[1]
    MP_rands[i] = MP_test(x_ini,x_propose)
    global x_ini       = MP_rands[i]
end

f(x) = exp(-0.5*x^2)/sqrt(2*pi)
xr   = range(-3,3,length=10001)

plot(MP_rands,st=:histogram,norm=:true,nbins=100,label="Metropolis method")
plot!(xr,f.(xr),lw=3,label="Gaussian") |> display

end

# %%
"""https://github.com/moruten/julia-code/blob/99b50bf329195ba102a9496602a22138f07d186b/Test-2021-8-12-no6.ipynb"""
module Original_plus

using Plots
using Random

function any_width_rand!(n_rand,width_l,width_r)
    rands=zeros(n_rand)
    L=width_r - width_l
    for i=1:n_rand
        rands[i] = width_l + L*rand()
    end
    return rands
end

function action_cal(x)
    exp(0.5*x^2)
end

function MP_test(x_ini,x_pro)
    S_ini = action_cal(x_ini)
    S_pro = action_cal(x_pro)

    if exp(S_ini - S_pro) > rand()
        return x_pro
    else
        return x_ini
    end
end

n_rand=1_000_000

MP_rands=zeros(n_rand)

x_propose  = 0.0
x_ini      = 0.0
width_rand = 2.0

@time for i=2:n_rand
    x_propose   = x_ini + any_width_rand!(1,-width_rand,width_rand)[1]
    MP_rands[i] = MP_test(x_ini,x_propose)
    global x_ini       = MP_rands[i]
end

f(x) = exp(-0.5*x^2)/sqrt(2*pi)
xr   = range(-3,3,length=10001)

using QuadGK
G(x) = exp(-exp(0.5*x^2))
Z = quadgk(G, -Inf, Inf)[1]
g(x) = G(x)/Z

histogram(MP_rands, norm=true, nbins=100, label="Metropolis method", alpha=0.3)
plot!(xr, f.(xr), lw=3, label="Gaussian")
plot!(xr, g.(xr), lw=3, label="exp(-exp(0.5*x^2))/Z", ls=:dash) |> display

end

# %%
module Corrected

using Plots
using Random

function any_width_rand!(n_rand,width_l,width_r)
    rands=zeros(n_rand)
    L=width_r - width_l
    for i=1:n_rand
        rands[i] = width_l + L*rand()
    end
    return rands
end

function action_cal(x)
    exp(0.5*x^2)
end

function MP_test(x_ini,x_pro)
    S_ini = action_cal(x_ini)
    S_pro = action_cal(x_pro)

    if S_ini / S_pro > rand() # <-- 変更はこの行だけ！
        return x_pro
    else
        return x_ini
    end
end

n_rand=1_000_000

MP_rands=zeros(n_rand)

x_propose  = 0.0
x_ini      = 0.0
width_rand = 2.0

@time for i=2:n_rand
    x_propose   = x_ini + any_width_rand!(1,-width_rand,width_rand)[1]
    MP_rands[i] = MP_test(x_ini,x_propose)
    global x_ini       = MP_rands[i]
end

f(x) = exp(-0.5*x^2)/sqrt(2*pi)
xr   = range(-5,5,length=10001)

plot(MP_rands,st=:histogram,norm=:true,nbins=100,label="Metropolis method", alpha=0.3)
plot!(xr,f.(xr),lw=3,label="Gaussian") |> display

end

# %%
module MH

using Plots
using Random

function any_width_rand!(n_rand,width_l,width_r)
    rands=zeros(n_rand)
    L=width_r - width_l
    for i=1:n_rand
        rands[i] = width_l + L*rand()
    end
    return rands
end

function action_cal(x)
    exp(0.5*x^2)
end

function MP_test(x_ini,x_pro)
    S_ini = action_cal(x_ini)
    S_pro = action_cal(x_pro)
    a = S_ini/S_pro
    if a > 1 || rand() ≤ a
        return x_pro
    else
        return x_ini
    end
end

n_rand=1_000_000

MP_rands=zeros(n_rand)

x_propose  = 0.0
x_ini      = 0.0
width_rand = 2.0

@time for i=2:n_rand
    x_propose   = x_ini + any_width_rand!(1,-width_rand,width_rand)[1]
    MP_rands[i] = MP_test(x_ini,x_propose)
    global x_ini       = MP_rands[i]
end

f(x) = exp(-0.5*x^2)/sqrt(2*pi)
xr   = range(-5,5,length=10001)

plot(MP_rands,st=:histogram,norm=:true,nbins=100,label="MH method", alpha=0.3)
plot!(xr,f.(xr),lw=3,label="Gaussian") |> display

end

# %%
using Plots
using Distributions

function update(H, prop, x)
    y = prop(x)
    a = exp(H(x) - H(y))
    a > 1 || rand() ≤ a ? y : x
end

function iter_update(H, prop, x0, N; warmup=min(N ÷ 10, 10^3))
    x = x0
    for _ in 1:warmup
        x = update(H, prop, x)
    end
    X = similar([x], N)
    X[1] = x
    for i in 2:N
        X[i] = update(H, prop, X[i-1])
    end
    X
end

H(x) = x^2/2
prop(x) = x + 0.2randn()
x0 = randn()
N = 10^6
X = iter_update(H, prop, x0, N)

histogram(X; norm=true, alpha=0.3, bin=100, label="MH sample")
plot!(x -> pdf(Normal(), x), -5, 5; ls=:dash, lw=2, label="normal dist.")

# %%
H(x) = x ≤ 0 ? Inf : x
prop(x) = x + 0.1randn()
x0 = abs(randn())
N = 10^6
X = iter_update(H, prop, x0, N)

histogram(X; norm=true, alpha=0.3, bin=100, label="MH sample")
plot!(x -> pdf(Exponential(), x), eps(), 10; label="exp. dist.")
plot!(xlim=(-0.1, 8))

# %%
H(x) = x ≤ 0 ? Inf : x/3 - (4 - 1)*log(x)
prop(x) = x + randn()
x0 = abs(randn())
N = 10^6
X = iter_update(H, prop, x0, N)

histogram(X; norm=true, alpha=0.3, bin=100, label="MH sample")
plot!(x -> pdf(Gamma(4, 3), x), eps(), 50; label="Gamma(4, 3)")

# %%


"""
nsolsc_solution_test

Test nsolsc with the atan function. Check answers and iteration stats.
"""
function nsolsc_solution_test()
    #
    # Local convergence with forward difference derivative
    #
    sdatal = nsolsc(atan, 1.0)
    solok = (abs(sdatal.solution) < 1.e-8)
    funok = (abs(sdatal.functionval) < 1.e-8)
    hs = size(sdatal.history)
    histok = (hs[1] == 5)
    locok = funok && solok && histok
    if ~locok
        println("local FD fails")
    end
    #
    # Local convergence with analytic derivative
    #
    sdataa = nsolsc(atan, 1.0, fpatan)
    solok = (abs(sdataa.solution) < 1.e-8)
    funok = (abs(sdataa.functionval) < 1.e-8)
    hs = size(sdataa.history)
    histok = (hs[1] == 5)
    analyticok = funok && solok && histok
    if ~analyticok
        println("failure with analytic derivative ")
        println(sdataa)
    end
    #
    # Global convergence
    #
    sdatag = nsolsc(atan, 10.0; maxit = 11, armfix = true)
    solok = (abs(sdatag.solution) < 1.e-8)
    funok = (abs(sdatag.functionval) < 1.e-8)
    hs = size(sdatag.history)
    histok = (hs[1] == 12)
    globok = funok && solok && histok
    if ~globok
        println("global FD fails")
    end
    #
    # Global convergence with parab3p
    #
    sdatap3p = nsolsc(atan, 30.0; rtol = 1.e-10, maxit = 11)
    solok = (abs(sdatap3p.solution) < 1.e-8)
    funok = (abs(sdatap3p.functionval) < 1.e-8)
    hs = size(sdatap3p.history)
    histok = (hs[1] == 12)
    p3pok = funok && solok && histok
    if ~p3pok
        println("parab3p fails")
    end

    #
    # Local convergence with secant method
    #
    #sdatas=nsolsc(atan,1.0; solver="secant")
    sdatas = secant(atan, 1.0)
    solok = (abs(sdatas.solution) < 1.e-10)
    funok = (abs(sdatas.functionval) < 1.e-10)
    hs = size(sdatas.history)
    histok = (hs[1] == 6)
    secantok = funok && solok && histok
    if ~secantok
        println("secant failure")
    end
    #
    # Initialize secant method when x0=0
    #
    #zedata=nsolsc(fcos,0.0;solver="secant",rtol=1.e-9)
    zedata = secant(fcos, 0.0; solver = "secant", rtol = 1.e-9)
    solution = 7.390851333858823e-01
    solok = (abs(zedata.solution - solution) < 1.e-9)
    funok = (abs(zedata.functionval) < 1.e-9)
    hs = size(zedata.history)
    histok = (hs[1] == 7)
    zecok = funok && solok && histok
    if ~zecok
        println("local FD fixup at zero fails")
    end
    #
    # Tricky line search problem
    #
    sdatal = nsolsc(
        linatan,
        200.0;
        sham = 5,
        maxit = 20,
        armmax = 10,
        armfix = true,
        rtol = 1.e-10,
    )
    solution = -100.0
    solok = (abs(sdatal.solution - solution) < 1.e-8)
    funok = (abs(sdatal.functionval) < 1.e-8)
    hs = size(sdatal.history)
    histok = (hs[1] == 6)
    shamfastok = funok && solok && histok
    if ~shamfastok
        println("Fast Shamanskii response FAILURE")
    end
    #
    # Test linesearch failure complaints.
    #
    armfail = nsolsc(atan, 10.0; armmax = 1, armfix = true, printerr = false)
    afok = false
    if armfail.idid == false && armfail.errcode == 1
        afok = true
    else
        println("Armijo failure test FAILED.")
    end
    #
    # Test residual failure mode and no history.
    #
    resok = false
    resfail = nsolsc(atan, 10.0; maxit = 3, armfix = true, keepsolhist = false)
    if resfail.idid == false && resfail.errcode == 10
        resok = true
    else
        println("Residual failure test FAILED.")
    end
    #
    # Test stagnation mode
    #
    stagdatan =
        nsolsc(ftanx, 4.5, ftanxp; rtol = 1.e-17, atol = 1.e-17, armfix = true, maxit = 14)
    fvals = stagdatan.history
    avals = stagdatan.stats.iarm
    ifvals = stagdatan.stats.ifun
    jvals = stagdatan.stats.ijac
    stagl = (length(fvals) == 6)
    stagf = (fvals[5] < 1.e-15)
    stags = (avals[6] == 5) && (ifvals[6] == 6) && (jvals[6] == 1)
    stagok = stagl && stags && stagf
    if ~stagok
        println("Stagnation test FAILED")
    end
    #
    #
    # Test chord method
    #
    lttest = nsolsc(atan, 0.5; solver = "chord")
    fvals = lttest.history
    chordl = (length(fvals) == 11)
    ratl = fvals[11] / fvals[10]
    chordr = (abs(ratl - 0.25) < 1.e-7)
    solok = (fvals[11] < 1.e-6)
    chordok = chordl && chordr && solok
    if ~chordok
        println("Chord test failed")
    end
    #
    # Make sure nsolsc knows if the solution and the intial iterate are
    # the same
    #
    lotout = nsolsc(flot, 0.0)
    lotok = (lotout.errcode == -1)
    if ~lotok
        println("Lottery test failed")
    end
    return locok &&
           globok &&
           secantok &&
           analyticok &&
           zecok &&
           shamfastok &&
           afok &&
           resok &&
           p3pok &&
           chordok &&
           lotok
end

function flot(x)
    flot = x * exp(x)
end


using Plots, GraphRecipes
import Plots: plot
export plot, spectrum

function plot(ϕ::FourierField{S, T}; ylims = ()) where {S, T <: FourierMetaData}
    dims = length(ϕ.metadata.grid.grid)
    if dims == 1
        x = ϕ.metadata.grid.grid[1][:]
        dd = ϕ.metadata.transform.backward * ϕ.data
        plot(x, real.(dd))
        plot!(xlabel = "x")
        plot!(ylabel = ϕ.metadata.name)
        plot!(legend = false, ylims = ylims)
    elseif dims == 2
        x = ϕ.metadata.grid.grid[1][:]
        y = ϕ.metadata.grid.grid[2][:]
        dd = ϕ.metadata.transform.backward * ϕ.data
        contourf(x, y, real.(dd)', linewidth = 0, color = :thermometer)
        contourf!(xlabel = "x")
        contourf!(ylabel = "y")
        contourf!(title =  ϕ.metadata.name)
    else
        print("Plotting is not supported for fields ")
        print("with dimensions greater ≥ 3")
    end
end

plot(ϕ::Field{S, T}; kwargs...) where {S <: FourierField, T} = plot(ϕ.data, kwargs)

function spectrum(ϕ::FourierField)
    dims = length(ϕ.metadata.grid.grid)
    if dims == 1
        f = ϕ.data
        g = log10.(abs.(f))[1:div(length(f),2)+1]
        ymax = maximum(g) * 1.1
        ymin = maximum(g)  - 16
        wvi = collect(1:div(length(f),2)+1) .- 1
        p1 = scatter(wvi, g, label = ϕ.metadata.name)
        scatter!(xlabel = "wavenumber index")
        scatter!(ylabel = "log10(spectral amplitude)")
        scatter!(ylims = (ymin, ymax))
        return p1
    elseif dims == 2
        g = log10.(abs.(ϕ.data))
        cmax = maximum(g) * 1.1
        cmin = maximum(g)  - 16
        clims = (cmin, cmax)
        p1 = heatmap(g, clims = clims, color = :thermometer)
        scatter!(xlabel = "x, wavenumber index")
        scatter!(ylabel = "y, wavenumber index")
        scatter!(title = "log10(spectral amplitude)")
        return p1
    end
    return nothing
end

plot(a::AbstractExpression) = plot(to_expr(a))
plot(a::Equation) = plot(to_expr(a))


using Pkg
Pkg.add("HTTP")
using Pkg
Pkg.add("Optim")
using Pkg
Pkg.add("GLM")
using Pkg
Pkg.add("LinearAlgebra")
using Pkg
Pkg.add("Random")
using Pkg
Pkg.add("Statistics")
using Pkg
Pkg.add("DataFrames")
using Pkg
Pkg.add("CSV")
using Pkg
Pkg.add("FreqTables")
using Optim, HTTP, GLM, LinearAlgebra, Random, Statistics, DataFrames, CSV, FreqTables
function ps3()
    #Question1
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS3-gev/nlsw88w.csv"
df = CSV.read(HTTP.get(url).body, DataFrame)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occupation

function mlogit(Beta, X, Z, y)

    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigBeta = [reshape(Beta,K,J-1) zeros(K)]

    bigZ = zeros(N,J)
    for j=1:J
        bigZ[:,j] = Z[:,j]-Z[:,J]
    end

    num = zeros(N,J)
    dem = zeros(N)
    for j=1:J
        XZ=cat(X,bigZ[:,j],dims=3)
        num[:,j] = exp.(XZ*bigBeta[:,j])
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)

    loglike = -sum( bigY.*log.(P) )

    return loglike
end

mlogit_hat_optim = optimize(b-> mlogit(b,X,Z, y), rand(7*(size(X,2))), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
println(mlogit_hat_optim.minimizer)
#Question2
#Gamma represents the change in utility in percentage terms#
#Question3#
function nested_logit(alpha, X, Z, y, nesting_structure)
    beta = alpha[1:end-3]
    lambda = alpha[end-2:end-1]
    gamma = alpha[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigBeta = [repeat(beta[1:K],1,length(nesting_structure[1])) repeat(beta[K+1:2K],1,length(nesting_structure[2])) zeros(K)]
end
ps3()


# LU decomposition of a band-diagonal matrix (ala Numerical Recipes)
#
function bandec(a::Array{Float64,2},n::Int64,m1::Int64,m2::Int64,al::Array{Float64,2},indx::Vector{Int64})
#      INTEGER m1,m2,mp,mpl,n,np,indx(n)
#      REAL d,a(np,mp),al(np,mpl),TINY
mm=m1+m2+1
@assert(size(a)  == (n,mm))
@assert(size(al) == (n,m1))
@assert(length(indx) == n)
TINY=1.e-20
#      INTEGER i,j,k,l,mm
#      REAL dum
l=m1
dum = 0.0
for i=1:m1
  for j=m1+2-i:mm
    a[i,j-l]=a[i,j]
  end
  l -= 1
  for j=mm-l:mm
    a[i,j]=0.0
  end
end
d=1.0
l=m1
for k=1:n
  dum=a[k,1]
  i=k
  if (l < n)
    l += 1
  end
  for j=k+1:l
    if (abs(a[j,1]) > abs(dum))
      dum=a[j,1]
      i=j
    end
  end
  indx[k]=i
  if (dum == 0.0)
    a[k,1]=TINY
  end
  if (i != k)
    d = -d
    for j=1:mm
      dum=a[k,j]
      a[k,j]=a[i,j]
      a[i,j]=dum
    end
  end
  for i=k+1:l
    dum=a[i,1]/a[k,1]
    al[k,i-k]=dum
    for j=2:mm
      a[i,j-1] = a[i,j] - dum*a[k,j]
    end
    a[i,mm]=0.
  end
end
return d
end

#  (C) Copr. 1986-92 Numerical Recipes Software

function bandec(a::Array{Complex{Float64},2},n::Int64,m1::Int64,m2::Int64,al::Array{Complex{Float64},2},indx::Vector{Int64})
#      INTEGER m1,m2,mp,mpl,n,np,indx(n)
#      REAL d,a(np,mp),al(np,mpl),TINY
mm=m1+m2+1
@assert(size(a)  == (n,mm))
@assert(size(al) == (n,m1))
@assert(length(indx) == n)
TINY=complex(1.e-20,0.0)
#      INTEGER i,j,k,l,mm
#      REAL dum
l=m1
dum = complex(0.0,0.0)
for i=1:m1
  for j=m1+2-i:mm
    a[i,j-l]=a[i,j]
  end
  l -= 1
  for j=mm-l:mm
    a[i,j]=complex(0.0,0.0)
  end
end
d=1.0
l=m1
for k=1:n
  dum=a[k,1]
  i=k
  if (l < n)
    l += 1
  end
  for j=k+1:l
    if (abs(a[j,1]) > abs(dum))
      dum=a[j,1]
      i=j
    end
  end
  indx[k]=i
  if (dum == 0.0)
    a[k,1]=TINY
  end
  if (i != k)
    d = -d
    for j=1:mm
      dum=a[k,j]
      a[k,j]=a[i,j]
      a[i,j]=dum
    end
  end
  for i=k+1:l
    dum=a[i,1]/a[k,1]
    al[k,i-k]=dum
    for j=2:mm
      a[i,j-1] = a[i,j] - dum*a[k,j]
    end
    a[i,mm]=0.
  end
end
return d
end

#  (C) Copr. 1986-92 Numerical Recipes Software


#cd("C:\\Users\\dcirk\\Documents\\CSCE689\\Project")

##### Facebook Analysis ######

# Load in Facebook Data
using DelimitedFiles, MatrixMarket, JSON
include("MaxRecipFunctions.jl")
edgelist = readdlm("Datasets\\facebook-wall.txt")
edgelist = convert(Matrix{Int64}, edgelist[:, 1:2])

# Convert to adjacency matrix
n_edges = size(edgelist)[1]
n_nodes = maximum(edgelist)
A = sparse(edgelist[:, 1], edgelist[:, 2], repeat([1], n_edges), n_nodes, n_nodes)
# Rewire algorithm
time_facebook = @timed begin
A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
facebook_recip = recip/n_edges
# reciprocity of rewired graph
facebook_rewire_recip = recip_rewire/n_edges
# upper bound
facebook_upper_bound = upperbound/n_edges
facebook_edges = n_edges
facebook_nodes = n_nodes

##### EU Email Analysis ######

A = MatrixMarket.mmread("Datasets\\email-Eu-core-temporal\\email-Eu-core-temporal.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]
# Rewire algorithm
time_email_EU = @timed begin
A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
email_EU_recip = recip/n_edges
# reciprocity of rewired graph
email_EU_rewire_recip = recip_rewire/n_edges
# upper bound
email_EU_upper_bound = upperbound/n_edges
email_EU_edges = n_edges
email_EU_nodes = n_nodes

##### Slashdot Analysis ######

A = MatrixMarket.mmread("Datasets\\soc-Slashdot0811\\soc-Slashdot0811.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]
# Rewire algorithm
time_slashdot = @timed begin
    A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
slashdot_recip = recip/n_edges
# reciprocity of rewired graph
slashdot_rewire_recip = recip_rewire/n_edges
# upper bound
slashdot_upper_bound = upperbound/n_edges
slashdot_edges = n_edges
slashdot_nodes = n_nodes

#### Mathoverflow Analysis ####

A = MatrixMarket.mmread("Datasets\\sx-mathoverflow\\sx-mathoverflow.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]
# Rewire algorithm
time_mathoverflow = @timed begin
    A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
mathoverflow_recip = recip/n_edges
# reciprocity of rewired graph
mathoverflow_rewire_recip = recip_rewire/n_edges
# upper bound
mathoverflow_upper_bound = upperbound/n_edges
mathoverflow_edges = n_edges
mathoverflow_nodes = n_nodes

#### College Message Analysis ####

A = MatrixMarket.mmread("Datasets\\CollegeMsg\\CollegeMsg.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]
# Rewire algorithm
time_collegemsg = @timed begin
    A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
collegemsg_recip = recip/n_edges
# reciprocity of rewired graph
collegemsg_rewire_recip = recip_rewire/n_edges
# upper bound
collegemsg_upper_bound = upperbound/n_edges
collegemsg_edges = n_edges
collegemsg_nodes = n_nodes

#### Epinions Analysis ####
A = MatrixMarket.mmread("Datasets\\soc-Epinions1\\soc-Epinions1.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]

# Rewire algorithm
time_epinions = @timed begin
    A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
epinions_recip = recip/n_edges
# reciprocity of rewired graph
epinions_rewire_recip = recip_rewire/n_edges
# upper bound
epinions_upper_bound = upperbound/n_edges
epinions_edges = n_edges
epinions_nodes = n_nodes

#### Twitter Analysis ####
A = MatrixMarket.mmread("Datasets\\higgs-twitter\\higgs-twitter_mention.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]

# Rewire algorithm
time_twitter = @timed begin
    A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
twitter_recip = recip/n_edges
# reciprocity of rewired graph
twitter_rewire_recip = recip_rewire/n_edges
# upper bound
twitter_upper_bound = upperbound/n_edges
twitter_edges = n_edges
twitter_nodes = n_nodes

#### PA Analysis ####
A = MatrixMarket.mmread("Datasets\\PAdata.mtx")
# Convert to adjacency matrix
n_edges = sum(A)
n_nodes = size(A)[1]

# Rewire algorithm
time_PA = @timed begin
    A_rewire = greedyRewire(A)
end

# Reciprocity before and after
recip = compute_reciprocity(A)
recip_rewire = compute_reciprocity(A_rewire)

out_deg = transpose(sum(A, dims = 1))
in_deg = sum(A, dims = 2)
min_deg = zeros(n_nodes)
for i in 1:n_nodes
    min_deg[i] = min(in_deg[i], out_deg[i])
end
upperbound = sum(min_deg)

# reciproicty of original graph
PA_recip = recip/n_edges
# reciprocity of rewired graph
PA_rewire_recip = recip_rewire/n_edges
# upper bound
PA_upper_bound = upperbound/n_edges

PA_edges = n_edges
PA_nodes = n_nodes

output_data =
["Facebook" facebook_recip facebook_rewire_recip facebook_upper_bound facebook_nodes facebook_edges time_facebook[2]
"EU Email" email_EU_recip email_EU_rewire_recip email_EU_upper_bound email_EU_nodes email_EU_edges time_email_EU[2]
"Slashdot" slashdot_recip slashdot_rewire_recip slashdot_upper_bound slashdot_nodes slashdot_edges time_slashdot[2]
"Mathoverflow" mathoverflow_recip mathoverflow_rewire_recip mathoverflow_upper_bound mathoverflow_nodes mathoverflow_edges time_mathoverflow[2]
"College Msg" collegemsg_recip collegemsg_rewire_recip collegemsg_upper_bound collegemsg_nodes collegemsg_edges time_collegemsg[2]
"Epinions" epinions_recip epinions_rewire_recip epinions_upper_bound epinions_nodes epinions_edges time_epinions[2]
"Twitter" twitter_recip twitter_rewire_recip twitter_upper_bound twitter_nodes twitter_edges time_twitter[2]
"Preferential" PA_recip PA_rewire_recip PA_upper_bound PA_nodes PA_edges time_PA[2]]

#writedlm("rewiringResults.txt", output_data)


using RobotZoo
using RobotDynamics
using StaticArrays
using BenchmarkTools
using FiniteDiff
using LinearAlgebra
using Test
using Random
Random.seed!(1)
const RD = RobotDynamics

# Continuous Linear model
n,m = 10,5
A,B = RobotZoo.RandomLinearModels.gencontrollable(n,m, :continuous)
d = randn(10)
model = RobotZoo.LinearModel(A,B,d)
@test RD.dims(model) == (n,m,n)
x,u = rand(model)
@test length(x) == n
@test length(u) == m
@test RD.dynamics(model, x, u) ≈ A*x + B*u + d
xn = similar(x)
RD.dynamics!(model, xn, x, u)
@test xn ≈ A*x + B*u + d

@test RobotZoo.iscontrollable(model)

J = zeros(n, n+m)
z = RD.KnotPoint(x, u, 0, NaN)
RD.jacobian!(RD.InPlace(), RD.UserDefined(), model, J, xn, z)
@test J[:,1:n] ≈ A
@test J[:,n+1:end] ≈ B

Jfd = similar(J)
FiniteDiff.finite_difference_jacobian!(
    Jfd,
    (y,z)->RD.dynamics!(model, y, z[1:n], z[n+1:end]),
    Vector(z.z)
)
@test Jfd ≈ J rtol=1e-6

# Linearized Continuous Model
Random.seed!(1)
model0 = RobotZoo.Cartpole()
xe = [0,pi,0,0]
ue = [0.]
model = RobotZoo.LinearModel(RD.InPlace(), RD.ForwardAD(), model0, xe, ue)
n,m = RD.dims(model)
@test n == 4
@test m == 1
dx = randn(4)
du = randn(1)
@test RD.dynamics(model, dx*0, du*0) ≈ RD.dynamics(model0, xe, ue)
@test norm(RD.dynamics(model, dx*1e-3, du*1e-3) - RD.dynamics(model0, xe, ue)) < 0.1

Jfd = zeros(4,5)
y = zeros(4)
z = RD.KnotPoint{n,m}(xe, ue, 0, NaN)
FiniteDiff.finite_difference_jacobian!(
    Jfd,
    (y,z)->RD.dynamics!(model0, y, z[1:n], z[n+1:end]),
    Vector(z.z)
)
@test Jfd[:,1:n] ≈ RobotZoo.getstatecoeffs(model)
@test Jfd[:,n+1:end] ≈ RobotZoo.getcontrolcoeffs(model)
@test RobotZoo.isstable(model) == false

# Linearize about stable equilibrium
xe = [0,0.0,0,0]
ue = [0.]
model_stable = RobotZoo.LinearModel(RD.InPlace(), RD.ForwardAD(), model0, xe, ue)
@test RobotZoo.isstable(model_stable)

# Use random generator on linear model
n,m = 12,6
model_rand = rand(RobotZoo.LinearModel, n, m)
@test RD.dims(model_rand) == (n,m,n)
@test RobotZoo.iscontrollable(model_rand)

## Discrete Linear Model
n,m = 10,5
h = 0.1
A,B = RobotZoo.RandomLinearModels.gencontrollable(n, m, :discrete)
d = randn(n)
model = RobotZoo.DiscreteLinearModel(h,A,B)
@test RobotZoo.getstatecoeffs(model) == A
@test RobotZoo.getcontrolcoeffs(model) == B
@test RobotZoo.getaffinecoeffs(model) == zeros(n)
@test RobotZoo.iscontrollable(model)

model = RobotZoo.DiscreteLinearModel(h,A,B,d)
@test RobotZoo.getstatecoeffs(model) == A
@test RobotZoo.getcontrolcoeffs(model) == B
@test RobotZoo.getaffinecoeffs(model) == d
@test RobotZoo.iscontrollable(model)

C = randn(n,n)
model = RobotZoo.DiscreteLinearModel(h,A,B,C,d)
@test RobotZoo.getstatecoeffs(model) == A
@test RobotZoo.getcontrolcoeffs(model) == B
@test RobotZoo.getnextstatecoeffs(model) == C
@test RobotZoo.getaffinecoeffs(model) == d

@test_nowarn RobotZoo.checktimestep(model, h)
@test_nowarn RobotZoo.checktimestep(model, h + eps(h))
@test_throws ArgumentError RobotZoo.checktimestep(model, h*1.1)

@test RD.dims(model) == (n,m,n)
x,u = rand(model)
z = RD.KnotPoint(x,u,0.,h)

# Explicit
model_explicit = RobotZoo.DiscreteLinearModel(h,A,B,d)
@test RobotZoo.getnextstatecoeffs(model_explicit) ≈ -I(n)
@test RD.discrete_dynamics(model_explicit, x, u, 0, h) ≈ A*x + B*u + d

J = zeros(n,n+m)
xn = zeros(n)
RD.jacobian!(RD.InPlace(), RD.UserDefined(), model_explicit, J, xn, z)
@test J[:,1:n] ≈ A
@test J[:,n+1:end] ≈ B

# Implicit
model_implicit = RobotZoo.DiscreteLinearModel(h,A,B,C,d)
@test RobotZoo.getnextstatecoeffs(model_implicit) ≈ C
@test RD.discrete_dynamics(model_implicit, x, u, 0, h) ≈ -C\(A*x + B*u + d)

RD.jacobian!(RD.InPlace(), RD.UserDefined(), model_implicit, J, xn, z)
@test J[:,1:n] ≈ -C\A
@test J[:,n+1:end] ≈ -C\B

## Discretize and Linearize
model0 = RobotZoo.Cartpole()
n,m = RD.dims(model0)

# Implicit
dmodel_explicit = RD.DiscretizedDynamics{RD.RK3}(model0)
xe = [0,pi,0,0]
ue = [0.]
model_explicit = RobotZoo.DiscreteLinearModel(
    RD.InPlace(),
    RD.ForwardAD(),
    dmodel_explicit,
    xe, ue, 0., h
)

J = zeros(n, n+m)
z = RD.KnotPoint{n,m}([xe;ue], 0.0, h)
xn = zeros(n)
RD.jacobian!(RD.StaticReturn(), RD.ForwardAD(), dmodel_explicit, J, xn, z)
@test J[:,1:n] ≈ RobotZoo.getstatecoeffs(model_explicit)
@test J[:,n+1:end] ≈ RobotZoo.getcontrolcoeffs(model_explicit)
@test RD.discrete_dynamics(dmodel_explicit, z) ≈
    RobotZoo.getaffinecoeffs(model_explicit)
@test RobotZoo.getnextstatecoeffs(model_explicit) == -I(n) * one(eltype(A))
@test RobotZoo.isstable(model_explicit) == false
@test RobotZoo.iscontrollable(model_explicit)

xe_stable = zeros(n)
model_explicit_stable = RobotZoo.DiscreteLinearModel(
    RD.InPlace(),
    RD.ForwardAD(),
    dmodel_explicit,
    xe_stable, ue, 0., h
)
@test RobotZoo.isstable(model_explicit_stable) == true

# Implicit
dmodel_implicit = RD.DiscretizedDynamics{RD.ImplicitMidpoint}(model0)
model_implicit= RobotZoo.DiscreteLinearModel(
    RD.InPlace(),
    RD.ForwardAD(),
    dmodel_implicit,
    xe, ue, 0., h
)
@test RD.dims(model_implicit) == (n,m,n)

Jc = zeros(n,n+m)
RD.jacobian!(RD.StaticReturn(), RD.ForwardAD(), model0, Jc, xn, z)
A = Jc[:,1:n]
B = Jc[:,n+1:end]
@test RobotZoo.getstatecoeffs(model_implicit) ≈ h/2*A + I
@test RobotZoo.getnextstatecoeffs(model_implicit) ≈ h/2*A - I
@test RobotZoo.getcontrolcoeffs(model_implicit) ≈ h*B

# Exponential integrator
model_linear = RobotZoo.LinearModel(RD.InPlace(), RD.ForwardAD(),
    model0, xe, ue, affine=false)
model_exp = RobotZoo.DiscreteLinearModel(model_linear, h)

# Compare the linear models
dx = randn(n) * 1e-2
du = randn(m) * 1e-2
xn_exp = RD.discrete_dynamics(model_exp, dx, du, 0.0, h)
xn_implicit = RD.discrete_dynamics(model_implicit, dx, du, 0, h) - model_implicit.d
xn_explicit = RD.discrete_dynamics(model_explicit, dx, du, 0, h) - model_explicit.d
@test norm(xn_exp - xn_implicit) < 1e-3
@test norm(xn_exp - xn_explicit) < 1e-3


#
# Function that computes the kinetic energy
#

function kinetic(n :: Int64, v :: Array{Float64})

  kinetic = 0.
  for i in 1:n
    kinetic = kinetic + 0.5*(v[i,1]^2 + v[i,2]^2)
  end

  return kinetic

end


@testset "Data" begin

@test artificialIn_SAR_image(2) == 2*π*ones(2,2)

@test artificial_S1_slope_signal(20,0.) == repeat([-π/2],20)

@test ismissing( artificial_S1_signal(-1.) )
@test ismissing( artificial_S1_signal(2.) )
@test artificial_S1_signal(2) == [-3*π/4,-3*π/4]

# for the remainder check data types only
@test length( artificial_S1_signal(20) ) == 20

@test size( artificial_S2_whirl_image(64) ) == (64,64)
@test length(artificial_S2_whirl_image(64)[1,1]) == 3

@test size( artificial_S2_rotation_image(64) ) == (64,64)
@test length(artificial_S2_rotation_image(64)[1,1]) == 3

@test size( artificial_S2_whirl_patch(8) ) == (8,8)
@test length(artificial_S2_whirl_patch(8)[1,1]) == 3

@test size( artificial_SPD_image(8) ) == (8,8)
@test size(artificial_SPD_image(8)[1,1]) == (3,3)

@test size( artificial_SPD_image2(8) ) == (8,8)
@test size(artificial_SPD_image2(8)[1,1]) == (3,3)
@test eltype( artificial_SPD_image2(8))  == Array{Float64,2}

@test length( artificial_S2_lemniscate([0.,0.,1.], 20))  == 20
@test length( artificial_S2_lemniscate([0.,0.,1.], 20)[1]) == 3

end


using PyCall
pygui(:qt)
using PyPlot
matplotlib[:rcdefaults]()
ion()

using Unwrap

x = linspace(-1,1,200)
y = x
r² = broadcast(+, x.^2, (y').^2)
z = exp(-r²)
φ = angle(exp(im*z*10π))

φ_old = copy(φ)
φ1 = unwrap2d(φ)
φ2 = unwrap2d(φ1) # Iterating improves result

function plot_map(args...;kwargs...)
    p = pcolormesh(args...;
                   rasterized=true,
                   cmap=plt[:cm][:get_cmap]("viridis"),
                   kwargs...)
    margins(0,0)
    p
end

function plot_phase(φ,gs,i,label,ylabels=false)
    fig[:add_subplot](get(gs, (1,i-1)))
    pm = plot_map(x,y,φ/2π)
    xlabel(L"x")
    if ylabels
        ylabel(L"y")
    else
        gca()[:set_yticklabels]([])
    end
    cax=fig[:add_subplot](get(gs, (0,i-1)))
    colorbar(pm,cax=cax,orientation="horizontal",label=label)
    cax[:set_xticklabels](cax[:get_xticklabels](),rotation=45)
    cax[:xaxis][:tick_top]()
    cax[:xaxis][:set_ticks_position]("both")
    cax[:xaxis][:set_label_position]("top")
end
fig = figure("2d",figsize=(10,4.5))
clf()
gs = matplotlib[:gridspec][:GridSpec](2,3,height_ratios=(1,10))
plot_phase(φ_old,gs,1,L"Wrapped [$\pi$]",true)
plot_phase(φ1,gs,2,L"1st unwrapped [$\pi$]")
plot_phase(φ2,gs,3,L"2nd unwrapped [$\pi$]")
gs[:tight_layout](fig)

savefig("2d.png",dpi=72,transparent=true)


points(S::Jacobi,n)=fromcanonical(S,gaussjacobi(n,S.a,S.b)[1])

immutable JacobiTransformPlan{DD,T,TT}
    space::Jacobi{TT,DD}
    points::Vector{T}
    weights::Vector{T}
end

plan_transform(S::Jacobi,v::Vector) = JacobiTransformPlan(S,gaussjacobi(length(v),S.a,S.b)...)
function transform(S::Jacobi,vals,plan::JacobiTransformPlan)
#    @assert S==plan.space
    x,w = plan.points, plan.weights
    V=jacobip(0:length(vals)-1,S.a,S.b,x)'
    w2=w.*(1-x).^(S.a-plan.space.a).*(1+x).^(S.b-plan.space.b)   # need to weight if plan is different
    nrm=(V.^2)*w2

    V*(w2.*vals)./nrm
end

immutable JacobiITransformPlan{DD,T}
    space::Jacobi{Float64,DD}
    points::Vector{T}
end

plan_itransform(S::Jacobi,cfs::Vector) = JacobiITransformPlan(S,points(S,length(cfs)))
function itransform(S::Jacobi,cfs,plan::JacobiITransformPlan)
#    @assert S==plan.space
    jacobip(0:length(cfs)-1,S.a,S.b,tocanonical(S,plan.points))*cfs
end

function coefficients(f::Vector,a::Jacobi,b::Chebyshev)
    if domain(a) == domain(b) && (!isapproxinteger(a.a-0.5) || !isapproxinteger(a.b-0.5))
        cjt(f,a.a,a.b)
    else
        defaultcoefficients(f,a,b)
    end
end
function coefficients(f::Vector,a::Chebyshev,b::Jacobi)
    if domain(a) == domain(b) && (!isapproxinteger(b.a-0.5) || !isapproxinteger(b.b-0.5))
        icjt(f,b.a,b.b)
    else
        defaultcoefficients(f,a,b)
    end
end

function coefficients(f::Vector,a::Jacobi,b::Jacobi)
    if domain(a) == domain(b) && (!isapproxinteger(a.a-b.a) || !isapproxinteger(a.b-b.b))
        jjt(f,a.a,a.b,b.a,b.b)
    else
        defaultcoefficients(f,a,b)
    end
end


using BenchmarkTools
using ForwardDiff
using Test
using Random
using StaticArrays
using LinearAlgebra
import TrajectoryOptimization: kinematics, ∇rotate, ∇composition1, ∇composition2, vector
const TO = TrajectoryOptimization

q1 = rand(UnitQuaternion)
q2 = rand(UnitQuaternion)
r = @SVector rand(3)
ω = @SVector rand(3)

ϕ = 0.1*@SVector [1,0,0]
@test logm(expm(ϕ)) ≈ ϕ
@test expm(logm(q1)) ≈ q1
@test angle(expm(ϕ)) ≈ 0.1
# @btime expm($ϕ)
# @btime logm($u1)
@test norm(q1 ⊕ ϕ) ≈ 1
@test q1 ⊖ q2 isa SVector{3}
@test (q1 ⊕ ϕ) ⊖ q1 ≈ ϕ

# Test inverses
q3 = q2*q1
@test q2\q3 ≈ q1
@test q3/q1 ≈ q2
@test inv(q1)*r ≈ q1\r
@test r ≈ q3\(q2*q1*r)

q = q1
rhat = UnitQuaternion(r)
@test q*r ≈ Vmat()*Lmult(q)*Rmult(q)'*Vmat()'r
@test q*r ≈ Vmat()*Lmult(q)*Lmult(rhat)*Tmat()*SVector(q)
@test q*r ≈ Vmat()*Rmult(q)'*Rmult(rhat)*SVector(q)

@test ForwardDiff.jacobian(q->UnitQuaternion{VectorPart}(q)*r,SVector(q)) ≈ ∇rotate(q,r)
# @btime ForwardDiff.jacobian(q->UnitQuaternion(q)*$r,SVector($q))
# @btime ∇rotate($q,$r)

@test ForwardDiff.jacobian(q->SVector(q2*UnitQuaternion{VectorPart}(q)),SVector(q1)) ≈ ∇composition1(q2,q1)
@test ForwardDiff.jacobian(q->SVector(UnitQuaternion{VectorPart}(q)*q1),SVector(q2)) ≈ ∇composition2(q2,q1)
# @btime ForwardDiff.jacobian(q->SVector($u2*UnitQuaternion(q)),SVector($u1))
# @btime ∇composition1($u2,$u1)
# @btime ForwardDiff.jacobian(q->SVector(UnitQuaternion(q)*$u1),SVector($u2))
# @btime ∇composition2($u2,$u1)

@test Lmult(q) ≈ ∇composition1(q,q2)

ϕ = @SVector zeros(3)
@test TO.∇differential(q) ≈ Lmult(q)*jacobian(VectorPart,ϕ)
@test TO.∇differential(q) ≈ Lmult(q)*jacobian(ExponentialMap,ϕ)
@test TO.∇differential(q) ≈ Lmult(q)*jacobian(CayleyMap,ϕ)
@test TO.∇differential(q) ≈ Lmult(q)*jacobian(MRPMap,ϕ)

############################################################################################
#                          ROLL, PITCH, YAW EULER ANGLES
############################################################################################

e0 = @SVector [deg2rad(45), deg2rad(60), deg2rad(20)]
e1 = RPY(e0...)
# e2 = RotXYZ(e0...)

# @test rotmat(e1) == RotX(e0[1])*RotY(e0[2])*RotZ(e0[3])
# @test rotmat(e1) ≈ e2

# @test e1*r ≈ e2*r
# @btime $e2*$r
# @btime $e1*$r

R = rotmat(e1)
e1_ = TO.from_rotmat(rotmat(e1))
@test e1_.θ ≈ e1.θ
@test e1_.ψ ≈ e1.ψ
@test e1_.ϕ ≈ e1.ϕ

# @test rotmat(e1*e1) ≈ RotXYZ(e2*e2)

e1 = RPY(rand(3)...)
e2 = RPY(rand(3)...)
R = rotmat(e2*e1)

# Test inverses
e1,e2 = rand(RPY), rand(RPY)
e3 = e2*e1
@test e2\e3 ≈ e1
@test e3/e1 ≈ e2
@test r ≈ e3\(e2*e1*r)

@test ∇rotate(e1,r) isa SMatrix{3,3}
@test ∇composition1(e2,e1) isa SMatrix{3,3}
@test ∇composition2(e2,e1) isa SMatrix{3,3}

# sϕ1,cϕ1 = sincos(e1.ϕ)
# sθ1,cθ1 = sincos(e1.θ)
# sψ1,cψ1 = sincos(e1.ψ)
#
# sϕ2,cϕ2 = sincos(e2.ϕ)
# sθ2,cθ2 = sincos(e2.θ)
# sψ2,cψ2 = sincos(e2.ψ)
#
# R11 = cθ2*cψ2*(cθ1*cψ1) - cθ2*sψ2*(sϕ1*sθ1*cψ1+cϕ1*sψ1) + sθ2*(-cϕ1*sθ1*cψ1+sϕ1*sψ1)
# R11 ≈ R[1,1]
# R12 = cθ2*cψ2*(-cθ1*sψ1) - cθ2*sψ2*(-sϕ1*sθ1*sψ1+cϕ1*cψ1) + sθ2*(cϕ1*sθ1*sψ1+sϕ1*cψ1)
# R12 ≈ R[1,2]
#
# dxϕ1 = -cθ2*sψ2*(cϕ1*sθ1*cψ1-sϕ1*sψ1) + sθ2*( sϕ1*sθ1*cψ1+cϕ1*sψ1)
# dxθ1 = -cθ2*cψ2*(sθ1*cψ1) - cθ2*sψ2*(sϕ1*cθ1*cψ1) + sθ2*(-cϕ1*cθ1*cψ1)
# dxψ1 = -cθ2*cψ2*(cθ1*sψ1) - cθ2*sψ2*(-sϕ1*sθ1*sψ1+cϕ1*cψ1) + sθ2*( cϕ1*sθ1*sψ1+sϕ1*cψ1)
#
# dyϕ1 = -cθ2*sψ2*(-sϕ1*sθ1*sψ1+cϕ1*cψ1) + sθ2*(cϕ1*sθ1*sψ1+sϕ1*cψ1)
#
# dψdϕ1 = -R11/(R11^2 + R12^2)*dϕ1 + (R12)/(R11^2 + R12^2)*dϕ1
#
# vals1 = @SVector [e1.ϕ, e1.θ, e1.ψ]
# R2 = rotmat(e2)
# f(e) = rotmat_to_rpy(R2*rotmat(e...))
# f(vals1)
# ForwardDiff.jacobian(f, vals1)

# @btime ∇rotation1($e2, $e1)
# @btime ∇rotation2($e2, $e1)

############################################################################################
#                              MODIFIED RODRIGUES PARAMETERS
############################################################################################
p1 = MRP(q1)
p2 = MRP(q2)
# @btime $p2*$p1
# @btime $p2*$r

# Test rotations and composition
@test p2*r ≈ q2*r
@test p2*p1*r ≈ q2*q1*r
@test rotmat(p2)*r ≈ p2*r
p3 = p2*p1
@test p2\p3 ≈ p1
@test p3/p1 ≈ p2

@test r ≈ p3\(p2*p1*r)

r = @SVector rand(3)
R = rotmat(q1)
@test R*r ≈ q1*r

p = MRP(p1)
R1 = rotmat(p)
@test R1*r ≈ R*r
@test p*r ≈ R*r

@test UnitQuaternion(p) ≈ q1

# Test composition jacobians
@test ForwardDiff.jacobian(x->SVector(p2*MRP(x)),SVector(p1)) ≈ ∇composition1(p2,p1)
@test ForwardDiff.jacobian(x->SVector(MRP(x)*p1),SVector(p2)) ≈ ∇composition2(p2,p1)
p0 = MRP(0,0,0)
@test TO.∇differential(p2) ≈ ∇composition1(p2,p0)

############################################################################################
#                              RODRIGUES PARAMETERS
############################################################################################

g1 = RodriguesParam(q1)
g2 = RodriguesParam(q2)
@test q2 ≈ -UnitQuaternion(g2)
@test g2 ≈ RodriguesParam(UnitQuaternion(g2))

# Test compostion and rotation
@test g1*r ≈ q1*r
@test g2*r ≈ q2*r
@test g2*g1*r ≈ q2*q1*r
@test rotmat(g2)*r ≈ g2*r

g3 = g2*g1
@test g2\g3 ≈ g1
@test g3/g1 ≈ g2
@test r ≈ g3\(g2*g1*r)
@test r ≈ (g3\g2*g1)*r

# Test Jacobians
@test ForwardDiff.jacobian(g->RodriguesParam(g)*r, SVector(g1)) ≈ ∇rotate(g1, r)

function compose(g2,g1)
    N = (g2+g1 + g2 × g1)
    D = 1/(1-g2'g1)
    return D*N
end
@test ForwardDiff.jacobian(g->compose(SVector(g2),g), SVector(g1)) ≈ ∇composition1(g2,g1)
@test ForwardDiff.jacobian(g->compose(g,SVector(g1)), SVector(g2)) ≈ ∇composition2(g2,g1)

g0 = RodriguesParam{Float64}(0,0,0)
@test ∇composition1(g2, g0) ≈ TO.∇differential(g2)

# Conversions
import TrajectoryOptimization: rotmat_to_quat
Random.seed!(1) # i = 3
q = rand(UnitQuaternion)
@test rotmat_to_quat(rotmat(q)) ≈ q

Random.seed!(2) # i = 4
q = rand(UnitQuaternion)
@test rotmat_to_quat(rotmat(q)) ≈ q

Random.seed!(3) # i = 2
q = rand(UnitQuaternion)
@test rotmat_to_quat(rotmat(q)) ≈ q

Random.seed!(5) # i = 1
q = rand(UnitQuaternion)
@test rotmat_to_quat(rotmat(q)) ≈ q

#
# # Quadrotors
# quad1 = Dynamics.Quadrotor()
# quad2 = Dynamics.Quadrotor2{Quat{VectorPart}}()
# x,u = rand(quad2)
# dynamics(quad1,x,u) ≈ dynamics(quad2,x,u)
# # @btime dynamics($quad1,$x,$u)
# # @btime dynamics($quad2,$x,$u)
# #
# # @btime state_diff($quad1,$x,$(1.1x))
# # @btime state_diff($quad2,$x,$(1.1x))
# # @btime state_diff_jacobian($quad1, $x)
# # @btime state_diff_jacobian($quad2, $x)
#
# a = @SVector rand(3)
# b = @SVector rand(3)
# @test skew(a)*b == cross(a,b)
# skew(b)*a == -skew(a)*b
#
# p = a
# r = b
# f(x) = MRP_to_DCM(x)*r
# @test ForwardDiff.jacobian(f,p) ≈ MRP_rotate_jacobian(p,r)
#
# # Test MRP compositio
# p1 = @SVector rand(3)
# p2 = @SVector rand(3)
# r = @SVector rand(3)
# R1 = MRP_to_DCM(p1)
# R2 = MRP_to_DCM(p2)
# R3 = R2*R1
# p3 = MRP_composition(p2,p1)
# @test MRP_rotate_vec(p1,r) ≈ R1*r
# @test MRP_rotate_vec(p2,r) ≈ R2*r
# @test MRP_rotate_vec(p3,r) ≈ R3*r
#
# ForwardDiff.jacobian(x->MRP_composition(x,p1), p2) ≈ MRP_composition_jacobian_p2(p2,p1)
# ForwardDiff.jacobian(x->MRP_composition(p2,x), p1) ≈ MRP_composition_jacobian_p1(p2,p1)


using LinearAlgebra

println("k_local(Elastic Modulus,thickness,height,length,number of nodes,row one, row two, column one, column two)")

function k_local(E,t,h,l,n,r_1,r_2,c_1,c_2)
    k_global = zeros(n,n)
    A = t*h
    k = ((E*A)/l)
    k_global[r_1,c_1] = k
    k_global[r_1,c_2] = -k
    k_global[r_2,c_1] = -k
    k_global[r_2,c_2] = k
    return k_global
end


@doc raw"""
```
Photometric(illumination_direction1::Vector{T} where T <: Real = [Inf, Inf, Inf], illumination_direction2::Vector{T} where T <: Real = [Inf, Inf, Inf], illumination_direction3::Vector{T} where T <: Real = [Inf, Inf, Inf], integration_scheme::AbstractIntegrator = Horn())
```
Creates a `Photometric` functor which attempts to produce the gradients and
integrate them to retrieve a height map using three-point photometric stereo.

This algorithm solves a system of linear equation formed by varying the lighting
conditions on an object while holding the viewing angle constant. Note that
the ilumination direction id defined negative of most other algorithms and as
such inouts may need to be corrected (see example).
# Output
Returns an M by N array (matching dimensions of original image) of Float `Z`
that represents the reconstructed height at the point and the gradients in
M by N arrays of Float `p` and `q`.
# Details
Let ñ₁, ñ₂ and ñ₃ define the ilumination direction let N be a 3 X 3 matrix
whos rows are formed by the row vectors of the illumination direction as per:
```math
N=\begin{bmatrix}&n_{11}&n_{12}&n_{13}\\&n_{21}&n_{22}&n_{23}\\&n_{31}&n_{32}&n_{33}\end{bmatrix}
```
and I₁, I₂ and I₃ be the intensity values at a point (x,y) and let Ĩ=[I₁,I₂,I₃]′
be the column vector formed by these values.

The surface normal at the point (x,y) can now be directly calculated using:
```math
ñ=N^{-1}Ĩ
```
where n is the surface normal.
# Arguments
The function arguments are described in more detail below.
## `illumination_direction1`
A `Vector{T} where T <: Real` that specifies the first illumination vector to be used by the
algorithm. The `illumination_direction` should be a vector of the form [x,y,z]
where x,y,z are int he range [0,1]. Defults to `[Inf,Inf,Inf]` which triggers
`estimate_img_properties` to be called.
## `illumination_direction2`
A `Vector{T} where T <: Real` that specifies the second illumination vector to be used by the
algorithm. The `illumination_direction` should be a vector of the form [x,y,z]
where x,y,z are int he range [0,1]. Defults to `[Inf,Inf,Inf]` which triggers
`estimate_img_properties` to be called.
## `illumination_direction3`
A `Vector{T} where T <: Real` that specifies the third illumination vector to be used by the
algorithm. The `illumination_direction` should be a vector of the form [x,y,z]
where x,y,z are int he range [0,1]. Defults to `[Inf,Inf,Inf]` which triggers
`estimate_img_properties` to be called.
## `integration_scheme`
A keywork argument of type `AbstractIntegrator` which specifies which integration
scheme is to be used to reconstruct the surface from the calculated normals.
Defutls to Horn().
!!! note
    If any `illumination_direction` are not defined (i.e. have defulted to `Inf`) they
    will be calculated at runtime using `estimate_img_properties`.
```julia
using Images, Makie, ShapeFromShading

#generate synthetic images
n₁ = [0,0,1.0]
n₂ = [1.0,0.2,1.0]
n₃ = [0.2,0.9,1.0]
img1, img2, img3 = generate_photometric(SynthSphere(50), 1.0, n₁, n₂, n₃)

#calculate the heightmap
photometric = Photometric(-n₁, -n₂, -n₃, Horn())
Z,p,q = photometric(img1, img2, img3)

#normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

#display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:2
surface(r, r, Z)
```
# Reference
1. R. Woodham, "Photometric Method For Determining Surface Orientation From Multiple Images", Optical Engineering, vol. 19, no. 1, 1980. [doi:10.1117/12.7972479](https://doi.org/10.1117/12.7972479)
"""
function (algorithm::Photometric)(img1::AbstractArray, img2::AbstractArray, img3::AbstractArray)

    illumination_direction1 = algorithm.illumination_direction1
    illumination_direction2 = algorithm.illumination_direction2
    illumination_direction3 = algorithm.illumination_direction3
    integration_scheme = algorithm.integration_scheme

    if illumination_direction1 == [Inf, Inf, Inf] || illumination_direction2 == [Inf, Inf, Inf] || illumination_direction3 == [Inf, Inf, Inf]
        ρ, illumination_direction1, σ, τ = estimate_img_properties(img1)
        ρ, illumination_direction2, σ, τ = estimate_img_properties(img2)
        ρ, illumination_direction3, σ, τ = estimate_img_properties(img3)
    end

    #setup illumination matrix
    N = zeros(Float64, 3, 3)
    for i = 1:3
        N[1,i] = illumination_direction1[i]
        N[2,i] = illumination_direction2[i]
        N[3,i] = illumination_direction3[i]
    end
    Ninv = inv(N)

    #setup image and gradient arrays
    I₁ = Float64.(img1)
    I₂ = Float64.(img2)
    I₃ = Float64.(img3)
    p = zeros(Float64, size(I₁))
    q = zeros(Float64, size(I₁))

    #calculate gradients
    for i in CartesianIndices(I₁)
        I = [I₁[i], I₂[i], I₃[i]]
        n = Ninv * I
        if n[3] == 0.0
            n = [0.0,0.0,-1.0]
        end
        p[i] = -n[1] / n[3]
        q[i] = -n[2] / n[3]
    end

    #reconstruct surface
    Z = integration_scheme(p, q)
    return Z, p, q
end


@doc raw"""
# Lotka-Volterra model in 2D

```math
\begin{aligned}
L (q, \dot{q}) &= \bigg( q_2 + \frac{\log q_2}{q_1} \bigg) \, \dot{q_1} + q_1 \, \dot{q_2} - H(q) , \\
H(q) &= a_1 \, q_1 + a_2 \, q_2 + b_1 \, \log q_1 + b_2 \, \log q_2
\end{aligned}
```

"""
module LotkaVolterra2d

    export lotka_volterra_2d_dg_gauge

    ϑ₁(t, q) = q[2] + log(q[2]) / q[1]
    ϑ₂(t, q) = q[1]

    dϑ₁dx₁(t, q) = - log(q[2]) / q[1]^2
    dϑ₁dx₂(t, q) = 1 + 1 / (q[1] * q[2])

    dϑ₂dx₁(t, q) = one(eltype(q))
    dϑ₂dx₂(t, q) = zero(eltype(q))

    include("lotka_volterra_2d_common.jl")
    include("lotka_volterra_2d_equations.jl")

    function d²ϑ₁d₁d₁(t, q)
        + 2 * log(q[2]) / q[1]^3
    end

    function d²ϑ₁d₁d₂(t, q)
        - 1 / (q[1]^2 * q[2])
    end

    function d²ϑ₁d₂d₁(t, q)
        - 1 / (q[1]^2 * q[2])
    end

    function d²ϑ₁d₂d₂(t, q)
        - 1 / (q[1] * q[2]^2)
    end

    function d²ϑ₂d₁d₁(t, q)
        zero(eltype(q))
    end

    function d²ϑ₂d₁d₂(t, q)
        zero(eltype(q))
    end

    function d²ϑ₂d₂d₁(t, q)
        zero(eltype(q))
    end

    function d²ϑ₂d₂d₂(t, q)
        zero(eltype(q))
    end

    function g̅₁(t, q, v)
        d²ϑ₁d₁d₁(t,q) * q[1] * v[1] + d²ϑ₁d₂d₁(t,q) * q[1] * v[2] + d²ϑ₂d₁d₁(t,q) * q[2] * v[1] + d²ϑ₂d₂d₁(t,q) * q[2] * v[2]
    end

    function g̅₂(t, q, v)
        d²ϑ₁d₁d₂(t,q) * q[1] * v[1] + d²ϑ₁d₂d₂(t,q) * q[1] * v[2] + d²ϑ₂d₁d₂(t,q) * q[2] * v[1] + d²ϑ₂d₂d₂(t,q) * q[2] * v[2]
    end

    function lotka_volterra_2d_ϑ_κ(κ, t, q, v, Θ, params)
        Θ[1] = (1-κ) * ϑ₁(t,q) - κ * f₁(t,q,q)
        Θ[2] = (1-κ) * ϑ₂(t,q) - κ * f₂(t,q,q)
        nothing
    end

    function lotka_volterra_2d_f_κ(κ::Real, t::Real, q::Vector, v::Vector, f::Vector, params)
        f[1] = (1-κ) * f₁(t,q,v) - κ * (g₁(t,q,v) + g̅₁(t,q,v)) - dHd₁(t, q, params)
        f[2] = (1-κ) * f₂(t,q,v) - κ * (g₂(t,q,v) + g̅₂(t,q,v)) - dHd₂(t, q, params)
        nothing
    end

    function lotka_volterra_2d_g_κ(κ::Real, t::Real, q::Vector, v::Vector, g::Vector, params)
        g[1] = (1-κ) * f₁(t,q,v) - κ * (g₁(t,q,v) + g̅₁(t,q,v))
        g[2] = (1-κ) * f₂(t,q,v) - κ * (g₂(t,q,v) + g̅₂(t,q,v))
        nothing
    end

    # function lotka_volterra_2d_g(κ::Real, t::Real, q::Vector, v::Vector, g::Vector)
    #     g[1] = (1-κ) * g₁(t,q,v) - κ * g̅₁(t,q,v) - κ * f₁(t,q,v)
    #     g[2] = (1-κ) * g₂(t,q,v) - κ * g̅₂(t,q,v) - κ * f₂(t,q,v)
    #     nothing
    # end

    function lotka_volterra_2d_dg_gauge(q₀=q₀, p₀=ϑ(0, q₀), κ=0; params=parameters)
        lotka_volterra_2d_ϑ = (t, q, v, p, params) -> lotka_volterra_2d_ϑ_κ(κ, t, q, v, p, params)
        lotka_volterra_2d_f = (t, q, v, f, params) -> lotka_volterra_2d_f_κ(κ, t, q, v, f, params)
        lotka_volterra_2d_g = (t, q, λ, g, params) -> lotka_volterra_2d_g_κ(κ, t, q, λ, g, params)

        IODE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
             lotka_volterra_2d_g, q₀, p₀;
             parameters=params, v̄=lotka_volterra_2d_v)
    end

end


include("hw1/prob4.jl")

using .hw1_4

const μ = 3.986004e5 #km^3/s^2
r1 = 8000. #km
r2 = 10r1

# println(hw1_4.oe_to_xyz([8000, 0, 0, 0, 0, 0], μ))
hw1_4.plot_maneuver(r1, r2, π/2, μ, π/2)
# hw1_4.plot_maneuver(r1, r2, π/2, μ, 0.1)


#=
    Model Predictive Control (feedback) approach for a mechanical oscillator with two masses, springs and attenuators

    Inspired by: https://github.com/rdeits/DynamicWalking2018.jl

    Author: Stephan Scholz
    Year: 2021
=#

using JuMP, Ipopt

my_optimizer = with_optimizer(Ipopt.Optimizer, print_level = 0)

const d₁ = 0.5; # Damping
const d₂ = 0.5;
const c₁ = 0.1; # Spring constant
const c₂ = 0.1;
const m₁ = 1.0; # Mass
const m₂ = 2.0;

# System matrix
const A = [0 0 1 0; 0 0 0 1; (-1/m₁)*(c₁ + c₂) c₂/m₁ (-1/m₁)*(d₁ + d₂) d₂/m₁; c₂/m₂ -c₂/m₂ d₂/m₂ -d₂/m₂]

const B = [0 0; 0 0; 1/m₁ 0; 0 1/m₂]  # Input matrix
const C = [1 0 0 0; 0 1 0 0]          # Output matrix

mech_osc(x,u) = A * x + B * u   # System dynamics

function mpc_emo(x_init)

    model = Model(my_optimizer)

    ΔT = 0.1            #  Sampling time
    n_horizon = 20;     # MPC horizon

    Nx = size(A)[1]
    Nu = size(B)[2]

    # Input constraints
    umin = -100.0
    umax = 100.0

    ref = [2.0, 4.0, 0.0, 0.0] # Reference

    @variables model begin
        x[1:Nx, 1:n_horizon]    # States
        e[1:Nx, 1:n_horizon]    # Errors = Reference - States
        umin <= u[1:Nu, 1:n_horizon] <= umax  # Input constraints
    end

    @constraint(model, x[1:Nx, 1] .== x_init )                          # Initial values
    @constraint(model, [k=1:n_horizon], e[:, k] .== ref[:] - x[:, k])   # Error = Reference - States
    @constraint(model, [k=1:n_horizon-1], x[:, k+1] .== x[:, k] + ΔT * mech_osc(x[:,k], u[:,k])) # System dynamics

    J = @NLexpression(model, 100*sum(e[i,end]^2 for i in 1:Nx) + sum(u[j,end]^2 for j in 1:Nu) )

    @NLobjective( model, Min, J)

    optimize!(model)
    return JuMP.value.(x), JuMP.value.(u), JuMP.value.(e)
end

xpos = rand(4)          # Recent states
x_hist = zeros(4,0)     # History of states

for i in 1:50

    # Run the MPC control optimization
    x_plan, u_plan, e_plan = mpc_emo(xpos)

    # Save states
    x_hist = hcat(x_hist, x_plan[:,1])

    # Apply the planned input signals and simulate one step in time
    ΔT = 0.1
    xpos = xpos + ΔT * mech_osc(xpos, u_plan[:,1])
end

using Plots
plot(x_hist', labels=["x1" "x2" "x3" "x4"], title="Model Predictive Control of a mechanical oscillator")


## Run Metropolis-Hastings using a Bayesian probit model with a Normal prior
## N(0, priorVar*I) on the vasoconstriction data

using Distributions, GeometricMCMC

include("vaso.jl")
include("probitNormalPrior.jl")

# Create Model instance
model =
  Model(nPars, data, logPrior, logLikelihood, gradLogPosterior, randPrior);

# Create instance mhOpts of Metropolis-Hastings options
mhOpts = MhOpts(55000, 5000, 0.1);

# Run Metropolis-Hastings simulation
mcmc, z = mh(model, mhOpts);

# Compute ZV-MH mean estimators based on linear polynomial
linearZvMcmc, linearCoef = linearZv(mcmc, z);

# Compute ZV-MH mean estimators based on quadratic polynomial
quadraticZvMcmc, quadraticCoef = quadraticZv(mcmc, z);


using Plots, Einsum, StaticArrays

# Simulation parameters
const Nx, Ny      = 401, 401    # grid points in x and y
const rho0        = 1000    # average density
const Nt          = 600   # number of timesteps
const NL          = 9       # D2Q9 Lattice
const tau         = 0.6    # collision timescale
const omega       = 1/tau # appears in the collision equation, precomputing saves a division.

const cxs         = @SVector[0, 0, 1, 1, 1, 0,-1,-1,-1]
const cys         = @SVector[0, 1, 1, 0,-1,-1,-1, 0, 1]
const opp         = @SVector[1,6,7,8,9,2,3,4,5] # bounce back array, opposite direction. Not used but will be when boundaries are added.
const weights     = @SVector[4/9,1/9,1/36,1/9,1/36,1/9,1/36,1/9,1/36] # sums to 1

##build the struct that will hold the state of the LB system##
struct LatticeState
    rho::Array{Float32, 2}    # macroscale density
    ux::Array{Float32, 2}    # macroscale velocity, x component
    uy::Array{Float32, 2}    # macroscale velocity, y component
    F::Array{Float32, 3}    # particle distribution function
    Feq::Array{Float32, 3}    # equillibrium term for collision
    tmp::Array{Float32, 2}   # place-holder array
end

function LatticeState(Nx,Ny)
    F = ones(Nx,Ny,NL)
    Feq = similar(F)
    ux = zeros(Nx,Ny)
    uy = similar(ux)
    rho = similar(ux)
    tmp = similar(ux)
    sum!(rho, F)

    return LatticeState(rho,ux,uy,F,Feq,tmp)
end

@fastmath @inbounds function apply_drift!(tmp::Array{<:Real},F::Array{<:Real})
    for i in 1:size(F)[3]
        A = @view F[:,:,i]
        circshift!(tmp,A, (cxs[i],0) )
        circshift!(A,tmp, (0,cys[i]) )
    end
end

@fastmath @inbounds function feq_point(rho::Real,ux::Real,uy::Real,cx::Int,cy::Int,weight::Real)
    cu = (cx * ux +  cy * uy)
    return rho*weight*(1. +  3. *cu + 4.5*cu*cu - 1.5*(ux*ux + uy*uy))
end

@fastmath @inbounds function calculate_feq!(Feq::Array{<:Real},rho::Array{<:Real},ux::Array{<:Real},uy::Array{<:Real})
    @einsum Feq[i,j,k] = feq_point(rho[i,j],ux[i,j],uy[i,j],cxs[k],cys[k],weights[k])
end

@fastmath @inbounds function apply_collision!(F::Array{<:Real},Feq::Array{<:Real})
    @. F -= omega .* (F .- Feq)
end

@fastmath @inbounds function calculate_u!(ux::Array{<:Real},uy::Array{<:Real},F::Array{<:Real},rho::Array{<:Real})
    @einsum ux[i,j] = F[i,j,k] * cxs[k] / rho[i,j]
    @einsum uy[i,j] = F[i,j,k] * cys[k] / rho[i,j]
end

@fastmath @inbounds function correct_F!(F::Array{<:Real},rho::Array{<:Real})
    @einsum F[i,j,k] = F[i,j,k] * rho0 / rho[i,j]
end

function iterate_lb!(f::LatticeState, Nt::Int)
    for it in 1:Nt
        apply_drift!(f.tmp,f.F) # Particle streaming step
        sum!(f.rho, f.F) # Calculate fluid density
        calculate_u!(f.ux, f.uy, f.F, f.rho) # Calculate macroscopic velocity
        calculate_feq!(f.Feq, f.rho, f.ux,f.uy) # Calculate Feq then apply the collision step
        apply_collision!(f.F, f.Feq)
    end
end

@fastmath @inbounds function set_initial_condition!(f::LatticeState)
    @. f.rho += [exp(-sqrt((x-Nx/2)^2 + (y-Ny/2)^2)) for x in 1:Nx, y in 1:Ny]; # Initial condition. Modify density with a pulse in the middle
    correct_F!(f.F,f.rho) #correct F to the new density
end

problem = LatticeState(Nx,Ny)
set_initial_condition!(problem)
iterate_lb!(problem,Nt) # Run simulation
Plots.heatmap(problem.rho, c=:redsblues, size=(650,640), clims=(rho0-.2, rho0+0.2), right_margin = 5Plots.mm,aspect_ratio=:equal) # Plot the final density variation
title!("Final density")
savefig("output.png")


using Surrogates
using Stheno
using Zygote
using Test

@testset "1D - 1D" begin

    lb = 0.0
    ub = 10.0
    f = x -> log(x)*exp(x)
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    my_k = SthenoKriging(x,y)
    x_new = 4.0
    y_new = f.(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [5.0, 6.0]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        Zygote.gradient(x -> my_k(x), 1.0)
    end

    @testset "Optimization" begin
        objective_function = x -> 2*x+1
        x = [2.0,4.0,6.0]
        y = [5.0,9.0,13.0]
        p = 2
        a = 2
        b = 6
        my_k_EI1 = SthenoKriging(x,y)
        surrogate_optimize(objective_function,EI(),a,b,my_k_EI1,UniformSample(),maxiters=200,num_new_samples=155)
    end
end

@testset "ND - 1D" begin

    lb = [0.0; 0.0]
    ub = [5.0; 10.0]
    f = x -> log(x[1])*exp(x[2])
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    my_k = SthenoKriging(x,y)
    x_new = (2.0, 1.0)
    y_new = f(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [(1.0, 2.0), (0.5, 2.5)]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        g1 = Zygote.gradient(x -> my_k(x), (1.0, 1.0))[1]
        g2 = Zygote.gradient(x -> my_k(x), [1.0, 1.0])[1]
        @test [g1...] ≈ g2
    end
    @testset "Optimization" begin
        objective_function_ND = z -> 3*hypot(z...)+1
        x = [(1.2,3.0),(3.0,3.5),(5.2,5.7)]
        y = objective_function_ND.(x)
        theta = [2.0,2.0]
        lb = [1.0,1.0]
        ub = [6.0,6.0]

        my_k_E1N = SthenoKriging(x,y)
        surrogate_optimize(objective_function_ND,EI(),lb,ub,my_k_E1N,UniformSample())
    end
end

@testset "1D - ND" begin

    lb = 0.0
    ub = 10.0
    f = x -> [x, x^2]
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)

    my_k = SthenoKriging(x,y,Stheno.GP(Stheno.EQ(), Stheno.GPC()))
    x_new = 4.0
    y_new = f.(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [5.0, 6.0]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    # separate Stheno models
    gpc = Stheno.GPC()
    gp1 = Stheno.GP(Stheno.EQ(), gpc)
    gp2 = Stheno.GP(Stheno.Exp(), gpc)
    my_k = SthenoKriging(x, y, (gp1, gp2))
    x_new = 0.2
    y_new = f.(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [5.5, 6.5]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        Zygote.gradient(x -> sum(my_k(x)), 1.0)
    end
end

@testset "ND - ND" begin

    lb = [0.0; 0.0]
    ub = [5.0; 10.0]
    f = x -> [x[1]*x[2], x[2]^2]
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)

    my_k = SthenoKriging(x,y,Stheno.GP(Stheno.EQ(), Stheno.GPC()))
    x_new = (1.0, 1.0)
    y_new = f(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [(1.0, 2.0), (0.5, 2.5)]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    # separate Stheno models
    gpc = Stheno.GPC()
    gp1 = Stheno.GP(Stheno.EQ(), gpc)
    gp2 = Stheno.GP(Stheno.Exp(), gpc)
    my_k = SthenoKriging(x, y, (gp1, gp2))
    x_new = (0.2, 0.1)
    y_new = f(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [(5.5, 0.5), (6.5, 0.5)]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        g1 = Zygote.gradient(x -> sum(my_k(x)), (1.1, 1.1))[1]
        g2 = Zygote.gradient(x -> sum(my_k(x)), [1.1, 1.1])[1]
        @test [g1...] ≈ g2
    end
end


using FirstOrderOptimization

## The state of this code:
# * seems like the algo works.
# 	* obj converges to a nice low number
# 	* i haven't checked it against an sdp solver yet
# 	* or tested for recovery
# * it does a super weird thing where
# 	the rank and obj shoot up in the 14th iteration. why?

srand(1)

type PrevRank
	r::Integer
end

### max norm regularized matrix completion problem

# minimize f(W) + g(W) + h(W)

# where W = [A X; X' B] and

# (data fitting)   f(W) = vecnorm(P.*(X - M))^2
# (regularization) g(W) = lambda*max(diag(X))
# (psd)            h(W) = (W \succeq 0) ? 0 : Inf

function generate_maxnorm_problem(m,n,lambda,k)
	# generate data
	X0 = randn(m,k)
	Y0 = randn(k,n)
	A = X0*Y0 + .1*randn(m,n)
	P = float(randn(m,n) .>= .5) # observed entries

	X(W) = W[1:m, m+1:end]

	function grad_f(W)
		gX = P.*(X(W)-A)
		return [zeros(m,m) gX; gX' zeros(n,n)]
	end

	function prox_g(W, alpha)
		Z = copy(W)
		oldmax = maximum(diag(W))
		newmax = oldmax - lambda*alpha/2
		for i=1:size(Z,1)
			if Z[i,i] > newmax
				Z[i,i] = newmax
			end
		end
		Z
	end

	# we're going to use a closure over prevrank
	# to remember what the rank of prox_h(W) was the last time we computed it
	# in order to avoid calculating too many eigentuples of W
	prevrank = PrevRank(k)

	function prox_h(W, alpha=0; TOL=1e-10)
		# debugging: @show prevrank.r
		while prevrank.r < size(W,1)
			l,v = eigs(Symmetric(W), nev = prevrank.r+1, which=:LR) # v0 = [v zeros(size(W,1), prevrank.r+1 - size(v,2))]
			if l[end] <= TOL
				prevrank.r = sum(l.>=TOL)
				return v*diagm(max(l,0))*v'
			else
				prevrank.r = 2*prevrank.r # double the rank and try again
			end
		end
		# else give up on computational cleverness
		l,v = eig(Symmetric(W))
		prevrank.r = sum(l.>=TOL)
		return v*diagm(max(l,0))*v'
	end

	# we're not going to bother checking whether W is psd or not
	# when evaluating the objective; in the course of the prisma
	# algo this makes no difference
	obj(W) = sum((P.*(X(W) - A)).^2) + lambda*maximum(diag(W))

	return grad_f, prox_g, prox_h, obj
end

function test_prisma(m,n,lambda,k)

	grad_f, prox_g, prox_h, obj = generate_maxnorm_problem(m,n,lambda,k)

	# initialize
	W = zeros(m+n,m+n)
	# lipshitz constant for f
	L_f = 2
	# orabona starts stepsize at
	# beta = lambda/sqrt((m+n)^2*mean(A.^2))
	beta   = lambda/sqrt(obj(W))
	ssr    = PrismaStepsize(beta)
	params = PrismaParams(ssr, 10, 1)

	# recover
	W = PRISMA(W, L_f,
		   grad_f,
		   prox_g,
		   prox_h,
		   obj,
		   params)
end

test_prisma(5,5,10,3)


@testset "Atomic measures" begin
    Mod.@polyvar x y
    δ1 = WeightedDiracMeasure([1, 0],     2.0)
    δ2 = WeightedDiracMeasure([1/2, 1/2], 3.0)
    η = AtomicMeasure([x, y], [δ1, δ2])
    @test η ≈ AtomicMeasure([x, y], [δ2, δ1])
    @test !(η ≈ AtomicMeasure([x, y], [WeightedDiracMeasure([1, 0],     3.0),
                                       WeightedDiracMeasure([1/2, 1/2], 2.0)]))
    @test string(η) == "Atomic measure on the variables x, y with 2 atoms:\n at [1.0, 0.0] with weight 2.0\n at [0.5, 0.5] with weight 3.0"
    p = x^2 + x + y + x*y^2 + 1
    @test dot(η, p) == 2 * p(x=>1, y=>0) + 3 * p(x => 1/2, y => 1/2)
    @test dot(η, p) == dot(p, η)
    μ = measure(η, monomials(p))
    @test dot(μ, p) == 2 * p(x=>1, y=>0) + 3 * p(x => 1/2, y => 1/2)
end


__precompile__()

module Cosmology

using QuadGK, Unitful
import Unitful: km, s
using UnitfulAstro: Mpc, Gpc, Gyr

export cosmology,
       age,
       angular_diameter_dist,
       comoving_radial_dist,
       comoving_transverse_dist,
       comoving_volume,
       comoving_volume_element,
       distmod,
       H,
       hubble_dist,
       hubble_time,
       luminosity_dist,
       lookback_time,
       scale_factor

abstract type AbstractCosmology end
abstract type AbstractClosedCosmology <: AbstractCosmology end
abstract type AbstractFlatCosmology <: AbstractCosmology end
abstract type AbstractOpenCosmology <: AbstractCosmology end

struct FlatLCDM{T<:Real} <: AbstractFlatCosmology
    h::T
    Ω_Λ::T
    Ω_m::T
    Ω_r::T
end
FlatLCDM(h::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
    FlatLCDM(promote(float(h), float(Ω_Λ), float(Ω_m), float(Ω_r))...)

a2E(c::FlatLCDM, a) = sqrt(c.Ω_r + c.Ω_m*a + c.Ω_Λ*a^4)

struct ClosedLCDM{T<:Real} <: AbstractClosedCosmology
    h::T
    Ω_k::T
    Ω_Λ::T
    Ω_m::T
    Ω_r::T
end
ClosedLCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
    ClosedLCDM(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
                       float(Ω_r))...)

struct OpenLCDM{T<:Real} <: AbstractOpenCosmology
    h::T
    Ω_k::T
    Ω_Λ::T
    Ω_m::T
    Ω_r::T
end
OpenLCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
    OpenLCDM(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
                     float(Ω_r))...)

function a2E(c::Union{ClosedLCDM,OpenLCDM}, a)
    a2 = a*a
    sqrt(c.Ω_r + c.Ω_m*a + (c.Ω_k + c.Ω_Λ*a2)*a2)
end

for c in ("Flat", "Open", "Closed")
    name = Symbol("$(c)WCDM")
    @eval begin
        struct $(name){T<:Real} <: $(Symbol("Abstract$(c)Cosmology"))
            h::T
            Ω_k::T
            Ω_Λ::T
            Ω_m::T
            Ω_r::T
            w0::T
            wa::T
        end
        function $(name)(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real,
                         w0::Real, wa::Real)
            $(name)(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
                            float(Ω_r), float(w0), float(wa))...)
        end
    end
end

function WCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real, w0::Real, wa::Real)
    if Ω_k < 0
        ClosedWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
    elseif Ω_k > 0
        OpenWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
    else
        FlatWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
    end
end

function a2E(c::Union{FlatWCDM,ClosedWCDM,OpenWCDM}, a)
    ade = exp((1 - 3*(c.w0 + c.wa))*log(a) + 3*c.wa*(a - 1))
    sqrt(c.Ω_r + (c.Ω_m + c.Ω_k*a)*a + c.Ω_Λ*ade)
end

function cosmology(;h=0.69,
                   Neff=3.04,
                   OmegaK=0,
                   OmegaM=0.29,
                   OmegaR=nothing,
                   Tcmb=2.7255,
                   w0=-1,
                   wa=0)

    if OmegaR === nothing
        OmegaG = 4.48131e-7*Tcmb^4/h^2
        OmegaN = Neff*OmegaG*(7/8)*(4/11)^(4/3)
        OmegaR = OmegaG + OmegaN
    end

    OmegaL = 1 - OmegaK - OmegaM - OmegaR

    if !(w0 == -1 && wa == 0)
        return WCDM(h, OmegaK, OmegaL, OmegaM, OmegaR, w0, wa)
    end

    if OmegaK < 0
        return ClosedLCDM(h, OmegaK, OmegaL, OmegaM, OmegaR)
    elseif OmegaK > 0
        return OpenLCDM(h, OmegaK, OmegaL, OmegaM, OmegaR)
    else
        return FlatLCDM(h, OmegaL, OmegaM, OmegaR)
    end
end

# hubble rate

scale_factor(z) = 1/(1 + z)
E(c::AbstractCosmology, z) = (a = scale_factor(z); a2E(c,a)/a^2)
H(c::AbstractCosmology, z) = 100 * c.h * E(c, z) * km / s / Mpc

hubble_dist0(c::AbstractCosmology) = 2997.92458/c.h * Mpc
hubble_dist(c::AbstractCosmology, z) = hubble_dist0(c)/E(c,z)

hubble_time0(c::AbstractCosmology) = 9.77814/c.h * Gyr
hubble_time(c::AbstractCosmology, z) = hubble_time0(c)/E(c,z)

# distances

Z(c::AbstractCosmology, z::Real; kws...) =
    QuadGK.quadgk(a->1/a2E(c,a), scale_factor(z), 1; kws...)[1]
Z(c::AbstractCosmology, z₁::Real, z₂::Real; kws...) =
    QuadGK.quadgk(a->1/a2E(c,a), scale_factor(z₂), scale_factor(z₁); kws...)[1]

comoving_radial_dist(c::AbstractCosmology, z; kws...) = hubble_dist0(c)*Z(c, z; kws...)
comoving_radial_dist(c::AbstractCosmology, z₁, z₂; kws...) = hubble_dist0(c)*Z(c, z₁, z₂; kws...)

comoving_transverse_dist(c::AbstractFlatCosmology, z; kws...) =
    comoving_radial_dist(c, z; kws...)
comoving_transverse_dist(c::AbstractFlatCosmology, z₁, z₂; kws...) =
    comoving_radial_dist(c, z₁, z₂; kws...)
function comoving_transverse_dist(c::AbstractOpenCosmology, z; kws...)
    sqrtok = sqrt(c.Ω_k)
    hubble_dist0(c)*sinh(sqrtok*Z(c, z; kws...))/sqrtok
end
function comoving_transverse_dist(c::AbstractOpenCosmology, z₁, z₂; kws...)
    sqrtok = sqrt(c.Ω_k)
    hubble_dist0(c)*sinh(sqrtok*Z(c, z₁, z₂; kws...))/sqrtok
end
function comoving_transverse_dist(c::AbstractClosedCosmology, z; kws...)
    sqrtok = sqrt(abs(c.Ω_k))
    hubble_dist0(c)*sin(sqrtok*Z(c,z; kws...))/sqrtok
end
function comoving_transverse_dist(c::AbstractClosedCosmology, z₁, z₂; kws...)
    sqrtok = sqrt(abs(c.Ω_k))
    hubble_dist0(c)*sin(sqrtok*Z(c, z₁, z₂; kws...))/sqrtok
end

angular_diameter_dist(c::AbstractCosmology, z; kws...) =
    comoving_transverse_dist(c, z; kws...)/(1 + z)
angular_diameter_dist(c::AbstractCosmology, z₁, z₂; kws...) =
    comoving_transverse_dist(c, z₁, z₂; kws...)/(1 + z₂)

luminosity_dist(c::AbstractCosmology, z; kws...) =
    comoving_transverse_dist(c, z; kws...)*(1 + z)

distmod(c::AbstractCosmology, z; kws...) =
    5 * log10(luminosity_dist(c, z; kws...) / Mpc) + 25

# volumes

comoving_volume(c::AbstractFlatCosmology, z; kws...) =
    (4pi/3)*(comoving_radial_dist(Gpc, c, z; kws...))^3
function comoving_volume(c::AbstractOpenCosmology, z; kws...)
    DH = hubble_dist0(Gpc, c)
    x = comoving_transverse_dist(Gpc, c, z; kws...)/DH
    sqrtok = sqrt(c.Ω_k)
    2pi*(DH)^3*(x*sqrt(1 + c.Ω_k*x^2) - asinh(sqrtok*x)/sqrtok)/c.Ω_k
end
function comoving_volume(c::AbstractClosedCosmology, z; kws...)
    DH = hubble_dist0(Gpc, c)
    x = comoving_transverse_dist(Gpc, c,z; kws...)/DH
    sqrtok = sqrt(abs(c.Ω_k))
    2pi*(DH)^3*(x*sqrt(1 + c.Ω_k*x^2) - asin(sqrtok*x)/sqrtok)/c.Ω_k
end

comoving_volume_element(c::AbstractCosmology, z; kws...) =
    hubble_dist0(Gpc, c)*angular_diameter_dist(Gpc, c, z; kws...)^2/a2E(c, scale_factor(z))

# times

T(c::AbstractCosmology, a0, a1; kws...) = QuadGK.quadgk(x->x/a2E(c,x), a0, a1; kws...)[1]
age(c::AbstractCosmology, z; kws...) = hubble_time0(c)*T(c, 0, scale_factor(z); kws...)
lookback_time(c::AbstractCosmology, z; kws...) = hubble_time0(c)*T(c, scale_factor(z), 1; kws...)

# Easily select a different unit
for f in (:hubble_dist0, :hubble_dist, :hubble_time0, :hubble_time, :comoving_radial_dist,
          :comoving_transverse_dist, :angular_diameter_dist, :luminosity_dist,
          :comoving_volume, :comoving_volume_element, :age, :lookback_time)
    @eval $f(u::Unitful.Unitlike, args...; kws...) = uconvert(u, $f(args...; kws...))
end

###############
# Deprecations
#

@deprecate hubble_dist_mpc0(c::AbstractCosmology) ustrip(hubble_dist0(c::AbstractCosmology))
@deprecate hubble_dist_mpc(c::AbstractCosmology, z) ustrip(hubble_dist(c::AbstractCosmology, z))
@deprecate hubble_time_gyr0(c::AbstractCosmology) ustrip(hubble_time0(c::AbstractCosmology))
@deprecate hubble_time_gyr(c::AbstractCosmology, z) ustrip(hubble_time(c::AbstractCosmology, z))
@deprecate comoving_radial_dist_mpc(c::AbstractCosmology, z; kws...) ustrip(comoving_radial_dist(c::AbstractCosmology, z; kws...))
@deprecate comoving_transverse_dist_mpc(c::AbstractCosmology, z; kws...) ustrip(comoving_transverse_dist(c::AbstractCosmology, z; kws...))
@deprecate angular_diameter_dist_mpc(c::AbstractCosmology, z; kws...) ustrip(angular_diameter_dist(c::AbstractCosmology, z; kws...))
@deprecate luminosity_dist_mpc(c::AbstractCosmology, z; kws...) ustrip(luminosity_dist(c::AbstractCosmology, z; kws...))
@deprecate comoving_volume_gpc3(c::AbstractCosmology, z; kws...) ustrip(comoving_volume(c::AbstractCosmology, z; kws...))
@deprecate comoving_volume_element_gpc3(c::AbstractCosmology, z; kws...) ustrip(comoving_volume_element(c::AbstractCosmology, z; kws...))
@deprecate age_gyr(c::AbstractCosmology, z; kws...) ustrip(age(c::AbstractCosmology, z; kws...))
@deprecate lookback_time_gyr(c::AbstractCosmology, z; kws...) ustrip(lookback_time(c::AbstractCosmology, z; kws...))

end # module


using ControlSystems

abstract type AbstractController end

struct TVLQRController <: AbstractController
    nom_x
    nom_u
    K
    xf
    uf
    L
    N
    Δt
    LQR_eps
end

"""
Solves the discrete-time algebraic Riccati equation (DARE)

$(SIGNATURES)
"""
function TVDARE(model, nom_x, nom_u, Q, R, Qf, dt)
    P = Qf
    n = model.n
    m = model.m
    N = length(nom_x)
    K = [zeros(m,n) for i = 1:N-1]

    model_d = discretize_model(model, :rk4)

    J = TrajectoryOptimization.PartedMatrix(model_d)

    for i = 1:length(nom_u)
        t = N - i
        xt = nom_x[t]
        ut = nom_u[t]

        jacobian!(J, model_d, xt, ut, dt)

        A = J.xx
        B = J.xu

        APA = A' * P * A
        APB = A' * P * B
        BPA = B' * P * A
        BPB = B' * P * B

        Kt = -(R + BPB)\BPA
        K[t] .= Kt

        P = Q + Kt' * R * Kt + (A + B * Kt)' * P * (A + B * Kt)
    end

    return K
end

struct ALTRO
    x
    u
    K
end

"""
Performs direct-collocation using Augmented Lagrangian method

$(SIGNATURES)

**Arguments**:

- `model :: TrajectoryOptimization.Model`
- `x0 :: Vector{Float64}`
- `xf :: Vector{Float64}`
- `Q :: Matrix{Float64}`
- `R :: Matrix{Float64}`
- `Qf :: Matrix{Float64}`
- `N :: Int64`
- `dt :: Float64`

**Returns**:

`struct` containing the nominal tracking trajectory and the gains
"""
function ALTRO(model, x0, xf, Q, R, Qf, N, dt;
               ilqr_opts=(iterations=200,
                          iterations_linesearch=10,
                          verbose=true, cost_tolerance=1.0e-6),
               al_opts=(verbose=false,
                        iterations=30,
                        penalty_scaling=10.0,
                        cost_tolerance=1.0e-6,
                        cost_tolerance_intermediate=1.0e-5,
                        constraint_tolerance=1.0e-4),
               )
    model_d = discretize_model(model, :rk4)
    T = eltype(x0)
    n = model_d.n;
    m = model_d.m;

    goal = goal_constraint(xf);

    #N = 251;

    U = [ones(m) for k = 1:N-1];

    #dt = 0.01;

    obj = LQRObjective(Q,R,Qf,xf,N);
    con = Constraints(N);
    con[N] += goal;

    prob = Problem(model_d,obj,constraints=con,x0=x0,N=N,dt=dt,xf=xf);
    initial_controls!(prob, U);

    opts_ilqr = iLQRSolverOptions(;ilqr_opts...);

    opts_al = AugmentedLagrangianSolverOptions{T}(;opts_uncon=opts_ilqr, al_opts...);

    opts_altro = ALTROSolverOptions{T}(verbose=false,
                                       opts_al=opts_al);

    prob, solver = TrajectoryOptimization.solve(prob, opts_altro); # solve with ALTRO
    prob.X[end] = xf

    return ALTRO(prob.X, prob.U, solver.solver_al.solver_uncon.K)
end

"""
Generates the nominal trajectory to track and a TVLQR follwer for the same trajectory

$(SIGNATURES)
"""
function TVLQRController(model, x0, xf, Q, R, Qf, N, dt;
                         TVLQR_Q = nothing,
                         TVLQR_R = nothing,
                         TVLQR_Qf = nothing,
                         LQR_eps = 1e-1,
                         ilqr_opts=(iterations=200,
                                    iterations_linesearch=10,
                                    verbose=true, cost_tolerance=1.0e-6),
                         al_opts=(verbose=false,
                                  iterations=30,
                                  penalty_scaling=10.0,
                                  cost_tolerance=1.0e-6,
                                  cost_tolerance_intermediate=1.0e-5,
                                  constraint_tolerance=1.0e-4),
                         )

    if isnothing(TVLQR_Q)
        TVLQR_Q = Q
    end
    if isnothing(TVLQR_R)
        TVLQR_R = R
    end
    if isnothing(TVLQR_Qf)
        TVLQR_Qf = Qf
    end

    nom_traj = ALTRO(model, x0, xf, Q, R, Qf, N, dt, ilqr_opts=ilqr_opts, al_opts=al_opts)
    return TVLQRController(model, x0, xf, nom_traj, TVLQR_Q, TVLQR_R, TVLQR_Qf, N, dt, LQR_eps)
end

"""
Generates a TVLQR trajectory follower given the nominal trajectory `nom_traj`

$(SIGNATURES)

`Q`, `R`, `Qf` refer to the TVLQR costs
"""
function TVLQRController(model, x0, xf, nom_traj, Q, R, Qf, N, dt, LQR_eps)
    model_d = discretize_model(model, :rk4)

    X = nom_traj.x
    U = nom_traj.u
    K = TVDARE(model, nom_traj.x, nom_traj.u, Q, R, Qf, dt)
    # K = nom_traj.K

    J = TrajectoryOptimization.PartedMatrix(model_d)
    uf = zeros(model.m)

    jacobian!(J, model_d, xf, uf, dt)

    A = J.xx
    B = J.xu
    L = dlqr(A, B, Q, R)
    return TVLQRController(X, U, K, xf, uf, -L, N, dt, LQR_eps)
end

function (ctrl::TVLQRController)(t, x, thetamask)
    i = convert(Int, ceil(t/ctrl.Δt))
    i_ = min(i, ctrl.N-1)
    # i = min(t, ctrl.N-1)

    dx = x - ctrl.xf
    @views dx .= wraptopi(dx, thetamask)
    #@show norm(dx)

    if (i_ < i || norm(dx) < ctrl.LQR_eps)
        x_ = ctrl.xf
        u_ = ctrl.uf
        K = ctrl.L
        #@show dx
    else
        x_ = ctrl.nom_x[i_]
        u_ = ctrl.nom_u[i_]
        K  = ctrl.K[i_]
    end

    dx = x - x_
    @views dx .= wraptopi(dx, thetamask)

    du = K * dx
    u = u_ + du

    return u
end

struct RandomController <: AbstractController
end

"""
Generates a controller with random actions

$(SIGNATURES)
"""
function RandomController(model, std)
    return (t, x, thetamask) -> std * randn(model.m)
end

function addnoise(ctrl::AbstractController, noisestd)
    function noisy(t, x, thetamask)
        u = ctrl(t, x, thetamask)
        return u .+ randn(size(u)...) .* noisestd
    end
    return noisy
end

"""
Clip the output of a controller to a maximum force value `maxu`.

$(SIGNATURES)
"""
function clip(ctrl::AbstractController, maxu)
    function cu(t, x, thetamask)
        u = ctrl(t, x, thetamask)
        return clamp.(u, -maxu, maxu)
    end
    return cu
end

function wraptopi(radians)
     return radians .- 2*pi*floor.((radians .+ pi)/(2*pi))
end

function wraptopi(radians, mask)
    oldrad = copy(radians)
    return mask .* wraptopi(radians) .+ (1.0 .- mask) .* oldrad
end


using LinearAlgebra
using SparseArrays

# using Flux

⊗ = kron

# LITTLE ENDIAN !
# 1st qubit = far right

# if IittleEndian then the last qubit is the control
function Controlled(gate::AbstractMatrix)
  sz = size(gate)
  n = sz[1]
  ret = spzeros(ComplexF64, 2*n, 2*n)
  ret[1:n,1:n] += I
  ret[n+1:2*n,n+1:2*n] = gate
  return ret
end

# returns sparse permutation matrix P of width 2^pad
# input indices of qubits are 1-based
# if G acts on qubits 1,2,... then P^(-1)*G*P acts on ind[1], ind[2],...
# LittleEndian!
function Permute(ind::Array{Int}; pad::Int=0)
  # assert maximum(ind) <= pad
  # assert ind[i] all different

  k = length(ind)
  pad = max(pad, maximum(ind))

  # reconstruct full permutation sigma
  sigma = append!(copy(ind), zeros(Int64, pad-k))
  # rest = setdiff(1:pad, ind)
  rest = filter(x -> !(x in ind), 1:pad)
  sigma[k+1:pad] = rest
  # todo: more stable sigma?

  # println(sigma)

  rows = zeros(Int64, 2^pad)
  cols = zeros(Int64, 2^pad)
  # vals = zeros(ComplexF64, 2^pad)
  vals = zeros(Int64, 2^pad)

  function permute_binary(b::Int64)
    ret = 0

    b -= 1

    for i=1:pad
      digit = b & 1
      ret += digit << (sigma[i]-1)
      b = b >> 1
    end

    return ret+1
  end

  for b in 1:2^pad
    rows[b] = b
    cols[b] = permute_binary(b)
    vals[b] = 1
  end

  # println("cols", cols)

  return sparse(rows, cols, vals)
end

function ApplyGateOn(gate::AbstractMatrix, ind::Array{Int}; pad::Int=0)
  # assert gate is acting on length(ind) qubits

  k = length(ind)
  pad = max(pad, maximum(ind))

  # w = size(gate)[1]
  P = Permute(ind, pad=pad)

  ret = gate
  for i in k+1:pad
    ret = [1 0; 0 1] ⊗ ret
  end

  return P'*ret*P
end

# https://docs.microsoft.com/en-us/qsharp/api/qsharp/microsoft.quantum.intrinsic.rx
function Rx(θ::Real)
  return [cos(θ/2) -im*sin(θ/2); -im*sin(θ/2) cos(θ/2)]
end

# https://docs.microsoft.com/en-us/qsharp/api/qsharp/microsoft.quantum.intrinsic.ry
function Ry(θ::Real)
  return [cos(θ/2) -sin(θ/2); sin(θ/2) cos(θ/2)]
end

# https://docs.microsoft.com/en-us/qsharp/api/qsharp/microsoft.quantum.intrinsic.rz
function Rz(θ::Real)
  return [exp(-im*θ/2) 0; 0 exp(im*θ/2)]
end

# https://docs.microsoft.com/en-us/qsharp/api/qsharp/microsoft.quantum.machinelearning.controlledrotation
# little endian!
# function ControlledRotation(ind::Tuple{Int, Array{}}, Pauli::Function,  param::Real)
function ControlledRotation(target::Int, control::Int, Pauli::Function, param::Real; pad::Int=0)
  # indices are 1 based
  # id = [1 0; 0 1]

  if target==control
    return ApplyGateOn( Pauli(param), [target], pad=pad)
    # return Rz(param)
  else
    C = Controlled(Pauli(param))
    return ApplyGateOn( C, [target, control], pad=pad)
  end

end