using RungeKutta
using Base.Test

## Solves the test problem dy/dt = -t^3 y^3
##
##   which has the exact solution y = 1/sqrt(C + t^4/2)
##   where C = 1/y_0^2 - t_0^4/2

summer = +
scaler = *

con = -0.4
deriv = (t,y) ->
begin
#    println (string("deriv: t = ", t, " y = ", y, " y' = ", con*y))
    con*y
end
t0 = 0.0
y0 = 1.75
exact = ((t) -> y0*exp(con*(t - t0)))

showst = ((t,y) -> string (y, "\t", (y - exact (t))))

function gen_soln1 (integrator,h,t,st)
    stn = st
    tn = t
    while tn < 5.0
        stn = integrator (tn,stn)
        tn  = tn+h
    end
    stn = integrator (tn,stn)
    tn  = tn+h
    println (showst (tn,stn))
end

function do_case1 (integrator, n)
    h = (n < 0) ? 2^(-n) : 0.5^(n)
    print (string (h, "\t\t"))
    gen_soln1 (integrator (h),h,t0,y0)
end

function solver1 (integrator,stats)
    println (stats)
    println ("# step yf err")
    for h in { x - 2 for x = 0:14 }
        do_case1 (integrator (scaler,summer,deriv), h)
    end
    println ("# All done!")
    end

solver1 (make_rkfe(), show_rkfe)
solver1 (make_rk3(), show_rk3)
solver1 (make_rk4a(), show_rk4a)
solver1 (make_rk4b(), show_rk4b)


# Whitening

## Solve whitening based on covariance
#
# finds W, such that W'CW = I
#
function cov_whitening(C::Cholesky{T}) where {T<:Real}
    cf = C.UL
    Matrix{T}(inv(istriu(cf) ? cf : cf'))
end

cov_whitening!(C::DenseMatrix{<:Real}) = cov_whitening(cholesky!(Hermitian(C, :U)))
cov_whitening(C::DenseMatrix{<:Real}) = cov_whitening!(copy(C))

cov_whitening!(C::DenseMatrix{<:Real}, regcoef::Real) = cov_whitening!(regularize_symmat!(C, regcoef))
cov_whitening(C::DenseMatrix{<:Real}, regcoef::Real) = cov_whitening!(copy(C), regcoef)

## Whitening type

struct Whitening{T<:Real}
    mean::Vector{T}
    W::Matrix{T}

    function Whitening{T}(mean::Vector{T}, W::Matrix{T}) where {T<:Real}
        d, d2 = size(W)
        d == d2 || error("W must be a square matrix.")
        isempty(mean) || length(mean) == d ||
        throw(DimensionMismatch("Sizes of mean and W are inconsistent."))
        return new(mean, W)
    end
end
Whitening(mean::Vector{T}, W::Matrix{T}) where {T<:Real} = Whitening{T}(mean, W)

indim(f::Whitening) = size(f.W, 1)
outdim(f::Whitening) = size(f.W, 2)
mean(f::Whitening) = fullmean(indim(f), f.mean)

transform(f::Whitening, x::AbstractVecOrMat) = transpose(f.W) * centralize(x, f.mean)

## Fit whitening to data

function fit(::Type{Whitening}, X::DenseMatrix{T};
             mean=nothing, regcoef::Real=zero(T)) where {T<:Real}
    n = size(X, 2)
    n > 1 || error("X must contain more than one sample.")
    mv = preprocess_mean(X, mean)
    Z = centralize(X, mv)
    C = rmul!(Z * transpose(Z), one(T) / (n - 1))
    return Whitening(mv, cov_whitening!(C, regcoef))
end

# invsqrtm

function _invsqrtm!(C::Matrix{<:Real})
    n = size(C, 1)
    size(C, 2) == n || error("C must be a square matrix.")
    E = eigen!(Symmetric(C))
    U = E.vectors
    evs = E.values
    for i = 1:n
        @inbounds evs[i] = 1.0 / sqrt(sqrt(evs[i]))
    end
    rmul!(U, Diagonal(evs))
    return U * transpose(U)
end

invsqrtm(C::DenseMatrix{<:Real}) = _invsqrtm!(copy(C))


# This file illustrates the simulation of MFSK signalling. The simulation is performed using
# waveform AWGN channel.

using Plots
using FFTW
using LinearAlgebra
using Statistics
using DigitalCommunications

# Settings
k = 2                                       # Bits per symbol
M = 2^3                                     # Constellation size
T = 1.                                      # Symbol duration
f = 1 / T                                   # Symbols per second
fs = 10                                     # Sampling frequency
ts = 1 / fs                                 # Sampling period
t = 0 : ts : T - ts                         # Time vector for symbol duration
l = length(t)                               # Number of samples per symbol
g(t) = 0 ≤ t ≤ T ? sqrt(2 * Eg / T) : 0     # Modulation pulse: rectangular pulse
Eg = 1                                      # Modulation pulse energy: rectangular pulse
nsymbols = Int(1e6)                         # Number of symbols
Δf = 1 / T                                  # 2 times the minimum frequency seperation for orthogonality

# Define basis
basis = map(m -> (t -> g(t) * exp(1im * 2π * m * Δf * t)), 1 : M)
basisvals = map(base -> base.(t), basis)
alphabet = map(m -> setindex!(zeros(M), 1, m), 1 : M)
references = map(1 : M) do mi
    Ai = alphabet[mi] * sqrt(2 * Eg)
    sum(Ai .* basisvals)
end

# Simulate the system
ebno = collect(0 : 1 : 10)
esno = ebno .+ 10 * log10(k)
m = rand(1 : M, nsymbols)
tx = map(m) do mi
    Ai = alphabet[mi] * sqrt(Eg)
    sum(Ai .* basisvals)
end
Es = mean(txi -> sum(abs.(txi).^2) * ts, tx)
symerr = zeros(length(esno))
for i = 1 : length(symerr)
    # Generate noise
    σ = sqrt(Es * fs / 2 / (10^(esno[i] / 10)))
    n = [σ * (randn(l) + 1im * randn(l)) for i = 1 : length(tx)]
    # Corrupt transmitted signal
    rx = tx + n
    # Detect symbols
    mr = map(rx) do rxi
        argmax(map(reference -> real(rxi ⋅ reference), references))
    end
    # Calculate symbol
    ber = sum(m .!= mr) / length(m)
    symerr[i] = ber
end

# Plots
plt = plot(title="$M-FSK", xlabel="esno [dB]", ylabel="Pe")
plot!(esno, berfsk.(esno, M), marker=:circle, yscale=:log10, label="theoretical")
plot!(esno, symerr, marker=:circle, yscale=:log10, label="montecarlo")


# # Basics
#
# ## Packages
#
# We load the following packages into our environment:

using OptimalTransport
using Distances
using Plots
using PythonOT: PythonOT
using Tulip

using LinearAlgebra
using Random

Random.seed!(1234)

const POT = PythonOT

# ## Problem setup
#
# First, let us initialise two random probability measures $\mu$ (source measure)
# and $\nu$ (target measure) in 1D:

M = 200
μ = fill(1 / M, M)
μsupport = rand(M)

N = 250
ν = fill(1 / N, N)
νsupport = rand(N);

# Now we compute the quadratic cost matrix $C_{ij} = \| x_i - x_j \|_2^2$:

C = pairwise(SqEuclidean(), μsupport', νsupport'; dims=2);

# ## Exact optimal transport
#
# The **earth mover's distance** is defined as the optimal value of the
# Monge-Kantorovich problem
# ```math
# \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle =
# \inf_{\gamma \in \Pi(\mu, \nu)} \sum_{i, j} \gamma_{ij} C_{ij},
# ```
# where $\Pi(\mu, \nu)$ denotes the set of couplings of $\mu$ and $\nu$, i.e., the set of joint
# distributions whose marginals are $\mu$ and $\nu$.
# If $C$ is the quadratic cost matrix, the earth mover's distance is known as the square of
# the **2-Wasserstein distance**.
#
# The function `emd` returns the optimal transport *plan* $\gamma$

γ = emd(μ, ν, C, Tulip.Optimizer());

# whilst using `emd2` returns the optimal transport *cost*:

emd2(μ, ν, C, Tulip.Optimizer())

# ## Entropically regularised optimal transport
#
# We may add an entropy term to the Monge-Kantorovich problem to obtain the
# **entropically regularised** optimal transport problem
# ```math
# \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle + \epsilon \Omega(\gamma),
# ```
# where $\Omega(\gamma) = \sum_{i, j} \gamma_{ij} \log(\gamma_{ij})$ is the negative
# entropy of the coupling $\gamma$ and $\epsilon$ controls the strength of the regularisation.
#
# This problem is *strictly convex* and admits a very efficient iterative scaling scheme for
# its solution known as the [Sinkhorn algorithm](https://doi.org/10.1214/aoms/1177703591).
#
# We compute the optimal entropically regularised transport plan:

ε = 0.01
γ = sinkhorn(μ, ν, C, ε);

# We can check that one obtains the same result with the Python Optimal Transport (POT) package:

γpot = POT.sinkhorn(μ, ν, C, ε)
norm(γ - γpot, Inf)

# We can compute the optimal cost (a scalar) of the entropically regularized optimal
# transport problem with `sinkhorn2`:

sinkhorn2(μ, ν, C, ε)

# ## Quadratically regularised optimal transport
#
# Instead of the common entropically regularised optimal transport problem, we can solve the
# [quadratically regularised optimal transport problem](https://doi.org/10.1007/s00245-019-09614-w)
# ```math
# \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle + \epsilon \frac{\| \gamma \|_F^2}{2}.
# ```
# One property of the quadratically regularised optimal transport problem is that the
# resulting transport plan $\gamma$ is *sparse*. We take advantage of this and represent it as
# a sparse matrix.

quadreg(μ, ν, C, ε; maxiter=100);

# ## Stabilized Sinkhorn algorithm
#
# When $\epsilon$ is very small, we can use a
# [log-stabilised version of the Sinkhorn algorithm](https://doi.org/10.1137/16M1106018).

ε = 0.005
γ = sinkhorn_stabilized(μ, ν, C, ε; maxiter=5_000);

# Again we can check that the same result is obtained with the POT package:

γ_pot = POT.sinkhorn(μ, ν, C, ε; method="sinkhorn_stabilized", numItermax=5_000)
norm(γ - γ_pot, Inf)

# ## Stabilized Sinkhorn algorithm with $\epsilon$-scaling
#
# In addition to log-stabilisation, we can use [$\epsilon$-scaling](https://doi.org/10.1137/16M1106018):

γ = sinkhorn_stabilized_epsscaling(μ, ν, C, ε; maxiter=5_000);

# The POT package yields the same result:

γpot = POT.sinkhorn(μ, ν, C, ε; method="sinkhorn_epsilon_scaling", numItermax=5000)
norm(γ - γpot, Inf)

# ## Unbalanced optimal transport
#
# [Unbalanced optimal transport](https://doi.org/10.1090/mcom/3303) deals with general
# positive measures which do not necessarily have the same total mass. For unbalanced
# source and target marginals $\mu$ and $\nu$ and a cost matrix $C$, entropically
# regularised unbalanced optimal transport solves
# ```math
# \inf_{\gamma \geq 0} \langle \gamma, C \rangle + \epsilon \Omega(\gamma) + \lambda_1 \mathrm{KL}(\gamma 1 | \mu) + \lambda_2 \mathrm{KL}(\gamma^{\mathsf{T}} 1 | \nu),
# ```
# where $\epsilon$ controls the strength of the entropic regularisation, and $\lambda_1$ and
# $\lambda_2$ control how strongly we enforce the marginal constraints.
#
# We construct two random measures, now with different total masses:

M = 100
μ = fill(1 / M, M)
μsupport = rand(M)

N = 200
ν = fill(1 / M, N)
νsupport = rand(N);

# We compute the quadratic cost matrix:

C = pairwise(SqEuclidean(), μsupport', νsupport'; dims=2);

# Now we solve the corresponding unbalanced, entropy-regularised transport problem with
# $\epsilon = 0.01$ and $\lambda_1 = \lambda_2 = 1$:

ε = 0.01
λ = 1
γ = sinkhorn_unbalanced(μ, ν, C, λ, λ, ε);

# We check that the result agrees with POT:

γpot = POT.sinkhorn_unbalanced(μ, ν, C, ε, λ)
norm(γ - γpot, Inf)

# ## Plots
#
# ### Entropically regularised transport
#
# Let us construct source and target measures again:

μsupport = νsupport = range(-2, 2; length=100)
C = pairwise(SqEuclidean(), μsupport', νsupport'; dims=2)
μ = normalize!(exp.(-μsupport .^ 2 ./ 0.5^2), 1)
ν = normalize!(νsupport .^ 2 .* exp.(-νsupport .^ 2 ./ 0.5^2), 1)

plot(μsupport, μ; label=raw"$\mu$", size=(600, 400))
plot!(νsupport, ν; label=raw"$\nu$")

# Now we compute the entropically regularised transport plan:

γ = sinkhorn(μ, ν, C, 0.01)
heatmap(
    μsupport,
    νsupport,
    γ;
    title="Entropically regularised optimal transport",
    size=(600, 600),
)

# ### Quadratically regularised transport
#
# Notice how the "edges" of the transport plan are sharper if we use quadratic regularisation
# instead of entropic regularisation:

γquad = quadreg(μ, ν, C, 5; maxiter=100);
heatmap(
    μsupport,
    νsupport,
    γquad;
    title="Quadratically regularised optimal transport",
    size=(600, 600),
)

# ### Sinkhorn barycenters
#
# For a collection of discrete probability measures $\{\mu_i\}_{i=1}^N \subset \mathcal{P}$,
# cost matrices $\{C_i\}_{i=1}^N$, and positive weights $\{\lambda_i\}_{i=1}^N$ summing to $1$,
# the entropically regularised **barycenter** in $\mathcal{P}$ is the discrete probability
# measure $\mu$ that solves
# ```math
# \inf_{\mu \in \mathcal{P}} \sum_{i = 1}^N \lambda_i \operatorname{OT}_{\epsilon}(\mu, \mu_i)
# ```
# where $\operatorname{OT}_\epsilon(\mu, \mu_i)$ denotes the entropically regularised
# optimal transport cost with marginals $\mu$ and $\mu_i$, cost matrix $C$, and entropic
# regularisation parameter $\epsilon$.
#
# We set up two measures and compute the weighted barycenters. We choose weights
# $\lambda_1 \in \{0.25, 0.5, 0.75\}$.

support = range(-1, 1; length=250)
mu1 = normalize!(exp.(-(support .+ 0.5) .^ 2 ./ 0.1^2), 1)
mu2 = normalize!(exp.(-(support .- 0.5) .^ 2 ./ 0.1^2), 1)

plt = plot(; size=(800, 400), legend=:outertopright)
plot!(plt, support, mu1; label=raw"$\mu_1$")
plot!(plt, support, mu2; label=raw"$\mu_2$")

mu = hcat(mu1, mu2)
C = pairwise(SqEuclidean(), support'; dims=2)
for λ1 in (0.25, 0.5, 0.75)
    λ2 = 1 - λ1
    a = sinkhorn_barycenter(mu, C, 0.01, [λ1, λ2], SinkhornGibbs())
    plot!(plt, support, a; label="\$\\mu \\quad (\\lambda_1 = $λ1)\$")
end
plt


using SeisPlot, SeisAcoustic, FFTW, LinearAlgebra

# velocity and density model
nz = 101; nx = 301
vel = 3000 * ones(nz, nx);  # m/s
vel[51:end,:] .= 3500;  # m/s
rho = 2000 * ones(nz, nx);  # kg/m^3

# top boundary condition
free_surface = false    #(pml or free_surface)

# grid size
h = 10

# organize these parameters into a structure
fparams = FdParams(rho, vel, free_surface, h; data_format=Float64);

# generate observations of one frequency slice by a point source
isz = 5; isx = 151; dt = 0.001;
src = Source(isz, isx, dt, fparams; amp=1/rho[isz,isx]/vel[isz,isx]^2*10^5);

# total modelling length
tmax = 2.0;
nt   = floor(Int64, tmax/dt) + 1

# working frequency
df = 1.0 / (dt * nt);
f_work = 20.0;
iw = floor(Int64, f_work/df) + 1;
f_work = (iw-1) * df;
omega  = 2.0 * pi * f_work;
wlet   = zeros(nt)
copyto!(wlet, src.it_min, src.p)
fwave  = fft(wlet)

# LU decomposition of Helmhotz operator
H = get_helmholtz_LU(fparams, omega);

# source vector
s = zeros(Complex{Float64}, fparams.Nz*fparams.Nx)
s[src.src2spt] = fwave[iw]

# produce modelling true wave field
u = zeros(Complex{Float64}, fparams.Nz * fparams.Nx);
ldiv!(u, H, s);
SeisPlotTX(real(reshape(u, fparams.Nz, fparams.Nx)), wbox=9, hbox=3)

# specify the monochromatic recordings
irx = collect(1:2:nx);
irz = 2 * ones(Int64, length(irx));
dobs= MonochromaticRecordings(irz, irx, omega, fparams);
sample_spt2rec!(dobs, u);

# initial homogenous velocity model
vel = 3000 * ones(nz, nx)
fparams = FdParams(rho, vel, free_surface, h; data_format=Float64);

# helmholtz operator of homogeneous model
H = get_helmholtz_LU(fparams, omega);

# source side wavefield
u = zeros(Complex{Float64}, fparams.Nz * fparams.Nx);
ldiv!(u, H, s);
SeisPlotTX(real(reshape(u, fparams.Nz, fparams.Nx)), wbox=9, hbox=3)

dsyn = MonochromaticRecordings(irz, irx, omega, fparams);
sample_spt2rec!(dsyn, u);

# compute the residue
dres = get_residue(dsyn, dobs);
g = velocity_gradient(dres, u, H, fparams);
g = reshape(g, fparams.nz, fparams.nx)
SeisPlotTX(real(g), wbox=9, hbox=3)

# compute the gradient numerically
vel0    = 3000 * ones(nz, nx);
delta_m = 1e-7;

# one model parameter's gradient
iz = 37; ix = 137;

# positive velocity model partubation
vel0[iz,ix] = vel0[iz,ix] + delta_m

# model parameter
fparams = FdParams(rho, vel0, free_surface, h; data_format=Float64);
H = get_helmholtz_LU(fparams, omega);

# source side wavefield
u = zeros(Complex{Float64}, fparams.Nz * fparams.Nx);
ldiv!(u, H, s);

d1 = MonochromaticRecordings(irz, irx, omega, fparams);
sample_spt2rec!(d1, u);
SeisPlotTX(real(reshape(u, fparams.Nz, fparams.Nx)), wbox=9, hbox=3)

# compute the residue
dres = get_residue(d1, dobs);
J_plus = 1/2.0 * (norm(dres.p))^2

# negative velocity model partubation
vel0[iz,ix] = vel0[iz,ix] - 2*delta_m

fparams = FdParams(rho, vel0, free_surface, h; data_format=Float64);
H = get_helmholtz_LU(fparams, omega);

# source side wavefield
u = zeros(Complex{Float64}, fparams.Nz * fparams.Nx);
ldiv!(u, H, s);

d2 = MonochromaticRecordings(irz, irx, omega, fparams);
sample_spt2rec!(d2, u);

# compute the residue
dres = get_residue(d2, dobs);
J_minus = 1/2.0 * (norm(dres.p))^2;

# compute gradient numerically
g_num = (J_plus - J_minus) / (2 * delta_m)
g[iz,ix]
(g[iz,ix] - g_num) / g_num


#=
A number chain is created by continuously adding the square of the digits in a number to form a new number until it has been seen before.

For example,

44 → 32 → 13 → 10 → 1 → 1
85 → 89 → 145 → 42 → 20 → 4 → 16 → 37 → 58 → 89

Therefore any chain that arrives at 1 or 89 will become stuck in an endless loop. What is most amazing is that EVERY starting number will eventually arrive at 1 or 89.

How many starting numbers below ten million will arrive at 89?
=#
cache = Int[]
function digitchain(n)
  while true
    if n < length(cache)
      return cache[n]
    end
    n = map((x)->x^2, digits(n)) |> sum
    if n == 1
      return 1
    elseif n == 89
      return 89
    end
  end
end

function calc()
  count = 0
  for i in 1:10_000_000
    d = digitchain(i)
    push!(cache, d)
    if d == 89
      count = count + 1
    end
  end
  count
end
@time println(calc())


# Lotka-Volterra, np = 1
id = :Sahlodin2011b

pL = [2.95]
pU = [3.05]
Sahlodin2011b_x0(p::Vector{T}) where T = T[1.2*one(T); 1.1*one(T)]
function Sahlodin2011b_f!(du::Vector{T}, u::Vector{T}, p::Vector{T}, t) where T
    du[1] = p[1]*u[1]*(1.0 - u[2])
    du[2] = p[1]*u[2]*(u[1] - 1.0)
    return
end
tspan = (0.0, 2.0)

prob = DBB.ODERelaxProb(Sahlodin2011b_f!, tspan, Sahlodin2011b_x0, pL, pU)
url = "https://www.sciencedirect.com/science/article/abs/pii/S0168927411000316"
source = "Sahlodin2011"
desc = "Lotka-Volterra, np = 1"


"""
    autocovariance(v::Vector,h::Int; corrected::Bool=true)
``\\hat{\\gamma}(h) =\\frac{1}{n}\\sum_{t=1}^{n-h}(v_{t+h}-\\bar{v})(v_t-\\bar{v})^*``
Calculate the autocovariance of dataset `v` with a delay `h`. If `corrected` is `true`
(the default) then the sum is scaled with `n-h`, whereas the sum
is scaled with `n` if corrected is `false` where `n = length(v)`.
"""
function autocovariance(v::AbstractVector,h::Int; corrected::Bool=true, mean = mean(v))
    n = length(v)
    covsum = zero(mean)
    for i in 1:n-h
        @inbounds covsum += (v[i]-mean)*conj(v[i+h]-mean)
    end
    gamma = covsum/(n - Int(corrected))
    return gamma
end

@doc raw"""
    pseudo_cov(x, y; xmean = mean(x), ymean = mean(y), corrected = true)
Compute the pseudo covariance between collections `x` and `y` returning a scalar:
```math
\frac{1}{n}\sum_{i=1}^{n} (x_i - \bar{x})(y_{i} - \bar{y})
```
Optionally,
precomputed means can be passed as keyword arguments. `pseudo_cov(x,y)` is functionally
equivalent to `Statistics.cov(x, conj(y); corrected = false)` but it is found to be
significantly faster and avoids allocations.
"""
@inline function pseudo_cov(x, y; xmean = mean(x), ymean = mean(y), corrected = true)
    n = length(x)
    @assert length(y) == n
    res = zero(promote_type(eltype(x), eltype(y))) / 1
    for i = 1:n
        @inbounds res += (x[i] - xmean) * (y[i] - ymean)
    end
    return res / (n - Int(corrected))
end


using TuringModels

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

# Turing model

@model m8_3(y) = begin
    α ~ Normal(1, 10)
    σ ~ Truncated(Cauchy(0, 1), 0, Inf)

    for i ∈ 1:length(y)
        y[i] ~ Normal(α, σ)
    end
end

y = [-1,1]

# Sample

chns = sample(m8_3(y), NUTS(0.95), 1000);

# Results rethinking

m83rethinking = "
      mean   sd  5.5% 94.5% n_eff Rhat
alpha 0.09 1.63 -2.13  2.39   959    1
sigma 2.04 2.05  0.68  4.83  1090    1
";

# Describe the posterior samples

describe(chns)

# End of `08/m8.3t.jl`


function isPrime(number::Int64)
    if number == 1
        return false
    elseif number == 2
        return true
    else
        for i in 2:sqrt(number)
            if mod(number, i) == 0  # Prime numbers are not divisible by the numbers that come before their square roots.
                return false
            end
        end
        return true  # if for loop finished, it's really a prime number :D
    end
end
while true
    number = parse(Int64, Base.prompt("Enter the number you want to check, press '0' to exit program")) #Taking input from user and creating while loop.
    if number == 0
        break
    else
        if isPrime(number)
            println(number, " is a prime number!")
        else
            println(number, " is NOT a prime number!")
        end
    end
end


using ModelingToolkit

### Define a differential equation system

@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
de = ODESystem(eqs, t, [x,y,z], [σ,ρ,β])
ode_f = ODEFunction(de)

### Use in DifferentialEquations.jl

using OrdinaryDiffEq
u₀ = ones(3)
tspan = (0.0,100.0)
p = [10.0,28.0,10/3]
prob = ODEProblem(ode_f,u₀,tspan,p)
sol = solve(prob,Tsit5())

using Plots
plot(sol,vars=(1,2,3))

generate_function(de)[1]

generate_function(de)[2]

jac = calculate_jacobian(de)

jac_expr = generate_jacobian(de)

ModelingToolkit.generate_factorized_W(de)[1]

@variables x y z
@parameters σ ρ β

# Define a nonlinear system
eqs = [0 ~ σ*(y-x),
       0 ~ x*(ρ-z)-y,
       0 ~ x*y - β*z]
ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β])
nlsys_func = generate_function(ns)

nl_f = @eval eval(nlsys_func[2])
# Make a closure over the parameters for for NLsolve.jl
f2 = (du,u) -> nl_f(du,u,(10.0,26.0,2.33))

using NLsolve
nlsolve(f2,ones(3))

@derivatives D3'''~t
@derivatives D2''~t
@variables u(t), x(t)
eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
       D2(x) ~ D(x) + 2]
de = ODESystem(eqs, t, [u,x], [])
de1 = ode_order_lowering(de)

de1.eqs

@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t Dx'~x Dy'~y Dz'~z
eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
J = [Dx(eqs[1].rhs) Dy(eqs[1].rhs) Dz(eqs[1].rhs)
 Dx(eqs[2].rhs) Dy(eqs[2].rhs) Dz(eqs[2].rhs)
 Dx(eqs[3].rhs) Dy(eqs[3].rhs) Dz(eqs[3].rhs)]

J = expand_derivatives.(J)

using LinearAlgebra
luJ = lu(J,Val(false))

luJ.L

invJ = inv(luJ)

function lorenz(du,u,p,t)
 du[1] = p[1]*(u[2]-u[1])
 du[2] = u[1]*(p[2]-u[3]) - u[2]
 du[3] = u[1]*u[2] - p[3]*u[3]
end

u = [x,y,z]
du = similar(u)
p = [σ,ρ,β]
lorenz(du,u,p,t)
du

J = [Dx(du[1]) Dy(du[1]) Dz(du[1])
     Dx(du[2]) Dy(du[2]) Dz(du[2])
     Dx(du[3]) Dy(du[3]) Dz(du[3])]
J = expand_derivatives.(J)

using SparseArrays
function SparseArrays.SparseMatrixCSC(M::Matrix{T}) where {T<:ModelingToolkit.Expression}
    idxs = findall(!iszero, M)
    I = [i[1] for i in idxs]
    J = [i[2] for i in idxs]
    V = [M[i] for i in idxs]
    return SparseArrays.sparse(I, J, V, size(M)...)
end
sJ = SparseMatrixCSC(J)

@parameters σ(..)
eqs = [D(x) ~ σ(t-1)*(y-x),
       D(y) ~ x*(σ(t^2)-z)-y,
       D(z) ~ x*y - β*z]

@derivatives Dₜ'~t
Dₜ(x*(σ(t^2)-z)-y)
expand_derivatives(Dₜ(x*(σ(t^2)-z)-y))

_f(x) = 2x + x^2
_f(x)

f(x) = 2x + x^2
@register f(x)

f(x)

function ModelingToolkit.derivative(::typeof(f), args::NTuple{1,Any}, ::Val{1})
    2 + 2args[1]
end
expand_derivatives(Dx(f(x)))

using SciMLTutorials
SciMLTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


"""
    lagrange(x, xd, yd)
adalah fungsi yang digunakan untuk mencari nilai interpolasi pada titik/vektor `x`, jika
diketahui suatu himpunan pasangan terurut `(xd,yd)`.

# Example
```jl
julia> xd = [1,2,3,5];

julia> yd = [1.06 1.12 1.34 1.78];

julia> lagrange(4,xd,yd)
1.6000000000000003

julia> lagrange([2.5,4,5.5],xd,yd)
3-element Array{Float64,1}:
 1.2175
 1.6000000000000003
 1.8025000000000002
```
return solusi hampiran interpolasi `y`.
"""
function lagrange(x, xd, yd)
  #% sedikit trik agar bisa menghitung nilai interpolasi pada banyak
  #% titik sekaligus. Intinya, tetap diproses satu titik demi titik.
  m=length(x);
  y=zeros(m)
  if m > 1
    for i=1:m
      # proses satu titik demi satu titik
      y[i]=lagrange(x[i], xd, yd);
    end
    return y
  end
  #% periksa jumlah titik dan tentukan derajat polinom
  ntitik = length(xd);
  n = ntitik-1;  #% derajat maksimum, sesuai jumlah titik yang ada
  #% hitung L_{n,k}(x) untuk k=1:(n+1)
  Ln=ones(1,ntitik);
  for i=1:ntitik
    for j=1:ntitik
      if i!=j
        Ln[i] = Ln[i] * (x-xd[j])/(xd[i]-xd[j]);;
      end
    end
  end
  #% hitung P_{N}(x)
  y = 0;
  for i=1:ntitik
    y = y + yd[i]*Ln[i];
  end
  return y
end


using ContinuedFractions, Test

@testset "Continued Fraction" begin
  @test sqrt(2) ≈ continuedfraction(i -> i == 0 ? 1 : 2, x->1, 20)
  # https://oeis.org/A003417
  arrayexp1 = [2, 1, 2, 1, 1, 4, 1, 1, 6, 1, 1, 8, 1, 1]
  fexp1(i) = arrayexp1[i + 1]
  @test exp(1) ≈ continuedfraction(fexp1, x->1, length(arrayexp1) - 1)
  # https://oeis.org/A001203
  arrayπ = [3, 7, 15, 1, 292, 1, 1, 1, 2, 1, 3]
  fπ(i) = arrayπ[i + 1]
  @test Float64(π) ≈ continuedfraction(fπ, x->1, length(arrayπ) - 1)

  for i ∈ 1:10
    atol, rtol =10.0^-rand(1:15), 10.0^-rand(1:15)
    answer = 1.61803398874989484820
    output = continuedfraction(i->1, x->1; atol=atol, rtol=rtol)[1]
    @test isapprox(answer, output, atol=atol, rtol=rtol)
    answer = sqrt(2)
    output = continuedfraction(i->i==0 ? 1 : 2, x->1; atol=atol, rtol=rtol)[1]
    @test isapprox(answer, output, atol=atol, rtol=rtol)
  end

end


mutable struct Lattice
    n::Int64 # number of rows in basis matrix
    m::Int64 # number of columns in basis matrix
    basis::Array{Float64} # matrix whose columns form basis of lattice
    μ::Array{Float64} # lower triangular matrix of gram-schmidt coefficients
    r::Array{Float64} # squared norm of gram-schmidt basis vectors
    gs::Array{Float64} # columns are gram-schmidt basis vectors

    """
        Lattice(basis)

    Construct a Lattice generated by the input basis matrix.
    """
    Lattice(basis) = gram_schmidt(basis);
    Lattice(n, m, basis, μ, r, gs) = new(n, m, basis, μ, r, gs);
end

"""
    gram_schmidt(M)

Compute Gram-Schmidt orthogonalization on columns of input matrix.
"""
function gram_schmidt(M::Array{<:Real})
    n = size(M, 1); # number of rows of matrix
    m = size(M, 2); # number of columns of matrix
    μ = zeros(Float64,m,m); # lower triangular matrix of gram-schmidt coefficients
    r = zeros(Float64,m); # squared norm of gram-schmidt basis vectors
    gs = zeros(Float64,n,m); # columns are gram-schmidt basis vectors
    for i = 1:m
        b = M[:,i];
        for j = 1:i-1 # compute gram-schmidt coefficients for j < i
            μ[i,j] = (transpose(b)*gs[:,j])/r[j];
        end
        for j = 1:i-1 # compute i-th gram-schmidt basis vector
            b -= μ[i,j]*gs[:,j];
        end
        gs[:,i] = b;
        r[i] = norm(b)^2; # squared norm of gram-schmidt vector
        μ[i,i] = (transpose(b)*M[:,i])/r[i];
    end
    return Lattice(n, m, M, μ, r, gs);
end

"""
    update_gram_schmidt!(L)

Update Gram-Schmidt coefficients and vectors.
"""
function update_gram_schmidt!(L::Lattice)
    for i = 1:L.m
        b = L.basis[:,i];
        for j = 1:i-1 # compute gram-schmidt coefficients for j < i
            L.μ[i,j] = (transpose(b)*L.gs[:,j])/L.r[j];
        end
        for j = 1:i-1 # compute i-th gram-schmidt basis vector
            b -= L.μ[i,j]*L.gs[:,j];
        end
        L.gs[:,i] = b;
        L.r[i] = norm(b)^2; # squared norm of gram-schmidt vector
        L.μ[i,i] = (transpose(b)*L.basis[:,i])/L.r[i];
    end
    return L;
end

"""
    size_reduce(L)

Compute size reduction of basis vectors of lattice, such that all Gram-Schmidt coefficients have absolute value at most 1/2. Modify the Gram-Schmidt coefficients of input lattice accordingly.
"""
function size_reduce!(L::Lattice)
    for i =2:L.m
        for j = i-1:-1:1
            L.basis[:,i] -= round(L.μ[i,j])*L.basis[:,j];
            for k = 1:j
                L.μ[i,k] -= round(L.μ[i,j])*L.μ[j,k];
            end
        end
    end
    return L
end

"""
    check_size_reduce(L, tol=1e-5)

Check whether basis vectors of input lattice are size-reduced.
"""
function check_size_reduce(L::Lattice, tol::Real=1e-5)
    for i = 1:L.m, j = 1:i-1
        if abs(L.μ[i,j]) > 0.5 + tol
            return false
        end
    end
    return true
end

"""
    lll!(L, δ=0.75)

Compute LLL reduction of input lattice, with parameter δ such that ``0.25 < δ < 1``.
"""
function lll!(L::Lattice, δ::Real=0.75)
    k = 2;
    while k <= L.m
        size_reduce!(L);
        if norm(L.gs[:,k])^2 >= (δ - L.μ[k,k-1]^2)*norm(L.gs[:,k-1])^2
            k += 1;
        else
            b = copy(L.basis[:,k]);
            L.basis[:,k] = copy(L.basis[:,k-1]);
            L.basis[:,k-1] = b;
            update_gram_schmidt!(L);
            k = max(2, k-1);
        end
    end
    return L;
end

"""
    lll(M, δ=0.75)

Compute LLL reduction of input matrix, with parameter δ such that ``0.25 < δ < 1``.
"""
function lll(M::Array{<:Real}, δ::Real=0.75)
    return lll(Lattice(M)).basis;
end

"""
    check_lll(L, δ=0.75)

Check whether input lattice has an LLL-reduced basis, with parameter δ such that ``0.25 < δ < 1``.
"""
function check_lll(L::Lattice, δ::Real=0.75)
    if !check_size_reduce(L)
        return false;
    end
    for i=1:L.m-1
        if L.r[i+1] < (δ-L.μ[i+1,i]^2)*L.r[i]
            return false;
        end
    end
    return true;
end

"""
    svp(M)

Compute shortest vector in lattice generated by M, using Schnorr-Euchner enumeration.
"""
function svp(M::Array{<:Real})
    L = Lattice(M);
    x = zeros(Int64, L.m);
    x[L.m] = 1;
    A = norm(M[:,L.m])^2
    sol, length =  enum_svp(L, x, L.m, A, 0.0);
    vect = sum([sol[j]*L.basis[:,j] for j=1:L.m]);
    return [vect, length, sol];
end

"""
    enum_svp(L, x, i, radius, l)

Enumerate all lattice points in subtree at height i, where parent nodes are defined by x. Helper method for [`svp`](@ref).
"""
function enum_svp(L::Lattice, x::Array{<:Integer}, i::Integer, radius::Real, l::Real)
    sol = copy(x);
    c = -sum(Float64[x[j]*L.μ[j,i] for j=i+1:L.m]);
    interval = sqrt((radius-l)/L.r[i]); # interval of possibilities for x[i]
    k = floor(Int64, c); # stores value of x[i]
    Δk = 0; # used to do search in zig-zag path
    Δ²k = -1; # used to do search in zig-zag path
    lbound = floor(Int64, c - interval);
    ubound = ceil(Int64, c + interval);
    while k>= lbound & k <= ubound
        l_new = l + (k-c)^2*L.r[i];
        if l_new <= radius # possible new solution found
            if i == 1 & ~(l_new == 0) # leaf of enumeration tree
                # l_new is approximation of length of new solution. Do exact computation of length to avoid errors
                old_sol = sol[i];
                sol[i] = k;
                new_length = norm(sum([sol[j]*L.basis[:,j] for j=1:L.m]))^2;
                if radius > new_length # prune enumeration tree
                    radius = new_length;
                else
                    sol[i] = old_sol;
                end
            elseif i > 1 # solve subtree enumeration recursively
                old_sol = sol[i];
                sol[i] = k;
                subtree_sol, subtree_length = enum_svp(L, sol, i-1, radius, l_new);
                if subtree_length < radius
                    sol[1:i-1] = subtree_sol[1:i-1];
                    radius = subtree_length;
                else
                    sol[i] = old_sol;
                end
            end
        end
        Δ²k = - Δ²k;
        Δk = Δ²k - Δk;
        k += Δk;
    end
    return [sol, radius];
end


using Printf

function sternbrocot(f::Function=(x) -> length(x) ≥ 20)::Vector{Int}
    rst = Int[1, 1]
    i = 3
    while !f(rst)
        append!(rst, Int[rst[i-1] + rst[i-2], rst[i-2]])
        i += 1
    end
    return rst
end

println("First 15 elements of Stern-Brocot series:\n", sternbrocot(x -> length(x) ≥ 15)[1:15], "\n")

for i in (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 100)
    occurr = findfirst(x -> x == i, sternbrocot(x -> i ∈ x))
    @printf("Index of first occurrence of %3i in the series: %4i\n", i, occurr)
end

print("\nAssertion: the greatest common divisor of all the two\nconsecutive members of the series up to the 1000th member, is always one: ")
sb = sternbrocot(x -> length(x) > 1000)
if all(gcd(prev, this) == 1 for (prev, this) in zip(sb[1:1000], sb[2:1000]))
    println("Confirmed.")
else
    println("Rejected.")
end


using Test
using NumericallyIntegrateArrays

@testset "simps" begin
    x = range(0,stop=π,length=10); dx = step(x)

    @testset "dx" begin
	    @testset "1D" begin
		    f = sin.(x)
		    @test isapprox(simps(f,step(x)),2,rtol=1e-3)
	    end
	    @testset "2D" begin
	        f = sin.(x) * cos.(x)'
	       	@test isapprox(simps(f,step(x)),2cos.(x),rtol=1e-3)
	    end
    end
    @testset "x array" begin
	    @testset "1D" begin
		    f = sin.(x)
		    @test isapprox(simps(f,x),2,rtol=1e-3)
	    end
	    @testset "2D" begin
	        f = sin.(x) * cos.(x)'
	       	@test isapprox(simps(f,x),2cos.(x),rtol=1e-3)
	    end
    end
end

@testset "trapz" begin
    x = range(0,stop=π,length=10); dx = step(x)

    @testset "dx" begin
	    @testset "1D" begin
		    f = sin.(x)
		    @test isapprox(trapz(f,step(x)),2,rtol=1e-1)
	    end
	    @testset "2D" begin
	        f = sin.(x) * cos.(x)'
	       	@test isapprox(trapz(f,step(x)),2cos.(x),rtol=1e-1)
	    end
    end
    @testset "x array" begin
	    @testset "1D" begin
		    f = sin.(x)
		    @test isapprox(trapz(f,x),2,rtol=1e-1)
	    end
	    @testset "2D" begin
	        f = sin.(x) * cos.(x)'
	       	@test isapprox(trapz(f,x),2cos.(x),rtol=1e-1)
	    end
    end
end


function hammingsequence(N)
    if N < 1
        throw("Hamming sequence exponent must be a positive integer")
    end
    ham = N > 4000 ? Vector{BigInt}([1]) : Vector{Int}([1])
    base2, base3, base5 = (1, 1, 1)
    for i in 1:N-1
        x = min(2ham[base2], 3ham[base3], 5ham[base5])
        push!(ham, x)
        if 2ham[base2] <= x
            base2 += 1
        end
        if 3ham[base3] <= x
            base3 += 1
        end
        if 5ham[base5] <= x
            base5 += 1
        end
    end
    ham
end

println(hammingsequence(20))
println(hammingsequence(1691)[end])
println(hammingsequence(1000000)[end])


#=
   Meeting Under the clock (Julian Simon 1994)

   """
   Meeting Under the Clock (This problem is posed by Julian Simon(1994))

   Two persons agree to arrive at the two clock sometime between 1 pm and 2 pm
   and to stay for 20 minutes. What is the probability that they will be there
   at the same time?
   """

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

          c1   30.1994   17.2942     0.1729    0.1723    9331.2798    0.9999     7267.3518
          c2   30.1059   17.3087     0.1731    0.1659   10124.1004    1.0000     7884.8134
        prob    0.5734    0.4946     0.0049    0.0054    9476.5095    1.0000     7380.4591

    According to the WebPPL model meeting_under_the_clock.wppl
    the exact probability is 0.5742602700373461.

    Cf ~/webppl/meeting_under_the_clock.wppl

=#
using Turing, StatsPlots, Distributions
include("jl_utils.jl")

@model function meeting_under_the_clock(w=20)
    c1 ~ DiscreteUniform(1,60)
    c2 ~ DiscreteUniform(1,60)
    # meeting within waiting time
    prob ~ Dirac(abs(c1-c2) <= w)
end

w = 20
model = meeting_under_the_clock(w)

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

display(chns)
# display(plot(chns))


"""
    `RBF_gauss(R;A=1.0,sigma=0.2,deriv=0)`

Kernel function for Gaussian RBF Interpolation. Returns the magnitude of evaluated
function.

The optional argument `deriv` is for the order of the derivative desired, meaning
0 for base, 1 for first derivative (gradient) and 2 for second derivative (Laplacian).

# ARGUMENTS
* `R::Float64` : Magnitude of radius vector -> R = |X_j - X_i|
* `A=1.0`      : Normalization value
* `sigma=0.2`  : Gaussian spreading
* `deriv=0`    : Order of desired derivative

# OUTPUTS
* `res`        : Calculated magnitude
"""
function RBF_gauss(R::Float64;
                   A=1.0,sigma=0.2,deriv=0)

    if deriv == 0
        # const1 = 1/(2*pi)^(3/2)
        # res = const1*exp(-R^2/2)

        res = SVPM.zeta_gauserf(R)
        #res = A*exp(-(R^2)/(2*sigma^2))
    elseif deriv == 1
        res = -((A*R)/(sigma^2))*exp(-(R^2)/(2*sigma^2))
    elseif deriv == 2
        res = ((A/sigma)^2)*((R/sigma)^2 - 3)*exp(-(R^2)/(2*sigma^2))
    end

    return res
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    ZScore()

The z-score (a.k.a. normal score) of `x` with mean `μ` and
standard deviation `σ` is the value `(x .- μ) ./ σ`.
"""
struct ZScore <: Colwise end

assertions(::Type{ZScore}) = [assert_continuous]

isrevertible(::Type{ZScore}) = true

function colcache(::ZScore, x)
  μ = mean(x)
  σ = std(x, mean=μ)
  (μ=μ, σ=σ)
end

colapply(::ZScore, x, c)  = @. (x - c.μ) / c.σ

colrevert(::ZScore, y, c) = @. c.σ * y + c.μ


# by LJR
using Plots
using LaTeXStrings

fileName = ARGS

function plotFile(file::String; method::Symbol=:no, order=5)
    fit     = true
    if method == :linear
        M = 1
    elseif method == :poly
        M = 2
    elseif method == :exp
        M = 3
    elseif method == :log
        M = 4
    elseif method == :no
        fit = false
    else
        error("Method not implemented")
    end
    # Get the filename without extention name like .csv
    fileRegex = match(r"(.+?)(\.[^.]*$|$)", file)
    csv = readdlm(file, ',', Float64)
    y = csv[:, 1]; x = csv[:, 2]
    # Extent plot domain by 20%
    x_max = maximum(x)
    x_min = 0
    @show    xlen = x_max-x_min
    x_plot = 0:xlen/100:(x_max+xlen/5)
    # Plot data
    scatter(x, y, fmt = :pdf, ylabel = "Voltage (V)", xlabel = "Current (A)", label="Data")
    # Plot the best fit line
    if fit
        plot!(x_plot, dataFit(x, y, x_plot, Val{M}, order), label="Best fit line")
        savefig(fileRegex[1]*"_$method"*".pdf")
    end
    savefig(fileRegex[1]*".pdf")
end

# Dispatch into the right argument
dataFit(x, y, x_plot, ::Type{Val{1}}, kwarg...) = dataFit(x, y, x_plot, Val{1})
dataFit(x, y, x_plot, ::Type{Val{3}}, kwarg...) = dataFit(x, y, x_plot, Val{3})
dataFit(x, y, x_plot, ::Type{Val{4}}, kwarg...) = dataFit(x, y, x_plot, Val{4})

# Linear fit by using linreg
function dataFit(x, y, x_plot, ::Type{Val{1}})
    b, a = linreg(x, y)
    map(x -> a*x + b, x_plot)
end

# Polynomial fit by QR factorization
function dataFit(x, y, x_plot, ::Type{Val{2}}, order)
    function horner(coef,x)
        sum = coef[end]
        for k = length(coef)-1:-1:1
            sum = coef[k] + x*sum
        end
        sum
    end
  A = [ float(x[i])^pow for i = 1:length(x), pow = 0:order ]
  co = A \ y
  map(x -> horner(co, x), x_plot)
end

function dataFit(x, y, x_plot, ::Type{Val{3}})
    fit = linreg(x, log(y))
    a1 = exp(fit[1]); a2 = fit[2]
    map(x -> a1*exp(x*a2), x_plot)
end

function dataFit(x, y, x_plot, ::Type{Val{4}})
    a1, a2 = linreg(log(x), y)
    map(x -> a1 + a2*log(x), x_plot)
end

plotFile.(fileName)


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

The radical of n, rad(n), is the product of the distinct prime factors of n. For
example, 504 = 2^3 × 3^2 × 7, so rad(504) = 2 × 3 × 7 = 42.

If we calculate rad(n) for 1 ≤ n ≤ 10, then sort them on rad(n), and sorting on
n if the radical values are equal, we get:

 Unsorted |   Sorted  |
 n rad(n) |  n rad(n) |  k
 ---------+-----------+---
 1   1    |  1   1    |  1
 2   2    |  2   2    |  2
 3   3    |  4   2    |  3
 4   2    |  8   2    |  4
 5   5    |  3   3    |  5
 6   6    |  9   3    |  6
 7   7    |  5   5    |  7
 8   2    |  6   6    |  8
 9   3    |  7   7    |  9
10  10    | 10  10    | 10

Let E(k) be the kth element in the sorted n column; for example, E(4) = 8 and
E(6) = 9.

If rad(n) is sorted for 1 ≤ n ≤ 100000, find E(10000).

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

using Primes
using Base.Test

function compute_table(limit)
    res = Tuple{Int,Int}[]
    for n in 1:limit
        push!(res, (prod(keys(factor(n))), n))
    end

    sort!(res)
    map(pair -> pair[2], res)
end

test_table = compute_table(10)

@test test_table[4] == 8
@test test_table[6] == 9

actual_table = compute_table(100_000)

println("Euler 124: $(actual_table[10_000])")


import RowEchelon
import LinearAlgebra

#=
Homework 7
Section 3.1
Problems 11
=#

A = [3 5 -8 4 ;
     0 -2 3 -7 ;
     0 0 1 5 ;
     0 0 0 2]

detA = LinearAlgebra.det(A)

println("Original Matrix:")
println(A)
println("Determinate:")
println(detA)


using ModelingToolkit, NonlinearSolve

@parameters t
@variables u1(t) u2(t) u3(t) u4(t) u5(t)

eqs = [
        0 ~ u1 - sin(u5),
        0 ~ u2 - cos(u1),
        0 ~ u3 - hypot(u1, u2),
        0 ~ u4 - hypot(u2, u3),
        0 ~ u5 - hypot(u4, u1)
]

sys = NonlinearSystem(eqs, [u1, u2, u3, u4, u5], [])

simple_sys = structural_simplify(sys)

equations(simple_sys)

prob = NonlinearProblem(simple_sys, [u5=>0.0])
sol = solve(prob, NewtonRaphson())

sol[u5]
sol[u1]

# stochastic differential equations
using ModelingToolkit, DifferentialEquations

@parameters t σ ρ β
@variables x(t) y(t) z(t)
D = Differential(t)

eqs = [
        D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z
       ]

noise_eqs = [0.1*x, 0.1*y, 0.1*z]

de = SDESystem(eqs, noise_eqs, t, [x, y, z], [σ, ρ, β])

u0map = [x=> 1.0, y=> 0.0, z=>0.0]
parammap = [σ => 10.0, ρ => 28.0, β => 8/3]

prob = SDEProblem(de, u0map, (0.0, 100.0), parammap)

sol = solve(prob)
plot(sol, vars=(x,y,z))

# ------------ optimization problem ----------------------
using ModelingToolkit, GalacticOptim, Optim

@variables x y
@parameters a b
loss = (a-x)^2 + b * (y -x^2)^2

sys = OptimizationSystem(loss, [x,y], [a,b])

# initial guess
u0 = [x=>1.0, y=>2.0]
p = [a => 6.0, b=> 7.0]

prob = OptimizationProblem(sys, u0, grad=true, hess=true)
sol = solve(prob) # doesn't work yet

# ------------ dsl demo via Catalyst ---------------------
using DifferentialEquations, Catalyst, Latexify, ModelingToolkit, Plots

repressilator = @reaction_network begin
    hillr(P₃,α,K,n), ∅ --> m₁
    hillr(P₁,α,K,n), ∅ --> m₂
    hillr(P₂,α,K,n), ∅ --> m₃
    (δ,γ), m₁ ↔ ∅
    (δ,γ), m₂ ↔ ∅
    (δ,γ), m₃ ↔ ∅
    β, m₁ --> m₁ + P₁
    β, m₂ --> m₂ + P₂
    β, m₃ --> m₃ + P₃
    μ, P₁ --> ∅
    μ, P₂ --> ∅
    μ, P₃ --> ∅
end α K n δ γ β μ;


using Plots
gr()

################################
### Algoritmo Newton-Raphson ###
################################

function NewtonRaphson(f, fd, x; tol = 1e-8, max_iter=1000, max_tempo=10.0)
    fx = f(x)               # primeira iteracao
    i = 0                   # contador de iteracoes
    t0 = time()             # tempo de execucao
    t1 = time() - t0
    estado = "Desconhecido" # estado do algoritmo

    # define erro absoluto e erro relativo e calcula o erro
    atol = rtol = tol
    erro = atol + rtol*abs(fx)

    # regras de parada
    rp1 = (abs(fx) ≤ erro)
    rp2 = (t1 ≥ max_tempo || i ≥ max_iter)

    while !(rp1 || rp2)
        fdx = fd(x)

        x = x - fx/fdx
        if abs(fdx) ≤ erro
            estado = "Falha na convergência. Derivada nula encontrada."
            break
        end
        fx = f(x)

        i += 1
        t1 = time() - t0

        rp1 = (abs(fx) ≤ erro)
        rp2 = (t1 ≥ max_tempo || i ≥ max_iter)
    end

    if rp1
        estado = "Convergência Obtida."
    elseif rp2
        i ≥ max_iter ? estado = "Falha na convergência. O Algoritmo atingiu iteração limite" : estado = "Falha na convergência. O algoritmo atingiu tempo limite"
    end
    println("\n",estado,"\n")
    println("Raiz em $x com $i iterações.\n")
    return x, fx, estado
end

Plots.plot(x -> x^2 - 2, -5, 5, leg=false)
f(x) = x^2 - 2
fd(x) = 2x
NewtonRaphson(f, fd, 500.0, tol = 1e-20)
NewtonRaphson(f, fd, -100.0)

f(x) = x * exp(x) - 1
fd(x) = exp(x) + x * exp(x)
Plots.plot(f, -2, 5, leg=false)
NewtonRaphson(f, fd, 0.0)

f(x) = x^2 + 3
fd(x) = 2x
Plots.plot(f, -2, 2, leg=false)
NewtonRaphson(f, fd, 10)

##################################################################
### Newton-Raphson aproximando derivada com diferenças finitas ###
##################################################################

function newton_fin_dif(f, x; h = 0.1, tol = 1e-10, N = 1000)
    c = 0
    for i = 1:N
        c += 1
        xnew = x - (2*h*f(x)) / (f(x+h) - f(x-h))
        if abs(xnew - x) <= tol
            print("root at $x with $(c) iterations.\nError: $(round(xnew-x, sigdigits=5))\n")
            return xnew, c
        end
        x = xnew
    end
    print("WARNING: Function did NOT converge.\nRoot at $x with $(c) iterations.\nError: $(round(xnew-x, sigdigits=5))\n")
    return xnew, c
end

f(x) = 2x^2 - 5x
plot(f, -2, 5, leg=false)
newton_fin_dif(f, 20)

newton_fin_dif(x -> 2x^2 - 5x, 100);
newton_fin_dif(x -> 2x^2 - 5x, -0.5)
newton_fin_dif(x -> 2x^2 - 5x, -500);

plot(x -> x^3 - 2x - 5, -2, 5, leg=false)
newton_fin_dif(x -> x^3 - 2x - 5, 5; h = 0.01, tol=0.0001);

##############################################
### Newton-Raphson Com multiplas dimensões ###
##############################################

function newton_mult_dim(f, G, x; alpha = 0.1, tol = 1e-10, N = 1000)
    c = 0
    for i = 1:N
        c += 1
        xnew = x - alpha*(f(x)./G(x))
        if abs(f(xnew)) <= tol
            print("roots at $xnew with $(c) iterations.\nError: $(round.(xnew-x, sigdigits=5))\n")
            return xnew, c
        end
        x = xnew
    end
    print("WARNING: Function did NOT converge.\nRoots at $x with $(c) iterations.\nError: $(round.(xnew-x, sigdigits=5))\n")
    return xnew, c
end

a1 = range(-20,20,step=0.1)
a2 = range(-20,20,step=0.1)
zplot(x,y) = x^2 + y^2 - 5^2
Plots.surface(a1, a2, zplot)

z(x) = x[1]^2 + x[2]^2 - 5^2
G(x) = [2x[1], 2x[2]]
x0 = [3.0; 5.0]
newton_mult_dim(z, G, x0)

#############################################
### Newton-Raphson usando pacote Roots.jl ###
#############################################

using Roots
f(x) = x^3 - 2x - 5
fp(x) = 3x^2 - 2
newton_fin_dif(f, 2);
NewtonRaphson(f, fp, 10)
x = Roots.newton(f, fp, 2, verbose=true)
x = find_zero((f, fp), 2, Roots.Newton(), verbose=true)
x, f(x)


function chebyshev_polynomial_derivative(order::S,x::T) where {T<:AbstractFloat,S<:Integer}

  polynomial = ones(T,1,order+1)
  poly_deriv = zeros(T,1,order+1)

  for i = 2:order+1
    if i == 2
      polynomial[i] = x
      poly_deriv[i] = 1.0
    else
      polynomial[i] = 2*x*polynomial[i-1]-polynomial[i-2]
      poly_deriv[i] = ((i-1)*polynomial[i-1]-(i-1)*x*polynomial[i])/(1-x^2)
    end
  end

  return poly_deriv

end

function chebyshev_polynomial_derivative(order::S,x::Array{T,1}) where {T<:AbstractFloat,S<:Integer}

  polynomial = ones(T,length(x),order+1)
  poly_deriv = zeros(T,length(x),order+1)

  for i = 1:length(x)
    poly_deriv[i,:] = chebyshev_polynomial_derivative(order,x[i])
  end

#=

  for i = 2:order+1
	for j = 1:length(x)
      if i == 2
        polynomial[j,i] = x[j]
        poly_deriv[j,i] = 1.0
      else
        polynomial[j,i] = 2*x[j]*polynomial[j,i-1]-polynomial[j,i-2]
        poly_deriv[j,i] = ((i-1)*polynomial[j,i-1]-(i-1)*x[j]*polynomial[j,i])/(1-x[j]^2)
      end
    end
  end

=#

  return poly_deriv

end


module CalcDist

import StatsFuns: normpdf, normcdf, norminvcdf,
                  tdistpdf, tdistcdf, tdistinvcdf,
                  chisqpdf, chisqcdf, chisqinvcdf,
                  fdistpdf, fdistcdf, fdistinvcdf

export normalpdf, normalcdf, invNorm,
       tpdf, tcdf, invT,
       X2pdf, X2cdf, invX2,
       Fpdf, Fcdf, invF,
       binompdf, binomcdf,
       poissonpdf, poissoncdf,
       geometpdf, geometcdf,
       nCr, nPr

##
# Continuous
##

"""
   normalpdf(x, μ = 0, σ = 1) \n
      x: value
      μ: mean
      σ: standard deviation
   probability density function (Normal)
"""
function normalpdf(x, μ = 0, σ = 1)
   return normpdf(μ, σ, x)
end

"""
   normalcdf(lower, upper, μ = 0, σ = 1) \n
      lower: lower boundary
      upper: upper boundary
          μ: mean
          σ: standard deviation
   cumulative distribution function (Normal)
"""
function normalcdf(lower, upper, μ = 0, σ = 1)
   return normcdf(μ, σ, upper) - normcdf(μ, σ, lower)
end

"""
   invNorm(area, μ = 0, σ = 1) \n
      area: probability area
         μ: mean
         σ: standard deviation
   inverse cumulative distribution function (Normal)
"""
function invNorm(area, μ = 0, σ = 1)
   return norminvcdf(μ, σ, area)
end

"""
   tpdf(x, ν) \n
      x: value
      ν: degrees of freedom
   probability density function (Student's t)
"""
function tpdf(x, ν)
   return tdistpdf(ν, x)
end

"""
   tcdf(lower, upper, ν) \n
      lower: lower boundary
      upper: upper boundary
          ν: degrees of freedom
   cumulative distribution function (Student's t)
"""
function tcdf(lower, upper, ν)
   return tdistcdf(ν, upper) - tdistcdf(ν, lower)
end

"""
   invT(area, ν) \n
      area: probability area
         ν: degrees of freedom
   inverse cumulative distribution function (Student's t)
"""
function invT(area, ν)
   return tdistinvcdf(ν, area)
end

"""
   X2pdf(x, ν) \n
      x: value
      ν: degrees of freedom
   probability density function (Chi-squared)
"""
function X2pdf(x, ν)
   return chisqpdf(ν, x)
end

"""
   X2cdf(lower, upper, ν) \n
      lower: lower boundary
      upper: upper boundary
          ν: degrees of freedom
   cumulative distribution function (Chi-squared)
"""
function X2cdf(lower, upper, ν)
   return chisqcdf(ν, upper) - chisqcdf(ν, lower)
end

"""
   invX2(area, ν) \n
      area: probability area
         ν: degrees of freedom
   inverse cumulative distribution function (Chi-squared)
"""
function invX2(area, ν)
   return chisqinvcdf(ν, area)
end

"""
   Fpdf(x, nν, dν) \n
       x: value
      nν: numerator degrees of freedom
      dν: denominator degrees of freedom
   probability density function (F)
"""
function Fpdf(x, nν, dν)
   return fdistpdf(nν, dν, x)
end

"""
   Fcdf(lower, upper, nν, dν) \n
      lower: lower boundary
      upper: upper boundary
         nν: numerator degrees of freedom
         dν: denominator degrees of freedom
   cumulative distribution function (F)
"""
function Fcdf(lower, upper, nν, dν)
   return fdistcdf(nν, dν, upper) - fdistcdf(nν, dν, lower)
end

"""
   invF(area, nν, dν) \n
      area: probability area
        nν: numerator degrees of freedom
        dν: denominator degrees of freedom
   inverse cumulative distribution function (F)
"""
function invF(area, nν, dν)
   return fdistinvcdf(nν, dν, area)
end

##
# Discrete
##

"""
   binompdf(n, p, x = nothing) \n
      n: trials
      p: probability of success
      x: value
   probability mass function (Binomial)
"""
function binompdf(n, p, x = nothing)
   if n ≤ 0
      throw(DomainError(n, "\t[n > 0] must be true."))
   end
   if p < 0 || p > 1
      throw(DomainError(p, "\t[0 ≤ p ≤ 1] must be true."))
   end

   if isnothing(x)
      x = range(0, stop=n)
   end

   probs = Vector{Real}()

   for i in x
      push!(probs, binomial(n, i) * p^i * (1-p)^(n-i))
   end

   return probs
end

"""
   binomcdf(n, p, x = nothing) \n
      n: trials
      p: probability of success
      x: value
   cumulative distribution function (Binomial)
"""
function binomcdf(n, p, x = nothing)
   if n ≤ 0
      throw(DomainError(n, "\t[n > 0] must be true."))
   end
   if p < 0 || p > 1
      throw(DomainError(p, "\t[0 ≤ p ≤ 1] must be true."))
   end

   if isnothing(x)
      x = range(0, stop=n)
   end

   probs = Vector{Real}()

   for i in x
      push!(probs, sum([ binomial(n, i) * p^i * (1-p)^(n-i) for i in 0:i ]))
   end

   return probs
end

"""
   poissonpdf(λ, x) \n
      λ: mean, variance
      x: value
   probability mass function (Poisson)
"""
function poissonpdf(λ, x)
   if λ ≤ 0
      throw(DomainError(λ, "\t[λ > 0] must be true."))
   end

   probs = Vector{Real}()

   for i in x
      push!(probs, ℯ^(-λ) * (λ^i / factorial(i)))
   end

   return probs
end

"""
   poissoncdf(λ, x) \n
      λ: mean, variance
      x: value
   cumulative distribution function (Poisson)
"""
function poissoncdf(λ, x)
   if λ ≤ 0
      throw(DomainError(λ, "\t[λ > 0] must be true."))
   end

   probs = Vector{Real}()

   for i in x
      push!(probs, ℯ^(-λ) * sum([ λ^i / factorial(i) for i in 0:i ]))
   end

   return probs
end

"""
   geometpdf(p, x) \n
      p: probability of success
      x: trial on which 1st success occurs
   probability density function (Geometric)
"""
function geometpdf(p, x)
   if p < 0 || p > 1
      throw(DomainError(p, "\t[0 ≤ p ≤ 1] must be true."))
   end

   probs = Vector{Real}()

   for i in x
      push!(probs, p * (1-p)^(i-1))
   end

   return probs
end

"""
   geometcdf(p, x) \n
      p: probability of success
      x: trial on which 1st success occurs
   cumulative distribution function (Geometric)
"""
function geometcdf(p, x)
   if p < 0 || p > 1
      throw(DomainError(p, "\t[0 ≤ p ≤ 1] must be true."))
   end

   probs = Vector{Real}()

   for i in x
      push!(probs, 1 - (1-p)^i)
   end

   return probs
end

##
# Counting
##

"""
   nCr(n, r) \n
      return number of combinations of n items taken r at a time. n and r must be
      non-negative. both n and r can be vectors.
"""
function nCr(n, r)
   if n < 0
      throw(DomainError(n, "\t[n ≥ 0] must be true."))
   end

   combs = Vector{Real}()

   for i in r
      if i < 0
         throw(DomainError(i, "\t[r ≥ 0] must be true."))
      end
      push!(combs, factorial(n) / ( factorial(i) * factorial(n-i) ))
   end

   return combs
end

function nCr(n::Vector{Int}, r)
   if r < 0
      throw(DomainError(r, "\t[r ≥ 0] must be true."))
   end

   combs = Vector{Real}()

   for i in n
      if i < 0
         throw(DomainError(i, "\t[n ≥ 0] must be true."))
      end
      push!(combs, factorial(i) / ( factorial(r) * factorial(i-r) ))
   end

   return combs
end

"""
   nPr(n, r)
      return number of permutations of n items taken r at a time. n and r must be
      non-negative. both n and r can be vectors.
"""
function nPr(n, r)
   if n < 0
      throw(DomainError(n, "\t[n ≥ 0] must be true."))
   end

   perms = Vector{Real}()

   for i in r
      if i < 0
         throw(DomainError(i, "\t[r ≥ 0] must be true."))
      end
      push!(perms, factorial(n) / factorial(n-i))
   end

   return perms
end

function nPr(n::Vector{Int}, r)
   if r < 0
      throw(DomainError(r, "\t[r ≥ 0] must be true."))
   end

   perms = Vector{Real}()

   for i in n
      if i < 0
         throw(DomainError(i, "\t[n ≥ 0] must be true."))
      end
      push!(perms, factorial(i) / factorial(i-r))
   end

   return perms
end

end # module


f = x -> begin
	N = lastindex(x)
    return 100sum((x[i+1] - x[i]^2)^2 for i = 1:N-1) + sum((x[i] - 1.0)^2 for i = 1:N-1)
end

g! = (g, x) -> begin
	N = lastindex(x)
	g[1] = -2*(1 - x[1]) - 400*x[1]*(-x[1]^2 + x[2])
	for i in 2:N-1
		g[i] = -2*(1 - x[i]) + 200*(-x[i-1]^2 + x[i]) - 400*x[i]*(-x[i]^2 + x[1 + i])
	end
	g[N] = 200 * (x[N] - x[N-1]^2)
	return g
end

fg! = (g, x) -> begin
	@warn "rosen fg! not implmented"
end

init = (n::Int=500) -> begin
	x0 = [j/(n+1) for j in 1:n]
	return n, x0
end

TestSet["rosen"] = UncProgram("rosen",  f, g!, fg!, init)


"""
`SymmetricDoubleWell` - The classic symmetric doublewell potential.

### Fields
* `x` - Position x
"""
function SymmetricDoubleWell(x)
    return (x^2 -1)^2;
end

"""
`AsymmetricDoubleWell` - An asymmetric double well potential requires |δ|< 1 for
multiple minima.

### Fields
* `x` - Position x
### Optional Fields
* `δ = 0.5` - Asymmetry parameter
"""
function AsymmetricDoubleWell(x; δ = 0.5)
    return x^4 - 1.5 * x^2 - δ * x;
end

"""
`FatSkinnyDoubleWell10` - The fat-skinny double well of degree 10. This introduces
entropic effects in dimension 1.

### Fields
* `x` - Position x
"""
function FatSkinnyDoubleWell10(x)
    return ((8-5*x)^8 * (2+5*x)^2)/(2^26);
end


"""
    fft_freqs(ts)

    Returns an array of DFT frequency points associated with the set of
    time-domain points.

    Assumes unit of argument `x` is sec. Output unit is in Hz (not rad/sec).
"""
function fft_freqs(ts)
    fs = (ts[2]-ts[1])^(-1) # sampling frequency
    N = Int(length(ts))
    deltaf = fs/N
    if N%2 ==0
        freqs = deltaf*(-N/2:N/2-1)
    else
        freqs = deltaf*(-(N-1)/2:(N-1)/2)
    end
    if length(ts) !== length(freqs)
        @printf("length(ts): %d, length(freqs): %d, delta_f: %f.\n",
                length(ts), length(freqs), deltaf)
        error("Number of time points should match the number of frequency
              points.")
    end
    return freqs
end


using LinearAlgebra,Statistics
include("plotshelper.jl")
x_vals = [0 0 0 ; 2 -3 -4]
y_vals = [0 0 0; 4 3 -3.50]
plt(x_vals,y_vals)
vals = [ 0 0 0; 2* [-2 1 2]]
plt(vals,vals)

f(x,y) = 0.2x + 0.1y

x_vec = [0 0; 3 3]
y_vec = [0 0; 4 -4]
z_vec = [0 0; f(x_vec[2,1],y_vec[2,1]) f(x_vec[2,2],y_vec[2,2])]

n=20
grid = range(-5,5,length=n)
z2 = [ f(grid[row],grid[col]) for row in 1:n, col in 1:n]
wireframe(grid,grid,z2,fill=:blue,gridalpha=1)
plot!(x_vec,y_vec,z_vec,color=[:blue,:green],linewidth=1,labels="", colorbar=false)

# linear equation

f(x) = 0.6cos(4x)+ 1.3
grid = range(-2,2,length=100)
y_min,y_max = extrema(f(x) for x in grid)
plt1 = plot(f,xlim=(-2,2),lable="f")
hline!(plt1,[f(0.5)],linestyle=:dot,linewidth=2, label="h line")
vline!(plt1,[-1.07, -0.5, 0.5, 1.07],linestyle=:dot,linewidth=2, label="v line")
plot!(plt1,fill(0,2),[y_min y_min;y_max y_max],lw=6,color=:blue,label="tange of f")

#eigen values and eigen vectors

A = [1 2;2 1]
evals, evacs = eigen(A)
a1, a2 = evals
eig_1 = [0 0;evacs[:,1]']


using PyPlot
A,B,C=1,2,3 # give variables a number
Nlevels=10

function demoSkill()
    A,B,C=1,2,3 # give variables a number
    Nlevels=10

    sigmoid(x)=1./(1+exp(-x))

    pLeftPlayerWins=zeros(Nlevels,Nlevels)
    for levelleft=1:10
        for levelright=1:10
            pLeftPlayerWins[levelleft,levelright]=sigmoid(levelleft-levelright)
        end
    end

    #prior_pot=exp.(-0.05*(Nlevels/2-linspace(0,Nlevels,Nlevels)).^2) # peaked around skill 5
    prior_pot=ones(Nlevels,1) # uniform prior on skill
    prior_pot=prior_pot./sum(prior_pot)
    priorA=PotArray(A,prior_pot)
    priorB=PotArray(B,prior_pot)
    priorC=PotArray(C,prior_pot)

    fAbeatsB=PotArray([A B],pLeftPlayerWins)
    fBbeatsC=PotArray([B C],pLeftPlayerWins)
    fAbeatsC=PotArray([A C],pLeftPlayerWins)

    Data = "A>B, B>C"
    pABC=fAbeatsB*fBbeatsC*priorA*priorB*priorC # p(A,B,C,Data)
    plotstuff(Data,pABC,fAbeatsC)

    Data = "A>B, A>B"
    pABC=fAbeatsB*fAbeatsB*priorA*priorB*priorC # p(A,B,C,Data)
    plotstuff(Data,pABC,fAbeatsC)

    Data = "(A>B)^6, (B>C)^6"
    pABC=fAbeatsB^6*fBbeatsC^6*priorA*priorB*priorC # p(A,B,C,Data)
    plotstuff(Data,pABC,fAbeatsC)

    Data = "A>B, C>B"
    pABC=fAbeatsB*(1-fBbeatsC)*priorA*priorB*priorC # p(A,B,C,Data)
    plotstuff(Data,pABC,fAbeatsC)

end

function plotstuff(Data,pABC,fAbeatsC)
    pABC=pABC./sum(pABC) # normalise to get p(A,B,C|Data)
    pAC=sum(pABC,B)
    pAbeatsC=sum(fAbeatsC*pAC)
    # What is the probability that A will beat C?
    println(Data*" : probability that A beats C = $(pAbeatsC.content)")
    pA=sum(pABC,[B C]); pB=sum(pABC,[A C]); pC=sum(pABC,[B A])
    figure(); title(Data)
    plot(pA.content,label="A skill")
    plot(pB.content,label="B skill")
    plot(pC.content,label="C skill")
    legend()
    # Can calculate mean skill levels using:
    println("mean skill for A=",sum(pA.content.*(1:Nlevels))) # mean skill level for A
    println("mean skill for B=",sum(pB.content.*(1:Nlevels))) # mean skill level for B
    println("mean skill for C=",sum(pC.content.*(1:Nlevels))) # mean skill level for C
end


### A Pluto.jl notebook ###
# v0.12.17

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ ebcb6ec6-61f9-11eb-05a3-817618350415
using PlutoUI

# ╔═╡ 2d1693e6-61fa-11eb-3cb6-9f5431449a41
begin
	using Plots
	pyplot()
end

# ╔═╡ f5d3f88c-620e-11eb-3bc8-e3e0e1598482
using Random

# ╔═╡ 6b69c792-62bd-11eb-0f1d-2d0081598dfa
using Distributions

# ╔═╡ b80b8e58-62bf-11eb-11fb-1397b06ae3e1
using QuadGK

# ╔═╡ ebd217ba-61f8-11eb-2337-9db63173ed34
md"# Module 2: _Probability Quick Refresher in Julia_
**CS6741: Statistical Foundations of Data Science**

In this module, we will quickly revise some aspects of Probability through exercises on Julia.

"

# ╔═╡ 4e0bacac-621c-11eb-17be-a968500dc6bc
md"### Dice Throws"

# ╔═╡ 48a0d2ac-6a1d-11eb-2441-897ca192e8ad
md"👉 How would you simulate a fair dice?"

# ╔═╡ 29fb09a2-61f9-11eb-0699-4dbf7daca0d7
faces = 1:6

# ╔═╡ bbae804a-61f9-11eb-03fe-89049682eae7
rand(faces)

# ╔═╡ e977219e-620d-11eb-2ca1-2141cc7c7d2a
md"What is the distribution of values obtained with throws?"

# ╔═╡ ee1bd118-61f9-11eb-1a18-df7d29a96c97
@bind throwdice Button("Throw dice")

# ╔═╡ 059c87a6-61fa-11eb-01b4-09c40af1be94
throws = []

# ╔═╡ 0a10f34e-61fa-11eb-365d-01ff16513530
begin
	throwdice
	for _ in 1:1000
		push!(throws, rand(faces))
	end
end

# ╔═╡ e00f4d02-61f9-11eb-2c5e-19ab03f4bdeb
begin
	throwdice
	histogram(throws, bins = [1, 2, 3, 4, 5, 6, 7], normalize=true, legend=false)
end

# ╔═╡ 20ea39f4-620e-11eb-05b3-716b88bcd205
md"👉 What about the distribution of throws of two dice?"

# ╔═╡ c0b6a868-61fb-11eb-2190-83d6d760ede4
begin
	N = 1000
	two_throws =  [(rand(faces), rand(faces)) for _ in 1:N]
end

# ╔═╡ 947c880a-6203-11eb-1bfe-95c8e630b395
begin
	freq_matrix = zeros(length(faces), length(faces))
	for i in faces
		for j in faces
			freq_matrix[i, j] = count(two_throws .== [(i, j)]) / N
		end
	end
end

# ╔═╡ fb898100-6204-11eb-1faf-cd75283f54d4
freq_matrix

# ╔═╡ b25c21be-6203-11eb-2f40-5db01b4beb89
heatmap(faces, faces, freq_matrix, clim=(0, 0.1))

# ╔═╡ 2b3a4f0c-620e-11eb-2021-bdaab9318fc1
md"👉 What is the probability that the sum of two dice is prime?"

# ╔═╡ c2ae2510-620a-11eb-07e0-91cf24657a55
function isPrime(x)
	if x == 1
		return false
	elseif x == 2
		return true
	end
	for i = 2:sqrt(x)
		if x%i == 0
			return false
		end
	end
	return true
end

# ╔═╡ 853a1ab6-6a1d-11eb-2e08-93e5bcc73c7c
isPrime.(2:12)

# ╔═╡ 33d3be8c-620e-11eb-3f7d-8f9a60cfdefc
md"We have 11 unique sums possible 2 through 12 out of which only 5 (2, 3, 5, 7, 11) are prime. Hence probability is:"

# ╔═╡ 6205d7aa-620d-11eb-3c75-83b491a84995
count(isPrime.(2:12)) / 11

# ╔═╡ 6f0eec76-6a1d-11eb-2585-17959367c09f
md"👉 Is the above correct? Why or why not?"

# ╔═╡ ba50d33e-6a1d-11eb-0687-47cb0b6f8242
md"👉 What visualiation would you use to help you with this?"

# ╔═╡ ec8b9cd0-6211-11eb-06fa-37880ae43e40
histogram([(rand(faces) + rand(faces)) for _ in 1:10000])

# ╔═╡ e37eefb6-6a1d-11eb-000b-89dbddabce5d
md"👉 How do we fix this?"

# ╔═╡ 5b2accf0-620e-11eb-06b9-ef354ebb3261
md"Correct theoretical approach assumes equal probability of the individual dice throws and not the sum"

# ╔═╡ 7dced6ec-6205-11eb-3e95-5f2fa0d5adf3
theorySol = sum([isPrime(i + j) for i in faces, j in faces]) / length(faces)^2

# ╔═╡ 7bbe53a6-620e-11eb-296a-dbf3eaab558d
md"Now we will compute the same thing by running lots of experiments with randomly sampled data. Usually called _Monte Carlo experiments_"

# ╔═╡ 014acff6-6206-11eb-0f3f-f17ef6edf8dd
@bind n_exp html"<input type=range min=1 max=6>"

# ╔═╡ 987d2bb2-620d-11eb-07b6-1736ac133029
n_exp

# ╔═╡ 0c205d04-6205-11eb-2d31-5f689e25a3f1
est = sum([isPrime(rand(faces) + rand(faces)) for i in 1:10^n_exp]) / 10^n_exp

# ╔═╡ beccfb34-620e-11eb-177a-711b597f5f6f
md"These random experiments give different results each time"

# ╔═╡ b0c7733e-620e-11eb-0fb6-f3a9f7631823
[sum([isPrime(rand(faces) + rand(faces)) for i in 1:10^5]) / 10^5 for j = 1: 10]

# ╔═╡ ce084cfc-620e-11eb-2881-8f30a7558bba
md"For reproducibility, we can set the seed of the random number generator"

# ╔═╡ d8b7f274-620e-11eb-13b1-3933af0f7157
begin
	Random.seed!(0)
	sum([isPrime(rand(faces) + rand(faces)) for i in 1:10^5]) / 10^5
end

# ╔═╡ 77fb84bc-621c-11eb-2cc3-71425a74f823
md"Notice that in the above we used the probability concepts of _independence_ of two random variables and _Monte Carlo_ simulations."

# ╔═╡ 59437fd2-621c-11eb-1e31-f1973639a0e2
md"### Picking balls"

# ╔═╡ 8fbbc5a6-6212-11eb-1020-7382ef578cea
md"Selection with and without replacement"

# ╔═╡ beeacf18-6212-11eb-1929-a5c830dddbcb
md"Let us suppose a bag contains 5 balls - 3 tennis balls, 2 cricket balls"

# ╔═╡ 6aa6b4ec-6212-11eb-39ea-db9d36b22cce
bag = ["🎾", "🥎", "🎾", "🎾", "🥎"]

# ╔═╡ bc4822e0-6212-11eb-3a49-e5307fe49156
rand(bag)

# ╔═╡ d5215c20-6212-11eb-32f5-c71343b60d22
md"👉 What is the probability that if I draw two balls both are tennis balls"

# ╔═╡ e6f2b97e-6212-11eb-0ef4-6593e840b93a
md"The easier case is with replacement. It is simply $P(🎾) \times P(🎾)$"

# ╔═╡ e49d22cc-6212-11eb-2732-27ad5e863ffc
3/5 * 3/5

# ╔═╡ f462afb0-6212-11eb-1a80-239d798e25f8
md"👉 What if we do not allow replacement?"

# ╔═╡ 383c0f2e-6213-11eb-355e-5b16663f1897
md"We are looking for $P(🎾|\phi) \times P(🎾|🎾)$"

# ╔═╡ 559fea86-6213-11eb-185e-b5fc917251fe
3/5 * 2/4

# ╔═╡ 73f6494e-6213-11eb-2a42-7b58fe47d597
md"👉 How do we implement this in Julia when running for more complex examples? Which one do you think is easier to implement - with or without replacements?"

# ╔═╡ 1ad550e8-6214-11eb-108b-0d6e08518018
md"With replacement:"

# ╔═╡ e455c886-6213-11eb-01b8-5df9e2bb7ea8
sum([(rand(bag) == "🎾" && rand(bag) == "🎾") for i in 1:100000])/100000

# ╔═╡ 33bc8ece-6a1f-11eb-20c9-fd32413cc2ac
md"👉 How would you implement this without replacement?"

# ╔═╡ 17ebe7ac-6214-11eb-38f1-3b757972e958
md"Without replacement:"

# ╔═╡ 3d281fc0-6214-11eb-16ae-89900d7f437f
begin
	num = 10000
	count_ = 0
	for i = 1:num
		bag1 = copy(bag)
		ball1 = rand(bag1)
		deleteat!(bag1, findfirst(x-> x == "🎾", bag1))
		ball2 = rand(bag1)
		if ball1 == "🎾" && ball2 == "🎾"
			count_ += 1
		end
	end
	count_ / num
end

# ╔═╡ 94d49374-621c-11eb-2985-b5ee5644c67e
md"Notice that in the above example we use the idea of _conditional probability_"

# ╔═╡ 5eff3d3a-621c-11eb-0824-f76604aed501
md"### Inclusion exclusion principle"

# ╔═╡ ba16f60e-623f-11eb-190a-af4758c32068
md"You have to inform $n$ people about some important news, but you have a fancy quantum phone that randomly dials a person. If you are allowed to make $r \geq n$ calls, what is the probability that all $n$ people are informed of the important news."

# ╔═╡ 51413a62-6a1f-11eb-1400-5918c088122c
md"👉 How would you do this? Think first of the mathematical approach, and then we will look at the computational approach."

# ╔═╡ eed974c0-623f-11eb-021e-ebfc3c85b934
md"Let $A_i$ denote the event that the $i$th person is not informed. Then the probability we seek is
```math
P(B) = 1 - P(A_1 \cup \ldots \cup A_n)
```"

# ╔═╡ 2cf52d58-6240-11eb-1e21-f3c0a71b4525
md"The probability of the union of those events is not as easy as it may seem. It requires application of the _inclusion exclusion_ principle repeatedly to obtain
```math
P(\large\cup_{i = 1}^n A_i) = \sum_{i = 1}^n P(A_i) - \sum_{\mbox{pairs}} P(A_i \cap A_j) + \sum_{\mbox{triplets}} P(A_i \cap A_j \cap A_k) + \ldots
```
"

# ╔═╡ e8c3415a-6240-11eb-328c-fd7e571ad7f9
md"So our probability becomes
```math
P(B) = 1 - \sum_{k = 1}^n (-1)^{k + 1} {n \choose k}p_k
```
where $p_k$ is the probability of at least $k$ people hearing about the news. From counting principles you can find out that
```math
p_k = \dfrac{(n - k)^r}{n^r} = \left(1 - \dfrac{k}{n}\right)^r
```
"

# ╔═╡ 74a09470-6241-11eb-00df-cb5032e755c3
function callprobabilty(n, r)
	return sum([(-1)^k * binomial(n, k) * (1 - k/n)^r for k in 0:n])
end

# ╔═╡ 9d189772-6241-11eb-16a3-25c54c8d4148
theorycurve = [callprobabilty(big(n), big(2*n)) for n in 1:100]

# ╔═╡ 024a8056-6242-11eb-23b6-a14e3fe68a74
plot(1:100, theorycurve)

# ╔═╡ 619109a4-6242-11eb-34eb-2527db9629bc
theorycurves = [ [callprobabilty(big(n), big(r_ratio*n)) for n in 1:200] for r_ratio in 1:3]

# ╔═╡ 85f02e2e-6242-11eb-2a27-7746c98beb47
plot(1:100, theorycurves, label=["r=n" "r=2n" "r=3n"])

# ╔═╡ ce1be85a-6242-11eb-026b-eb7215cc8c18
function simulatecalls(n, r, N)
	failed = float(0)
	for exp in 1:N
		counts = zeros(n)
		for k in 1:r
			person = rand(1:n)
			counts[person] += 1
		end
		if minimum(counts) == 0
			failed += 1
		end
	end
	return 1.0 - failed / N
end

# ╔═╡ 363c8308-6249-11eb-3580-8b721c00a424

# ╔═╡ e0efb4ea-6245-11eb-2423-b7d55a99bf0e
empiricalcurves = [[simulatecalls(n, r_ratio*n, 1000) for n in 1:100] for r_ratio in 1:3]

# ╔═╡ 1aaa0faa-6246-11eb-1f5c-69cbf277efea
begin
	scatter(1:100, empiricalcurves, label="")
	plot!(1:100, theorycurves, label=["r=n" "r=2n" "r=3n"])
end

# ╔═╡ 55b868f4-62bd-11eb-16a8-e14408e3e4bd
md"### Testing independence in random number generators"

# ╔═╡ a1c8ec8c-6a1f-11eb-0a85-fb5cac062c56
md"Randomness is a big requirement in statistical analysis and computational methods that depend on sampling."

# ╔═╡ 95bd11e8-6a1f-11eb-3287-0b32ba568406
md"👉 Do you know of any simple way to generate pseudo random numbers?"

# ╔═╡ abe1778a-6249-11eb-3b17-0de274e8e3fe
md"
A pseudo-random generator is given by the following sequence of outputs:
```math
X_{n + 1} = (a X_n + c) \mbox{ mod } m
```
"

# ╔═╡ 277b4f60-624a-11eb-2e11-9b6c5aeb3b4a
function prng(n)
	c = 7
	m = 13
	a = 5
	out = zeros(Int, n)
	out[1] = floor(rand() * m)
	for i in 2:n
		out[i] = (a * out[i - 1] + c) % m
	end
	return out
end

# ╔═╡ 9645b9a6-624a-11eb-09e9-c399cb136a5c
prng(5)

# ╔═╡ f6e2aff6-624c-11eb-16b5-318c5f230a56
n_vals = 1000

# ╔═╡ b97aa654-6a1f-11eb-11da-f91221057751
md"👉 Are two consecutive random numbers generated with `prng` independent? How would you check this computationally?"

# ╔═╡ ff30f85c-624b-11eb-3cf2-c170108d7092
vals1 = [prng(2) for i in 1:n_vals]

# ╔═╡ e59bc748-6a1f-11eb-2463-7dc3313f4b35
md"Remember that two random variables $X$ and $Y$ are indpendent if
```math
P(X, Y) = P(X) P(Y)
```"

# ╔═╡ b197bcf2-624c-11eb-3b99-efbedf0c6f3b
begin
	a = rand(1:13)
	b = rand(1:13)
	p_a = sum([vals1[i][1] == a for i in 1:n_vals]) / n_vals
	p_b = sum([vals1[i][2] == b for i in 1:n_vals]) / n_vals
	p_ap_b = p_a * p_b
	p_ab = sum([vals1[i][1] == a && vals1[i][2] == b for i in 1:n_vals]) / n_vals
	"a=$a b=$b p(x = a)=$p_a p(y = b)=$p_b p(x = a, y = b)=$p_ab p(x = a)p(y = b)=$p_ap_b"
end

# ╔═╡ 6b0245fe-624f-11eb-38f1-e702814d57c3
vals2 = [[rand(1:13), rand(1:13)] for i in 1:n_vals]

# ╔═╡ 5193bb62-62bc-11eb-1df7-1b7a7e9983ed
begin
	p_a2 = sum([vals2[i][1] == a for i in 1:n_vals]) / n_vals
	p_b2 = sum([vals2[i][2] == b for i in 1:n_vals]) / n_vals
	p_a2p_b2 = p_a2 * p_b2
	p_ab2 = sum([vals2[i][1] == a && vals2[i][2] == b for i in 1:n_vals]) / n_vals
	"a=$a b=$b p(x = a)=$p_a2 p(y = b)=$p_b2 p(x = a, y = b)=$p_ab2 p(x = a)p(y = b)=$p_a2p_b2"
end

# ╔═╡ 76fffb4e-62bd-11eb-153b-690eeaed1160
md"### Distributions and their plots"

# ╔═╡ fae8040a-62be-11eb-18a4-9fe8dee70bed
D_n = Normal(0, 1)

# ╔═╡ 0ea764ea-62bf-11eb-0e98-437744ee9fa9
plot(-5:0.1:5, [pdf(D_n, x) for x in -5:0.1:5], label="Normal PDF")

# ╔═╡ 3771b0ec-62bf-11eb-0300-6d9503c5716b
plot(-5:0.1:5, [cdf(D_n, x) for x in -5:0.1:5], label="CDF")

# ╔═╡ 66abd1d6-62c3-11eb-2e4e-1f67f89c3526
begin
	plot(-5:0.1:5, [cdf(D_n, x) for x in -5:0.1:5], label="CDF")
	plot!(-5:0.1:5, [quadgk(x -> pdf(D_n, x), (-5, y)...)[1] for y in -5:0.1:5], line = (4, :dash), label="From integration")
end

# ╔═╡ c420bab4-62c0-11eb-13c5-51ef6be81102
D_b = Binomial(5, 0.5)

# ╔═╡ ee11743a-62c0-11eb-04ef-8f30696c46ee
plot(0:5, [pdf(D_b, x) for x in 0:5], line=:stem, marker=:circle, label="Binomial")

# ╔═╡ 66a1bfe2-62bf-11eb-36d4-cfcaff920522
mean(D_n)

# ╔═╡ 93a82e9c-62cc-11eb-1b9c-cbc0e30b623e
md"### Moments of distributions"

# ╔═╡ f5b41e14-62bf-11eb-3257-6f5f5d6ceee6
md"What is the mean?
For a discrete distribution:
```math
\mathbb{E}[x]  = \sum_x xp(x)
```
For a continuous distribution:
```math
\mathbb{E}[x]  = \int_{-\infty}^{\infty} x f(x) dx,
```
where $f(x)$ is the pdf of the distribution.
"

# ╔═╡ 2be88b12-62c2-11eb-14de-6912598bc6bb
area = quadgk(x -> pdf(D_n, x), (-5, 5)...)

# ╔═╡ ec35489c-62bf-11eb-0311-2fed6b29d0c6
computed_mean = quadgk(x -> x * pdf(D_n, x), (-5, 5)...)

# ╔═╡ 5dffb66c-62c1-11eb-32f6-4b0b17d2aed1
mean(Binomial(5, 0.5))

# ╔═╡ 35406a36-62c2-11eb-347e-9b03c368119a
tot_length = sum([pdf(D_b, x) for x in 0:5])

# ╔═╡ 653515e4-62c1-11eb-25b0-930f782be996
computed_mean_binominal = sum([x * pdf(D_b, x) for x in 0:5])

# ╔═╡ 6a1d4ba4-62c2-11eb-3b71-431ee0c1fdf5
md"
```math
\mbox{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2
```
"

# ╔═╡ 565657a8-62c2-11eb-22ad-9fbf7e7b44e0
var(D_n)

# ╔═╡ 641557fe-62c2-11eb-1ffd-0b1486231fa0
quadgk(x -> x^2 * pdf(D_n, x), (-5, 5)...)[1] - mean(D_n)^2

# ╔═╡ a46d7728-62cc-11eb-3486-d799a3d29298
md"### Sampling and sums of distributions"

# ╔═╡ 5811e2d8-62cc-11eb-266a-3d1ed4d9270b
n_samples = 1000

# ╔═╡ 2d0191b4-62c3-11eb-1479-21830ac879eb
x1 = rand(D_n, n_samples)

# ╔═╡ c845fae0-62c6-11eb-2a29-bdbc0c07d6d4
histogram(x1, normalize=true)

# ╔═╡ bd903520-62c6-11eb-1f71-49b662373de7
x2 = rand(D_b, n_samples)

# ╔═╡ 485386b2-62cc-11eb-1dd4-31857e128d7f
histogram(x2, normalize=true)

# ╔═╡ c4f908fa-62c6-11eb-259d-576dee85efb9
x = x1 + x2

# ╔═╡ 505c1658-62cc-11eb-1068-659d4286a4da
histogram(x, normalize=true)

# ╔═╡ 6e365fee-62cc-11eb-28d4-cd498c21895d
begin
	histogram(x, normalize=true)
	plot!(-5:0.1:5, [pdf(D_n, x)+pdf(D_b, x) for x in -5:0.1:5])
end

# ╔═╡ ad175b5a-62cc-11eb-2913-8741eca3d89c
md"The random variable which is the sum of two independent random variable does not necessarily have a pdf as the sum of the pdfs of the two random variables."

# ╔═╡ c19c5abc-62cc-11eb-3bcd-d9ee89e4619d
begin
	conv(x) = sum(pdf(D_n,x-k)*pdf(D_b,k) for k=-5:0.1:8)
	plot(-5:0.1:8, conv.(-5:0.1:8))
end

# ╔═╡ e6d9958c-62cd-11eb-164e-4b52aedd1584
md"The convolution operation on the individual pdfs creates the pdf of the sum of two random variables. It is defined as:
```math
h(z) = (p1 \otimes p2)(z) = \int_{-\infty}^{\infty} f(t) g(z - t) dt
```
This needs to be changed to discrete versions if distribution/s are discrete.
"

# ╔═╡ b68f9fdc-62cd-11eb-1852-0116fe475f36
begin
	histogram(x, normalize=true)
	plot!(-5:0.1:8, conv.(-5:0.1:8))
end

# ╔═╡ 7b02ff44-62ce-11eb-2c40-b5a734f869cb
md"### Inverse"

# ╔═╡ 7e5765f6-62ce-11eb-222e-9786c5fc0911
md"When we say that the $\mbox{CDF}(x) = u$, we mean that the probability of the random number being less than or equal to $x$ is $u$. We can now ask an opposite question: For what value of $x$ is the probability of the random number being less than $x$ is a given $u$. This is usually referred to as the _inverse function_."

# ╔═╡ beea6cda-62ce-11eb-0ef1-253f1a40504b
md"Though not mathematically tight, we can define the inverse function $\mbox{CDF}^{-1}$ as:
```math
\mbox{CDF}^{-1}(u) = \inf\{x\ :\ \mbox{CDF}(x) \geq u\}
```
"

# ╔═╡ 0047b980-62cf-11eb-3ca2-ad2f18f0e091
cdf_n(x) = cdf(Normal(0, 1), x)

# ╔═╡ 1676307e-62cf-11eb-39e5-bd732850ab2f
cdf_n_inv(u) = minimum(filter((x) -> (cdf_n(x) >= u), -5:0.01:5))

# ╔═╡ 343f9c62-62cf-11eb-2d09-cff6967a1d65
plot(-5:0.1:5, [cdf_n(x) for x in -5:0.1:5])

# ╔═╡ 5fab01a2-62cf-11eb-1c61-f9da64f9a75f
plot(0.1:0.01:0.9, [cdf_n_inv(u) for u in 0.1:0.01:0.9])

# ╔═╡ 9a083902-62d2-11eb-2570-bdb7ea1e1f5e
md"Let us relate the inverse function to area under the PDF curve"

# ╔═╡ c483fc1c-62d1-11eb-14b6-f50474a58fb2
cdf_n_inv(0.8)

# ╔═╡ a9a39152-62d2-11eb-3b4a-5d37f7195f3a
cdf_n_inv(0.95)

# ╔═╡ f437e702-62d1-11eb-1ac8-63c7f6ec16c4
xslider = @bind x_slider html"<input type=range min=-5 max=5 step=0.01>"

# ╔═╡ 8439e2ec-62d2-11eb-3772-d17e099debb5
x_slider

# ╔═╡ 794c86e6-62d2-11eb-3e27-ebd8a1b3d4e3
u_cdf = cdf_n(x_slider)

# ╔═╡ cb5e2858-62d1-11eb-1489-037703d411ec
begin
	plot(x->x, x->pdf(Normal(0, 1), x), -10, 10, label="")
	plot!(x->x, x->pdf(Normal(0, 1), x), x_slider, 10, fill=(0, :orange), label="")
end

# ╔═╡ 2a8682e2-62d3-11eb-119e-cbd804e1b0b2
md"### Fitting given data to distributions"

# ╔═╡ a089b976-62d3-11eb-2702-8945a47bf858
n_slider = @bind n_samples_exp html"<input type=range min=1 max=5 step=1>"

# ╔═╡ 3151c634-62d3-11eb-2698-95e4e14aab2d
samples = rand(Normal(10, 2.5), 10^n_samples_exp) + rand(Normal(0, 1), 10^n_samples_exp)

# ╔═╡ 62e55e4a-62d3-11eb-21e3-7b02354ddd02
histogram(samples)

# ╔═╡ 786f3d9c-62d3-11eb-28a7-b387242e3760
fit(Normal, samples)

# ╔═╡ 9531fffe-62d3-11eb-124c-0d2cccea92ae
md"Notice how the estimate of μ is reasonably unbiased even for small sizes of the samples, but how σ seems to be higher."

# ╔═╡ Cell order:
# ╟─ebd217ba-61f8-11eb-2337-9db63173ed34
# ╠═ebcb6ec6-61f9-11eb-05a3-817618350415
# ╠═2d1693e6-61fa-11eb-3cb6-9f5431449a41
# ╟─4e0bacac-621c-11eb-17be-a968500dc6bc
# ╟─48a0d2ac-6a1d-11eb-2441-897ca192e8ad
# ╠═29fb09a2-61f9-11eb-0699-4dbf7daca0d7
# ╠═bbae804a-61f9-11eb-03fe-89049682eae7
# ╟─e977219e-620d-11eb-2ca1-2141cc7c7d2a
# ╠═ee1bd118-61f9-11eb-1a18-df7d29a96c97
# ╠═059c87a6-61fa-11eb-01b4-09c40af1be94
# ╠═0a10f34e-61fa-11eb-365d-01ff16513530
# ╠═e00f4d02-61f9-11eb-2c5e-19ab03f4bdeb
# ╟─20ea39f4-620e-11eb-05b3-716b88bcd205
# ╠═c0b6a868-61fb-11eb-2190-83d6d760ede4
# ╠═947c880a-6203-11eb-1bfe-95c8e630b395
# ╠═fb898100-6204-11eb-1faf-cd75283f54d4
# ╠═b25c21be-6203-11eb-2f40-5db01b4beb89
# ╟─2b3a4f0c-620e-11eb-2021-bdaab9318fc1
# ╟─c2ae2510-620a-11eb-07e0-91cf24657a55
# ╠═853a1ab6-6a1d-11eb-2e08-93e5bcc73c7c
# ╟─33d3be8c-620e-11eb-3f7d-8f9a60cfdefc
# ╠═6205d7aa-620d-11eb-3c75-83b491a84995
# ╟─6f0eec76-6a1d-11eb-2585-17959367c09f
# ╟─ba50d33e-6a1d-11eb-0687-47cb0b6f8242
# ╠═ec8b9cd0-6211-11eb-06fa-37880ae43e40
# ╟─e37eefb6-6a1d-11eb-000b-89dbddabce5d
# ╟─5b2accf0-620e-11eb-06b9-ef354ebb3261
# ╠═7dced6ec-6205-11eb-3e95-5f2fa0d5adf3
# ╟─7bbe53a6-620e-11eb-296a-dbf3eaab558d
# ╠═014acff6-6206-11eb-0f3f-f17ef6edf8dd
# ╠═987d2bb2-620d-11eb-07b6-1736ac133029
# ╠═0c205d04-6205-11eb-2d31-5f689e25a3f1
# ╟─beccfb34-620e-11eb-177a-711b597f5f6f
# ╠═b0c7733e-620e-11eb-0fb6-f3a9f7631823
# ╟─ce084cfc-620e-11eb-2881-8f30a7558bba
# ╠═f5d3f88c-620e-11eb-3bc8-e3e0e1598482
# ╠═d8b7f274-620e-11eb-13b1-3933af0f7157
# ╟─77fb84bc-621c-11eb-2cc3-71425a74f823
# ╟─59437fd2-621c-11eb-1e31-f1973639a0e2
# ╟─8fbbc5a6-6212-11eb-1020-7382ef578cea
# ╟─beeacf18-6212-11eb-1929-a5c830dddbcb
# ╠═6aa6b4ec-6212-11eb-39ea-db9d36b22cce
# ╠═bc4822e0-6212-11eb-3a49-e5307fe49156
# ╟─d5215c20-6212-11eb-32f5-c71343b60d22
# ╟─e6f2b97e-6212-11eb-0ef4-6593e840b93a
# ╠═e49d22cc-6212-11eb-2732-27ad5e863ffc
# ╟─f462afb0-6212-11eb-1a80-239d798e25f8
# ╟─383c0f2e-6213-11eb-355e-5b16663f1897
# ╠═559fea86-6213-11eb-185e-b5fc917251fe
# ╟─73f6494e-6213-11eb-2a42-7b58fe47d597
# ╟─1ad550e8-6214-11eb-108b-0d6e08518018
# ╠═e455c886-6213-11eb-01b8-5df9e2bb7ea8
# ╟─33bc8ece-6a1f-11eb-20c9-fd32413cc2ac
# ╟─17ebe7ac-6214-11eb-38f1-3b757972e958
# ╠═3d281fc0-6214-11eb-16ae-89900d7f437f
# ╟─94d49374-621c-11eb-2985-b5ee5644c67e
# ╟─5eff3d3a-621c-11eb-0824-f76604aed501
# ╟─ba16f60e-623f-11eb-190a-af4758c32068
# ╟─51413a62-6a1f-11eb-1400-5918c088122c
# ╟─eed974c0-623f-11eb-021e-ebfc3c85b934
# ╟─2cf52d58-6240-11eb-1e21-f3c0a71b4525
# ╟─e8c3415a-6240-11eb-328c-fd7e571ad7f9
# ╠═74a09470-6241-11eb-00df-cb5032e755c3
# ╠═9d189772-6241-11eb-16a3-25c54c8d4148
# ╠═024a8056-6242-11eb-23b6-a14e3fe68a74
# ╠═619109a4-6242-11eb-34eb-2527db9629bc
# ╠═85f02e2e-6242-11eb-2a27-7746c98beb47
# ╠═ce1be85a-6242-11eb-026b-eb7215cc8c18
# ╠═363c8308-6249-11eb-3580-8b721c00a424
# ╠═e0efb4ea-6245-11eb-2423-b7d55a99bf0e
# ╠═1aaa0faa-6246-11eb-1f5c-69cbf277efea
# ╟─55b868f4-62bd-11eb-16a8-e14408e3e4bd
# ╟─a1c8ec8c-6a1f-11eb-0a85-fb5cac062c56
# ╟─95bd11e8-6a1f-11eb-3287-0b32ba568406
# ╟─abe1778a-6249-11eb-3b17-0de274e8e3fe
# ╠═277b4f60-624a-11eb-2e11-9b6c5aeb3b4a
# ╠═9645b9a6-624a-11eb-09e9-c399cb136a5c
# ╠═f6e2aff6-624c-11eb-16b5-318c5f230a56
# ╟─b97aa654-6a1f-11eb-11da-f91221057751
# ╠═ff30f85c-624b-11eb-3cf2-c170108d7092
# ╟─e59bc748-6a1f-11eb-2463-7dc3313f4b35
# ╠═b197bcf2-624c-11eb-3b99-efbedf0c6f3b
# ╠═6b0245fe-624f-11eb-38f1-e702814d57c3
# ╠═5193bb62-62bc-11eb-1df7-1b7a7e9983ed
# ╟─76fffb4e-62bd-11eb-153b-690eeaed1160
# ╠═6b69c792-62bd-11eb-0f1d-2d0081598dfa
# ╠═fae8040a-62be-11eb-18a4-9fe8dee70bed
# ╠═0ea764ea-62bf-11eb-0e98-437744ee9fa9
# ╠═3771b0ec-62bf-11eb-0300-6d9503c5716b
# ╠═66abd1d6-62c3-11eb-2e4e-1f67f89c3526
# ╠═c420bab4-62c0-11eb-13c5-51ef6be81102
# ╠═ee11743a-62c0-11eb-04ef-8f30696c46ee
# ╠═66a1bfe2-62bf-11eb-36d4-cfcaff920522
# ╠═b80b8e58-62bf-11eb-11fb-1397b06ae3e1
# ╟─93a82e9c-62cc-11eb-1b9c-cbc0e30b623e
# ╟─f5b41e14-62bf-11eb-3257-6f5f5d6ceee6
# ╠═2be88b12-62c2-11eb-14de-6912598bc6bb
# ╠═ec35489c-62bf-11eb-0311-2fed6b29d0c6
# ╠═5dffb66c-62c1-11eb-32f6-4b0b17d2aed1
# ╠═35406a36-62c2-11eb-347e-9b03c368119a
# ╠═653515e4-62c1-11eb-25b0-930f782be996
# ╟─6a1d4ba4-62c2-11eb-3b71-431ee0c1fdf5
# ╠═565657a8-62c2-11eb-22ad-9fbf7e7b44e0
# ╠═641557fe-62c2-11eb-1ffd-0b1486231fa0
# ╟─a46d7728-62cc-11eb-3486-d799a3d29298
# ╠═5811e2d8-62cc-11eb-266a-3d1ed4d9270b
# ╠═2d0191b4-62c3-11eb-1479-21830ac879eb
# ╠═c845fae0-62c6-11eb-2a29-bdbc0c07d6d4
# ╠═bd903520-62c6-11eb-1f71-49b662373de7
# ╠═485386b2-62cc-11eb-1dd4-31857e128d7f
# ╠═c4f908fa-62c6-11eb-259d-576dee85efb9
# ╠═505c1658-62cc-11eb-1068-659d4286a4da
# ╠═6e365fee-62cc-11eb-28d4-cd498c21895d
# ╟─ad175b5a-62cc-11eb-2913-8741eca3d89c
# ╠═c19c5abc-62cc-11eb-3bcd-d9ee89e4619d
# ╟─e6d9958c-62cd-11eb-164e-4b52aedd1584
# ╠═b68f9fdc-62cd-11eb-1852-0116fe475f36
# ╟─7b02ff44-62ce-11eb-2c40-b5a734f869cb
# ╟─7e5765f6-62ce-11eb-222e-9786c5fc0911
# ╟─beea6cda-62ce-11eb-0ef1-253f1a40504b
# ╠═0047b980-62cf-11eb-3ca2-ad2f18f0e091
# ╠═1676307e-62cf-11eb-39e5-bd732850ab2f
# ╠═343f9c62-62cf-11eb-2d09-cff6967a1d65
# ╠═5fab01a2-62cf-11eb-1c61-f9da64f9a75f
# ╟─9a083902-62d2-11eb-2570-bdb7ea1e1f5e
# ╠═c483fc1c-62d1-11eb-14b6-f50474a58fb2
# ╠═a9a39152-62d2-11eb-3b4a-5d37f7195f3a
# ╠═f437e702-62d1-11eb-1ac8-63c7f6ec16c4
# ╠═8439e2ec-62d2-11eb-3772-d17e099debb5
# ╠═794c86e6-62d2-11eb-3e27-ebd8a1b3d4e3
# ╠═cb5e2858-62d1-11eb-1489-037703d411ec
# ╟─2a8682e2-62d3-11eb-119e-cbd804e1b0b2
# ╠═a089b976-62d3-11eb-2702-8945a47bf858
# ╠═3151c634-62d3-11eb-2698-95e4e14aab2d
# ╠═62e55e4a-62d3-11eb-21e3-7b02354ddd02
# ╠═786f3d9c-62d3-11eb-28a7-b387242e3760
# ╟─9531fffe-62d3-11eb-124c-0d2cccea92ae


using StatisticalRethinking
#gr(size=(600,600));

N = 20
p_grid = range( 0 , stop=1 , length=N )

prior = ones( 20 );

likelihood = [pdf(Binomial(9, p), 6) for p in p_grid]
likelihood[1:5]

unstd_posterior = likelihood .* prior;

posterior = unstd_posterior  ./ sum(unstd_posterior);

p1 = plot( p_grid , posterior ,
    xlab="probability of water" , ylab="posterior probability",
    lab = "interpolated", title="20 points" )
p2 = scatter!( p1, p_grid , posterior, lab="computed" )

prior1 = [p < 0.5 ? 0 : 1 for p in p_grid]
prior2 = [exp( -5*abs( p - 0.5 ) ) for p in p_grid]

p3 = plot(p_grid, prior1,
  xlab="probability of water" , ylab="posterior probability",
  lab = "semi_uniform", title="Other priors" )
scatter!(p3, p_grid, prior1, lab = "semi_uniform grid point")
plot!(p3, p_grid, prior2,  lab = "double_exponential" )
scatter!(p3, p_grid, prior2,  lab = "double_exponential grid point" )

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


f(x::Vector) = (x[1] - 5.0)^2
function g!(x::Vector, storage::Vector)
    storage[1] = 2.0 * (x[1] - 5.0)
end

initial_x = [0.0]

store_trace, show_trace = false, false
results = Optim.naive_gradient_descent(f, g!, initial_x,
	                                   0.1, 10e-8, 1_000,
	                                   store_trace, show_trace)
@assert length(results.trace.states) == 0
@assert results.converged
@assert norm(results.minimum - [5.0]) < 0.01

results = Optim.naive_gradient_descent(f, g!, initial_x)
@assert length(results.trace.states) == 0
@assert results.converged
@assert norm(results.minimum - [5.0]) < 0.01

eta = 0.9

f(x::Vector) = (1.0 / 2.0) * (x[1]^2 + eta * x[2]^2)

function g!(x::Vector, storage::Vector)
    storage[1] = x[1]
    storage[2] = eta * x[2]
end

store_trace, show_trace = true, false
results = Optim.naive_gradient_descent(f, g!, [1.0, 1.0],
		                               0.1, 10e-8, 1_000,
	                                   store_trace, show_trace)
@assert length(results.trace.states) > 0
@assert results.converged
@assert norm(results.minimum - [0.0, 0.0]) < 0.01


using DifferentialEquations
using Plots

function sis_ode(du,u,p,t)
        SH,IH,SL,IL = u
        betaHH,betaHL,betaLH,betaLL,gamma=p
        du[1]=-(betaHH*IH+betaHL*IL)*SH+gamma*IH
        du[2]=+(betaHH*IH+betaHL*IL)*SH-gamma*IH
        du[3]=-(betaLH*IH+betaLL*IL)*SL+gamma*IL
        du[4]=+(betaLH*IH+betaLL*IL)*SL-gamma*IL
    end

parms =[10,0.1,0.1,1,1]
init=[0.19999,0.00001,0.799,0.001]
tspan=tspan = (0.0,15.0)

sis_prob = ODEProblem(sis_ode,init,tspan,parms)
sis_sol = solve(sis_prob,saveat=0.1)

plot(sis_sol,xlabel="Time (Years)",ylabel="Proportion of Population")


# OLS estimation using (\)
# https://github.com/giob1994/Alistair.jl/blob/3a11c19150169695581b46e4d1895f0641a4c29d/src/linregress.jl#L38-L41
function ols(X, Y)
    # Ax = B
    # x = B \ A
    # https://github.com/JuliaLang/julia/blob/d789231e9985537686052db9b2314c0d51656308/stdlib/LinearAlgebra/src/generic.jl#L827-L855
    return (X' * X) \ (X' * Y)
end

# OLS estimation using inv()
# inv() is a more computationally-intensive procedure than \ operator
# https://github.com/mcreel/Econometrics/blob/b29cffa8ce7f845fa73788522f9c7753d40684f4/Examples/OLS/OlsFit.jl#L13-L14
function ols_(X, Y)
    return inv(X'*X)*X'*Y
end


using Symbolics, SymbolicUtils
using SymbolicUtils: istree, operation, arguments
using Symbolics: value, get_variables, solve_for, derivative
using SymbolicUtils.Rewriters
using SymbolicUtils.Code

"""
    sym_taylor generates a Taylor series on variable x to order n and at
    point x₀, i.e., Σ aᵢ*(x - x₀)ⁱ
"""
function sym_taylor(x, n, x₀=0)
    x = value(x)
    @variables a[0:n]
    a = value.(a)
    a[1] + sum([a[i+1]*(x-x₀)^i for i = 1:n])
end

"""
    sym_frobenius generates a Frobenius series on variable x to order n and at
    point x₀, i.e., (x - x₀)^α Σ aᵢ*(x - x₀)ⁱ
"""
function sym_frobenius(x, n, x₀=0)
    x = value(x)
    @variables a[0:n]
    a = value.(a)
    @syms α
    (x-x₀)^α * sym_taylor(x, n, x₀)
end

# pox (power-of-x) is a symbolic function to keep track of the powers of x
# pox(k,n) means k*x^n
@syms pox(k, n)

is_pox(x) = istree(x) && operation(x)==pox
is_not_pox(x) = !is_pox(x)

get_coef(p) = is_pox(p) ? arguments(p)[1] : p
get_power(p) = is_pox(p) ? arguments(p)[2] : 0

replace_x(eq, x) = substitute(eq, Dict(x => pox(1,1)))

count_rule1 = @rule ^(pox(~k, ~n1), ~n2) => isequal(~k,1) ? pox(1, ~n1 * ~n2) : pox(^(~k,~n2), ~n1 * ~n2)
count_rule2 = @rule pox(~k1, ~n1) * pox(~k2, ~n2) => pox(~k1 * ~k2, ~n1 + ~n2)
count_rule3 = @acrule pox(~k, ~n) * ~u::is_not_pox => pox(~k * ~u, ~n)

"""
    collect_powers separates the powers of x in eq (a polynomial) and returns
    a dictionary of power => term
"""
function collect_powers(eq, x)
    eq = expand(expand_derivatives(eq))
    eq = replace_x(eq, x)
    #eq = Prewalk(PassThrough(count_rule1))(eq)
    eq = Fixpoint(Prewalk(PassThrough(Chain([count_rule1, count_rule2, count_rule3]))))(eq)

    if !istree(eq)
        return Dict{Any, Any}(0 => eq)
    elseif is_pox(eq)
        return Dict{Any, Any}(get_power(eq) => get_coef(eq))
    else
        eqs = Dict{Any, Any}()

        for term in arguments(eq)
            n = get_power(term)
            if haskey(eqs, n)
                eqs[n] = eqs[n] + get_coef(term)
            else
                eqs[n] = get_coef(term)
            end
        end

        return eqs
    end
end

"""
    strip_powers removes the powers from the output of collect_powers
    and returns a sorted list of terms/sub-equations
"""
function strip_powers(eqs)
    @syms α
    d = Dict(α => 0)
    eqs = Dict(substitute(power, d) => eq for (power, eq) in eqs)
    collect(values(sort(eqs)))
end

"""
    solve_rational gets a triangular system of linear equations and solve
    for the variables one by one starting from the first equation

    eqs is a list of terms as prepared by strip_powers

    solve_rational tries to generate rational results but falls back to reals
    if not possible
"""
function solve_rational(eqs)
    eqs = Num.(eqs)
    vals = Dict()

    for eq in eqs
        vars = get_variables(eq)
        vars = filter(x -> !haskey(vals,x), vars)

        if length(vars) > 0
            dep = vars[end]
            eq = substitute(eq, vals)
            # the last variable of each eq is considered the dependent variable
            A = value(substitute(eq, Dict(dep => 0)))
            B = value(substitute(eq, Dict(dep => 1)))
            try
                vals[dep] = A // (B - A)
            catch
                vals[dep] = A / (B - A)
            end
        end
    end

    vals
end

"""
    finalize_taylor substitutes the calculated coefficients back into the
    original Taylor or Frobenius series
"""
function finalize_taylor(x, y, vals)
    u = substitute(y, vals)
    vars = [v for v in get_variables(y) if !isequal(v,x)]
    factor(u, vars)
end

"""
    A main entry point
    solve_taylor solves a Taylor series expansion of a linear differential equation

    x is the independent variable (a Symbolic Sym)
    y is the dependent variable (a Symbolic Sym)
    diffeq is the differential equation, assuming diffeq ~ 0
"""
function solve_taylor(x, y, diffeq)
    eqs = strip_powers(collect_powers(diffeq, x))
    vals = solve_rational(eqs)
    finalize_taylor(x, y, vals)
end

"""
    A main entry point
    solve_frobenius solves a Frobenius series expansion of a linear differential equation

    x is the independent variable (a Symbolic Sym)
    y is the dependent variable (a Symbolic Sym)
    diffeq is the differential equation, assuming diffeq ~ 0
"""
function solve_frobenius(x, y, diffeq; abstol=1e-8, rationalize_coef=false)
    eqs = strip_powers(collect_powers(diffeq, x))
    eq = eqs[1]
    @syms α
    d = Dict(v => 1 for v in get_variables(eq) if !isequal(v, α))

    rs, zs = find_roots(substitute(eq, d), α)
    isempty(rs) && error("The indicial equation has no real root")
    α₁ = rs[end]
    @info "α is $α₁"

    d = Dict(α => α₁)
    eqs = [substitute(eq,d) for eq in eqs]
    vals = solve_rational(eqs[2:end])
    if rationalize_coef
        vals = rationalize(vals)
    end
    vals[α] = α₁
    finalize_taylor(x, y, vals)
end

###############################################################################

"""
    factor factors a polynomial eq over a list of variables (vars)
    it assumes trivial separatibilty
"""
function factor(eq, vars)
    f = []
    for v in vars
        d = Dict(u => (isequal(v,u) ? 1 : 0) for u in vars)
        push!(f, v * substitute(eq, d))
    end

    +(f...)
end

"""
    solve_newton is a symbolic Newton-Ralphson solver
    f is a symbolic equation to be solved (f ~ 0)
    x is the variable to solve
    x₀ is the initial guess
"""
function solve_newton(f, x, x₀; abstol=1e-8, maxiter=50)
    xₙ = Complex(x₀)
    ∂f = value(derivative(f, x))

    for i = 1:maxiter
        d = Dict(x => xₙ)
        xₙ₊₁ = xₙ - substitute(f, d) / substitute(∂f, d)
        if abs(xₙ₊₁ - xₙ) < abstol
            return xₙ₊₁
        else
            xₙ = xₙ₊₁
        end
    end
    return xₙ
end

"""
    poly_degree returns the degree of a polynomial equation eq
"""
function poly_degree(poly, x)
    eqs = collect_powers(poly, x)
    round(Int, maximum(keys(eqs)))
end

"""
    find_roots returns all the real and complex roots of a polynomial (poly)
    for variable x

    the output is rs, sz, where rs is a list of real roots and zs is a list of
    complex roots

    if poly is not a polynomial, the number of desired roots (n) needs to be
    provided. For example, find_roots(sin(x), x, 10) returns 10 different, but
    not necessarily sequential, multiples of π
"""
function find_roots(poly, x, n=-1; abstol=1e-8)
    n = (n == -1 ? poly_degree(poly, x) : n)

    rs = Float64[]
    zs = Complex[]

    while n > 0
        z = solve_newton(poly, x, Complex(rand(),rand()))
        if abs(imag(z)) < abstol
            r = real(z)
            push!(rs, r)
            poly = poly / (x - r)
            n -= 1
        else
            append!(zs, [z, conj(z)])
            poly = poly / ((x-z)*(x-conj(z)))
            n -= 2
        end
    end
    sort(rs), zs
end

##############################################################################

function continued_fraction(x, n; abstol=1e-8)
    try
        l = zeros(Int, n)
        for i = 1:n
            l[i] = floor(Int, x)
            Δx = x - l[i]
            Δx < abstol && break
            x = 1 / Δx
        end
        y = 0
        for i = n:-1:1
            if l[i] + y > 0
                y = 1 // (l[i] + y)
            end
        end
        return 1 // y
    catch
        return x
    end
end

function float_to_rational(x; n=10, abstol=1e-8)
    s = x < 0
    y = continued_fraction(abs(x), n) * (s ? -1 : +1)
    if abs(x - y) < abstol
        return y
    else
        return x
    end
end

# function rationalize(vals)
#     r = Dict()
#     for p in vals
#         eq = last(p)
#         vars = get_variables(eq)
#         if length(vars) == 1
#             var = vars[1]
#             eq = eq / var
#             if !(eq isa Int) || !(eq isa Rational)
#                 eq = float_to_rational(eq)
#             end
#             r[first(p)] = eq * var
#         else
#             r[first(p)] = eq
#         end
#     end
#     r
# end

################################### Examples! #################################

function test_harmonic()
    @syms x
    y = sym_taylor(x, 10)
    D = Differential(x)
    solve_taylor(x, y, D(D(y)) + y)
end

function test_airy()
    @syms x
    y = sym_taylor(x, 10)
    D = Differential(x)
    solve_taylor(x, y, D(D(y)) + x*y)
end

function test_bessel(ν = sqrt(2); rationalize_coef=false)
    @syms x
    y = sym_frobenius(x, 10)
    D = Differential(x)
    solve_frobenius(x, y, x^2 * D(D(y)) + x * D(y) - (x^2 + ν^2)*y; rationalize_coef=rationalize_coef)
end


using IterativeSolvers

A = Float64[
  16 4 8;
  4 5 -4;
  8 -4 22;
]
a = Float64[16; 18; -22]
x = Float64[1; 1; 1]

maxiters = 100
println("\nJacobi:")
for i in 1:100
  res = jacobi!(x, A, a, tol=1.0e-8, maxiter=1)
  if i < 5 || i in [1; 10:10:maxiters]
    println([i x'])
  end
  if res[2].isconverged
    println([i x'])
    break
  end
end

println("\nGauss-Seidel:")
@show  A * x
println()

x = Float64[1; 1; 1]
for i in 1:100
  res = gauss_seidel!(x, A, a, tol=1.0e-8, maxiter=1)
  if i < 5 || i in [1; 10:10:maxiters]
    println([i x'])
  end
  if res[2].isconverged
    println([i x'])
    break
  end
end

println("\nSequential Succesive Over-Relaxation:")
@show A * x
println()

x = Float64[1; 1; 1]
ω = 1.0
for i in 1:100
  res = ssor!(x, A, a, ω, tol=1.0e-8, maxiter=1)
  if i < 5 || i in [1; 10:10:maxiters]
    println([i x'])
  end
  if res[2].isconverged
    println([i x'])
    break
  end
end

println()
A * x


using DataInterpolations, Test
using QuadGK
using DataInterpolations: integral

function test_integral(func, tspan, name::String)
  t1 = minimum(tspan)
  t2 = maximum(tspan)
  @testset "$name" begin
    qint, err = quadgk(func, t1, t2)
    aint = integral(func, t1, t2)
    @test isapprox(qint, aint, atol=1e-8)
  end
end

# Linear Interpolation
u = 2.0collect(1:10)
t = 1.0collect(1:10)
A = LinearInterpolation(u,t)

test_integral(A, t, "Linear Interpolation (Vector)")

# Quadratic Interpolation
u = [1.0, 4.0, 9.0, 16.0]
t = [1.0, 2.0, 3.0, 4.0]
A = QuadraticInterpolation(u,t)

test_integral(A, t, "Quadratic Interpolation (Vector)")

# QuadraticSpline Interpolation
u = [0.0, 1.0, 3.0]
t = [-1.0, 0.0, 1.0]

A = QuadraticSpline(u,t)

test_integral(A, t, "Quadratic Interpolation (Vector)")

# CubicSpline Interpolation
u = [0.0, 1.0, 3.0]
t = [-1.0, 0.0, 1.0]

A = CubicSpline(u,t)

test_integral(A, t, "Cubic Spline Interpolation (Vector)")


struct InvertibleScalarFunction{S,T}
    f::S
    f⁻¹::T
end

function Base.:∘(isf1::InvertibleScalarFunction, isf2::InvertibleScalarFunction)
    f = isf1.f ∘ isf2.f
    f⁻¹ = isf2.f⁻¹ ∘ isf1.f⁻¹
    return InvertibleScalarFunction(f, f⁻¹)
end

(isf::InvertibleScalarFunction)(x) = isf.f(x)

function Base.inv(isf::InvertibleScalarFunction)
    f, f⁻¹ = isf.f, isf.f⁻¹
    return InvertibleScalarFunction(f⁻¹, f)
end

const RealVector = AbstractVector{<:Real}

is_constant(column::RealVector) = maximum(column) == minimum(column)

function binner(column::RealVector, bins::RealVector)
    rtol = 1.0e-5
    atol = 1.0e-8
    eps = @. atol + rtol * abs(column)
    return searchsortedfirst.(Ref(bins), column + eps) .- 1
end

function uniform_binner(column::RealVector, num_bins::Integer)
    if is_constant(column)
        return ones(Int46, length(column))
    end
    col_min = minimum(column)
    col_max = maximum(column)
    bins = range(col_min, stop = col_max, length = num_bins)
    return binner(column, bins)
end

function quantile_binner(column::RealVector, num_bins::Integer)
    bins = range(0, stop=1, length=num_bins)
    return binner(column, quantile(column, bins))
end

function log_scaler(_::RealVector)
    return InvertibleScalarFunction(log, exp)
end

function logistic_scaler(_::RealVector)
    sigmoid(x::T) where {T<:Real} = one(T) / (one(T) + exp(-x))
    logit(x::T) where {T<:Real} = log(x / (one(T) - x))
    return InvertibleScalarFunction(sigmoid, logit)
end

function min_max_scaler(column::RealVector, range::RealVector = [0, 1])
    max_val = maximum(column)
    min_val = minimum(column)
    min_range, max_range = range[1], range[2]
    f(x) = (x - min_val) / (max_val - min_val) * (max_range - min_range) + min_range
    f⁻¹(y) = (y - min_range) * (max_val - min_val) / (max_range - min_range) + min_val
    return InvertibleScalarFunction(f, f⁻¹)
end

function max_abs_scaler(column::RealVector)
    max_val = maximum([maximum(column), abs(minimum(column))])
    f(x) = x / max_val
    f⁻¹(y) = y * max_val
    return InvertibleScalarFunction(f, f⁻¹)
end

function standardizer(column::RealVector)
    m = mean(column)
    s = std(column)
    f(x) = (x - m) / s
    f⁻¹(y) = y * s + m
    return InvertibleScalarFunction(f, f⁻¹)
end


using SymbolicUtils
using Metatheory

sym_eq(x, y) = simplify(x == y)

@testset "Basic optimization" begin
    @syms a b c x y z
    ex = 2a + 2b - (a*(a + b))
    res = optimize(ex)
    @test sym_eq(res, (a+b)*(2-a))

    ex = sin(a^2)/cos(a^2)
    res = optimize(ex)
    @test isequal(res, tan(a^2)) # sym_eq does not work
    # TODO report issue

    ex = sin(1/x * (a * b + a * c))^2 + cos((a*(b+c))/x)^2
    res = optimize(ex)
    @test res == 1
end