using Plots,DifferentialEquations
Plots.PlotlyJSBackend()
#=
Objective: This script intends to reproduce the Fast Offshore Wells Model.
Paper: Fast Offshore Wells Model (FOWM): A practical dynamic model for multiphase oil production systems in deepwater and ultra-deepwater scenarios
Authors: Fabio C.Diehl Thiago K.Anzai Cristina S.Almeida Oscar F.Von Meien Saul S.Neto Vinicius R.Rosa Mario C.M.M.Campos Filipe Reolon Giovani Gerevini Cassiano Ranza MarceloFarenzena Jorge O.Trierweiler
=#
function fowm(du,u,p,t)
    #=
    States:
        01 - m_ga:      Gas Annular
        02 - m_gt:      Gas Tubing
        03 - m_lt:      Liq Tubing
        04 - m_gb:      Gas Bubble
        05 - m_gr:      Gas Flowline
        06 - m_lr:      Liq Flowline
    =#
    m_ga,m_gt,m_lt,m_gb,m_gr,m_lr = u
    #=
    Parameters:
    - Operation Parameter:
        01 - z: Choke Valve Opening Fraction
    - Well Parameters
        01 - ρ_l:       Liquid Density
        02 - P_r:       Reservoir's Pressure
        03 - P_s:       After Choke Pressure
        04 - α_gw:      Gas Mass Fraction at Reservoir's Condition
        05 - ρ_mres:    Reservoir's Density
        06 - M:         Gas Molecular Weight
        07 - T:         Average Temperature
        08 - L_r:       Riser Length
        09 - L_fl:      Flowline Length
        10 - L_t:       Tubing Length
        11 - L_a:       Annular Length
        12 - H_t:       Distance Xmas Tree - Bottom Hole
        13 - H_pdg:     Distance Xmas Tree - PDG Point
        14 - H_vgl:     Distance Xmas Tree - Gas Lift Injection Point
        15 - D_ss:      Subsea Pipeline Diameter
        16 - D_t:       Tubing Diameter
        17 - D_a:       Annular Diameter
        18 - θ:         Annular Diameter
    - Estimated Parameters in Diehl et al.
        01 - m_still:   Minimum Liquid Mass at Subsea Pipeline
        02 - C_g:       Virtual Valve Flow Constant
        03 - C_out:     Choke Valve Flow Constant
        04 - V_eb:      Bubble Volume
        05 - E:         Mass Fraction of Gas Bypassing Virtual Valve
        06 - K_w:       Xmas Tree Flow Constant
        07 - K_a:       Annular-Tubing Flow Constant
        08 - K_r:       Reservoir's Flow Constant
        09 - ωᵤ:        Bubble Location (Assistant Parameter)
    =#
    z,ρ_l,P_r,P_s,α_gw,ρ_mres,M,T,L_r,L_fl,L_t,L_a,H_t,
    H_pdg,H_vgl,D_ss,D_t,D_a,θ,m_still,C_g,C_out,V_eb,E,K_w,K_a,K_r,ωᵤ = p

    # Universal Gas Constant
    R      = 8314 # J/kmolK
    g      = 9.81
    #=
    First Rank Variables:
        01 - V_t:       Tubing Volume
        02 - V_ss:      Bubble Downstream Volume
        03 - V_a:       Annular Volume
        04 - A_ss:      Subsea Pipeline Cross Sectional Area
        05 - V_gt:      Tubing Gas Volume
        06 - ρ_mt:      Mixture Density
        07 - ρ_gt:      Gas Density
    =#
    V_t    = π*D_t^2/4*L_t
    V_ss   = π*D_ss^2/4*L_r + π*D_ss^2/4*L_fl
    V_a    = π*D_a^2/4*L_a
    A_ss   = π*D_ss^2/4
    V_gt   = V_t - m_lt/ρ_l
    ρ_mt   = (m_gt+m_lt)/V_t
    ρ_gt   = m_gt/V_gt
    #=
    Second Rank Variables:
        01 - P_ai:      Annular Gas Injection Point Pressure
        02 - P_tt:      Tubing's Topside Pressure
        03 - P_tb:      Tubing Gas Injection Point Pressure
        04 - P_pdg:     Pressure at PDG Point
        05 - P_bh:      Pressure at Bottom Hole
        06 - P_rt:      Pressure at the Riser's Topside
        07 - P_rb:      Pressure before the Bubble
        08 - P_eb:      Bubble Pressure
        09 - ρ_ai:      Gas Density at Annular
        10 - α_gt:      Gas Mass Fraction at Tubing
        11 - α_gr:      Gas Mass Fraction at Subsea Pipeline
        12 - α_lr:      Liquid Mass Fraction at Subsea Pipeline

    =#
    P_ai   = m_ga*(R*T/V_a/M+g*L_a/V_a)
    P_tt   = ρ_gt*R*T/M
    P_tb   = P_tt + ρ_mt*g*H_vgl
    P_pdg  = P_tb +ρ_mres*g*(H_pdg-H_vgl)
    P_bh   = P_pdg + ρ_mres*g*(H_t-H_pdg)
    P_rt   = m_gr*R*T/M/(ωᵤ*V_ss-(m_lr+m_still)/ρ_l)
    P_rb   = P_rt + (m_lr+m_still)*g*sin(θ)/A_ss
    P_eb   = m_gb*R*T/M/V_eb
    ρ_ai   = M*P_ai/R/T
    α_gt   = m_gt/(m_gt+m_lt)
    α_gr   = m_gr/(m_gr+m_lr)
    α_lr   = 1-α_gr
    #=
    Mass Flow Rates:
        01 - W_iv:      Gas Annular to tubing
        02 - W_r:       Reservoir to Bottom Hole
        03 - W_whg:     Gas at Xmas Tree
        04 - W_whl:     Liquid at Xmas Tree
        05 - W_gc:      Gas Annular
        06 - W_g:       Gas at Virtual Valve
        07 - W_gout:    Gas through Choke
        08 - W_lout:    Liquid through Choke
    =#
    W_iv   = K_a*√(ρ_ai*max(0,P_ai-P_tb))
    W_r    = K_r*(1 - 0.2*P_bh/P_r-(0.8*P_bh/P_r)^2)
    W_whg  = K_w*α_gt*√(ρ_l*max(0,P_tt-P_rb))
    W_whl  = K_w*(1-α_gt)*√(ρ_l*max(0,P_tt-P_rb))
    W_gc   = W_iv
    W_g    = max(0,C_g*(P_eb-P_rb))
    W_gout = α_gr*C_out*z*√(ρ_l*max(0,P_rt-P_s))
    W_lout = α_lr*C_out*z*√(ρ_l*max(0,P_rt-P_s))
    # States' Derivative
    du[1]  = δm_ga = W_gc - W_iv
    du[2]  = δm_gt = W_r*α_gw+W_iv-W_whg
    du[3]  = δm_lt = W_r*(1-α_gw) - W_whl
    du[4]  = δm_gb = (1-E)*W_whg - W_g
    du[5]  = δm_gr = E*W_whg + W_g - W_gout
    du[6]  = δm_lr = W_whl - W_lout
end

function simulate_fowm(u₀,tf,z,Δt=[])
    #=
    Args:
        - u₀: Initial Condition
        - tf: Final Simulation Time (hours)
        - z:  Choke Valve Opening Fraction
        - Δt: Time Step to Save Results (hours)
    =#
    tspan = (0,tf*3600)
    p = [z,ρ_l,P_r,P_s,α_gw,ρ_mres,M,T,L_r,L_fl,L_t,L_a,H_t,H_pdg,H_vgl,D_ss,D_t,D_a,θ,m_still,C_g,C_out,V_eb,E,K_w,K_a,K_r,ωᵤ]
    prob = ODEProblem(fowm,u₀,tspan,p)
    sol = solve(prob,Tsit5(),saveat = Δt)
    return sol
end

# Parameters for Simulation (Diehl et al. Table 2 Well A and Table 3 Real Well A)
ρ_l         =   900.0
P_r         =   225e5
P_s         =   10e5
α_gw        =   0.0188
ρ_mres      =   892.0
M           =   18.0
T           =   298.0
L_r         =   1569.0
L_fl        =   2928.0
L_t         =   1639.0
L_a         =   1118.0
H_t         =   1279.0
H_pdg       =   1117.0
H_vgl       =   916.0
D_ss        =   0.15
D_t         =   0.15
D_a         =   0.14
θ           =   π*60/180
m_still     =   6.222e1
C_g         =   1.137e-3
C_out       =   2.039e-3
V_eb        =   6.098e1
E           =   1.545e-1
K_w         =   6.876e-4
K_a         =   2.293e-5
K_r         =   1.269e2
ωᵤ          =   2.780e0

# Initial Conditions
x₀ = repeat([10000.0],6)
# Simulation Time and Time-step in Hours
tf = 30.0
Δt = 1/60*15*3600
# Choke Opening Fraction
z = 0.5
# Simulating
sol = simulate_fowm(x₀,tf,z,Δt)

m_ga = sol[1,:]
m_gt = sol[2,:]
m_lt = sol[3,:]
m_gb = sol[4,:]
m_gr = sol[5,:]
m_lr = sol[6,:]

V_t    = π*D_t^2/4*L_t
V_ss   = π*D_ss^2/4*L_r + π*D_ss^2/4*L_fl
V_a    = π*D_a^2/4*L_a
A_ss   = π*D_ss^2/4
V_gt   = V_t .- m_lt./ρ_l
ρ_mt   = (m_gt+m_lt)/V_t
ρ_gt   = m_gt./V_gt

R      = 8314 # J/kmolK
g      = 9.81

P_ai   =  m_ga.*(R.*T./V_a./M+g.*L_a./V_a)
P_tt   =  ρ_gt.*R.*T./M
P_tb   =  P_tt .+ ρ_mt.*g.*H_vgl
P_pdg  =  P_tb .+ρ_mres.*g.*(H_pdg.-H_vgl)
P_bh   =  P_pdg .+ ρ_mres.*g.*(H_t-H_pdg)
P_rt   =  m_gr.*R.*T./M./(ωᵤ.*V_ss.-(m_lr.+m_still)./ρ_l)
P_rb   =  P_rt .+ (m_lr.+m_still).*g.*sin.(θ)./A_ss
P_eb   =  m_gb.*R.*T./M./V_eb
ρ_ai   =  M.*P_ai./R./T
α_gt   =  m_gt./(m_gt.+m_lt)
α_gr   =  m_gr./(m_gr.+m_lr)
α_lr   =  1 .-α_gr

t = sol.t

plotly()
plot(t,P_pdg)


mutable struct Points
	position::Array{Float64,2}
	velocity::Array{Float64,2}
	acceleration::Array{Float64,2}
	mass::Array{Float64,1}
	size::Int
end

function acceleration_update(points::Points, force_field::Array{Float64,3})
	for i = 1:points.size
		points.acceleration[:,i] = force_field[:,ceil(Int,points.position[1,i]/dx),ceil(Int,points.position[2,i]/dx)] / points.mass[i]
	end
end

function velocity_update(points::Points)
	for i = 1:points.size
		points.velocity[:,i] += points.acceleration[:,i] * dt #(accelration+old_acc)/2
	end
end

function position_update(points::Points)
	for i = 1:points.size
		points.position[:,i] += points.velocity[:,i] * dt  #+ acceleration*dt*dt/mass
	end
end

function assign_mass(density::Array{Float64,2}, points::Points, dx)
	for i = 1:points.size
		density[ceil(Int, points.position[1,i]/dx),ceil(Int, points.position[2,i]/dx)] = points.mass[i]
	end
end

function gauss_seidel(potential::Array{Float64,2}, density::Array{Float64,2})
	for i = 2:grid-1
		for j = 2:grid-1
			potential[i,j] = -1/4 * (dx^2*4*pi*GravitationalConstant*density[i,j] + (- potential[i-1,j] - potential[i+1,j] - potential[i,j-1] - potential[i,j+1]))
		end
	end
end

function residuum(potentail::Array{Float64,2}, density::Array{Float64,2})
	r = 0.
	for i = 2:grid-1
		for j = 2:grid-1
			r += (4*pi*GravitationalConstant*density[i,j] - (potentail[i+1,j] + potentail[i-1,j] + potentail[i,j+1] + potentail[i,j-1] - 4potentail[i,j]) * 1/dx^2)^2
		end
	end
	return sqrt(r/(grid-2)^2)
end

using WriteVTK

N = 20
grid = 200
ϵ = 0.0001 #residual
ω = 1.8 #SOR
γ = 0.9 #upwind
itermax = 2000
dims = 2
GravitationalConstant = 6.673e-8
cells = WriteVTK.MeshCell[MeshCell(VTKCellTypes.VTK_VERTEX, [i]) for i = 1:N]

points = Points(rand(dims,N), zeros(dims,N), zeros(dims,N), rand(N)*100, N)
density = zeros(Float64, grid, grid)
potential = zeros(Float64, grid, grid)
force_field = zeros(Float64, 2, grid, grid)

dx = 1/grid
dt = 1

for i = 1:60
	outfile = vtk_grid(string("my_vtr_file",i), [i*dx for i=0:grid],[i*dx for i=0:grid]) do vtk
		vtk_cell_data(vtk, density, "Density")
		vtk_cell_data(vtk, potential, "Potential")
		vtk_cell_data(vtk, vcat(force_field, zeros(1,grid,grid)), "Force")
	end

	outfile = vtk_grid(string("my_vtu_file",i), points.position, cells) do vtk
		vtk_cell_data(vtk, points.mass, "Mass")
		vtk_cell_data(vtk, vcat(points.velocity, zeros(1,N)), "Velocity")

		vtk_point_data(vtk, points.mass, "Pointmass")
		vtk_point_data(vtk, vcat(points.velocity, zeros(1,N)), "Pointvelocity")
	end

	density = zeros(Float64, grid, grid)
	assign_mass(density, points, dx)
	for n=1:itermax
		gauss_seidel(potential, density)
		r = residuum(potential, density)
		println(r)
		if r < ϵ
			break
		end
	end

	for i=2:grid-1
		for j=2:grid-1
			force_field[1,i,j] = -(potential[i+1,j] - potential[i-1,j])/dx
			force_field[2,i,j] = -(potential[i,j+1] - potential[i,j-1])/dx
		end
	end

	acceleration_update(points,force_field)
	velocity_update(points)
	position_update(points)
end


using JuMP
using CPLEX
include("instances/exemple-lecture-graphe.jl")

function Benders(file::String, b::Int64)
  n, m, adj, demande = readGraph(file::String)
  #### probleme maitre ####
  # Create the model
  m0 = Model(CPLEX.Optimizer)

  ## Variables
  @variable(m0, y[1:m] >= 0, Int)

  ## Constraints

  ## Objective
  @objective(m0, Min, sum(y[i] for i in 1:m))

  #resolution
  optimize!(m0)

  y_star = JuMP.getvalue.( m0[:y] )

  valid = false

  while valid == false
    #### sous-probleme ####
    # Create the model
    m1 = Model(CPLEX.Optimizer)

    ## Variables
    @variable(m1, v[1:m+n] >= 0)

    ## Constraints
    @constraint(m1, [i in 1:m], - v[i] - v[adj[2*i-1]] + v[adj[2*i]] <= 0)
    @constraint(m1, [i in 1:m], - v[i] + v[adj[2*i-1]] - v[adj[2*i]] <= 0)
    @constraint(m1, v[1] == 0) # on saute la source
    @constraint(m1, sum(v[i] for i in 1:m+n) == 1)

    ## Objective
    @objective(m1, Max, - sum(b*y_star[ij]*v[n+ij] for ij in 1:m) + sum(demande[i]*v[i] for i in 1:n))

    #resolution
    optimize!(m1)

    v_star = JuMP.getvalue.( m1[:v] )

    if - sum(b*y_star[ij]*v_star[n+ij] for ij in 1:m) + sum(demande[i]*v_star[i] for i in 1:n) > 0 + 1e-8

      #### adding Constraints
      @constraint(m0, - sum(b*y[ij]*v_star[n+ij] for ij in 1:m) + sum(demande[i]*v_star[i] for i in 1:n) <= 0) # (4b)

      optimize!(m0)
      y_star = JuMP.getvalue.( m0[:y] )

    else
      valid = true
    end
  end
  println(y_star)
  println(JuMP.getobjectivevalue.( m0 ))
end


# %%
using Tmatrix

# inputs
rx, rz = 1e-6, 1.3e-6
n_max = 3
k1_r = 1e7;
k1_i = 0.0;
k2_r = 1.5e7;
k2_i = 1e3;

n_θ_points = 10
n_ϕ_points = 20

rotationally_symmetric = false

θ_array, ϕ_array = meshgrid_θ_ϕ(
    n_θ_points,
    n_ϕ_points;
    min_θ = 1e-16,
    min_ϕ = 1e-16,
    rotationally_symmetric = rotationally_symmetric,
)

r_array, n̂_array = ellipsoid(rx, rz, θ_array)


include("distribution.jl")

export TE

ϵ = 0.0000001

function TE_ML(data::Matrix{Int64}; base=2) # D_KL(p(x|y,z)||p(x|y))

  x2         = data[2:end,1]
  x1         = data[1:end-1,1]
  y1         = data[1:end-1,2]
  x_max      = maximum(data[:,1])
  y_max      = maximum(data[:,2])
  px2x1y1    = Shannon.fe3p(hcat(x2,x1,y1))
  px1y1      = sum(px2x1y1, 1)
  px2x1      = sum(px2x1y1, 3)
  px1        = sum(px2x1y1, (1,3))

  px2_c_x1y1 = zeros(x_max, x_max, y_max)
  px2_c_x1   = zeros(x_max, x_max)

  for x2_index = 1:x_max
    for x1_index = 1:x_max
      px2_c_x1[x2_index, x1_index] = px2x1[x2_index, x1_index, 1] / px1[1, x1_index, 1]
      for y1_index = 1:y_max
        px2_c_x1y1[x2_index, x1_index, y1_index] = px2x1y1[x2_index, x1_index, y1_index] / px1y1[1, x1_index, y1_index]
      end
    end
  end

  r = 0
  for x2_index = 1:x_max
    for x1_index = 1:x_max
      for y1_index = 1:y_max
        if abs(px2x1y1[x2_index, x1_index, y1_index]) > ϵ && abs(px2_c_x1y1[x2_index, x1_index, y1_index]) > ϵ && abs(px2_c_x1[x2_index, x1_index]) > ϵ
          r = r + px2x1y1[x2_index, x1_index, y1_index] * (log(base, px2_c_x1y1[x2_index, x1_index, y1_index]) - log(base, px2_c_x1[x2_index, x1_index]))
        end
      end
    end
  end
  return r
end

  # Transfer Entropy
  function TE(data::Matrix{Int64}; base=2, mode="ML", pseudocount=0)
    @assert (size(data)[2] == 2) "TE required two columns (X,Y) to calcualte TE(X->Y)."

    modes      = ["ML", "Maximum Likelihood"]
    umodes     = map(x->uppercase(x), modes)
    known_mode = uppercase(mode) in umodes
    pmodes     = map(x->", $x",modes)
    pmodes     = foldl(*, pmodes[1][3:end], pmodes[2:end])
    @assert known_mode "Mode may be any of the following: [$pmodes]."

    if uppercase(mode) in umodes[1:2]
      return TE_ML(data, base)
    end
  end


import RRSTExperiments: InputEdge

"""
    solve_inc_st_with_model(n::Int, E::Vector{InputEdge}, x::Vector{Tuple{Int, Int}}, k::Int)

The incremental minimum spanning tree problem
"""
function solve_inc_st_with_model(n::Int, S::Vector{Float64}, E::Vector{InputEdge}, x::Vector{Tuple{Int, Int}}, k::Int)
    @assert length(S) == length(E)
    V = collect(1:n) # set of nodes
    Vminus1 = setdiff(V, [1]) # commodity nodes
    A = [(e.i, e.j) for e ∈ E] ∪ [(e.j, e.i) for e ∈ E]
    X = Dict([(e.i, e.j) => ((e.i, e.j) ∈ x || (e.j, e.i) ∈ x ? 1 : 0) for e ∈ E])
    L = n-1-k

    # Model
    model = Model(CPLEX.Optimizer)
    set_silent(model)
    set_optimizer_attribute(model, "CPX_PARAM_EPINT", 1e-8)

    # Variables
    @variable(model, y[E] ≥ 0) # yₑ=1 if e∈E belongs to the spanning tree; 0 otherwise
    @variable(model, w[A] ≥ 0)
    @variable(model, f[A, Vminus1] ≥ 0) # flow variables

    # Objective
    @objective(model, Min, sum(S[i] * y[e] for (i, e) ∈ enumerate(E)))

    #Constraints
    for k ∈ Vminus1 # sources
        @constraint(model, sum(f[(j,1),k] for j ∈ filter(j -> (j,1) ∈ A, V)) -
        sum(f[(1,j),k] for j ∈ filter(j -> (1,j) ∈ A,V)) == -1)
    end

    for k ∈ Vminus1, i ∈ Vminus1 # balances
        if i ≠ k
            @constraint(model, sum(f[(j,i),k] for j ∈ filter(j -> (j,i) ∈ A, V)) -
            sum(f[(i,j),k] for j ∈ filter(j -> (i,j) ∈ A, V)) == 0)
        end
    end

    for k ∈ Vminus1 # sinks
        @constraint(model, sum(f[(j,k),k] for j ∈ filter(j -> (j,k) ∈ A,V)) -
        sum(f[(k,j),k] for j ∈ filter(j -> (k,j) ∈ A,V)) == 1)
    end

    for k ∈ Vminus1, a ∈ A # capacity
        @constraint(model, f[a,k] ≤ w[a])
    end

    # tree
    @constraint(model, sum(w[a] for a ∈ A) == n - 1)

    for e ∈ E
        @constraint(model, y[e] == w[(e.i, e.j)] + w[(e.j, e.i)])
    end

    @constraint(model, sum(y[e] * X[(e.i, e.j)] for e ∈ E) ≥ L)

    # Solve
    optimize!(model)

    status=termination_status(model)
    if status == MOI.OPTIMAL
        obj_value = objective_value(model)
        return status, obj_value, Array(value.(y))
    else
        return status, missing, missing
    end
end

function solve_inc_st_with_algorithm(n::Int, S::Vector{Float64}, E::Vector{InputEdge}, x::Vector{Tuple{Int, Int}}, k::Int)
    Cₘ = maximum(map(e -> e.C, E))
    A = [InputEdge(e.i, e.j, (e.i, e.j) in x || (e.j, e.i) in x ? 0.0 : 10 * Cₘ, S[i]) for (i, e) in enumerate(E)]

    return RRSTExperiments.solve_rec_st_with_algorithm(n, A, k)
end


function ssdDistance(referenceImage::regImage,templateImage::regImage,
                     transformedGrid::Array{Float64,1};
                     doDerivative::Bool=false,doHessian::Bool=false,options::regOptions=regOptions(),centeredGrid::Array{Float64,1}=zeros(1))

  # get relevant options
  parametricOnly = options.parametricOnly
  matrixFree = options.matrixFree

  # check centered grid
  if(size(centeredGrid)[1]==1)
    centeredGrid = getCellCenteredGrid(referenceImage).data
  end

  # check grid
  if(checkForOddNumberOfGridPoints(referenceImage.data,transformedGrid))
    error("StaggeredGrids are not supported. Please transform to cell centered first.")
  end

  # interpolation of the template image at transformed grid points
  transformedImage, dX_transformedImage, dY_transformedImage =
      interpolateImage(templateImage,transformedGrid,doDerivative=true)

  # measure the ssd distance
  N = prod(size(referenceImage.data)) # product of sizes

  residual = Array(transformedImage .- referenceImage.data)[:]
  prodH = prod(referenceImage.voxelsize)
  functionValue = 0.5 * prodH * residual' * residual
  # calculate ssd derivatives matrix free?
  if(matrixFree)
    dFunctionValue,d2FunctionValue = ssdDerivativesMatrixFree(dX_transformedImage,dY_transformedImage,residual,centeredGrid,prodH,N,parametricOnly,doDerivative,doHessian)
  else
    dFunctionValue,d2FunctionValue = ssdDerivativesMatrixBased(dX_transformedImage,dY_transformedImage,residual,centeredGrid,prodH,N,parametricOnly,doDerivative,doHessian)
  end

  return [functionValue,dFunctionValue,d2FunctionValue,(dX_transformedImage,dY_transformedImage)]

end

function ssdDerivativesMatrixBased(dX_transformedImage::Array{Float64,1},
                                   dY_transformedImage::Array{Float64,1},
                                   residual::Array{Float64,1},
                                   centeredGrid::Array{Float64,1},
                                   pixelSize::Number,N::Number,
                                   parametricOnly::Bool,doDerivative::Bool,doHessian::Bool)

  if(doDerivative)
    dTransformedImage = spdiagm((dX_transformedImage,dY_transformedImage),[0,N])
    if(parametricOnly)
      Q = sparse([centeredGrid[1:N] centeredGrid[N+1:end] ones(N)])
      Q = [Q spzeros(size(Q)[1],size(Q)[2]); spzeros(size(Q)[1],size(Q)[2]) Q]
      dTransformedImage = dTransformedImage * Q
    end
    dFunctionValue = pixelSize .* dTransformedImage' * residual
  else
    dFunctionValue = 0.0
    dTransformedImage = 0.0
  end

  if(doHessian)
    d2FunctionValue = pixelSize .* dTransformedImage' * dTransformedImage
  else
    d2FunctionValue = 0.0
  end

  return dFunctionValue, d2FunctionValue

end

function ssdDerivativesMatrixFree(dX_transformedImage::Array{Float64,1},
                                  dY_transformedImage::Array{Float64,1},
                                  residual::Array{Float64,1},
                                  centeredGrid::Array{Float64,1},
                                  pixelSize::Number,N::Number,
                                  parametricOnly::Bool,doDerivative::Bool,doHessian::Bool)

  if(doDerivative)
    dFunctionValue = BLAS.scal(2*N,pixelSize,dTtransposedMultiplication(N,dX_transformedImage,dY_transformedImage,residual),1)
    if(parametricOnly)
      dFunctionValue = QtransposedMultiplication(N,dFunctionValue,centeredGrid)
    end
  else
    dFunctionValue = 0.0
    dTransformedImage = 0.0
  end

  if(doHessian)
    if(parametricOnly)
      d2FunctionValue = computeParametricHessian(N,pixelSize,dX_transformedImage,dY_transformedImage,centeredGrid)
    else
      d2FunctionValue(x) = hessianFunction(N,pixelSize,dX_transformedImage,dY_transformedImage,x)
    end
  else
    d2FunctionValue = 0.0
  end

  return dFunctionValue,d2FunctionValue

end


using Turing
using Distributions
using Base.Test

# Test the @assume and @predict macros on a model without conditioning.

@model test begin
  @assume z ~ Normal(0,1)
  @assume x ~ Bernoulli(1)
  @observe 1 ~ Bernoulli(x / 2)
  @observe 0 ~ Bernoulli(x / 2)
  @predict x
end

s = SMC(10000)
p = PG(100,10)

res = sample(test, s)

@test reduce(&, res[:x]) == 1  #check that x is always 1
@test res[:logevidence] ≈ 2 * log(0.5)

res = sample(test, p)

@test reduce(&, res[:x]) == 1  #check that x is always 1
# PG does not provide logevidence estimate


using BetaDispersion, Distances, StatsBase,BenchmarkTools, RCall
using Test
R"library(vegan)"

   x = rand(1000,100)
   y =rand(1:20,1000)
   j = @benchmark dispersion($x,$y,BrayCurtis)
   r = @benchmark R"betadisper(vegdist($x),as.factor($y))"
   @test mean(j.times) < mean(r.times)

   Dj = pairwise(BrayCurtis(),x,dims = 1)
   Dr =R"vegdist($x)"
   dispj = dispersion(Dj,y)
   dispr = R"betadisper($Dr,as.factor($y))"
   j = @benchmark permutest(dispj,9999)
   r = @benchmark R"permutest($dispr,pairwise = TRUE, permutations = 9999)"
   @test mean(j.times) < mean(r.times)


using JuMP
using GLPK
using ModelGraphs

m1 = Model(with_optimizer(GLPK.Optimizer))

@variable(m1, xm[i in 1:2],Bin)
@constraint(m1, xm[1] + xm[2] <= 1)
@objective(m1, Max, 16xm[1] + 10xm[2])

m2 = Model(solver=GLPKSolverMIP())
@variable(m2, xs[i in 1:2],Bin)
@variable(m2, y[i in 1:2], Bin)
@constraint(m2, y[1] + y[2] <= 1)
@constraint(m2, 8xs[1] + 2xs[2] + y[1] + 4y[2] <= 10)
@objective(m2, Max, 4y[2])

## Model Graph
graph = ModelGraph()
heur(g) = 16
setsolver(graph, LagrangeSolver(update_method=:subgradient,max_iterations=30,lagrangeheuristic=heur))
n1 = add_node(graph)
setmodel(n1,m1)
n2 = add_node(graph)
setmodel(n2,m2)

## Linking
# m1[x] = m2[x]  ∀i ∈ {1,2}
@linkconstraint(graph, [i in 1:2], n1[:xm][i] == n2[:xs][i])

solution = solve(graph)

true


using Statistics
N=100000;
val=zeros(N)
for i=1:N
    gamma=-2.1 # minimum activation probability
    ActMax = 1
    agent_act =
        (
            (ActMax .^ (gamma + 1) - ActMin .^ (gamma + 1)) .* rand(N) .+
            ActMin .^ (gamma + 1)
        ) .^ (1.0 / (gamma + 1))
    val[i]=mean(agent_act)+sqrt(mean(agent_act .* agent_act))
end
# mmmean=
# s=1.5104924655094145

# s1=-0.019221243887172085

println(mean(val))


using JuMP, Gurobi, StatsFuns

# problem parameters
NumAssets = [1000] # number of variables minus one
epsilon = logspace(-4.5,-0.3011,50)
epsilon = reverse(epsilon) # constraint violation tolerance

# iterate over the number of assets
for ass = 1:size(NumAssets,1)

	# random variable distribution parameters
	mean = Array{Float64}(NumAssets[ass])
	stdev = Array{Float64}(NumAssets[ass])
	for i = 1:NumAssets[ass]
		mean[i] = 1.05 + 0.3*(NumAssets[ass] - i)/(NumAssets[ass]-1)
		stdev[i] = (0.05 + 0.6*(NumAssets[ass] - i)/(NumAssets[ass]-1))/3
	end

	# matrix with stdevs on the diagonal
	sigma = zeros(Float64, NumAssets[ass], NumAssets[ass])
	for i = 1:NumAssets[ass]
		sigma[i,i] = stdev[i]
	end

	# iterate over the epsilons
	for iter_e = 1:size(epsilon,1)

		println("Iter: (",ass,",",iter_e,")  ","NumAssets: ",NumAssets[ass],"  epsilon: ",epsilon[iter_e])

		# create folder for storing results
		dirName = "C:/Users/rkannan/Desktop/SA for CCP/portfolio_analytical/" *
					string(NumAssets[ass]) * "/" * string(iter_e) * "/"
		mkpath(dirName)

		# write details of case study to text file
		details_file = dirName * "portfolio1.txt"
		open(details_file, "w") do f
			write(f,"Model: portfolio1 \n")
			write(f,"Method: EXACT \n")
			write(f,"Number of assets: $(NumAssets[ass]) \n")
			write(f,"epsilon: $(epsilon[iter_e]) \n")
		end
		# write epsilon to text file
		epsilon_file = dirName * "epsilon.txt"
		open(epsilon_file, "w") do f
			write(f,"$(epsilon[iter_e]) \n")
		end

		# pass params as keyword arguments to GurobiSolver
		m = Model(solver=GurobiSolver(Presolve=0,OutputFlag=0))

		@variable(m, objvar)
		@variable(m, 0 <= y[1:NumAssets[ass]] <= 1)

		@objective(m, Max, objvar)
		@constraint(m, sum(y[i] for i = 1:NumAssets[ass]) == 1)
		@constraint(m, soc, sum(mean[i]*y[i] for i = 1:NumAssets[ass]) >= objvar + norminvcdf(1-epsilon[iter_e])*norm(sigma*y))

		status = solve(m)

		objvar_opt = getvalue(objvar)
		time_to_opt = getsolvetime(m)

		y_opt = getvalue(y)
		soln_file = dirName * "soln.txt"
		open(soln_file, "w") do f
			for i = 1:NumAssets[ass]
				write(f,"$(y_opt[i]) ")
			end
		end

		# write objective
		obj_file = dirName * "objvals.txt"
		writeMode = "w"
		open(obj_file, writeMode) do f
			write(f,"$objvar_opt \n")
		end
		# write solution time
		time_file = dirName * "solntime.txt"
		open(time_file, writeMode) do f
			write(f,"$time_to_opt \n")
		end

		println("Objective: ",objvar_opt,"  Time: ",time_to_opt)
		println("")

		sleep(2)

	end

	sleep(5)
end


include("./params.jl")

function fitgamma(isi)

    epsilon = 1e-6;

    n = length(isi);
    T = sum(isi);

    lambda = n / T;

    alpha = 0.2
    beta = 1.0

    eta = 0.001

    dll_dalpha = epsilon * 10
    dll_dbeta = epsilon * 10

    while max(abs(dll_dalpha), abs(dll_dbeta)) > epsilon

        dll_dalpha = - n / alpha + n * log(beta) + sum(log(isi)) - n * digamma(alpha);
        dll_dbeta = n * (1+alpha) / beta - T;

        alpha = alpha + eta*dll_dalpha;
        beta = beta + eta*dll_dbeta;

    end

    k = alpha + 1
    lambda = beta / k

    lambda, k

end


module TestUtils
    using ManifoldLearning
    using Test

    @testset "Utils" begin

	k = 3
	X = rand(3, 15)

	# test k-nn graph
	D, E = ManifoldLearning.find_nn(X, k, excluding=false)
	@test size(X,2) == size(D,2) && size(D, 1)-1 == k
	@test size(X,2) == size(E,2) && size(E, 1)-1 == k

	D, E = ManifoldLearning.find_nn(X, k)
	@test size(X,2) == size(D,2) && size(D, 1) == k
	@test size(X,2) == size(E,2) && size(E, 1) == k

	# test connected components
	CC = ManifoldLearning.components(E)
	@test length(CC) > 0

	# test shortest path
        D = transpose([7 9 14; 7 10 15; 9 10 2; 15 6 0; 6 9 11; 14 2 9])
        E = transpose([2 3 6; 1 3 4; 1 2 6; 2 5 4; 3 4 6; 1 3 5])
	P, PD = ManifoldLearning.dijkstra(D, E, 1)
	@test PD[5] == 20. && PD[6] == 11.
    end
end


export Disk,𝔻

# T is Real or Complex
# TT is (Real,Real) or Complex
"""
    Disk(c,r)

represents the disk centred at `c` with radius `r`.
"""
struct Disk{T} <: Domain{T}
    center::T
    radius::Float64
end

Disk(c,r) = Disk{typeof(c)}(c,r)
Disk(c::Real,r) = Disk(complex(c),r)
Disk(c::Integer,r) = Disk(float(c),r)
Disk(c::Complex{<:Integer},r) = Disk(float(c),r)

Disk() = Disk(Vec(0.,0.),1.)
Disk(::AnyDomain) = Disk(NaN,(NaN,NaN))

const 𝔻 = Disk()

isambiguous(d::Disk) = isnan(d.radius) && all(isnan,d.center)

#canonical is rectangle [r,0]x[-π,π]
# we assume radius and centre are zero for now
fromcanonical(D::Disk{T},x,t) where {T<:Real} =
    Vec(D.radius*x*cos(t)+D.center[1],D.radius*x*sin(t)+D.center[2])
tocanonical(D::Disk{T},x,y) where {T<:Real} =
    Vec(sqrt((x-D.center[1])^2+(y-D.center[2])^2)/D.radius,
        atan(y-D.center[2],x-D.center[1]))
checkpoints(d::Disk) = [fromcanonical(d,(.1,.2243));fromcanonical(d,(-.212423,-.3))]

# function points(d::Disk,n,m,k)
#     ptsx=0.5*(1-gaussjacobi(n,1.,0.)[1])
#     ptst=points(PeriodicSegment(),m)
#
#     Float64[fromcanonical(d,x,t)[k] for x in ptsx, t in ptst]
# end

boundary(d::Disk) = Circle(d.center,d.radius)


using Plots
using Interact
# plotlyjs()
pyplot()

unique_vehicle_ids(high_level) = sort(unique(high_level[:vehicle_ID]))

function plot_number_of_trips_per_vehicle(high_level)
    vehicles = unique_vehicle_ids(high_level)
    counts = map(id->sum(high_level[:vehicle_ID] .== id), vehicles)
    bar(collect(1:length(vehicles)), counts, size=(800,400), xlabel="vehicle", ylabel="number of trips")
end

function plot_histogram_per_vehcile(high_level, target)
    vehicles = unique_vehicle_ids(high_level)
    xlims = extrema(high_level[target])
    p1 = histogram(high_level[find(high_level[:vehicle_ID] .== vehicles[1]), target], title=1, xlims=xlims)
    p2 = histogram(high_level[find(high_level[:vehicle_ID] .== vehicles[2]), target], title=2, xlims=xlims)
    p3 = histogram(high_level[find(high_level[:vehicle_ID] .== vehicles[3]), target], title=3, xlims=xlims)
    p4 = histogram(high_level[find(high_level[:vehicle_ID] .== vehicles[4]), target], title=4, xlims=xlims)
    p5 = histogram(high_level[find(high_level[:vehicle_ID] .== vehicles[5]), target], title=5, xlims=xlims)
    plot(p1, p2, p3, p4, p5, size=(950, 600))
end

function scatter_targets(high_level, A, B)
    vehicles = unique_vehicle_ids(high_level)
    p = plot(xlabel=string(A), ylabel=string(B))
    for (i,id) in enumerate(vehicles)
        x = mean(high_level[find(high_level[:vehicle_ID] .== id), A])
        y = mean(high_level[find(high_level[:vehicle_ID] .== id), B])
        scatter!(p, Float64[x], Float64[y], lab=string(i))
    end
    p
end

function plot_drive(index::Int; scatter::Bool=false)
    utmx = Float64[]
    utmy = Float64[]
    path = anon_paths[index]["utm"]
    for pt in path
        push!(utmx, pt[1])
        push!(utmy, pt[2])
    end
    lox, hix = extrema(utmx)
    loy, hiy = extrema(utmy)
    max_sep = max(hix - lox, hiy - loy)
    midx = (hix + lox)/2
    midy = (hiy + loy)/2
    half_sep = max_sep/2*1.1
    xlims = (midx-half_sep, midx+half_sep)
    ylims = (midy-half_sep, midy+half_sep)
    p = plot(utmx, utmy, xlabel="x position [m]", ylabel="y position [m]", xlims=xlims, ylims=ylims, size=(800,800))

    if scatter == true
        scatter!(p, utmx, utmy)
    end

    p
end


"""
    lcparray

Kasai's algorithm for linear-time construction of LCP array from Suffix Array
"""
function lcparray(sa::Vector{Int}, data::Vector{T}) where T<:Integer
    n = length(sa)
    lcps = similar(sa)
    rank = similar(sa)
    for i = 1:n
        rank[sa[i]] = i
    end

    lcp = 0
    for i = 1:n
        rank[i] > 1 || continue
        j = sa[rank[i]-1]
        while i+lcp <= n && j+lcp <= n && data[i+lcp] == data[j+lcp]
            lcp += 1
        end
        lcps[rank[i]] = lcp
        lcp > 0 && (lcp -= 1)
    end
    lcps[1] = 0
    lcps
end


# # Replication of Mackowiak, Matejka and Wiederholt (2018)

# This example replicates [Mackowiak, Matejka and Wiederholt (2018)](https://www.sciencedirect.com/science/article/abs/pii/S002205311830139X) using the [DRIPs](https://github.com/afrouzi/DRIPs.jl) package.

# [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/afrouzi/DRIPs.jl/binder?filepath=examples) to run and modify the following code (no software is needed on the local machine).

# See [Afrouzi and Yang (2020)](http://www.afrouzi.com/dynamic_inattention.pdf) for background on the theory.

# ## Contents
# * [Ex. 1: AR(2) Process](@ref ex1)
#     * [Initialization](@ref ex1_initialize)
#     * [Replication of Figures (2) and (3)](@ref ex1_solution)
#     * [Measure Performance](@ref ex1_performance)
# * [Ex. 2: AR(3) and ARMA(2,1) Processes](@ref ex2)
#     * [Solve and Measure Performance](@ref ex2_solve)
#     * [Replication of Figure (5)](@ref ex2_fig5)
# * [Ex. 3: Price-setting with Rational Inattention](@ref ex3)
#     * [The Case with No Strategic Complementarity](@id ex3_no_strcomp)
#     * [The Case with Strategic Complementarity](@id ex3_strcomp)
#     * [Replication of Figure (6)](@ref ex3_fig6)
# * [Ex. 4: Business Cycle Model with News Shocks](@id ex4)
#     * [Setup](@ref ex4_setup)
#     * [Initialization](@ref ex4_initialize)
#     * [Solve and Measure Performance](@ref ex4_solve)
#     * [Replication of Figure (7)](@ref ex4_fig7)

# ## [Ex. 1: AR(2) Process](@id ex1)
# In this example, authors assume that the optimal action follows an AR(2) process,
# ```math
# \begin{aligned}
#     x_t = \phi_1 x_{t-1} + \phi_2 x_{t-2} + \theta_0 \varepsilon_t
# \end{aligned}
# ```
# Then, we can write a state-space form as:
# ```math
# \begin{aligned}
#         \left[\begin{array}{c} x_t \\ x_{t-1} \end{array}\right]
#         & =
#        \underset{\mathbf{A}}{\underbrace{\left[\begin{array}{cc}
#             \phi_1 & \phi_2 \\
#             1 & 0\end{array}\right]}}\,
#             \left[\begin{array}{c} x_{t-1} \\ x_{t-2} \end{array}\right]
#         + \underset{\mathbf{Q}}{\underbrace{\left[\begin{array}{c}
#             \theta_0 \\
#             0\end{array}\right]}}\, \varepsilon_t
# \end{aligned}
# ```
# We now characterize the optimal signal, $S_{t} = h_1 X_t + h_2 X_{t-1} + \psi_t$, as a function of $\phi_2$ and the capacity for information processing ($\kappa$).

# ### [Initialization](@id ex1_initialize)
# Include the package:
using DRIPs, LinearAlgebra;

# Assign value to deep parameters and define the structure of the problem

β   = 1.0   ; # Time preference
θ0  = 1.0   ;

ϕ2_seq  = 0:0.02:0.9        ;      # sequence of values for the AR(2) parameter
n_ϕ     = length(ϕ2_seq)    ;

κ_seq   = 0.008:0.01:3.5    ;      # sequence of values for the information processing capacity
n_κ     = length(κ_seq)     ;

H   = [1; 0]    ;
Q   = [θ0 ; 0]    ;

# ### [Replication of Figures (2) and (3)](@id ex1_solution)
# Solve for different values of $\phi_2$:

κ   = 0.2   ; # fix κ for this exercise so we can vary ϕ2

h1  = zeros(n_ϕ,1)    ;
h2  = zeros(n_ϕ,1)    ;
h2norm = zeros(n_ϕ,1) ;

for i in 1:n_ϕ
    ϕ2  = ϕ2_seq[i] ;
    ϕ1  = 0.9 - ϕ2 ;
    A   = [ϕ1 ϕ2 ; 1.0 0.0] ;

    ex1a = Drip(κ,β,A,Q,H,fcap=true);

    h1[i] = ex1a.ss.Y[1]
    h2[i] = ex1a.ss.Y[2]
    h2norm[i] = ex1a.ss.Y[2]/ex1a.ss.Y[1] # normalize the first signal weight h1 to one
end

# Here is the replication of Figure (2):
using Printf, LaTeXStrings, Plots; pyplot();
plot(ϕ2_seq,h2norm[:,1],
    title       = L"Optimal Signal, AR(2) processes with $\phi_1 + \phi_2 = 0.9$, $\kappa$ constant.",
    xlabel      = L"$\phi_2$",
    ylabel      = L"$h_2$(with $h_1$ normalized to 1)",
    label       = L"Signal Weight on $X_{t-1}$, $h_2$ (with $h_1$ normalized to 1)",
    legend      = :topleft,
    lw          = 2,
    color       = :black,
    xlim        = (0,0.9),
    xticks      = (0:0.1:0.9),
    titlefont   = font(10), legendfont = font(7), tickfont = font(7),
    size        = (650,300),
    grid        = :off, framestyle = :box)

# Now, we solve for the signal weight and standard deviation of noise, for different values of $\kappa$

ϕ2   = 0.4           ;  # fix ϕ2 in this exercise so we can vary κ
ϕ1   = 0.99 - ϕ2     ;

A    = [ϕ1 ϕ2 ; 1.0 0.0] ;

ω_sol   = zeros(n_κ,1) ;
σz_sol  = zeros(n_κ,1) ;

for i in 1:n_κ
    κ = κ_seq[i] ;

    ex1b = Drip(κ,β,A,Q,H,fcap=true);

    h1temp = ex1b.ss.Y[1]
    h2temp = ex1b.ss.Y[2]

    ω_sol[i]  = 1.0 - (h2temp/h1temp)/(ϕ2 + (1.0-ϕ1)*(h2temp/h1temp)) ;
    ρ = (ω_sol[i] + (1.0-ω_sol[i])*ϕ1)/h1temp
    σz_sol[i] = ρ*ex1b.ss.Σ_z[1,1] ;
end

# Below is the replication of Figure (3):
a = fill(NaN, n_κ, 1)

plot(κ_seq,[ω_sol[:,1] a],
    xlabel      = L"$\kappa$",
    label       = [L"Signal weight on $X_{t}$, $\omega$ (left axis)" L"St. dev. of noise  $\sigma_{\psi}$ (right axis)"],
    title       = L"Optimal signal as function of $\kappa$, AR(2) example.",
    linestyle   = [:solid :dash], color = [:black :gray40],
    xticks      = (0:1:3),
    xlim        = (-0.02,3.02),
    yticks      = (0:0.5:1),
    ylim        = (-0.01,1.01),
    lw          = 2, grid = :off, legend = :right,
    titlefont   = font(10), legendfont = font(8), tickfont = font(8),
    framestyle  = :box)
plot!(twinx(),κ_seq,σz_sol[:,1],
    linestyle   = :dash, color = :gray40,
    label       = "",
    xlim        = (-0.02,3.02),
    xticks      = false,
    yticks      = (0:6:12),
    ylim        = (-0.12,12.12),
    lw          = 2, tickfont = font(7),
    size        = (650,300),
    grid        = :off, framestyle  = :box)

# ### [Measure Performance](@id ex1_performance)
# Benchmark the solution for random values of $\phi_2$:
using BenchmarkTools;
@benchmark Drip(κ,β,[0.9-ϕ2 ϕ2; 1.0 0.0],Q,H,fcap = true) setup = (ϕ2 = 0.9*rand())

# Benchmark the solution for random values of $κ$:
@benchmark Drip(κ,β,[0.5 0.4; 1.0 0.0],Q,H,fcap = true) setup = (κ = rand())

# ## [Ex. 2: AR(3) and ARMA(2,1) Processes](@id ex2)

# Here, we replicate the AR(3) and ARMA(2,1) examples in Mackowiak, Matejka and Wiederholt (2018).
# Using a similar method to the AR(2) case, we can write the law of motion of optimal actions as a state-space form.

# ### [Solve and Meausre Performance](@id ex2_solve)

# #### The AR(3) Example

# Initialize:

β       = 1.0           ;
θ0      = 1.0           ;
κ       = 10.8          ;

ϕ1      = 1.5           ;
ϕ2      = -0.9          ;
ϕ3      = 0.1           ;

A   = [ϕ1 ϕ2 ϕ3 ; 1 0 0 ; 0 1 0] ;
Q   = [θ0 ; 0; 0]   ;
H   = [1  ; 0; 0]   ;

# Solve:

ex2a  = Drip(κ,β,A,Q,H;tol=1e-8);

# Measure performance for different values of $\kappa$:

@benchmark Drip(κ,β,A,Q,H;tol=1e-8) setup = (κ = 10*rand())

# To be consistent with MMW(2018), scale the signal vector so that the weight on the first element is one ($h_1 = 1$):

h1 = 1;
h2 = ex2a.ss.Y[2]/ex2a.ss.Y[1]
h3 = ex2a.ss.Y[3]/ex2a.ss.Y[1]

# Print the weights:

s = @sprintf("  h1 = %5.3f, h2 = %5.3f, h3 = %5.3f", h1, h2, h3)  ;
println(s) ;

# Since we scaled the signal vector, we also need to adjust the noise in the signal accordingly:
AdjNoise   = 3.879
AdjPara_ex2a= AdjNoise*ex2a.ss.Y[1]

# Calculate IRFs:

irf_ex2a = irfs(ex2a; T=30) ;
xirf_ex2a   = irf_ex2a.x[1,1,:]
xhatirf_ex2a= irf_ex2a.x_hat[1,1,:]

x     = zeros(3,30);
xhat_noise_ex2a = zeros(3,30);

for ii in 1:30
    if ii==1
        xhat_noise_ex2a[:,ii] = ex2a.ss.K;
    else
        xhat_noise_ex2a[:,ii] = A*xhat_noise_ex2a[:,ii-1]+(ex2a.ss.K*ex2a.ss.Y')*(x[:,ii]-A*xhat_noise_ex2a[:,ii-1]);
    end
end
xhat_noise_ex2a = xhat_noise_ex2a[1,:]*AdjPara_ex2a ;

# #### The ARMA(2,1) Example:

# Initialize:

β       = 1.0           ;
θ0      = 0.5           ;
θ1      = -0.1          ;
κ       = 1.79          ;

ϕ1      = 1.3           ;
ϕ2      = -0.4          ;

A   = [ϕ1 ϕ2 θ1 ; 1.0 0.0 0.0 ; 0.0 0.0 0.0] ;
Q   = [θ0 ; 0.0; 1.0]   ;
H   = [1.0; 0.0; 0.0]   ;

# Solve:

ex2b = Drip(κ,β,A,Q,H,tol=1e-8);

# Measure performance for different values of $\kappa$:

@benchmark Drip(κ,β,A,Q,H;tol=1e-8) setup = (κ = 2*rand())

# To be consistent with MMW(2018), scale the signal vector so that the weight on the first element is one ($h_1 = 1$):

h1 = 1;
h2 = ex2b.ss.Y[2]/ex2b.ss.Y[1]
h3 = ex2b.ss.Y[3]/ex2b.ss.Y[1]

# Print the weights:

s = @sprintf("  h1 = %5.3f, h2 = %5.3f, h3 = %5.3f", h1, h2, h3)  ;
println(s) ;

# Since we scaled the signal vector, we also need to adjust the noise in the signal accordingly:

AdjNoise    = 1.349
AdjPara_ex2b= AdjNoise*ex2b.ss.Y[1]

# Calculate the IRFs

irf_ex2b     = irfs(ex2b; T=30) ;
xirf_ex2b    = irf_ex2b.x[1,1,:]
xhatirf_ex2b = irf_ex2b.x_hat[1,1,:]

x     = zeros(3,30);
xhat_noise_ex2b = zeros(3,30);

for ii in 1:30
    if ii==1
        xhat_noise_ex2b[:,ii] = ex2b.ss.K;
    else
        xhat_noise_ex2b[:,ii] = A*xhat_noise_ex2b[:,ii-1]+(ex2b.ss.K*ex2b.ss.Y')*(x[:,ii]-A*xhat_noise_ex2b[:,ii-1]);
    end
end
xhat_noise_ex2b = xhat_noise_ex2b[1,:]*AdjPara_ex2b ;

# ### [Replication of Figure (5)](@id ex2_fig5)

p1 = plot(1:30,[xirf_ex2a xhatirf_ex2a xhat_noise_ex2a],
    title       = "AR(3) example, optimal signal with i.i.d. noise",
    color             = [:black :gray20 :gray40], markerstrokecolor = [:black :gray20 :gray40],
    yticks      = (-0.8:0.4:1.6),
    ylim        = (-0.82,1.62))
p2 = plot(1:30,[xirf_ex2b xhatirf_ex2b xhat_noise_ex2b],
    title       = "ARMA(2,1) example, optimal signal with i.i.d. noise",
    color       = [:black :gray20 :gray40], markerstrokecolor = [:black :gray20 :gray40],
    yticks      = (0:0.1:0.6),
    ylim        = (-0.05,0.65))
Plots.plot(p1,p2,
    layout      = (1,2),
    label       = [L"$X_{t}$ to $\varepsilon_t$" L"$Y_t=E[X_t|I_t]$ to $\varepsilon_t$" L"$Y_t=E[X_t|I_t]$ to $\psi_t$"],
    marker      = [:circle :circle :star8], markercolor = [:black :false :gray40], markersize  = [3 5 5],
    legend      = :topright,
    xticks      = (0:2:30),
    xlim        = (0,30),
    lw          = 1.5,
    legendfont  = font(8), guidefont=font(9), titlefont=font(9), tickfont=font(8),
    size        = (700,300),
    grid        = :off, framestyle = :box)

# ## [Ex. 3: Price-setting with Rational Inattention](@id ex3)

# We now replicate the price-setting exercise in MMW (2018) and its comparison with the Woodford (2002) example. This corresponds to Figure (6) in their paper.
# The model structure is identifcal to our [Example 1](https://github.com/afrouzi/DRIPs.jl/blob/master/examples/notebooks/ex1_pricing_pe_nofeedback.ipynb) (without strategic complementarity) and [Example 2](https://github.com/afrouzi/DRIPs.jl/blob/master/examples/notebooks/ex2_pricing_pe_with_feedback.ipynb) (with strategic complementarity).

# ### [The Case with No Strategic Complementarity](@id ex3_no_strcomp)

# Initialize:
ρ   = 0.9;        #persistence of money growth
σ_u = 0.1;          #std. deviation of shocks to money growth

# Primitives of the DRIP:
κ   = 0.62;
β   = 1.0;
A   = [1 ρ; 0 ρ];
Q   = σ_u*[1; 1];
H   = [1; 0];
ω   = 0.1

# Solve:

ex_opt   = Drip(κ,β,A,Q,H,fcap=true);
capa_opt = DRIPs.capacity(ex_opt) #returns capacity utilized in bits

# Measure performance for different values of $\kappa$:

@benchmark Drip(κ,β,A,Q,H,fcap=true) setup = (κ = 2*rand())

# Calculate IRFs:

irfs_ex_opt = irfs(ex_opt, T = 12);
output_opt  = (irfs_ex_opt.x[1,1,:] - irfs_ex_opt.a[1,1,:]) ;
output_opt  = [0;output_opt] ;

# Now, to compare with Woodford (2002), assume that firms observe a noisy signal, $S_t = q_t + \zeta_t$ where $\zeta_t$ is an idiosyncratic noise. We first define a function to solve the corresponding Kalman filtering problem.

function K_filtering(A,Q,Ysignal,Σz,Σ0 ; maxit=10000,tol=1e-10,w=1)
    err = 1
    iter = 0
    while (err > tol) & (iter < maxit)
        global Knew = Σ0*Ysignal*inv(Ysignal'*Σ0*Ysignal .+ Σz)
        global Σp_temp = Σ0 - Knew*Ysignal'*Σ0
        global Σ1 = A*Σp_temp*A' + Q*Q'

        err     = norm(Σ1 - Σ0,2)/norm(Σ0,2)
        Σ0 = w*Σ1 + (1-w)*Σ0

        #println("Iteration $iter. Difference: $err")
        iter += 1
    end

    return(Knew,Σ0,Σp_temp)
end;

# Now find the capacity utilized under Woodford’s formulation such that it yields the same information flow as the optimal signal under ratinoal inattention.

Ywoodford = [1;0]
Σ1_init = ex_opt.ss.Σ_1

Σz_new_b  = 0.01
Σz_new_u  = 0.1
Σz_new    = (Σz_new_b+Σz_new_u)/2
for i in 1:10000
    (Knew,Σ1_new,Σp_temp) = K_filtering(A,Q,Ywoodford,Σz_new,Σ1_init;w=0.5)
    capa_woodford = 0.5*log(det(Σ1_new)/det(Σp_temp))/log(2)

    if capa_woodford > capa_opt
        global Σz_new_b  = Σz_new
    else
        global Σz_new_u  = Σz_new
    end
    global Σz_new    = (Σz_new_b+Σz_new_u)/2
    err = abs(capa_woodford - capa_opt)
    if err < 1e-5
        break
    end
end

# Calculate impulse responses under Woodford's formulation
e_k = 1;
x   = zeros(2,12);
xhat= zeros(2,12);
a   = zeros(2,12);

for ii in 1:12
    if ii==1
        x[:,ii]     = Q*e_k;
        xhat[:,ii]  = (Knew*Ywoodford')*(x[:,ii]);
    else
        x[:,ii]     = A*x[:,ii-1];
        xhat[:,ii]  = A*xhat[:,ii-1]+(Knew*Ywoodford')*(x[:,ii]-A*xhat[:,ii-1]);
    end
    a[:,ii]  .= H'*xhat[:,ii];
end

output_woodford  = (x[1,:] - a[1,:]) ;
output_woodford  = [0;output_woodford] ;

# Before plotting the IRFs, we also solve the case with strategic complementarity.

# ### [The Case with Strategic Complementarity](@id ex3_strcomp)

# We now turn to the example with strategic complementarity.
# As in our [Example 2](https://github.com/afrouzi/DRIPs.jl/blob/master/examples/notebooks/ex2_pricing_pe_with_feedback.ipynb), we first define a function to solve the fixed point with endogenous feedback.

function ge_drip(ω,β,A,Q,          #primitives of drip except for H because H is endogenous
                 α,                #strategic complementarity
                 Hq,               #state space rep. of Δq
                 L;                #length of truncation
                 H0       = Hq,    #optional: initial guess for H (Hq is the true solution when α=0)
                 maxit    = 200,   #optional: max number of iterations for GE code
                 tol      = 1e-4)  #optional: tolerance for iterations
    err   = 1;
    iter  = 0;
    M     = [zeros(1,L-1) 0; Matrix(I,L-1,L-1) zeros(L-1,1)];
    while (err > tol) & (iter < maxit)
        if iter == 0
            global ge  = Drip(ω,β,A,Q,H0; w=0.9);
        else
            global ge  = Drip(ω,β,A,Q,H0; Ω0=ge.ss.Ω, Σ0=ge.ss.Σ_1, maxit=100);
        end

        XFUN(jj) = ((I-ge.ss.K*ge.ss.Y')*ge.A)^jj * (ge.ss.K*ge.ss.Y') * (M')^jj
        X = DRIPs.infinitesum(XFUN; maxit=L, start = 0);  #E[x⃗]=X×x⃗

        XpFUN(jj) = α^jj * X^(jj)
        Xp = DRIPs.infinitesum(XpFUN; maxit=L, start = 0);

        H1 = (1-α)*Xp'*Hq;
        err= 0.5*norm(H1-H0,2)/norm(H0)+0.5*err;
        H0 = H1;

        iter += 1;
        if iter == maxit
            print("GE loop hit maxit\n")
        elseif mod(iter,10) == 0
            println("Iteration $iter. Difference: $err")
        end

    end
    print(" Iteration Done.\n")
    return(ge)
end;

# Now, we solve for the optimal signal structure under rational inattention.

# Initialize:

ρ   = 0.9;        #persistence of money growth
σ_u = 0.1;        #std. deviation of shocks to money growth
α   = 0.85;        #degree of strategic complementarity
L   = 40;         #length of truncation
Hq  = ρ.^(0:L-1); #state-space rep. of Δq

ω   = 0.08;
β   = 1 ;
A   = [1 zeros(1,L-2) 0; Matrix(I,L-1,L-1) zeros(L-1,1)];
M   = [zeros(1,L-1) 0; Matrix(I,L-1,L-1) zeros(L-1,1)]; # shift matrix
Q   = [σ_u; zeros(L-1,1)];

# Solve:

ex_ge   = ge_drip(ω,β,A,Q,α,Hq,L) ;

# Print capacity utilized in with strategic complementarity:

capa_ge = DRIPs.capacity(ge)

# Measure performance for random values of ω

using Suppressor
@suppress @benchmark ge_drip(ω,β,A,Q,α,Hq,L) setup = (ω = 0.1*rand())

# Calculate IRFs
geirfs  = irfs(ex_ge,T = L) ;

dq = diagm(Hq)*geirfs.x[1,1,:];
q  = inv(I-M)*dq ;
output_ge_opt  = q - geirfs.a[1,1,:] ;
output_ge_opt  = [0;output_ge_opt] ;

# Finally, to compare with the IRFs under Woodford (2002)’s specification, we find signal noise such that it yields the same information flow as the optimal signal structure.

Ywoodford_ge = Hq
Σ1_init  = ex_ge.ss.Σ_1

Σz_new_b = 0.05
Σz_new_u = 0.12
Σz_new   = (Σz_new_b+Σz_new_u)/2

for i in 1:10000
    (Knew,Σ1_new,Σp_temp) = K_filtering(A,Q,Ywoodford_ge,Σz_new,Σ1_init;w=0.5)
    capa_woodford = 0.5*log(det(Σ1_new)/det(Σp_temp))/log(2)

    if capa_woodford > capa_ge
        global Σz_new_b  = Σz_new
    else
        global Σz_new_u  = Σz_new
    end
    global Σz_new    = (Σz_new_b+Σz_new_u)/2
    err = abs(capa_woodford - capa_ge)
    if err < 1e-5
        break
    end
end

XFUN(jj) = ((I-Knew*Ywoodford_ge')*A)^jj * (Knew*Ywoodford_ge') * (M')^jj
X = DRIPs.infinitesum(XFUN; maxit=L, start = 0);  #E[x⃗]=X×x⃗

XpFUN(jj) = α^jj * X^(jj)
Xp = DRIPs.infinitesum(XpFUN; maxit=L, start = 0);

H1 = (1-α)*Xp'*Hq;

# Calculate impurse responses under Woodford's signal
e_k = 1;
x   = zeros(40,40);
xhat= zeros(40,40);
a   = zeros(1,40);

for ii in 1:40
    if ii==1
        x[:,ii]     = Q*e_k;
        xhat[:,ii]  = (Knew*Ywoodford_ge')*(x[:,ii]);
    else
        x[:,ii]     = A*x[:,ii-1];
        xhat[:,ii]  = A*xhat[:,ii-1]+(Knew*Ywoodford_ge')*(x[:,ii]-A*xhat[:,ii-1]);
    end
    a[:,ii]  .= H1'*xhat[:,ii];
end
dq = diagm(Hq)*x[1,:];
q  = inv(I-M)*dq ;
output_ge_woodford  = q - a[1,:] ;
output_ge_woodford = [0;output_ge_woodford] ;

# We now have all the IRFs to replicate Figure (6)

# ### [Replication of Figure (6)](@id ex3_fig6)

p1 = plot(0:12,[output_woodford output_opt],
    title       = L"The case of $\xi=1$",
    color       = [:black :gray40], markerstrokecolor = [:black :gray40],
    ylim        = (-0.003,0.053),
    ytick       = (0:0.01:0.05))
p2 = plot(0:12,[output_ge_woodford[1:13] output_ge_opt[1:13]],
    title       = L"The case of $\xi=0.15$",
    color       = [:black :gray40], markerstrokecolor = [:black :gray40],
    ylim        = (-0.005,0.105),
    ytick       = (0:0.02:0.1))
Plots.plot(p1,p2,
    layout      = (2,1),
    label       = ["Woodford Model" "Model with optimal signals"],
    legend      = :topright,
    marker      = [:circle :star8], markersize = [5 5], markercolor = [:false :gray40],
    xlim        = (0,12),
    xtick       = (0:1:12),
    xlabel      = "Time",
    lw          = 1.5,
    legendfont  = font(8), titlefont = font(10), guidefont = font(9), tickfont = font(8),
    size        = (700,400),
    grid        = :off, framestyle = :box)

# ## [Ex. 4: Business Cycle Model with News Shocks](@id ex4)
# In this section, we replicate the business cycle model with news shocks in Section 7 in Mackowiak, Matejka and Wiederholt (2018).
#
# ### [Setup](@id ex4_setup)
#
# #### Full-Information
# The techonology shock follows AR(1) process:
# ```math
# \begin{aligned}
# z_{t} = \rho z_{t-1} + \sigma \varepsilon_{t-k}
# \end{aligned}
# ```
# and the total labor input is:
# ```math
# \begin{aligned}
# n_{t} = \int_0^1 n_{i,t} di .
# \end{aligned}
# ```
#
# Under perfect information, the households chooses the utility-maximizing labor supply, all firms choose the profit-maximizing labor input, and the labor market clearing condition is:
# ```math
# \begin{aligned}
# \frac{1-\gamma}{\psi + \gamma}w_{t} = \frac{1}{\alpha}(z_t - w_t).
# \end{aligned}
# ```
# Then, the market clearing wages and the equilibrium labor input are:
# ```math
# \begin{aligned}
#     w_{t} & = \frac{\frac{1}{\alpha}}{\frac{1-\gamma}{\psi+\gamma} + \frac{1}{\alpha}} z_t \equiv \xi z_t \\
#     n_t & = \frac{1}{\alpha}(1-\xi) z_t.
# \end{aligned}
# ```
#
#
# #### Rational Inattention
#
# Firms wants to keep track of their ideal price,
# ```math
# \begin{aligned}
#     n_{t}^* = \frac{1}{\alpha} z_t - \frac{1}{\alpha} \frac{\psi + \gamma}{1 - \gamma} n_t
# \end{aligned}
# ```
# where $n_{t} = \int_0^1 n_{i,t} di$. Then, firm $i$'s choice depends on its information set at time $t$:
# ```math
# \begin{aligned}
#     n_{i,t} = E_{i,t} [n_{t}^*].
# \end{aligned}
# ```
#
# Note that now the state space representation for $n_{t}^*$ is determined in the equilibrium. However, we know that this is a Guassian process and by Wold's theorem we can decompose it to its $MA(\infty)$ representation:
# ```math
# \begin{aligned}
#     n_{t}^*=\Phi(L)\varepsilon_t
# \end{aligned}
# ```
# where $\Phi(.)$ is a lag polynomial and $\varepsilon_t$ is the shock to technology. Here, we have basically guessed that the process for $p_{i,t}^*$ is determined uniquely by the history of monetary shocks which requires that rational inattention errors of firms are orthogonal (See [Afrouzi (2020)](http://www.afrouzi.com/strategic_inattetion.pdf)). Our objective is to find $\Phi(.)$.
#
# Now, as in our [Example 2](https://github.com/afrouzi/DRIPs.jl/blob/master/examples/notebooks/ex2_pricing_pe_with_feedback.ipynb), we can represent the problem in a matrix notation.
#
# ### [Initialization](@id ex4_initialize)

β   = 1     ;   # Time preference
γ   = 1/3   ;   # Inverse of intertemporal elasticity of substitution
ψ   = 0     ;   # Inverse of Frisch elasticity
α   = 3/4   ;   # Labor share in production function
θ   = -1/α*(ψ+γ)/(1-γ)  ;
ξ   = θ/(θ-1)   ;

ρ   = 0.9   ;   #persistence of technology shocks
σ   = 1     ;   #std. deviation of technology shocks
ω   = 6.5   ;   # Information cost

L   = 40    ;   #length of truncation
k   = 8     ;   #news horizon

M   = [zeros(1,L-1) 0; Matrix(I,L-1,L-1) zeros(L-1,1)]; # shift matrix
Hz  = ρ.^(0:L-1)
Hz = (M^k)*Hz

A   = M     ;
Q   = [σ; zeros(L-1,1)] ;

# Also, define a function that solves the GE problem and returns the solution in a `Drip` structure:

function ge_drip(ω,β,A,Q,          #primitives of drip except for H because H is endogenous
                 α,                #strategic complementarity
                 θ,
                 Hz,               #state space rep. of z
                 L;                #length of truncation
                 w_out    = 0.5,    #optional: initial guess for H (Hq is the true solution when α=0)
                 H0       = Hz,    #optional: initial guess for H (Hq is the true solution when α=0)
                 maxit    = 200,   #optional: max number of iterations for GE code
                 tol      = 1e-6)  #optional: tolerance for iterations
    err   = 1;
    iter  = 0;
    M     = [zeros(1,L-1) 0; Matrix(I,L-1,L-1) zeros(L-1,1)];
    eye   = Matrix(I,L,L)
    while (err > tol) & (iter < maxit)
        if iter == 0
            global ge  = Drip(ω,β,A,Q,H0;w=0.5, tol=1e-8);
        else
            global ge  = Drip(ω,β,A,Q,H0;w=0.9, tol=1e-8, Ω0=ge.ss.Ω, Σ0=ge.ss.Σ_1, maxit=1000);
        end

        XFUN(jj) = ((eye-ge.ss.K*ge.ss.Y')*ge.A)^jj * (ge.ss.K*ge.ss.Y') * (M')^jj
        X = DRIPs.infinitesum(XFUN; maxit=L, start = 0);  #E[x⃗]=X×x⃗

        H1 = (1/α)*Hz + θ*X'*H0 ;

        err= 0.5*norm(H1-H0,2)/norm(H0)+0.5*err;

        H0 = w_out*H1 + (1.0-w_out)*H0 ;

        iter += 1;
        if iter == maxit
            print("GE loop hit maxit\n")
        elseif mod(iter,10) == 0
            println("Iteration $iter. Difference: $err")
        end

    end
    print(" Iteration Done.\n")
    return(ge)
end;

# ### [Solve and Measure Performance](@id ex4_solve)

# Solve:

ge      = ge_drip(ω,β,A,Q,α,θ,Hz,L)  ;

# Measure performance by solving the model for different values of ω:

@suppress @benchmark ge_drip(ω,β,A,Q,α,θ,Hz,L) setup = (ω = 6.5*rand())

# Calculate IRFs and profit loss:

geirfs  = irfs(ge,T = L)         ;
profit_loss = sum((geirfs.x[1,1,:]/100 - geirfs.x_hat[1,1,:]/100).^2) ;

s = @sprintf("==: Profit loss from rational inattention = %6.5f", profit_loss)  ;
println(s) ;

# ### [Replication of Figure (7)](@id ex4_fig7)

n_opt   = geirfs.a[1,1,:]   ; # Optimal labor input under rational inattention
n_fullinfo = σ*1/α*(1-ξ)*Hz ; # Optimal labor input under full information

plot(1:30,[n_fullinfo[1:30] n_opt[1:30]],
    title       = "The impulse response of labor input to a productivity shock.",
    ylabel      = "Percent",
    label       = ["Equilibrium under perfect information" "Equilibirum when firms are subject to rational inattention"],
    legend      = :topright,
    color       = [:black :gray40], markerstrokecolor = [:black :gray40],
    marker      = [:circle :star8], markercolor = [:false :gray40], markersize = [5 5],
    ylim        = (-0.05,0.85),
    ytick       = (0:0.1:0.8),
    xlim        = (0,30),
    xtick       = (0:2:30),
    lw          = 1.5,
    legendfont  = font(8), titlefont = font(10), tickfont = font(8), guidefont = font(9),
    size        = (650,400),
    grid        = :off, framestyle = :box)


#=
julia> fib(n) = n < 3 ? 1 : fib(n-1) + fib(n-2)

julia> fib.(1:10)
10-element Array{Int64,1}:
  1
  1
  2
  3
  5
  8
 13
 21
 34
 55
 =#

function fib(n)
    if n < 3
        return (result = 1, counter = 1)
    else
        result1, counter1 = fib(n - 1)
        result2, counter2 = fib(n - 2)
        return (result = result1 + result2, counter = 1 + counter1 + counter2)
    end
end

# How many times does the function get called?
#=
julia> fib(6)
(result = 8, counter = 15)

julia> fib(10)
(result = 55, counter = 109)

julia> fib(20)
(result = 6765, counter = 13529)
=#

# How well does it perform?
#=
julia> @btime fib(40)
  569.906 ms (0 allocations: 0 bytes)
(result = 102334155, counter = 204668309)
=#

const fib_cache = Dict()

_fib(n) = n < 3 ? 1 : fib(n-1) + fib(n-2)

function fib(n)
    if haskey(fib_cache, n)
        return fib_cache[n]
    else
        value = _fib(n)
        fib_cache[n] = value
        return value
    end
end

# Benchmark performance. Make sure cache is emptied at each
#=
julia> @btime fib(40) setup=(empty!(fib_cache))
  38.731 ns (0 allocations: 0 bytes)
102334155
=#


if haskey(ENV, "CI")
    ENV["PLOTS_TEST"] = "true"
    ENV["GKSwstype"] = "100" # gr segfault workaround
end

using FluxOptTools, Optim, Zygote, Flux, Plots, Test, Statistics, Random
##
@testset "FluxOptTools" begin
@info "Testing FluxOptTools"
@testset "copy" begin
@info "Testing copy"

m = Chain(Dense(1,5,tanh), Dense(5,5,tanh) , Dense(5,1))
x = collect(LinRange(-pi,pi,100)')
y = sin.(x)
sp = sortperm(x[:])

loss() = mean(abs2, m(x) .- y)
Zygote.refresh()
pars = Flux.params(m)
pars0 = deepcopy(pars)
npars = veclength(pars)
@test npars == 46

copy!(pars, zeros(pars))
@test all(all(iszero, p) for p in pars)

p = zeros(pars)
copy!(pars, 1:npars)
copy!(p, pars)
@test p == 1:npars

grads = Zygote.gradient(loss, pars)
grads0 = deepcopy(grads)

copy!(grads, zeros(grads))
@test all(all(iszero,grads[k]) for k in keys(grads.grads))

p = zeros(grads)
copy!(grads, 1:npars)
copy!(p, grads)
@test p == 1:npars
end

## Test optimization ============================================

end

# NOTE: tests below fail if they are in a testset, probably Zygote's fault

m = Chain(Dense(1,5,tanh), Dense(5,5,tanh) , Dense(5,1))
x = collect(LinRange(-pi,pi,100)')
y = sin.(x)
sp = sortperm(x[:])

loss() = mean(abs2, m(x) .- y)
@show loss()
Zygote.refresh()
pars = Flux.params(m)

opt = ADAM(0.01)
@show loss()
for i = 1:500
    grads = Zygote.gradient(loss, pars)
    Flux.Optimise.update!(opt, pars, grads)
end
@show loss()
@test loss() < 1e-1
plot(x[sp], [y[sp] m(x)[sp]])  |> display
plot(loss, pars, l=0.5, npoints=50, seriestype=:contour) |> display

lossfun, gradfun, fg!, p0 = optfuns(loss, pars)
res = Optim.optimize(Optim.only_fg!(fg!), p0, BFGS())
@test loss() < 1e-3
plot(loss, pars, l=0.1, npoints=50) |> display

plot(x[sp], [y[sp] m(x)[sp]]) |> display

## Benchmark Optim vs ADAM
losses_adam = map(1:10) do i
    @show i
    Random.seed!(i)
    m = Chain(Dense(1,5,tanh), Dense(5,5,tanh) , Dense(5,1))
    x = collect(LinRange(-pi,pi,100)')
    y = sin.(x)
    loss() = mean(abs2, m(x) .- y)
    Zygote.refresh()
    pars = Flux.params(m)
    opt = Flux.ADAM(0.2)
    trace = [loss()]
    for i = 1:500
        l,back = Zygote.pullback(loss, pars)
        push!(trace, l)
        grads = back(l)
        Flux.Optimise.update!(opt, pars, grads)
    end
    trace
end

res_lbfgs = map(1:10) do i
    @show i
    Random.seed!(i)
    m = Chain(Dense(1,5,tanh), Dense(5,5,tanh) , Dense(5,1))
    x = LinRange(-pi,pi,100)'
    y = sin.(x)
    loss() = mean(abs2, m(x) .- y)
    Zygote.refresh()
    pars = Flux.params(m)
    lossfun, gradfun, fg!, p0 = optfuns(loss, pars)
    res = Optim.optimize(Optim.only_fg!(fg!), p0, BFGS(), Optim.Options(iterations=500, store_trace=true))
    res
end

losses_SLBFGS = map(1:10) do i
    @show i
    Random.seed!(i)
    m = Chain(Dense(1,5,tanh), Dense(5,5,tanh) , Dense(5,1))
    x = LinRange(-pi,pi,100)'
    y = sin.(x)
    loss() = mean(abs2, m(x) .- y)
    Zygote.refresh()
    pars = Flux.params(m)
    lossfun, gradfun, fg!, p0 = optfuns(loss, pars)
    opt = SLBFGS(lossfun,p0; m=3, ᾱ=1., ρ=false, λ=.0001, κ=0.1)
    function train(opt, p0, iters=20)
        p = copy(p0)
        g = zeros(veclength(pars))
        trace = [loss()]
        for i = 1:iters
            g = gradfun(g,p)
            p = apply(opt, g, p)
            push!(trace, opt.fold)
        end
        trace
    end
    trace = train(opt,p0, 500)
end

##
valuetrace(r) = getfield.(r.trace, :value)
valuetraces = valuetrace.(res_lbfgs)
plot(valuetraces, yscale=:log10, xscale=:identity, lab="", c=:red)
plot!(losses_adam, lab="", c=:blue, xlabel="Epochs", ylabel="Loss")
plot!(losses_SLBFGS, lab="", c=:green)


sp = SplineParams(linspace(0, 5, 10), 0, 3)
cp = ChebParams(6, 2.0, 5.0)
lp = LinParams(linspace(0, 5, 10), 0)

@testset "Test subtype structure" begin
    @testset "Basis families" begin
        for T in [Cheb, Lin, Spline]
            @test T <: BasisFamily
        end
    end

    @testset "params" begin
        for T in [ChebParams, SplineParams, LinParams]
            @test T <: BasisParams
        end
    end

    @testset "basis structure representations" begin
        for T in [Tensor, Direct, Expanded]
            @test T <: AbstractBasisMatrixRep
        end
    end
end

@testset "test _param method" begin
    for (T, TP) in [(Cheb, ChebParams), (Lin, LinParams), (Spline, SplineParams)]
        @test BasisMatrices._param(T) == TP
        @test BasisMatrices._param(T()) == TP
    end
end

@testset "test more Param methods" begin

    @testset "Test extra outer constructors" begin
        @test sp == SplineParams(10, 0, 5, 3)
        @test lp == LinParams(10, 0, 5)
    end

    @testset "test equality of params" begin
        @test (  ==(sp, lp) ) == false
        @test (  ==(sp, cp) ) == false
        @test (  ==(cp, lp) ) == false

        @test   ==(sp, sp)
        @test   ==(cp, cp)
        @test   ==(lp, lp)
    end
end


using LinearAlgebra, StatsBase, CoupledFields
# using Base.Test

# write your own tests here
# @test 1 == 1

function simfields(t::Vector{Float64}, p::Int64, σₑ::Float64)
    n = length(t)
    Xa = sin.(2π*repeat(t, inner=(1,p))/5.0) * diagm(0=>rand([-1, 1], p))
    Xb = sin.(2π*repeat(t, inner=(1,p))/2.5 .+ 1.5π) * diagm(0=>rand([-1, 1], p))
    Xa = StatsBase.zscore(Xa,1) .+ 0.05*randn(n, p)
    Xb = StatsBase.zscore(Xb,1) .+ σₑ*randn(n, p)
    return Xa, Xb
end

t1 = collect((2001+1/24):(1/12):2015)
lpfield, hpfield = simfields(t1, 10, 1.0)
Z = InputSpace(lpfield, hpfield, [1.1,1.1])
kpars = GaussianKP(Z.X)
model = gKCCA([1.0 -5.0 2], Z.X, Z.Y, kpars)


### ChainStandard

# Stores the output of a Monte Carlo run
type ChainStandard{N<:Number,T<:AbstractFloat,V<:AbstractVector,A<:AbstractArray} <: AbstractChain
    values::A
    loglikelihood::V
    numburnin::Int
    accepted::Int
    proposed::Int
    runtime::Real
    ChainStandard(v::AbstractArray{N},ll::AbstractVector{T},b::Int,a::Int,p::Int,r::Real) = new(v,ll,b,a,p,r)
end

###Factory function
@inline function _chain{N<:Number,T<:AbstractFloat}(::Type{Val{:standard}},nparas::Int,nsamples::Int,::Type{N},::Type{T};numburnin::Int =0)
    v = zeros(N,_valuestuple(nparas,nsamples))
    ChainStandard{N,T,Vector,Array}(v,Vector{T}(nsamples),numburnin,0,0,0.0)
end

function store!(c::ChainStandard,s::AbstractSample,j::Int)
    nparas = numparas(s)
    @assert numparas(c) == nparas && j <= numsamples(s)
    if c.proposed < numsamples(c)
        c.proposed += 1
        copy!(c.values,(c.proposed-1)*nparas+1,s.values,(j-1)*nparas+1,nparas)
        c.loglikelihood[c.proposed] = s.loglikelihood[j]
    else
        warn("Chain is full. Additional results cannot be stored")
    end
end

function show(io::IO,c::ChainStandard)
  println("ChainStandard with numparas = $(numparas(c)) and numsamples = $(numsamples(c))")
  println("Samples proposed = $(c.proposed), samples accepted = $(c.accepted), acceptance rate = $(c.accepted/c.proposed)")
  println("Total runtime = $(c.runtime)")
  println("Additional fields: :values, :loglikelihood")
end


using GlobalDiffEq, OrdinaryDiffEq, LinearAlgebra, Random
using GlobalDiffEq: vec_expt_full, orth_vec
using Test

Random.seed!(123)

@testset "GlobalDiffEq Tests" begin

@testset "Vector Expectation" begin

  @testset "Vector Expectation Full" begin
    @test vec_expt_full(2) ≈ 2/π
  end;

  @testset "Random Orthogonal Vectors" begin
    function compute_vec_series(k, n)
      for _ ∈ 1:k
        z₁, z₂ = orth_vec(n)
        @test norm(z₁) ≈ 1
        @test norm(z₂) ≈ 1
        @test dot(z₁, z₂) ≈ 0 atol=1e-10
      end
    end
    k = 100 # number of random iters
    n = 10 # dimension of the vector
    compute_vec_series(k, n)
  end;
end;

@testset "Global Pendulum" begin
  l = 1.0                             # length [m]
  m = 1.0                             # mass[m]
  g = 9.81                            # gravitational acceleration [m/s²]

  function pendulum!(du,u,p,t)
      du[1] = u[2]                    # θ'(t) = ω(t)
      du[2] = -3g/(2l)*sin(u[1]) + 3/(m*l^2)*p(t) # ω'(t) = -3g/(2l) sin θ(t) + 3/(ml^2)M(t)
  end

  θ₀ = 0.01                           # initial angular deflection [rad]
  ω₀ = 0.0                            # initial angular velocity [rad/s]
  u₀ = [θ₀, ω₀]                       # initial state vector
  tspan = (0.0,10.0)                  # time interval

  M = t->0.1sin(t)                    # external torque [Nm]

  prob = ODEProblem(pendulum!,u₀,tspan,M)

  v0 = solve(prob, Tsit5(), dt=0.1, reltol=1e-12, abstol=0).(1:10)

  @testset "Global Pendulum Richardson" begin
    ve = solve(prob, GlobalRichardson(SSPRK33()), dt=0.2, reltol=1e-12, abstol=0).(1:10)
    @test norm(ve - v0) / norm(v0) < 1e-10
  end;

  @testset "Global Pendulum Adjoint" begin
    # The adjoint tests with pendulum! fail because the parameters are a function rather than a constant term
    sol_adjoint = solve(prob, GlobalAdjoint(SSPRK33()), dt=0.2, reltol=1e-4, abstol=1e-4)
    ve_adjoint = sol_adjoint.(1:10)
    @test norm(ve_adjoint - v0) / norm(v0) < 1e-4
  end;
end;

@testset "Global Lotka Volterra" begin
  function Lotka_Volterra!(du,u,p,t)
    du[1] = p[1]*u[1] - p[2]*u[1]*u[2]
    du[2] = -p[3]*u[2] + u[1]*u[2]
  end

  u₀_LV = [1.0;1.0]
  p_LV = [1.5,1.0,3.0]
  tspan_LV = (0.0,10.0)
  prob_LV = ODEProblem(Lotka_Volterra!,u₀_LV,tspan_LV,p_LV)

  v0_LV = solve(prob_LV, Tsit5(), dt=0.1, reltol=1e-12, abstol=0).(1:10)

  @testset "Global Lotka Volterra Richardson" begin
    sol_LV_Richardson = solve(prob_LV, GlobalRichardson(SSPRK33()), dt=0.2, reltol=1e-12, abstol=0)
    ve_LV_Richardson = sol_LV_Richardson.(1:10)
    @test norm(ve_LV_Richardson - v0_LV) / norm(v0_LV) < 1e-10
  end;

  @testset "Global Lotka Volterra Adjoint" begin
    sol_LV_adjoint = solve(prob_LV, GlobalAdjoint(Tsit5()), dt=0.2, reltol=1e-10, abstol=1e-10)
    ve_LV_adjoint = sol_LV_adjoint.(1:10)
    @test norm(ve_LV_adjoint - v0_LV) / norm(v0_LV) < 1e-10
  end;
end;

end;


"""
Linearly interpolates two vector x and y on a linear grid consisting of `nsteps`.
"""
function interp(x::Vector, y::Vector, nsteps::Int)
    # Interpolate
    itp = interpolate((x,), y, Gridded(Linear()))
    # Interpolate at the given resolution
    x_fills = LinRange(minimum(x), maximum(x), nsteps)
    y_fills = itp(x_fills)

    return collect(x_fills), y_fills
end

export interp


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

fint_comp = load_op_and_grad("./build/libFintComp", "fint_comp")
################## End Load Operator ##################

neqs_ = 20
θ = constant(rand(3,3))
A_ = rand(10,3)
fints = A_*θ
el = constant(rand(0:neqs_, 10, 3), dtype=Int32)
neqs = constant(neqs_, dtype=Int32)
# TODO: specify your input parameters
u = fint_comp(fints,el,neqs)
sess = tf.Session()
init(sess)
run(sess, u)
error()

# TODO: change your test parameter to `m`
# gradient check -- v
function scalar_function(θ)
    m = A_*θ
    return sum(tanh(fint_comp(m,el,neqs)))
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand(3,3))
v_ = rand(3,3)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session()
init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


module DifferenceEquations

using ChainRulesCore
using Distributions
using LinearAlgebra
using CommonSolve
using UnPack
using PDMats
using RecursiveArrayTools
using SciMLBase: @add_kwonly, NullParameters, promote_tspan, AbstractRODESolution
using DiffEqBase
using DiffEqBase: __solve

include("matrix_vector_of_vectors.jl")
include("utilities.jl")
include("problems/state_space_problems.jl")
include("solutions/state_space_solutions.jl")
include("solve.jl")
include("algorithms/linear.jl")
include("algorithms/quadratic.jl")

# Exports
export MatrixVectorOfArray

export AbstractStateSpaceProblem, LinearStateSpaceProblem, QuadraticStateSpaceProblem
export StateSpaceSolution, DirectIteration, KalmanFilter

export solve

end # module


import Zygote
import JuMP

using Test: @test, @testset
using PartiallyObservedInverseGames.ForwardOptimalControl: solve_lqr
using PartiallyObservedInverseGames.InverseOptimalControl:
    solve_inverse_lqr, lqr_lagrangian, lqr_lagrangian_grad_x, lqr_lagrangian_grad_u
using LinearAlgebra: I

#=========================================== Forward LQR ===========================================#

T = 100
A = repeat([[
    1 1
    0 1
]], T)
B = repeat([[0, 1][:, :]], T)
Q = I
R = 100I
x0 = [10.0, 10.0]

forward_converged, forward_solution, forward_model = solve_lqr(A, B, Q, R, x0, T)

@testset "Forward LQR" begin
    @test forward_converged
end

#====================== Inverse LQR as nested constrained optimization problem =====================#

@testset "Inverse LQR" begin
    Q̃ = [
        [
            1//3 0
            0 0
        ],
        [
            0 0
            0 2//3
        ],
    ]
    R̃ = [R]
    inverse_converged, inverse_solution, inverse_model, Q_est, R_est =
        solve_inverse_lqr(forward_solution.x, Q̃, R̃; A, B)

    @testset "Gradient Check" begin
        grad_args = (inverse_solution.x, inverse_solution.u, inverse_solution.λ)
        grad_kwargs = (; Q = Q_est, R = R_est, A, B)

        ∇ₓL_ad, ∇ᵤL_ad = Zygote.gradient(
            (x, u) -> lqr_lagrangian(x, u, inverse_solution.λ; grad_kwargs...),
            inverse_solution.x,
            inverse_solution.u,
        )
        ∇ₓL_manual = lqr_lagrangian_grad_x(grad_args...; grad_kwargs...)
        ∇ᵤL_manual = lqr_lagrangian_grad_u(grad_args...; grad_kwargs...)
        atol = 1e-10

        @test isapprox(∇ₓL_ad, ∇ₓL_manual; atol = atol)
        @test isapprox(∇ₓL_ad[:, 2:end], inverse_solution.∇ₓL; atol = atol)
        @test isapprox(∇ᵤL_ad, ∇ᵤL_manual; atol = atol)
        @test isapprox(∇ᵤL_ad, inverse_solution.∇ᵤL; atol = atol)
    end

    @testset "Solution Sanity" begin
        @test inverse_converged
        @test Q_est[1, 1] / Q_est[2, 2] ≈ Q[1, 1] / Q[2, 2]
        @test Q_est[1, 1] / R_est[1, 1] ≈ Q[1, 1] / R[1, 1]
        @test Q_est[2, 2] / R_est[1, 1] ≈ Q[2, 2] / R[1, 1]
        @test inverse_solution.x ≈ solve_lqr(A, B, Q_est, R_est, x0, T)[2].x
    end
end


#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

# Some image manipulation functions

module Image

export downSample, interpolateFlow, warpMotion, readImage, writeFlo

# Down sample an image to a smaller size
#   We can think of the image as dividing up into a grid of equal sized pixels
#   The pixels of the desired smaller image overlap the pixels of the larger original image
#   We want the value of the output pixel to be the average of the input pixels overlapped weighted by
#   area of overlap
# Pre: size(oi)>=(nw,nh)
function downSample(oi::Matrix{Float32}, nw::Int, nh::Int)
# Returns Matrix{Float32} size=(nh,nw)
  (oh, ow) = size(oi)
  if ow==nw && oh==nh
    return oi
  end
  fx = Float32(ow)/Float32(nw)   # width of old pixels per new pixel
  fy = Float32(oh)/Float32(nh)   # height of old pixels per new pixel
  ni = zeros(Float32, nh, nw)
  for y in 1:nh
    for x in 1:nw
      r = 0.0f0
      fi = Float32(y-1)*fy     # start vertical in old pixels of new pixel
      fli = Float32(y)*fy      # end vertical in old pixels of new pixel
      li = round(Int, floor(fli))+1  # end vertical index in old pixels of new pixel
      fj = Float32(x-1)*fx     # start horozontal in old pixels of new pixel
      flj = Float32(x)*fx      # end horozontal in old pixels of new pixel
      lj = round(Int, floor(flj))+1  # end horozontal index in old pixels of new pixel
      for i in round(Int, floor(fi))+1:min(li, oh)
        for j in round(Int, floor(fj))+1:min(lj, ow)
          c = 1.0f0  # weight each full overlapped old pixel by one, adjusted below
          if i-fi<1  # adjust for start vertical overlap
            c = i-fi
          end
          if j-fj<1  # adjust for start horozontal overlap
            c *= j-fj
          end
          if i==li   # adjust for end vertical overlap
            c *= fli-li+1
          end
          if j==lj   # adjust for end horozontal overlap
            c *= flj-lj+1
          end
          r += c*oi[i,j]
        end
      end
      ni[y,x] = r/(fx*fy)  # adjust for total area of old pixels overlapped by new pixel
    end
  end
  return ni
end

# Linearly interpolate between array values
# Pre: 1<=x<=size(a,2) && 1<=y<=size(a,1) && 1<=j<=size(a,3)
function interpolate(a::Array{Float32, 3}, x::Float32, y::Float32, j)
# Returns Float32
  xx = round(Int, floor(x))
  yy = round(Int, floor(y))
  alpha = x-xx
  beta = y-yy
  if alpha>0.0f0
    if beta>0.0f0
      return        beta *(alpha*a[yy+1,xx+1,j]+(1.0f0-alpha)*a[yy+1,xx,j])+
             (1.0f0-beta)*(alpha*a[yy,  xx+1,j]+(1.0f0-alpha)*a[yy,  xx,j])
    else
      return alpha*a[yy,xx+1,j]+(1.0f0-alpha)*a[yy,xx,j]
    end
  else
    if beta>0.0f0
      return beta*a[yy+1,xx,j]+(1.0f0-beta)*a[yy,xx,j]
    else
      return a[yy,xx,j]
    end
  end
end

# Linearly interpolate between array values
# Pre: 1<=x<=size(a,2) && 1<=y<=size(a,1)
function interpolate(a::Matrix{Float32}, x::Float32, y::Float32)
# Returns Float32
  xx = round(Int, floor(x))
  yy = round(Int, floor(y))
  alpha = x-xx
  beta = y-yy
  if alpha>0.0f0
    if beta>0.0f0
      return        beta *(alpha*a[yy+1,xx+1]+(1.0f0-alpha)*a[yy+1,xx])+
             (1.0f0-beta)*(alpha*a[yy,  xx+1]+(1.0f0-alpha)*a[yy,  xx])
    else
      return alpha*a[yy,xx+1]+(1.0f0-alpha)*a[yy,xx]
    end
  else
    if beta>0.0f0
      return beta*a[yy+1,xx]+(1.0f0-beta)*a[yy,xx]
    else
      return a[yy,xx]
    end
  end
end

# Given a flow at a coarser grain, interpolate it to a finer grain
# Pre: size(ou,1)<=nh && size(ou,2)<=nw && size(ou,3)=2
function interpolateFlow(ou::Array{Float32, 3}, nw::Int, nh::Int)
# Returns Array{Float32, 3} size=(nh,nw,2)
  oh = size(ou, 1)
  ow = size(ou, 2)
  if ow==nw && oh==nh
    return ou
  end
  nu = zeros(Float32, nh, nw, 2)
  s = Float32(Float64(nw)/Float64(ow))
  for y in 1:nh
    for x in 1:nw
      nx = Float32(x-1)/Float32(nw)*Float32(ow)+1
      ny = Float32(y-1)/Float32(nh)*Float32(oh)+1
      if nx>=ow
        nx = Float32(ow)
      end
      if ny>=oh
        ny = Float32(oh)
      end
      nu[y,x,1] = interpolate(ou, nx, ny, 1)/s
      nu[y,x,2] = interpolate(ou, nx, ny, 2)/s
    end
  end
  return nu
end

# Given a flow and an image, move the image according to the flow
# Pre: size(i,1)==size(u,1)==size(ii,1) && size(i,2)=size(u,2)==size(ii,2) && size(u,3)=2
function warpMotion(i::Matrix{Float32}, u::Array{Float32, 3}, ii::Matrix{Float32})
# Returns Matrix{Float32} size=size(i)
  (h, w) = size(i)
  wi = zeros(Float32, h, w)
  for y in 1:h
    for x in 1:w
      nx = x+u[y,x,1]
      ny = y+u[y,x,2]
      wi[y,x] = (nx<1 || nx>w || ny<1 || ny>h ? ii[y,x] : interpolate(i, nx, ny))
    end
  end
  return wi
end

function readImage(fn::AbstractString)
  f = open(fn, "r")
  w = read(f, UInt32)
  h = read(f, UInt32)
  img = zeros(Float32, h, w)
  for i in 1:h
    for j in 1:w
      img[i,j] = Float32(read(f, UInt8))/255.0
    end
  end
  close(f)
  return img
end

function writeFlo(u::Array{Float32, 3}, fn::AbstractString)
  h = size(u, 1)
  w = size(u, 2)
  f = open(fn, "w")
  write(f, w)
  write(f, h)
  for i in 1:h
    for j in 1:w
      write(f, u[i, j, 1])
      write(f, u[i, j, 2])
    end
  end
  close(f)
end

end


using Rotations
using CoordinateTransformations
using StaticArrays
using LinearAlgebra

struct Quaternion{T}
    w::T
    x::T
    y::T
    z::T
end

Quaternion(v::AbstractVector) = Quaternion(v[1],v[2],v[3],v[4])
Quaternion(w::Real, v::AbstractVector) = Quaternion(w,v[1],v[2],v[3])
Quaternion(q::Quat) = Quaternion(q.w, q.x, q.y, q.z)
Quaternion(r::Rotations.Rotation{3,T}) where T = Quaternion(Quat(r))
CoordinateTransformations.Quat(q::Quaternion) = Quat(SVector(q)...)
Base.vec(q::Quaternion) = SVector(q.x, q.y, q.z)
StaticArrays.SVector(q::Quaternion) = SVector(q.w, q.x, q.y, q.z)
scalar(q::Quaternion) = q.w

function Base.:*(q2::Quaternion, q1::Quaternion)
    w = scalar(q1)*scalar(q2) - vec(q1)'vec(q2)
    v = scalar(q1)*vec(q2) + scalar(q2)*vec(q1) + vec(q2) × vec(q1)
    Quaternion(w,v)
end

Base.:*(A::AbstractMatrix, q::Quaternion) = A*SVector(q)
Base.:*(q::Quaternion, A::AbstractMatrix) = SVector(q)'A
Base.:*(a::Real, q::Quaternion) = Quaternion(a*scalar(q), a*vec(q))
function Base.:*(q::Quaternion, r::AbstractVector)
    @assert length(r) == 3
    # conj(q)*r
    vec(q*Quaternion(zero(eltype(r)),r)*inv(q))
    # s,v = scalar(q), vec(q)
    # [0; v*v'r + s^2*r + 2s*skew(v)*r + skew(v)^2*r]
end

Base.inv(q::Quaternion) = Quaternion(q.w, -q.x, -q.y, -q.z)
LinearAlgebra.normalize(q::Quaternion) = Quaternion(normalize(SVector(q)))

function skew(v::AbstractVector)
    @assert length(v) == 3
    @SMatrix [0   -v[3]  v[2];
              v[3] 0    -v[1];
             -v[2] v[1]  0]
end

function Lmult(q::Quaternion)
    @SMatrix [q.w -q.x -q.y -q.z
              q.x  q.w -q.z  q.y;
              q.y  q.z  q.w -q.x;
              q.z -q.y  q.x  q.w];
end

function Rmult(q::Quaternion)
    @SMatrix [q.w -q.x -q.y -q.z
              q.x  q.w  q.z -q.y;
              q.y -q.z  q.w  q.x;
              q.z  q.y -q.x  q.w];
end

function Base.conj(q::Quaternion)
    s,v = scalar(q), vec(q)
    vhat = skew(v)
    # C = [s^2 + v'v v'vhat; vhat*v v*v' + s^2*I + 2s*vhat + vhat*vhat]
    C = v*v' + s^2*I + 2s*vhat + vhat*vhat
    return C
end

function deriv_conj(q::Quaternion, r::AbstractVector)
    s,v = scalar(q), vec(q)
    j2 = 2(s*I + skew(v))*r
    j3 = v*r' + I*v'r - 2s*skew(r) - skew(v)*skew(r) - skew(skew(v)*r)
    return [j2 j3]
end

function deriv(q::Quaternion, ω::AbstractVector)
    SVector(0.5*q*Quaternion(zero(eltype(ω)), ω))
end


module SlideBlock

export  solveSlideBlockProblemRBFS,
        solveSlideBlockProblemGSU,
        solveSlideBlockProblemGSBF,
        solveSlideBlockProblemGSBF2,
        solveSlideBlockProblemGSAS,
        solveSlideBlockProblemGSAS2

using AIMACore

import AIMACore: search, goal_test, step_cost, actions, result, heuristic
import Base: ==, show, hash

const E_INVALID_INPUT_LENGTH = "The input parameter is of a wrong length"
const E_INVALID_INPUT = "The input parameter is invalid"
const E_INVALID_STATE = "The state is invalid"
const E_INVALID_ACTION = "The action is invalid"

XY(n, N) = mod(n, N) + 1, mod(div(n,N), N) + 1

dist(t1, t2)    = abs(t1[1] - t2[1]) + abs(t1[2] - t2[2])
dist(n1, n2, N) = dist(XY(n1, N), XY(n2, N))

struct Grid{T} <: State
    loc::Vector{Tuple{Int, Int}}
    h::UInt
end

function Grid{T}(v::Vector{Int}) where T
    N = T*T
    @assert length(v) == N E_INVALID_INPUT
    loc = Vector{Tuple{Int, Int}}(undef, N)
    sum = UInt(0)
    for i = 1:N
        n = v[i]
        @assert n < N E_INVALID_INPUT
        loc[n+1] = XY(i-1, T)
        sum += UInt(n)
        sum *= UInt(N)
    end
    return Grid{T}(loc, sum)
end

Grid(T::Int) = Grid{T}([x for x = 0:T*T-1])

dist(g1::Grid{T}, g2::Grid{T}) where T = sum(dist.(g1.loc, g2.loc))
h2(g1::Grid{T}, g2::Grid{T})   where T = count(g1.loc .!= g2.loc)

==(g1::Grid, g2::Grid) = hash(g1) == hash(g2)

hash(g::Grid, h::UInt) = xor(g.h, h)

function convert(::Type{Matrix{Int}}, g::Grid{T}) where T
    N = T*T
    m = Matrix{Int}(undef, T, T)
    for i = 1:N
        t = g.loc[i]
        m[t[2], t[1]] = i - 1
    end
    return m
end

Matrix{Int}(g::Grid{T}) where T = convert(Matrix{Int}, g)

function show(io::IO, g::Grid{T}) where T
    m = Matrix{Int}(g)
    for j = 1:T-1
        show(io, m[j,:])
        println(io, "")
    end
    show(io, m[T,:])
end

struct Move <: Action
    d::Symbol
end
Move(d::AbstractString) = Move(Symbol(d))

mutable struct SlideBlockProblem{SA <: SearchAlgorithm, T} <: Problem
    initial_state::Grid{T}
    goal_state::Grid{T}
    search_algorithm::SA
    h::Function
end

SlideBlockProblem(init::Grid{T},
                  sa::SA, fh::Function) where {SA <: SearchAlgorithm, T} =
    SlideBlockProblem(init, Grid(T), sa, fh)

function result(problem::SlideBlockProblem{SA,T},
                state::Grid{T}, action::Move) where {SA <: SearchAlgorithm, T}
    N = UInt(T*T)
    loc = deepcopy(state.loc)
    x1, y1 = loc[1]
    xn, yn = action.d == :left  ? (x1 - 1, y1) :
             action.d == :right ? (x1 + 1, y1) :
             action.d == :up    ? (x1, y1 - 1) :
             action.d == :down  ? (x1, y1 + 1) :
             throw(ErrorException(E_INVALID_ACTION))

    i = findfirst(x-> x == (xn,yn), loc)
    if i != 0 && 1 <= xn <= T && 1 <= yn <= T
        loc[1] = (xn, yn)
        loc[i] = (x1, y1)
        #Readjusting the hash
        hn = UInt(i-1)*(N ^ (N-UInt(xn-1+T*(yn-1))))
        h1 = UInt(i-1)*(N ^ (N-UInt(x1-1+T*(y1-1))))
        h = state.h - hn + h1
        return Grid{T}(loc, h)
    else
        throw(ErrorException(E_INVALID_STATE))
    end
end

step_cost(problem::SlideBlockProblem, state::Grid, action::Move) = 1

heuristic(problem::SlideBlockProblem, state::Grid) = problem.h(state, problem.goal_state)

goal_test(problem::SlideBlockProblem, state::Grid) = (state == problem.goal_state)

function actions(problem::SlideBlockProblem, state::Grid{T}) where T
    ret=[]
    x1, y1 = state.loc[1]
    x1 != 1 && push!(ret, Move(:left))
    x1 != T && push!(ret, Move(:right))
    y1 != 1 && push!(ret, Move(:up))
    y1 != T && push!(ret, Move(:down))
    return ret
end

function solveSlideBlockProblem(obj::SA, h::Function) where SA
    vinit = [7, 2, 4, 5, 0, 6, 8, 3, 1]
    problem = SlideBlockProblem(Grid{3}(vinit), obj, h)
    path = search(problem)
    path isa Symbol && return path
    ret=[]
    count = 0
    for iter in path
        push!(ret, iter.state)
    end
    return ret
end

GSD = GraphSearchDepth(Grid(3))
GSU = GraphSearchUniformCost(Grid(3))
GSBF = GraphSearchBestFirst(Grid(3))
GSAS = GraphSearchAStar(Grid(3))
RBFS = RecursiveBestFirstSearch(Grid(3))

solveSlideBlockProblemGSU() = solveSlideBlockProblem(GSU, dist)
solveSlideBlockProblemGSBF() = solveSlideBlockProblem(GSBF, dist)
solveSlideBlockProblemGSAS() = solveSlideBlockProblem(GSAS, dist)
solveSlideBlockProblemRBFS() = solveSlideBlockProblem(RBFS, dist)
solveSlideBlockProblemGSAS2() = solveSlideBlockProblem(GSAS, h2)
solveSlideBlockProblemGSBF2() = solveSlideBlockProblem(GSBF, h2)

end


using PyPlot

"""
    mag, pbr, sbr = frespHBF(f, f1, f2, phi, fp, msg)

Compute the frequency response, the passband ripple and the stopband ripple
of a Saramaki HBF. If `msg` is non-null, a plot is made. `fp` is the passband
edge. `phi` is used by designHBF.
"""
function frespHBF(f, f1, f2, phi=1, fp=0.2, msg="")
    if isempty(f)
        f = range(0, 0.5, length=1024)
    end
    if f1 isa Tuple # Presume that f1 is a (val, csd) tuple
        f1 = f1[1]
    end
    if f2 isa Tuple # Presume that f2 is a (val, csd) tuple
        f2 = f2[1]
    end

    Npts = length(f)
    w = 2 * pi * f
    z = exp.(im * w)
    cos_w = real.(z)

    n2 = length(f2)
    F2 = zeros(size(w))
    for i = 1:n2
        F2 = F2 + f2[i] * cos.(w * (2*i - 1))
    end
    F2 = F2 * 2
    mag = evalF1(f1, F2)

    passband = 1:floor(Int, 2*fp*(Npts - 1) + 1)
    stopband = Npts .+ 1 .- passband
    pbr = maximum(abs.(abs.(mag[passband]) .- 1))
    sbr = maximum(abs.(mag[stopband]))

    if !isempty(msg)
        clf()
        gcf().canvas.set_window_title("HBF Frequency Response")

        subplot(211)

        F1 = evalF0(f1, z, phi)
        plot(f, abs.(F1), linestyle="--", linewidth=1)
        plot(f, phi*abs.(F2), linestyle=":", linewidth=1)
        plot(f, abs.(mag), linewidth=1)
        legend(["F1", "F2", "HBF"])
        title(msg)
        grid()
        axis([0, 0.5, 0, 1.1])

        subplot(212)

        plot(f, dbv(mag))
        axis([0, 0.5, -150, 10])
        grid()

        msg = @sprintf("pbr=%.1e", pbr)
        text(0.0, -10, msg, verticalalignment="top")
        msg = @sprintf("sbr=%.0fdB", amp2db(sbr))
        text(0.5, dbv(sbr), msg, horizontalalignment="right", verticalalignment="bottom")

        sleep(0.5) # slow down to see graphs
    end

    return mag, pbr, sbr
end


#L1DistLoss

using LaTeXStrings, LossFunctions, Measures, Latexify

l1fn = "L(r)=r"
l1dist = latexify(l1fn)
Plots.plot(L1DistLoss(), c=:olive, linewidth=2, annotation=(1.5, .4, L"L\left( r \right) = |r|", 18))


"""
    $(SIGNATURES)

Calculate the ΔG'⁰ of formation for all `mets`, as well as the covariance matrix of the
errors.

# Example
```
equilibrator = Equilibrator()

mets = bigg.(["atp", "adp", "amp", "pi", "h2o"])

standard_dgf_prime_mu, standard_dgf_cov = standard_dg_formation(eq, mets)
```
"""
function standard_dg_formation(equilibrator::Equilibrator, mets::Vector{String})

    pymets = [equilibrator.cc.get_compound(cname) for cname in mets]

    met_formations = [equilibrator.cc.standard_dg_formation(met) for met in pymets]
    dgs = [met_formation[1] for met_formation in met_formations]

    sigmas_fin = hcat([met_formation[2] for met_formation in met_formations]...)
    sigmas_inf = hcat([met_formation[3] for met_formation in met_formations]...)

    leg_trans = [
        pymet.transform(
            equilibrator.cc.p_h,
            equilibrator.cc.ionic_strength,
            equilibrator.cc.temperature,
            equilibrator.cc.p_mg,
        ).m_as(
            "kJ/mol",
        ) for pymet in pymets
    ]

    standard_dgf_prime_mu = dgs + leg_trans
    standard_dgf_cov = sigmas_fin' * sigmas_fin + 1e6 .* sigmas_inf' * sigmas_inf

    return (standard_dgf_prime_mu)u"kJ/mol", (standard_dgf_cov)u"kJ/mol"
end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

function fem_advection(u,v,m,n,h)
    fem_advection_ = load_op_and_grad("./build/libFemAdvection","fem_advection", multiple=true)
    u,v,m_,n_,h = convert_to_tensor(Any[u,v,m,n,h], [Float64,Float64,Int64,Int64,Float64])
    ii, jj, vv = fem_advection_(u,v,m_,n_,h)
    SparseTensor(ii+1, jj+1, vv, (m+1)*(n+1), (m+1)*(n+1))
end

m = 10
n = 10
h = 0.1
u = rand(4*m*n)
v = rand(4*m*n)

# TODO: specify your input parameters
U = fem_advection(u,v,m,n,h)
sess = Session(); init(sess)
# @show run(sess, U)

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(x)
    return sum(values(fem_advection(u,x,m,n,h))^2)
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand(4*m*n))
v_ = rand(4*m*n)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")