#=

  Jobs puzzle in ConstraintSolver.jl

  This is a standard problem in Automatic Reasoning.

  From http://www-unix.mcs.anl.gov/~wos/mathproblems/jobs.html
  """
  Jobs Puzzle

  There are four people:  Roberta, Thelma, Steve, and Pete.
  Among them, they hold eight different jobs.
  Each holds exactly two jobs.
  The jobs are chef, guard, nurse, clerk, police officer (gender
  not implied), teacher, actor, and boxer.
  The job of nurse is held by a male.
  The husband of the chef is the clerk.
  Roberta is not a boxer.
  Pete has no education past the ninth grade.
  Roberta, the chef, and the police officer went golfing together.

  Question:  Who holds which jobs?
  """

  The answer:
  Chef       Thelma
  Guard      Roberta
  Nurse      Steve
  Clerk      Pete
  Police     Steve
  Teacher    Roberta
  Actor      Pete
  Boxer      Thelma

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function jobs_puzzle(print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
   num_people = 4
   num_jobs = 8
   people = 1:num_people
   Roberta, Thelma, Steve, Pete = people

   @variable(model, 1 <= Jobs[1:num_jobs] <= num_people, Int)
   Chef, Guard, Nurse, Clerk, PoliceOfficer, Teacher, Actor, Boxer = Jobs

   # Each holds exactly two jobs.
   for i in 1:4
        count_ctr(model, Jobs, :(==), i, 2)
   end

   # The job of nurse is held by a male.
   # (Nurse == Steve \/ Nurse == Pete),
   # either_eq(model, Nurse,Steve, Nurse,Pete)
   # is_member_of(model,Nurse, [Steve,Pete])
   b1 = @variable(model, binary=true)
   @constraint(model, b1 := {Nurse == Steve || Nurse == Pete})
   @constraint(model, b1 == 1)

   # @either_eq(model, :(Nurse==Steve), :(Nurse==Pete))

   # The husband of the chef is the clerk.
   # (Clerk == Steve   \/ Clerk == Pete),
   # either_eq(model, Clerk,Steve, Clerk,Pete)
   # is_member_of(model,Clerk, [Steve,Pete])
   b2 = @variable(model, binary=true)
   @constraint(model, b2 := {Clerk == Steve || Clerk == Pete})
   @constraint(model, b2 == 1)

   # (Chef  #= Roberta #\/ Chef #= Thelma),
   # either_eq(model, Chef,Roberta, Chef,Thelma)
   # is_member_of(model,Chef, [Roberta,Thelma])
   b3 = @variable(model, binary=true)
   @constraint(model, b3 := {Chef  == Roberta || Chef == Thelma})
   @constraint(model, b3 == 1)

   @constraint(model, Chef != Clerk)

   # Roberta is not a boxer.
   @constraint(model, Roberta != Boxer)

   # Pete has no education past the ninth grade.
   @constraint(model, Pete != Teacher)
   @constraint(model, Pete != PoliceOfficer)
   @constraint(model, Pete != Nurse)

   # Roberta, [and] the chef, and the police officer
   # went golfing together.
   @variable(model,Roberta <= RobertaVar <= Roberta, Int)
   @constraint(model, [RobertaVar,Chef,PoliceOfficer] in CS.AllDifferent())

   # From the name of the job
   # (Actor == Steve \/ Actor == Pete)
   # either_eq(model, Actor,Steve, Actor,Pete)
   # is_member_of(model,Actor, [Steve,Pete])
   b4 = @variable(model, binary=true)
   @constraint(model, b4 := {(Actor == Steve || Actor == Pete)})
   @constraint(model, b4 == 1)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                jobs_val = convert.(Integer,JuMP.value.(Jobs; result=sol))
                println("jobs:$jobs_val")

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

@time jobs_puzzle(true,true)


#=
    Cheating model.

    The model is a port of the Cheating model (the first version)
    https://github.com/CamDavidsonPilon/Probabilistic-Programming-and-Bayesian-Methods-for-Hackers/blob/master/Chapter2_MorePyMC/Ch2_MorePyMC_PyMC3.ipynb
    """
    Social data has an additional layer of interest as people are not always honest with responses, which
    adds a further complication into inference. For example, simply asking individuals "Have you ever cheated
    on a test?" will surely contain some rate of dishonesty. What you can say for certain is that the true
    rate is less than your observed rate (assuming individuals lie only about not cheating; I cannot
    imagine one who would admit "Yes" to cheating when in fact they hadn't cheated).

    To present an elegant solution to circumventing this dishonesty problem, and to demonstrate Bayesian
    modeling, we first need to introduce the binomial distribution.

    ...

    We will use the binomial distribution to determine the frequency of students cheating during an exam.
    If we let $N$ be the total number of students who took the exam, and assuming each student is interviewed
    post-exam (answering without consequence), we will receive integer $X$ "Yes I did cheat" answers. We
    then find the posterior distribution of $p$, given $N$, some specified prior on $p$, and observed data $X$.

    This is a completely absurd model. No student, even with a free-pass against punishment, would admit to
    cheating. What we need is a better algorithm to ask students if they had cheated. Ideally the algorithm
    should encourage individuals to be honest while preserving privacy. The following proposed algorithm
    is a solution I greatly admire for its ingenuity and effectiveness:

        In the interview process for each student, the student flips a coin, hidden from the interviewer.
        The student agrees to answer honestly if the coin comes up heads. Otherwise, if the coin comes
        up tails, the student (secretly) flips the coin again, and answers "Yes, I did cheat" if the
        coin flip lands heads, and "No, I did not cheat", if the coin flip lands tails. This way, the
        interviewer does not know if a "Yes" was the result of a guilty plea, or a Heads on a second
        coin toss. Thus privacy is preserved and the researchers receive honest answers.
    """

    This model:

    Summary Statistics
    parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

            p    0.2273    0.1249     0.0040    0.0087   290.4227    1.0058      121.5666
     obs_prop    0.3559    0.0460     0.0015    0.0023   403.9539    0.9991      169.0891

    Quantiles
    parameters      2.5%     25.0%     50.0%     75.0%     97.5%
        Symbol   Float64   Float64   Float64   Float64   Float64

            p    0.0255    0.1364    0.2166    0.3005    0.4916
     obs_prop    0.2700    0.3200    0.3500    0.3800    0.4500

   Cf cheating2.jl for a port of a variant of the PyMC3 model (also from the "Hacker" book).
=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function cheating_model(N,X)
    p ~ Uniform(0,1) # freq cheating

    truths       ~ filldist(Bernoulli(p),N)
    first_flips  ~ filldist(Bernoulli(0.5), N)
    second_flips ~ filldist(Bernoulli(0.5), N)

    val = tzeros(N) # We must defined val to be used as LHS in @.
    @. val = first_flips*truths + (1 - first_flips)*second_flips
    obs_prop ~ Dirac(sum(val) / N)
    X ~ Binomial(N, obs_prop )

    p_prior ~ Uniform(0,1)

end

N = 100
yes_response = 35
model = cheating_model(N, yes_response)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5),  1_000)
# chns = sample(model, PG(5),  MCMCThreads(), 10_000, 4)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns[[:p,:obs_prop,:p_prior]])
# display(plot(chns[[:p,:obs_prop,:p_prior]]))

# Another way:
println("mean(truths): $(mean(group(chns, :truths).value.data,dims=1) |> mean)")


##########################
# FunctionCell Interface #
##########################
"""
	affinityprop(fs [;kwargs])

Constructs an untrained cell that when piped data inside, clusters the data using the
affinity propagation algorithm; `fs` is a function that is used to calculate the
similarity matrices during training as well as when assigning new data to existing clusters.
By default, `fs(x,y) = 1 - pairwise(CosineDist(),x,y)`

# Keyword arguments (same as in `Clustering`)
  * `maxiter` maximum number of iterations (default `200`)
  * `tol` tolerable change in objective at convergence (default `1.0e-6`)
  * `damp` dampening coefficient; values should be in [0.0, 1.0] (default `0.5`); larger values indicate
slower and probably more stable updates
  * `display` Level of information; can be `:none`, `:final`, `:iter` (default `:none`)

!!! note
	The implementation differes from the one in `Clustering.jl` by allowing the user to
	input directly data and not a similarity matrix to the algorithm.

Read the `Clustering.jl` documentation for more information.
"""
affinityprop(fs::Function=(x,y)->one(eltype(x)) .- Distances.pairwise(Distances.CosineDist(),x,y); kwargs...) =
	FunctionCell(affinityprop, (fs,), ModelProperties(), kwtitle("Affinity propagation", kwargs); kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	affinityprop(x, fs::Function=(x,y)->1-Distances.pairwise(CosineDist(),x,y) [;kwargs])

Performs affinity propagation clustering on the dataset `x` using `fs` to calculate the self-similarity matrix;
When piping data to the resulting clustering (e.g. trained cell), the similarities between each input observation
and all exemplars are returned. One can apply `targets(indmax, ...)` to obtain the corresponding exemplar
index (e.g. cluster) for each input observation.
"""
# Training
affinityprop(x::T where T, fs::Function=(x,y)->1-Distances.pairwise(Distances.CosineDist(),x,y); kwargs...) =
	affinityprop(getx!(x), fs; kwargs...)
affinityprop(x::T where T<:AbstractVector, fs::Function=(x,y)->1-Distances.pairwise(Distances.CosineDist(),x,y); kwargs...) =
	affinityprop(mat(x, LearnBase.ObsDim.Constant{2}()), fs; kwargs...)
affinityprop(x::T where T<:AbstractMatrix, fs::Function=(x,y)->1-Distances.pairwise(Distances.CosineDist(),x,y); kwargs...) =
begin
	# Perform clustering
	affpdata = Clustering.affinityprop(fs(x,x); kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x),length(affpdata.exemplars))

	FunctionCell(affinityprop, Model((fs,x[:,sort(affpdata.exemplars)],affpdata), modelprops), kwtitle("Affinity propagation", kwargs));
end

# Execution
affinityprop(x::T where T<:CellData, model::Model{<:Tuple{<:Function, <:AbstractArray, <:Clustering.AffinityPropResult}}) =
	datacell(affinityprop(getx!(x), model), gety(x))
affinityprop(x::T where T<:AbstractVector, model::Model{<:Tuple{<:Function, <:AbstractArray, <:Clustering.AffinityPropResult}}) =
	affinityprop(mat(x, LearnBase.ObsDim.Constant{2}()), model)
affinityprop(x::T where T<:AbstractMatrix, model::Model{<:Tuple{<:Function, <:AbstractArray, <:Clustering.AffinityPropResult}}) =
	model.data[1](model.data[2], x) #fs(exemplars,x)


# using TuringModels
using CSV, DataFrames, Query, Turing, Optim, StatsBase

# define precis from StatisticalRethinking module
using NamedArrays
function precis(df::DataFrame; digits=3, depth=Inf, alpha=0.11)
  m = zeros(length(names(df)), 5)
  for (indx, col) in enumerate(names(df))
    m[indx, 1] = mean(df[:, col])
    m[indx, 2] = std(df[:, col])
    q = quantile(df[:, col], [alpha/2, 0.5, 1-alpha/2])
    m[indx, 3] = q[1]
    m[indx, 4] = q[2]
    m[indx, 5] = q[3]
  end
  NamedArray(
    round.(m, digits=digits),
    (names(df), ["mean", "sd", "5.5%", "50%", "94.5%"]),
    ("Rows", "Cols"))
end

# ### snippet 5.3

df = CSV.read(joinpath(@__DIR__, "data", "Howell1.csv"))

# create categorical sex variable
df = df |> @mutate(sexcat = _.male == 1 ? "male" : "female") |> DataFrame
categorical!(df, :sexcat)

```
get levels code to create index variable:
this could have been done like this:
    df = df |> @mutate(sexcat = _.male == 1 ? 2 : 1) |> DataFrame
for the same result, but this code will do the job if you already
have a categorial variable.
```
df = df |> @mutate(sex = levelcode(_.sexcat)) |> DataFrame

@model model5_8(y, group) = begin
    n_groups = length(unique(group))
    σ ~ Uniform(0, 50)
    α ~ filldist(Normal(178, 20), n_groups)
    μ = α[group]
    @. y ~ Normal(μ, σ)
end

y = df.height
group = df.sex

model = model5_8(y, group)
m5_8 = optimize(model, MAP())

m5_8 = sample(model, NUTS(), MCMCThreads(), 2000, 4)
describe(m5_8)

post = DataFrame(m5_8)

# post |> @rename("α[1]" => "male")

# retrieve category names
dict = zip(["α[1]", "α[2]"], levels(df.sexcat)) |> Dict
rename!(post, dict)

post = post |> @mutate(diff = _.female - _.male) |> DataFrame

precis(post)


using IntervalArithmetic
using Test

setformat(:full)

# NB: Due to "world age" problems, do *not* wrap everything in another @testset
# The `setrounding` commands must be at top level

setrounding(Interval, :slow)
x = Interval(0.5)
@testset "Tight, slow rounding by changing rounding mode" begin
    @test rounding(Interval) == :slow
    @test sin(x) == Interval(0.47942553860420295, 0.479425538604203)
end

setrounding(Interval, :accurate)
@testset "Accurate, fast rounding using prevfloat and nextfloat" begin
    @test rounding(Interval) == :accurate
    @test sin(x) == Interval(0.47942553860420295, 0.47942553860420306)
end

setrounding(Interval, :none)
@testset "No rounding" begin
    @test rounding(Interval) == :none
    @test sin(x) == Interval(0.479425538604203, 0.479425538604203)
    w = 0 ± big(1)
    @test isa(w+w, Interval)
end

setrounding(Interval, :tight)
@testset "Tight, fast rounding using FastRounding.jl" begin
    @test rounding(Interval) == :tight
    @test sin(x) == Interval(0.47942553860420295, 0.479425538604203)
end

setformat(:standard)

# end
# end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMQuad.jl/blob/master/LICENSE

# Tensorial quadrature rules in 1, 2, 3 dimensions

function tensor_product(w::Tuple, p::Tuple, dim::Int)
    @assert length(w) == length(p)

    N = length(w)
    weights = Float64[]
    points = NTuple{dim, Float64}[]
    for i in CartesianIndices(ntuple(i -> 1:N, dim))
        push!(weights, prod(w[k] for k in Tuple(i)))
        push!(points, ntuple(k -> p[i[k]], dim))
    end
    return zip(Tuple(weights), Tuple(points))
end

names = [:GLSEG, :GLQUAD, :GLHEX]
names2 = ["segment", "quadrilateral", "hexahedron"]

for n in 1:length(QUAD_DATA)
    points, weights = QUAD_DATA[n]
    order = 2(n-1)+1
    for dim in (1, 2, 3)
        n_points = length(points)^dim
        quadname = QuoteNode(Symbol(string(names[dim], n_points)))
        z = tensor_product(weights, points, dim)
        @eval begin
            @doc """
                get_quadrature_points(::Type{Val{:$($(quadname))})

            Gauss-Legendre quadrature, $($(n_points)) point rule on $($(names2[dim]))."""
            get_quadrature_points(::Type{Val{$(quadname)}}) = $z
        end
        @eval get_order(::Type{Val{$(quadname)}}) = $order
    end
end


# title: Fitzhugh-Nagumo Bayesian Parameter Estimation Benchmarks
# author: Vaibhav Dixit, Chris Rackauckas

ProjDir = @__DIR__
using DiffEqBayesStan, BenchmarkTools, MCMCChains, StatsPlots
using Tables, AxisKeys, DataFrames, Random

using OrdinaryDiffEq, RecursiveArrayTools, Distributions
using ParameterizedFunctions, StanSample
using Plots
gr(fmt=:png)
Random.seed!(123)

### Defining the problem

# The [FitzHugh-Nagumo model](https://en.wikipedia.org/wiki/FitzHugh%E2%80%93Nagumo_model) is a simplified version of [Hodgkin-Huxley model](https://en.wikipedia.org/wiki/Hodgkin%E2%80%93Huxley_model) and is used to describe an excitable system (e.g. neuron).

fitz = @ode_def FitzhughNagumo begin
  dv = v - v^3/3 -w + l
  dw = τinv*(v +  a - b*w)
end a b τinv l

prob_ode_fitzhughnagumo = ODEProblem(fitz,[1.0,1.0],(0.0,10.0),[0.7,0.8,1/12.5,0.5])
sol = solve(prob_ode_fitzhughnagumo, Tsit5())

# Data is genereated by adding noise to the solution obtained above.

t = collect(range(1,stop=10,length=10))
sig = 0.20
data = convert(Array, VectorOfArray([(sol(t[i]) + sig*randn(2)) for i in 1:length(t)]))

### Plot of the data and the solution.

scatter(t, data[1,:])
scatter!(t, data[2,:])
plot!(sol)

### Priors for the parameters which will be passed for the Bayesian Inference

priors = [
  truncated(Normal(1.0,0.5),0,1.5),truncated(Normal(1.0,0.5),0,1.5),
  truncated(Normal(0.0,0.5),0.0,0.5),truncated(Normal(0.5,0.5),0,1)
]

### Benchmarks

#### Stan.jl backend

tmpdir = joinpath(ProjDir, "tmp")

diffeq_string = "
  vector sho(
    real t,vector internal_var___u,
    real internal_var___p_1, real internal_var___p_2,
    real internal_var___p_3, real internal_var___p_4) {

  vector[2] internal_var___du;

  internal_var___du[1] = internal_var___p_4 - 0.3333 * internal_var___u[1] ^ 3 + -1 * internal_var___u[2] + internal_var___u[1];
  internal_var___du[2] = internal_var___p_3 * (internal_var___p_1 + -1 * internal_var___p_2 * internal_var___u[2] + internal_var___u[1]);

  return internal_var___du;
}
"

bayesian_result_stan =
  stan_inference(prob_ode_fitzhughnagumo,t,data,priors;
    num_threads=4, num_cpp_chains=4, num_chains=2, num_samples=1250, output_format = :dataframe,
    diffeq_string, tmpdir);

describe(bayesian_result_stan.chains) |> display
println()

st =  read_samples(bayesian_result_stan.model, :table)
describe(DataFrame(st)) |> display
println()

ka = read_samples(bayesian_result_stan.model, :keyedarray)
ka |> display
println()

chns = read_samples(bayesian_result_stan.model, :mcmcchains)
plot(chns)


"""
Linear-Stability-Calculators: Shallow Water Model
=================================================
"""

include("src/shallow_water_setup.jl")

             ks = Dict()
ks[Cartesian()] = collect(2e-5:2e-5:2.5e-3)
ks[Spherical()] = collect(2e-5:2e-5:2.5e-3) * 6.3781e6 * cos(π/4)

        Uj = 1.0                 # Jet Parameters
   Ljscale = 20
    Nmodes = 2
         σ = Dict()
         ω = Dict()
    σmodes = Dict()
      phys = Dict()
background = Dict()

geometrys = (Cartesian(), Spherical())

for geometry in geometrys

     @info string("Setting Geometry = ", geometry)

                    grid = Grid(geometry, N=250)
          phys[geometry] = Physics(geometry; grid=grid)
    background[geometry] = Bickley_Jet(geometry; phys=phys[geometry], Uj, Ljscale)

     plot_basic_state(
          grid,
          background,
          geometry,
          string("basic_state_",typeof(geometry),".png")
          )

     for (index, k) in enumerate(ks[geometry])
          A = build_matrix(geometry; k, background, phys)
          σ[(k,geometry)], ω[(k,geometry)], σmodes[(k,geometry)] = compute_spectrum(A, k, phys, Nmodes);
          @info string("     k = ", k, " and max growth = ", maximum(σ[(k,geometry)]), "\n")
     end

     plot_growth_rates(
          ks,
          σ,
          phys,
          geometry,
          Nmodes,
          string("growth_rates_",typeof(geometry),"_Bickley_Jet.png")
          )

     mode_number = 1;
           σ_max = maximum( σ[(k,geometry)][mode_number] for k in ks[geometry]);
         k_index = sortperm([σ[(k,geometry)][mode_number] for k in ks[geometry]],rev=true)[1];
               k = ks[geometry][k_index];

     print("σ_max = ",σ_max, " with k = ", k, " at k_index =", k_index, "\n")
     plot_1D_streamfunction(
          k_index,
          k,
          phys,
          σmodes,
          geometry,
          mode_number,
          string("modes_1D_streamfunction_",typeof(geometry),".png")
          )

     plot_2D_streamfunction(
          k_index,
          k,
          phys,
          σmodes,
          geometry,
          mode_number,
          string("modes_2D_streamfunction_",typeof(geometry),".png")
          )

     plot_2D_vorticity(
          phys,
          k_index,
          k,
          σmodes,
          geometry,
          mode_number,
          string("modes_2D_vorticity_",typeof(geometry),".png")
          )

end


using Sundials, Test, LinearAlgebra

# create mass matrix problems
function make_mm_probs(mm_A, ::Type{Val{iip}}) where iip
  # iip
  mm_b = vec(sum(mm_A; dims=2))
  function mm_f(du,u,p,t)
    mul!(du,mm_A,u)
    du .+= t * mm_b
    nothing
  end
  mm_g(du,u,p,t) = (@. du = u + t; nothing)

  # oop
  mm_f(u,p,t) = mm_A * (u .+ t)
  mm_g(u,p,t) = u .+ t

  mm_analytic(u0, p, t) = @. 2 * u0 * exp(t) - t - 1

  u0 = ones(3)
  tspan = (0.0, 1.0)

  prob = ODEProblem(ODEFunction{iip,true}(mm_f, analytic=mm_analytic, mass_matrix=mm_A), u0, tspan)
  prob2 = ODEProblem(ODEFunction{iip,true}(mm_g, analytic=mm_analytic), u0, tspan)

  prob, prob2
end

mm_A = Float64[-2 1 4; 4 -2 1; 2 1 3]
prob, prob2 = make_mm_probs(mm_A, Val{true})

sol = solve(prob,  ARKODE(),abstol=1e-8,reltol=1e-8)
sol2 = solve(prob2,ARKODE(),abstol=1e-8,reltol=1e-8)

@test norm(sol .- sol2) ≈ 0 atol=1e-7


module PrecompileFuncs

using TuringDash
using TuringDash.Turing
using TuringDash.Random
using TuringDash.CSV, TuringDash.DataFrames

function precompile_turing()
    chn = Base.Channel{Vector{Float64}}(Inf)
    df1 = nothing
    global df1
    df1 = CSV.read(download("https://raw.githubusercontent.com/efmanu/TuringDashApp.jl/master/TuringDash/datasets/data1.csv"), DataFrame)
    model_str1 = """
    @model regression_model(x, y) = begin
        a ~ Normal()
        b ~ Normal()
        for i in 1:10
            y[i] ~ Normal(a + b * x[i], 1.0)
        end
    end
    """
    model_str2 = "regression_model(x, y)"
    for colname in Symbol.(names(df1))
        @eval $colname = df1.$colname
    end

    eval(Meta.parse(model_str1))
    model = eval(Meta.parse(model_str2))

    alg = MH(
        [0.25 0.05;
         0.05 0.50]
    )
    rng = Random.GLOBAL_RNG

    nperiteration = 50
    nsamples = 20

    r = sample(rng, model, alg, nperiteration; chain_type=MCMCChains.Chains, save_state=true, progress=false)
    put!(chn, Array(r)[end,1:2])
    for i in (nperiteration + 1):nperiteration:nsamples
        r = Turing.Inference.resume(r, nperiteration, save_state=true, progress=false)
        put!(chn, Array(r)[end,1:2])
    end
    return "success"
end

end


using Archibald
using Test
using Distributions
using DataFrames
using LinearAlgebra

@testset "Over Sampling MCMC" begin

    function generate_data(;n = 10,
                            T = 100,
                            β = 0.25,
                            family = Bernoulli,
                            τ = 1,
                            link = x -> cdf(Normal(), x))
        i_index = repeat(1:n, inner = T)
        t_index = repeat(1:T, n)
        μ_x = rand(Normal(0, 1), n)
        α_i = rand(Normal(0, τ), n)
        x = vcat(rand.(Normal.(μ_x, 1), T )...)
        μ_vec = α_i[i_index] .+ x * β

        y = rand.(family.(link.(μ_vec)))

        return y, x, i_index, t_index, α_i
    end

    function linear_target_dist_oversample(y,
                                           x,
                                           i_index,
                                           single_parameters,
                                           over_parameters)
        β = single_parameters[1]
        τ = exp(single_parameters[2])
        α_matrix = over_parameters
        S = size(α_matrix, 2)
        ll = Vector{Float64}(undef, S)
        i_index = Int.(i_index)
        for s in 1:S
            resid_vec = y .- α_matrix[i_index, s] .- x*β
            ll[s] = sum(log.(pdf.(Normal.(0, 1), resid_vec)))
        end
        τ_prior = TruncatedNormal(0, 3, 0, Inf)
        τ_prior_dens = pdf(τ_prior, τ)
        param_probability = sum(ll) + log(τ_prior_dens)
        return param_probability
    end

    sim_n = 5
    sim_t = 100
    sim_τ = 0.5
    linear_data = generate_data(n = sim_n, T = sim_t, τ=sim_τ, link = x -> Normal(x, 1), family = identity)

    oversample_parameters = randn(sim_n, 3)
    chain_over = Archibald.MarkovChainOverSample(
        [0.1, 0.2],
        3,
        oversample_parameters,
        [linear_data[1], linear_data[2], linear_data[3]],
        (x_single, x_over, data) -> linear_target_dist_oversample(data[1],
                                                   data[2],
                                                   data[3],
                                                   x_single, x_over)
    )
    MetropolisHastings!(
        chain_over,
        (x) -> MvNormal(x, 0.1),
        (x) -> MvNormal(sim_n, x[2]),
        logs = true
        )

    mvn_draws = MetropolisHastings(
        1_000,
        chain_over,
        x -> MvNormal(x, 0.1), # jeez what an elegant design right? right?
        x -> MvNormal(sim_n, x[2]),
        logs = true,
        burn = 1_000)
    # now "for real" with lower proposal variance
    mvn_df = tidybayes(mvn_draws...)
    mvn_df
    @test isa(mvn_df, DataFrame)
    @test isa(mvn_draws[1][1], Vector)
    @test isa(mvn_draws[1][1][1], Float64)
    @test !isnan(mvn_draws[1][1][2])

end

@testset "Basic Metropolis Hastings" begin
    function generate_data(; N = 100)
        sigma_b = log(rand(InverseGamma(2, 2)))
        sigma_e = log(rand(InverseGamma(1, 1)))
        β = [1, 2, 3, 4]

        X = randn(2, 4, N)
        W = randn(2, 1, N)
        b = randn(1, N).*sigma_b
        e = randn(2, N).*sigma_e

        Y = Matrix{Float64}(undef, 2, N)
        for i in 1:N
            Y[:, i] = X[:, :, i] * β + W[:, :, i]*b[:, i] + e[:, i]
        end
        return Y, X, W, b, β, e, sigma_b, sigma_e
    end

    """
    Create target distribution for regression problem
    """
    function mvn_lik(y, X, W, parameters)
        N = size(X, 3)
        β = parameters[1:4]
        # constrain sigmas to be positive using exp(.)
        σ_b = exp(parameters[5])
        σ_e = exp(parameters[6])

        ll_i = Vector{Float64}(undef, N)
        for i in 1:N
            # Account for variance from Wb_i
            Σ_sqrt = σ_b^2 * W[:, :, i]*W[:, :, i]' + σ_e^2 .* Matrix(I,2,2)
            Σ = (Σ_sqrt + Σ_sqrt')./2 # make it Hermitian
            resid = y[:, i] - X[:, :, i]*β
            ll_i[i] = log(pdf(MvNormal([0, 0], Σ), resid))
        end

        # sum log lik contribution
        ll = sum(ll_i)
        # priors
        σ_b_prior = pdf(InverseGamma(2, 2), σ_b)
        σ_e_prior = pdf(InverseGamma(1, 1), σ_e)

        P = ll + log(sum(pdf.(Normal(0, 10), β))) + log(σ_b_prior) + log(σ_e_prior)
        return P
    end

    test_Y, test_X, test_W, test_b, test_β, test_e, test_σ_b, test_σ_e = generate_data(N = 1_00)

    # wrapper function to match arguments with previous P()
    mvn_target_pdf(x, data) =  mvn_lik(data[1], data[2], data[3], x)

    # Save true values for comparison later
    true_values = vcat(test_β, exp(test_σ_b), exp(test_σ_e), test_b[:])[:]
    true_params = DataFrame(
        [
            :true_value => true_values,
            :term => "parameter_" .* string.(1:size(true_values, 1)),
            :term_name => [
                "β_1",
                "β_2",
                "β_3",
                "β_4",
                "σ_b",
                "σ_e",
                "b_" .* string.(1:(size(true_values, 1) - 6))...
            ]
        ]
    )

    # Initialise chain
    mvn_markov_state = MarkovChain(
        fill(1.0, size(true_params, 1)),
        [test_Y, test_X, test_W],
        mvn_target_pdf
    )
    # burn in ad hoc for 20k
    mvn_draws = MetropolisHastings(
        1_000,
        mvn_markov_state,
        x -> MvNormal(x, 0.1), # jeez what an elegant design right? right?
        logs = true,
        burn = 1_000)
    # now "for real" with lower proposal variance
    mvn_df = tidybayes(mvn_draws...)
    @test isa(mvn_df, DataFrame)
    @test size(mvn_df, 1) == 106000
    @test isa(mvn_draws[1][1], Vector)
    @test isa(mvn_draws[1][1][1], Float64)
    @test !isnan(mvn_draws[1][1][2])
end


#---------------------------------------------------------------------------------------------------
# iTEBD
#---------------------------------------------------------------------------------------------------
# gate
function gate!(
    mat::AbstractMatrix,
    ΓA::AbstractArray{<:Number, 3},
    ΓB::AbstractArray{<:Number, 3},
    λl::AbstractVector{<:Real},
    d1::Integer,
    d2::Integer,
    bound::Integer,
    cutoff::Real
)
    α, β = size(ΓA, 1), size(ΓB, 3)
    tlmul!(λl, ΓA)
    Γ = tcmul(mat, tgroup(ΓA, ΓB))
    ΓA2, λc, ΓB2 = tsvd(Γ, α, d1, d2, β, bound=bound, cutoff=cutoff, renormalize=true)
    trmul!(ΓA2, λc)
    tlmul!(1.0 ./ λl, ΓA2)
    ΓA2, λc, ΓB2
end
#---------------------------------------------------------------------------------------------------
# 2-site iTEBD
export itebd2!
function itebd2!(
    mat::AbstractMatrix,
    ΓA::AbstractArray{<:Number, 3},
    ΓB::AbstractArray{<:Number, 3},
    λl::AbstractVector{<:Real};
    bound::Integer=100,
    cutoff::Real=1e-7
)
    d1, d2 = size(ΓA, 2), size(ΓB, 2)
    ΓA2, λ1, ΓB2 = gate!(mat,  ΓA,  ΓB, λl, d1, d2, bound, cutoff)
    ΓB3, λ2, ΓA3 = gate!(mat, ΓB2, ΓA2, λ1, d2, d1, bound, cutoff)
    ΓA3, ΓB3, λ2
end

#---------------------------------------------------------------------------------------------------
# 3-site iTEBD
export itebd3!
function itebd3!(
    mat::AbstractMatrix,
    ΓAB1::AbstractArray{<:Number, 3},
    ΓC1::AbstractArray{<:Number, 3},
    λl::AbstractVector{<:Real};
    bound::Integer=100,
    cutoff::Real=1e-7
)
    d1, d2, d3 = size(ΓA, 2), size(ΓB, 2),  size(ΓC, 2)
    ΓA2, λ1, ΓBC2 = gate!(mat, ΓAB1, ΓC1, λl, d1, d2 * d3, bound, cutoff)
    ΓB3, λ2, ΓCA3 = gate!(mat, ΓBC2, ΓA2, λ1, d2, d3 * d1, bound, cutoff)
    ΓC4, λ3, ΓAB4 = gate!(mat, ΓBC2, ΓA2, λ2, d3, d1 * d2, bound, cutoff)
    ΓAB4, ΓC4, λ3
end


# Bi-objective set partitionning problem (2-SPA)

# ---- Packages to use
using vOptGeneric, JuMP, GLPK

# ---- Parser reading an instance of 2-SPA (format of instances compliant with vOptLib)
function loadInstance2SPA(fname::String)

    f = open(fname)
    nbctr, nbvar = parse.(Int, split(readline(f))) # number of constraints , number of variables
    A = zeros(Int, nbctr, nbvar)                   # matrices of constraints
    c1 = zeros(Int, nbvar)                         # 1st vector of costs
    c2 = zeros(Int, nbvar)                         # 2nd vector of costs
    nb = zeros(Int, nbvar)
    for i in 1:nbvar
        flag = 1
        for valeur in split(readline(f))
            if flag == 1
                c1[i] = parse(Int, valeur)
                flag +=1
            elseif flag == 2
                c2[i] = parse(Int, valeur)
                flag +=1
            elseif flag == 3
                nb[i] = parse(Int, valeur)
                flag +=1
            else
                j = parse(Int, valeur)
                A[j,i] = 1
            end
        end
    end
    close(f)
    return c1, c2, A
end

# ---- Compute the set of non-dominated points Y_N of a 2SPA with vOptGeneric
function computeYNfor2SPA(  nbvar::Int,
                            nbctr::Int,
                            A::Array{Int,2},
                            c1::Array{Int,1},
                            c2::Array{Int,1}
                         )

    # ---- setting the model
    model = vModel( GLPK.Optimizer )
    JuMP.set_silent(model)

    @variable(model, x[1:nbvar], Bin)
    @constraint(model, [i=1:nbctr],(sum((x[j]*A[i,j]) for j in 1:nbvar)) == 1)
    @addobjective(model, Min, sum(c1[i]*x[i] for i in 1:nbvar))
    @addobjective(model, Min, sum(c2[i]*x[i] for i in 1:nbvar))

    # ---- Invoking the solver (epsilon constraint method)
    vSolve(model, method=:epsilon, step = 0.5)

    # ---- Querying the results
    start = time()
    Y_N = getY_N(model)
    elapsedTime = time() - start
    return Y_N, elapsedTime
end

# ---- Save an instance of 2-SPA (format of instances compliant with vOptLib)
function saveInstance2SPA(fname::String, Y_N, elapsedTime::Float64)
    way = "Y_N_"*fname
    open(way, "w") do f
        write(f, "$elapsedTime \n")
        write(f, "$(length(Y_N)) \n")
        for i in 1:length(Y_N)
            write(f, "$(Y_N[i][1]) $(Y_N[i][2]) \n")
        end
    end
end

# ---- Entry point
function main(fname::String)

    print("Compute Y_N with vOPtGeneric for a set partitionning problem with 2 objectives (2-SPA) \n\n")

    # load a numerical instance of 2SPA ----------------------------------------
    c1, c2, A = loadInstance2SPA(fname)
    nbctr = size(A,1)
    nbvar = size(A,2)
    nbobj = 2

    # compute the set of non-dominated points Y_N ------------------------------
    Y_N, elapsedTime = computeYNfor2SPA(nbvar, nbctr, A, c1, c2)

    # save the result on a file
    saveInstance2SPA(fname, Y_N, elapsedTime)
end

# ---- Example

# main("biodidactic5.txt")


using GLMdiagnostics
using Test
using Distributions
using Random
using DataFrames
using StatsBase
using GLM
using Gadfly
using Statistics
using Compose

###########################################
# - Create simulated data to test functions
###########################################

# Set reproducible seed
Random.seed!(123)

# Simulate data from a Normal distribution
d = Normal(1, 2)
y = rand(d, 100)

# Assign grouping variables (5 groups of 20)
x = repeat([1, 2, 3, 4, 5],
            inner = 20,
            outer = 1)

df = hcat(y, x)
df = convert(DataFrame, df)
colnames = ["y","x"]
rename!(df, Symbol.(colnames))

##########################################
# - Run linear model (one-way ANOVA)
##########################################

# Run model
modAov = lm(@formula(y ~ 1 + x), df)

# Make qqplot for ANOVA model
anova_qqplot = plot_qq(modAov)


function msa(network::AbstractNetwork,
              trips::AbstractMatrix{T},
              costfn::CostFunction;
              basedon=:link,
              errtol=1e-4) where {T<:Real, U<:Integer}
    # initialize
    flows = allornothing(network, trips, costfn; basedon=basedon)

    # start iteration
    err = 1.
    iterno = 2 # first iteration @ initialization
    while err > errtol
        ## find target solution
        linkcosts = costfn(flows).costs
        shortflows = allornothing(network, trips, linkcosts; basedon=basedon)

        ## find stepsize
        μ = 1/iterno

        ## calculate new link travel times
        flows += μ * (shortflows - flows)

        ## calculate error
        err = sum(flows .* linkcosts)/sum(shortflows .* linkcosts) - 1.
        iterno += 1
    end
    return (flows, err)
end


# Import seeding generator for reproducibility
using StableRNGs

"""
    Sigmoid activation function
"""
function sigmoid(Z)
    A = 1 ./ (1 .+ exp.(.-Z))
    return (A = A, Z = Z)
end

"""
    ReLU activation function
"""
function relu(Z)
    A = max.(0, Z)
    return (A = A, Z = Z)
end

"""
    Funtion to initialise the parameters or weights of the desired network.
"""
function initialise_model_weights(layer_dims, seed)
    params = Dict()

    # Build a dictionary of initialised weights and bias units
    for l=2:length(layer_dims)
        params[string("W_", (l-1))] = rand(StableRNG(seed), layer_dims[l], layer_dims[l-1]) * sqrt(2 / layer_dims[l-1])
        params[string("b_", (l-1))] = zeros(layer_dims[l], 1)
    end

    return params
end

"""
    Make a linear forward calculation
"""
function linear_forward(A, W, b)
    # Make a linear forward and return inputs as cache
    Z = (W * A) .+ b
    cache = (A, W, b)

    @assert size(Z) == (size(W, 1), size(A, 2))

    return (Z = Z, cache = cache)
end

"""
    Make a forward activation from a linear forward.
"""
function linear_forward_activation(A_prev, W, b, activation_function="relu")
    @assert activation_function ∈ ("sigmoid", "relu")
    Z, linear_cache = linear_forward(A_prev, W, b)

    if activation_function == "sigmoid"
        A, activation_cache = sigmoid(Z)
    end

    if activation_function == "relu"
        A, activation_cache = relu(Z)
    end

    cache = (linear_step_cache=linear_cache, activation_step_cache=activation_cache)

    @assert size(A) == (size(W, 1), size(A_prev, 2))

    return A, cache
end

"""
    Forward the design matrix through the network layers using the parameters.
"""
function forward_propagate_model_weights(DMatrix, parameters)
    master_cache = []
    A = DMatrix
    L = Int(length(parameters) / 2)

    # Forward propagate until the last (output) layer
    for l = 1 : (L-1)
        A_prev = A
        A, cache = linear_forward_activation(A_prev,
                                             parameters[string("W_", (l))],
                                             parameters[string("b_", (l))],
                                             "relu")
        push!(master_cache , cache)
    end

    # Make predictions in the output layer
    Ŷ, cache = linear_forward_activation(A,
                                         parameters[string("W_", (L))],
                                         parameters[string("b_", (L))],
                                         "sigmoid")
    push!(master_cache, cache)

    return Ŷ, master_cache
end

"""
    Computes the log loss (binary cross entropy) of the current predictions.
"""
function calculate_cost(Ŷ, Y)
    m = size(Y, 2)
    epsilon = eps(1.0)

    # Deal with log(0) scenarios
    Ŷ_new = [max(i, epsilon) for i in Ŷ]
    Ŷ_new = [min(i, 1-epsilon) for i in Ŷ_new]

    cost = -sum(Y .* log.(Ŷ_new) + (1 .- Y) .* log.(1 .- Ŷ_new)) / m
    return cost
end

"""
    Derivative of the Sigmoid function.
"""
function sigmoid_backwards(∂A, activated_cache)
    s = sigmoid(activated_cache).A
    ∂Z = ∂A .* s .* (1 .- s)

    @assert (size(∂Z) == size(activated_cache))
    return ∂Z
end

"""
    Derivative of the ReLU function.
"""
function relu_backwards(∂A, activated_cache)
    return ∂A .* (activated_cache .> 0)
end

"""
    Partial derivatives of the components of linear forward function
    using the linear output (∂Z) and caches of these components (cache).
"""
function linear_backward(∂Z, cache)
    # Unpack cache
    A_prev , W , b = cache
    m = size(A_prev, 2)

    # Partial derivates of each of the components
    ∂W = ∂Z * (A_prev') / m
    ∂b = sum(∂Z, dims = 2) / m
    ∂A_prev = (W') * ∂Z

    @assert (size(∂A_prev) == size(A_prev))
    @assert (size(∂W) == size(W))
    @assert (size(∂b) == size(b))

    return ∂W , ∂b , ∂A_prev
end

"""
    Unpack the linear activated caches (cache) and compute their derivatives
    from the applied activation function.
"""
function linear_activation_backward(∂A, cache, activation_function="relu")
    @assert activation_function ∈ ("sigmoid", "relu")

    linear_cache , cache_activation = cache

    if (activation_function == "relu")

        ∂Z = relu_backwards(∂A , cache_activation)
        ∂W , ∂b , ∂A_prev = linear_backward(∂Z , linear_cache)

    elseif (activation_function == "sigmoid")

        ∂Z = sigmoid_backwards(∂A , cache_activation)
        ∂W , ∂b , ∂A_prev = linear_backward(∂Z , linear_cache)

    end

    return ∂W , ∂b , ∂A_prev
end

"""
    Compute the gradients (∇) of the parameters (master_cache) of the constructed model
    with respect to the cost of predictions (Ŷ) in comparison with actual output (Y).
"""
function back_propagate_model_weights(Ŷ, Y, master_cache)
    # Initiate the dictionary to store the gradients for all the components in each layer
    ∇ = Dict()

    L = length(master_cache)
    Y = reshape(Y , size(Ŷ))

    # Partial derivative of the output layer
    ∂Ŷ = (-(Y ./ Ŷ) .+ ((1 .- Y) ./ ( 1 .- Ŷ)))
    current_cache = master_cache[L]

    # Backpropagate on the layer preceeding the output layer
    ∇[string("∂W_", (L))], ∇[string("∂b_", (L))], ∇[string("∂A_", (L-1))] = linear_activation_backward(∂Ŷ,
                                                                                current_cache,
                                                                                "sigmoid")
    # Go backwards in the layers and compute the partial derivates of each component.
    for l=reverse(0:L-2)
        current_cache = master_cache[l+1]
        ∇[string("∂W_", (l+1))], ∇[string("∂b_", (l+1))], ∇[string("∂A_", (l))] = linear_activation_backward(∇[string("∂A_", (l+1))],
                                                                                current_cache,
                                                                                "relu")
    end

    # Return the gradients of the network
    return ∇
end

"""
    Update the paramaters of the model using the gradients (∇)
    and the learning rate (η).
"""
function update_model_weights(parameters, ∇, η)

    L = Int(length(parameters) / 2)

    # Update the parameters (weights and biases) for all the layers
    for l = 0: (L-1)
        parameters[string("W_", (l + 1))] -= η .* ∇[string("∂W_", (l + 1))]
        parameters[string("b_", (l + 1))] -= η .* ∇[string("∂b_", (l + 1))]
    end

    return parameters
end

"""
    Check the accuracy between predicted values (Ŷ) and the true values(Y).
"""
function assess_accuracy(Ŷ , Y)
    @assert size(Ŷ) == size(Y)
    return sum((Ŷ .> 0.5) .== Y) / length(Y)
end

"""
    Train the network using the desired architecture that best possible
    matches the training inputs (DMatrix) and their corresponding ouptuts(Y)
    over some number of iterations (epochs) and a learning rate (η).
"""
function train_network(layer_dims , DMatrix, Y;  η=0.001, epochs=1000, seed=2020, verbose=true)
    # Initiate an empty container for cost, iterations, and accuracy at each iteration
    costs = []
    iters = []
    accuracy = []

    # Initialise random weights for the network
    params = initialise_model_weights(layer_dims, seed)

    # Train the network
    for i = 1:epochs

        Ŷ , caches  = forward_propagate_model_weights(DMatrix, params)
        cost = calculate_cost(Ŷ, Y)
        acc = assess_accuracy(Ŷ, Y)
        ∇  = back_propagate_model_weights(Ŷ, Y, caches)
        params = update_model_weights(params, ∇, η)

        if verbose
            println("Iteration -> $i, Cost -> $cost, Accuracy -> $acc")
        end

        # Update containers for cost, iterations, and accuracy at the current iteration (epoch)
        push!(iters , i)
        push!(costs , cost)
        push!(accuracy , acc)
    end
        return (cost = costs, iterations = iters, accuracy = accuracy, parameters = params)
end


# This script implements a partial differential equation describing changes in
# concentration with respect to longitudinal distance along a stream channel
# due to advection. I try using both the symbolic PDE method (which works)
# and the finite-difference operators method (which doesn't work yet)

##### ----- Define Parameters, Initial Conditions, Boundary Conditions ----- #####
v = 1.0 # Velocity
p = [v] # parameter vector
tspan = (0.0, 10.0) # simulation timespan
xₙ = 10 # physical length of simulation
Δx = 0.01 # grid spacing along x dimension (longitudinal)
N = convert(Int64, xₙ/Δx - 1) # number of interior points
x = range(0, step = Δx, length = (N+2)) # x-points of solution, including boundaries
h(p,t) = zeros(N) # history function describing solute concentrations before the start of the simulation
C₀ = zeros(N) # initial condition
for i in 1:N
  u₀[i] = 1.0
end

##### ----- Automated Finite-Difference for Symbolic PDES Method ----- #####
using ModelingToolkit
@parameters t x v
@variables(u(..))
@derivatives Dt'~t
@derivatives Dx'~x

#1D PDE and boundary conditions
eq = Dt(u(t,x)) ~ -v*Dx(u(t,x))
bcs = [u(0,x) ~ 0,
        u(t,0) ~ 1.0,
        Dx(u(t,10.0)) ~ 0.0]

# Space and time domains
domains = [t ∈ IntervalDomain(0.0,convert(Float64, xₙ)),
           x ∈ IntervalDomain(tspan[1],tspan[2])]

# PDE system
pdesys = PDESystem(eq,bcs,domains,[t,x],[u])

# Method of lines discretization
discretization = MOLFiniteDifference(Δx,3) #Third order finite difference approximation with nodes spaced by 0.01

# Convert the PDE problem into an ODE problem
prob = discretize(pdesys,discretization)

# Solve ODE problem
sol = solve(prob,Tsit5(),saveat=1.0)

# Plot results and compare with exact solution
t = sol.t
x = prob.space[2]

plt = plot()
for i in 1:length(t)
    plot!(x,Array(prob.extrapolation[1](t[i])*sol.u[i]),label="t=$(t[i])")
end
display(plt)

##### ----- Finite Difference Operator Method ----- #####
ΔUpwind = DiffEqOperators.UpwindDifference(1,4,Δx,xₙ) # first derivative operator, fourth-order approximation
Dirichletbc = DiffEqOperators.DirichletBC(1.0,0.0) # boundary condition operator

@advect(ΔUpwind,Dirichletbc)
prob = ODEProblem(advect!, u₀, tspan, p)
alg = Rosenbrock23(autodiff=false)
maxdt = Δx/v
sol = solve(prob, alg, saveat = 1.0, dtmax = 0.8*maxdt)
t = sol.t
u = sol.u

plt = plot()
for i in 1:length(t)
    plot!(x, bc*u[[i]][1] ,label="t=$(t[i])")
end
display(plt)


using MultivariateStats
using LinearAlgebra
using Test
import Statistics: mean, cov
import Random

@testset "CCA" begin

    Random.seed!(34568)

    dx = 5
    dy = 6
    p = 3

    # CCA with zero means

    X = rand(dx, 100)
    Y = rand(dy, 100)

    Px = qr(randn(dx, dx)).Q[:, 1:p]
    Py = qr(randn(dy, dy)).Q[:, 1:p]

    M = CCA(Float64[], Float64[], Px, Py, [0.8, 0.6, 0.4])

    @test xindim(M) == dx
    @test yindim(M) == dy
    @test xmean(M) == zeros(dx)
    @test ymean(M) == zeros(dy)
    @test xprojection(M) == Px
    @test yprojection(M) == Py
    @test correlations(M) == [0.8, 0.6, 0.4]

    @test xtransform(M, X) ≈ Px'X
    @test ytransform(M, Y) ≈ Py'Y

    ## CCA with nonzero means

    ux = randn(dx)
    uy = randn(dy)

    M = CCA(ux, uy, Px, Py, [0.8, 0.6, 0.4])

    @test xindim(M) == dx
    @test yindim(M) == dy
    @test xmean(M) == ux
    @test ymean(M) == uy
    @test xprojection(M) == Px
    @test yprojection(M) == Py
    @test correlations(M) == [0.8, 0.6, 0.4]

    @test xtransform(M, X) ≈ Px' * (X .- ux)
    @test ytransform(M, Y) ≈ Py' * (Y .- uy)

    ## prepare data

    n = 1000
    dg = 10
    G = randn(dg, n)

    X = randn(dx, dg) * G + 0.2 * randn(dx, n)
    Y = randn(dy, dg) * G + 0.2 * randn(dy, n)
    xm = vec(mean(X, dims=2))
    ym = vec(mean(Y, dims=2))
    Zx = X .- xm
    Zy = Y .- ym

    Cxx = cov(X, dims=2)
    Cyy = cov(Y, dims=2)
    Cxy = cov(X, Y, dims=2)
    Cyx = Cxy'

    ## ccacov

    # X ~ Y
    M = fit(CCA, X, Y; method=:cov, outdim=p)
    Px = xprojection(M)
    Py = yprojection(M)
    rho = correlations(M)
    @test xindim(M) == dx
    @test yindim(M) == dy
    @test outdim(M) == p
    @test xmean(M) == xm
    @test ymean(M) == ym
    @test issorted(rho; rev=true)

    @test Px' * Cxx * Px ≈ Matrix(I, p, p)
    @test Py' * Cyy * Py ≈ Matrix(I, p, p)
    @test Cxy * (Cyy \ Cyx) * Px ≈ Cxx * Px * Diagonal(rho.^2)
    @test Cyx * (Cxx \ Cxy) * Py ≈ Cyy * Py * Diagonal(rho.^2)
    @test Px ≈ MultivariateStats.qnormalize!(Cxx \ (Cxy * Py), Cxx)
    @test Py ≈ MultivariateStats.qnormalize!(Cyy \ (Cyx * Px), Cyy)

    # Y ~ X
    M = fit(CCA, Y, X; method=:cov, outdim=p)
    Py = xprojection(M)
    Px = yprojection(M)
    rho = correlations(M)
    @test xindim(M) == dy
    @test yindim(M) == dx
    @test outdim(M) == p
    @test xmean(M) == ym
    @test ymean(M) == xm
    @test issorted(rho; rev=true)

    @test Px' * Cxx * Px ≈ Matrix(I, p, p)
    @test Py' * Cyy * Py ≈ Matrix(I, p, p)
    @test Cxy * (Cyy \ Cyx) * Px ≈ Cxx * Px * Diagonal(rho.^2)
    @test Cyx * (Cxx \ Cxy) * Py ≈ Cyy * Py * Diagonal(rho.^2)
    @test Px ≈ MultivariateStats.qnormalize!(Cxx \ (Cxy * Py), Cxx)
    @test Py ≈ MultivariateStats.qnormalize!(Cyy \ (Cyx * Px), Cyy)

    ## ccasvd

    # n > d
    M = fit(CCA, X, Y; method=:svd, outdim=p)
    Px = xprojection(M)
    Py = yprojection(M)
    rho = correlations(M)
    @test xindim(M) == dx
    @test yindim(M) == dy
    @test outdim(M) == p
    @test xmean(M) == xm
    @test ymean(M) == ym
    @test issorted(rho; rev=true)

    @test Px' * Cxx * Px ≈ Matrix(I, p, p)
    @test Py' * Cyy * Py ≈ Matrix(I, p, p)
    @test Cxy * (Cyy \ Cyx) * Px ≈ Cxx * Px * Diagonal(rho.^2)
    @test Cyx * (Cxx \ Cxy) * Py ≈ Cyy * Py * Diagonal(rho.^2)
    @test Px ≈ MultivariateStats.qnormalize!(Cxx \ (Cxy * Py), Cxx)
    @test Py ≈ MultivariateStats.qnormalize!(Cyy \ (Cyx * Px), Cyy)

    # different input types
    XX = convert.(Float32, X)
    YY = convert.(Float32, Y)

    MM = fit(CCA, view(XX, :, 1:400), view(YY, :, 1:400); method=:svd, outdim=p)

    # test that mixing types doesn't error
    xtransform(M, XX)
    ytransform(M, YY)
    xtransform(MM, XX)
    ytransform(MM, YY)

    # type stability
    for func in (xmean, ymean, xprojection, yprojection, correlations)
        @test eltype(func(M)) == Float64
        @test eltype(func(MM)) == Float32
    end
end


PJane = 25000
PJoe = 40000
P = PJane + PJoe
i = 0.1
m = 1
N = 6
iJane = 0.09

F = P * (1 + i / m)^(N * m)

iJoe = ((F - PJane * (1 + i)^N) / PJoe)^(1/N) - 1


Random.seed!(0)

# Linear Dynamic model
# = x = [p; v]
dt = 0.01
A = [1.0 dt; 0.0 1.0]
B = zeros(2, 2)
W = 0.01*Matrix{Float64}(I, 2, 2)

dyn = LinearDynamicModel(A, B, W)

# Linear Observation Model
# observe postition
H = [1 0]
obs = LinearObservationModel(H)

# define linear kalman filter
kf = KalmanFilter(dyn, obs)

# simulation
R_gt = 0.5*Matrix{Float64}(I, 1, 1) # ground truth noise covariance
time_step = 0.1:dt:1000
s0 = State([0.0, 1.0], Matrix{Float64}(I, 2, 2)) # initial state belief


"""
Calculates the Gaunt coefficients, defined by:

    a(m,n,μ,ν,q) = (2(n+ν-2q)+1)/2 (n+ν-2q-m-μ)!/(n+ν-2q+m+μ)! ∫₋₁⁺¹ P_m^n(x) P_ν^μ(x) P_{n+ν-2q}^{m+μ}(x) dx.

or defined by:

    P_n^m(x) P_ν^μ(x) = ∑_{q=0}^{q_{max}} a(m,n,μ,ν,q) P_{n+ν-2q}^{m+μ}(x)

This is a Julia implementation of the stable recurrence described in:

    Y.-l. Xu, "Fast evaluation of Gaunt coefficients: recursive approach", J. Comp. Appl. Math., 85:53–65, 1997.
"""
function gaunt{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int;normalized::Bool=false)
    if normalized
        normalizedgaunt(T,m,n,μ,ν)
    else
        scale!(normalization(T,m,n,μ,ν),gaunt(T,m,n,μ,ν;normalized=true))
    end
end
"""
Calculates the Gaunt coefficients in 64-bit floating-point arithmetic.
"""
gaunt(m::Int,n::Int,μ::Int,ν::Int;kwds...) = gaunt(Float64,m,n,μ,ν;kwds...)

function normalization{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int)
    pochhammer(n+one(T),n)*pochhammer(ν+one(T),ν)/pochhammer(n+ν+one(T),n+ν)*gamma(n+ν-m-μ+one(T))/gamma(n-m+one(T))/gamma(ν-μ+one(T))
end

normalization(::Type{Float64},m::Int,n::Int,μ::Int,ν::Int) = normalization1(Float64,n,ν)*normalization2(Float64,n-m,ν-μ)

function normalization1(::Type{Float64},n::Int,ν::Int)
    if n ≥ 8
        if ν ≥ 8
            return exp((n+0.5)*log1p(n/(n+1))+(ν+0.5)*log1p(ν/(ν+1))+(n+ν+0.5)*log1p(-(n+ν)/(2n+2ν+1))+n*log1p(-2ν/(2n+2ν+1))+ν*log1p(-2n/(2n+2ν+1)))*stirlingseries(2n+1.0)*stirlingseries(2ν+1.0)*stirlingseries(n+ν+1.0)/stirlingseries(n+1.0)/stirlingseries(ν+1.0)/stirlingseries(2n+2ν+1.0)
        else
            return pochhammer(ν+1.0,ν)/(2n+2ν+1.0)^ν*exp(ν+(n+0.5)*log1p(n/(n+1))+(n+ν+0.5)*log1p(-(n+ν)/(2n+2ν+1))+n*log1p(-2ν/(2n+2ν+1)))*stirlingseries(2n+1.0)*stirlingseries(n+ν+1.0)/stirlingseries(n+1.0)/stirlingseries(2n+2ν+1.0)
        end
    elseif ν ≥ 8
        return pochhammer(n+1.0,n)/(2n+2ν+1.0)^n*exp(n+(ν+0.5)*log1p(ν/(ν+1))+(n+ν+0.5)*log1p(-(n+ν)/(2n+2ν+1))+ν*log1p(-2n/(2n+2ν+1)))*stirlingseries(2ν+1.0)*stirlingseries(n+ν+1.0)/stirlingseries(ν+1.0)/stirlingseries(2n+2ν+1.0)
    else
        return pochhammer(n+1.0,n)*pochhammer(ν+1.0,ν)/pochhammer(n+ν+1.0,n+ν)
    end
end

function normalization2(::Type{Float64},nm::Int,νμ::Int)
    if nm ≥ 8
        if νμ ≥ 8
            return edivsqrt2pi*exp((nm+0.5)*log1p(νμ/(nm+1))+(νμ+0.5)*log1p(nm/(νμ+1)))/sqrt(nm+νμ+1.0)*stirlingseries(nm+νμ+1.0)/stirlingseries(nm+1.0)/stirlingseries(νμ+1.0)
        else
            return (nm+νμ+1.0)^νμ*exp(-νμ+(nm+0.5)*log1p(νμ/(nm+1)))*stirlingseries(nm+νμ+1.0)/stirlingseries(nm+1.0)/gamma(νμ+1.0)
        end
    elseif νμ ≥ 8
        return (nm+νμ+1.0)^nm*exp(-nm+(νμ+0.5)*log1p(nm/(νμ+1)))*stirlingseries(nm+νμ+1.0)/stirlingseries(νμ+1.0)/gamma(nm+1.0)
    else
        return gamma(nm+νμ+1.0)/gamma(nm+1.0)/gamma(νμ+1.0)
    end
end

function normalizedgaunt{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int)
    qmax = min(n,ν,(n+ν-abs(m+μ))÷2)
    a = Vector{T}(qmax+1)
    a[1] = one(T)
    if μ == m && ν == n # zero class (i) of Aₚ
        if μ == m == 0
            for q = 1:qmax
                p = n+ν-2q
                a[q+1] = α(T,n,ν,p+2)/α(T,n,ν,p+1)*a[q]
            end
        else
            for q = 1:qmax
                p = n+ν-2q
                p₁,p₂ = p-m-μ,p+m+μ
                a[q+1] = (p+1)*(p₂+2)*α(T,n,ν,p+2)/(p+2)/(p₁+1)/α(T,n,ν,p+1)*a[q]
            end
        end
    else
        qmax > 0 && (a[2] = secondinitialcondition(T,m,n,μ,ν))
        q = 2
        if qmax > 1
            p = n+ν-2q
            p₁,p₂ = p-m-μ,p+m+μ
            if A(T,m,n,μ,ν,p+4) != 0
                a[q+1] = (c₁(T,m,n,μ,ν,p,p₁,p₂)*a[q] + c₂(T,m,n,μ,ν,p,p₂)*a[q-1])/c₀(T,m,n,μ,ν,p,p₁)
            else
                a[q+1] = thirdinitialcondition(T,m,n,μ,ν)
            end
            q+=1
        end
        while q ≤ qmax
            p = n+ν-2q
            p₁,p₂ = p-m-μ,p+m+μ
            if A(T,m,n,μ,ν,p+4) != 0
                a[q+1] = (c₁(T,m,n,μ,ν,p,p₁,p₂)*a[q] + c₂(T,m,n,μ,ν,p,p₂)*a[q-1])/c₀(T,m,n,μ,ν,p,p₁)
            elseif A(T,m,n,μ,ν,p+6) != 0
                a[q+1] = (d₁(T,m,n,μ,ν,p,p₁,p₂)*a[q] + d₂(T,m,n,μ,ν,p,p₁,p₂)*a[q-1] + d₃(T,m,n,μ,ν,p,p₂)*a[q-2])/d₀(T,m,n,μ,ν,p,p₁)
            else
                a[q+1] = (p+1)*(p₂+2)*α(T,n,ν,p+2)/(p+2)/(p₁+1)/α(T,n,ν,p+1)*a[q]
            end
            q+=1
        end
    end
    a
end

function secondinitialcondition{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int)
    n₄ = n+ν-m-μ
    mn = m-n
    μν = μ-ν
    temp = 2n+2ν-one(T)
    return (temp-2)/2*(1-temp/n₄/(n₄-1)*(mn*(mn+one(T))/(2n-1)+μν*(μν+one(T))/(2ν-1)))
end

function thirdinitialcondition{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int)
    n₄ = n+ν-m-μ
    mn = m-n
    μν = μ-ν
    temp = 2n+2ν-one(T)
    temp1 = mn*(mn+one(T))*(mn+2)*(mn+3)/(2n-1)/(2n-3) + 2mn*(mn+one(T))*μν*(μν+one(T))/(2n-1)/(2ν-1) + μν*(μν+one(T))*(μν+2)*(μν+3)/(2ν-1)/(2ν-3)
    temp2 = (temp-4)/(2(n₄-2)*(n₄-3))*temp1 - mn*(mn+one(T))/(2n-1)-μν*(μν+one(T))/(2ν-1)
    return temp*(temp-6)/4*( (temp-2)/n₄/(n₄-1)*temp2 + one(T)/2 )
end

α{T}(::Type{T},n::Int,ν::Int,p::Int) = (p^2-(n+ν+1)^2)*(p^2-(n-ν)^2)/(4p^2-one(T))
A{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int) = p*(p-one(T))*(m-μ)-(m+μ)*(n-ν)*(n+ν+one(T))

c₀{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₁::Int) = (p+2)*(p+3)*(p₁+1)*(p₁+2)*A(T,m,n,μ,ν,p+4)*α(T,n,ν,p+1)
c₁{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₁::Int,p₂::Int) = A(T,m,n,μ,ν,p+2)*A(T,m,n,μ,ν,p+3)*A(T,m,n,μ,ν,p+4) + (p+1)*(p+3)*(p₁+2)*(p₂+2)*A(T,m,n,μ,ν,p+4)*α(T,n,ν,p+2) + (p+2)*(p+4)*(p₁+3)*(p₂+3)*A(T,m,n,μ,ν,p+2)*α(T,n,ν,p+3)
c₂{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₂::Int) = -(p+2)*(p+3)*(p₂+3)*(p₂+4)*A(T,m,n,μ,ν,p+2)*α(T,n,ν,p+4)

d₀{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₁::Int) = (p+2)*(p+3)*(p+5)*(p₁+2)*(p₁+4)*A(T,m,n,μ,ν,p+6)*α(T,n,ν,p+1)
d₁{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₁::Int,p₂::Int) = (p+5)*(p₁+4)*A(T,m,n,μ,ν,p+6)*( A(T,m,n,μ,ν,p+2)*A(T,m,n,μ,ν,p+3) + (p+1)*(p+3)*(p₁+2)*(p₂+2)*α(T,n,ν,p+2) )
d₂{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₁::Int,p₂::Int) = (p+2)*(p₂+3)*A(T,m,n,μ,ν,p+2)*( A(T,m,n,μ,ν,p+5)*A(T,m,n,μ,ν,p+6) + (p+4)*(p+6)*(p₁+5)*(p₂+5)*α(T,n,ν,p+5) )
d₃{T}(::Type{T},m::Int,n::Int,μ::Int,ν::Int,p::Int,p₂::Int) = -(p+2)*(p+4)*(p+5)*(p₂+3)*(p₂+5)*(p₂+6)*A(T,m,n,μ,ν,p+2)*α(T,n,ν,p+6)


using Plots
using DifferentialEquations
using DiffEqUncertainty
#----------------------------------------------
# Specify and Solve model:

function fitzhugh_nagumo!(du, u, p, t)

    v, w = u
    a, b, I, R, τ = p

    du[1] = v - (v^3)/3 - w  + R * I

    du[2] = (v + a - (b * w)) / τ
end

u0 = [-2.0, 1.0]
p = [0.7, 0.8, 0.34, 1.0, 10]
tspan = (0.0, 100.0)
prob = ODEProblem(fitzhugh_nagumo!,u0,tspan,p)
sol = solve(prob)

plot(sol, vars =(1))

cb = ProbIntsUncertainty(0.1,1)
cb2 = ProbIntsUncertainty(0.05,1)
ensemble_prob = EnsembleProblem(prob)
sim = solve(ensemble_prob,Euler(),trajectories=100,callback=cb,dt=1/10)
sim2 = solve(ensemble_prob,Euler(),trajectories=50,callback=cb2,dt=1/100)

plot(sim, vars=(0,1), color="grey")
plot!(sim2, vars=(0,1), color="lightsalmon2")
plot!(sol, vars =(1), color="red3", linewidth =3)


"""
    Basic implementation of Weibull function with 2-parameters
    κ: shape parameter, κ > 0
    λ: scale parameter, λ > 0
    t: current time (generally calculated with wb_t function)

"""
weibull_f(λ,κ,t) = t >= 0 ? (κ / λ) * ((t/ λ)^(κ - 1))*(ℯ^-(t/λ)^κ) : 0

"""
 Weibull scaling function
 ttf: expected time-to-failure
 γ: factor to tune max y (yₙ) value, e.g. for max y=1, with ttf =100,
 λ=0.5, κ=05 => γ=8

 TODO: check if γ might be calculated from ttf, λ and κ
"""
wb_t(γ,ttf,t) = (t%ttf)*γ/ttf

"""
 Basic logistic function
 http://reliawiki.org/index.php/Logistic

"""

log_f(b,k,t) = 1 / ( 1+ b*(ℯ^(-k*t)))

"""
Add random normally distributed  noise to a given value.
Ensures that always:  0 ⋜  v + ϵ ⋜ y, for a random noise ϵ.

    v: original value without noise
    ϵₛ: noise standard deviation
    ϵ₀: noise μ
    y:  max value the noised value can take
"""
function nnoised(v,ϵₛ,ϵ₀,y)
    ϵ = 0
    #force only noised values within range
    while true
        ϵ = rand(Normal(ϵ₀,ϵₛ),1)[1]
        (( v + ϵ ) > y || ( v + ϵ ) < 0)  || break
    end
    return v + ϵ
end

"""
Add random normally distributed noise between 0 and y=1.
See nnoised(v,ϵₛ,ϵ₀,y).
"""

nnoised(v,ϵₛ) = nnoised(v,ϵₛ,0,1)

"""
    Calculates remaining ttf(time-to-failure) for a given expected ttf
    and current t
    ttf: expected time-to-failure
    t: current time

"""
exp_t(ttf,t) = ttf - (t%ttf)

"""
    Basic implementation of exponential deterioration function
    according to:

    * T. Wang, J. Yu, D. Siegel, and J. Lee, “2008 International Conference on Prognostics and Health Management, PHM2008,”
      2008 Int. Conf. Progn. Heal. Manag. PHM 2008, 2008.

      * A. Salvador Palau, M. Dhada, and A. Parlikad, “Multi-Agent System architectures for collaborative prognostics,”
      J. Intell. Manuf., 2019.

    a: max expected value of condition/health indicator (i.e. 1 if normalised)
    b: curvature parameter
    t: remaining time-to-failure (generally calculated with exp_t function)

"""
exp_f(a,b,t) = a*(1.0 - ℯ^(-(b)*t))

"""
    Scales weibull input t value considering shifting
    for long term behaviour

    ttf: expected time-to-failure
    γ: factor to tune max y (yₙ) value, e.g. for max y=1, with ttf =100,
    λ=0.5, κ=05 => γ=8
    d: units to shift to the right (near TTF)

"""
wb_ts(γ,d,ttf,t) = ((t+d)%ttf)*γ/ttf

"""
    Scales exponential function considering shifting for long term
    behaviour.

    ttf: expected time-to-failure
    t: value to scale
    d: units to shift to the right (near TTF)
"""
exp_ts(d,ttf,t) = ttf - (t%ttf) -d

"""
    Scales logistic function considering shifting for long term
    behaviour.

    ttf: expected time-to-failure
    t: value to scale
    d: units to shift to the right (near TTF)
"""

log_ts(d,ttf,t) = (t+d)%ttf

"""
Introduces a variation on the x value according to the phase (ϕ)
Used concretely for long-term TTF, we assume:
    * An asset has expected TTF (x)
    * After TTF the asset can be repaired or maintained and then will
      operate for a number of cycles/phases.
    * The current cycle/phase is given by ϕ
    * After every phase ϕ, the asset reduces its TTF by Δᵩ
"""
xᵩ(x,ϕ,Δᵩ) =  x * (1 -  (ϕ * Δᵩ))

"""
    Determines criteria for valid values to TTF for
    the exponential function.

        t₀: unused, defined to keep uniformity of function
            (TODO: check if this can be removed)
        t: value to evaluate
"""
exp_c(t₀,t) = t > 0

"""
    Determines criteria for valid values to TTF for
    the weibull function.

        t₀: initial value for the cycle/phase
        t: value to evaluate
"""

wb_c(t₀,t) = t >= t₀

"""
    Determines criteria for valid values to TTF for
    the logistics function.

        t₀: initial value for the cycle/phase
        t: value to evaluate
"""
log_c(t₀,t) = t >= t₀

"""
Linear degradation
m: degradation slope
b: expected life
t: current time
"""
lineal_d(m,b,t) = begin
    v = -m * t + b
    return v > 0 ? v : 0
end

"""
    Generates the time series until failure of an asset, given
    the provided parameters:

    ttf: expected time-to-failure
    n: maximum number of cycles the asset can last
    Δᵩ: Delta of number of time steps that are reduced as the asset
        reaches a new cycle/phase
    ϵₛ: Standard deviation of the gaussian noise
    f: function to generate values until ttf
    params: params required by the f
    tsf: function to scale t for the given function
    tse: extra parameter required by tsf
    f_c: function that defines criteria for a valid value in the time series.
"""
function values_f(ttf,n,Δᵩ,ϵₛ,f,params,tsf,tse_params,f_c)
    ϕ = 0
    t₀ = 0
    vs = []
    for i=0:n-1
        t = i
        ϕ₀ = ϕ
        ϕ = floor(t/ttf)
        rttf = xᵩ(ttf,ϕ,Δᵩ)
        diff = ttf - rttf
        t = tsf(tse_params...,diff,ttf,t)
        t₀ = i == 1 ? t : ϕ > ϕ₀ ? t : t₀
        remt = ϕ > ϕ₀ ? (n - i) >= rttf && rttf > 0 : true
        if f_c(t₀,t) && remt
            #push!(vs_f,(t,t₀,ϕ,ϕ₀,rmng,rttf,(n - i),f(params...,t)))
            push!(vs,nnoised(f(params...,t),ϵₛ))
        end
    end
    l = length(vs) + 1
    for i=l:n
        push!(vs,0.0)
    end
    return vs
end

"""
 It illustrates how to use the other TTE functions
"""
function generate_ttf_series()
    n=1000
    ttf = 200
    λ = 1.0 #0.5
    γ = 6.0 #16
    k = 1.0 #0.5
    #p = plot(title="",legend=false)#,xlims=[0,1000])
    ϵₛ = 0.05
    a = 1.0
    b = 0.05
    Δᵩ = 0.05#0.009
    maxᵩ = 9

    vs = []
    funs = [
            (exp_f,(a,b),exp_ts,(),exp_c),
            (weibull_f,(λ,k),wb_ts,(γ),wb_c),
            (log_f,(50,0.1),log_ts,(),log_c)
            ]

    for fs in funs
        push!(vs,values_f(ttf,n,Δᵩ,ϵₛ,fs...))
    end

    # for d=1:length(vs)
    #     p = plot!(p, vs[d])
    # end
    # for c=ttf:ttf:n
    #     p = plot!(p,[c,c],[0,1.0], c=:red, ls=:dash, w=2)
    # end
    # p
end

"""
    It adds downtime to a series
        ot: downtime (time steps)

"""

function add_downtime_to_series(ttf,ot,s_series)

    #log_info(" RECEIVED adding downtime ==>$(s_series)<==")

    # split s_series in phases/cycles of ttf
    phased_s = [s_series[i:min(i + ttf - 1, end)] for i in 1:ttf:length(s_series)]

    ph_ot = Vector{Float64}()

    for ph in phased_s
        if length(ph) < ttf
            push!(ph_ot,ph...)
        else
            push!(ph_ot,vcat(ph,zeros(Float64,ot))...)
        end
    end

    # log_info(" adding downtime $(size(ph_ot)) --- $(phased_s)")
    return ph_ot
end

"""
 It generates the sensor time series for a given asset
    ttf:time-to-failure
    ot: off-time
"""
function generate_sensor_series(ttf,n,Δᵩ,ϵₛ,ot,funs)
    vs = zeros(Float64,length(funs),n)
    #log_info("received generate sensor ==> $ot ")
    dim_dtvs = n + Int64(floor(n/ttf)) * ot  #n%ttf > 0 && n >= ttf ? n+ot+1 : n+ot

    dtvs = zeros(Float64,length(funs),dim_dtvs)

    for i=1:length(funs)
        vs[i,:] = values_f(ttf,n,Δᵩ,ϵₛ,funs[i]...)
    end

    for si=1:length(funs)
        s = vs[si,:]
        dtvs[si,:] = add_downtime_to_series(ttf,ot,s)
    end

    return dtvs
end

function generate_rul_series(ttf,Δᵩ,n,ot)
    rul = [  i%xᵩ(ttf,floor((i-1)/ttf),Δᵩ) == 0 ? 0 : round(xᵩ(ttf,floor((i-1)/ttf),Δᵩ) - i%xᵩ(ttf,floor((i-1)/ttf),Δᵩ)) for i=1:n ]
    #rul = [ xᵩ(ttf,floor(i/ttf),Δᵩ) - i%xᵩ(ttf,floor(i/ttf),Δᵩ) + 1 for i=1:n ]
    # print("PRE RUL: $(size(rul))")
    rul = add_downtime_to_series(ttf,ot,rul)
    # print("RUL: $(size(rul))")
    return rul
end


module Madelung

using SpecialFunctions: erfc
using LinearAlgebra: dot, det, inv

function integer_mesh(mesh)
  spaces = (-n1:n1 for n1 in mesh)
  gvecs = Iterators.product(spaces...)
  return gvecs
end

function sum_lattice(vsr_of_r::Function, axes::Matrix, nr::Int)
  val = 0.0
  ndim = size(axes, 1)
  mesh = nr*ones(Int, ndim)
  for g in integer_mesh(mesh)
    a = axes*collect(g)
    r = sqrt(dot(a, a))
    if (r<eps(Float64))
      continue
    end
    val += vsr_of_r(r)
  end
  return val
end

function madelung_constant(axes::Matrix, alpha::Float64, nr::Int, nk::Int)
  ndim = size(axes, 1)
  # real-space part is independent of spatial dimensions
  vsr_of_r(r) = erfc(alpha*r)/r
  vlr_r0 = 2*alpha/sqrt(pi)
  # reciprocal-space part changes with dimension
  vlr_of_k2d(k) = 2*pi/k*erfc(k/(2*alpha))
  vlr_of_k3d(k) = 4*pi/(k*k)*exp(-k*k/(4*alpha*alpha))
  if ndim == 2
    vsr_k0 = 2*sqrt(pi)/alpha
    vlr_of_k = vlr_of_k2d
  elseif ndim == 3
    vsr_k0 = pi/(alpha*alpha)
    vlr_of_k = vlr_of_k3d
  else
    msg = "ndim=$ndim not implemented"
    err = ErrorException(msg)
    throw(err)
  end
  # perform Ewald sums
  raxes = 2*pi*transpose(inv(axes))
  vsr = sum_lattice(vsr_of_r, axes, nr)
  vlr = sum_lattice(vlr_of_k, raxes, nk)
  # collect results
  volume = abs(det(axes))
  vsr = (vsr - vlr_r0)/2
  vlr = (vlr - vsr_k0)/volume/2
  vmad = vsr+vlr
  return vmad
end

end


a_wgs84=6378137.
f_wgs84 = 1/298.257223563
b_wgs84 = a_wgs84*(1-f_wgs84)
e_wgs84 = sqrt((a_wgs84^2-b_wgs84^2)/a_wgs84^2)

#export ECEFtoGeodetic, GeodetictoECEF
export GeodetictoECEF

function GeodetictoECEF2(lat::Number, lon::Number, h::Number)
    # Auxiliary variables.
    sin_lat, cos_lat = sincos(lat)
    sin_lon, cos_lon = sincos(lon)

    # Radius of curvature [m].
    N = a_wgs84/sqrt(1 - e_wgs84^2*sin_lat^2 )

    # Compute the position in ECEF frame.
    SVector{3}((                     N + h)*cos_lat*cos_lon,
               (                     N + h)*cos_lat*sin_lon,
               ( (b_wgs84/a_wgs84)^2*N + h)*sin_lat)

    println(SVector)
end