using LinearAlgebra, DataFrames,Optim, ForwardDiff, BenchmarkTools,Distributions,
    Expectations, QuantEcon, Statistics, GLM
using Distributions: invsqrt2π, log2π, sqrt2, invsqrt2
using DCDC
using Test
using Distributed, Suppressor
using Plots

begin "Dynamic Decision Process"
    σ₀ = 2;
    β = 0.8;
    nM = 50;
    nT = 5;
    ddc=DynamicDecisionProcess(2,0.8);
    computeEquilibrium(ddc);
    checked = [check_ee(ddc) for i = 1:100];
    plot(ddc.ValueFn.xdata,ddc.ValueFn.y);
    plot!(ddc.PolicyFn.xdata,ddc.PolicyFn.y);
    data = simulate_ddc(nM,nT,ddc);
end

# The moment condition is
# u'(a_t) = β d_Trans(a_t,x_t) u'(a_t+1)

reshaped_data = convert_data(data);

using Flux #Get derivatives
using Flux.Tracker
using Flux.Tracker: update!

θ₀ = ones(2)  #θ = (β,σ)
Transition = ddc.trans;
dtrans = (x,c) -> Tracker.gradient(Transition,x,c);
dtransx = (x,c) -> Tracker.gradient(Transition,x,c)[1];
R = dtransx.(reshaped_data.action_t,reshaped_data.state_t);
θ = θ₀

lb = [0,0];
ub = [1,100];
function _moment(θ)
    (β₀,σ₀ ) = θ₀;
    util = Utility(σ₀);
    dutil = (x) -> Tracker.gradient(util,x)[1].data;
    dudc  = dutil.(reshaped_data.action_t);
    dudc1 = dutil.(reshaped_data.action_t_1);
    η = dudc - β₀ .* dudc1 .* R;
    return(η'*η)
end

function _moment2(σ₀)
    dutil = (x) -> Tracker.gradient(Utility(σ₀);,x)[1].data;
    dudc  = dutil.(reshaped_data.action_t);
    dudc1 = dutil.(reshaped_data.action_t_1);
    η = dudc - 0.8 .* dudc1 .* R;
    return(η'*η)
end

optimize(_moment,zeros(2),BFGS(),autodiff=:forward)


using MPI
using StaticArrays
using ClimateMachine
using ClimateMachine.VariableTemplates
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.MPIStateArrays
using ClimateMachine.DGMethods
using Printf

using ClimateMachine.BalanceLaws

import ClimateMachine.BalanceLaws: vars_state, nodal_init_state_auxiliary!

using ClimateMachine.Mesh.Geometry: LocalGeometry

if !@isdefined integration_testing
    const integration_testing = parse(
        Bool,
        lowercase(get(ENV, "JULIA_CLIMA_INTEGRATION_TESTING", "false")),
    )
end

struct GradTestModel{dim, dir} <: BalanceLaw end

vars_state(m::GradTestModel, ::Auxiliary, T) = @vars begin
    a::T
    ∇a::SVector{3, T}
end
vars_state(::GradTestModel, ::Prognostic, T) = @vars()

function nodal_init_state_auxiliary!(
    ::GradTestModel{dim, dir},
    aux::Vars,
    tmp::Vars,
    g::LocalGeometry,
) where {dim, dir}
    x, y, z = g.coord
    if dim == 2
        aux.a = x^2 + y^3 - x * y
        if dir isa EveryDirection
            aux.∇a = SVector(2 * x - y, 3 * y^2 - x, 0)
        elseif dir isa HorizontalDirection
            aux.∇a = SVector(2 * x - y, 0, 0)
        elseif dir isa VerticalDirection
            aux.∇a = SVector(0, 3 * y^2 - x, 0)
        end
    else
        aux.a = x^2 + y^3 + z^2 * y^2 - x * y * z
        if dir isa EveryDirection
            aux.∇a = SVector(
                2 * x - y * z,
                3 * y^2 + 2 * z^2 * y - x * z,
                2 * z * y^2 - x * y,
            )
        elseif dir isa HorizontalDirection
            aux.∇a = SVector(2 * x - y * z, 3 * y^2 + 2 * z^2 * y - x * z, 0)
        elseif dir isa VerticalDirection
            aux.∇a = SVector(0, 0, 2 * z * y^2 - x * y)
        end
    end
end

using Test
function test_run(mpicomm, dim, direction, Ne, N, FT, ArrayType)
    connectivity = dim == 3 ? :full : :face
    brickrange = ntuple(j -> range(FT(0); length = Ne[j] + 1, stop = 3), dim)
    topl = StackedBrickTopology(
        mpicomm,
        brickrange,
        periodicity = ntuple(j -> false, dim),
        connectivity = connectivity,
    )

    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
    )

    model = GradTestModel{dim, direction}()
    dg = DGModel(
        model,
        grid,
        nothing,
        nothing,
        nothing;
        state_gradient_flux = nothing,
    )

    exact_aux = copy(dg.state_auxiliary)

    auxiliary_field_gradient!(
        model,
        dg.state_auxiliary,
        (:∇a,),
        dg.state_auxiliary,
        (:a,),
        grid,
        direction,
    )

    # Wrapping in Array ensure both GPU and CPU code use same approx
    approx = Array(dg.state_auxiliary.∇a) ≈ Array(exact_aux.∇a)
    err = euclidean_distance(exact_aux, dg.state_auxiliary)
    return approx, err
end

let
    ClimateMachine.init()
    ArrayType = ClimateMachine.array_type()

    mpicomm = MPI.COMM_WORLD

    numelem = (5, 5, 5)

    expected_result = Dict()
    expected_result[Float64, 2, 1] = 6.2135207410935696e+00
    expected_result[Float64, 2, 2] = 2.3700094936518794e+00
    expected_result[Float64, 2, 3] = 8.7105082013050261e-01
    expected_result[Float64, 2, 4] = 3.1401219279927611e-01

    expected_result[Float64, 3, 1] = 7.9363467666175236e+00
    expected_result[Float64, 3, 2] = 2.8059223082616098e+00
    expected_result[Float64, 3, 3] = 9.9204334582727094e-01
    expected_result[Float64, 3, 4] = 3.5074028853276679e-01

    expected_result[Float32, 2, 1] = 6.2135176658630371e+00
    expected_result[Float32, 2, 2] = 2.3700129985809326e+00
    expected_result[Float32, 3, 1] = 7.9363408088684082e+00
    expected_result[Float32, 3, 2] = 2.8059237003326416e+00

    @testset for FT in (Float64, Float32)
        lvls =
            integration_testing || ClimateMachine.Settings.integration_testing ?
            (FT === Float32 ? 2 : 4) : 1

        @testset for polynomialorder in ((4, 4), (4, 0))
            @testset for dim in 2:3
                @testset for direction in (
                    EveryDirection(),
                    HorizontalDirection(),
                    VerticalDirection(),
                )
                    err = zeros(FT, lvls)
                    for l in 1:lvls
                        approx, err[l] = test_run(
                            mpicomm,
                            dim,
                            direction,
                            ntuple(j -> 2^(l - 1) * numelem[j], dim),
                            polynomialorder,
                            FT,
                            ArrayType,
                        )
                        if polynomialorder[end] != 0 ||
                           direction isa HorizontalDirection
                            @test approx
                        else
                            @test err[l] ≈ expected_result[FT, dim, l]
                        end
                    end
                    if polynomialorder[end] != 0 ||
                       direction isa HorizontalDirection
                        @info begin
                            msg = "Polynomial order = $polynomialorder, direction = $direction\n"
                            for l in 1:(lvls - 1)
                                rate = log2(err[l]) - log2(err[l + 1])
                                msg *= @sprintf(
                                    "\n  rate for level %d = %e\n",
                                    l,
                                    rate
                                )
                            end
                            msg
                        end
                    end
                end
            end
        end
    end
end

nothing


function igyp00()
#
#  generalization of Igarash and Ypma
#
    p = poly([2.35;2.37;2.39])
    z = [2.35 1;2.37 1;2.39 1];
    p, PolyZeros(z)
end


""" Cálculo de movimento retilíneo
# Prepare um algoritmo para calcular
# o espaço percorrido por um móvel em movimento retilíneo uniforme dada a seguinte fórmula:
# S=So+V*T.
# Inicialmente o algoritmo deverá perguntar ao usuário
# os valores do espaço inicial(So),
# da velocidade(V)
# e do tempo(T).
# Após a entrada dos devidos valores estes devem ser substituídos na fórmula,
# resultando assim no espaço percorrido(S) que deverá ser informado na tela do usuário.
"""

println("informe o valor da velocidade")
mru = parse(Float32, readline())

println("informe o valor do tempo")
mru *= parse(Float32, readline())

println("informe o valor do espaço inicial")
mru += parse(Float32, readline())

println("espaço percorrido = ", mru)


# Extract of Combinatorics.jl
#
# You may (or may not!) want to call the function `permutations(a, t)` in your
# solution. If you would like to use it, include this file with
# `include("permutations.jl")`.
#
# License:
#
#    Copyright (c) 2013-2015: Alessandro Andrioni, Jiahao Chen and other
#    contributors.
#
#    Permission is hereby granted, free of charge, to any person obtaining a
#    copy of this software and associated documentation files (the "Software"),
#    to deal in the Software without restriction, including without limitation
#    the rights to use, copy, modify, merge, publish, distribute, sublicense,
#    and/or sell copies of the Software, and to permit persons to whom the
#    Software is furnished to do so, subject to the following conditions:
#
#    The above copyright notice and this permission notice shall be included in
#    all copies or substantial portions of the Software.
#
#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
#    THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#    DEALINGS IN THE SOFTWARE.

struct Permutations{T}
    a::T
    t::Int
end

Base.eltype(::Type{Permutations{T}}) where {T} = Vector{eltype(T)}

Base.length(p::Permutations) = (0 <= p.t <= length(p.a)) ? factorial(length(p.a), length(p.a)-p.t) : 0

"""
    permutations(a)

Generate all permutations of an indexable object `a` in lexicographic order. Because the number of permutations
can be very large, this function returns an iterator object.
Use `collect(permutations(a))` to get an array of all permutations.
"""
permutations(a) = Permutations(a, length(a))

"""
    permutations(a, t)

Generate all size `t` permutations of an indexable object `a`.
"""
function permutations(a, t::Integer)
    if t < 0
        t = length(a) + 1
    end
    Permutations(a, t)
end

function Base.iterate(p::Permutations, s = collect(1:length(p.a)))
    (!isempty(s) && max(s[1], p.t) > length(p.a) || (isempty(s) && p.t > 0)) && return
    nextpermutation(p.a, p.t ,s)
end

function nextpermutation(m, t, state)
    perm = [m[state[i]] for i in 1:t]
    n = length(state)
    if t <= 0
        return(perm, [n+1])
    end
    s = copy(state)
    if t < n
        j = t + 1
        while j <= n &&  s[t] >= s[j]; j+=1; end
    end
    if t < n && j <= n
        s[t], s[j] = s[j], s[t]
    else
        if t < n
            reverse!(s, t+1)
        end
        i = t - 1
        while i>=1 && s[i] >= s[i+1]; i -= 1; end
        if i > 0
            j = n
            while j>i && s[i] >= s[j]; j -= 1; end
            s[i], s[j] = s[j], s[i]
            reverse!(s, i+1)
        else
            s[1] = n+1
        end
    end
    return (perm, s)
end


using Distributions
using Base.Test

mu = [1.0, 2.0, 3.0]
C = [4. -2. -1.; -2. 5. -1.; -1. -1. 6.]
h = mu
J = C
mv_test_distributions = [
	                 Dirichlet(3, 2.0),
	                 Dirichlet([2.0, 1.0, 3.0]),
	                 IsoNormal(mu, 2.0),
	                 DiagNormal(mu, [1.5, 2.0, 2.5]),
	                 MvNormal(mu, C),
	                 IsoNormalCanon(h, 2.0),
	                 DiagNormalCanon(h, [1.5, 2.0, 1.2]),
	                 MvNormalCanon(h, J)]

for d in mv_test_distributions
    @test size(d) == size(rand(d))
    @test length(d) == length(rand(d))
end


# 2D FWI on Overthrust model using minConf library
# Author: Philipp Witte, pwitte@eoas.ubc.ca
# Date: May 2017
#

using PyCall, PyPlot, HDF5, opesciSLIM.TimeModeling, opesciSLIM.SLIM_optim, SeisIO

# Load starting model
~isfile("overthrust_2D_initial_model.h5") && run(`wget ftp://ftp.slim.gatech.edu/data/SoftwareRelease/WaveformInversion.jl/2DFWI/overthrust_2D_initial_model.h5`)
n,d,o,m0 = read(h5open("overthrust_2D_initial_model.h5","r"), "n", "d", "o", "m0")

# Set up model structure
model0 = Model((n[1],n[2]), (d[1],d[2]), (o[1],o[2]), m0)

# Bound constraints
v0 = sqrt.(1./model0.m)
vmin = ones(Float32,model0.n) + 0.3f0
vmax = ones(Float32,model0.n) + 5.5f0
vmin[:,1:21] = v0[:,1:21]	# keep water column fixed
vmax[:,1:21] = v0[:,1:21]

# Slowness squared [s^2/km^2]
mmin = vec((1f0./vmax).^2)
mmax = vec((1f0./vmin).^2)

# Load data
block = segy_read("/scratch/slim/pwitte/overthrust2D/overthrust_2D.segy")
dobs = joData(block)

# Set up wavelet
src_geometry = Geometry(block; key="source", segy_depth_key="SourceDepth")
wavelet = source(src_geometry.t[1],src_geometry.dt[1],0.008f0)	# 8 Hz wavelet
q = joData(src_geometry,wavelet)

############################### FWI ###########################################

# Optimization parameters
srand(1)	# set seed of random number generator
niterations = 20
batchsize = 20
proj(x) = reshape(median([vec(mmin) vec(x) vec(mmax)],2),model0.n)	# bound projection operator
fTerm = 1f-2
gradTerm = 1f-4
num_feval = 0
max_feval = 100
fval = 0
fvals = zeros(Float32,niterations)
dsub = []
qsub = []

# evaluate FWI misfit function for current model
function misfit(x)
	global num_feval; num_feval += 1
	model_new = deepcopy(model0)
	model_new.m = proj(reshape(x,model0.n))
	return f,g = fwi_objective(model_new,qsub,dsub)
end

# Backtracking line search with Armijo-Goldstein condition
function backtracking_linesearch(objective, x, grad; f_curr=[], alpha=1f0, tau=.5f0, c=.5f0, maxiter=10)

	# Parameters
	p = -grad/norm(grad)	# normalized descent direction
	m = dot(p,grad)
	t = -c*m
	isempty(f_curr) && (f_curr = objective(x)[1])	# current function value
	f_new = objective(x + alpha*p)[1]	# function value with initial step size
	iteration = 0

	# Line search
	while(f_curr - f_new < alpha*t && iteration < maxiter)
		println("	Iter LS: ", iteration, "; ",f_curr - f_new," >= ", alpha*t, "; alpha: ",alpha)
		alpha *= tau
		f_new = objective(x + alpha*p)[1]
		iteration += 1
	end
	return alpha*p
end

# Line search with armijo rule and curvature
function linesearch_wolfe(objective, x, grad_curr; f_curr=[], alpha=1f-4, tau=10f0, c1=1f-4, c2=.1f0, maxiter=10)

	# Parameters
	p = -grad_curr/norm(grad_curr)	# normalized descent direction
	isempty(f_curr) && ((f_curr, grad_curr) = objective(x))	# current function value and gradient
	f_new, grad_new = objective(x + alpha*p)	# function value and gradient with initial step size
	iteration = 0

	# line search with wolfe conditions
	while(f_new > f_curr + c1*alpha*dot(p,grad_curr) || -dot(p,grad_new) > -c2*dot(p,grad_curr) && iteration < maxiter)

		println("	Iter LS: ", iteration, "; ",f_new," <= ", f_curr + c1*alpha*dot(p,grad_curr),"; ",
		        -dot(p,grad_new), " <= ", -c2*dot(p,grad_curr), "; alpha: ",alpha)
		alpha *= 5f0
		f_new, grad_new = objective(x + alpha*p)
		iteration += 1
	end
	return alpha*p
end

# Main loop
for j=1:niterations

	# select current batch
	idx = randperm(dobs.nsrc)[1:batchsize]
	dsub = subsample(dobs,idx)
	qsub = subsample(q,idx)

	# get fwi objective function value and gradient
	fval, gradient = fwi_objective(model0,qsub,dsub)
	println("FWI iteration no: ",j,"; function value: ",fval)

	# linesearch
	step = backtracking_linesearch(misfit, vec(model0.m), gradient; f_curr=fval, alpha=10.f0, tau=.1f0, c=1e-4)
	#step = linesearch_wolfe(misfit, vec(model0.m), gradient; f_curr=fval, alpha=1f-4, tau=5f0, c1=1f-4, c2=.9f0)

	# Update model and bound projection
	model0.m = proj(model0.m + reshape(step,model0.n))

	# termination criteria
	fvals[j] = fval
	#if fval <= fTerm || norm(gradient) <= gradTerm || num_feval > max_feval
	#	break
	#end

end

# Save results
write(h5open("results_linesearch_armijo.h5", "w"), "x", sqrt.(1./reshape(x,model0.n)), "fsave", fsave, "fhistory",fvals)
# write(h5open("results_linesearch_wolfe.h5", "w"), "x", sqrt.(1./reshape(x,model0.n)), "fsave", fsave, "fhistory",fvals)


module RulesWishartMarginalsTest

using Test
using ReactiveMP
using Random
using Distributions

import ReactiveMP: @test_marginalrules

@testset "marginalrules:Wishart" begin

    @testset ":out_ν_S (m_out::Wishart, m_ν::PointMass, m_S::PointMass)" begin

        @test_marginalrules [ with_float_conversions = true, float32_atol = 1e-5 ] Wishart(:out_ν_S) [
            (
                input = (m_out = Wishart(3.0, [ 3.0 -1.0; -1.0 4.0 ]), m_ν = PointMass(2.0), m_S = PointMass([ 1.0 0.0; 0.0 1.0 ])),
                output = (out = Wishart(2.0, [ 14/19 -1/19; -1/19 15/19 ]), ν = PointMass(2.0), S = PointMass([ 1.0 0.0; 0.0 1.0 ]))
            ),
            (
                input = (m_out = Wishart(7.0, [ 9.0 -2.0; -2.0 1.0 ]), m_ν = PointMass(4.0), m_S = PointMass([ 4.0 -2.0; -2.0 4.0 ])),
                output = (out = Wishart(8.0, [ 128/49 -34/49; -34/49 32/49 ]), ν = PointMass(4.0), S = PointMass([ 4.0 -2.0; -2.0 4.0 ]))
            ),
            (
                input = (m_out = Wishart(4.0, [ 9.0 -2.0 1.0; -2.0 5.0 -2.0; 1.0 -2.0 11.0 ]), m_ν = PointMass(3.0), m_S = PointMass([ 11.0 -2.0 1.0; -2.0 5.0 -2.0; 1.0 -2.0 9.0 ])),
                output = (out = Wishart(3.0, [ 2092/423 -1/1 211/423; -1/1 5/2 -1/1; 211/423 -1/1 2092/423 ]), ν = PointMass(3.0), S = PointMass([ 11.0 -2.0 1.0; -2.0 5.0 -2.0; 1.0 -2.0 9.0 ]))
            ),
        ]

    end

end

end


abstract type AtomicStructure end

struct BCC <: AtomicStructure end
struct FCC <: AtomicStructure end

# TODO: introduce type Slipsystem instead of using NTuple

#slip planes and slip directions for BCC
function _get_slipsystems(::FCC)
    slip_planes = [ [ 1., 1., 1.],
                    [ 1., 1.,-1.],
                    [ 1.,-1.,-1.],
                    [ 1.,-1., 1.]]
    slip_directions = [ ([ 1., 0.,-1.], [ 0., 1.,-1.], [ 1.,-1., 0.]);
                        ([ 0., 1., 1.], [ 1., 0., 1.], [ 1.,-1., 0.]);
                        ([ 1., 0., 1.], [ 0., 1.,-1.], [ 1., 1., 0.]);
                        ([ 0., 1., 1.], [ 1., 0.,-1.], [ 1., 1., 0.]) ]
    return slip_planes, slip_directions
end

function _get_slipsystems(::BCC)
    slip_planes = [ [ 1., 1., 0.],
                    [ 1.,-1., 0.],
                    [ 1., 0., 1.],
                    [-1., 0.,-1.],
                    [ 0., 1., 1.],
                    [ 0., 1.,-1.]]
    slip_directions = [ ([ 1.,-1., 1.], [ 1.,-1.,-1.]);
                        ([ 1., 1., 1.], [ 1., 1.,-1.]);
                        ([-1., 1., 1.], [-1.,-1., 1.]);
                        ([ 1., 1., 1.], [ 1.,-1., 1.]);
                        ([ 1., 1.,-1.], [-1., 1.,-1.]);
                        ([-1., 1., 1.], [ 1., 1., 1.]) ]
    return slip_planes, slip_directions
end

function slipsystems(as::AtomicStructure, R::Rotations.Rotation)
    # generate rotated slip systems
    slip_systems = NTuple{2, Tensor{1,3}}[]
    slip_planes, slip_directions = _get_slipsystems(as)
    for (i, sp) in enumerate(slip_planes)
        for sd in slip_directions[i]
            push!(slip_systems, (Tensor{1,3}(R*sp), Tensor{1,3}(R*sd)))
        end
    end

    normalized_slip_systems = Vector{NTuple{2, Tensor{1,3,Float64,3}}}(undef, length(slip_systems))
    for i=1:length(slip_systems)
        nss = (slip_systems[i][1] / norm(slip_systems[i][1]), slip_systems[i][2] / norm(slip_systems[i][2]))
        normalized_slip_systems[i] = nss
    end
    return normalized_slip_systems
end


using LinearAlgebra
using QuasiTriangular
using Random
using Test
using KroneckerTools

Random.seed!(123)
#a = rand(2,3)
#b = rand(3,4)
#c = zeros(8)

#gemm!('N', 'N', 1.0, vec(a), 2, 3, vec(b), 4, 1.0, c)
#@test reshape(c,2,4) ≈ a*b

for m in [1, 3]
    for n in [3, 4]
        local a = randn(n,n)
        t = QuasiTriangular.QuasiUpperTriangular(schur(a).T)
        depth = 4
        for p = 0:4
            for q = depth - p
                d = randn(m*n^(p+q+1))
                d_orig = copy(d)
                w = similar(d)
                println("m = $m, p = $p, q = $q")
                d = copy(d_orig)
                @time KroneckerTools.kron_mul_elem_t!(w, a, d, n^p, n^q*m)
                @test w ≈ kron(kron(Matrix{Float64}(I(n^p)),a'),Matrix{Float64}(I(m*n^q)))*d_orig
                d = copy(d_orig)
                @time KroneckerTools.kron_mul_elem_t!(w, t, d, n^p, n^q*m)
                @test w ≈ kron(kron(Matrix{Float64}(I(n^p)),t'),Matrix{Float64}(I(m*n^q)))*d_orig
            end
        end
    end
end

ma = 2
na = 2
a = randn(ma,na)
mc = 3
order = 4
c = randn(mc,mc)
b = randn(na,mc^order)
b_orig = copy(b)
d = randn(ma,mc^order)
w1 = Vector{Float64}(undef, ma*mc^order)
w2 = Vector{Float64}(undef, ma*mc^order)
KroneckerTools.a_mul_b_kron_c!(d, a, b, c, order, w1, w2)
cc = c
for i = 2:order
    global cc = kron(cc,c)
end
@test d ≈ a*b_orig*cc

b = copy(b_orig)
KroneckerTools.a_mul_kron_b!(d,b,c,order)
cc = c
for i = 2:order
    global cc = kron(cc,c)
end
@test d ≈ b_orig*cc

order = 3
ma = 2
na = 4
a = randn(ma,na)
mb1 = 2
nb1 = 4
b1 = randn(mb1,nb1)
mb2 = 2
nb2 = 2
b2 = randn(mb2,nb2)
b = [b1, b2]
c = randn(ma,nb1*nb2)
work = zeros(16)
KroneckerTools.a_mul_kron_b!(c,a,b,work)
@test c ≈ a*kron(b[1],b[2])

order = 3
ma = 2
na = 4
a = randn(ma,na)
mb = 4
nb = 8
b = randn(mb,nb)
c = randn(2,2)
d = randn(2,2)
work1 = zeros(mb*nb)
work2 = zeros(mb*nb)
e = randn(ma,8)
KroneckerTools.a_mul_b_kron_c_d!(e, a, b, c, d, order, work1, work2)
@test e ≈ a*b*kron(c,kron(d,d))

order = 2
ma = 2
na = 4
a = randn(ma,na)
mb = 4
nb = 8
b = randn(mb,nb)
c = randn(ma*ma*nb)
d = randn(na*na*mb)
work1 = rand(na*na*mb)
work2 = rand(na*na*mb)
@time KroneckerTools.kron_at_kron_b_mul_c!(d, a, order, b, c, work1, work2)
@time KroneckerTools.kron_at_kron_b_mul_c!(d, a, order, b, c, work1, work2)

@test d ≈ kron(kron(a',a'),b)*c

println("test1")
order = 2
ma = 2
na = 4
q = 2
a = rand(ma, na)
b = rand(q*na^order)
c = rand(q*ma^order)
work1 = rand(q*max(ma, na)^order)
work2 = similar(work1)
KroneckerTools.kron_a_mul_b!(c, a, order, b, q, work1, work2)
@time KroneckerTools.kron_a_mul_b!(c, a, order, b, q, work1, work2)
@test c ≈ kron(kron(a,a),Matrix{Float64}(I(q)))*b

println("test2")
b = rand(q*ma^order)
c = rand(q*na^order)
KroneckerTools.kron_at_mul_b!(c, a, order, b, q, work1, work2)
@time KroneckerTools.kron_at_mul_b!(c, a, order, b, q, work1, work2)
@test c ≈ kron(kron(a',a'),Matrix{Float64}(I(q)))*b

println("test3")
order = 2
mc = 2
nc = 3
c = rand(mc,nc)
ma = 2
na = 4
a = rand(ma,na)
nb = nc^order
b = rand(na,nb)
d = rand(ma,mc^order)
work1 = rand(ma*max(mc, nc)^order)
work2 = similar(work1)
KroneckerTools.a_mul_b_kron_ct!(d, a, b, c, order, work1, work2)
@time KroneckerTools.a_mul_b_kron_ct!(d, a, b, c, order, work1, work2)
@test d ≈ a*b*kron(c',c')

println("test4")
nb = mc^order
b = rand(ma,nb)
d = rand(na,nc^order)
work1 = rand(na*max(mc, nc)^order)
work2 = similar(work1)
KroneckerTools.at_mul_b_kron_c!(d, a, b, c, order, work1, work2)
@time KroneckerTools.at_mul_b_kron_c!(d, a, b, c, order, work1, work2)
@test d ≈ a'*b*kron(c,c)


using Revise

using BNPRegressionGGA2021
using Distributions
using Gadfly
using Random

const BNP = BNPRegressionGGA2021

#=================================#
# Example 1 - 1 trivial predictor #
#=================================#

function simulate_sample(N0, N1)
    dy = MixtureModel(Beta, [(1.0, 5.0), (10.0, 2.0)], [0.2, 0.8])
    X0 = ones(N0, 1)
    X1 = ones(N1, 1)
    y0 = rand(dy, N0)
    y1 = LinRange(0.0, 1.0, N1) |> collect
    return dy, y0, X0, y1, X1
end;

Random.seed!(1);
N0, N1 = 500, 50;
dy, y0, X0, y1, X1 = simulate_sample(N0, N1);
m = BNP.DGPMBeta(; y0, X0, y1, X1);
chainf, chainβ = BNP.sample!(m; mcmcsize = 10000, burnin = 5000);
plot(
    layer(x = y1, y = mean(chainf), Geom.line, color = ["bnp"]),
    layer(x = y1, y = pdf.(dy, y1), Geom.line, color = ["true"]),
    layer(x = y0, Geom.histogram(density = true, bincount = 50), color = ["hist"]),
)

#===================================================#
# Example 2 - 4 trivial and 1 non-trivial predictor #
#===================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    dy1 = MixtureModel(Beta, [(1.0, 5.0), (10.0, 2.0)], [0.2, 0.8])
    dy2 = MixtureModel(Beta, [(1.0, 5.0), (10.0, 2.0)], [0.8, 0.2])
    X0 = [ones(N0) rand([0, 1], N0, 5)]
    X1 = kron(ones(N1), [ones(2, 5) [1, 0]])
    y0 = [X0[i, end] == 1 ? rand(dy2) : rand(dy1) for i in 1:N0]
    y1 = LinRange(0, 1, N1) |> collect |> x -> repeat(x, size(X1, 1) ÷ N1)
    return dy1, dy2, y0, X0, y1, X1
end

Random.seed!(1);
N0, N1 = 500, 50;
dy1, dy2, y0, X0, y1, X1 = simulate_sample(N0, N1);
ygrid = LinRange(0, 1, N1) |> collect;
plot(
    layer(x = ygrid, y = pdf.(dy2, ygrid), Geom.line, color=["pdf2"]),
    layer(x = ygrid, y = pdf.(dy1, ygrid), Geom.line, color=["pdf1"]),
)
m = BNP.DGPMBeta(; y0, X0, y1, X1);
chainf, chainβ, chaing = BNP.sample!(m; mcmcsize = 2000, burnin = 1000);
plot(x = y1, y = mean(chainf), color = string.(X1[:, end]), Geom.line)


mutable struct SymmetricTensor{T} <: AbstractArray{T,4}
    w::Int
    h::Int
    f::Int
    contents::Array{T,3}

    function SymmetricTensor{T}(w::Int, h::Int, f::Int) where T
        contents = Array{T}(undef, w, h, div(f * (f + 1), 2))
        new{T}(w, h, f, contents)
    end
end

function Base.size(o::SymmetricTensor)
    return (o.w, o.h, o.f, o.f)
end

function Base.getindex(o::SymmetricTensor, i::Int, j::Int, k::Int, l::Int)
    o.contents[i, j, convert_indices(o.f, order_indices(k, l)...)]
end

function Base.setindex!(o::SymmetricTensor, x, i::Int, j::Int, k::Int, l::Int)
    o.contents[i, j, convert_indices(o.f, order_indices(k, l)...)] = x
end

function order_indices(i::Int, j::Int)
    return j < i ? (j, i) : (i, j)
end

function convert_indices(n::Int, i::Int, j::Int)
    return n * (i - 1) + j - div(i * (i - 1), 2)
end

"""
```
descriptor = RegionCovariance(feature_img)
```
"""
struct RegionCovariance
    P::Array{Float64,3}
    Q::SymmetricTensor{Float64}
    nrows::Int
    ncols::Int
    nfeatures::Int

    function RegionCovariance(F::AbstractArray{Float64,3})
        nrows, ncols, nfeatures = length.(axes(F))
        P = featurewise_integral_images(F, nrows, ncols, nfeatures)
        Q = feature_product_integral_images(F, nrows, ncols, nfeatures)
        new(P, Q, nrows, ncols, nfeatures)
    end

    # Calculate the integral image of each feature layer of `F`. This is equivalent to
    # calling `integral_image(Fᵢ)` in a loop, where `Fᵢ` is a view of ith feature layer of
    # `F`, but reproducing the functionality here results in significantly fewer
    # allocations.
    function featurewise_integral_images(F::AbstractArray{Float64,3},
                                         nrows::Int, ncols::Int, nfeatures::Int)
        P = Array{Float64,3}(undef, nrows, ncols, nfeatures)

        @inbounds for f in 1:nfeatures
            for r in 1:nrows, c in 1:ncols
                P[r,c,f] = F[r,c,f]
                P[r,c,f] += r > 1 ? P[r-1,c,f] : 0.0
                P[r,c,f] += c > 1 ? P[r,c-1,f] : 0.0
                P[r,c,f] -= r > 1 && c > 1 ? P[r-1,c-1,f] : 0.0
            end
        end

        return P
    end

    # Given a feature image with 𝑛 features, where each pixel is a corresponds to a feature
    # vector 𝐱 ∈ ℝⁿ, calculate the 𝑛 × 𝑛 matrix 𝐴 for each 𝐱 such that 𝐴[𝑖,𝑗] = 𝐱ᵢ ⋅ 𝐱ⱼ, and
    # then the integral image of that matrix. Note that 𝐴 in this case is symmetric: a
    # custom data structure is used to leverage this symmetry and save space.
    function feature_product_integral_images(F::AbstractArray{Float64,3},
                                             nrows::Int, ncols::Int, nfeatures::Int)
        Q = SymmetricTensor{Float64}(nrows, ncols, nfeatures)

        @inbounds for fᵢ in 1:nfeatures
            for fⱼ in fᵢ:nfeatures
                for r in 1:nrows, c in 1:ncols
                    Q[r,c,fᵢ,fⱼ] = F[r,c,fᵢ] * F[r,c,fⱼ]
                    Q[r,c,fᵢ,fⱼ] += r > 1 ? Q[r-1,c,fᵢ,fⱼ] : 0.0
                    Q[r,c,fᵢ,fⱼ] += c > 1 ? Q[r,c-1,fᵢ,fⱼ] : 0.0
                    Q[r,c,fᵢ,fⱼ] -= r > 1 && c > 1 ? Q[r-1,c-1,fᵢ,fⱼ] : 0.0
                end
            end
        end

        return Q
    end
end

function covariance_matrix(d::RegionCovariance, rows::AbstractRange, cols::AbstractRange)
    𝐂 = Array{Float64}(undef, d.nfeatures, d.nfeatures)
    return covariance_matrix!(𝐂, d, rows, cols)
end

function covariance_matrix!(𝐂::AbstractArray{Float64,2}, d::RegionCovariance,
                            rows::AbstractRange, cols::AbstractRange)
    row₀, col₀ = first.((rows, cols))
    row₁, col₁ = last.((rows, cols))
    return covariance_matrix!(𝐂, d, row₀, col₀, row₁, col₁)
end

function covariance_matrix(d::RegionCovariance,
                           row₀::Integer, col₀::Integer,
                           row₁::Integer, col₁::Integer)
    𝐂 = Array{Float64}(undef, d.nfeatures, d.nfeatures)
    return covariance_matrix!(𝐂, d, row₀, col₀, row₁, col₁)
end

function covariance_matrix!(𝐂::AbstractArray{Float64,2}, d::RegionCovariance,
                            row₀::Integer, col₀::Integer,
                            row₁::Integer, col₁::Integer)
    region_size = (row₁-row₀+1)*(col₁-col₀+1)

    𝐐 = sum_feature_products_over_region(d.Q, Val{d.nfeatures}(), row₀, col₀, row₁, col₁)
    𝐩 = sum_over_region_by_feature(d.P, Val{d.nfeatures}(), row₀, col₀, row₁, col₁)
    result = (𝐐 - ((𝐩 * 𝐩') / region_size)) / (region_size-1)

    𝐂 .= result
end

function sum_feature_products_over_region(Q::SymmetricTensor{Float64}, ::Val{nfeatures},
                                          row₀::Integer, col₀::Integer, row₁::Integer,
                                          col₁::Integer) where nfeatures
    𝐐 = zeros(MMatrix{nfeatures,nfeatures,Float64})

    @inbounds for fᵢ in 1:nfeatures
        for fⱼ in fᵢ:nfeatures
            𝐐[fᵢ,fⱼ] = Q[row₁,col₁,fᵢ,fⱼ]
            𝐐[fᵢ,fⱼ] -= row₀ > 1 ? Q[row₀-1,col₁,fᵢ,fⱼ] : 0.0
            𝐐[fᵢ,fⱼ] -= col₀ > 1 ? Q[row₁,col₀-1,fᵢ,fⱼ] : 0.0
            𝐐[fᵢ,fⱼ] += row₀ > 1 && col₀ > 1 ? Q[row₀-1,col₀-1,fᵢ,fⱼ] : 0.0
            𝐐[fⱼ,fᵢ] = 𝐐[fᵢ,fⱼ]
        end
    end

    return 𝐐
end

function sum_over_region_by_feature(P::AbstractArray{Float64,3}, ::Val{nfeatures},
                                    row₀::Integer, col₀::Integer, row₁::Integer,
                                    col₁::Integer) where nfeatures
    𝐩 = zeros(MVector{nfeatures,Float64})

    @inbounds for f in 1:nfeatures
        𝐩[f] = P[row₁,col₁,f]
        𝐩[f] -= row₀ > 1 ? P[row₀-1,col₁,f] : 0.0
        𝐩[f] -= col₀ > 1 ? P[row₁,col₀-1,f] : 0.0
        𝐩[f] += row₀ > 1 && col₀ > 1 ? P[row₀-1,col₀-1,f] : 0.0
    end

    return 𝐩
end


@testset "CUCB" begin
  @testset "Algorithm" begin
    n = 3
    ε = .1
    reward = Distribution[Bernoulli(.5 + ((i == j) ? ε : 0.)) for i in 1:n, j in 1:n]
    instance = UncorrelatedPerfectBipartiteMatching(reward, PerfectBipartiteMatchingMunkresSolver())

    Random.seed!(1)
    n_rounds = 20
    s = simulate(instance, CUCB(), n_rounds)
    @test s.round == n_rounds
    @test sum(values(s.arm_counts)) == n_rounds * n
    @test s.regret ≈ 3.7
    @test s.reward ≈ 36.0
    @test s.regret + s.reward ≈ 39.7

    for arm in keys(s.arm_counts)
      @test s.arm_average_reward[arm] ≈ s.arm_reward[arm] / s.arm_counts[arm]
    end
  end

  @testset "Trace" begin
    n = 3
    ε = .1
    reward = Distribution[Bernoulli(.5 + ((i == j) ? ε : 0.)) for i in 1:n, j in 1:n]
    instance = UncorrelatedPerfectBipartiteMatching(reward, PerfectBipartiteMatchingMunkresSolver())

    n_rounds = 20
    test_policy_trace(instance, CUCB(), n_rounds, n_rounds * n, 3.7, 36.0, 39.7)
  end
end


#!/usr/bin/env julia

const OUTPUT = haskey(ENV, "OUTPUT")

function usage()
    prog = basename(Base.source_path())
#    println(STDERR,"Usage ",prog," <rows> <cols> <y1> <y2> <x1> <x2> <no. of threads><lamda> <no. of iter>")
    println(STDERR,"Usage ",prog," <rows> <cols> <y1> <y2> <x1> <x2> <lamda> <no. of iter>")
    println(STDERR,"\t<rows>   - number of rows")
    println(STDERR, "\t<cols>    - number of cols\n")
    println(STDERR, "\t<y1> 	 - y1 value of the speckle\n")
    println(STDERR, "\t<y2>      - y2 value of the speckle\n")
    println(STDERR, "\t<x1>       - x1 value of the speckle\n")
    println(STDERR, "\t<x2>       - x2 value of the speckle\n")
#    println(STDERR, "\t<no. of threads>  - no. of threads\n")
    println(STDERR, "\t<lamda>   - lambda (0,1)\n")
    println(STDERR, "\t<no. of iter>   - number of iterations\n")
    exit(1)
end

function main(args)
    if length(args) != 8
      usage()
    end
    rows = parse(Int64,args[1])
    cols = parse(Int64,args[2])
    if ((rows % 16) != 0) | ((cols % 16) != 0)
        println(STDERR,"rows and cols must be mutiples of 16")
        exit(1)
    end
    r1 = parse(Int32,args[3]) + 1
    r2 = parse(Int32,args[4]) + 1
    c1 = parse(Int32,args[5]) + 1
    c2 = parse(Int32,args[6]) + 1
    lambda = parse(Float32,args[7])
    niter = parse(Int32,args[8])

    size_R = (r2 - r1 + 1) * (c2 - c1 + 1);

    iN::Array{Int32,1} = [i-1 for i in 1:rows]
    iS::Array{Int32,1} = [i+1 for i in 1:rows]
    jW::Array{Int32,1} = [j-1 for j in 1:cols]
    jE::Array{Int32,1} = [j+1 for j in 1:cols]
    iN[1] = 1
    iS[rows] = rows
    jW[1] = 1
    jE[cols] = cols

    dN::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))
    dS::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))
    dW::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))
    dE::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))

    println("Randomizing the input matrix")

    srand(7)
    I = rand(Float32,rows,cols)
    J::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols),init = J -> J[Base.localindexes(J)] = exp.(I[Base.localindexes(J)]))
    c::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))

    println("Start the SRAD main loop")

    tic()

    sum::Float32 = 0
    sum2::Float32 = 0
    for iter in 1:niter
        sum = 0
        sum2 = 0
        for i in r1:r2
            for j in c1:c2
                tmp = J[i,j]
                sum += tmp
                sum2 += tmp * tmp
            end
        end
    end

    meanROI::Float32 = sum / size_R
    varROI::Float32 = (sum / size_R) - meanROI * meanROI
    q0sqr::Float32 = varROI / (meanROI * meanROI)

    @sync @parallel for i in 1:size(J,1)
        for j in size(J,2)
            Jc = J[i,j]

            # directional derivates
            dN[i,j] = J[iN[i],j] - Jc
            dS[i,j] = J[iS[i],j] - Jc
            dW[i,j] = J[i,jW[j]] - Jc
            dE[i,j] = J[i,jE[j]] - Jc

            G2 = (dN[i,j] * dN[i,j] + dS[i,j] * dS[i,j] + dW[i,j] * dW[i,j] +
                  dE[i,j] * dE[i,j]) /
                 (Jc * Jc)

            L = (dN[i,j] + dS[i,j] + dW[i,j] + dE[i,j]) / Jc

            num = (0.5 * G2) - ((1.0 / 16.0) * (L * L))
            den = 1 + (0.25 * L)
            qsqr = num / (den * den)

            # diffusion coefficent (equ 33)
            den = (qsqr - q0sqr) / (q0sqr * (1 + q0sqr))
            c[i,j] = 1.0 / (1.0 + den)

            # saturate diffusion coefficent
            if (c[i,j] < 0)
                c[i,j] = 0
            elseif (c[i,j] > 1)
                c[i,j] = 1
            end
        end
    end

    @sync @parallel for i in 1:size(J,1)
        for j in 1:size(J,2)
            # diffusion coefficient
            cN = c[i,j]
            cS = c[iS[i],j]
            cW = c[i,j]
            cE = c[i,jE[j]]

            # divergence (equ 58)
            D = cN * dN[i,j] + cS * dS[i,j] + cW * dW[i,j] + cE * dE[i,j]

            # image update (equ 61)
            J[i,j] += 0.25 * lambda * D
        end
    end

    toc()

    println("Computation done")

    if OUTPUT
        f = open("output.txt","w")
        for i in 1:size(J,1)
            for j in 1:size(J,2)
               print(f,@sprintf("%.5f ",J[i,j]))
            end
            println(f)
        end
        close(f)
    end
end


#=
demo_HDP_MultinomialDirichlet.jl

A demo for Hierarchical Dirichlet Process mixture models with MultinomialDirichlet
Bayesian components. This demo uses auxiliary variable method for for inference.

28/07/2015
Adham Beyki, odinay@gmail.com
=#

using BNP
srand(123)

true_gg = 1.0
true_aa = 0.5
n_groups = 100
n_group_j   = 100 * ones(Int, n_groups)
KK_truncation = 25

beta_tilde = rand(Distributions.Beta(1, true_gg), KK_truncation)
beta_tilde[KK_truncation] = 1.0
my_beta = BNP.stick_breaking(beta_tilde)

pi_tilde = zeros(Float64, n_groups, KK_truncation)
my_pi = zeros(Float64, n_groups, KK_truncation)
for jj = 1:n_groups
    for kk = 1:KK_truncation-1
        gamma1 = true_aa * my_beta[kk]
        gamma2 = true_aa * (1 - sum(my_beta[1:kk]))
        pi_tilde[jj, kk] = rand(Distributions.Beta(gamma1, gamma2))
    end
    pi_tilde[jj, KK_truncation] = 1.0
    my_pi[jj, :] = BNP.stick_breaking(pi_tilde[jj, :][:])
end

true_zz = Array(Vector{Int}, n_groups)
for jj = 1:n_groups
    true_zz[jj] = zeros(Int, n_group_j[jj])
    for ii = 1:n_group_j[jj]
        kk = sample(my_pi[jj, :][:])
        true_zz[jj][ii] = kk
    end
end

zz_flat = Int[]
for jj = 1:n_groups
	append!(zz_flat, true_zz[jj])
end

zz_unq = sort(unique(zz_flat))
true_KK = length(zz_unq)
for jj = 1:n_groups
	for kk = 1:true_KK
		idx = find(x -> x==zz_unq[kk], true_zz[jj])
		true_zz[jj][idx] = kk
	end
end

n_groups   = 50
n_group_j  = 100 * ones(Int, n_groups)
vocab_size = 25

# constructing the topics
true_topics = BNP.gen_bars(true_KK, vocab_size, 0)

# constructing the observations and labels
xx = Array(Vector{Int}, n_groups)
true_nn = zeros(Int, n_groups, true_KK)

for jj = 1:n_groups
    xx[jj] = zeros(Int, n_group_j[jj])
    for ii = 1:n_group_j[jj]
        kk = true_zz[jj][ii]
        xx[jj][ii] = sample(true_topics[kk, :][:])
        true_nn[jj, kk] += 1
    end
end

## ------- inference -------- ##
# constructing the Bayesian component
dd = vocab_size
aa = 0.5*dd
q0 = MultinomialDirichlet(dd, aa)

# constructing the  model
dhdp_KK_init = 1
KK_truncation = 15
dhdp_gg = 1.0
dhdp_g1 = 0.1
dhdp_g2 = 0.1
dhdp_aa = 0.5
dhdp_a1 = 0.1
dhdp_a2 = 0.1
dhdp_aw = 2.0
dhdp_bw = 3.0
dhdp = dHDP(q0, dhdp_KK_init, dhdp_gg, dhdp_g1, dhdp_g2, dhdp_aa, dhdp_a1, dhdp_a2, dhdp_aw, dhdp_bw)

# sampling
zz = Array(Vector{Int}, n_groups)
for jj = 1:n_groups
    zz[jj] = ones(Int, n_group_j[jj])
    zz[jj] = rand(1:dhdp.KK, n_group_j[jj])
end

n_burnins   = 500
n_lags      = 1
n_samples   = 500
sample_hyperparam = false
n_internals = 10
store_every = 1000
results_path = ""
filename    = "demo_dHDP_MultinomialDirichlet"

KK_list, KK_dict, w_tilde = BNP.truncated_gibbs_sampler(dhdp, KK_truncation, xx, zz,
    n_burnins, n_lags, n_samples, sample_hyperparam, n_internals, store_every, results_path, filename)

KK_hist = hist(KK_list, 0.5:maximum(KK_list)+0.5)[2]
candidate_KK = indmax(KK_hist)
posterior_components, nn = posterior(dhdp, xx, KK_dict[candidate_KK], candidate_KK)


# Forward/Backward

# {forward,backward}(a, A, L)
for f in (:forward, :backward)
    f!  = Symbol("$(f)!")    # forward!
    fl! = Symbol("$(f)log!") # forwardlog!

    @eval begin
        """
            $($f)(a, A, L; logl) -> (Vector, Float)

        Compute $($f) probabilities using samples likelihoods.
        See [Forward-backward algorithm](https://en.wikipedia.org/wiki/Forward–backward_algorithm).

        **Output**
        - `Vector{Float64}`: $($f) probabilities.
        - `Float64`: log-likelihood of the observed sequence.
        """
        function $(f)(a::AbstractVector, A::AbstractMatrix, L::AbstractMatrix; logl = false)
            m = Matrix{Float64}(undef, size(L))
            c = Vector{Float64}(undef, size(L)[1])
            if logl
                $(fl!)(m, c, a, A, L)
            else
                warn_logl(L)
                $(f!)(m, c, a, A, L)
            end
            m, sum(log.(c))
        end
    end
end

# {forward,backward}(hmm, observations)
for f in (:forward, :backward)
    f!  = Symbol("$(f)!")    # forward!
    fl! = Symbol("$(f)log!") # forwardlog!

    @eval begin
        """
            $($f)(hmm, observations; logl, robust) -> (Vector, Float)

        Compute $($f) probabilities of the `observations` given the `hmm` model.

        **Output**
        - `Vector{Float64}`: $($f) probabilities.
        - `Float64`: log-likelihood of the observed sequence.

        **Example**
        ```julia
        using Distributions, HMMBase
        hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
        y = rand(hmm, 1000)
        probs, tot = $($f)(hmm, y)
        ```
        """
        function $(f)(hmm::AbstractHMM, observations; robust = false, kwargs...)
            L = likelihoods(hmm, observations; robust = robust, kwargs...)
            $(f)(hmm.a, hmm.A, L; kwargs...)
        end
    end
end

# Posteriors

"""
    posteriors(α, β) -> Vector

Compute posterior probabilities from `α` and `β`.

**Arguments**
- `α::AbstractVector`: forward probabilities.
- `β::AbstractVector`: backward probabilities.
"""
function posteriors(α::AbstractMatrix, β::AbstractMatrix)
    γ = Matrix{Float64}(undef, size(α))
    posteriors!(γ, α, β)
    γ
end

"""
    posteriors(a, A, L; logl) -> Vector

Compute posterior probabilities using samples likelihoods.
"""
function posteriors(a::AbstractVector, A::AbstractMatrix, L::AbstractMatrix; kwargs...)
    α, _ = forward(a, A, L; kwargs...)
    β, _ = backward(a, A, L; kwargs...)
    posteriors(α, β)
end

"""
    posteriors(hmm, observations; logl, robust) -> Vector

Compute posterior probabilities using samples likelihoods.

**Example**
```julia
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
y = rand(hmm, 1000)
γ = posteriors(hmm, y)
```
"""
function posteriors(hmm::AbstractHMM, observations; robust = false, kwargs...)
    L = likelihoods(hmm, observations; robust = robust, kwargs...)
    posteriors(hmm.a, hmm.A, L; kwargs...)
end

# Likelihood

"""
    loglikelihood(hmm, observations; logl, robust) -> Float64

Compute the log-likelihood of the observations under the model.
This is defined as the sum of the log of the normalization coefficients in the forward filter.

**Output**
- `Float64`: log-likelihood of the observations sequence under the model.

**Example**
```jldoctest
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
loglikelihood(hmm, [0.15, 0.10, 1.35])
# output
-4.588183811489616
```
"""
function loglikelihood(hmm::AbstractHMM, observations; logl = false, robust = false)
    forward(hmm, observations, logl = logl, robust = robust)[2]
end


# ---
# title: 476. Number Complement
# id: problem476
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Bit Manipulation
# link: <https://leetcode.com/problems/number-complement/description/>
# hidden: true
# ---
#
# Given a **positive** integer `num`, output its complement number. The
# complement strategy is to flip the bits of its binary representation.
#
#
#
# **Example 1:**
#
#
#
#     Input: num = 5
#     Output: 2
#     Explanation: The binary representation of 5 is 101 (no leading zero bits), and its complement is 010. So you need to output 2.
#
#
# **Example 2:**
#
#
#
#     Input: num = 1
#     Output: 0
#     Explanation: The binary representation of 1 is 1 (no leading zero bits), and its complement is 0. So you need to output 0.
#
#
#
#
# **Constraints:**
#
#   * The given integer `num` is guaranteed to fit within the range of a 32-bit signed integer.
#   * `num >= 1`
#   * You could assume no leading zero bit in the integer's binary representation.
#   * This question is the same as 1009: <https://leetcode.comhttps://leetcode.com/problems/complement-of-base-10-integer/>
#
#
## @lc code=start
using LeetCode

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


# TODO:  TO FIX !!!
#/////////////////////////////////////////////////
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using SparseArrays, LinearAlgebra
using ViewerGL
GL = ViewerGL

function randomshapes()
	V,EV,FV = Lar.randomcubes(5, .8)
	model3d = V,EV,FV
	VV = [[k] for k=1:size(V,2)]

	Sigma =  Lar.spaceindex((V,FV))
	GL.VIEW( GL.numbering(.15)((V,[VV, EV, FV])) )
	V,EVs = Lar.biconnectedComponent((V,EV::Lar.Cells)) # 2-connected components (H & T)
	GL.VIEW(GL.GLExplode(V,EVs,1.5,1.5,1.5,12,0.99));

	W,bases,coboundaries = Lar.chaincomplex(V,FV,EV)
	EV, FV, CV = bases
	copEV, copFE, copCF = coboundaries # OK

#
EV = Lar.cop2lar(copEV)
FE = [findnz(copFE[k,:])[1] for k=1:size(copFE,1)]
FV = [collect(Set(cat(EV[e] for e in FE[f]))) for f=1:length(FE)]
FV = convert(Lar.Cells, FV)
W = convert(Lar.Points, V')
WW = [[k] for k=1:size(W,2)]

GL.VIEW(GL.numbering(.2)((W,[WW,EV])));
#

	W = convert(Lar.Points, W')
	triangulated_faces = Lar.triangulate(W, [copEV, copFE])
	FVs = convert(Array{Lar.Cells}, triangulated_faces)

	Plasm.viewcolor(V::Lar.Points, FVs::Array{Lar.Cells})
end

randomshapes()


"""
	iterate_until_collisionless(f)

Sample `f` until the result is collisionless.
"""
function iterate_until_collisionless(f)

    """samples f until the result is collisionless

	as it takes a function, example usage :

		random_mode_occupation_collisionless(n::Int,m::Int) = iterate_until_collisionless(() -> random_mode_occupation(n,m))
	"""

    while true

        state = f()
        if is_collisionless(state)
            return state
        end

    end

end

"""
	fill_arrangement(occupation_vector)
	fill_arrangement(r::ModeOccupation)
	fill_arrangement(input::Input)

Convert a mode occupation list to a mode assignement.
"""
function fill_arrangement(occupation_vector)

    """from a mode occupation list to a mode assignment list (Tichy : s to d(s))

    For instance if vect{s} = (2,0,1) then vect{d} = (1,1,3) """

    arrangement = zeros(eltype(occupation_vector), sum(occupation_vector))

    position = 1
    for i in 1:length(occupation_vector)
        for j in 1:occupation_vector[i]
            arrangement[position] = i
            position += 1
        end
    end

    arrangement

end

fill_arrangement(r::ModeOccupation) = fill_arrangement(r.state)

fill_arrangement(inp::Input) = fill_arrangement(inp.r)

"""
	random_occupancy(n::Int, m::Int)

Return a vector of size `m`	with `n` randomly placed ones.
"""
function random_occupancy(n::Int, m::Int)

	""" returns a vector of size m with n randomly placed ones """
	if n > m
		throw(ArgumentError("Implemented at most one photon per mode"))
	else
		occupancy_vector = shuffle(append!(ones(n), zeros(m-n)))
		occupancy_vector = Int.(occupancy_vector)
	end

	return occupancy_vector
end

"""
	random_mode_occupation(n::Int, m::Int)

Create a [`ModeOccupation`](@ref) from a mode occupation list of `n` ramdomly placed ones
among `m` sites.
"""
random_mode_occupation(n::Int, m::Int) = ModeOccupation(random_occupancy(n,m))

"""
	random_mode_occupation_collisionless(n::Int, m::Int)

Create a [`ModeOccupation`](@ref) from a random mode occupation that is likely collisionless.
"""
function random_mode_occupation_collisionless(n::Int, m::Int)

	n<=m ? iterate_until_collisionless(() -> random_mode_occupation(n,m)) : error("n>m cannot make for collisionless mode occupations")

end

"""
	at_most_one_photon_per_bin(occupancy_vector::Vector{Int})
	check_at_most_one_particle_per_mode(occ)

Check wether `occupancy_vector` contains more than one photon per site.
"""
function at_most_one_photon_per_bin(occupancy_vector::Vector{Int})
	all(in([0,1]).(occupancy_vector))
end

check_at_most_one_particle_per_mode(occ) = at_most_one_photon_per_bin(occ) ? nothing : error("more than one input per mode")

"""
	occupancy_vector_to_partition(occupancy_vector)
	occupancy_vector_to_mode_occupancy(occupancy_vector)

Return a partition of occupied modes from an `occupancy_vector`.
"""
function occupancy_vector_to_partition(occupancy_vector)

	partition = []

	check_at_most_one_particle_per_mode(occupancy_vector)
	# a partition cannot be defined by taking twice the same mode, it would lead to
	# hidden errors in the results (double counting)

	for (mode, selected) in enumerate(occupancy_vector)
		if selected == 1
			append!(partition, mode)
		end
	end
	partition

end

"""
	occupancy_vector_to_mode_occupancy(occupancy_vector)

Return a partition of occupied modes from an `occupancy_vector`.
"""
function occupancy_vector_to_mode_occupancy(occupancy_vector)
	occupancy_vector_to_partition(occupancy_vector)
end

"""
	scattering_matrix(U::Matrix, input_state::Vector{Int}, output_state::Vector{Int})
	scattering_matrix(U::Interferometer, r::ModeOccupation, s::ModeOccupation)
	scattering_matrix(U::Interferometer, i::Input, o::FockDetection)

Return the submatrix of `U` whose rows and columns are respectively defined by
`input_state` and `output_state`.

!!! note "Reference"
	[http://arxiv.org/abs/quant-ph/0406127v1](http://arxiv.org/abs/quant-ph/0406127v1)
"""
function scattering_matrix(U::Matrix, input_state::Vector{Int}, output_state::Vector{Int})

    """
    U = interferometer matrix, size m*m
    input_state = input occupation number vector (s[i] is the number of photons in the ith input)
    output_state = output occupation number vector (r[i] is the number of photons in the ith output)
    n photons

    follows http://arxiv.org/abs/quant-ph/0406127v1
    """

    m = size(U,1)
    n = sum(input_state)

    if length(input_state) != m || length(output_state) != m
        throw(DimensionMismatch())
    elseif sum(input_state) != sum(output_state)
        error("photon number not the same at input and output")
    end

    index_input = fill_arrangement(input_state)
    index_output = fill_arrangement(output_state)

	try
		U[index_input,index_output]
	catch err
		if isa(MethodError, err)
			copy(U)[index_input,index_output]
		else
			println(err)
		end
	end

end

scattering_matrix(interf::Interferometer, r::ModeOccupation, s::ModeOccupation) = scattering_matrix(interf.U,r.state,s.state)
scattering_matrix(interf::Interferometer, i::Input, o::FockDetection) = scattering_matrix(interf.U,i.r,o.s)

function vector_factorial(occupancy_vector)

    """returns the function mu at the denominator of event probabilities"""

    prod(factorial.(occupancy_vector))

end

vector_factorial(r::ModeOccupation) = vector_factorial(r.state)
vector_factorial(i::Input) = vector_factorial(i.r)
vector_factorial(o::FockDetection) = vector_factorial(o.s)

"""
	bosonic_amplitude(U, input_state, output_state, permanent=ryser)
	process_amplitude(U, input_state, output_state, permanent=ryser)

Compute the probability amplitude to go from `input_state` to `output_state`
through the interferomter `U` in the [`Bosonic`](@ref) case.
"""
function bosonic_amplitude(U, input_state, output_state, permanent = ryser)

    """event amplitude"""

    permanent(scattering_matrix(U, input_state, output_state))/sqrt(vector_factorial(input_state) * vector_factorial(output_state))
end

function process_amplitude(U, input_state, output_state, permanent = ryser)

	bosonic_amplitude(U, input_state, output_state, permanent)
end

"""
	bosonic_probability(U, input_state, output_state)
	process_probability(U, input_state, output_state)

Compute the probability to go from `input_state` to `output_state`
through the interferometer `U` in the [`Bosonic`](@ref) case.
"""
function bosonic_probability(U, input_state, output_state)

	"""bosonic process_probability"""

    abs(process_amplitude(U, input_state, output_state))^2

end

function process_probability(U, input_state, output_state)
	#@warn "obsolete function, use bosonic_probability or probability"
	bosonic_probability(U, input_state, output_state)
end

"""
	distinguishable_probability(U, input_state, output_state, permanent=ryser)
	process_probability_distinguishable(U, input_state, output_state, permanent=ryser)

Compute the probability to go from `input_state` to `output_state` through
the interferomter `U` in the [`Distinguishable`](@ref) case.
"""
function distinguishable_probability(U, input_state, output_state, permanent = ryser)

	"""distinguishable (or classical) process_probability"""

	permanent(abs.(scattering_matrix(U, input_state, output_state)).^2)/sqrt(vector_factorial(input_state) * vector_factorial(output_state))

end

function process_probability_distinguishable(U, input_state, output_state, permanent = ryser)

	#@warn "obsolete function, use distinguishable_probability or probability"

	distinguishable_probability(U, input_state, output_state, permanent)
end

### need to implement partial distinguishability processs probabilities ###

"""
	process_probability_partial(U, S, input_state, output_state)
	process_probability_partial(interf::Interferometer, input_state::Input{TIn} where {TIn<:PartDist},output_state::FockDetection)

Compute the probability to go from `input_state` to `output_state` through the
interferometer `U` in the [`PartDist`](@ref) case where partial distinguishable is described
by the [`GramMatrix`](@ref) `S`.

!!! note "Reference"
    [https://arxiv.org/abs/1410.7687](https://arxiv.org/abs/1410.7687)
"""
function process_probability_partial(U, S, input_state,output_state)

    """computes the partially distinguishable process probability according to Tichy's tensor permanent https://arxiv.org/abs/1410.7687"""

    n = size(S,1)

    @argcheck sum(input_state) == sum(output_state) "particles not conserved"

    @argcheck sum(input_state) == n "S matrix doesnt have the same number of photons as the input"

	M = scattering_matrix(U, input_state, output_state)
	W = Array{eltype(U)}(undef, (n,n,n))

	for ss in 1:n
	    for rr in 1:n
	        for j in 1:n
	            W[ss,rr,j] = M[ss, j] * conj(M[rr, j]) * S[rr, ss]
	        end
	    end
	end

    1/(vector_factorial(input_state) * vector_factorial(output_state)) * ryser_tensor(W)

end

process_probability_partial(interf::Interferometer, input_state::Input{TIn} where {TIn<:PartDist},output_state::FockDetection) = process_probability_partial(interf.U, input_state.G.S, input_state.r.state,output_state.s.state)

"""
	compute_probability(ev::Event{TIn, TOut}) where {TIn<:InputType, TOut<:FockDetection}

Given an [`Event`](@ref), gives the probability to get the outcome `TOut` when `TIn`
passes though the interferometer `ev.interferometer`.
"""
function compute_probability!(ev::Event{TIn,TOut}) where {TIn<:InputType, TOut<:FockDetection}

	check_probability_empty(ev)

	ev.proba_params.precision = eps()
	ev.proba_params.failure_probability = 0

	if TIn == Bosonic

		ev.proba_params.probability = bosonic_probability(ev.interferometer.U, ev.input_state.r.state, ev.output_measurement.s.state)

	elseif TIn == Distinguishable

		ev.proba_params.probability = distinguishable_probability(ev.interferometer.U, ev.input_state.r.state, ev.output_measurement.s.state)

	else
		ev.proba_params.probability = process_probability_partial(ev.interferometer, ev.input_state, ev.output_measurement)
	end

	ev.proba_params.probability = clean_proba(ev.proba_params.probability)

end

"""
	output_mode_occupation(n::Int, m::Int)

Return all possible configurations of `n` photons among `m` modes.
"""
function output_mode_occupation(number_photons, number_modes)

	nlist = collect(1:number_modes)

	for i = 1:number_photons-1
		sublist = []
		for j = 1:length(nlist)
			for k = 1:number_modes
				push!(sublist, vcat(nlist[j], k))
			end
		end
		nlist = sublist
	end

	nlist

end

"""
	check_suppression_law(event)

Check if the event is suppressed according to the [rule](https://arxiv.org/pdf/1002.5038.pdf).
"""
function check_suppression_law(event)

	if mod(sum(event), length(event)) != 0
		return true
	else
		return false
	end

end

### this is an old function that needs to be cleaned

# function partition_probability_distribution_distinguishable(part, U)
# 	"""generates a vector giving the probability to have k photons in
# 	the partition part for the interferometer U by the random walk method
# 	discussed in a 22/02/21 email
# 	part is written like (1,2,4) if it is the output modes 1, 2 and 4"""
#
# 	function proba_photon_in_partition(i, part, U)
# 		"""returns the probability that the photon i ends up in the set of outputs part, for distinguishable particles only"""
#
# 		sum([abs(U[i, j])^2 for j in part]) # note the inversion line column
# 	end
#
#
# 	function walk(probability_vector, walk_number, part, U)
#
# 		new_probability_vector = similar(probability_vector)
# 		n = size(U)[1]
# 		proba_this_walk = proba_photon_in_partition(walk_number, part, U)
#
# 		for i in 1 : n+1
# 			new_probability_vector[i] = (1-proba_this_walk) * probability_vector[i] + proba_this_walk * probability_vector[i != 1 ? i-1 : n+1]
# 		end
#
# 		new_probability_vector
# 	end
#
# 	n = size(U)[1]
# 	probability_vector = zeros(n+1)
# 	probability_vector[1] = 1
#
# 	for walk_number in 1 : n
# 		probability_vector = walk(probability_vector, walk_number, part, U)
# 	end
#
# 	if !(isapprox(sum(probability_vector), 1., rtol = 1e-8))
# 		println("WARNING : probabilites do not sum to one ($(sum(probability_vector)))")
# 	end
#
# 	probability_vector
# end

function is_collisionless(r)
    all(r .<= 1)
end

is_collisionless(r::ModeOccupation) = is_collisionless(r.state)
is_collisionless(i::Input) = is_collisionless(i.r.state)

first_modes_array(n::Int,m::Int) = first_modes(n,m).state


@testset "test_simplification" begin
## SMINREAL ##
G = ss([-5 0 0 0; 0 -1 -2.5 0; 0 4 0 0; 0 0 0 -6], [2 0; 0 1; 0 0; 0 2],
       [0 3 0 0; -2 0 0 1], [0 0; 1 0])
@test sminreal(G) == G
@test sminreal(G[1, 1]) == ss(0)
@test sminreal(G[1, 2]) == ss([-1 -2.5; 4 0], [1; 0], [3 0], [0])
@test sminreal(G[2, 1]) == ss([-5], [2], [-2], [1])
@test sminreal(G[2, 2]) == ss([-6], [2], [1], [0])

## MINREAL ##

s = tf("s")
P = [1/(s+1) 2/(s+3); 1/(s+1) 1/(s+1)]
sys = ss(P)
sysmin = minreal(sys)

@test size(sysmin.A,1) == 3 # Test that the reduction of sys worked

@test norminf(sys - sysmin)[1] < 1e-15 # And that the answer is correct

@test_broken balreal(sys-sysmin)

@test all(sigma(sys-sysmin, [0.0, 1.0, 2.0])[1] .< 1e-15)  # Previously crashed because of zero dimensions in tzero

t = 0:0.1:10
y1,x1 = step(sys,t)[[1,3]]
y2,x2 = step(sysmin,t)[[1,3]]
@test sum(abs2,y1.-y2) < 1e-6 # Test that the output from the two systems are the same

end


using NLopt
using Plots
using SparseArrays

struct ModelParameters
    γ; # RRA
    μ; # risk premium
    σ; # vol of risky asset
    ν; # vol of int rate
    ρ; # discount rate
    η; # correlation between risky return and risk free rate
    rgrid; # grid for interest rate
    rmin;  # min and max interest rates in the grid
    rmax;  # -//-
    Δr; # step in interest rates grid
    rbar; # [For the geometric OU]
    θ; # [For the geometric OU]
end

function solve_optimal_risky_share(A, params)
    # This function takes a vector of A (investment opportuinities)
    # for a given value of interest rate and a struct of model
    # parameters. It then numerically solves for optimal risky
    # asset allocation.
    rgrid = params.rgrid
    γ = params.γ
    σ = params.σ
    μ = params.μ
    η = params.η
    ν = params.ν;

    # Modifying for non-zero correlation
    αopt=zeros(size(A))
    for i = 1:size(rgrid)[1]
        ri = rgrid[i]
        Ai = A[i]

        if i == 1
            Ailower = A[1]
            Aihigher = A[2]
        elseif i == I
            Ailower = A[I-1]
            Aihigher = A[I]
        else
            Ailower = A[i-1]
            Aihigher = A[i+1]
        end

        # Forming function to solve numerically
        function to_minimize(x::Vector, grad::Vector)
            ((ri + x[1]*μ - 0.5*x[1]^2*σ^2)^(-γ) * (μ - x[1]*σ^2) + Ai*(1-γ)*x[1]*σ^2 + (Aihigher - Ailower)*σ*ri*ν*η/(2*Δr))^2
        end

        function positive_consumption_constraint(x::Vector, grad::Vector)
            -(ri + x[1]*μ - 0.5*x[1]^2*σ^2)
        end

        opt = Opt(:LN_COBYLA, 1)
        lower_bounds!(opt, [-5])
        upper_bounds!(opt, [Inf])
        ftol_abs!(opt, 1e-20)

        min_objective!(opt, to_minimize)
        inequality_constraint!(opt, (x,g) -> positive_consumption_constraint(x,g), 1e-8)

        (minf, minx, ret) = optimize(opt, [0.7])

        αopt[i] = minx[1]
    end

    return αopt
end

function form_matrix_B(α, params)
    # This function forms matrix B that is used in the iteration
    # step to solve for vector of investment opportunities A
    # given a vector of optimal risky shares

    # Getting parameters:
    Δr = params.Δr;
    γ = params.γ;
    ν = params.ν;
    σ = params.σ;
    γ = params.γ;
    η = params.η;
    I = size(params.rgrid)[1];
    rbar = params.rbar;
    θ = params.θ;

    # Modifying for nonzero correlation between risky and risk free asset.
    # Log risk free rate follows a random walk (implies a positive drift for
    # the level of interest rate). Using central differences
    xvec = ν^2 * rgrid.^2/(2 *Δr^2) .- (1-γ)*α.*rgrid*σ*ν*η/(2*Δr) .- rgrid .* 0.5 * ν^2 *0.5/Δr;
    yvec = 0.5*(1-γ)^2 .* α.^2 .* σ^2 .- rgrid.^2 .*ν^2/Δr^2;
    zvec = ν^2 * rgrid.^2/(2 *Δr^2) .+ (1-γ)*α.*rgrid*σ*ν*η/(2*Δr) .+ rgrid .* 0.5 * ν^2 *0.5/Δr;
    v1 = xvec[1];
    vI = zvec[I];
    yvec[1] = yvec[1] + v1;
    yvec[I] = yvec[I] + vI;

    # Forming matrix B using
    B = spdiagm(-1 => xvec[2:I], 0=>yvec, 1=>zvec[1:(I-1)])
    return B
end

function calculate_utility_flow(α, params)
    μ = params.μ;
    γ = params.γ;
    σ = params.σ;
    rgrid = params.rgrid;
    return (rgrid .+ α.*μ .- 0.5.*α.^2 .* σ^2).^(1-γ)
end

function update_value_function(Acurrent, params)
    # 1. Given current A vector calculate optimal risky share
    αopt = solve_optimal_risky_share(Acurrent, params)

    # 2. Given optimal risky share calculate utility flow
    u = calculate_utility_flow(αopt, params)

    # 3. Given optimal alpha calculate "transition" matrix B
    B = form_matrix_B(αopt, params)

    # 4. Given old A vector, flow utility and "transition"
    # matrix update A vector
    I = size(params.rgrid)[1];
    ρ = params.ρ;
    Δ = 1;
    γ = params.γ;

    Aupdated = ((1/Δ + ρ) .* spdiagm(0 => ones(I)) .- B) \ (u .+ Acurrent./Δ);

    return Aupdated, αopt, u
end

function iterate_value_function(params, A_init, tol)
    A_prev = A_init;
    i = 0;
    dist = Inf
    while dist > tol
        i += 1;
        A_next,_ = update_value_function(A_prev, params);

        # Calculating the difference between value functions:
        dist = sum(abs.(A_next .- A_prev));
        A_prev = A_next;
        if i % 1000 == 0
            @show i
            @show dist
            print("\n")
        end
    end

    # Iterate one more step to get all parameters:
    return update_value_function(A_prev, params)
end

# Start with a guess of negative exponential declining A
rmin = 0;
rmax = 1;
I = 1001;
rgrid = range(rmin, stop = rmax, length = I);
Δr = (rmax - rmin)/(I-1);

I = size(rgrid)[1];
A_init = 10000 * exp.(-rgrid);

# params = ModelParameters(4, 0.04, 0.18, 0.01, 0.05, -0.25, rgrid,rmin,rmax,Δr)
params = ModelParameters(3, 0.04, 0.18, 0.01, 0.075, -0.0, rgrid,rmin,rmax,Δr,0.1,1)

Anew, αopt, u = update_value_function(A_init, params)
Plots.plot(Anew)
Plots.plot(αopt)
Plots.plot(u)

Anew, αopt, u = update_value_function(Anew, params)
Plots.plot(rgrid,Anew)
Plots.plot(rgrid,αopt)
Plots.plot(u*(1-params.γ))

# 1. Given current A vector calculate optimal risky share
αopt = solve_optimal_risky_share(A_init, params)

# 2. Given optimal risky share calculate utility flow
u = calculate_utility_flow(αopt, params)

# 3. Given optimal alpha calculate "transition" matrix B
B = form_matrix_B(αopt, params)
sum(B)

mathematicaColor1 = RGBA(0.37,0.51,0.71,1)
mathematicaColor2 = RGBA(0.88,0.61,0.14,1)
mathematicaColor3 = RGBA(0.56,0.69,0.19,1)

# Solving for different value of correlation:
params1 = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, 0.99, rgrid,rmin,rmax,Δr,0.02,1)
params2 = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, 0.0, rgrid,rmin,rmax,Δr,0.02,1)
params3 = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, -0.99, rgrid,rmin,rmax,Δr,0.02,1)

A1, α1, flow1 = iterate_value_function(params1, A_init, 1e-6)
A2, α2, flow2 = iterate_value_function(params2, A_init, 1e-6)
A3, α3, flow3 = iterate_value_function(params3, A_init, 1e-6)

Plots.plot(rgrid, α1, label = string("Corr = ", params1.η),
    xaxis = ("Risk Free Rate", (0,0.1)), yaxis = "Risky Share",legend=:bottomleft, lw = 2.75,grid=false,
    yguidefontsize=12,guidefontsize=12,legendfontsize=12,tickfontsize=12, color=mathematicaColor1)
Plots.plot!(rgrid, α2, label = string("Corr = ", params2.η), lw = 2.75, color=mathematicaColor2)
Plots.plot!(rgrid, α3, label = string("Corr = ", params3.η), lw = 2.75, color=mathematicaColor3)
Plots.savefig("../figures/share_dynamic_vary_corr.pdf")

# Plotting hedging demand: difference between α under non-zero correlation and
# α under zero correlation.

# Solving for different value of Risk Premium:
params4 = ModelParameters(3, 0.02, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)
params5 = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)
params6 = ModelParameters(3, 0.10, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)

A4, α4, flow4 = iterate_value_function(params4, A_init, 1e-4)
A5, α5, flow5 = iterate_value_function(params5, A_init, 1e-4)
A6, α6, flow6 = iterate_value_function(params6, A_init, 1e-7)

Plots.plot(rgrid, α4, label = string("RP = ", params4.μ*100,"%"),
    xaxis = ("Risk Free Rate", (0,0.1)), yaxis = "Risky Share",legend=:bottomleft, lw = 2.75,grid=false,
    yguidefontsize=12,guidefontsize=12,legendfontsize=12,tickfontsize=12, color=mathematicaColor1)
Plots.plot!(rgrid, α5, label = string("RP = ", params5.μ*100,"%"), lw = 2.75, color=mathematicaColor2)
Plots.plot!(rgrid, α6, label = string("RP = ", params6.μ*100,"%"), lw = 2.75, color=mathematicaColor3)
Plots.savefig("../figures/share_dynamic_vary_rp.pdf")

# Solving for different value of Discount Rate:
params7 = ModelParameters(3, 0.06, 0.18, 0.01, 0.050, 0, rgrid,rmin,rmax,Δr,0.02,1)
params8 = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)
params9 = ModelParameters(3, 0.06, 0.18, 0.01, 0.100, 0, rgrid,rmin,rmax,Δr,0.02,1)

A7, α7, flow7 = iterate_value_function(params7, A_init, 1e-7)
A8, α8, flow8 = iterate_value_function(params8, A_init, 1e-7)
A9, α9, flow9 = iterate_value_function(params9, A_init, 1e-7)

Plots.plot(rgrid, α7, label = string("Discount Rate  = ", params7.ρ*100,"%"),
    xaxis = ("Risk Free Rate", (0,0.1)), yaxis = "Risky Share",legend=:bottomleft, lw = 2.75,grid=false,
    yguidefontsize=12,guidefontsize=12,legendfontsize=12,tickfontsize=12, color=mathematicaColor1)
Plots.plot!(rgrid, α8, label = string("Discount Rate  = ", params8.ρ*100,"%"), lw = 2.75, color=mathematicaColor2)
Plots.plot!(rgrid, α9, label = string("Discount Rate  = ", params9.ρ*100,"%"), lw = 2.75, color=mathematicaColor3)
Plots.savefig("../figures/share_dynamic_vary_discount_rate.pdf")

# Solving for different value of Risk Aversion:
params10 = ModelParameters(2, 0.06, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)
params11 = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)
params12 = ModelParameters(4, 0.06, 0.18, 0.01, 0.075, 0, rgrid,rmin,rmax,Δr,0.02,1)

A10, α10, flow10 = iterate_value_function(params10, A_init, 1e-7)
A11, α11, flow11 = iterate_value_function(params11, A_init, 1e-7)
A12, α12, flow12 = iterate_value_function(params12, A_init, 1e-7)

Plots.plot(rgrid, α10, label = string("RRA = ", params10.γ),
    xaxis = ("Risk Free Rate", (0,0.1)), yaxis = "Risky Share",legend=:bottomleft, lw = 2.75,grid=false,
    yguidefontsize=12,guidefontsize=12,legendfontsize=12,tickfontsize=12, color=mathematicaColor1)
Plots.plot!(rgrid, α11, label = string("RRA = ", params11.γ), lw = 2.75, color=mathematicaColor2)
Plots.plot!(rgrid, α12, label = string("RRA = ", params12.γ), lw = 2.75, color=mathematicaColor3)
Plots.savefig("../figures/share_dynamic_vary_rra.pdf")

################################################################
# Numerically calculating drift of log consumption

params = ModelParameters(3, 0.06, 0.18, 0.01, 0.075, 0.0, rgrid,rmin,rmax,Δr,0.1,1)
A, α, flow=iterate_value_function(params, A_init, 1e-6)
γ = params.γ
σ = params.σ
μ = params.μ
ν = params.ν
ρ = params.ρ

dα = (α[2:(I-1)] .- α[1:(I-2)])./Δr;
ddα = (α[3:I] - 2*α[2:(I-1)] + α[1:(I-2)])./Δr;
α = α[2:(I-1)]
r = rgrid[2:(I-1)]
f = r .+ α * μ .- 0.5*α.^2 * σ^2;
df = 1 .+ dα * μ .- α .* dα * σ^2;
ddf = ddα * μ .- (dα.^2 + α .* ddα) * σ^2;
dlogc = 0.5*ν^2 * r .* (df .* f .+ ddf .* f - df.^2)./(f.^2);

Plots.plot(r,dlogc,xaxis = ("Risk Free Rate", (0,0.25)),label="",
    yaxis = "Log Consumption Drift",lw = 2.75,grid=false,
    yguidefontsize=12,guidefontsize=12,legendfontsize=12,tickfontsize=12, color=mathematicaColor1)
Plots.savefig("../figures/dynamic_model_log_consumption_drift.pdf")

# Calculating instantenous volatility of log(c):
voldlogc = ν * r .* df./f

Plots.plot(r,dlogc,xaxis = ("Risk Free Rate", (0,0.25)),label="Drift",
    yaxis = "Log Consumption Drift",lw = 2.75,grid=false,legend=:topleft,
    yguidefontsize=12,guidefontsize=12,legendfontsize=12,tickfontsize=12, color=mathematicaColor1)
Plots.plot!(r,voldlogc,lw = 2.75, color=mathematicaColor2, label = "Volatility")


#2-dim version

include("grid.jl")
include("boundaries.jl")
include("visualize.jl")
include("eos.jl")
include("reconstruct.jl")
include("rsolvers.jl")
include("gravity.jl")
include("viscosity.jl")
include("integrators.jl")

#using Winston
using Plots

#Primary variables to conserved variables
function prim2con(rho::AbstractMatrix,
                  velx::AbstractMatrix,
                  vely::AbstractMatrix,
                  eps::AbstractMatrix)

    q = zeros(size(rho, 1), size(rho, 2), 4)
    q[:, :, 1] = rho
    q[:, :, 2] = rho .* velx
    q[:, :, 3] = rho .* vely
    q[:, :, 4] = rho .* eps .+ 0.5rho .* (velx.^2.0 .+ vely.^2.0)

    return q
end

#Conserved variables to primary variables
function con2prim(q)
    rho = q[:, :, 1]
    velx = q[:, :, 2] ./ rho
    vely = q[:, :, 3] ./ rho
    eps = q[:, :, 4] ./ rho .- 0.5(velx.^2.0 .+ vely.^2.0)

    press = eos_press(rho, eps, gamma)

    return rho, eps, press, velx, vely
end

#time step calculation
function calc_dt(hyd, dtp)
    cs = sqrt(eos_cs2(hyd.rho, hyd.eps, gamma))
    dtnew = 1.0
    for j = (hyd.g+1):(hyd.ny-hyd.g+1), i = (hyd.g+1):(hyd.nx-hyd.g+1)
        dtnew = min(dtnew, (hyd.x[i+1] - hyd.x[i]) / max(abs(hyd.velx[j, i]+cs[j, i]), abs(hyd.velx[j, i]-cs[j, i])))
        dtnew = min(dtnew, (hyd.y[j+1] - hyd.y[j]) / max(abs(hyd.vely[j, i]+cs[j, i]), abs(hyd.vely[j, i]-cs[j, i])))
    end

    dtnew = min(cfl*dtnew, 1.05*dtp)

    return dtnew
end

#snapshot of the simulation
function snapshot(hyd, var, name, it)

    #cm = Uint32[Color.convert(Color.RGB24,c) for c in flipud(Color.colormap("RdBu"))]
    xs = hyd.g+1
    xe = hyd.nx-hyd.g
    ye = hyd.ny-hyd.g

    hdata = var[xs:ye, xs:xe]
    #pf=FramedPlot()
    clims = (0., 2.5)
    #img = Winston.data2rgb(hdata, clims, cm)
	  #img = Winston.imagesc(hdata)
    #add(pf, Image((hyd.x[xs], hyd.x[xe]), (hyd.y[xs], hyd.y[ye]), img;))
		pf = heatmap(hdata,ratio=:equal)
		#add(pf, img)
    #setattr(pf, xrange=(hyd.x[xs], hyd.x[xe]))
    #setattr(pf, yrange=(hyd.y[xs], hyd.y[ye]))
    #setattr(pf, title="rho")

    #setattr(pf.frame,
    #        draw_spine=false,
    #        draw_ticks=false,
    #        draw_ticklabels=false)

    numb = lpad(it, 3, '0')
    savefig(pf,"$(name)_$(numb).png")
		closeall()

    return nothing
end

#Construct the right hand side of the group of equations
function calc_rhs(hyd, dt)

    #reconstruction and prim2con
    hyd = reconstruct(hyd)

    #apply source terms for the interfaces
    #hyd = source_terms(hyd, dt)

    #unsplit flux
    fx, fy = uflux(hyd, dt)

    #add artificial viscosity
    #vfx, vfy = artificial_viscosity(hyd, dt)
    #fx += vfx
    #fy += vfy

    #flux difference
    fluxdiff = zeros(hyd.ny, hyd.nx, 4)
    for j = (hyd.g-1):(hyd.ny-hyd.g+1), i = (hyd.g-1):(hyd.nx-hyd.g+1)
        dx = abs(hyd.x[i+1]-hyd.x[i])
        dy = abs(hyd.y[j+1]-hyd.y[j])

        for k = 1:4
            fluxdiff[j, i, k] = (fx[j, i, k] - fx[j, i-1, k])/dx + (fy[j, i, k] - fy[j-1, i, k])/dy
        end
    end

    return -fluxdiff
end

###############
# main program
###############

function evolve(hyd, tend, gamma, cfl, nx, ny)

    time = range(0.0, length=100, stop=tend)
    it = 1

    dt = 1.0e-8
    dtp = dt

    #get initial timestep
    dt = calc_dt(hyd, dt)

    #initial prim2con
    hyd.q = prim2con(hyd.rho, hyd.velx, hyd.vely, hyd.eps)

    t = 0.0
    i = 0

    visualize(hyd)

    while t < tend

        if i % 10 == 0
            println("$i $t $dt")
            #visualize(hyd)
        end

        if time[it] <= t
            snapshot(hyd, hyd.rho, "density", it)
            snapshot(hyd, hyd.press, "pressure", it)
            it += 1
        end

        #calculate new timestep
        dt = calc_dt(hyd, dt)

        #time integration
        hyd = RK2(hyd, dt)

        #update time
        t += dt
        i += 1

    end

    return hyd
end

#basic parameters
gamma = 1.4
cfl = 0.4

nx = 50
ny = 50
tend = 1.0e-1

#initialize
hyd = data2d(nx, ny)
#r32x, r32y = r32kernel(hyd.x, hyd.y)

#set up grid

include("../tests/1dim_hd.jl")
include("../tests/2dim_hd.jl")

hyd = grid_setup(hyd, 0.0, 1.0, 0.0, 1.0)
#hyd = grid_setup(hyd, 0.0, 1.0, 0.0, 0.5)
#hyd = grid_setup(hyd, -0.25, 0.25, -0.75, 0.75)
#hyd = grid_setup(hyd, 0., 1., -0.5, 0.5)

#set up initial data
#hyd = setup_taylor2(hyd)
hyd = setup_blast(hyd)
#hyd = setup_tubexy(hyd)
#hyd = setup_tubey(hyd)
#hyd = setup_tubex(hyd)
#hyd = setup_collision(hyd)
#hyd = setup_fall(hyd)
#hyd = setup_kh(hyd)

#visualize(hyd)

#main integration loop
hyd = evolve(hyd, tend, gamma, cfl, nx, ny)
#hyd = @profile evolve(hyd, tend, gamma, cfl, nx, ny)


function CZ(ALPHA,BETA,EL)
      A= [0.770;   0.241;  -0.100;  -0.416;  -0.731; -1.053; -1.366;
          -1.646; -1.917; -2.120; -2.248; -2.229]
      #
      S= 0.2 * ALPHA
      K= floor(S)+3
      if(K <= 1)
          K= 2
      end
      if(K >=  12)
          K=  11
      end
      DA= S - K + 3
      L = K + fix( 1.1*sign(DA) )
      K=convert(Int64, K)
      L=convert(Int64, L)
      S= A[K] + abs(DA) * (A[L] - A[K])
      return czv = S*(1-(BETA/57.3)^2) - .19*(EL/25.0)
end


################################
###    Summary statistics    ###
################################
"""
	poisson_fixation(empirical_values,λds, λdn)

Divergence sampling from Poisson distribution. The expected neutral and selected fixations are subset through their relative expected rates ([`fixNeut`](@ref), [`fixNegB`](@ref), [`fixPosSim`](@ref)). Empirical values are used are used to simulate the locus *L* along a branch of time *T* from which the expected *Ds* and *Dn* raw count estimated given the mutation rate (``\\mu``). Random number generation is used to subset samples arbitrarily given the success rate ``\\lambda`` in the distribution.

```math
\\mathbb{E}[D_N] = X \\in Poisson\\left(\\lambda = D \\times \\left[\\frac{\\mathbb{E}[D_+] + \\mathbb{E}[D_-]}{\\mathbb{E}[D_+] + \\mathbb{E}[D_-] + \\mathbb{E}[D_0]}\\right]\\right)
```
```math
\\mathbb{E}[D_S] = X \\in Poisson\\left(\\lambda = D \\times \\left[\\frac{\\mathbb{E}[D_0]}{\\mathbb{E}[D_+] + \\mathbb{E}[D_-] + \\mathbb{E}[D_0]}\\right]\\right)
```
# Arguments
 - `observedValues::Array`: Array containing the total observed divergence.
 - ` λds::Float64`: expected neutral fixations rate.
 - ` λdn::Float64`: expected selected fixations rate.
# Returns
 - `Array{Int64,1}` containing the expected count of neutral and selected fixations.

"""
function poisson_fixation(;empirical_values::Array, λds::Array, λdn::Array,λweak::Array,λstrong::Array)

    ds             = @. λds / (λds + λdn) * empirical_values
    dn             = @. λdn / (λds + λdn) * empirical_values
    dn_weak        = @. λweak / (λds + λdn) * empirical_values
    dn_strong      = @. λstrong / (λds + λdn) * empirical_values

    sampled_Ds     = pois_rand.(ds)
    sampled_Dn     = pois_rand.(dn)
    sampled_weak   = pois_rand.(dn_weak)
    sampled_strong = pois_rand.(dn_strong)

    alphas         = @. [sampled_weak/sampled_Dn sampled_strong/sampled_Dn (sampled_weak+sampled_strong)/sampled_Dn]

	out            = alphas,sampled_Dn, sampled_Ds
	return out
end

"""
	poisson_polymorphism(empirical_values,λps,λpn)

Polymorphism sampling from Poisson distributions. The total expected neutral and selected polimorphism are subset through the relative expected rates at the frequency spectrum ([`fixNeut`](@ref), [`DiscSFSNeutDown`](@ref),). Empirical sfs are used to simulate the locus *L* along a branch of time *T* from which the expected *Ps* and *Pn* raw count are estimated given the mutation rate (``\\mu``). Random number generation is used to subset samples arbitrarily from the whole sfs given each frequency success rate ``\\lambda`` in the distribution.

The success rate managing the Poisson distribution by the observed count each frequency.  We considered both sampling variance and process variance is affecting the number of variable alleles we sample from SFS. This variance arises from the random mutation-fixation process along the branch. To incorporate this variance we do one sample per frequency-bin and use the total sampled variation and the SFS for the summary statistics.

```math
\\mathbb{E}[P_N] = \\sum_{x=0}^{x=1} X \\in Poisson\\left(\\lambda = SFS_{(x)} \\times \\left[\\frac{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}]}{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}] + \\mathbb{E}[P_{0(x)}]}\\right]\\right)
```

```math
\\mathbb{E}[P_S] = \\sum_{x=0}^{x=1} X \\in Poisson\\left(\\lambda = SFS_{(x)} \\times \\left[\\frac{\\mathbb{E}[P_{0(x)}]}{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}] + \\mathbb{E}[P_{0(x)}]}\\right]\\right)
```

# Arguments
 - `observedValues::Array{Int64,1}`: Array containing the total observed divergence.
 - ` λps::Array{Float64,1} `: expected neutral site frequency spectrum rate.
 - ` λpn::Array{Float64,1} `: expected selected site frequency spectrum rate.
# Returns
 - `Array{Int64,2}` containing the expected total count of neutral and selected polymorphism.

"""
function poisson_polymorphism(;empirical_values::Array, λps::Array, λpn::Array,pol::Array)

	# Neutral λ;156948
	λ1 = @. λps / (λps + λpn) * empirical_values'
	# Selected λ;
	λ2 = @. λpn / (λps + λpn) * empirical_values'

	# Neutral λ;
	#=λ1 = @. λps / pol[:,1] * observedValues[:,2]'
	# Selected λ;
	λ2 = @. λpn / pol[:,2] * observedValues[:,1]'=#

	# Replace negative and 0 rates values. Strange behaviour depending on θ and Γ values
	# Relative rates output NaN values due to 0 divisons.

	replace!(λ1,NaN=>1)
	replace!(λ2,NaN=>1)

	sampled_Ps = pois_rand.(λ1)
	sampled_Pn = pois_rand.(λ2)

	return (sampled_Pn, sampled_Ps)
end

"""

	Ouput the expected values from the Poisson sampling process. Please check [`poisson_fixation`](@ref) and [`poisson_polymorphism`](@ref) to understand the samplingn process. α(x) is estimated through the expected values of Dn, Ds, Pn and Ps.

	# Arguments
	 - `afs::Vector`: Array containing SFS subset
	 - `divergence::Vector`: Array containing divergence counts
	 - `observed::Matrix`: Array containing SFS
	 - `m::Array`: Array containing modeled parameters
	 - `nt::Array`: Array containing neutral polymorphic rates at frequencies dac
	 - `sl::Array`: Array containing selective polymorphic rates at frequencies dac
	 - `d::Array`: Array containing fixation rates
	 - `p::Array`: Array containing total polymorphic rates
	 - `dac::Array`: Array containing Derived Allel Counts to subset
	# Returns
	Expected α_x values
	 - `Array{Int64,2}` true α, γ, shape parameter and α_x values.

	sampling_from_rates(afs,divergence,observed,m,nt,d,p,dac)
"""
function sampling_from_rates(afs::Vector,divergence::Vector,observed::Matrix,m::Array,nt::Array,sl::Array,d::Array,p::Array,dac::Array)

	ds             = d[:,1]
	dn             = d[:,2]
	dn_weak        = d[:,3]
	dn_strong      = d[:,4]
	gn             = abs.(m[:,4])
	sh             = round.(m[:,end-1],digits=5)

	#alxSummStat, alphasDiv, expectedDn, expectedDs, expectedPn, expectedPs = sampledAlpha(divergence=divergence,afs=s,observed=observed,λdiv=[ds,dn,dn_weak,dn_strong],λpol=[permutedims(nt),permutedims(sl)],dac=dnac)

	## Outputs
	alphas, expDn, expDs = poisson_fixation(empirical_values=divergence,λds=ds,λdn=dn,λweak=dn_weak,λstrong=dn_strong)
	expPn, expPs         = poisson_polymorphism(empirical_values=afs,λps=nt,λpn=sl,pol=p)

	θ_pn = @. observed[1,1] / (expPn[:,1]);
	θ_ps = @. observed[1,2] / (expPs[:,1]);

	r_pn = ones(size(expPn));
	r_ps = ones(size(expPs));
	for j=2:size(dac,1)
		@.  r_pn[:,j] = (observed[j,1]) / (expPn[:,j]);
		@. r_pn[:,j] = r_pn[:,j] / θ_pn;

		@. r_ps[:,j] = (observed[j,2]) / (expPs[:,j]);
		@. r_ps[:,j] = r_ps[:,j] / θ_ps;
	end

	expPn_rj = zeros(size(expPn))
	expPs_rj = zeros(size(expPs))

	expPn_rj[:,1] = @. expPn[:,1] * θ_pn;
	expPs_rj[:,1] = @. expPs[:,1] * θ_ps;
	for j = 2:size(dac,1)
		@. expPn_rj[:,j] =  expPn[:,j] * θ_pn * r_pn[:,j];
		@. expPs_rj[:,j] =  expPs[:,j] * θ_ps * r_ps[:,j];
	end

	## Alpha from expected values. Used as summary statistics
    α_x    = @. round(1 - ((expDs/expDn) * (expPn/expPs)),digits=5)
    α_x_rj	 = @. round(1 - ((expDs/expDn) * (expPn_rj/expPs_rj)),digits=5)

	expected_values = hcat(round.(alphas,digits=5),gn,sh,α_x)

	return(expected_values)
end

"""
	summary_statistics(param::parameters,rates::JLD2.JLDFile,analysis_folder::String,summstat_size::Int64,replicas::Int64,bootstrap::Bool)

Estimate summary statistics using observed data and analytical rates. *analysis_folder* will check for the SFS and divergence file and will be used to output summary statistics

# Arguments
 - `param::parameters` : Mutable structure containing the models
 - `rates::JLD2.JLDFile` : HDF5 containing solved models, fixation and polymorphic rates
 - `analysis_folder::String` : Folder containing the SFS and divergence files. It will be used to output the observed data and summary estatistics.
 - `summstat_size::Int64` : Number of summary statistics
 - `replicas::Int64` : Number of bootstrap replicas
 - `bootstrap::Bool` : Boolean to perform or not bootstrapping
# Output
 - Obserded data and summary statistics to ABC inference
"""
function summary_statistics(;param::parameters,h5_file::String,analysis_folder::String,summstat_size::Int64,replicas::Int64=1,bootstrap::Bool=false)

	# param = parameters(n=50,dac=[1,2,4,5,10,20,25,30,50,70]);
	# analysis_folder="/home/jmurga/test/simulations/prueba/";
	# # h5_file="/home/jmurga/test/simulations/prueba/prueba.jld2";
	# h5_file="/home/jmurga/test/simulations/rates.jld2";
	# replicas=1;
	# bootstrap=true;
	# summstat_size=10^5

	#Opening files
	sfs_files        = filter(x -> occursin("sfs",x), readdir(analysis_folder,join=true));
	divergence_files = filter(x -> occursin("div",x), readdir(analysis_folder,join=true));

	sfs,divergence,α,observed = open_sfs_div(sfs_files,divergence_files,param.dac,replicas,bootstrap);

	#Open rates
	h5    = jldopen(h5_file);
	h5_subset   = h5[string(param.N) * "/" * string(param.n)];

	#Subset index
	idx    = sample(1:size(h5_subset["models"],1),summstat_size,replace=false);
	models = Array(view(h5_subset["models"],idx,:));
	dsdn   = Array(view(h5_subset["dsdn"],idx,:));
	pol    = Array(view(h5_subset["pol"],idx,:));

	# Filtering polymorphic rate by dac
	n    = hcat(map(x -> view(h5_subset["neut"][x],:),param.dac)...);
	s    = hcat(map(x -> view(h5_subset["sel"][x],:),param.dac)...);
	neut = Array(view(n,idx,:));
	sel  = Array(view(s,idx,:));

	#Making summaries
	f(afs,divergence,observed,m=models,nt=neut,sl=sel,d=dsdn,p=pol,dac=param.dac) = sampling_from_rates(afs,divergence,observed,m,nt,sl,d,p,dac);

	expected_values = [hcat(models[:,1:5],@. round(1 - (dsdn[:,1]/dsdn[:,2] * sel/neut),digits=5))]

	expected_values = progress_pmap(f,sfs,divergence,observed;progress = Progress(size(sfs,1),desc="Sampling αₓ "));

	# Controling outlier cases
	expected_values = pmapbatch(filter_expected,expected_values);

	# nuisance =  [[1.16  0.62  0.32  0.04  1.  1.  1. 1.  1.  1.],[1.08  0.7   0.2   0.41  1.51  1.3   1.18  1.01  0.94  1.04],[1.16  0.62  0.32  0.04  1.11  1.12  1.03  0.92  0.93  1.02],[0.76  0.91  1.44  1.98  0.22  0.46  0.58  0.7  0.77  1.01],[0.8   1.03  1.91  3.02  1.6   0.45  0.16  0.05  0.14  0.26]];

	# expected_values = progress_pmap(y -> vcat(y,vcat(map(x-> nuisance_alpha(y,x),nuisance)...)),expected_values)

	# flt_b(x) = any(x .> 1)
	# flt = pmapbatch(y -> .!mapslices(flt_b,y[:,6:end],dims=2)[:,1],expected_values)

	# flt_v = (x,y) -> @view x[y,:]

	# expected_values = pmapbatch(flt_v,expected_values,flt)

	w(x,name) = CSV.write(name,DataFrame(x,:auto),delim='\t',header=false);

	progress_map(w, α, analysis_folder * "/alphas_" .* string.(1:size(sfs,1)) .* ".txt";progress= Progress(size(sfs_files,1),desc="Writting α "));

	progress_pmap(w, expected_values,  analysis_folder * "/summstat_" .* string.(1:size(sfs,1)) .* ".txt";progress= Progress(size(sfs,1),desc="Writting summaries "));

	return(expected_values)
end

function filter_expected(x::Matrix{Float64})

	replace!(x, -Inf=>NaN)
	x = x[vec(.!any(isnan.(x),dims=2)),:]
	x = x[(x[:,3] .< 1),:]

	return(x)
end

# function nuisance_alpha(a::Matrix{Float64},b::Matrix{Float64})

# 		tmp_param = @view a[:,1:5];
# 		tmp_summs = @view a[:,6:end];

# 		tmp = round.(tmp_summs .* b,digits=5)

# 		out = hcat(tmp_param,tmp);
# 		return(out)
# end


module Ramnit
mutable struct RANDINT
    seed
end

function rand_int_modulus!(r, modulus)
    ix = r.seed
    ix = 16807*(ix % 127773) - 2836*trunc(Int64, ix/127773) & 0xFFFFFFFF
    r.seed = ix
    return ix % modulus
end

function gen_one(seed=0x123abc12, tld="")
    domain = ""
    r = RANDINT(seed)
    seed_a = r.seed
    domain_len = rand_int_modulus!(r,12) + 8
    seed_b = r.seed
    for i in 1:domain_len
        domain *= 'a' + rand_int_modulus!(r,25)
    end
    domain *= tld
    return domain
end

function gen_many(cnt=100)
    res = []
    r = RANDINT(0x123abc12)
    for _ in 1:cnt
        seed_a = r.seed
        domain_len = rand_int_modulus!(r,12) + 8
        seed_b = r.seed
        domain = ""
        for i in 1:domain_len
            domain *= 'a' + rand_int_modulus!(r,25)
        end
        m = seed_a*seed_b
        r.seed = (m + trunc(Int64, m/(2^32))) % 2^32
        res =[res..., domain]
    end
    return res
end

end  # module Ramnit


using ProgressMeter

"""
    likelihood(yhat, experimental_ecdf, experimental, I, ν0)

Find the likelihood of a specific set of parameters (yhat), which is based on
the ordinary least squares of the experimental vs. theoretical CDFs

# Arguments
- `yhat`: an array containing parameters for qcc, η, and the standard deviation
- `experimental_ecdf`: the estimated ECDF of the experimental data
- `experimental`: the experimental data
- `I`: spin (3 in the case of Boron-10)
- `ν0`: the Larmor frequency

"""
function likelihood(yhat, experimental_ecdf, experimental, I, ν0)
    # We must choose a distribution for the error!
    yhat[5] < 0 && return 0
    likelihood_dist = Normal(0, yhat[5])

    # Here we're generating the theoretical sample and converting it to a CDF
    quad = Quadrupolar(yhat[1:4])
    ismissing(quad) && return 0
    powder_pattern = estimate_powder_pattern(quad, 100_000, ν0, I)
    theoretical_ecdf = ecdf(powder_pattern).(experimental[:, 1])
    theoretical_ecdf .-= theoretical_ecdf[1]
    theoretical_ecdf ./= theoretical_ecdf[end]
    # Then we return the likelihood, based on two CDFs' differences
    return sum(logpdf.(likelihood_dist, experimental_ecdf .- theoretical_ecdf))
end

function likelihood(yhat, experimental_ecdf, experimental; sites = 1)
    if sites == 1
        # We must choose a distribution for the error!
        yhat[7] < 0 && return 0
        likelihood_dist = Normal(0, yhat[7])

        # Here we're generating the theoretical sample and converting it to a CDF
        csa = ChemicalShift(vcat(yhat[1:6], [1.0]))
        ismissing(csa) && return 0
        powder_pattern = estimate_powder_pattern(csa, 100_000)
        theoretical_ecdf = ecdf(powder_pattern).(experimental[:, 1])

        # Then we return the likelihood, based on two CDFs' differences
        return sum(logpdf.(likelihood_dist, experimental_ecdf .- theoretical_ecdf))
    elseif sites == 2
        # We must choose a distribution for the error!
        yhat[14] < 0 && return 0
        likelihood_dist = Normal(0, yhat[14])

        # Here we're generating the theoretical sample and converting it to a CDF
        csa = ChemicalShift(vcat(yhat[1:6], [yhat[13]], yhat[7:12], [1 - yhat[13]]))
        ismissing(csa) && return 0
        powder_pattern = estimate_powder_pattern(csa, 100_000)
        theoretical_ecdf = ecdf(powder_pattern).(experimental[:, 1])

        # Then we return the likelihood, based on two CDFs' differences
        return sum(logpdf.(likelihood_dist, experimental_ecdf .- theoretical_ecdf))
    end
end

"""
    metropolis_hastings(experimental)

An implementation of the Metropolis-Hastings algorithm for Monte Carlo Markov
Chain (MCMC) Bayesian estimates, with optional keyword arguments for the number
of samples (default 1,000,000) and tolerances for each parameter (default 0.1
for qcc, 0.1 for η, 0.05 for σ)

"""
function metropolis_hastings(
    experimental;
    interaction = "quadrupolar",
    N = 1_000_000,
    sites = 1,
    tol = interaction == "quadrupolar" ? [0.1, 0.2, 0.1, 0.1, 0.05] : sites == 1 ? [100.0, 20.0, 30.0, 5.0, 0.1, 0.05, 0.05] : [10.0, 5.0, 10.0, 5.0, 0.1, 0.05, 10.0, 5.0, 10.0, 5.0, 0.1, 0.05, 0.01, 0.05],
    I = 3
)
    experimental_ecdf = get_experimental_ecdf(experimental)
    if interaction == "quadrupolar"
        ν0 = get_ν0(experimental, experimental_ecdf)

        samples = zeros(N, 5)  # initialize zero array for samples

        #We need to define prior distributions for each parameter
        prior_dist_qcc = Uniform(0, 9)
        prior_dist_σqcc = Uniform(0, 3)
        prior_dist_η = Uniform(0, 1)
        prior_dist_ση = Uniform(0, 1)
        prior_dist_σ = Uniform(0, 1)

        prior_quad(x) = logpdf(prior_dist_qcc, x[1]) + logpdf(prior_dist_σqcc, x[2]) +
            logpdf(prior_dist_η, x[3]) + logpdf(prior_dist_ση, x[4]) +
            logpdf(prior_dist_σ, x[5]) # Assume variables are independent

        a = [rand(prior_dist_qcc), rand(prior_dist_σqcc), rand(prior_dist_η),
            rand(prior_dist_ση),rand(prior_dist_σ)]
        samples[1, :] = a

        # Repeat N times
        @showprogress for i = 2:N
            b = a + tol .* (rand(5) .- 0.5)  # Compute new state randomly
            # Calculate density
            prob_old = likelihood(a, experimental_ecdf, experimental, I, ν0) +
                       prior_quad(a)
            prob_new = likelihood(b, experimental_ecdf, experimental, I, ν0) +
                       prior_quad(b)
            r = prob_new - prob_old # Compute acceptance ratio
            if log(rand()) < r
                a = b  # Accept new state and update
            end
            samples[i, :] = a  # Update state
        end
    else
        if sites == 1
            samples = zeros(N, 7)  # initialize zero array for samples

            #We need to define prior distributions for each parameter
            prior_dist_σᵢₛₒ = Uniform(minimum(experimental[:, 1]), maximum(experimental[:, 1]))
            prior_dist_σσᵢₛₒ = Uniform(0.00001, 800)
            prior_dist_Δσ = Uniform(-4000, 4000)
            prior_dist_σΔσ = Uniform(0.00001, 400)
            prior_dist_ησ = Uniform(0, 1)
            prior_dist_σησ = truncated(Normal(0.00001, 0.2), 0.0, Inf)
            prior_dist_σ = truncated(Normal(0.00001, 0.5), 0.0, Inf)

            prior_csa(x) = logpdf(prior_dist_σᵢₛₒ, x[1]) + logpdf(prior_dist_σσᵢₛₒ, x[2]) +
                logpdf(prior_dist_Δσ, x[3]) + logpdf(prior_dist_σΔσ, x[4]) +
                logpdf(prior_dist_ησ, x[5]) + logpdf(prior_dist_σησ, x[6]) +
                logpdf(prior_dist_σ, x[7]) # Assume variables are independent

            a = [rand(prior_dist_σᵢₛₒ), rand(prior_dist_σσᵢₛₒ), rand(prior_dist_Δσ),
                rand(prior_dist_σΔσ), rand(prior_dist_ησ), rand(prior_dist_σησ),
                rand(prior_dist_σ)]
            samples[1, :] = a

            # Repeat N times
            @showprogress for i = 2:N
                b = a + tol .* (rand(7) .- 0.5)  # Compute new state randomly
                # Calculate density
                prob_old = likelihood(a, experimental_ecdf, experimental) + prior_csa(a)
                prob_new = likelihood(b, experimental_ecdf, experimental) + prior_csa(b)
                r = prob_new - prob_old # Compute acceptance ratio
                if log(rand()) < r
                    a = b  # Accept new state and update
                end
                samples[i, :] = a  # Update state
            end
        elseif sites == 2
            samples = zeros(N, 14)  # initialize zero array for samples

            #We need to define prior distributions for each parameter
            prior_dist_σᵢₛₒ1 = Uniform(minimum(experimental[:, 1]), maximum(experimental[:, 1]))
            prior_dist_σσᵢₛₒ1 = Uniform(0.00001, 800)
            prior_dist_Δσ1 = Uniform(-4000, 4000)
            prior_dist_σΔσ1 = Uniform(0.00001, 400)
            prior_dist_ησ1 = Uniform(0, 1)
            prior_dist_σησ1 = truncated(Normal(0.00001, 0.2), 0.0, Inf)
            prior_dist_σᵢₛₒ2 = Uniform(minimum(experimental[:, 1]), maximum(experimental[:, 1]))
            prior_dist_σσᵢₛₒ2 = Uniform(0.00001, 800)
            prior_dist_Δσ2 = Uniform(-4000, 4000)
            prior_dist_σΔσ2 = Uniform(0.00001, 400)
            prior_dist_ησ2 = Uniform(0, 1)
            prior_dist_σησ2 = truncated(Normal(0.00001, 0.2), 0.0, Inf)
            prior_dist_w = Uniform(0, 1)
            prior_dist_σ = truncated(Normal(0.00001, 0.5), 0.0, Inf)

            a = [rand(prior_dist_σᵢₛₒ1), rand(prior_dist_σσᵢₛₒ1), rand(prior_dist_Δσ1),
                rand(prior_dist_σΔσ1), rand(prior_dist_ησ1), rand(prior_dist_σησ1),
                rand(prior_dist_σᵢₛₒ2), rand(prior_dist_σσᵢₛₒ2), rand(prior_dist_Δσ2),
                rand(prior_dist_σΔσ2), rand(prior_dist_ησ2), rand(prior_dist_σησ2),
                rand(prior_dist_w),
                rand(prior_dist_σ)]
            samples[1, :] = a

            # Repeat N times
            @showprogress for i = 2:N
                b = a + tol .* (rand(14) .- 0.5)  # Compute new state randomly

                # Calculate density
                prob_old = likelihood(a, experimental_ecdf, experimental, sites = 2) +
                logpdf(prior_dist_σᵢₛₒ1, a[1]) + logpdf(prior_dist_σσᵢₛₒ1, a[2]) +
                   logpdf(prior_dist_Δσ1, a[3]) + logpdf(prior_dist_σΔσ1, a[4]) +
                   logpdf(prior_dist_ησ1, a[5]) + logpdf(prior_dist_σησ1, a[6]) +
                   logpdf(prior_dist_σᵢₛₒ2, a[7]) + logpdf(prior_dist_σσᵢₛₒ2, a[8]) +
                   logpdf(prior_dist_Δσ2, a[9]) + logpdf(prior_dist_σΔσ2, a[10]) +
                   logpdf(prior_dist_ησ2, a[11]) + logpdf(prior_dist_σησ2, a[12]) +
                   logpdf(prior_dist_w, a[13]) +
                   logpdf(prior_dist_σ, a[14])
                prob_new = likelihood(b, experimental_ecdf, experimental, sites = 2) +
                logpdf(prior_dist_σᵢₛₒ1, b[1]) + logpdf(prior_dist_σσᵢₛₒ1, b[2]) +
                   logpdf(prior_dist_Δσ1, b[3]) + logpdf(prior_dist_σΔσ1, b[4]) +
                   logpdf(prior_dist_ησ1, b[5]) + logpdf(prior_dist_σησ1, b[6]) +
                   logpdf(prior_dist_σᵢₛₒ2, b[7]) + logpdf(prior_dist_σσᵢₛₒ2, b[8]) +
                   logpdf(prior_dist_Δσ2, b[9]) + logpdf(prior_dist_σΔσ2, b[10]) +
                   logpdf(prior_dist_ησ2, b[11]) + logpdf(prior_dist_σησ2, b[12]) +
                   logpdf(prior_dist_w, b[13]) +
                   logpdf(prior_dist_σ, b[14])
                r = prob_new - prob_old # Compute acceptance ratio
                if log(rand()) < r
                    a = b  # Accept new state and update
                end
                samples[i, :] = a  # Update state
            end
        end
    end

    return samples
end;


using Gen

@gen function my_model(xs)
    slope = @trace(normal(0, 2), :slope)
    intercept = @trace(normal(0, 10), :intercept)
    for (i, x) in enumerate(xs)
        @trace(normal(slope * x + intercept, 1), "y-$i")
    end
end

function my_inference_program(xs, ys, num_iters)
    constraints = choicemap()
    for (i, y) in enumerate(ys)
        constraints["y-$i"] = y
    end
    (trace, _) = generate(my_model, (xs,), constraints)
    slope_selection = select(:slope)
    intercept_selection = select(:intercept)
    for iter=1:num_iters
        (trace, _) = metropolis_hastings(trace, slope_selection)
        (trace, _) = metropolis_hastings(trace, intercept_selection)
    end
    return (trace[:slope], trace[:intercept])
end

xs = [1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]
ys = [8.23, 5.87, 3.99, 2.59, 0.23, -0.66, -3.53, -6.91, -7.24, -9.90]
(slope, intercept) = my_inference_program(xs, ys, 100)
println("slope: $slope, intercept: $intercept")


module PhysCon

using Parameters
#-----------------------------------------------------------------------
# DESCRIPTION:
# Calculate saturation vapor pressure and latent heat of vaporization
#
export physcon
@with_kw struct phys{T<:Number}
    grav::T = 9.80665;               # Gravitational acceleration (m/s2)
    tfrz::T = 273.15;                # Freezing point of water (K)
    σ::T = 5.67e-08;                 # Stefan-Boltzmann constant (W/m2/K4)
    mmdry::T = 28.97 / 1000;         # Molecular mass of dry air (kg/mol)
    mmh2o::T = 18.02 / 1000;         # Molecular mass of water (kg/mol)
    Cpd::T = 1005;                   # Specific heat of dry air at constant pressure (J/kg/K)
    Cpw::T = 1846;                   # Specific heat of water vapor at constant pressure (J/kg/K)
    Rgas::T = 8.31446;               # Universal gas constant (J/K/mol)
    visc_0::T = 13.3e-06;            # Kinematic viscosity at 0C and 1013.25 hPa (m2/s)
    Dh_0::T = 18.9e-06;              # Molecular diffusivity (heat) at 0C and 1013.25 hPa (m2/s)
    Dv_0::T = 21.8e-06;              # Molecular diffusivity (H2O) at 0C and 1013.25 hPa (m2/s)
    Dc_0::T = 13.8e-06;              # Molecular diffusivity (CO2) at 0C and 1013.25 hPa (m2/s)
end

physcon = phys{Float32}()

end