#
# Maneuver calculations for orbits within the Twobody Problem
#

"""
Provides the radius of escape.
"""
escape_radius(r₀, v₀, aₜ) = r₀ * v₀ / (20 * aₜ^2 * r₀^2)^(1/4)
escape_radius(orbit::T, m::ConstantManeuver) where T <: RestrictedTwoBodySystem = escape_radius(radius(orbit), velocity(orbit), m.aₜ)

"""
Provides the velocity at escape.
"""
escape_velocity(r₀, v₀, aₜ, μ) = √(2 * μ / escape_radius(r₀, v₀, aₜ))
escape_velocity(orbit::T, m::ConstantManeuver) where T <: RestrictedTwoBodySystem = escape_velocity(radius(orbit), velocity(orbit), m.aₜ, orbit.body.μ)

"""
Provides time delta from the provided initial orbit to escape.
"""
escape_time(r₀, v₀, aₜ) = upreferred(v₀ / aₜ) * (1 - (upreferred(20aₜ^2 * r₀^2) / upreferred(v₀^4))^(1/8))
escape_time(orbit::T, m::ConstantManeuver) where T <: RestrictedTwoBodySystem = escape_time(radius(orbit), velocity(orbit), m.aₜ)

"""
Provides the path length from the initial condition to escape.
"""
escape_path_length(r₀, v₀, aₜ) = upreferred(v₀^2 / 2aₜ) * (1 - upreferred((1/v₀) * (20aₜ^2 * r₀^2)^(1/4)))
escape_path_length(orbit::T, m::ConstantManeuver) where T <: RestrictedTwoBodySystem = escape_path_length(radius(orbit), velocity(orbit), m.aₜ)


using Plots, OrdinaryDiffEq, DiffEqPhysics, RecursiveArrayTools
gr()

G = 2.95912208286e-4
M = [1.00000597682, 0.000954786104043, 0.000285583733151, 0.0000437273164546, 0.0000517759138449, 1/1.3e8]
planets = ["Sun", "Jupiter", "Saturn", "Uranus", "Neptune", "Pluto"]

pos_x = [0.0,-3.5023653,9.0755314,8.3101420,11.4707666,-15.5387357]
pos_y = [0.0,-3.8169847,-3.0458353,-16.2901086,-25.7294829,-25.2225594]
pos_z = [0.0,-1.5507963,-1.6483708,-7.2521278,-10.8169456,-3.1902382]
pos = ArrayPartition(pos_x,pos_y,pos_z)

vel_x = [0.0,0.00565429,0.00168318,0.00354178,0.00288930,0.00276725]
vel_y = [0.0,-0.00412490,0.00483525,0.00137102,0.00114527,-0.00170702]
vel_z = [0.0,-0.00190589,0.00192462,0.00055029,0.00039677,-0.00136504]
vel = ArrayPartition(vel_x,vel_y,vel_z)

tspan = (0.,200_000)

const ∑ = sum
const N = 6
potential(p, t, x, y, z, M) = -G*∑(i->∑(j->(M[i]*M[j])/sqrt((x[i]-x[j])^2 + (y[i]-y[j])^2 + (z[i]-z[j])^2), 1:i-1), 2:N)

nprob = NBodyProblem(potential, M, pos, vel, tspan)
sol = solve(nprob,Yoshida6(), dt=100);

orbitplot(sol,body_names=planets)

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


#! /usr/bin/julia

# Rosetta Code, Arithmetic-geometric mean

function agm{T<:FloatingPoint,U<:Integer}(x::T, y::T, e::U=5)
    0 < y && 0 < y && 0 < e || throw(DomainError())
    err = e*eps(x)
    (g, a) = extrema([x, y])
    while err < (a - g)
        ap = a
        a = 0.5*(a + g)
        g = sqrt(ap*g)
    end
    return a
end

x = 1.0
y = 1.0/sqrt(2.0)

println("Using literal-precision float numbers:")
println("    agm(", x, ",", y, ") = ", agm(x, y))

println()
println("Using half-precision float numbers:")
x = float16(x)
y = float16(y)
println("    agm(", x, ",", y, ") = ", agm(x, y))

println()
println("Using ", get_bigfloat_precision(), "-bit float numbers:")
x = BigFloat(1.0)
y = x/sqrt(BigFloat(2.0))
println("    agm(", x, ",", y, ") = \n    ", agm(x, y))


## Exercise 6-7
## A number, a, is a power of b if it is divisible by b and a/b is a power of b. Write a function called ispower that takes parameters a and b and returns true if a is a power of b.

## TIP
## You will have to think about the base case.
println("Ans: ")

function ispower(a, b)
    if a == b
        return true
    end

    if a % b == 0 && ispower(a / b, b)
        return true
    end

    return false
end

println(ispower(1000, 10))
println(ispower(343, 7))
println(ispower(31, 2))

println("End.")


# This file contains the functions to compute the emissivities of bulks and multilayers

# Factors
"""
    bose_einstein(w,T)

Bose-Einstein factor using the definition:
```math
\\Theta(\\omega,T) = \\frac{\\hbar\\omega}{e^{\\frac{\\hbar\\omega}{k_B T}}-1}
```

"""
function bose_einstein(w,T)
    u = w*ħ/kb/T
    if T==0
        return 0.0
    end
    if w==0.0
        return kb*T
    end

    return u*kb*T/(exp(u)-1.0)
end

"""
    planck(w,T)

Planck's Distribution using the definition:

```math
q^{\\text{BB}}(\\omega,T) =  \\Theta(\\omega,T) \\frac{k_0^2}{4\\pi^2}
```
"""
function planck(w,T)
    return  bose_einstein(w,T)*(w/c0)^2/(2.0*pi)^2
end

"""
    planck_fraction(w1,w2,T)

fraction of the energy between pulsations w1 and w2

```math
q^{\\text{BB}}(\\omega_1,\\omega_2,T) =  \\frac{1}{\\sigma T^4}\\int_{\\omega_1}^{\\omega_2}q^{\\text{BB}}(\\omega)d\\omega
```
"""
function planck_fraction(w1,w2,T)
  pl(w) = planck(w,T)
  (val, err) = quadgk(pl, w1, w2 ;rtol=1e-8)
  return val/sigma/T^4
end

"""
    wien(T)

Wien frequency: return the frequency for which Planck distribution is maximum.
T is given in Kelvin and the result is given in radHz
"""
function wien(T)
    return   2.8214393721220787*kb*T/ħ
end

"""
    invwien(w)

Inverse of Wien law : return the temperature (K)
for a given frequency in radHz
"""
function invwien(w)
    return   w*ħ/(2.8214393721220787*kb)
end

" Unit conversion from rad/s to wavelength in m"
unitconv(w) = 2.0*pi*c0/w

" Wien's wavelength obtained from wien's law in m"
lambda_wien(T) = 2.8977729e-3/T

"Monocromatic directional emissivity"
function emissivity_kx_w(structure, kx, w)

    (Rte,Tte)=power_rt(structure, te(), kx ,w)
    (Rtm,Ttm)=power_rt(structure, tm(), kx ,w)

    if imag(permittivity(substrate(structure),w)) == 0.0
        integr1 = 1.0 - Rte - Tte
        integr2 = 1.0 - Rtm - Ttm
    else
        integr1 = 1.0 - Rte
        integr2 = 1.0 - Rtm
    end

    return integr1 + integr2
end

"Total directional emissivity"
function emissivity_kx(structure,T, kx, wi,wf; rtol=1e-8)

    e(u) = emissivity_kx_w(structure, kx, u*kb*T/ħ)*u^3/(exp(u)-1.0)
    val :: Float64  = 0.0
    err :: Float64  = 0.0
    (val,err) = quadgk(e, wi*ħ/kb/T , wf*ħ/kb/T ; rtol=rtol)

    return val*kb^4/ħ^3/c0^2/(2.0*pi)^2/sigma
end

" Monocromatic hemispherical emissivity"
function emissivity_w(structure, w; rtol = 1e-8)
    e_kx(kx) = kx*emissivity_kx_w(structure ,kx, w)

    val :: Float64  = 0.0
    err :: Float64  = 0.0
    (val,err) = quadgk(e_kx, 0.0, w/c0 ; rtol=rtol)

    return val/(w/c0)^2
end

" total hemispherical emissivity"
function emissivity(structure,T;rtol = 1e-8)
    e(u) = emissivity_w(structure, u*kb*T/ħ;rtol = rtol)*u^3/(exp(u)-1.0)
    e2(t) = e(t/(1.0-t))/(1.0-t)^2

    val :: Float64  = 0.0
    err :: Float64  = 0.0
    (val,err) = quadgk(e2, 0.0, 1.0 ; rtol=rtol)

    return val*kb^4/ħ^3/c0^2/(2.0*pi)^2/sigma
end

" total hemispherical emissivity with integration bounds for frequency"
function emissivity(structure,T,wi,wf;rtol = 1e-8)
    e(u) = emissivity_w(structure, u*kb*T/ħ; rtol = rtol)*u^3/(exp(u)-1.0)
    val :: Float64  = 0.0
    err :: Float64  = 0.0
    (val,err) = quadgk(e, wi*ħ/kb/T , wf*ħ/kb/T ; rtol=rtol)

    return val*kb^4/ħ^3/c0^2/(2.0*pi)^2/sigma
end

function emissivity_fraction(structure,T,wi,wf;rtol = 1e-8)
    return emissivity_kx(structure,T,0.0,wi,wf;rtol = rtol)/planck_fraction(wi,wf,T)
end


using Optim: optimize, Fminbox, LBFGS
using Zygote: gradient

include("./kernels.jl")

"""
Following (15.14) in MLAPP, the posterior predictive for noise free
observation ``\\mathbf{f}_*`` with new data ``\\mathbf{X}_*`` is

``
p(\\mathbf{f}_* | \\mathbf{X}_*, \\mathbf{X}, \\mathbf{y}) =
\\mathscr{N}(\\mathbf{f}_* | \\mathbf{\\mu}_*, \\mathbf{\\Sigma}_*)
``

To derive posterior predictive distribution for ``\\mathbf{y}_*``,
we need to add the observation noise to ``\\mathbf{f}_*``.
"""
function predict(kernel, X, y, Xs; decomp=false)
    Ky = covariancematrix(kernel, X)
    Ks = covarianceblock(kernel, X, Xs)
    Kss = covariancematrix(kernel, Xs)
    invKy = inv(Ky)
    μs = Ks'*invKy*y
    Σs = Kss - Ks'*invKy*Ks
    return μs, Σs
end

"""
Negative log likelihood for the marginal distribution ``p(y|X)``.
The numerically stable version follows **Algorithm 15.1** in MLAPP.
"""
function negloglike(X, y, Ky; stable=false)
    N = size(X, 1)
    if stable
        L = cholesky(Ky).L
        invL = inv(L)
        α = invL'*invL*y
        loglike = -0.5*y'*α - sum(log.(diag(L))) - 0.5*N*log(2π)
        return -loglike
    else
        logdetKy = log(det(Ky))
        loglike = -0.5*y'*inv(Ky)*y - 0.5*logdetKy - 0.5*N*log(2π)
        return -loglike
    end
end

"""
Negative log-likelihood for squared exponential kernel (log parameters).
"""
function sqexp_negloglike_log(X, y, logl, logσf, logσy; stable=false)
    l = exp(logl)
    σf = exp(logσf)
    σy = exp(logσy)
    sekernel = SquaredExponentialKernel(l, σf, σy)
    # Ky = covariancematrix_alt(sekernel, X)
    Ky = covariancematrix(sekernel, X)
    return negloglike(X, y, Ky, stable=stable)
end

function sqexp_negloglike(X, y, l, σf, σy; stable=false)
    sekernel = SquaredExponentialKernel(l, σf, σy)
    Ky = covariancematrix(sekernel, X)
    return negloglike(X, y, Ky, stable=stable)
end

"""
Analytic gradient of the log marginal likelihood for squared exponential kernel.
"""
function sqexp_gradient_analytic!(G, X, y, paramvals...)
    sekernel = SquaredExponentialKernel(paramvals...)
    grad = -logmarginalgradient(sekernel, X, y)
    for i = 1:size(G, 1)
        G[i] = grad[i]
    end
    nothing
end

"""
Autodiff gradient of the log marginal gradient using Zygote.
"""
function sqexp_gradient_autodiff!(G, X, y, paramvals...)
    grad = gradient(params -> sqexp_negloglike(X, y, params...), paramvals)[1]
    for i = 1:size(G, 1)
        G[i] = grad[i]
    end
    nothing
end

function optimize_sekernel(
    X, y;
    init_params=ones(3),
    lower=[1e-6, 1e-6, 1e-6],
    upper=[Inf, Inf, Inf],
    stable=true,
    optimizer=LBFGS()
    )
    result = optimize(
        params -> sqexp_negloglike(X, y, params..., stable=stable),
        lower,
        upper,
        init_params,
        Fminbox(optimizer)
    )
    return result
end

"""
Optimizing squared exponential kernel hyperparameters with provided gradient.
The second input to `optimize` is the function to calculate gradient.
Its first argument is inserted by the `optimize` function to keep track
of the current gradient.
"""
function optimize_sekernel(
    X, y, g!;
    init_params=ones(3),
    lower=[1e-6, 1e-6, 1e-6],
    upper=[Inf, Inf, Inf],
    stable=true,
    optimizer=LBFGS()
    )
    result = optimize(
        params -> sqexp_negloglike(X, y, params..., stable=stable),
        (G, params) -> g!(G, X, y, params...),
        lower,
        upper,
        init_params,
        Fminbox(optimizer)
    )
    return result
end

function mixture_negloglike(X, y, logl, logσf, logσy, θ0, θ1)
    l = exp(logl)
    σf = exp(logσf)
    σy = exp(logσy)
    mixturekernel = MixtureKernel(l, σf, σy, θ0, θ1)
    Ky = covariancematrix(mixturekernel, X)
    return negloglike(X, y, Ky)
end

"""
    logmarginalpartial(kernel, X, y, ∂Ky∂θ)

Partial derivative of log marginal likelihood ``\\log{P(y|X)}`` wrt
parameter ``θ``.

# Parameters
- `X`: input matrix
- `y`: labels
- `∂Ky∂θ`: partial derivative of covariance matrix ``K_y`` wrt ``θ``.

"""
function logmarginalpartial(kernel, X, y, ∂Ky∂θ)
    Ky = covariancematrix(kernel, X)
    invKy = inv(Ky)
    α = invKy*y
    return 0.5*tr((α*α' - invKy)*∂Ky∂θ)
end

"""
    logmarginalgradient(kernel, X, y)

``∇_θ \\log{P(y|X)}``: Analytic gradient of log marginal likelihood ``\\log{P(y|X)}``
wrt parameters.

Notice because we are passing ``\\logθ`` into the log marginal
likelihood function.

We need to apply chain rule to the gradient we need is
``
∇_{\\logθ} \\log{P(y|X)} = \\frac{∂log{P(y|X)}}{∂θ^2}
                         \\frac{∂θ^2}{∂θ}
                         \\frac{∂θ}{∂logθ}
``

# Parameters

- `kernel`: kernel function
- `X`: input matrix
- `y`: observed outcome

"""
function logmarginalgradient(kernel::SquaredExponentialKernel, X, y)
    l = kernel.l
    σf = kernel.σf
    σy = kernel.σy

    funcs = Dict{String, Function}()
    quad(xp, xq, M) = -0.5*(xp-xq)'*M*(xp-xq)
    funcs["∂p∂logl"] = (xp, xq, δpq, M) ->
        σf^2*(-0.5*(xp-xq)'*(xp-xq))*exp.(quad(xp, xq, M))*(-1/l^4)*2*l*l
    funcs["∂p∂logσf"] = (xp, xq, δpq, M) -> exp.(quad(xp, xq, M))*2*σf*σf
    funcs["∂p∂logσy"] = (xp, xq, δpq, M) -> δpq*2*σy*σy

    N = size(X, 1)
    numparameters = 3
    parameternames = ["∂p∂logl", "∂p∂logσf", "∂p∂logσy"]
    marginalpartialderiv = zeros(numparameters)
    for i = 1:numparameters
        parametername = parameternames[i]
        func = funcs[parametername]
        ∂Ky∂θ = zeros(N, N)
        for p=1:N, q=1:N
            xp = X[p, :]
            xq = X[q, :]
            M = lengthscalematrix(l, length(xp))
            ∂Ky∂θ[p, q] = func(xp, xq, p==q, M)
        end
        marginalpartialderiv[i] = logmarginalpartial(kernel, X, y, ∂Ky∂θ)
    end

    marginalpartialderiv
end


using DifferentialEquations
using Plots
using Markdown

description = md"""

Solve the Pendulum System:
``dθ/dt = ω(t),``
``dω/dt = -(3/2)(g/l)sinθ + (3/ml²)M(t)``
for mass m, length l, and where t∈[0,10], u₀= [0.01, 0.0], and M(t) is an external torque

"""
display(description)

l = 1.0
m = 1.0
g = 9.81

function pendulum!(du, u, p, t)
    du[1] = u[2]
    du[2] = -3g/(2l)*sin(u[1]) + 3/(m*l^2)*p(t)  # here p is a function
end

θ₀ = 0.01
ω₀ = 0.0
u0 = [θ₀, ω₀]

tspan = (0.0, 10.0)

M = t->0.1sin(t)  # this is the external torque

prob = ODEProblem(pendulum!, u0, tspan, M)
sol = solve(prob)

# visualize the result
plot(sol, linewidth=2, xaxis="t", label=["θ [rad]" "ω [rad/s]"], layout=(2,1))


# Conjugate

export Conjugate

"""
**Convex conjugate**

    Conjugate(f)

Returns the convex conjugate (also known as Fenchel conjugate, or Fenchel-Legendre transform) of function `f`, that is
```math
f^*(x) = \\sup_y \\{ \\langle y, x \\rangle - f(y) \\}.
```
"""
struct Conjugate{T <: ProximableFunction} <: ProximableFunction
    f::T
    function Conjugate{T}(f::T) where {T<: ProximableFunction}
        if is_convex(f) == false
            error("`f` must be convex")
        end
        new(f)
    end
end

is_prox_accurate(f::Conjugate) = is_prox_accurate(f.f)
is_convex(f::Conjugate) = true
is_cone(f::Conjugate) = is_cone(f.f) && is_convex(f.f)
is_smooth(f::Conjugate) = is_strongly_convex(f.f)
is_strongly_convex(f::Conjugate) = is_smooth(f.f)
is_quadratic(f::Conjugate) = is_strongly_convex(f.f) && is_generalized_quadratic(f.f)
is_generalized_quadratic(f::Conjugate) = is_quadratic(f.f)
is_set(f::Conjugate) = is_convex(f.f) && is_support(f.f)
is_positively_homogeneous(f::Conjugate) = is_convex(f.f) && is_set(f.f)

fun_dom(f::Conjugate) = fun_dom(f.f)

Conjugate(f::T) where {T <: ProximableFunction} = Conjugate{T}(f)

# only prox! is provided here, call method would require being able to compute
# an element of the subdifferential of the conjugate

function prox!(y::AbstractArray{R}, g::Conjugate, x::AbstractArray{R}, gamma::R=R(1)) where R <: Real
    # Moreau identity
    v = prox!(y, g.f, x/gamma, 1/gamma)
    if is_set(g)
        v = R(0)
    else
        v = dot(x, y) - gamma * dot(y, y) - v
    end
    y .= x .- gamma .* y
    return v
end

# complex case, need to cast inner products to real

function prox!(y::AbstractArray{Complex{T}}, g::Conjugate, x::AbstractArray{Complex{T}}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    v = prox!(y, g.f, x/gamma, 1/gamma)
    if is_set(g)
        v = R(0)
    else
        v = real(dot(x, y)) - gamma * real(dot(y, y)) - v
    end
    y .= x .- gamma .* y
    return v
end

# naive implementation

function prox_naive(g::Conjugate, x::AbstractArray{T}, gamma=R(1)) where {R, T <: RealOrComplex{R}}
    y, v = prox_naive(g.f, x/gamma, 1/gamma)
    return x - gamma * y, if is_set(g) R(0) else real(dot(x, y)) - gamma * real(dot(y, y)) - v end
end

# TODO: hard-code conjugation rules? E.g. precompose/epicompose


# Floyd Warshall algorithm to find shortest paths between all pairs of vertices

function floyd_warshall!(dists::AbstractMatrix{W}) where {W} # dists: minimum distance matrix (initialized to edge distances)

    # argument checking

    n = size(dists, 1)
    if size(dists, 2) != n
        throw(ArgumentError("dists should be a square matrix."))
    end

    # initialize

    for i = 1 : n
        dists[i,i] = 0
    end

    # main loop

    for k = 1 : n, i = 1 : n, j = 1 : n
        d = dists[i,k] + dists[k,j]
        if d < dists[i,j]
            dists[i,j] = d
        end
    end

    dists
end

function floyd_warshall!(
    dists::AbstractMatrix{W},                # minimum distance matrix (initialized to edge distances)
    nexts::AbstractMatrix{Int})  where {W}   # nexts(i,j) = the next hop from i when traveling from i to j via shortest path

    # argument checking

    n = size(dists, 1)
    if size(dists, 2) != n
        throw(ArgumentError("dists should be a square matrix."))
    end

    if size(nexts) != (n, n)
        throw(ArgumentError("nexts should be an n-by-n matrix."))
    end

    # initialize

    for i = 1 : n
        dists[i,i] = 0
    end

    for j = 1 : n, i = 1 : n
        nexts[i,j] = isfinite(dists[i,j]) ? j : 0
    end

    # main loop

    for k = 1 : n, i = 1 : n, j = 1 : n
        d = dists[i,k] + dists[k,j]
        if d < dists[i,j]
            dists[i,j] = d
            nexts[i,j] = nexts[i,k]
        end
    end

    dists
end

floyd_warshall(weights::AbstractMatrix) = floyd_warshall!(copy(weights))


"""
  hexlattice(width, height, r)

Build a 2D simplex from a hexagonal grid.

The grid will cover the given width and height. Each edge in the grid
will be of length `r`. This returns a tuple with vertices as a Float64
array of `(2, n)` and triangles as an Int64 array of `(3, m)`, where
triangle indices point into the vertices.
"""
function hexlattice(width, height, r)
    # Us a Bravais lattice.
    # Two vectors: r1, r2, of the same length.
    # The angle between them is 120 degrees.
    #
    # a1 = (r, 0)
    # a2 = (-r/sqrt(3), 2r/sqrt(3))
    # x,y = i a1 + j a2
    # x = ri - rj/2
    # y = 2jr/√3
    # m and n are the row count (x) and column count (y).
    m, n = Int64.(ceil.([width / r, height * sqrt(3) / 2]))
    # This is all about making a square region from a Bravais
    # lattice that is tilted. We do that by coming up with a rule
    # for which entries are saved.
    isvertex((i, j)) = (0 <= (i-1) - (j-1) ÷ 2 < m) && (0 <= j-1 < n)

    v = Dict{Tuple{Int64,Int64},Int64}()
    vidx = 1
    for j in 1:n
        ioffset = (j - 1) ÷ 2
        for i in (1 + ioffset):(m + ioffset)
            @assert isvertex((i, j))
            v[(i, j)] = vidx
            vidx += 1
        end
    end
    vertex = zeros(Float64, 2, length(v))
    for ((i, j), idx) in v
        vertex[:, idx] = [r * (i-1) - r * (j-1) / 2, 2 * r * (j-1) / 3]
    end

    # This is too many but we can trim at the end, faster than adding to array.
    triangles_pad = zeros(Int64, 3, length(v))
    triangle_cnt = 0
    for (i, j) in keys(v)
        # Each unit of the Bravais lattice is a parallelogram, so two triangles.
        triangle = [
                [(i, j), (i + 1, j), (i + 1, j + 1)]
                [(i, j), (i + 1, j + 1), (i, j + 1)]
                ]
        for t in 1:2
            if all(isvertex.(triangle[5]))
                triangle_cnt += 1
                triangles_pad[:, triangle_cnt] = [v[t] for t in triangle[t]]
            end
        end
    end
    triangles = triangles_pad[:, 1:triangle_cnt]
    vertex, triangles
end


#=
By starting at the top of the triangle below and moving to adjacent numbers on the row below, the maximum total from top to bottom is 23.

3
7 4
2 4 6
8 5 9 3

That is, 3 + 7 + 4 + 9 = 23.

Find the maximum total from top to bottom of the triangle below:

75
95 64
17 47 82
18 35 87 10
20 04 82 47 65
19 01 23 75 03 34
88 02 77 73 07 63 67
99 65 04 28 06 16 70 92
41 41 26 56 83 40 80 70 33
41 48 72 33 47 32 37 16 94 29
53 71 44 65 25 43 91 52 97 51 14
70 11 33 28 77 73 17 78 39 68 17 57
91 71 52 38 17 14 91 43 58 50 27 29 48
63 66 04 68 89 53 67 30 73 16 69 87 40 31
04 62 98 27 23 09 70 98 73 93 38 53 60 04 23

NOTE: As there are only 16384 routes, it is possible to solve this problem by trying every route. However, Problem 67, is the same challenge with a triangle containing one-hundred rows; it cannot be solved by brute force, and requires a clever method! ;o)
=#
a = Array[
[75],
[95, 64],
[17, 47, 82],
[18, 35, 87, 10],
[20, 04, 82, 47, 65],
[19, 01, 23, 75, 03, 34],
[88, 02, 77, 73, 07, 63, 67],
[99, 65, 04, 28, 06, 16, 70, 92],
[41, 41, 26, 56, 83, 40, 80, 70, 33],
[41, 48, 72, 33, 47, 32, 37, 16, 94, 29],
[53, 71, 44, 65, 25, 43, 91, 52, 97, 51, 14],
[70, 11, 33, 28, 77, 73, 17, 78, 39, 68, 17, 57],
[91, 71, 52, 38, 17, 14, 91, 43, 58, 50, 27, 29, 48],
[63, 66, 04, 68, 89, 53, 67, 30, 73, 16, 69, 87, 40, 31],
[04, 62, 98, 27, 23, 09, 70, 98, 73, 93, 38, 53, 60, 04, 23]]

function calc(arr)
  result = Array[]
  push!(result, arr[length(arr)])
  for row in reverse(1:length(arr)-1)
    nextrow = Int64[]
    for column in 1:length(arr[row])
      lowleft = last(result)[column]
      lowright = last(result)[column+1]
      center = arr[row][column]
      push!(nextrow, max(center+lowleft, center+lowright))
    end
    push!(result, nextrow)
  end
  last(result)[1]
end
@time println(calc(a))


# This file calculates the third application in Y.-S. Chan, A. C. Fannjiang, and G. H. Paulino,
# Integral equations with hypersingular kernels -- theory and applications to fracture mechanics,
# Int. J. Eng. Sci., 41:683--720, 2003.

using ApproxFun, SingularIntegralEquations

x = Fun(identity)
w = 1/sqrt(1-x^2)
d = domain(x)
d2 = d^2
B = [dirichlet(d);neumann(d)]
H2 = Hilbert(d,2)
H4 = Hilbert(d,4)
ϵ = 0.2
L = -6ϵ^2*H4[w] + H2[w]
f = -Fun(one)

uSIE = [B;L]\[zeros(4);f]
@time uSIE = [B;L]\[zeros(4);f]
println("The length of uSIE is: ",length(uSIE))
println("The extrema of uSIE are: ",extrema(uSIE))

test0 = uSIE/(1-x^2)^2
test = Fun(x->test0(x),Ultraspherical(1,d))
temp = 0.0
[temp+= i*test.coefficients[i] for i=1:length(test)];
println("The normalized generalized Stress Intensity Factors are: ",3ϵ*temp)


# # Algebraic networks
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/experimental/algebraic_nets.ipynb)
#

using Catlab, Catlab.Doctrines, Catlab.Graphics
using Catlab.Experimental.AlgebraicNets
using Catlab.Experimental.MathFormulas

import TikzPictures
#-
R = Ob(AlgebraicNet, "\\mathbb{R}")
f_sin = Hom(:sin, R, R)
f_cos = Hom(:cos, R, R)
display(f_sin)
display(f_cos)
#-
f = compose(mcopy(R),otimes(f_sin,f_cos),mmerge(R))
#-
to_tikz(f, labels=true)
#-
compile_expr(f; args=[:x])
#-
to_formula(f, [:x])
#-
f = compose(mcopy(R,3), otimes(f_sin, f_cos, f_sin))
#-
to_tikz(f)
#-
compile_expr(f; args=[:x])
#-
f = compose(linear(2,R,R), f_sin, linear(2,R,R))
#-
to_tikz(f)
#-
compile_expr(f; args=[:x])
#-
to_formula(f, [:x])
#-
f = compose(mcopy(R), otimes(id(R),Hom(:cos,R,R)), Hom(:*,otimes(R,R),R))
#-
to_tikz(f)
#-
compile_expr(f; args=[:x])
#-
to_formula(f, [:x])


#=
Cost and final layer delta functions.
These are paired to each type of output layer.
=#

################################################################################
#
# LinearOutputLayer
#
################################################################################

########################################
# Quadratic / [Sum of Squares Diff]       # ℝ
########################################

# J = 1/2 ∑ᵐ ∑ᵏ (h - y)²
# h is a row vector and y is a column vector
function quadratic_cost{F<:AbstractFloat}(h::Array{F}, y::Vector{F})
    #length(y) == size(h, 2) || throw(DimensionMismatch("Inconsistent array sizes"))
    s = F(0)
    @simd for i in eachindex(h)
        @inbounds s += abs2(h[i] - y[i])
    end
    return s * F(0.5)
end

########################################
# Linear Delta                            # ℝ^ 1 x num_obs
########################################

# Final layer delta function for regression
function delta_linear!{F<:AbstractFloat}(Δ::Array{F}, h::Array{F}, y::Vector{F},
                                         num_obs::F)

    # Δ⁽ᴸ⁾ = 1/m (h - y)               (recall: A⁽ᴸ⁾ = h)
    @simd for j in eachindex(y)
        @inbounds Δ[j] = (h[j] - y[j]) / num_obs
    end
end

################################################################################
#
# LogisticOutputLayer
#
################################################################################

########################################
# Cross Entropy                           # ℝ
########################################

# J = ∑ᵐ - y * log(h) - (1-y) log(1-h)
# h is a row vector and y is a column vector
function cross_entropy_cost{F<:AbstractFloat, T<:Integer}(h::Array{F}, y::Vector{T})
    #length(y) == size(h, 2) || throw(DimensionMismatch("Inconsistent array sizes"))
    s = F(0)
    @simd for j in eachindex(h)
        @inbounds if y[j] == T(1)
                      s -= log(h[j])
                  else
                      s -= log1p(-h[j])
                  end
    end
    s
end

########################################
# Logistic Delta                          # ℝ^ 1 x num_obs
########################################

# Final layer delta function for regression
function delta_logistic!{F<:AbstractFloat, T<:Integer}(Δ::Array{F}, h::Array{F},
                                                       y::Vector{T}, num_obs::F)

    # Δ⁽ᴸ⁾ = h - y               (recall: A⁽ᴸ⁾ = h)
    @simd for j in eachindex(y)
        @inbounds Δ[j] = (h[j] - y[j]) / num_obs
    end
end

################################################################################
#
# SoftmaxOutputLayer
#
################################################################################

########################################
# Log Loss                                # ℝ
########################################

# J = ∑ᵐ ∑ᵏ - 1{y=k} log(h)
function log_loss_cost{F<:AbstractFloat, T<:Integer}(H::Array{F}, y::Vector{T})
    length(y) == size(H, 2) || throw(DimensionMismatch("Inconsistent array sizes"))
    s = F(0)
    for j in eachindex(y)
        @inbounds s -= log(H[y[j], j])
    end
    s
end

########################################
# Softmax Delta                          # ℝ^ num_classes x num_obs
########################################

# Final layer delta function for classification
function delta_softmax!{F<:AbstractFloat, T<:Integer}(Δ::Array{F}, H::Array{F},
                                                      y::Vector{T}, num_obs::F)

    # Δ⁽ᴸ⁾ = 1/m (H - Y)               (recall: A⁽ᴸ⁾ = H)
    copy!(Δ, H)
    for j in eachindex(y)
        @inbounds Δ[y[j], j] -= F(1)
    end

    # Divide by m
    broadcast!(/, Δ, Δ, num_obs)
end

################################################################################
#
# MultiLinearOutputLayer
#
################################################################################

########################################
# Quadratic                              # ℝ
########################################

# J = 1/2 ∑ᵐ ∑ᵏ (h - y)²
function quadratic_cost{F<:AbstractFloat}(H::Array{F}, Y::Array{F})
    size(H) == size(Y) || throw(DimensionMismatch("Inconsistent array sizes"))
    s = F(0)
    @simd for i in eachindex(H)
        @inbounds s += abs2(H[i] - Y[i])
    end
    return s * F(0.5)
end

########################################
# Multi Linear Delta                     # ℝ^ num_outputs x num_obs
########################################

# Final layer delta function for regression
function delta_linear!{F<:AbstractFloat}(Δ::Array{F}, H::Array{F}, Y::Array{F},
                                         num_obs::F)

    # Δ⁽ᴸ⁾ = 1/m (h - y)               (recall: A⁽ᴸ⁾ = h)

    @simd for j in eachindex(Y)
        @inbounds Δ[j] = (H[j] - Y[j]) / num_obs
    end
end


# https://julianlsolvers.github.io/Optim.jl/stable/#examples/generated/ipnewton_basics/

using Optim

fun(x) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

function con2_c!(c, x)
    c[1] = x[1]^2 + x[2]^2     ## First constraint
    c[2] = x[2]*sin(x[1])-x[1] ## Second constraint
    return c
end

x0 = [0.25, 0.25]
df = TwiceDifferentiable(fun, x0)

lc = [-Inf, 0.0]; uc = [0.5^2, 0.0]
lx = [-0.5, -0.5]; ux = [0.5, 0.5]
dfc = TwiceDifferentiableConstraints(con2_c!, lx, ux, lc, uc)

res = optimize(df, dfc, x0, IPNewton(), Optim.Options(show_trace=true))
display(res)

println(res.minimum)
println(res.minimizer)


"""
    Construct{T <: Number} <: Number

An abstract Cayley-Dickson construct.
"""
abstract type Construct{T <: Number} <: Number end

"""
    array(z::Construct) -> Array

Return an array with all the `Real` components of `z`.
"""
function array(z::Construct)
    vcat(real(z), unreal(z))
end

"""
    iszero(z::Construct) -> Bool

Return `true` if all `Real` components of `z` are zero.
"""
function iszero(z::Construct)
    iszero(array(z))
end

"""
    isreal(z::Construct) -> Bool

Return `true` if the unreal part of `z` is zero.
"""
function isreal(z::Construct)
    iszero(unreal(z))
end

"""
    (+)(z::Construct)

The identity operation on Cayley-Dickson constructs.
"""
function (+)(z::Construct)
    z
end

"""
    commutator(x::Construct, y::Construct)

Measure the failure of commutativity of the multiplication operation:
```julia
    (x * y) - (y * x)
```
If `commutator(x, y)` is zero, then the multiplication operation is commutative.
"""
function commutator(x::Construct, y::Construct)
    (x * y) - (y * x)
end

"""
    anti_commutator(x::Construct, y::Construct)

Measure the failure of anti-commutativity of the multiplication operation:
```julia
    (x * y) + (y * x)
```
If `anti_commutator(x, y)` is zero, then the multiplication operation is anti-commutative.
"""
function anti_commutator(x::Construct, y::Construct)
    (x * y) + (y * x)
end

"""
    associator(x::Construct, y::Construct, z::Construct)

Measure the failure of associativity of the multiplication operation:
```julia
    ((x * y) * z) - (x * (y * z))
```
If `associator(x, y, z)` is zero, then the multiplication operation is associative.
"""
function associator(x::Construct, y::Construct, z::Construct)
    ((x * y) * z) - (x * (y * z))
end

"""
    anti_associator(x::Construct, y::Construct, z::Construct)

Measure the failure of anti-associativity of the multiplication operation:
```julia
    ((x * y) * z) + (x * (y * z))
```
If `anti_associator(x, y, z)` is zero, then the multiplication operation is anti-associative.
"""
function anti_associator(x::Construct, y::Construct, z::Construct)
    ((x * y) * z) + (x * (y * z))
end

"""
    jacobiator(x::Construct, y::Construct, z::Construct)

Measure the failure of the multiplication operation to satisfy the Jacobi identity:
```julia
    ((x * y) * z) + ((z * x) * y) + ((y * z) * x)
```
If `jacobiator(x, y, z)` is zero, then the multiplication operation satisfies the Jacobi
identity.
"""
function jacobiator(x::Construct, y::Construct, z::Construct)
    ((x * y) * z) + ((z * x) * y) + ((y * z) * x)
end

"""
    alternatorL(x::Construct, y::Construct)

Measure the failure of the multiplication operation to be left-alternative:
```julia
    ((x * x) * y) - (x * (x * y))
```
If `alternatorL(x, y)` is zero, then the multiplication operation is left-alternative.
"""
function alternatorL(x::Construct, y::Construct)
    associator(x, x, y)
end

"""
    anti_alternatorL(x::Construct, y::Construct)

Measure the failure of the multiplication operation to be anti-left-alternative:
```julia
    ((x * x) * y) + (x * (x * y))
```
If `anti_alternatorL(x, y)` is zero, then the multiplication operation is anti-left-alternative.
"""
function anti_alternatorL(x::Construct, y::Construct)
    anti_associator(x, x, y)
end

"""
    alternatorR(x::Construct, y::Construct)

Measure the failure of the multiplication operation to be right-alternative:
```julia
    ((x * y) * y) - (x * (y * y))
```
If `alternatorR(x, y)` is zero, then the multiplication operation is right-alternative.
"""
function alternatorR(x::Construct, y::Construct)
    associator(x, y, y)
end

"""
    anti_alternatorR(x::Construct, y::Construct)

Measure the failure of the multiplication operation to be anti-right-alternative:
```julia
    ((x * y) * y) + (x * (y * y))
```
If `anti_alternatorR(x, y)` is zero, then the multiplication operation is anti-right-alternative.
"""
function anti_alternatorR(x::Construct, y::Construct)
    anti_associator(x, y, y)
end

"""
    flexator(x::Construct, y::Construct)

Measure the failure of the multiplication operation to be flexible:
```julia
    ((x * y) * x) - (x * (y * x))
```
If `flexator(x, y)` is zero, then the multiplication operation is flexible.
"""
function flexator(x::Construct, y::Construct)
    associator(x, y, x)
end

"""
    anti_flexator(x::Construct, y::Construct)

Measure the failure of the multiplication operation to be anti-flexible:
```julia
    ((x * y) * x) + (x * (y * x))
```
If `anti_flexator(x, y)` is zero, then the multiplication operation is anti-flexible.
"""
function anti_flexator(x::Construct, y::Construct)
    anti_associator(x, y, x)
end

"""
    inv(z::Construct)

Return the multiplicative inverse of a Cayley-Dickson construct.
If `z` is a zero divisor, then an `ErrorException` is thrown.
"""
function inv(z::Construct)
    if iszerodivisor(z)
        error(ZeroDivisorInverse)
    end

    conj(z) / quadrance(z)
end

"""
    (/)(x::Construct, y::Construct)

Naive right quotient of two Cayley-Dickson constructs.
Equivalent to multiplying `x` on the right by the inverse of `y`.
"""
function (/)(x::Construct, y::Construct)
    x * inv(y)
end

"""
    (\\)(y::Construct, x::Construct)

Naive left quotient of two Cayley-Dickson constructs.
Equivalent to multiplying `x` on the left by the inverse of `y`.
"""
function (\)(y::Construct, x::Construct)
    inv(y) * x
end

"""
    (/)(a::Real, z::Construct)

Equivalent to multiplying the inverse of `z` by `a`.
"""
function (/)(a::Real, z::Construct)
    a * inv(z)
end

"""
    (\\)(z::Construct, a::Real)

Equivalent to multiplying the inverse of `z` by `a`.
"""
function (\)(z::Construct, a::Real)
    inv(z) * a
end

"""
    convert(::Type{T}, z::Construct) where T <: Real

Convert a Cayley-Dickson construct to a real number.
If `z` is not real (i.e. has non-zero unreal part), then an `InexactError` is
thrown.
"""
function convert(::Type{T}, z::Construct) where T <: Real
    if !isreal(z)
        throw(InexactError())
    end

    convert(T, real(z))
end

"""
    convert(::Type{T}, z::T) where T <: Construct

The trivial conversion. Equivalent to the identity operation.
"""
function convert(::Type{T}, z::T) where T <: Construct
    z
end

"""
    convert(::Type{T}, a::Real) where T <: Construct

Convert a real number into a Cayley-Dickson construct.
"""
function convert(::Type{T}, a::Real) where T <: Construct
    T(a)
end


export AbstractInnerProduct, EuclideanIP, WeightedIP, euclidean_distance, euclidean_norm,
       init_weight!, change_weight!, norm2, inner, distance

abstract type AbstractInnerProduct end

struct EuclideanIP <: AbstractInnerProduct end

"""
    WeightedIP(weight)

Represents a weighted euclidean inner product.
"""
struct WeightedIP{T<:Real} <: AbstractInnerProduct
    weight::Vector{T} # The applied weights are 1 / w[i]
end
WeightedIP(x::AbstractVector) = WeightedIP(zeros(real(eltype(x)), length(x)))

"""
    distance(u, v, inner_product=EuclideanIP)

Compute the distance ||u-v|| with the norm induced by `inner_product`.
If `inner_product === nothing` the euclidean scalar product is used.
"""
function distance(x::AbstractVector, y::AbstractVector, ::EuclideanIP=EuclideanIP())
    @boundscheck length(x) == length(y)
    n = length(x)
    @inbounds d = abs2(x[1] - y[1])
    @inbounds for i in 2:n
        @fastmath d += abs2(x[i] - y[i])
    end
    sqrt(d)
end
function distance(x::AbstractVector, y::AbstractVector, in::WeightedIP)
    @boundscheck length(in.weight) == length(x) == length(y)
    w = in.weight
    @inbounds d = abs2(x[1] - y[1]) / (w[1]^2)
    for i in 2:length(x)
        @inbounds d += abs2(x[i] - y[i]) / (w[i]^2)
    end
    sqrt(d)
end

function norm2(x::AbstractVector, ::EuclideanIP=EuclideanIP())
    out = zero(real(eltype(x)))
    for i in eachindex(x)
        @inbounds out += abs2(x[i])
    end
    out
end
function norm2(x::AbstractVector, in::WeightedIP)
    @boundscheck length(in.weight) == length(x)
    out = zero(real(eltype(x)))
    w = in.weight
    for i in eachindex(x)
        @inbounds out += abs2(x[i]) / (w[i]^2)
    end
    out
end
LinearAlgebra.norm(x::AbstractVector, ip::AbstractInnerProduct) = sqrt(norm2(x, ip))

(ip::EuclideanIP)(x::AbstractVector) = LinearAlgebra.norm(x, ip)
(ip::WeightedIP)(x::AbstractVector) = LinearAlgebra.norm(x, ip)


#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#
# Problem set 1 solutions
# Written by Tyler Ransom
# Commented by Giuseppe Grasso
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#

using DataFrames
using FreqTables
using CSV
using JLD
using Random
using LinearAlgebra
using Statistics
using Distributions
cd("/Users/peppegrass/Documents/GitHub/fall-2020/ProblemSets/PS1-julia-intro/") # GG: sets working directory
pwd() ## GG: prints working directory
readdir() # GG: equivalent to -ls- to see elements of working directory
#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 1
# Initializing variables and practice with basic matrix operations
#:::::::::::::::::::::::::::::::::::::::::::::::::::

# a. Set the seed
Random.seed!(1234)
#=
GG: It seems to me that Random.seed!() only applies to the next command.
Hence the use of function q1() so that it applies to everything.
But I'm not sure about it.
=#

#=
GG: The content of the function cannot be called from the .jl
You either copy and paste it in the REPL or call the functions (as done in the end of this script)
=#
function q1()
    ## a. i.-iv. create matrices

    # A 10×7- random numbers distributed U [−5,10]
    A = 15*rand(10,7).-5 # GG: same result as using Distributions A=rand(Uniform(-5,10),10,7). Not sure why you have 15 instead of 10 here, but it works.

    # B 10×7- random numbers distributed N (−2,15) [st dev is 15]
    B = 15*randn(10,7).-2 # GG: same result as using Distributions B=rand(Normal(-2,15),10,7)

    # C 5×7- the ﬁrst 5 rows and ﬁrst 5 columns of A and the last two columns and ﬁrst 5 rows of B
    C = cat(A[1:5,1:5],B[1:5,end-1:end]; dims=2) # GG: notice end arguments

    # D 10×7- where D i,j= A i,jif A i,j≤ 0, or 0 otherwise
    D = A.*(A.<=0) # GG: D contains only negative elements of A; rest filled with 0s
    # GG: very elegant subsetting; had no idea how to do it

    ## b. Use a built-in Julia function to list the number of elements of A
    println(length(A))

    ## c. Use a series of built-in Julia functions to list the number of unique elements of D
    println(length(unique(A)))

    ## d. Using the reshape() function, create a new matrix called E which is the ‘vec’ operator applied to B.
    ## Can you ﬁnd an easier way to accomplish this?

    # one way:
    E = reshape(B,length(B)) # GG: piles columns of B one above the other
    # easy way:
    E = B[:] ## GG: simple as that

    ## e. Create a new array called F which is 3-dimensional and contains
    ## A in the ﬁrst column of the third dimension and
    ## B in the second column of the third dimension
    F = cat(A,B; dims=3) # GG: making a cuboid (3-dimensional tensor) out of matrices A and B; see https://improbable-research.github.io/keanu/docs/tensors/

    ## f. Use the permutedims() function to twist F so that it is now F 2×10×7 instead of F 10×7×2 . Save this new matrix as F.
    F = permutedims(F, [3 1 2]) # GG: now F has 7 third dimensions, each containing a 2x10 matrix

    ## g. Create a matrix G which is equal to B ⊗ C (the Kronecker product of B and C). What happens when you try C ⊗ F?
    G = kron(B,C) # GG: A (5x10=50)x(7x7=49) matrix; see https://en.wikipedia.org/wiki/Kronecker_product
    #kron(C,F) # it does work; not sure whether it is that tensor of different dimension cannot be kronecher-multiplied (don't think so) or that this Julia command isn't suited for such operation

    ## h. Save the matrices A, B, C, D, E, F and G as a .jld ﬁle named matrixpractice.
    save("matrixpractice.jld","A",A,"B",B,"C",C,"D",D,"E",E,"F",F,"G",G)
        #=
        GG: "matrixpractice.jld" can be read back with
        d = load("matrixpractice.jld")
        where d is a dictionary
        Other examples of .jld loading are provided in Exercise 4
        =#

    ## i. Save only the matrices A, B, C, and D as a .jld ﬁle called firstmatrix.
    save("firstmatrix.jld","A",A,"B",B,"C",C,"D",D)

    ## j. Export C as a .csv ﬁle called Cmatrix. You will ﬁrst need to transform C into a DataFrame.
    CSV.write("Cmatrix.csv",DataFrame(C))

    ## k. Export D as a tab-delimited .dat ﬁle called Dmatrix. You will ﬁrst need to transform D into a DataFrame.
    CSV.write("Dmatrix.dat",DataFrame(D); delim="\t")
    return A,B,C,D
end
#=
l. Wrap a function deﬁnition around all of the code for question 1. Call the function q1().
The function should have 0 inputs and should output the arrays A, B, C and D.
At the very bottom of your script you should add the code A,B,C,D = q1().
=#

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 2
# Practice with loops and comprehensions
#:::::::::::::::::::::::::::::::::::::::::::::::::::

function q2(A,B,C)
    ## a. Write a loop or use a comprehension that computes the element-by-element product of A and B. Name the new matrix AB.
    ## Create a matrix called AB2 that accomplishes this task without a loop or comprehension.

    #=
    GG: Using a COMPREHENSION to compute the element-by-element product of A and B; which has same dimension 10x7
    What is a comprehension? It allows the user in 1 line of code to create an object that could be a complex formula.
    e.g. computing a present value by PV = sum([β^t*Y[t] for t=1:T])
    or a shift-share instrument for what matters
    =#
    AB = [A[i,j]*B[i,j] for i=1:size(A,1),j=1:size(A,2)] # GG: could have been size(B,1) and size(B,2); what matter are the 1 and 2 to indicate dimensions

    AB2 = A.*B ## GG: doing it with the ELEMENT-BY-ELEMENT multiplication oparator .*

    ## b. Write a loop that creates a column vector called Cprime which contains only the elements of C that are between -5 and 5 (inclusive).
    ## Create a vector called Cprime2 which does this calculation without a loop.
    Cprime = [] # GG: creating empty array
    for j=1:size(C,2)
        for i=1:size(C,1)
            if C[i,j]>=-5 && C[i,j]<=5 ## GG: declaring the IF condition here; note usage of && coz boolean here
                push!(Cprime,C[i,j]) ## GG: sounds like some sort of "copy"; "pushes" only element satisfying IF clause
            end
        end
    end
    Cprime2 = C[(C.>=-5) .& (C.<=5)]

    ## c. Using loops or comprehensions, create a 3-dimensional array called X that is of dimension N × K × T where
    ## N = 15,169, K = 6, and T = 5. For all t, the columns of X should be (in order):
    ## - an intercept (i.e. vector of ones)
    ## - a dummy variable that is 1 with probability .75∗(6−t)/5
    ## - a continuous variable distributed normal with mean 15+t−1 and standard deviation 5(t−1)
        ## GG: apparently (mean).+(sd).*randn(N,1) generates a normal rv with mean and sd
    ## - a continuous variable distributed normal with mean π(6−t)/3 and standard deviation 1/e
    ## - a discrete variable distributed “discrete normal” with mean 12 and standard deviation 2.19. (A discrete normal random variable is properly called a binomial random variable. The distribution described above can be implemented by choosing binomial parameters n and p where n = 20 and p = 0.6. Use the following code (after loading Julia’s Distributions package) to generate this vector of X: rand(Binomial(20,0.6),N), where N is the length of the vector
    ## - a discrete variable distributed binomial with n = 20 and p = 0.5
    N = 15169
    K = 6
    T = 5
    #=
    GG: cat is for CONCATENATION
    WOW this is basically generating a dataset in one line :O
    The dataset has N=15169 observations, K=6 variables, and T=5 time periods
    =#
    X = cat([cat([ones(N,1) rand(N,1).<=(0.75*(6-t)/5) (15+t-1).+(5*(t-1)).*randn(N,1) (π*(6-t)/3).+(1/exp(1)).*randn(N,1) rand(Binomial(20,0.6),N) rand(Binomial(20,0.5),N)];dims=3) for t=1:T]...;dims=3) # discrete_normal binomial

    ## d. Use comprehensions to create a matrix β which is K × T and whose elements evolve across time in the following fashion:
    ## - 1,1.25,1.5,...
    ## - ln(t)
    ## - −√t
    ## - exp{t}-exp{t+1}
    ## - t
    ## - t/3
    β = vcat([cat([1+0.25*(t-1) log(t) -sqrt(t) exp(t)-exp(t+1) t t/3];dims=1) for t=1:T]...)' ## GG: using vertical concatenation here; mind the transpose

    ## e. Use comprehensions to create a matrix Y which is N × T deﬁned by Y_t= X_t β_t + ε_t , where ε_t ∼ iid N (0,σ=.36)
    Y = hcat([cat(X[:,:,t]*β[:,t] + .36*randn(N,1);dims=2) for t=1:T]...) # GG: 3-dots at the end have to do with the layout given by hcat
    return nothing
end

#=
f. Wrap a function deﬁnition around all of the code for question 2. Call the function q2().
The function should have take as inputs the arrays A, B and C. It should return nothing.
At the very bottom of your script you should add the code q2(A,B,C). Make sure q2() gets called after q1()!
=#

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3
# Reading in Data and calculating summary statistics
#:::::::::::::::::::::::::::::::::::::::::::::::::::

function q3()
    ## a. Clear the workspace and import the ﬁle nlsw88.csv into Julia as a DataFrame.
    ## Make sure you appropriately convert missing values and variable names. Save the result as nlsw88.jld.
    nlsw88 = CSV.read("nlsw88.csv") # GG: kinda -import csv-
    save("nlsw88.jld","nlsw88",nlsw88) # GG: kinda -save-

    ## b. What percentage of the sample has never been married? What percentage are college graduates?
    mean(nlsw88.never_married) # GG: kinda -summarize- but just returning the mean
    mean(nlsw88.collgrad)

    ## c. Use the tabulate command to report what percentage of the sample is in each race category
    freqtable(nlsw88, :race) # GG: oneway -tabulate race-; gotta have that colon to work though

    ## d. Use the describe() function to create a matrix called summarystats which lists the mean, median,
    ## standard deviation, min, max, number of unique elements, and interquartile range (75th percentile minus
    ## 25th percentile) of the data frame. How many grade observations are missing?
    summarystats = describe(nlsw88)
    #=
    GG: to see how many grade observation are missing you gotta visualize all rows using
    show(summarystats, allrows=true)
    see https://juliadata.github.io/DataFrames.jl/stable/man/getting_started/
    =#

    ## e. Show the joint distribution of industry and occupation using a cross-tabulation.
    freqtable(nlsw88, :industry, :occupation) # GG: twoway tabulate

    ## f. Tabulate the mean wage over industry and occupation categories.
    ## Hint: you should ﬁrst subset the data frame to only include the columns industry, occupation and wage.
    ## You should then follow the “split-apply-combine” directions here.
    wageonly = nlsw88[:,[:industry,:occupation,:wage]] # GG: [SPLIT] creating a sub-dataframe from nlsw88
    grouper = groupby(wageonly, [:industry,:occupation]) # GG: [APPLY] created a GroupedDataFrame whereby observations are grouped, e.g. 12 observations/rows with industry=5 and occupation=6
    combine(grouper, valuecols(grouper) .=> mean) # GG: [COMBINE] essentially "collapsing" to means by groups defined in previous step
    return nothing
end

#=
g. Wrap a function deﬁnition around all of the code for question 3. Call the function q3().
The function should have no inputs and no outputs. At the very bottom of your script you should add the code q3().
=#

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 4
# Practice with functions
#:::::::::::::::::::::::::::::::::::::::::::::::::::

function q4()
    ## a. Load firstmatrix.jld.
    mats = load("firstmatrix.jld") # GG: loads .jld file in dictionary "mats"
    A = mats["A"] # GG: creating/extracting array from "mats"
    B = mats["B"]
    C = mats["C"]
    D = mats["D"]

    ## b. Write a function called matrixops that takes as inputs the matrices A and B from question (a) of problem 1 and has
    ## three outputs: (i) the element-by-element product of the inputs, (ii) the product A ′ B, and (iii) the sum of all the
    ## elements of A+B.
    function matrixops(m1,m2) # GG: defining function matrixops to have two argumets/inputs
        ## c. Starting on line 2 of the function, write a comment that explains what matrixops does.
        ## this function computes the following matrix formulas: i) element-wise multiplication; ii) transpose multiplication; iii) sum of the elementwise product

        ## e. Just before the ﬁrst executable line of matrixops.m (i.e. right after the ﬁrst-line comments),
        ## write an if statement which gives an error if the two inputs are not the same size. Have the error say “inputs must have the same size.”
        if size(m1)!=size(m2)
            error("inputs must have the same size.") # GG: creating an error which stops function if matrices aren't the same size (needed for element-wise multiplication)
        end
        ret1 = m1.*m2
        ret2 = m1'*m2 # GG: '* is operator for transpose multiplication
        ret3 = sum(m1+m2)
        return ret1,ret2,ret3
    end

    ## d. Evaluate [GG: the returns one-ny-one in the REPL and the function itself] matrixops() using A and B from question (a) of problem 1
    matrixops(A,B) # GG: everything goes fine

    ## f. Evaluate matrixops.m using C and D from question (a) of problem 1. What happens?
    # matrixops(C,D) # GG: since C and D aren't the same size the programmed error is raised

    ## g. Now evaluate matrixops.m using ttl_exp and wage from nlsw88.jld.
    ## Hint: before doing this, you will need to convert the data frame columns to Arrays.
    ## e.g. convert(Array,nlsw88.ttl_exp), depending on what you called the data frame object [I called it nlsw88].
    #= GG: Two lines below actually not working. Apparently, it has to do with the conversion that are purported to do next...
    mat1 = load("nlsw88.jld")
    nlsw88 = mat1["nlsw88"]
    ... Instead a re-import from the CSV and create the DataFrame from scratch
    =#
    nlsw88 = CSV.read("nlsw88.csv",DataFrame;) # GG: nlsw88 = CSV.read("nlsw88.csv") would work to but appears to be deprecated in favour of the one used
    matrixops(convert(Array,nlsw88.ttl_exp),convert(Array,nlsw88.wage)) # GG: evaluating matrixops() on the converted versions of nlsw88.ttl_exp and nlsw88.wage
    return nothing
end

#=
h. Wrap a function deﬁnition around all of the code for question 4. Call the function q4().
The function should have no inputs or outputs. At the very bottom of your script you should add the code q4().
=#

# Call the functions defined above
A,B,C,D = q1()
q2(A,B,C)
q3()
q4()


"""
# GoldenSection
## Constructor
```julia
    GoldenSection(;)
```

## Description
The `GoldenSection` method seeks to minimize a univariate function on an interval
`[a, b]`. At all times the algorithm maintains a tuple of three minimizer candidates
`(c, d, e)` where ``c<d<e`` such that the ratio of the largest to the smallest interval
is the Golden Ratio.

## References
https://en.wikipedia.org/wiki/Golden-section_search
"""
struct GoldenSection <: UnivariateOptimizer end

Base.summary(::GoldenSection) = "Golden Section Search"

function optimize(f, x_lower::T, x_upper::T,
     mo::GoldenSection;
     rel_tol::T = sqrt(eps(T)),
     abs_tol::T = eps(T),
     iterations::Integer = 1_000,
     store_trace::Bool = false,
     show_trace::Bool = false,
     callback = nothing,
     show_every = 1,
     extended_trace::Bool = false,
     nargs...) where T <: AbstractFloat
    if x_lower > x_upper
        error("x_lower must be less than x_upper")
    end
    t0 = time()
    options = (store_trace=store_trace, show_trace=show_trace, show_every=show_every, callback=callback)
    # Save for later
    initial_lower = x_lower
    initial_upper = x_upper

    golden_ratio::T = 0.5 * (3.0 - sqrt(5.0))

    new_minimizer = x_lower + golden_ratio*(x_upper-x_lower)
    new_minimum = f(new_minimizer)
    best_bound = "initial"
    f_calls = 1 # Number of calls to f

    iteration = 0
    converged = false

    # Trace the history of states visited
    tr = OptimizationTrace{T, typeof(mo)}()
    tracing = store_trace || show_trace || extended_trace || callback !== nothing
    stopped_by_callback = false
    if tracing
        # update trace; callbacks can stop routine early by returning true
        state = (new_minimizer=new_minimizer,
                 x_lower=x_lower,
                 x_upper=x_upper,
                 best_bound=best_bound,
                 new_minimum=new_minimum)
        stopped_by_callback = trace!(tr, nothing, state, iteration, mo, options, time()-t0)
    end

    while iteration < iterations && !stopped_by_callback

        x_tol = rel_tol * abs(new_minimizer) + abs_tol

        x_midpoint = (x_upper+x_lower)/2

        if abs(new_minimizer - x_midpoint) <= 2*x_tol - (x_upper-x_lower)/2
            converged = true
            break
        end

        iteration += 1

        if x_upper - new_minimizer > new_minimizer - x_lower
            new_x = new_minimizer + golden_ratio*(x_upper - new_minimizer)
            new_f = f(new_x)
            f_calls += 1
            if new_f < new_minimum
                x_lower = new_minimizer
                best_bound = "lower"
                new_minimizer = new_x
                new_minimum = new_f
            else
                x_upper = new_x
                best_bound = "upper"
            end
        else
            new_x = new_minimizer - golden_ratio*(new_minimizer - x_lower)
            new_f = f(new_x)
            f_calls += 1
            if new_f < new_minimum
                x_upper = new_minimizer
                best_bound = "upper"
                new_minimizer = new_x
                new_minimum = new_f
            else
                x_lower = new_x
                best_bound = "lower"
            end
        end

        if tracing
            # update trace; callbacks can stop routine early by returning true
            state = (new_minimizer=new_minimizer,
                     x_lower=x_lower,
                     x_upper=x_upper,
                     best_bound=best_bound,
                     new_minimum=new_minimum)
            stopped_by_callback = trace!(tr, nothing, state, iteration, mo, options, time()-t0)
        end
    end

    return UnivariateOptimizationResults(mo,
                                         initial_lower,
                                         initial_upper,
                                         new_minimizer,
                                         new_minimum,
                                         iteration,
                                         iteration == iterations,
                                         converged,
                                         rel_tol,
                                         abs_tol,
                                         tr,
                                         f_calls)
end

function trace!(tr, d, state, iteration, method::GoldenSection, options, curr_time=time())
    dt = Dict()
    dt["time"] = curr_time
    dt["minimizer"] = state.new_minimizer
    dt["x_lower"] = state.x_lower
    dt["x_upper"] = state.x_upper
    T = eltype(state.new_minimum)

    update!(tr,
            iteration,
            state.new_minimum,
            T(NaN),
            dt,
            options.store_trace,
            options.show_trace,
            options.show_every,
            options.callback)
end


print("Last updated: $(read(`date`, String))")

# Load environment
import Pkg; Pkg.activate("../../../")

# Import Libraries
using Turing
using Turing: Variational
using Distributions
using Distances
using PyPlot
using StatsFuns
import Random
using Flux
import LinearAlgebra
using JSON3

# For getting quantiles along array dimensions
function quantiles(X, q; dims, drop=false)
  Q = mapslices(x -> quantile(x, q), X, dims=dims)
  out = drop ? dropdims(Q, dims=dims) : Q
  return out
end

# To extract parameters from ADVI model.
function make_extractor(m, q, nsamples=1000)
    qsamples = rand(q, nsamples)
    _, sym2range = Variational.bijector(m; sym_to_ranges = Val(true));
    return sym -> qsamples[collect(sym2range[sym][1]), :]
end
;

# Squared-exponential covariance function
function sqexp_cov_fn(D, alpha, rho; d_is_squared=false)
    if d_is_squared
        return alpha ^ 2 * exp.(-0.5 * D/(rho^2))
    else
        return alpha ^ 2 * exp.(-0.5 * (D/rho) .^ 2)
    end
end

@model function GP(y, X, m_alpha=0.0, s_alpha=1.0, m_rho=0.0, s_rho=1.0, m_sigma=0.0, s_sigma=1.0)
    # Distance matrix.
    D2 = pairwise(Distances.SqEuclidean(), X, dims=1)

    # Priors.
    alpha ~ LogNormal(m_alpha, s_alpha)
    rho ~ LogNormal(m_rho, s_rho)
    sigma ~ LogNormal(m_sigma, s_sigma)

    # Realized covariance function
    K = sqexp_cov_fn(D2, alpha, rho, d_is_squared=true)

    # Sampling Distribution.
    y ~ MvNormal(K + LinearAlgebra.I * sigma^2)  # mean=0, covariance=K.
end;

# Read data.

# Directory where all simulation data are stored.
data_path = joinpath(@__DIR__, "../data/gp-data-N30.json")

# Load data in JSON format.
data = let
    x = open(f -> read(f, String), data_path)
    JSON3.read(x)
end

# Reshape data if needed.
y = Float64.(data[:y])
X = Float64.(reshape(data[:x], length(y), 1))

f = Float64.(data[:f])
x_grid = Float64.(data[:x_grid])

N = size(X, 1);

# Plot data
plt.scatter(vec(X), y, label="Data")
plt.plot(x_grid, f, c="grey", ls=":", label="True f(x)")
plt.xlabel("x")
plt.ylabel("y")
plt.legend();

# Create model.
m = begin
    m_alpha = 0.0
    s_alpha = 0.1
    m_rho = 0.0
    s_rho = 1.0
    m_sigma = 0.0
    s_sigma = 1.0
    GP(y, X, m_alpha, s_alpha, m_rho, s_rho, m_sigma, s_sigma)
end;

# Fit via ADVI. You can also use HMC.
Random.seed!(0)

# initialize variational distribution (optional)
q0 = Variational.meanfield(m)

# NOTE: ADVI(num_elbo_samples, max_iters)

# Compile
@time q = vi(m, ADVI(1, 1), q0, optimizer=Flux.ADAM(1e-1));

# RUN
@time q = vi(m, ADVI(1, 2000), q0, optimizer=Flux.ADAM(1e-1));

# Get posterior samples
extract_gp = make_extractor(m, q)
alpha = vec(extract_gp(:alpha));
rho = vec(extract_gp(:rho));
sigma = vec(extract_gp(:sigma));
advi_samples = Dict(:alpha => alpha, :rho => rho, :sigma => sigma);

# Fit via HMC.
Random.seed!(0)
burn = 1000
nsamples = 1000

# Compile
@time _ = sample(m, HMC(0.01, 1), 1)

# Run
@time chain = sample(m, HMC(0.01, 100), burn + nsamples)

# Get posterior samples
alpha = vec(group(chain, :alpha).value.data[end-nsamples:end, :, 1]);
rho = vec(group(chain, :rho).value.data[end-nsamples:end, :, 1]);
sigma = vec(group(chain, :sigma).value.data[end-nsamples:end, :, 1]);
hmc_samples = Dict(:alpha => alpha, :rho => rho, :sigma => sigma);

# Fit via NUTS.
Random.seed!(7)

# Compile
@time _ = sample(m, NUTS(4, 0.8), 9);

# Run
@time chain = begin
    nsamples = 1000  # number of MCMC samples
    nadapt = 1000  # number of iterations to adapt tuning parameters in NUTS
    iterations = nsamples + nadapt
    target_accept_ratio = 0.8

    sample(m, NUTS(nadapt, target_accept_ratio, max_depth=10), iterations);
end

# Get posterior samples
alpha = vec(group(chain, :alpha).value.data[:, :, 1]);
rho = vec(group(chain, :rho).value.data[:, :, 1]);
sigma = vec(group(chain, :sigma).value.data[:, :, 1]);
nuts_samples = Dict(:alpha => alpha, :rho => rho, :sigma => sigma);

# This funciton returns a function for predicting at new points given parameter values.
function make_gp_predict_fn(Xnew, y, X)
    N = size(X, 1)
    M = size(Xnew, 1)
    Q = N + M
    Z = [Xnew; X]
    D2 = pairwise(SqEuclidean(), Z, dims=1)

    return (alpha, rho, sigma) -> let
        K = sqexp_cov_fn(D2, alpha, rho, d_is_squared=true) + LinearAlgebra.I * sigma^2
        Koo_inv = inv(K[(M+1):end, (M+1):end])
        Knn = K[1:M, 1:M]
        Kno = K[1:M, (M+1):end]
        C = Kno * Koo_inv
        m = C * y
        S = Matrix(LinearAlgebra.Hermitian(Knn - C * Kno'))
        mvn = MvNormal(m, S)
        rand(mvn)
    end
end

# Function for plotting parameter posterior.
function plot_post(samples, name; bins=nothing, suffix="")
    param = samples[name]
    plt.hist(param, density=true, bins=bins)
    plt.xlabel(name)
    plt.ylabel("density")
    if suffix !== ""
        suffix = "($(suffix))"
    end
    plt.title("Histogram of $(name) $(suffix)");
end

function plot_fn_posterior(samples; figsize=(12, 4), figsize_f=figsize, suffix="")
    # Get parameters
    alpha = samples[:alpha]
    rho = samples[:rho]
    sigma = samples[:sigma]
    nsamps = length(alpha)

    # Plot parameters posterior.
    plt.figure(figsize=figsize)
    plt.subplot(1, 3, 1)
    plot_post(samples, :alpha, bins=30, suffix=suffix)
    plt.subplot(1, 3, 2)
    plot_post(samples, :rho, bins=30, suffix=suffix)
    plt.subplot(1, 3, 3)
    plot_post(samples, :sigma, bins=30, suffix=suffix)
    plt.axvline(data[:sigma], ls=":", label="truth", color="red")
    plt.legend()
    plt.tight_layout()

    # Make predictions at new locations.
    X_new = reshape(collect(range(-3.5, 3.5, length=100)), 100, 1)
    gp_predict = make_gp_predict_fn(X_new, y, X)

    ynew = [gp_predict(alpha[m], rho[m], sigma[m]) for m in 1:nsamps]
    ynew = hcat(ynew...)

    # Summary statistics of posterior function.
    pred_mean = mean(ynew, dims=2)
    pred_lower = quantiles(ynew, 0.025, dims=2)
    pred_upper = quantiles(ynew, 0.975, dims=2)

    if suffix !== ""
        suffix = "($(suffix))"
    end

    # Plot predictions.
    plt.figure(figsize=figsize_f)
    plt.plot(vec(X_new), vec(pred_mean), color="blue", label="Posterior predictive mean")
    plt.plot(x_grid, f, color="red", ls=":", label="True f(x)")

    plt.scatter(vec(X), vec(y), color="black", label="Data")
    plt.fill_between(vec(X_new), vec(pred_upper), vec(pred_lower), color="blue",
                     alpha=0.2, label="95% Credible Interval")
    plt.legend(loc="upper left", fontsize=8)
    plt.title("GP Posterior predictive with 95% credible interval $(suffix)");

    plt.xlabel("x")
    plt.ylabel("y")

    plt.ylim(-2, 2)
    plt.tight_layout()
end

plot_fn_posterior(advi_samples, figsize=(9, 3), suffix="ADVI")

plot_fn_posterior(hmc_samples, figsize=(9, 3), suffix="HMC");

plot_fn_posterior(nuts_samples, figsize=(9, 3), suffix="NUTS");


# example from
# Barraud, A.Y., "A numerical algorithm to solve A'XA - X = Q",
# IEEE Trans. Auto. Contr., AC-22, pp. 883-885, 1977.

A = [1.5 2.5 3.5 4.5 5.5 6.5 7.5;
		1 -1.5 2.5 3.5 4.5 5.5 6.5;
		0 0 1.5 -2.5 3.5 4.5 5.5;
		0 0 0 1.5 2.5 -3.5 4.5;
		0 0 0 1 1.5 2.5 -3.5;
		0 0 0 0 0 1.5 -2.5;
		0 0 0 0 0 1 1.5]

q = [5.25 1.5 19.25 22.5 55.75 60 49.25;
		0 0 3.25 -0.25 -6 -11 11.25;
		0 0 62.75 83 183.25 222.75 129.75;
		0 0 0 56.5 208.25 232.25 250.75;
		0 0 0 0 532.25 739.5 476.75;
		0 0 0 0 0 631 745.75;
		0 0 0 0 0 0 133.75]

Q = (q + q') - diagm(0 => diag(q))

x = [1 1 1 1 1 1 1;
		0 1 2 3 4 5 6;
		0 0 1 3 6 10 15;
		0 0 0 1 4 10 20;
		0 0 0 0 1 5 15;
		0 0 0 0 0 1 6;
		0 0 0 0 0 0 1];

X = (x + x') - diagm(0 => diag(x)) # theoretical solution

A = Array(A') # A' gets Adjoint{Float64,Array{Float64,2}}, but sylvester needs Array
Q = -Q

X1 = dlyap(A,Q) # computed solution
##################################################

print("Testing dlyap: ")
if sum((X-X1).^2) < 1e-6
	println("OK")
end


module kuramoto_model

using DifferentialEquations, Flux, Optim, DiffEqFlux, Plots, Random

rng=MersenneTwister(123) #make the results consistent across runs
N=10000
ω=randn(rng,N)
norder=2 #OA=1, versus Hannay Ansatz=2

#ω0,γ,K,β
model_params=[2π, 0.1,0.8,0.1]
ic=rand(rng,N)*2π;

tend=20.0
tstart=0.0
sampling=0.05

function full_model(dϕ,ϕ,p,t)
    ω0, γ, K, β=p
    ψ,R=OrderParameter(ϕ)
    for i in 1:N
        dϕ[i]=ω0+γ*ω[i]+R*K*sin(ψ-ϕ[i]+β)
    end
end

function reduced_model(du,u,p,t)
    ω0, γ, K, β=p
    ψ,R=u
    n=norder
    du[1]=ω0+K/2.0*sin(β)*(1+R^(2*norder))
    du[2]=-γ*R+K/2.0*cos(β)*R*(1-R^(2*norder))
end

function OrderParameter(ϕ;m=1.0)
    Z=1.0/N*sum(exp.(im.*ϕ*m))
    R=abs(Z)
    ψ=angle(Z)
    return ψ,R
end

function OrderParameterComplex(ϕ;m=1.0)
    Z=1.0/N*sum(exp.(im.*ϕ*m))
    return Z
end

icr=[angle(OrderParameterComplex(ic)), abs(OrderParameterComplex(ic))]

function InitPlot()

    lay = @layout [ a{0.6w} [b; c]]

    t=tstart:0.01:tend
    sol=solve(ODEProblem(full_model,ic,(tstart,tend),model_params), Tsit5(), saveat=t)

    num_tps=length(t)
    R=zeros(num_tps)
    ψ=zeros(num_tps)
    for tp in 1:num_tps
        psi, rr=OrderParameter(sol[:,tp])
        R[tp]=rr
        ψ[tp]=psi
    end

    pl1=plot(ψ,R,lw=2, legend=false, color=:blue, proj = :polar)

    #radius plot
    pl2=plot(t,R,lw=2, legend=false, color=:blue)

    #angle plot
    pl3=plot(t,sin.(ψ), lw=2, legend=false, color=:blue)

    pl=plot(pl1, pl2, pl3, layout=lay)
    display(pl)

end

function predict_adjoint(param) # Our 1-layer neural network
    prob=ODEProblem(reduced_model,icr,(tstart,tend), model_params)
    Array(concrete_solve(prob,Tsit5(),icr,param,saveat=tstart:sampling:tend,abstol=1e-8,reltol=1e-6))
end

function generate_data(ic; param=model_params)
    prob=ODEProblem(full_model,ic,(tstart,tend), param)
    tps=tstart:sampling:tend
    num_tps=length(tps)
    d1=Array(concrete_solve(prob,Tsit5(),ic,param,saveat=tps))
    # Get a time series of the order parameter
    d2=mapslices(OrderParameterComplex, d1, dims=1)
    data=Float64[angle.(d2); abs.(d2)]
    data=reshape(data,2,num_tps)
    return data
end

# Generate some data to fit, and add some noise to it
data=generate_data(ic)
#σN=0.05
#data+=σN*randn(size(data))

function loss_adjoint(param)
    prediction = predict_adjoint(param)
    loss=0.0
    for k in 1:length(prediction[1,:])
        Rp=prediction[2,k]
        Rd=data[2,k]
        θp=prediction[1,k]
        θd=data[1,k]
        loss+=sqrt((Rp*cos(θp)-Rd*cos(θd))^2+(Rp*sin(θp)-Rd*sin(θd))^2)
    end
    loss,prediction
end

function train_model(;pguess=[6.1, 0.15,1.1,0.15])
    println("The initial loss function is $(loss_adjoint(pguess)[1])")
    res = DiffEqFlux.sciml_train(loss_adjoint,pguess,BFGS(initial_stepnorm=0.0001), maxiters=4000)
    println("The parameters are $(res.minimizer) with final loss value $(res.minimum)")
    return(res)
end

function plotFit(param; tend=20.0)

    validationPlot(param,ic; tend=tend)

end

function validationPlot(param,ic; tend=10.0)

    lay = @layout [ a{0.6w} [b; c]]

    t=tstart:0.01:tend
    icr=[angle(OrderParameterComplex(ic)), abs(OrderParameterComplex(ic))]

    sol=solve(ODEProblem(full_model,ic,(tstart,tend),model_params), Tsit5(), saveat=t)
    sol_fit=solve(ODEProblem(reduced_model,icr,(tstart,tend),param), Tsit5(), saveat=t)

    num_tps=length(t)
    R=zeros(num_tps)
    ψ=zeros(num_tps)
    for tp in 1:num_tps
        psi, rr=OrderParameter(sol[:,tp])
        R[tp]=rr
        ψ[tp]=psi
    end

    pl1=plot(ψ,R,lw=2, legend=false, color=:blue, proj = :polar)
    plot!(pl1, sol_fit, color=:green, proj=:polar, vars=(1,2))

    #radius plot
    pl2=plot(t,R,lw=2, legend=false, color=:blue)
    plot!(pl2, sol_fit, color=:green, lw=2.0, vars=(0,2))

    #angle plot
    pl3=plot(t,sin.(ψ), lw=2, legend=false, color=:blue)
    plot!(pl3, sol_fit.t, sin.(sol_fit[1,:]), color=:green)

    pl=plot(pl1, pl2, pl3, layout=lay)
    display(pl)

end

end #module

#d=kuramoto_model.generate_data(randn(kuramoto_model.N))
#kuramoto_model.predict_adjoint(kuramoto_model.model_params)

kuramoto_model.InitPlot()
results_kuramoto=kuramoto_model.train_model()
kuramoto_model.plotFit(results_kuramoto.minimizer)
kuramoto_model.validationPlot(results_kuramoto.minimizer, π*rand(kuramoto_model.N))
kuramoto_model.validationPlot(results_kuramoto.minimizer, 1.5*π*rand(kuramoto_model.N))


using Plots, LinearAlgebra, Random, Statistics
pyplot()

# compute gradient
function df(X::Matrix{Float64}, y::Vector{Float64}, theta::Vector{Float64})
    return X' * (X * theta - y)
end

# gradient descent
function lr_gd(X::Matrix{Float64}, y::Vector{Float64}, alpha::Float64, num_iterations::Int64)
    (num_samples, num_features) = size(X)
    theta = zeros(num_features)
    for i = 1:num_iterations
        theta += df(X, y, theta)
    end
    return theta
end

# stochastic gradient descent, batch size of 1
function lr_sgd(X::Matrix{Float64}, y::Vector{Float64}, alpha::Float64, batch_size::Int64, num_iterations::Int64)
    (num_samples, num_features) = size(X)
    theta = zeros(num_features)
    for i = 1:num_iterations
        # use num_samples minibatches of size 1 to compute gradient to update theta
        # using k:k in order to keep type as a matrix
        theta += mean(df(X[k:k,:], y[k:k], theta) for k in (rand(1:num_samples) for b = 1:batch_size));
    end
    return theta
end
function main()
    num_features = 1024;
    dataset_sizes = 2 .^ [7,8,9,10,11,12,13,14,15,16];
    num_iterations = 1024;
    alpha = 0.001;
    # we will do this for the vanilla sgd case, without varying minibatch size
    batch_size = 1;
    gd_times = Float64[];
    sgd_times = Float64[];
    for num_samples in dataset_sizes
        X = randn(num_samples, num_features);
        y = randn(num_samples);
        gd_time = (@elapsed lr_gd(X, y, alpha, num_iterations)) / num_iterations;
        sgd_time = (@elapsed lr_sgd(X, y, alpha, batch_size, num_iterations)) / num_iterations;
        push!(gd_times, gd_time);
        push!(sgd_times, sgd_time);
        println("dataset size $num_samples --- GD time = $gd_time, SGD time = $sgd_time");
    end
    plot(dataset_sizes, gd_times, label="Gradient Descent (GD)", color=:blue,
        markershape=:circle, xaxis=:log, yaxis=:log);
    plot!(dataset_sizes, sgd_times, label="Stochastic Gradient Descent (SGD, with minibatch size=$batch_size)",
        markershape=:rect, color=:orange, marker=:, xaxis=:log, yaxis=:log,
        xticks = 10 .^ [2, 3, 4, 5, 6],
        xlabel="Number of data points", ylabel="Time per iteration (seconds)",
        title="Comparing Runtime per Iteration for GD and SGD");
    savefig("timesgd.pdf")
end
main()