using TuringModels, StatsFuns

Turing.setadbackend(:reverse_diff)
#nb Turing.turnprogress(false);

d = DataFrame(CSV.read(joinpath(@__DIR__, "..", "..", "data", "reedfrogs.csv"),
  delim=';'));

size(d) |> display # Should be 48x5

# Set number of tanks

d[:tank] = 1:size(d,1);

# Thanks to Kai Xu!

@model m12_2(density, tank, surv) = begin

    # Separate priors on α and σ for each tank
    σ ~ Truncated(Cauchy(0, 1), 0, Inf)
    α ~ Normal(0, 1)

    # Number of unique tanks in the data set
    N_tank = length(tank)

    # vector for the priors for each tank
    a_tank = Vector{Real}(undef, N_tank)

    # For each tank [1,..,48] set a prior. Note the [] around Normal().
    a_tank ~ [Normal(α, σ)]

    # Observation
    logitp = [a_tank[tank[i]] for i = 1:N_tank]
    surv ~ VecBinomialLogit(density, logitp)

end

# Sample

chns = sample(m12_2(Vector{Int64}(d[:density]), Vector{Int64}(d[:tank]),
    Vector{Int64}(d[:surv])), Turing.NUTS(0.8), 1000);

# CmdStan result

m122rethinking = "
                mean   sd  5.5% 94.5% n_eff Rhat
a               1.30 0.25  0.90  1.70 11662    1
sigma       1.62 0.22  1.30  1.99  6556    1
a_tank[1]   2.12 0.88  0.84  3.60 16091    1
a_tank[2]   3.05 1.10  1.52  4.92 10962    1
a_tank[3]   1.00 0.66 -0.02  2.10 18175    1
a_tank[4]   3.05 1.11  1.47  4.96 10181    1
a_tank[5]   2.13 0.87  0.85  3.62 13720    1
a_tank[6]   2.12 0.86  0.86  3.59 11628    1
a_tank[7]   3.07 1.13  1.47  5.03 10315    1
a_tank[8]   2.13 0.87  0.86  3.60 13754    1
a_tank[9]  -0.18 0.60 -1.14  0.76 18218    1
a_tank[10]  2.11 0.86  0.83  3.58 15121    1
a_tank[11]  1.00 0.67 -0.04  2.09 17390    1
a_tank[12]  0.58 0.62 -0.41  1.60 17209    1
a_tank[13]  0.99 0.66 -0.04  2.09 15225    1
a_tank[14]  0.19 0.62 -0.79  1.20 18293    1
a_tank[15]  2.13 0.89  0.83  3.63 12445    1
a_tank[16]  2.11 0.87  0.87  3.61 12385    1
a_tank[17]  2.89 0.80  1.76  4.29 12583    1
a_tank[18]  2.38 0.66  1.43  3.49 14437    1
a_tank[19]  2.00 0.58  1.12  2.99 13959    1
a_tank[20]  3.67 1.03  2.20  5.44 10629    1
a_tank[21]  2.38 0.65  1.42  3.47 15309    1
a_tank[22]  2.39 0.66  1.42  3.49 13614    1
a_tank[23]  2.40 0.67  1.41  3.53 11868    1
a_tank[24]  1.69 0.52  0.90  2.55 18468    1
a_tank[25] -1.00 0.45 -1.74 -0.30 18153    1
a_tank[26]  0.16 0.40 -0.47  0.81 21895    1
a_tank[27] -1.44 0.50 -2.28 -0.69 16718    1
a_tank[28] -0.47 0.41 -1.15  0.17 20160    1
a_tank[29]  0.15 0.40 -0.48  0.80 19401    1
a_tank[30]  1.44 0.49  0.70  2.24 15407    1
a_tank[31] -0.64 0.42 -1.33 -0.01 15356    1
a_tank[32] -0.31 0.40 -0.95  0.32 19130    1
a_tank[33]  3.18 0.78  2.06  4.55 10894    1
a_tank[34]  2.70 0.66  1.75  3.84 13573    1
a_tank[35]  2.69 0.64  1.74  3.78 13876    1
a_tank[36]  2.06 0.53  1.26  2.92 16329    1
a_tank[37]  2.06 0.51  1.29  2.91 14672    1
a_tank[38]  3.88 0.97  2.52  5.57  9349    1
a_tank[39]  2.70 0.64  1.77  3.78 13444    1
a_tank[40]  2.34 0.56  1.49  3.31 14966    1
a_tank[41] -1.82 0.48 -2.61 -1.10 14214    1
a_tank[42] -0.58 0.36 -1.16 -0.02 17203    1
a_tank[43] -0.46 0.35 -1.02  0.08 17762    1
a_tank[44] -0.34 0.34 -0.90  0.20 16740    1
a_tank[45]  0.58 0.35  0.02  1.14 18946    1
a_tank[46] -0.57 0.34 -1.13 -0.03 19761    1
a_tank[47]  2.05 0.51  1.30  2.90 15122    1
a_tank[48]  0.00 0.33 -0.53  0.53 18236    1
";

# Draw summary

describe(chns)

# End of `12/m12.2t.jl`


export wrand, wrandn, wrandnp, wsobol, sobol, soboln, sobolnp, wgrid, xgrid, near

##### random

"""
    wrand(d, k) = Weighted(rand(d, k))
Uniformly distributed vectors in `[0,1]^d`, as columns of a `Weighted{Matrix}`
which knows to `clamp` them to this box. Keyword `weights=true` gives weights ∝ `1 .+ rand(k)` rather than constant.
Default is now `k=1`, making a one-column matrix.

    wrandn(d, k) = π .* Weighted(randn(d, k))
Normally distributed `d`-vectors, of mean zero and std. dev. `scale=π` by default.
Keyword `max=10` clamps absolute values to be less than this.

    wrandnp(d, k)
Absolute value of normally distributed...
"""
wrand(d::Int, k::Int=1; weights=true) = Weighted( rand(d,k), weights ? 1 .+ rand(k) : ones(k) ,0,1)

@doc @doc(wrand)
function wrandn(d::Int, k::Int=1; scale=π, weights=true, max=Inf)
    out = Weighted( scale .* randn(d,k), weights ? 1 .+ rand(k) : ones(k) )
    max==Inf ? out : clamp!(out, -max, max)
end

@doc @doc(wrand)
function wrandnp(d::Int, k::Int=1; scale=π, weights=true, max=Inf)
    out = Weighted( abs.(scale .* randn(d,k)), weights ? 1 .+ rand(k) : ones(k), 0,Inf)
    max==Inf ? out : clamp!(out, 0, max)
end

##### sub-random

import Sobol

"""    sobol(d, k)
First `k` entries from the `d`-dimensional Sobol sequence, as columns of a `Weighted{Matrix}`. """
function sobol(d::Int, k::Int=1)
    s = Sobol.SobolSeq(d)
    cols = hcat([Sobol.next!(s) for i=1:k]...)
    Weighted(cols, ones(k), 0,1)
end

using SpecialFunctions: erfinv

"""    soboln(d, k)
Normally distributed `d`-vectors, of mean zero and std. dev. π, generated from Sobol sequence,
as columns of a `Weighted{Matrix}`. """
function soboln(d::Int, k::Int=1; scale=π)
    boxed = sobol(d,k)
    Weighted( scale .* erfinv.( -1 .+ 2 .*array(boxed)), boxed.weights)
end

function sobolnp(d::Int, k::Int=1; scale=π)
    boxed = sobol(d,k)
    Weighted( scale .* erfinv.(boxed.array), boxed.weights, 0, Inf)
end

##### ordered!

wgrid(d::Int, range::AbstractRange) = Weighted(xgrid(d,range), extrema(range)...)

##### not Weighted

"""    xgrid(d, 0:0.1:5)
Gives the matrix whose colums are `d`-vectors, forming a grid of the given range in all dimensions. """
xgrid(d::Int, n::Int) = xgrid(d, 0:n)

function (xgrid(d::Int, range::AbstractRange{T}; big=false)::Matrix{T}) where T
    d==1 && return Matrix(collect(range)')

    n = length(range)
    !big && n^d > 4_000_000 && @error "xgrid is happiest with n^d = < 4×10^6, you have $n^$d... keyword big=true will disable this check"

    intcols = zeros(Int,d,n^d)
    @inbounds for i=1:n^d
        @views digits!(intcols[:,i], i-1; base=n)
    end

    range[1] .+ step(range) .* intcols
end

function near(x::Matrix, y, dist::Real; verbose=false)
    yesno = minimum( pairwise2(x,array(y)) ,dims=2) .<= (dist^2)
    sum(yesno)==0 && begin @warn("no nearby points, returning complete matrix"); return x end
    verbose && @info string("kept ",sum(yesno)," / ",size(x,2)," columns: ",round(100*sum(yesno)/size(x,2),2)," percent")
    x[:,vec(yesno)]
end

using Distances

"""
    pairwise2(x, y=x) = Distances.pairwise(SqEuclidean(), x, y)
Resulting `mat[i,j]` is distance sqared from `x[:,i]` to `y[:,j]`. Implementation varies.
"""
pairwise2(x::Matrix, y::Matrix) = pairwise(SqEuclidean(), x, y, dims=2)
pairwise2(x::Matrix) = pairwise(SqEuclidean(), x)

pairwise2(x::AbsMat, y::AbsMat) = diag(x'*x) .+ diag(y'*y)' .- 2x'*y
function pairwise2(x::AbsMat) # = diag(x'*x) .+ diag(x'*x)' .- 2x'*x
    mat = x'*x
    vec = diag(mat)
    vec .+ vec' .- 2 .* mat
end

# using Requires
#
# function init_create_req()
#     @require Distances = "???" begin
#
# 	pairwise2(x::Matrix, y::Matrix) = pairwise(SqEuclidean(), x, y)
# 	pairwise2(x::Matrix) = pairwise(SqEuclidean(), x)
#
# 	end # @require
# end


# ---
# title: 152. Maximum Product Subarray
# id: problem152
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Dynamic Programming
# link: <https://leetcode.com/problems/maximum-product-subarray/description/>
# hidden: true
# ---
#
# Given an integer array `nums`, find the contiguous subarray within an array
# (containing at least one number) which has the largest product.
#
# **Example 1:**
#
#
#
#     Input: [2,3,-2,4]
#     Output: 6
#     Explanation:  [2,3] has the largest product 6.
#
#
# **Example 2:**
#
#
#
#     Input: [-2,0,-1]
#     Output: 0
#     Explanation:  The result cannot be 2, because [-2,-1] is not a subarray.
#
#
## @lc code=start
using LeetCode

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


# Wrap everything with a module to allow redefition of type
module MC

"""
Composite type to represent a spin state
"""
mutable struct SpinState
    num_spins::Int
    s::Array{Int8,1}
    energy::Int
    tot_mag::Int
end

"""
Energy
"""
function energy(s)
    n = length(s)
    - sum((s[i] * s[ifelse(i == n, 1, i+1)] for i in 1:n))
end

"""
Total magnetization
"""
total_magnetization(s) = sum(s)

"""
Constructor
"""
function SpinState(s)
    ss = SpinState(length(s), copy(s), energy(s), total_magnetization(s))
    sanity_check(ss)
    ss
end

"""
Sanity check
"""
function sanity_check(ss)
    @assert energy(ss.s) == ss.energy
    @assert total_magnetization(ss.s) == ss.tot_mag
end

"""
Take an object of SpinState as an input and update it in place.
"""
function update!(ss, β, niters, rng)
    min_h = -2
    max_h = 2
    s = ss.s
    n = ss.num_spins
    prob = [1/(1+exp(-2*β*h)) for h in min_h:max_h]
    for iter in 1:niters, i in 1:n
        sl = s[ifelse(i == 1, n, i-1)]
        sr = s[ifelse(i == n, 1, i+1)]
        # h = -2, 0, 2
        h = sl + sr
        si_old = s[i]
        s[i] = ifelse(rand(rng) < prob[h-min_h+1], +1, -1)

        # Update observables with O(1) operations
        ss.energy += (si_old - s[i]) * h
        ss.tot_mag += (s[i] - si_old)
    end
end

end


"""
    AbstractLinRange

Supertype for mutable or static ranges with `len` linearly spaced elements
between its `start` and `stop`.
"""
abstract type AbstractLinRange{T} <: AbstractRange{T} end

Base.firstindex(::AbstractLinRange) = 1

"""
    LinSRange

A static range with len linearly spaced elements between its start and stop.
The size of the spacing is controlled by len, which must be an Int.
"""
struct LinSRange{T,B,E,L,D} <: AbstractLinRange{T}
    function LinSRange{T}(start, stop, len) where T
        len >= 0 || throw(ArgumentError("srange($start, stop=$stop, length=$len): negative length"))
        if len == 1
            start == stop || throw(ArgumentError("srange($start, stop=$stop, length=$len): endpoints differ"))
            return new{T,start, stop, 1, 1}()
        end
        return new{T, T(start), T(stop), len, max(len-1,1)}()
    end
end

function LinSRange(start, stop, len::Integer)
    return LinSRange{typeof((stop-start)/len)}(start, stop, len)
end

function Base.getproperty(r::LinSRange, s::Symbol)
    if s === :start
        return first(r)
    elseif s === :stop
        return last(r)
    elseif s === :len
        return length(r)
    elseif s === :lendiv
        return lendiv(r)
    else
        error("type $(typeof(r)) has no property $s")
    end
end

LinSRange{T}(r::AbstractRange) where {T} = LinSRange{T}(first(r), last(r), length(r))

"""
    LinMRange

A mutable range with len linearly spaced elements between its start and stop.
The size of the spacing is controlled by len, which must be an Int.
"""
mutable struct LinMRange{T} <: AbstractLinRange{T}
    start::T
    stop::T
    len::Int
    lendiv::Int

    function LinMRange{T}(start, stop, len) where T
        len >= 0 || throw(ArgumentError("mrange($start, stop=$stop, length=$len): negative length"))
        if len == 1
            start == stop || throw(ArgumentError("mrange($start, stop=$stop, length=$len): endpoints differ"))
            return new(start, stop, 1, 1)
        end
        return new(start, stop, len, max(len-1,1))
    end
end

function LinMRange(start, stop, len::Integer)
    return LinMRange{typeof((stop-start)/len)}(start, stop, len)
end

LinMRange{T}(r::AbstractRange) where {T}= LinMRange{T}(first(r), last(r), length(r))

function Base.setproperty!(r::LinMRange, s::Symbol, val)
    if s === :start
        return set_first!(r, val)
    elseif s === :stop
        return set_last!(r, val)
    elseif s === :len
        return set_length!(r, val)
    elseif s === :lendiv
        return set_lendiv!(r, val)
    else
        error("type $(typeof(r)) has no property $s")
    end
end

for (F,f) in ((:M,:m), (:S,:s))
    LR = Symbol(:Lin, F, :Range)
    frange = Symbol(f, :range)

    @eval begin
        Base.:(-)(r::$(LR)) = $(LR)(-first(r), -last(r), length(r))

        $(LR){T}(r::$(LR){T}) where {T} = r
        #$(LR){T}(r::AbstractRange) where {T} = $(LR){T}(first(r), last(r), length(r))
        $(LR)(r::AbstractRange{T}) where {T} = $(LR){T}(r)

        function Base.:(-)(r1::$(LR){T}, r2::$(LR){T}) where T
            len = length(r1)
            (len == length(r2) ||
             throw(DimensionMismatch("argument dimensions must match")))
            return $(LR){T}(
                convert(T, -(first(r1), first(r2))),
                convert(T, -(last(r1), last(r2))),
                len
               )
        end

        function Base.:(+)(r1::$(LR){T}, r2::$(LR){T}) where T
            len = length(r1)
            (len == length(r2) ||
             throw(DimensionMismatch("argument dimensions must match")))
            return $(LR){T}(
                convert(T, +(first(r1), first(r2))),
                convert(T, +(last(r1), last(r2))),
                len
               )
        end
    end
end


include("Cosmology.jl")
using Cosmology
using Base.Test

function test_approx_eq_rtol(va, vb, rtol, astr, bstr)
    diff = maximum(abs(va - vb))
    tol = rtol*max(maximum(abs(va)), maximum(abs(vb)))
    if diff > tol
        sdiff = string("|", astr, " - ", bstr, "| <= ", tol)
        error("assertion failed: ", sdiff,
              "\n  ", astr, " = ", va,
              "\n  ", bstr, " = ", vb,
              "\n  difference = ", diff, " > ", tol)
    end
end

macro test_approx_eq_rtol(a, b, c)
    :(test_approx_eq_rtol($(esc(a)), $(esc(b)), $(esc(c)), $(string(a)), $(string(b))))
end

# values from http://icosmos.co.uk/

dist_rtol = 1e-6
age_rtol = 2e-4

c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1651.9145 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3303.829 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 151.0571 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6607.6579 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.4694 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.7527 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.4694-5.7527 age_rtol

c = cosmology(h=0.7, OmegaK=0.1, OmegaM=0.3, OmegaR=0)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1619.9588 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3209.784 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 140.0856 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6479.8352 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.064 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.5466 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.064-5.5466 age_rtol

c = cosmology(h=0.7, OmegaK=-0.1, OmegaM=0.3, OmegaR=0)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1686.5272 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3408.937 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 163.8479 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6746.1088 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.925 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.9868 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.925-5.9868 age_rtol

c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1612.0585 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3224.1169 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 140.3851 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6448.2338 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.1915 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.6464 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.1915-5.6464 age_rtol

c = cosmology(h=0.7, OmegaK=0.1, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1588.0181 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3147.6227 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 132.0466 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6352.0723 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 12.8488 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.4659 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 12.8488-5.4659 age_rtol

c = cosmology(h=0.7, OmegaK=-0.1, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
@test_approx_eq_rtol angular_diameter_dist_mpc(c,1) 1637.5993 dist_rtol
@test_approx_eq_rtol comoving_radial_dist_mpc(c,1) 3307.9932 dist_rtol
@test_approx_eq_rtol comoving_volume_gpc3(c,1) 149.8301 dist_rtol
@test_approx_eq_rtol luminosity_dist_mpc(c,1) 6550.3973 dist_rtol
@test_approx_eq_rtol age_gyr(c,0) 13.5702 age_rtol
@test_approx_eq_rtol age_gyr(c,1) 5.8482 age_rtol
@test_approx_eq_rtol lookback_time_gyr(c,1) 13.5702-5.8482 age_rtol


# MCMC.jl
"""
My MCMC stuff

Currently the only method is gibbs. Eventually, I would
like to add a Metropolis Hastings method.
"""
module MCMC

import Distributions
export gibbs, metropolis
"""
Generalized Gibbs sampler.

# Arguments
* `init`: Initial values for gibbs sampler
* `update`: function which takes a state::State as the only input and returns the update
* `B`: number of MCMC samples desired
* `burn`: burn-in
* `printFreq`: shows progress bar every printFreq%. (default: 0 => don't print progress)

# Example:

```julia
using Distributions

# Generate data
const n = 1000
const (mu,sig2) = (5,2)
const y = rand(Normal(mu,sqrt(sig2)),n)

# precomputes
const ybar = mean(y)
const (sig2_a, sig2_b) = (2.0, 1.0)
rig(shape::Float64,rate::Float64) = 1 / rand(Gamma(shape,1/rate))

# Define State
immutable State
  mu::Float64
  sig2::Float64
end

# define an update function for State
function update(state::State)
  const newMu = rand(Normal(ybar,sqrt(state.sig2/n)))
  const newSig2 = rig(sig2_a+n/2, sig2_b + sum((y-newMu).^2)/2)
  State(newMu,newSig2)
end

# run gibbs
@time out = gibbs(State(0,1),update,10000,1000);

## post processing
const postMu = map(o -> o.mu, out)
const postSig2 = map(o -> o.sig2, out)

```
"""

function gibbs{T}(init::T, update, B::Int, burn::Int; printFreq::Int=0)
  const out = Array{T,1}( (B+burn) )
  out[1] = init
  for i in 2:(B+burn)
    out[i] = update(out[i-1])
    if printFreq > 0 && i % printFreq == 0
      print("\rProgress: ",i,"/",B+burn)
    end
  end
  out[ (burn+1):end ]
end

"""
Univariate Metropolis step (with Normal proposal)

    metropolis(curr, loglike_plus_logprior, candSig::Real;
               inbounds = x-> -Inf<x<Inf)

# Arguments
* `curr`: current value of parameter
* `loglike_plus_logprior`: log-likelihoodd plus log-prior as function of the parameter to be updated
* `candSig`: The sd of the normal proposal
* `inbounds`: A function which checks if the current parameter is in the support
"""
function metropolis(curr, loglike_plus_logprior, candSig::Real;
                    inbounds = x-> -Inf<x<Inf)
  # need to do autotune

  const cand = rand(Distributions.Normal(curr,candSig))

  if inbounds(cand)
    p = loglike_plus_logprior(cand) - loglike_plus_logprior(curr)
    p > log(rand()) ?  cand : curr
  else
    curr
  end

end

end # module


@testset "Simulated Annealing" begin
    Random.seed!(1)

    function f_s(x::Vector)
        (x[1] - 5.0)^4
    end
    options = Optim.Options(iterations=100_000)
    results = Optim.optimize(f_s, [0.0], SimulatedAnnealing(), options)
    @test norm(Optim.minimizer(results) - [5.0]) < 0.1

    function rosenbrock_s(x::Vector)
        (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
    end
    options = Optim.Options(iterations=100_000)
    results = Optim.optimize(rosenbrock_s, [0.0, 0.0], SimulatedAnnealing(), options)
    @test norm(Optim.minimizer(results) - [1.0, 1.0]) < 0.1

    options = Optim.Options(iterations=10, show_trace=true, store_trace=true, extended_trace=true)
    results = Optim.optimize(rosenbrock_s, [0.0, 0.0], SimulatedAnnealing(), options)
end


using Test
using CartesianMesh
#using Revise
using CutCell
include("useful_routines.jl")

mesh = UniformMesh([0.,0.],[1.,1.],[2,2])
@test allequal(CutCell.elements_per_mesh_side(mesh),[2,2])
cellmaps = CutCell.cell_maps(mesh)
cm1 = CutCell.CellMap([0.,0.],[0.5,0.5])
cm2 = CutCell.CellMap([0.,0.5],[0.5,1.])
cm3 = CutCell.CellMap([0.5,0.],[1.,0.5])
cm4 = CutCell.CellMap([0.5,0.5],[1.,1.])
@test allequal(cellmaps,[cm1,cm2,cm3,cm4])

coords = CutCell.nodal_coordinates(mesh.x0,mesh.widths,mesh.nelements,[5,5])
xrange = range(0.0,stop=1.,length=5)
ycoords = repeat(xrange,5)
xcoords = repeat(xrange,inner=5)
@test allapprox(coords,vcat(xcoords',ycoords'))

connectivity = CutCell.nodal_connectivity([5,5],3,9,mesh.nelements)
c1 = [1,2,3,6,7,8,11,12,13]
c2 = [3,4,5,8,9,10,13,14,15]
c3 = [11,12,13,16,17,18,21,22,23]
c4 = [13,14,15,18,19,20,23,24,25]
testconn = hcat(c1,c2,c3,c4)
@test allequal(connectivity,testconn)

cellconnectivity = CutCell.cell_connectivity(mesh)
testcellconn = [0  1  0  3
                3  4  0  0
                2  0  4  0
                0  0  1  2]
@test allequal(testcellconn,cellconnectivity)

mesh = UniformMesh([0.,0.],[1.,1.],[3,2])
connectivity = CutCell.nodal_connectivity([10,7],4,16,mesh.nelements)
testnodeconn = vcat(43:46,50:53,57:60,64:67)
@test allequal(connectivity[:,5],testnodeconn)

cellconnectivity = CutCell.cell_connectivity(mesh)
c1 = [0,3,2,0]
c2 = [1,4,0,0]
c3 = [0,5,4,1]
c4 = [3,6,0,2]
c5 = [0,0,6,3]
c6 = [5,0,0,4]
testconn = hcat(c1,c2,c3,c4,c5,c6)
@test allequal(cellconnectivity,testconn)

mesh = UniformMesh([0.,0.],[1.,1.],[3,2])
femesh = CutCell.Mesh(mesh,16)
testnodeconn = vcat(43:46,50:53,57:60,64:67)
@test allequal(testnodeconn,femesh.nodalconnectivity[:,5])
@test CutCell.nodes_per_element(femesh) == 16
@test CutCell.nodes_per_mesh_side(femesh) == [10,7]
@test CutCell.number_of_nodes(femesh) == 70

@test allequal(CutCell.bottom_boundary_node_ids(femesh),1:7:64)
@test allequal(CutCell.right_boundary_node_ids(femesh),64:70)
@test allequal(CutCell.top_boundary_node_ids(femesh),7:7:70)
@test allequal(CutCell.left_boundary_node_ids(femesh),1:7)

mesh = CutCell.Mesh([0.,0.],[4.,3.],[4,3],9)
isboundarycell = CutCell.is_boundary_cell(mesh)
testboundarycell = ones(Int,12)
testboundarycell[5] = testboundarycell[8] = 0
@test allequal(isboundarycell,testboundarycell)

isboundarycell = CutCell.is_boundary_cell(mesh.cellconnectivity)
@test allequal(isboundarycell,testboundarycell)

mesh = CutCell.Mesh([2.,1.],[2.,1.],[3,2],4)
@test CutCell.cell_id(mesh,[3.,1.25]) == 3
@test CutCell.cell_id(mesh,[10/3+eps(),1.5+eps()]) == 6


@testset "kernelproduct" begin
    rng = MersenneTwister(123456)
    v1 = rand(rng, 3)
    v2 = rand(rng, 3)

    k1 = LinearKernel()
    k2 = SqExponentialKernel()
    k3 = RationalQuadraticKernel()

    k = KernelProduct([k1, k2])
    @test length(k) == 2
    @test k(v1, v2) == (k1 * k2)(v1, v2)
    @test (k * k)(v1, v2) ≈ k(v1, v2)^2
    @test (k * k3)(v1, v2) ≈ (k3 * k)(v1, v2)

    @testset "kernelmatrix" begin
        rng = MersenneTwister(123456)

        Nx = 5
        Ny = 4
        D = 3

        w1 = rand(rng) + 1e-3
        w2 = rand(rng) + 1e-3
        k1 = SqExponentialKernel()
        k2 = LinearKernel()
        k = k1 * k2

        @testset "$(typeof(x))" for (x, y) in [
            (ColVecs(randn(rng, D, Nx)), ColVecs(randn(rng, D, Ny))),
            (RowVecs(randn(rng, Nx, D)), RowVecs(randn(rng, Ny, D))),
        ]
            @test kernelmatrix(k, x, y) ≈ kernelmatrix(k1, x, y) .* kernelmatrix(k2, x, y)

            @test kernelmatrix(k, x) ≈ kernelmatrix(k1, x) .* kernelmatrix(k2, x)

            K_diag_manual = kerneldiagmatrix(k1, x) .* kerneldiagmatrix(k2, x)
            @test kerneldiagmatrix(k, x) ≈ K_diag_manual

            tmp = Matrix{Float64}(undef, length(x), length(y))
            @test kernelmatrix!(tmp, k, x, y) ≈ kernelmatrix(k, x, y)

            tmp_square = Matrix{Float64}(undef, length(x), length(x))
            @test kernelmatrix!(tmp_square, k, x) ≈ kernelmatrix(k, x)

            tmp_diag = Vector{Float64}(undef, length(x))
            @test kerneldiagmatrix!(tmp_diag, k, x) ≈ kerneldiagmatrix(k, x)
        end
    end
    test_ADs(x->SqExponentialKernel() * LinearKernel(c= x[1]), rand(1), ADs = [:ForwardDiff, :ReverseDiff])
    @test_broken "Zygote issue"
end


using OrdinaryDiffEq, DiffEqBase, Base.Test, DiffEqProblemLibrary, DiffEqDevTools
srand(100)
u = [0.5,0.2]
prob = prob_ode_2Dlinear

## Convergence Testing
println("Convergence Test on Linear")
dts = 1.//2.^(4:-1:2)
testTol = 1

println("Feagin RKs")
sol =solve(prob,Feagin10(),dt=dts[1])

const linear_bigα = parse(BigFloat,"1.01")
f = (t,u,du) -> begin
  for i in 1:length(u)
    du[i] = linear_bigα*u[i]
  end
end
analytic = (t,u0) -> u0*exp(linear_bigα*t)
prob_ode_bigfloat2Dlinear = ODETestProblem(f,map(BigFloat,rand(4,2)).*ones(4,2)/2,analytic)

prob = prob_ode_bigfloat2Dlinear

sim = test_convergence(dts,prob,Feagin10())
@test abs(sim.𝒪est[:final]-8) < testTol #Lowered due to low test dt

sim = test_convergence(dts,prob,Feagin12())
@test abs(sim.𝒪est[:final]-12) < testTol

sim = test_convergence(dts,prob,Feagin14())
@test abs(sim.𝒪est[:final]-15) < testTol #Upped to 15 for test

f = (t,u) -> (linear_bigα*u)
prob_ode_bigfloatlinear = ODETestProblem(f,parse(BigFloat,"0.5"),analytic)
prob = prob_ode_bigfloatlinear

dts = 1.//2.^(6:-1:3)
sim = test_convergence(dts,prob,Feagin10())
@test abs(sim.𝒪est[:final]-10) < testTol

dts = 1.//2.^(4:-1:2)
sim = test_convergence(dts,prob,Feagin12())
@test abs(sim.𝒪est[:final]-12) < testTol

sim = test_convergence(dts,prob,Feagin14())
@test abs(sim.𝒪est[:final]-15) < testTol #Upped to 15 for test

prob = prob_ode_bigfloat2Dlinear

#compile
sol =solve(prob,Feagin10(),dt=dts[1])
sol =solve(prob,Feagin12(),dt=dts[1])
sol =solve(prob,Feagin14(),dt=dts[1])

#test
@time sol =solve(prob,Feagin10(),dt=dts[1])
@time sol =solve(prob,Feagin12(),dt=dts[1])
@time sol =solve(prob,Feagin14(),dt=dts[1])


function optimize(f::Function, lower::T, upper::T;
        alg::UnivariateAlgorithm=GoldenSection(),
        reltol=sqrt(eps(T)),
        abstol=eps(T),
        maxiter::Integer=1_000) where {T<:AbstractFloat}
    lower < upper || error("Variable lower bound has to be smaller than upper bound.")
    _optimize(f, lower, upper, alg; reltol=T(reltol), abstol=T(abstol), maxiter)
end

"""
    optimize(f::Function, lower::Real, upper::Real; kwargs...)

Compute the minimizer of the function `f` using the bracket (`lower`, `upper`).

# Keywords
- `alg::UnivariateAlgorithm=GoldenSection()`: the algorithm choice
- `reltol=sqrt(eps(T))`: relative tolerance
- `abstol=eps(T)`: absolute tolerance
- `maxiter::Integer=1_000`: maximum numer of iterations
"""
function optimize(f::Function, lower::Real, upper::Real; kwargs...)
    T = promote_type(typeof(lower/1), typeof(upper/1))
    optimize(f, T(lower), T(upper); kwargs...)
end


# ---
# title: 808. Soup Servings
# id: problem808
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/soup-servings/description/>
# hidden: true
# ---
#
# There are two types of soup: type A and type B. Initially we have `N` ml of
# each type of soup. There are four kinds of operations:
#
#   1. Serve 100 ml of soup A and 0 ml of soup B
#   2. Serve 75 ml of soup A and 25 ml of soup B
#   3. Serve 50 ml of soup A and 50 ml of soup B
#   4. Serve 25 ml of soup A and 75 ml of soup B
#
# When we serve some soup, we give it to someone and we no longer have it.  Each
# turn, we will choose from the four operations with equal probability 0.25. If
# the remaining volume of soup is not enough to complete the operation, we will
# serve as much as we can.  We stop once we no longer have some quantity of both
# types of soup.
#
# Note that we do not have the operation where all 100 ml's of soup B are used
# first.
#
# Return the probability that soup A will be empty first, plus half the
# probability that A and B become empty at the same time.
#
#
#
#
#
#     **Example:**
#     Input: N = 50
#     Output: 0.625
#     Explanation:
#     If we choose the first two operations, A will become empty first. For the third operation, A and B will become empty at the same time. For the fourth operation, B will become empty first. So the total probability of A becoming empty first plus half the probability that A and B become empty at the same time, is 0.25 * (1 + 1 + 0.5 + 0) = 0.625.
#
#
#
# **Notes:**
#
#   * `0 <= N <= 10^9`.
#   * Answers within `10^-6` of the true value will be accepted as correct.
#
#
## @lc code=start
using LeetCode

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


# Transformations between types
using SparseArrays

export taylor_expansion_pep
export shift_and_scale
export mobius_transform

"""
    struct ShiftScaledNEP <: NEP
    ShiftScaleNEP(orgnep::NEP[,shift=0][,scale=1])

Transforms a nep (orgnep) M(λ)v to a new nep T(λ)=M(scale*λ+shift). This can be used if the method does not have an easy implementation of shift and scaling. Usage of this transformation can slow down the algorithm.
"""
struct ShiftScaledNEP <: NEP
    shift::Number
    scale::Number
    orgnep::NEP
    function ShiftScaledNEP(orgnep;shift=0,scale=1)
        return new(shift,scale,orgnep)
    end
end
"""
    shift_and_scale(orgnep::NEP;shift=0,scale=1)

Transforms the orgnep by defining a new NEP from the relation
T(λ)=M(scale * λ+shift) where M is the orgnep. This function tries
 to preserve the NEP type, e.g., a shift_and_scale operation on
an SPMF-object, return an SPMF object. If it cannot preserve
the type, it will return a nep of the struct `ShiftScaledNEP`.

# Example
```julia-repl
julia> nep0=nep_gallery("pep0")
julia> σ=3; α=10;
julia> nep1=shift_and_scale(nep0,shift=σ,scale=α)
julia> opnorm(compute_Mder(nep0,α*(4+4im)+σ)-compute_Mder(nep1,4+4im))
8.875435870738592e-12
```
"""
function shift_and_scale(orgnep::NEP;shift=0,scale=1)
    return ShiftScaledNEP(orgnep,shift=shift,scale=scale);
end

# Native implementation for PEPs by transformation of the coefficients. This is not necessarily fast.
function shift_and_scale(orgnep::PEP;shift=0,scale=1)
    At=copy(shift*scale*orgnep.A); # Allocate correct type (not so fast)
    m=size(At,1)-1
    for j=0:m
        AA = zero(At[j+1])
        for i=j:m
            factor=((scale^j)*(shift^(i-j))*factorial(i)/factorial(i-j))/factorial(j)
            AA+=orgnep.A[i+1]*factor;

        end
        At[j+1]=AA;
    end
    return PEP(At)
end

# A shift and rescale of a DEP is DEP where the matrix coefficients and the delays are rescaled and another matrix coefficient with delay zero is added.
function shift_and_scale(orgnep::DEP;shift=0,scale=1)
    J = Matrix{eltype(orgnep.A[1])}(I, size(orgnep.A[1]))
    return DEP([broadcast(*,orgnep.A,exp.(-orgnep.tauv*shift)/scale); [-shift/scale * J] ],[orgnep.tauv*scale; 0])
end

# Native implementation for SPMF. Only for generic SPMF (not AbstractSPMF)
function shift_and_scale(orgnep::SPMF_NEP;shift=0,scale=1)
    orgfv=get_fv(orgnep);
    m=size(orgfv,1);
    fv = Vector{Function}(undef, m)
    # create anonymous functions corresponding to the
    # shift and scaled problem
    for i=1:m
        fv[i] = S -> orgfv[i](scale*S + shift*one(S))::((S isa Number) ? Number : Matrix)
    end
    # create a new SPMF with transformed functions
    return SPMF_NEP(get_Av(orgnep),fv;Schur_fact=orgnep.Schur_factorize_before);
end

# Size and issparse implemented by delegation
function size(nep::ShiftScaledNEP,dim=-1)
    return size(nep.orgnep,dim)
end
function issparse(nep::ShiftScaledNEP)
    return issparse(nep.orgnep)
end

function compute_Mder(nep::ShiftScaledNEP,λ::Number,i::Integer=0)
    # Just scale the orgnep call. Chain rule for differentiation.
    return (nep.scale^i)*compute_Mder(nep.orgnep,nep.scale*λ+nep.shift,i);
end

function compute_MM(nep::ShiftScaledNEP,S,V)
    # Just call orgnep with a different S
    return compute_MM(nep.orgnep,S*nep.scale+eye(S)*nep.shift,V)
end

function compute_Mlincomb(nep::ShiftScaledNEP,λ::Number,V::AbstractVecOrMat,
                          a::Vector=ones(size(V,2)))
    # Multiply the coefficient matrix V with a diagonal matrix
    # corresponding to the chain rule
    p=size(V,2);
    z=nep.scale.^Array{eltype(V),1}(0:p-1)
    W=V * diagm(0 => z) # not very fast
    return compute_Mlincomb(nep.orgnep,nep.scale*λ+nep.shift,W,a);
end

"""
    mobius_transform(orgnep::NEP,[,a=1][,b=0][,c=0][,d=1])

Transforms a nep (orgnep) ``M(λ)v`` to a new nep ``T(λ)=M((aλ+b)/(cλ+d))``.
This function tries to preserve the type such that `T`
and `M` are of the same NEP-type (see `shift_and_scale()`).
If it cannot be preserved it will return a `MobiusTransformedNEP`.
It is in general advised to try to preserve the type, and
the use of `MobiusTransformedNEP` can
considerably slow down NEP-access.

# Example

```julia-repl
julia> nep0=nep_gallery("pep0")
julia> a=1; b=3; c=4; d=5;
julia> nep1=mobius_transform(nep0,a=a,b=b,c=c,d=d);
julia> s=3;
julia> opnorm(compute_Mder(nep0,(a*s+b)/(c*s+d))-compute_Mder(nep1,s))
0.0
```
"""
function mobius_transform(orgnep::NEP;a=1,b=0,c=0,d=1)
    return MobiusTransformedNEP(orgnep;a=a,b=b,c=c,d=d)
end

function mobius_transform(orgnep::SPMF_NEP;a=1,b=0,c=0,d=1)
    orgfv=get_fv(orgnep);
    m=size(orgfv,1);
    fv = Vector{Function}(undef, m)
    # create anonymous functions corresponding to the
    # möbius transformed problem
    for i=1:m
        fv[i] = S -> orgfv[i]((a*S + b*one(S)) / (c*S + d*one(S)))::((S isa Number) ? Number : Matrix)
    end
    # create a new SPMF with transformed functions
    return SPMF_NEP(get_Av(orgnep),fv;Schur_fact=orgnep.Schur_factorize_before);
end

"""
Hidden type representing a transformed NEP
"""
struct MobiusTransformedNEP <: NEP
    a::Number
    b::Number
    c::Number
    d::Number
    orgnep::NEP
    function MobiusTransformedNEP(orgnep;a=1,b=0,c=0,d=1)
        return new(a,b,c,d,orgnep)
    end
end

# Size and issparse implemented by delegation
function size(nep::MobiusTransformedNEP,dim=-1)
    return size(nep.orgnep,dim)
end
function issparse(nep::MobiusTransformedNEP)
    return issparse(nep.orgnep)
end

function compute_Mder(nep::MobiusTransformedNEP,λ::Number,i::Integer=0)
    # I did not found a better way
    return compute_Mder_from_MM(nep,λ,i);
end

function compute_MM(nep::MobiusTransformedNEP,S,V)
    # Just call orgnep with a different S
    return compute_MM(nep.orgnep, (nep.a*S + nep.b*I) / (nep.c*S + nep.d*I), V)
end

function compute_Mlincomb(nep::MobiusTransformedNEP,λ::Number,V::AbstractVecOrMat
                          ,a::Vector=ones(size(V,2)))
    # I did not found a better way
    return compute_Mlincomb_from_MM(nep,λ,V,a)
end

# consider renaming this function
"""
    taylor_expansion_pep(orgnep::NEP[,d=2])

Compute the truncated (with `d` terms) Taylor series of the NEP.
The output is a [`PEP`](@ref).
"""
function taylor_expansion_pep(nep::NEP,d::Integer=2)
    A=Vector{Matrix}(undef,d+1)
    for i=0:d
        A[i+1]=compute_Mder(nep,0,i)/factorial(i);
    end
    return PEP(A)
end


mutable struct FTestResult{N}
    ssr::NTuple{N, Float64}
    dof::NTuple{N, Int}
    dof_resid::NTuple{N, Int}
    r2::NTuple{N, Float64}
    fstat::Tuple{Vararg{Float64}}
    pval::Tuple{Vararg{PValue}}
end

"""A helper function to determine if mod1 is nested in mod2"""
function issubmodel(mod1::LinPredModel, mod2::LinPredModel; atol=0::Real)
    mod1.rr.y != mod2.rr.y && return false # Response variables must be equal

    # Test that models are nested
    pred1 = mod1.pp.X
    npreds1 = size(pred1, 2)
    pred2 = mod2.pp.X
    npreds2 = size(pred2, 2)
    # If model 1 has more predictors, it can't possibly be a submodel
    npreds1 > npreds2 && return false
    # Test min norm pred2*B - pred1 ≈ 0
    rtol = Base.rtoldefault(typeof(pred1[1,1]))
    nresp = size(pred2, 1)
    return norm(view(qr(pred2).Q'pred1, npreds2 + 1:nresp, :)) <= max(atol, rtol*norm(pred1))
end

_diffn(t::NTuple{N, T}) where {N, T} = ntuple(i->t[i]-t[i+1], N-1)

_diff(t::NTuple{N, T}) where {N, T} = ntuple(i->t[i+1]-t[i], N-1)

dividetuples(t1::NTuple{N}, t2::NTuple{N}) where {N} = ntuple(i->t1[i]/t2[i], N)

"""
    ftest(mod::LinearModel...; atol=0::Real)

For each sequential pair of linear predictors in `mod`, perform an F-test to determine if
the first one fits significantly better than the next.

A table is returned containing residual degrees of freedom (DOF), degrees of freedom,
difference in DOF from the preceding model, sum of squared residuals (SSR), difference in
SSR from the preceding model, R², difference in R² from the preceding model, and F-statistic
and p-value for the comparison between the two models.

!!! note
    This function can be used to perform an ANOVA by testing the relative fit of two models
    to the data

Optional keyword argument `atol` controls the numerical tolerance when testing whether
the models are nested.

# Examples

Suppose we want to compare the effects of two or more treatments on some result. Because
this is an ANOVA, our null hypothesis is that `Result ~ 1` fits the data as well as
`Result ~ 1 + Treatment`.

```jldoctest
julia> dat = DataFrame(Treatment=[1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2.],
                       Result=[1.1, 1.2, 1, 2.2, 1.9, 2, .9, 1, 1, 2.2, 2, 2],
                       Other=[1, 1, 2, 1, 2, 1, 3, 1, 1, 2, 2, 1]);

julia> model = lm(@formula(Result ~ 1 + Treatment), dat);

julia> nullmodel = lm(@formula(Result ~ 1), dat);

julia> bigmodel = lm(@formula(Result ~ 1 + Treatment + Other), dat);

julia> ft = ftest(model.model, nullmodel.model)
        Res. DOF DOF ΔDOF    SSR    ΔSSR     R²    ΔR²       F* p(>F)
Model 1       10   3      0.1283         0.9603
Model 2       11   2   -1 3.2292 -3.1008 0.0000 0.9603 241.6234 <1e-7

julia> ftest(bigmodel.model, model.model, nullmodel.model)
        Res. DOF DOF ΔDOF    SSR    ΔSSR     R²    ΔR²       F*  p(>F)
Model 1        9   4      0.1038         0.9678
Model 2       10   3   -1 0.1283 -0.0245 0.9603 0.0076   2.1236 0.1790
Model 3       11   2   -1 3.2292 -3.1008 0.0000 0.9603 241.6234  <1e-7
```
"""
function ftest(mods::LinearModel...; atol=0::Real)
    nmodels = length(mods)
    for i in 2:nmodels
        issubmodel(mods[i], mods[i-1], atol=atol) ||
        throw(ArgumentError("F test $i is only valid if model $i is nested in model $(i-1)"))
    end

    SSR = deviance.(mods)

    nparams = dof.(mods)

    df2 = _diffn(nparams)
    df1 = Int.(dof_residual.(mods))

    MSR1 = dividetuples(_diff(SSR), df2)
    MSR2 = dividetuples(SSR, df1)[1:nmodels-1]

    fstat = dividetuples(MSR1, MSR2)
    pval = PValue.(ccdf.(FDist.(df2, df1[1:nmodels-1]), fstat))
    return FTestResult(SSR, nparams, df1, r2.(mods), fstat, pval)
end

function show(io::IO, ftr::FTestResult{N}) where N
    Δdof = _diffn(ftr.dof_resid)
    Δssr = _diffn(ftr.ssr)
    ΔR² = _diffn(ftr.r2)

    nc = 10
    nr = N
    outrows = Matrix{String}(undef, nr+1, nc)

    outrows[1, :] = ["", "Res. DOF", "DOF", "ΔDOF", "SSR", "ΔSSR",
                     "R²", "ΔR²", "F*", "p(>F)"]

    outrows[2, :] = ["Model 1", @sprintf("%.0d", ftr.dof_resid[1]),
                     @sprintf("%.0d", ftr.dof[1]), " ",
                     @sprintf("%.4f", ftr.ssr[1]), " ",
                     @sprintf("%.4f", ftr.r2[1]), " ", " ", " "]

    for i in 2:nr
        outrows[i+1, :] = ["Model $i", @sprintf("%.0d", ftr.dof_resid[i]),
                           @sprintf("%.0d", ftr.dof[i]), @sprintf("%.0d", Δdof[i-1]),
                           @sprintf("%.4f", ftr.ssr[i]), @sprintf("%.4f", Δssr[i-1]),
                           @sprintf("%.4f", ftr.r2[i]), @sprintf("%.4f", ΔR²[i-1]),
                           @sprintf("%.4f", ftr.fstat[i-1]), string(ftr.pval[i-1]) ]
    end
    colwidths = length.(outrows)
    max_colwidths = [maximum(view(colwidths, :, i)) for i in 1:nc]

    for r in 1:nr+1
        for c in 1:nc
            cur_cell = outrows[r, c]
            cur_cell_len = length(cur_cell)

            padding = " "^(max_colwidths[c]-cur_cell_len)
            if c > 1
                padding = " "*padding
            end

            print(io, padding)
            print(io, cur_cell)
        end
        print(io, "\n")
    end
end


module DWDDynam

using Catlab.WiringDiagrams.DirectedWiringDiagrams
using  Catlab.CategoricalAlgebra
using  Catlab.CategoricalAlgebra.FinSets
import Catlab.WiringDiagrams: oapply

import ..UWDDynam: nstates, eval_dynamics, euler_approx

using OrdinaryDiffEq, DynamicalSystems
import OrdinaryDiffEq: ODEProblem
import DynamicalSystems: DiscreteDynamicalSystem

export AbstractMachine, ContinuousMachine, DiscreteMachine,
nstates, ninputs, noutputs, eval_dynamics, readout, euler_approx

using Base.Iterators
import Base: show, eltype

"""

A directed open dynamical system operating on information fo type `T`.
A machine  `m` has type signature  (`m.ninputs`, `m.outputs`).
"""
abstract type AbstractMachine{T} end

"""

An undirected open continuous system. The dynamics function `f` defines an ODE ``\\dot u(t) = f(u(t),x(t),p,t)`` where ``u`` is the state and ``x`` captures the exogenous variables.

The readout function may only depend on the state, so it must be of the form ``r(u(t))``.
"""
struct ContinuousMachine{T} <: AbstractMachine{T}
    ninputs::Int
    nstates::Int
    noutputs::Int
    dynamics::Function
    readout::Function
end

"""

A directed open discrete dynamical system. The dynamics function `f` defines a discrete update rule ``u_{n+1} = f(u_n, x_n, p, t)`` where ``u_n`` is the state and ``x_n`` is the value of the exogenous variables at the ``n``th time step.

The readout function may only depend on the state, so it must be of the form ``r(u_n)``.
"""
struct DiscreteMachine{T} <: AbstractMachine{T}
    ninputs::Int
    nstates::Int
    noutputs::Int
    dynamics::Function
    readout::Function
end

show(io::IO, vf::ContinuousMachine) = print("ContinuousMachine(ℝ^$(vf.nstates) × ℝ^$(vf.ninputs) → ℝ^$(vf.nstates))")
show(io::IO, vf::DiscreteMachine) = print("DiscreteMachine(ℝ^$(vf.nstates) × ℝ^$(vf.ninputs) → ℝ^$(vf.nstates))")
eltype(m::AbstractMachine{T}) where T = T

nstates(f::AbstractMachine) = f.nstates
ninputs(f::AbstractMachine) = f.ninputs
noutputs(f::AbstractMachine) = f.noutputs
readout(f::AbstractMachine, u::AbstractVector) = f.readout(u)

"""    eval_dynamics(m::AbstractMachine, u::AbstractVector, xs::AbstractVector, p, t)

Evaluates the dynamics of the machine `m` at state `u`, parameters `p`, and time `t`. The exogenous variables are set by `xs` which may either be a collection of functions ``x(t)`` or a collection of constant values.

The length of `xs` must equal the number of inputs to `m`.
"""
eval_dynamics(f::AbstractMachine, u::AbstractVector, xs::AbstractVector{T}, p, t::Real) where T <: Function =
    eval_dynamics(f, u, [x(t) for x in xs], p, t)

eval_dynamics(f::AbstractMachine, u::AbstractVector, xs::AbstractVector, p, t::Real)  = begin
    ninputs(f) == length(xs) || error("$xs must have length $(ninputs(f)) to set the exogenous variables.")
    f.dynamics(u, xs, p, t)
end

eval_dynamics(f::AbstractMachine, u::AbstractVector, xs::AbstractVector) =
    eval_dynamics(f, u, xs, nothing, 0)
eval_dynamics(f::AbstractMachine, u::AbstractVector, xs::AbstractVector, p) =
    eval_dynamics(f, u, xs, p, 0)

"""    euler_approx(m::ContinuousMachine, h)

Transforms a continuous machine `m` into a discrete machine via Euler's method with step size `h`. If the dynamics of `m` is given by ``\\dot{u}(t) = f(u(t),x(t),p,t)`` the the dynamics of the new discrete system is given by the update rule ``u_{n+1} = u_n + h f(u_n, x_n, p, t)``.
"""
euler_approx(f::ContinuousMachine{T}, h::Float64) where T = DiscreteMachine{T}(
    ninputs(f), nstates(f), noutputs(f),
    (u, x, p, t) -> u + h*eval_dynamics(f, u, x, p, t),
    f.readout
)

"""    euler_approx(m::ContinuousMachine)

Transforms a continuous machine `m` into a discrete machine via Euler's method where the step size is introduced as a new parameter, the last in the list of parameters.
"""
euler_approx(f::ContinuousMachine{T}) where T = DiscreteMachine{T}(
    ninputs(f), nstates(f), noutputs(f),
    (u, x, p, t) -> u + p[end]*eval_dynamics(f, u, x, p[1:end-1], t),
    f.readout
)
euler_approx(fs::Vector{ContinuousMachine{T}}, args...) where T =
    map(f->euler_approx(f,args...), fs)

euler_approx(fs::AbstractDict{S, ContinuousMachine{T}}, args...) where {S, T} =
    Dict(name => euler_approx(f, args...) for (name, f) in fs)

# Integration with ODEProblem in OrdinaryDiffEq.jl

"""    ODEProblem(m::ContinuousMachine, xs::Vector, u0::Vector, tspan)

Constructs an ODEProblem from the vector field defined by `(u,p,t) -> m.dynamics(u, x, p, t)`. The exogenous variables are determined by `xs`.
"""
ODEProblem(m::ContinuousMachine, u0::AbstractVector, xs::AbstractVector, tspan::Tuple{Real, Real}, p=nothing)  =
    ODEProblem((u,p,t) -> eval_dynamics(m, u, xs, p, t), u0, tspan, p)

ODEProblem(m::ContinuousMachine, u0::AbstractVector, x, tspan::Tuple{Real, Real}, p=nothing) =
    ODEProblem(m, u0, collect(repeated(x, ninputs(m))), tspan, p)

ODEProblem(m::ContinuousMachine{T}, u0::AbstractVector, tspan::Tuple{Real, Real}, p=nothing) where T =
    ODEProblem(m, u0, T[], tspan, p)

"""    DiscreteDynamicalSystem(m::DiscreteMachine, xs::Vector, u0::Vector, p)

Constructs an DiscreteDynamicalSystem from the equation of motion defined by
`(u,p,t) -> m.dynamics(u, x, p, t)`. The exogenous variables are determined by `xs`. Pass `nothing` in place of `p` if your system does not have parameters.
"""
DiscreteDynamicalSystem(m::DiscreteMachine, u0::AbstractVector, x,  p; t0::Int = 0) =
  DiscreteDynamicalSystem(m, u0, collect(repeated(x, ninputs(m))), p; t0=t0)

DiscreteDynamicalSystem(m::DiscreteMachine{T}, u0::AbstractVector, xs::AbstractVector, p; t0::Int = 0) where T = begin
  if nstates(m) == 1
    DiscreteDynamicalSystem1d(m, u0[1], xs, p; t0=t0)
  else
    !(T <: AbstractFloat) || error("Cannot construct a DiscreteDynamicalSystem if the type is a float")
    DiscreteDynamicalSystem(
        (u,p,t) -> SVector{nstates(m)}(eval_dynamics(m, u, xs, p, t)),
        u0, p; t0=t0)
  end
end

DiscreteDynamicalSystem(m::DiscreteMachine, u0::Real, xs::AbstractVector, p; t0::Int = 0) =
  DiscreteDynamicalSystem1D(m, u0, xs, p; t0=t0)

# if the system is 1D then the state must be represented by a number NOT by a 1D array
DiscreteDynamicalSystem1d(m::DiscreteMachine{T}, u0::Real, xs::AbstractVector, p; t0::Int = 0) where T = begin
  nstates(m) == 1 || error("The machine must have exactly 1 state")
  !(T <: AbstractFloat) || error("Cannot construct a DiscreteDynamicalSystem if the type is a float")
  DiscreteDynamicalSystem(
      (u,p,t) -> eval_dynamics(m, [u], xs, p, t)[1],
      u0, p; t0=t0)
end

DiscreteDynamicalSystem(m::DiscreteMachine, u0, p; t0::Int = 0) =
    DiscreteDynamicalSystem(m, u0, [], p; t0 = t0)

"""    fills(m::AbstractMachine, d::WiringDiagram, b::Int)

Checks if `m` is of the correct signature to fill box `b` of the  wiring diagram `d`.
"""
function fills(m::AbstractMachine, d::WiringDiagram, b::Int)
    b <= nboxes(d) || error("Trying to fill box $b, when $d has fewer than $b boxes")
    b = box_ids(d)[b]
    return ninputs(m) == length(input_ports(d,b)) && noutputs(m) == length(output_ports(d,b))
end

colimitmap!(f::Function, output, C::Colimit, input) = begin
    for (i,x) in enumerate(input)
        y = f(i, x)
        I = legs(C)[i](1:length(y))
        # length(I) == length(y) || error("colimitmap! attempting to fill $(length(I)) slots with $(length(y)) values")
        output[I] .= y
    end
    return output
end

@inline fillreadouts!(y, d, xs, Outputs, statefun) = colimitmap!(y, Outputs, xs) do i,x
    return x.readout(statefun(i))
end

@inline fillstates!(y, d, xs, States, statefun, inputfun, p, t) = colimitmap!(y, States, xs) do i, x
    return x.dynamics(statefun(i), inputfun(i), p, t)
end

@inline fillwire(w, d, readouts, Outputs) = readouts[legs(Outputs)[w.source.box - 2](w.source.port)] # FIX - box re-indexing

fillreadins!(readins, d, readouts, Outputs, Inputs, ins) = begin
    for (i,w) in enumerate(wires(d))
        if w.target.box == output_id(d)
            continue
        elseif w.source.box == input_id(d)
            readins[legs(Inputs)[w.target.box - 2](w.target.port)] += ins[w.source.port]
        else
            readins[legs(Inputs)[w.target.box - 2](w.target.port)] += fillwire(w, d, readouts, Outputs)
        end

    end
    return readins
end

"""    oapply(d::WiringDiagram, ms::Vector)

Implements the operad algebras for directed composition of dynamical systems given a
composition pattern (implemented by a directed wiring diagram `d`)
and primitive systems (implemented by a collection of
machines `ms`).

Each box of the composition pattern `d` is filled by a machine with the
appropriate type signature. Returns the composite machine.
"""
function oapply(d::WiringDiagram, xs::Vector{Machine}) where {T, Machine <: AbstractMachine{T}}
    isempty(wires(d, input_id(d), output_id(d))) || error("d is not a valid composition syntax because it has pass wires")
    nboxes(d) == length(xs)  || error("there are $nboxes(d) boxes but $length(xs) machines")
    for box in 1:nboxes(d)
        fills(xs[box], d, box) || error("$xs[box] does not fill box $box")
    end

    S = coproduct((FinSet∘nstates).(xs))
    Inputs = coproduct((FinSet∘ninputs).(xs))
    Outputs = coproduct((FinSet∘noutputs).(xs))
    ys = zeros(T, length(apex(S)))

    states(u::AbstractVector, b::Int) = u[legs(S)[b](1:xs[b].nstates)]

    v = (u::AbstractVector, ins::AbstractVector, p, t::Real) -> begin
        readouts = zeros(T, length(apex(Outputs)))
        readins = zeros(T, length(apex(Inputs)))

        get_states(b) = states(u,b)
        get_inputs(b) = view(readins, legs(Inputs)[b](:))

        fillreadouts!(readouts, d, xs, Outputs, get_states)
        fillreadins!(readins, d, readouts, Outputs, Inputs, ins)
        fillstates!(ys, d, xs, S, get_states, get_inputs, p, t)
        return ys
    end

    function readout(u::AbstractVector)
        readouts = zeros(T, length(apex(Outputs)))
        get_states(b) = states(u,b)

        fillreadouts!(readouts, d, xs, Outputs, get_states)
        r = zeros(T, length(d.output_ports))
        for w in in_wires(d, output_id(d))
            r[w.target.port] += fillwire(w, d, readouts, Outputs)
        end
        return r
    end

    return Machine(length(d.input_ports), length(apex(S)), length(d.output_ports), v, readout)

end

"""    oapply(d::WiringDiagram, m::AbstractMachine)

A version of `oapply` where each box of `d` is filled with the machine `m`.
"""
function oapply(d::WiringDiagram, x::AbstractMachine)
    oapply(d, collect(repeated(x, nboxes(d))))
end

"""    oapply(d::WiringDiagram, generators::Dict)

A version of `oapply` where `generators` is a dictionary mapping the name of each box to its corresponding machine.
"""
function oapply(d::WiringDiagram, xs::AbstractDict)
    oapply(d, [xs[box.value] for box in boxes(d)])
end

end #module


@symbol_func function tau_E(cur_reactor::AbstractReactor)
  cur_scaling = cur_reactor.mode_scaling

  cur_tau_E = cur_scaling[:constant]

  cur_tau_E *= cur_reactor.H

  cur_tau_E *= cur_reactor.I_P ^ cur_scaling[:I]

  cur_tau_E *= cur_reactor.R_0 ^ cur_scaling[:R]

  cur_tau_E *= a(cur_reactor) ^ cur_scaling[:a]

  cur_tau_E *= kappa_tau(cur_reactor) ^ cur_scaling[:kappa]

  cur_tau_E *= cur_reactor.n_bar ^ cur_scaling[:n]

  cur_tau_E *= cur_reactor.B_0 ^ cur_scaling[:B]

  cur_tau_E *= cur_reactor.A ^ cur_scaling[:A]

  cur_tau_E /= P_L(cur_reactor) ^ cur_scaling[:P]

  cur_tau_E
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

module SEIRmodel
using OrdinaryDiffEq

#Susceptible-exposed-infected-recovered model function
function seir_ode(dY,Y,p,t)
    #Infected per-Capita Rate
    β = p[1]
    #Incubation Rate
    σ = p[2]
    #Recover per-capita rate
    γ = p[3]
    #Death Rate
    μ = p[4]

    #Susceptible Individual
    S = Y[1]
    #Exposed Individual
    E = Y[2]
    #Infected Individual
    I = Y[3]
    #Recovered Individual
    #R = Y[4]

    dY[1] = μ-β*S*I-μ*S
    dY[2] = β*S*I-(σ+μ)*E
    dY[3] = σ*E - (γ+μ)*I
end

function main()
    #Pram (Infected Rate, Incubation Rate, Recover Rate, Death Rate)
    pram=[520/365,1/60,1/30,774835/(65640000*365)]
    #Initialize Param(Susceptible Individuals, Exposed Individuals, Infected Individuals)
    init=[0.8,0.1,0.1]
    tspan=(0.0,365.0)

    seir_prob = ODEProblem(seir_ode,init,tspan,pram)

    sol=solve(seir_prob, alg=Tsit5());
    return sol
end

# using Plots

# va = VectorOfArray(sol.u)
# y = convert(Array,va)
# R = ones(size(sol.t))' - sum(y,dims=1);

# plot(sol.t,[y',R'],xlabel="Time",ylabel="Proportion")

end


const Mean = let
    init() = (0, 0.0)
    op((n, sum), val) = n + 1, sum + val
    value((n, sum)) = sum / n
    (; init, op, value)
end

struct ExpectationAnalysis end

(e::ExpectationAnalysis)(input::ProcessedLayer) = groupreduce(Mean, input)

"""
    expectation()

Compute the expected value of the last argument conditioned on the preceding ones.
"""
expectation() = transformation(ExpectationAnalysis())


using ForwardDiff,SparseArrays

include(string(pwd(),"/src/gensys.jl"))
include(string(pwd(),"/misc/aux_matrix.jl"))

function dF1_theta(A,dGamma_0,dGamma_1,dGamma_2)
    m = size(A,1)
    return kron(A'*I,I(m))*dGamma_0 - kron(A'^2,I(m))*dGamma_1 - dGamma_2
end

function dF1_A(A,Gamma_0,Gamma_1)
    m = size(A,1)
    b1 = Gamma_0 - Gamma_1*A
    return kron(I(m),b1) - kron(A'*I,Gamma_1)
end

function dA_theta(A,Gamma_0,Gamma_1,dGamma_0,dGamma_1,dGamma_2)
    dvecA = dF1_A(A,Gamma_0,Gamma_1)
    dTheta = dF1_theta(A,dGamma_0,dGamma_1,dGamma_2)
    return -dvecA\Array(dTheta)
end

function dB_theta(Gamma_0,Gamma_1,Gamma_3,A,dGamma_0,dGamma_1,dGamma_2,dGamma_3)
    dA = dA_theta(A,Gamma_0,Gamma_1,dGamma_0,dGamma_1,dGamma_2)
    dA[abs.(dA) .< 10*eps()] .= 0
    dA = sparse(dA)

    m = size(A,1)
    n = size(Gamma_3,2)

    b0 = Array(Gamma_0 - Gamma_1*A)
    b0[abs.(b0) .< 10*eps()] .= 0
    b0 = pinv(b0)

    b1 = kron((b0*Gamma_3)', b0)
    b2 = dGamma_0 - kron(A',I(m))*dGamma_1 - kron(I(m),Gamma_1)*dA

    return -b1*b2 + kron(I(n),b0)*dGamma_3
end

## These are the gensys matrices

#order to par
#1 alfa
#2 beta
#3 epsilon
#4 theta
#5 sig
#6 sigma: this is the std dev of the innovation
#7 phi
#8 phi_pi
#9 phi_y
#10 rho_v

GAMMA_0_foo(par) = [par[2]  0    0  0;
           1    par[5]  0  0;
           0    0    0  0;
           0    0    0  1]

function GAMMA_1_foo(par)
    THETA = (1-par[1])/(1-par[1]+par[1]*par[3])
    lamb = (1-par[4])*(1-par[2]*par[4])/par[4]*THETA
    kappa = lamb*(par[5]+(par[7]+par[1])/(1-par[1]))
     res = [ 1       -kappa  0   0;
            0        par[5]    1   0;
           -par[8]  -par[9]  1  -1;
            0        0      0   par[10]]
end

PSI_foo(par) = [0; 0; 0; par[6]]

PI_foo(par) = [par[2]  0;
      1    par[5];
      0    0;
      0    0]

## These are the matrices for the same model written as iskrev (2008)

#order to par
#1 alfa
#2 beta
#3 epsilon
#4 theta
#5 sig
#6 sigma: this is the std dev of the innovation
#7 phi
#8 phi_pi
#9 phi_y
#10 rho_v

function Gamma_0(par)
    THETA = (1-par[1])/(1-par[1]+par[1]*par[3])
    lamb = (1-par[4])*(1-par[2]*par[4])/par[4]*THETA
    kappa = lamb*(par[5]+(par[7]+par[1])/(1-par[1]))
    res = [1 -kappa 0 0;
            0       1   1/par[5] 0;
            -par[8] -par[9] 1  -1;
            0           0    0   1
            ]
    return res
end

Gamma_1(par) = [par[2] 0 0 0;
                1/par[5] 1 0 0;
                0        0 0 0;
                0        0 0 0]

Gamma_2(par) =[0 0 0 0;
               0 0 0 0;
               0 0 0 0;
               0 0 0 par[10]]

Gamma_3(par) = [0;
                0;
                0;
               par[6]]

## Diff for the model
# Now it will be a #estrut par x ## reduce form par
# In which the reduced form depends of 3 matrices, in the following order
# A, that governs the auto regressive part of the reduce form: first 16 cols
# C that govens the observation matrix in the Kalman Filter (usually should be all zeros): 4 columns
# Omega, that is BB' in which B governs the shock transmission: remaining 10 columns (it is a vech!)

function diff_mod(par;spar=true)
    Gamma0 = Gamma_0(par)
    Gamma1 = Gamma_1(par)
    Gamma2 = Gamma_2(par)
    Gamma3 = Gamma_3(par)

    G0 = GAMMA_0_foo(par)
    G1 = GAMMA_1_foo(par)
    Psi = PSI_foo(par)
    Pi = PI_foo(par)

    dGamma0 = ForwardDiff.jacobian(Gamma_0,par)
    dGamma1 = ForwardDiff.jacobian(Gamma_1,par)
    dGamma2 = ForwardDiff.jacobian(Gamma_2,par)
    dGamma3 = ForwardDiff.jacobian(Gamma_3,par)

    gen_sol = gensys(G0,G1,Psi,Pi, verbose = false)
    A = gen_sol.Theta1

    if spar == true #use sparse arrays
        dGamma0 = sparse(dGamma0)
        dGamma1 = sparse(dGamma1)
        dGamma2 = sparse(dGamma2)
        dGamma3 = sparse(dGamma3)

        Gamma0 = sparse(Gamma0)
        Gamma1 = sparse(Gamma1)
        Gamma2 = sparse(Gamma2)
        Gamma3 = sparse(Gamma3)

        A[abs.(A) .< 10*eps()] .= 0
        A = sparse(A)
    end

    dA = dA_theta(A,Gamma0,Gamma1,dGamma0,dGamma1,dGamma2)
    dB = dB_theta(Gamma0,Gamma1,Gamma3,A,dGamma0,dGamma1,dGamma2,dGamma3)

    if spar == true
        dA = sparse(dA)
    end

    return dA,dB
end


import Base.in
export in, Perm, fenchel

# the set of permutations of 1,...,n
mutable struct Perm <: ContiSet
  dim::Int
end

# check if w is in the Permutation set
function in(w::AbstractArray, p::Perm)
  n = p.dim
  @assert length(w) == n
  return Set(w) == Set(collect(1:n))
end

# compute min w^x st x in p
function fenchel(p::Perm, w::AbstractArray)
  n = p.dim
  @assert length(w) == n
  i = sortperm(w, rev = true) # biggest index of c is first element of i, etc.
  h = zeros(n)
  for ii=1:n
    h[i[ii]] = ii
  end
  return h
end


using IterativeSolvers

""" Convex Factorization Machine with Hazan's algorithm (squared loss)
    Yamada, M. et. al., http://arxiv.org/abs/1507.01073

Model:

f(x;w,W) = w^t z + 0.5 trace(W (xx^t - diag(x.*x)))

Parameters
----------

num_iter: int
   The number of iteration for Hazan's algorithm: (default: 100)

reg_W: double
   The regularization parameter for interaction term (2-way factors) (default 100)

w: double
   d + 1 dimensional vector the global bias (w[0]) and user-item bias (w[1:])

U: double
   d by T dimensional matrix such that W = U U^t

"""
mutable struct CFM
  num_iter::Int64
  reg_W::Float64
  w::Array{Float64, 1}
  U::Array{Float64, 2}

  CFM(num_iter, reg_W) = new(num_iter, reg_W, Float64[], Array{Float64}(0,0))
end

function train(model::CFM, X, Y)
  X = convert(SparseMatrixCSC{Float64,Int64}, X)
  Y = Array(Y)

  n, d = size(X)

  #Add bias for training data
  Z = hcat(ones(n), X)

  T = model.num_iter
  η = model.reg_W

  model.U = zeros(d, T)
  P = zeros(d, T)
  λ = ones(T)
  w = zeros(d+1, 1)
  global fval = zeros(T)

  for t in 1:T
    tmp = X * model.U[:, 1:t]
    fQ = 0.5 * (vec(sum(tmp .* tmp, 2)) - (X .* X) * vec(sum(model.U[:, 1:t] .* model.U[:, 1:t], 2)))

    ZY = Z' * (Y - fQ)

    #Conjugate Gradient: Solve Zw = Y
    wout = cg!(w, Z'*Z, ZY, tol=1e-6, verbose = false, maxiter=100)
    model.w = vec(wout)

    tr_err = Y - (Z*model.w + fQ)

    #Frank-Wolfe update: eigs(X diag(tr_err) X^t, 1)
    pout = eigs(X'*spdiagm(sparsevec(tr_err))*X, nev=1, which=:LR, maxiter = 300, tol = 1e-1)[2]

    p = vec(real(pout))

    #Optimal step size
    err = η*((X*p).^2 - (X .* X) * (p .* p)) - fQ
    α = ((tr_err' * err) / (err' * err))

    #Update
    P[:, t] = √η * p
    λ[1:t-1] = (1 - α) * λ[1:t-1]
    λ[t] = max(1e-10, α)

    model.U[:, 1:t] = P[:, 1:t] .* sqrt.(λ[1:t])'

    #Traning RMSE
    global fval[t] = sqrt(mean(tr_err.^2))

    #fval: training RMSE
    println("iteration:", t, " Training RMSE:", fval[t])

		tmp = nothing
		fQ = nothing
		ZY = nothing
		wout = nothing
		tr_err = nothing
		pout = nothing
		p = nothing
		err = nothing
		α= nothing
	end
end

function predicter(model::CFM, X)
  X = convert(SparseMatrixCSC{Float64,Int64}, X)

  #Compute fQ (2-way factor)
  tmp = X * model.U
  fQ = 0.5 * (sum(tmp .* tmp, 2) - (X .* X) * sum(model.U .* model.U, 2))

  Ŷ = model.w[1] + X * model.w[2:end] + fQ

  return Ŷ
end


#=
Derivatives of (r, l, P, T) with respect to the mass coordinate.
=#

include("constants.jl")
include("physics.jl")
include("opacity.jl")

# TODO: implement an opacity that changes with composition
const logkappa_spl = opal_test_spl()

function drdm(m, r, l, P, T, rho, mu)
    #=
    Derivative of r with respect to m, from mass conservation.
    =#
    return 1 / (4 * pi * r .^ 2 .* rho)
end # drdm

function dldm(m, r, l, P, T, rho, X, Y)
    #=
    Derivative of l with respect to m, from energy conservation.
    =#
    return epsilon_pp(rho, T, X) + epsilon_cno(rho, T, X, Y)
end # dldm

function dPdm(m, r, l, P, T)
    #=
    Derivative of P with respect to m, from hydrostatic equilibrium.
    =#
    return -G * m / (4 * pi * r .^ 4)
end # dPdm

function dTdm(m, r, l, P, T, mu, kappa)
    #=
    Derivative of T with respect to m, from... thermodynamics.
    =#
    grad_rad = del_rad(m, l, P, T, kappa)
    grad_ad = del_ad()
    stable =  grad_rad < grad_ad
    if stable
        # @debug("dTdm: stable to convection")
        grad = grad_rad
    else
        # @debug("dTdm: unstable to convection")
        grad = grad_ad
    end # if
    factor = - G .* m .* T ./ (4 * pi .* r .^ 4 .* P)
     return factor .* grad
end # dTdm

function deriv(m, r, l, P, T, X, Y)
    #=
    Subroutine for shootf.
    Calulates the derivatives of (r, l, P, T) with respect to mass.
    See Kippenhahn chapter 10, page 89.
    =#
    # @debug("deriv: m, y = ", m, ", ", [r, l, P, T])
    mu = mu_from_composition(X, Y)
    rho = rho_ideal(P, T, mu)
    kappa = try
        opacity(logkappa_spl, rho, T)
    catch
        @debug("deriv: Opacity interpolation failed!")
        @debug("deriv: y = ", [r, l, P, T])
        @debug("deriv: logT = ", log10(T))
        @debug("deriv: logR = ", log10(rho / (T / 1e6)^3))
        throw(DomainError)
    end
    dy = [drdm(m, r, l, P, T, rho, mu),
          dldm(m, r, l, P, T, rho, X, Y),
          dPdm(m, r, l, P, T),
          dTdm(m, r, l, P, T, mu, kappa)]
    # @debug("deriv: dy = ", dy)
    return dy
end # deriv


### Motion/Measurement Models ###

"""
`DynamicsModel` is an abstract type to encapsulate linear
and nonlinear dynamics models
"""
abstract type DynamicsModel end

"""
`ObservationModel` is an abstract type to encapsulate linear
and nonlinear observation models
"""
abstract type ObservationModel end

"""
    LinearDynamicsModel(A::AbstractMatrix,B::AbstractMatrix,W::Symmetric)
    LinearDynamicsModel(A::AbstractMatrix,B::AbstractMatrix,W::AbstractMatrix)

Construct linear dynamics model with transition matrix A,
control matrix B, and symmetric zero-mean process noise with
symmetric covariance matrix W
"""
struct LinearDynamicsModel{a<:AbstractMatrix, b<:AbstractMatrix,
                c<:Symmetric} <: DynamicsModel
    A::a
    B::b
    W::c
    #=
    function LinearDynamicsModel{a,b,c}(A::Matrix{a},B::Matrix{b},W::Symmetrix{c})
        where{a<:Number,b<:Number,c<:Number}

        (size(A,1) == size(B,1) == size(W,1)) ?
        new(A,B,W) : error("first dimensions not matching")
    end
    =#
end

function LinearDynamicsModel(A::AbstractMatrix, B::AbstractMatrix,
                             W::AbstractMatrix)

    return LinearDynamicsModel(A,B,Symmetric(W))
end

"""
    predict(m::LinearDynamicsModel, x::AbstractVector{<:Number}, u::AbstractVector{<:Number})
    predict(m::LinearDynamicsModel, x::AbstractVector{<:Number}, u::AbstractVector{<:Number}, rng::AbstractRNG)

Uses the linear dynamics model to propagate the state x one step forward in time with control input u.
If rng is given, it adds process noise.
"""
function predict(m::LinearDynamicsModel, x::AbstractVector{<:Number},
                 u::AbstractVector{<:Number})
    return m.A * x + m.B * u
end
function predict(m::LinearDynamicsModel, x::AbstractVector{T},
                 u::AbstractVector{T},
                 rng::AbstractRNG) where T<:Number
    return m.A * x + m.B * u + cholesky(m.W).L * randn(rng, T, size(m.W, 1))
end

function jacobian(m::LinearDynamicsModel, x, u)
    return m.A
end

"""
    LinearObservationModel(C::AbstractMatrix,D::AbstractMatrix,V::Symmetric)
    LinearObservationModel(C::AbstractMatrix,D::AbstractMatrix,V::AbstractMatrix)
    LinearObservationModel(C::AbstractMatrix,V::Symmetric)
    LinearObservationModel(C::AbstractMatrix,V::AbstractMatrix)

Construct linear observation dynamics model with output matrix C,
feedthrough matrix D, and symmetric zero-mean measurement noise with
symmetric covariance matrix V
"""
struct LinearObservationModel{a<:AbstractMatrix,b<:AbstractMatrix,c<:Symmetric} <: ObservationModel
    C::a
    D::b
    V::c
    #=
    LinearObservationModel(C,D,V) = (size(C,1) == size(D,1) == size(V,1)) ?
        new(C,D,V) : error("first dimensions not matching")
    =#
end

function LinearObservationModel(C::AbstractMatrix, D::AbstractMatrix,
                                V::AbstractMatrix)
    return LinearObservationModel(C,D,Symmetric(V))
end

function LinearObservationModel(C::AbstractMatrix, V::Symmetric)
    n = size(C,1)
    return LinearObservationModel(C,zeros(Bool,n,n),V)
end

function LinearObservationModel(C::AbstractMatrix, V::AbstractMatrix)
    n = size(C,1)
    return LinearObservationModel(C,zeros(Bool,n,n),Symmetric(V))
end

"""
    measure(m::LinearObservationModel, x::AbstractVector{<:Number}, u::AbstractVector{<:Number})
    measure(m::LinearObservationModel, x::AbstractVector{T}, u::AbstractVector{T}, rng::AbstractRNG) where T<:Number

Returns an observation of state x according to the linear observation model m, with control inputs u.
If rng is passed, adds additive Gaussian noise to the observation.
"""
function measure(m::LinearObservationModel, x::AbstractVector{<:Number},
                 u::AbstractVector{<:Number})
    return m.C * x + m.D * u
end
function measure(m::LinearObservationModel, x::AbstractVector{T},
                 u::AbstractVector{T},
                 rng::AbstractRNG) where T<:Number
    return m.C * x + m.D * u + cholesky(m.V).L * randn(rng, T, size(m.V, 1))
end

function jacobian(m::LinearObservationModel, x, u)
    m.C
end

"""
    NonlinearDynamicsModel(f::Function,W::Symmetric)
    NonlinearDynamicsModel(f::Function,W::AbstractMatrix)

Construct nonlinear dynamics model with transition function f
and symmetric zero-mean process noise with symmetric covariance matrix W
"""
struct NonlinearDynamicsModel{F<:Function, c<:Symmetric} <: DynamicsModel
    f::F
    W::c
end

function NonlinearDynamicsModel(f::Function, W::AbstractMatrix)
    return NonlinearDynamicsModel(f,Symmetric(W))
end

"""
    predict(m::NonLinearDynamicsModel, x::AbstractVector{<:Number}, u::AbstractVector{<:Number})
    predict(m::NonLinearDynamicsModel, x::AbstractVector{<:Number}, u::AbstractVector{<:Number}, rng::AbstractRNG)

Uses the non linear dynamics model to propagate the state x one step forward in time with control input u.
If rng is given, it adds process noise.
"""
function predict(m::NonlinearDynamicsModel, x::AbstractVector{<:Number},
                 u::AbstractVector{<:Number})
    return m.f(x, u)
end
function predict(m::NonlinearDynamicsModel, x::AbstractVector{T},
                 u::AbstractVector{T},
                 rng::AbstractRNG) where T<:Number
    return m.f(x, u) + cholesky(m.W).L * randn(rng, T, size(m.W, 1))
end

function jacobian(m::NonlinearDynamicsModel, x::AbstractVector, u::AbstractVector)
    return ForwardDiff.jacobian(μ -> m.f(μ, u), x)
end

"""
    NonlinearObservationModel(h::Function,V::Symmetric)
    NonlinearObservationModel(h::Function,V::AbstractMatrix)

Construct nonlinear observation dynamics model with measurement function h
and symmetric zero-mean measurement noise with symmetric covariance matrix V
"""
struct NonlinearObservationModel{F<:Function, c<:Symmetric} <: ObservationModel
    h::F
    V::c
end

function NonlinearObservationModel(h::Function, V::AbstractMatrix)
    return NonlinearObservationModel(h,Symmetric(V))
end

"""
    measure(m::NonlinearObservationModel, x::AbstractVector{<:Number}, u::AbstractVector{<:Number})
    measure(m::NonlinearObservationModel, x::AbstractVector{T}, u::AbstractVector{T}, rng::AbstractRNG) where T<:Number

Returns an observation of state x according to the non linear observation model m, with control inputs u.
If rng is passed, adds additive Gaussian noise to the observation.
"""
function measure(m::NonlinearObservationModel, x::AbstractVector{<:Number},
                 u::AbstractVector{<:Number})
    return m.h(x, u)
end
function measure(m::NonlinearObservationModel, x::AbstractVector{T},
                 u::AbstractVector{T},
                 rng::AbstractRNG) where T<:Number
    return m.h(x, u) + cholesky(m.V).L * randn(rng, T, size(m.V, 1))
end

function jacobian(m::NonlinearObservationModel, x::AbstractVector, u::AbstractVector)
    return ForwardDiff.jacobian(μ -> m.h(μ, u), x)
end


#===============================================================================

	Code for finding the terminal state of the HJB for a Hugget Model

===============================================================================#

using LinearAlgebra, SparseArrays, Plots

function terminal()
	# Define model parameters

	global s = 2
	global ρ = 0.05
	z1 = .1
	z2 = .2
	global z = [z1 z2]
	la1 = 0.6
	la2 = 0.8
	λ = [la1 la2]

	r0 = 0.03
	rmin = 0.01
	rmax = 0.04

	H= 1000
	global amin = -0.15
	global amax = 4
	a = LinRange(amin,amax,H)'
	global a = convert(Array, a)'
	global da = (amax-amin)/(H-1);

	global aa = [a a]
	global zz = ones(H,1)*z

	maxit= 100
	crit = 1e-6
	Delta = 100

	dVf, dVb, c = [zeros(H,2) for i in 1:3]

	global Aswitch = [-sparse(I, H, H)*λ[1] sparse(I, H, H)*λ[1];sparse(I, H, H)*λ[2] -sparse(I, H, H)*λ[2]]

	Ir = 40
	crit_S = 1e-8

	# INITIAL GUESS
	r = r0

	v0=zeros(H,2)
	v0[:,1] = (z[1] .+ r.*a).^(1-s)/(1-s)/ρ
	v0[:,2] = (z[2] .+ r.*a).^(1-s)/(1-s)/ρ
	dist =[]
	r_r=[]
	rmin_r=[]
	rmax_r=[]
	V_n =[]
	g_r=[]
	adot=[]
	V_r=[]
	S=[]

	for ir=1:Ir
		push!(r_r,r)
		push!(rmin_r, rmin)
		push!(rmax_r, rmax)

		global v = v0

		for n=1:maxit
		    V = v
		    push!(V_n,V)
		    #forward difference
		    dVf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/da
		    dVf[H,:] = (z .+ r.*amax).^(-s) #will never be used, but impose state constraint a<=amax just in case
		    #backward difference
		    dVb[2:H,:] = (V[2:H,:]-V[1:H-1,:])/da
		    dVb[1,:] = (z .+ r.*amin).^(-s) #state constraint boundary condition

		    I_concave = dVb .> dVf #indicator whether value function is concave (problems arise if this is not the case)

		    # consumption and savings with forward difference
		    cf = dVf.^(-1/s)
		    ssf = zz + r.*aa - cf
		    # consumption and savings with backward difference
		    cb = dVb.^(-1/s)
		    ssb = zz + r.*aa - cb
		    # consumption and derivative of value function at steady state
		    c0 = zz + r.*aa
		    dV0 = c0.^(-s)

		    #= dV_upwind makes a choice of forward or backward differences based on
		     the sign of the drift    =#
		    If = ssf .> 0 #positive drift → forward difference
		    Ib = ssb .< 0 #negative drift → backward difference
		    I0 = (1.0 .-If-Ib) #at steady state

		    dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0  #important to include third term
		    c = dV_Upwind.^(-1/s)
		    u = c.^(1-s)/(1-s)

		    # CONSTRUCT MATRIX
		    X = - min.(ssb,0)/da
		    Y = - max.(ssf,0)/da + min.(ssb,0)/da
		    Z = max.(ssf,0)/da

		    A1=spdiagm(0 => Y[:,1], -1 => X[2:H,1], 1 =>Z[1:H-1,1])
		    A2=spdiagm(0=> Y[:,2], -1 =>X[2:H,2], 1 => Z[1:H-1,2])

			global A = [A1 spzeros(H,H);spzeros(H,H) A2] + Aswitch

		    B = (1/Delta + ρ)*sparse(I, 2*H, 2*H) - A

		    u_stacked = [u[:,1];u[:,2]]
		    V_stacked = [V[:,1];V[:,2]]

		    b = u_stacked + V_stacked/Delta
		    V_stacked = B\b #SOLVE SYSTEM OF EQUATIONS

		    V = [V_stacked[1:H] V_stacked[H+1:2*H]]

		    V_change = V - v
		    v = V

		    push!(dist,findmax(abs.(V_change))[1])
		    if dist[n]<crit
		        println("Value Function Converged, Iteration = ")
		        println(n)
		        break
		    end
		end

	#===============================================================================

		Code for finding the terminal state of the KFE

	===============================================================================#
		AT = A'
		b = zeros(2*H,1)

		# need to fix one value, otherwise matrix is singular
		i_fix = 1
		b[i_fix]=.1
		row = [zeros(1,i_fix-1) 1 zeros(1,2*H-i_fix)]
		AT[i_fix,:] = row

		# Solve linear system
		gg = AT\b
		g_sum = gg'*ones(2*H,1)*da
		gg = gg./g_sum;

		global g = [gg[1:H] gg[H+1:2*H]]

		check1 = g[:,1]'*ones(H,1)*da
		check2 = g[:,2]'*ones(H,1)*da

		push!(g_r, g)
		push!(adot,zz + r.*aa - c)
		push!(V_r,V)

		push!(S,g[:,1]'*a*da + g[:,2]'*a*da)

			# UPDATE INTEREST RATE
			if S[ir][1] > crit_S
			    println("Excess Supply")
			    rmax = r
			    r = 0.5*(r+rmin)
			elseif S[ir][1]< -crit_S
			    println("Excess Demand")
			    rmin = r
			    r = 0.5*(r+rmax)
			elseif abs(S[ir][1])<crit_S
			    println("Equilibrium Found, Interest rate =")
			    println(r)
			    break
			end
	end

return g, v, r

end


using Interpolations
using Polynomials

function ConcreteStressStrainModel(eps_co,eps_cu,f_c,f_t)

    #define concrete model strain range
    eps_c=range(0,eps_cu,length=100)

    #fill initial stress vector
    sig_c=zeros(Float64,length(eps_c))

    #define model for stress-strain relationship
    pre=findall(x->(x<=eps_co), eps_c)
    sig_c[pre]=f_c*(2*eps_c[pre]/eps_co-(eps_c[pre]./eps_co).^2)
    post=findall(x->(x>eps_co), eps_c)
    sig_c[post].=f_c.+((0.15.*f_c)./(eps_co.-eps_cu)).*(eps_c[post].-eps_co)

    #calculate initial elastic modulus
    E_ci=(sig_c[2]-sig_c[1])/(eps_c[2]-eps_c[1])

    #convert signs to represent compression
    eps_c=-eps_c
    sig_c=-sig_c

    #define concrete tension stress-strain curve
    f_cr=f_t  #tensile modulus of rupture
    eps_cr=f_cr/E_ci #strain associated with modulus of rupture

    eps_c=vcat(eps_cr,eps_c)
    # eps_c=vcat(eps_cr*1.000001,eps_c)
    # eps_c=vcat(0.003,eps_c)
    sig_c=vcat(f_cr,sig_c)
    # sig_c=vcat([0, 0],sig_c)

    return eps_c,sig_c,E_ci,eps_cr

end

function SteelStressStrainModel(E_si,sig_sy,sig_su,sig_sf,eps_sy,eps_sy1,eps_su,esp_sf)

    # define elastic part of steel stress-strain curve
    sig_s=[0 sig_sy]
    eps_s=[0 eps_sy]

    #define plateau and hardening part of steel stress-strain curve
    p = fit([eps_sy1, eps_su, eps_sf],[sig_sy, sig_su, sig_sf],2)

    sig_fit = zeros(11)
    eps_range = eps_sy1:(eps_sf-eps_sy1)./10:eps_sf

    for i=1:11
    sig_fit[i] = p(eps_range[i])
    end

    eps_s=vcat(vec(eps_s),collect(range(eps_sy1,length=11,stop=eps_sf)))
    sig_s=vcat(vec(sig_s), sig_fit)

    # add compression part of steel stress-strain curve
    sig_s=vcat(-reverse(sig_s[2:end]),sig_s)
    eps_s=vcat(-reverse(eps_s[2:end]),eps_s)

    return eps_s, sig_s

end

function Response5000(M,P,h,s,w,A_s,z_s,E_ci,E_si,eps_cr,SteelStrain,SteelStress,ConcreteStrain,ConcreteStress)

    # define initial concrete locations in cross-section (from bottom, z=0)
    A_c=ones(s,1)*h/s*w
    z_c=collect(range(h/s*0.5,h-h/s*0.5,length=s))

    # initialize elastic modulus
    E_c=ones(s,1)*E_ci
    E_s=ones(length(A_s),1)*E_si

    SteelFit = LinearInterpolation(SteelStrain, SteelStress)
    ConcreteFit = LinearInterpolation(reverse(ConcreteStrain), reverse(ConcreteStress))

    #initialize stress, strain, and curvature
    eps_P=0
    eps_M_c=zeros(length(A_c),1)
    eps_M_s=zeros(length(A_s),1)
    rho_inv_hist=zeros(length(M))
    rho_inv=0

    sig_c_hist=zeros(length(A_c),1)
    sig_s_hist=zeros(length(z_s),1)
    eps_c_hist=zeros(length(A_c),1)
    eps_s_hist=zeros(length(z_s),1)

    for i=2:length(M)

        # E_c
        # E_s
        # eps_P
        # rho_inv

        # define load increment
        dP=P[i]-P[i-1]
        dM=M[i]-M[i-1]

        # calculate cross-section centroid
        z_bar=(sum(A_c.*z_c.*E_c)+sum(A_s.*z_s.*E_s))./((sum(A_c.*E_c)+sum(A_s.*E_s)))

        # calculate concrete and steel strains from P
        eps_P=eps_P+dP./(sum(E_s.*A_s)+sum(E_c.*A_c))

        # calculate 1 over the radius of curvature from M
        rho_inv=rho_inv.+dM./(sum(A_s.*E_s.*(z_s.-z_bar).^2).+sum(A_c.*E_c.*(z_c.-z_bar).^2))

        rho_inv_hist[i]=rho_inv

        # calculate strains from M
        eps_M_c=-(z_c.-z_bar).*rho_inv   #strain in concrete
        eps_M_s=-(z_s.-z_bar).*rho_inv   #strain in steel

        # combine axial and flexural strain, calculate strain history
        eps_c_hist=eps_M_c.+eps_P  #total strain in concrete
        eps_s_hist=eps_M_s.+eps_P  #total strain in steel

        index=findall(x->x<eps_cr,eps_c_hist)
        antindex=findall(x->x>=eps_cr,eps_c_hist)

        if abs(minimum(eps_c_hist))>abs(minimum(ConcreteStrain))
            println("The concrete has crushed, P=",P[i]," M=",M[i])
            return eps_c_hist,eps_s_hist,sig_c_hist,sig_s_hist,z_bar,z_c,rho_inv_hist
        end
        steelfailure=0
        for i=1:length(z_s)
            if abs(eps_s_hist[i])>maximum(SteelStrain)
                steelfailure=1
            else
                steelfailure=0
            end
        end

        if steelfailure==1
            println("The steel has reached its ultimate strain, P=",P[i], " M=",M[i])
            return eps_c_hist,eps_s_hist,sig_c_hist,sig_s_hist,z_bar,z_c,rho_inv_hist
        end

        #store previous stress and strain history
        sig_s_hist_old = sig_s_hist
        sig_c_hist_old = sig_c_hist
        eps_s_hist_old = eps_s_hist
        eps_c_hist_old = eps_c_hist

        # calculate concrete and steel stresses
        sig_c_hist[index]=ConcreteFit(eps_c_hist[index])
        if isempty(antindex)==false
             sig_c_hist[antindex].=0
        end
        sig_s_hist=SteelFit(eps_s_hist)

        #tangent modulus
        E_c=(sig_c_hist .- sig_c_hist_old)./(eps_c_hist .- eps_c_hist_old)
        E_s=(sig_s_hist .- sig_s_hist_old)./(eps_s_hist .- eps_s_hist_old)
        # update elastic modulus
        # E_c=sig_c_hist./eps_c_hist
        # E_s=sig_s_hist./eps_s_hist

    end
    return eps_c_hist,eps_s_hist,sig_c_hist,sig_s_hist,z_bar,z_c,rho_inv_hist, E_c, E_s
end

# define concrete cross-section
w=50  #width
h=50  #depth/height
s=50*4  #number of cross-section divisions

# define rebar area and location (measured from bottom, z=0) in
# cross-section
     # number size location
     Rebar=[10 8 2.5]

RebarQuantity=Rebar[:,1]
RebarDiameter=Rebar[:,2]/8
RebarLocation=Rebar[:,3]

RebarSpacing=w./RebarQuantity

A_s=RebarQuantity.*pi.*(RebarDiameter./2).^2
z_s=RebarLocation

#define concrete compression engineering stress-strain model
#modified Hognestad
eps_co=0.002        #strain at f'c
eps_cu=0.0035       #ultimate strain at brittle failure
f_c=6.00            #concrete compressive strength
f_t=1/10*f_c        #concrete tensile strength

ConcreteStrain,ConcreteStress,E_ci,eps_cr=ConcreteStressStrainModel(eps_co,eps_cu,f_c,f_t)

E_si=29000                  #steel elastic modulus
sig_sy=60                   #steel yield stress
sig_su=80                   #steel ultimate stress
sig_sf=sig_sy*(1+0.10)      #steel fracture stress#
eps_sy=sig_sy/E_si          #steel yield strain
eps_sy1=sig_sy/E_si*(10)    #steel strain at end of yield plateau
eps_su=0.20                 #steel ultimate strain
eps_sf=0.30                 #steel fracture strain

SteelStrain,SteelStress=SteelStressStrainModel(E_si,sig_sy,sig_su,sig_sf,eps_sy,eps_sy1,eps_su,eps_sf)

M=range(0,38798.0,length=100)
P=range(0, 0,length=100)

eps_c_hist,eps_s_hist,sig_c_hist,sig_s_hist,z_bar,z_c,rho_inv_hist, E_c, E_s=Response5000(M,P,h,s,w,A_s,z_s,E_ci,E_si,eps_cr,SteelStrain,SteelStress,ConcreteStrain,ConcreteStress)

plot(eps_c_hist, z_c)


using ChromNet

# project_groupgm
IC = [1 0.1 0.4 0.2; 0.1 1 0.3 0.5; 0.4 0.3 1 0.6; 0.2 0.5 0.6 1]
header = ["A1", "A2", "B1", "B2"]
groups = build_groups(inv(IC), header)
G,headerG = build_groupgm(IC, header, groups)

ans = [
    1.0 0.9 0.6 0.6;
    0.9 1.0 0.9 0.9;
    0.6 0.9 1.0 0.6;
    0.6 0.9 0.6 1.0
]
@test all((ChromNetPaper.project_groupgm(G, header, groups, x->true) .- ans) .< 1e-8)

ans = [
    1.0 0.1 0.6 0.6;
    0.1 1.0 0.8 0.8;
    0.6 0.8 1.0 0.6;
    0.6 0.8 0.6 1.0
]
@test all((ChromNetPaper.project_groupgm(G, header, groups, x->x[1] < 0.1) .- ans) .< 1e-8)

# run with default group filter
ChromNetPaper.project_groupgm(G, header, groups)


############################################################################################
## This file is part of RLinearAlgebra.jl
##
## Overview: abstractions and methods for routines used in solving linear systems of
## equations.
##
## Contents
## - Abstract Types
## - `rsubsolve!` function documentation
## - Vector Row Projection Solvers
## - Vector Column Projection Solvers
## - Block Row Projection Solvers
## - Block Column Projection Solvers
## - Preconditioned Krylov Solvers
## - Export Statements
##
############################################################################################

# Dependencies: LinearAlgebra

#############################################
# Abstract Types
#############################################
"""
    LinSysSolveRoutine

Abstract supertype that specifies the type of linear system solver routine being deployed.
"""
abstract type LinSysSolveRoutine end

"""
    LinSysVecRowProjection <: LinSysSolveRoutine

Abstract supertype for vector row action projection methods.
"""
abstract type LinSysVecRowProjection <: LinSysSolveRoutine end

"""
    LinSysVecColProjection <: LinSysSolveRoutine

Abstract supertype for vector column action projection methods.
"""
abstract type LinSysVecColProjection <: LinSysSolveRoutine end

"""
    LinSysBlkRowProjection <: LinSysSolveRoutine

Abstract supertype for block row action projection methods.
"""
abstract type LinSysBlkRowProjection <: LinSysSolveRoutine end

"""
    LinSysBlkColProjection <: LinSysSolveRoutine

Abstract supertype for block column action projection methods.
"""
abstract type LinSysBlkColProjection <: LinSysSolveRoutine end

"""
    LinSysPreconKrylov <: LinSysSolveRoutine

Abstract supertype for preconditioned Krylov solver.
"""
abstract type LinSysPreconKrylov <: LinSysSolveRoutine end

#############################################
# `rsubsolve!` Function Documentation`
#############################################
"""
    rsubsolve!(
        type::LinSysSolveRoutine,
        x::AbstractVector,
        samp::Tuple,
        iter::Int64)

A common interface for specifying different strategies for solving a subproblem generated
    by a sampling, selecting or sketching operation on a linear system. The `type` argument
    is used to select the appropriately defined strategy. The argument `x` is the current
    iterate value for the solution. The argument `samp` depends on the subtype
    of `LinSysSolveRoutine` that is being deployed and is described below. The `iter`
    argument is the iteration counter.

The value of `samp` depends on the type of subtype of `LinSysSolveRoutine` being deployed.
To describe this, let `A` be the coefficient matrix of the system, and `b` its constant
vector.
- For `T<:LinSysVecRowProjection`, `samp` is a two-dimensional tuple where the first
    entry is a vector in the row space of `A`; and the second entry is a scalar value
    corresponding to a linear combination of the elements of `b`.
- For `T<:LinSysVecColProjection`, `samp` is a three-dimensional tuple where the first
    entry is a vector of `length(x)` corresponding to the search direction; the second
    entry is a matrix with the same number of rows as `A` (usually it is `A`);
    and the third entry is a scalar-valued residual for the normal system corresponding
    to `samp[1]'* A' * (A * x - b)`.

The function `rsubsolve!` updates the quantity `x` and any fields of `type` that must be
    updated.
"""
function rsubsolve!(
    type::Nothing,
    x::AbstractVector,
    samp::Tuple,
    iter::Int64,
)
    return nothing
end

#############################################
# Vector Row Projection Solvers
#############################################
include("linear_solver_routines/vec_row_projection_std.jl")
include("linear_solver_routines/vec_row_projection_portho.jl")
include("linear_solver_routines/vec_row_projection_fortho.jl")

#############################################
# Vector Column Projection Solvers
#############################################
include("linear_solver_routines/vec_col_projection_std.jl")
include("linear_solver_routines/vec_col_projection_portho.jl")
include("linear_solver_routines/vec_col_projection_fortho.jl")

#############################################
# Block Row Projection Solvers
#############################################
# Exact Solvers

# Approximate Solvers

#############################################
# Block Column Projection Solvers
#############################################
# Exact Solvers

# Approximate Solvers

#############################################
# Precondition Krylov Solvers
#############################################
#Blendenpik

#############################################
# Export Statements
#############################################
# See RLinearAlgebra.jl
#export LinSysSolveRoutine, LinSysVecRowProjection, LinSysVecColProjection,
#    LinSysBlkRowProjection, LInSysBlkColProjection, LinSysPreconKrylov


using Test

using iLQGames:
    ProductSystem,
    Unicycle4D,
    n_states,
    n_controls,
    dx,
    linearize

using StaticArrays

# create two separate dynamics systems
unicycle1 = Unicycle4D{0.1}()
unicycle2 = Unicycle4D{0.1}()
xid1 = SVector{4}(1:4)
xid2 = SVector{4}(5:8)
uid1 = SVector{2}(1:2)
uid2 = SVector{2}(3:4)

# and combine them in the product system
dyn = ProductSystem((unicycle1, unicycle2))

# test some operations on the joint system
x = @SVector rand(n_states(dyn))
u = @SVector rand(n_controls(dyn))
t = 0.0;

@testset "dx" begin
    dx_dyn = dx(dyn, x, u, t)
    dx_unicycle1 = dx(unicycle1, x[xid1], u[uid1], t)
    dx_unicycle2 = dx(unicycle2, x[xid2], u[uid2], t)

    @test isapprox(dx_dyn[xid1], dx_unicycle1)
    @test isapprox(dx_dyn[xid2], dx_unicycle2)
end

@testset "linearize" begin
    lin_dyn = linearize(dyn, x, u, t)
    lin_unicycle1 = linearize(unicycle1, x[xid1], u[uid1], t)
    lin_unicycle2 = linearize(unicycle2, x[xid2], u[uid2], t)

    # linearizing the combined system should be the same as linearizing the
    # subsystems separately
    @test isapprox(lin_dyn.A[xid1, xid1], lin_unicycle1.A)
    @test isapprox(lin_dyn.A[xid2, xid2], lin_unicycle2.A)
    @test isapprox(lin_dyn.B[xid1, uid1], lin_unicycle1.B)
    @test isapprox(lin_dyn.B[xid2, uid2], lin_unicycle2.B)
end


using Unitful, UnitfulAngles
import Dates
using Test

units = (u"doubleTurn", u"turn", u"halfTurn", u"quadrant", u"sextant", u"octant", u"clockPosition", u"hourAngle", u"compassPoint", u"hexacontade", u"brad", u"°", u"grad", u"arcminute", u"arcsecond", u"rad", u"diameterPart")
quantities = (0.5, 1, 2, 4, 6, 8, 12, 24, 32, 60, 256, 360, 400, 21600, 1296000, 2π, 120π)

@test all(1u"turn" ≈ q*u for (q, u) in zip(quantities, units))
for _f in (:sin, :cos, :tan, :sec, :csc, :cot), (q, u) in zip(quantities, units), a in 13:17
    @test @eval $_f(2π/$a) ≈ $_f($q*$u/$a)
end

for (_f, _x) in zip((:asin, :acos, :atan, :asec, :acsc, :acot), (.5, √3/2, √3/3, 2/√3, 2, √3))
    @test @eval $_f(u"°", $_x) ≈ 30u"°"
end

@test atan(u"°", 1,1) == 45u"°"

@test convert(Dates.Time, 45u"°") == Dates.Time(3,0,0)

@test convert(u"°", Dates.Time(12,0,0)) == 180u"°"

####################
# for _f in (:sin, :cos), (q, u) in zip(qrats, rats)
    # @test @eval $(Symbol("$(_f)pi"))(2) === $_f($q*$u)
# end


# Hidden Markov models

# Copyright (c) 2016 Jeffrey W. Miller
# This software is released under the MIT License.
# If you use this software in your research, please cite:
#
#   Jeffrey W. Miller (2016). Lecture Notes on Advanced Stochastic Modeling. Duke University, Durham, NC.

# HMM with discrete emission distributions, assumed to be on {1,...,K}.
module DiscreteHMM
	using Distributions

	# Emission distribution specification
	log_q(x,p) = log(p[x])  # log density
	qrnd(p) = rand(Categorical(p))  # random sample

	# Initialize an array of emission distribution parameters
	function phi_init(m,x)
		K = maximum(x)
	    return [rand(Dirichlet(K,1)) for s = 1:m]
	end

	# Maximum likelihood estimate of array of emission distribution
	# parameters using data x with weights gamma.
	function phi_max!(phi,x,gamma)
	    m = length(phi)
	    n = length(x)
	    for s = 1:m
	    	ph = phi[s]
	    	ph[:] = 0.0
	        for j = 1:n
	            ph[x[j]] += gamma[j,s]
	        end
	        ph[:] /= sum(ph)
	    end
	end

	# Core HMM code
	include("hmm.jl")
end # module

# HMM with univariate normal emission distributions.
module NormalHMM
	using Distributions

	# Emission distribution specification
	log_q(x,D) = logpdf(D,x)  # log density
	qrnd(D) = rand(D)  # random sample

	# Initialize an array of emission distribution parameters
	function phi_init(m,x)
		D_m = Normal(mean(x),std(x))
		D_v = InverseGamma(1.0,var(x))
	    return [Normal(rand(D_m),sqrt(rand(D_v))) for s = 1:m]
	end

	# Maximum likelihood estimate of array of emission distribution
	# parameters using data x with weights gamma.
	function phi_max!(phi,x,gamma)
	    m = length(phi)
	    n = length(x)
	    for s = 1:m
	         weights = gamma[:,s] / sum(gamma[:,s])
             mu = dot(weights,x)
             sigma = sqrt(dot(weights,(x-mu).*(x-mu)))
             phi[s] = Normal(mu,sigma)
	    end
	end

	# Core HMM code
	include("hmm.jl")
end # module

# HMM with multivariate normal emission distributions,
# with diagonal covariance matrices.
module MVNdiagHMM
	using Distributions

	# Emission distribution specification
	log_q(x,D) = logpdf(D,x)  # log density
	qrnd(D) = rand(D)  # random sample

	# Initialize an array of emission distribution parameters
	function phi_init(m,x)
		d = length(x[1])  # dimension
		mu = mean(x)
		sigma = Float64[std(Float64[xi[j] for xi in x]) for j=1:d]
		randsig(d) = sqrt(rand(InverseGamma(1,1),d))
		phi = [MvNormal(randn(d).*sigma+mu, randsig(d).*sigma) for s=1:m]
	end

	# Maximum likelihood estimate of array of emission distribution
	# parameters using data x with weights gamma.
	function phi_max!(phi,x,gamma)
	    m = length(phi)
	    n = length(x)
		d = length(x[1])  # dimension
	    for s = 1:m
	         weights = gamma[:,s] / sum(gamma[:,s])
	         mu = zeros(d)
	         for j = 1:n
	         	mu += weights[j]*x[j]
	         end
	         sigma2 = zeros(d)
	         for j = 1:n
	         	sigma2 += weights[j]*(x[j]-mu).*(x[j]-mu)
	         end
             phi[s] = MvNormal(mu,sqrt(sigma2))
	    end
	end

	# Core HMM code
	include("hmm.jl")
end # module


using DelimitedFiles

function neighbours(mat, i, j)
	t = vec(mat[max(i-1,1):min(i+1,size(mat,1)), max(j-1,1):min(j+1,size(mat,2))])
	count(c->c=='#', t) - ((mat[i,j] == '#') ? 1 : 0)
end

function visible(mat, i, j)
	n = size(mat,1)
	m = size(mat,2)
	dirs = [[1,0],[1,1],[0,1],[-1,1],[-1,0],[-1,-1],[0,-1],[1,-1]]
	c = 0
	for dir in dirs
		pos = [i,j]
		while true
			pos += dir
			if pos[1] < 1 || pos[2] < 1 || pos[1] > n || pos[2] > m
				break
			elseif mat[pos...] == '#'
				c += 1
				break
			elseif mat[pos...] == 'L'
				break
			end
		end
	end
	c
end

function x(func, num)
	t = vec(readdlm("11.txt", String))
	n = length(t)
	m = length(t[1])
	mat = Array{Char,2}(undef, n,m)
	for i=1:n, j=1:m
		mat[i,j] = t[i][j]
	end

	while true
		mat2 = copy(mat)
		for i=1:n, j=1:m
			if mat[i,j] == 'L' && func(mat, i, j) == 0
				mat2[i,j] = '#'
			elseif mat[i,j] == '#' && func(mat, i, j) >= num
				mat2[i,j] = 'L'
			end
		end
		if mat == mat2
			break
		end
		mat=copy(mat2)
		for i=1:n, j=1:m
			if mat2[i,j] == 'L' && func(mat2, i, j) == 0
				mat[i,j] = '#'
			elseif mat2[i,j] == '#' && func(mat2, i, j) >= num
				mat[i,j] = 'L'
			end
		end
		if mat == mat2
			break
		end
	end
	count(c->c=='#', mat)
end
#part 1
println(x(neighbours, 4))
#part 2
println(x(visible, 5))


#= note: to run this in parallel, from the system
prompt, and before starting julia, do
export JULIA_NUM_THREADS=X
where X is the desired number of threads. On
my 32 core system, 10 is good (cache contention issues)
Remember to re-set threads to 1 before using MPI
=#
using Statistics, Econometrics, LinearAlgebra, Plots
include("QIVmodel.jl")
function main()
LNW, X, Z = getdata()
n = size(LNW,1)
mcmcreps = 100000
burnin = 10000
# use st. errs. from ordinary QR as quide to setting tuning.
basetuning = [0.1, 0.05, 0.03, 0.005, 0.3, 0.2, 0.2]
# adjust this up or down to achieve decent acceptance rate
scale = 0.1 # higher for higher taus
tuning = scale*basetuning
EducEffect = zeros(9,3)
Constant = zeros(9,3)
for i = 1:9
    τ = round(i/10,digits=1)
    Σ = τ*(1.0-τ)*Z'Z/n
    Σinv = inv(Σ)
    θ = X\LNW  # OLS start values
    # to do ordinary QR via MCMC, set Z=X (just to verify that MCMC works!)
    lnL = θ -> likelihood(θ, LNW, X, Z, τ, Σinv)
    Proposal = θ -> proposal(θ, tuning)
    Prior = θ -> prior(θ)
    chain = mcmc(θ, mcmcreps, burnin, Prior, lnL, Proposal)
    V = cov(chain[:,1:end-1])
    d = dstats(chain)
    Constant[i,:] = d[1,[1,7,8]] # save posterior mean, 5% and 95% quantile
    EducEffect[i,:] = d[2,[1,7,8]] # save posterior mean, 5% and 95% quantile
    #scatter(chain[:,1],chain[:,2], show=true, reuse=false)
end

# plot results
# educ
τ = (1:9)/10
plot(reuse=false)
educ = EducEffect[:,1]
lb = EducEffect[:,2]
ub = EducEffect[:,3]
plot!(τ, [educ educ], fillrange=[lb ub], fillalpha=0.3, c=:green, legend=:none)
xticks!((1:9)/10)
#savefig("EducJUNK.png")
# constant
τ = (1:9)/10
plot(reuse=false)
c = Constant[:,1]
lb = Constant[:,2]
ub = Constant[:,3]
plot!(τ, [c c], fillrange=[lb ub], fillalpha=0.3, c=:green, legend=:none)
xticks!((1:9)/10)
#savefig("ConstantJUNK.png")
return
end
main()


# EXCLUDE FROM TESTING
# NOTE: This is temporarily not working until we merge the DMSTAG routines with the new Clang branch
#
# This implements src/snes/examples/tutorials/ex2.c from PETSc using the PETSc.jl package, using SNES
#
# This is the same as SNES_ex2b.j, except that we show how automatic differentiation can be used to
# compute the jacobian.
#
# Newton method to solve u'' + u^{2} = f, sequentially.

using PETSc, MPI, LinearAlgebra, SparseArrays, Plots, ForwardDiff

if ~MPI.Initialized()
    MPI.Init()
end

PETSc.initialize()

```
    Computes initial guess
```
function FormInitialGuess!(dm,x)

    x_Local = PETSc.DMCreateLocalVector(dm);
    x_array = PETSc.DMStagVecGetArray(dm,x_Local);
    x_array[1:2,1:2] .= 1;
    PETSc.DMLocalToGlobal(dm, x_Local, PETSc.INSERT_VALUES, x)
end

```
    Computes the residual f, given solution vector x
```
function FormResidual!(f,x)
    n       =   length(x);
    xp      =   LinRange(0.0,1.0, n);
    F       =   6.0.*xp .+ (xp .+1.e-12).^6.0;      # define source term function

    dx      =   1.0/(n-1.0);
    f[1]    =   x[1] - 0.0;
    for i=2:n-1
        f[i] = (x[i-1] - 2.0*x[i] + x[i+1])/dx^2 + x[i]*x[i] - F[i]
    end
    f[n]    =   x[n] - 1.0;

end

```
    Wrapper which makes it easier to compute the jacobian using automatic differntiation
```
function  ForwardDiff_routine(x)

    f   = zero(x)               # vector of zeros, of same type as e
    FormResidual!(f,x);

    return f;
end

```
    This copies a julia sparse matrix to PETSc MatSeqAIJ format
```
function Mat_JuliaToPETSc!(J::PETSc.MatSeqAIJ, J_julia::SparseMatrixCSC)

    for i = 1:size(J_julia,1)
        col = J_julia[i,:];
        row = ones(Int32,length(col.nzind))*i;
        for j=1:length(col.nzind)
            J[i, col.nzind[j]] = col.nzval[j];
        end
    end
    PETSc.assemble(J);  # finalize assembly

end

```
    Computes the jacobian, given solution vector x
```
function FormJacobian!(x, args...)

    J        =   args[1];        # preconditioner = args[2], in case we want it to be different from J

    # Use AD to compute jacobian; by transferring x into sparse, the output will be sparse
    J_julia  =   ForwardDiff.jacobian(ForwardDiff_routine,sparse(x));

    if typeof(J) <: PETSc.AbstractMat
        Mat_JuliaToPETSc!(J, J_julia);  # transfer julia sparse matrix 2 petsc
    else
        J .= J_julia;
    end
end

# ==========================================
# Main code

# Compute initial solution
nx   =   5;
nz   =   5;
x0   =   -5;
xend =    5;
z0   = - 10;
zend =    0;

dm  = PETSc.DMStagCreate2d(MPI.COMM_SELF,PETSc.DM_BOUNDARY_NONE,PETSc.DM_BOUNDARY_NONE,nx,nz,1,1,0,0,1,PETSc.DMSTAG_STENCIL_BOX,2);
PETSc.DMStagSetUniformCoordinatesExplicit(dm, x0, xend, z0, zend);
x   =   PETSc.DMCreateGlobalVector(dm);

#FormInitialGuess!(dm,x);

# Compute initial jacobian using a julia structure to obtain the nonzero structure
# Note that we can also obtain this structure in a different manner
#Jstruct  = zeros(n,n);
#FormJacobian!(x, Jstruct);                              # jacobian in julia form
#Jsp      =   sparse(Float64.(abs.(Jstruct) .> 0))       # sparse julia, with 1.0 in nonzero spots
#PJ       =   PETSc.MatSeqAIJ(Jsp);                      # transfer to PETSc (initialize matrix with correct nonzero pattern)

# Setup SNES
#x_s = PETSc.VecSeq(x);                  # solution vector
#res = PETSc.VecSeq(zeros(size(x)));     # residual vector

#S = PETSc.SNES{Float64}(MPI.COMM_SELF;
#        snes_rtol=1e-12,
#        snes_monitor=true,
#        snes_converged_reason=true);
#PETSc.setfunction!(S, FormResidual!, res)
#PETSc.setjacobian!(S, FormJacobian!, PJ, PJ)

# solve
#PETSc.solve!(x_s, S);

# Extract & plot solution
#x_sol = x_s.array;                  # convert solution to julia format
#FormResidual!(res.array,x_sol)      # just for checking, compute residual
#@show norm(res.array)

#plot(LinRange(0,1,n),x_sol,xlabel="width",ylabel="solution")

#PETSc.finalize()