#!/usr/bin/env julia
#
# usage:
#  script......... plaint.jl /path/to/system.yml
#  interactively.. julia -qi plaint.jl

using YAML              #read data file
using DelimitedFiles    #output results
using LinearAlgebra

"""
    ksnr(e, M, eps, max_iterations)

Solve Kepler's equation with NR method.
"""
function ksnr(e::Float64, M::Float64,
              eps::Float64 = 1e-16, nm::Int64 = 16)
    x = M + e * sin(M)
    d = (x - e * sin(x) - M) / (1 - e * cos(x))
    while abs(d) > eps && nm > 0
        x -= d
        d = (x - e * sin(x) - M) / (1 - e * cos(x))
        nm -= 1
    end
    return x
end

"""
    euler_rotd(Ω, i, ω)

Return Eulerian rotation.
"""
function euler_rotd(Ω::Float64, i::Float64, ω::Float64)
    return [ cosd(Ω) -sind(Ω) 0; sind(Ω) cosd(Ω) 0; 0 0 1 ] *
           [ 1 0 0; 0 cosd(i) -sind(i); 0 sind(i) cosd(i) ] *
           [ cosd(ω) -sind(ω) 0; sind(ω) cosd(ω) 0; 0 0 1 ]
end

"""
    oe_to_xv!(mu, j, orbital_elements, positions, velocities)

Convert orbital elements to cartesian vectors.
"""
function oe_to_xv!(mu::Float64, j::Int64,
                   oe::Array{Float64,2},
                   vr::Array{Float64,2}, vv::Array{Float64,2})
    a, e, i, Ω, ω, M = oe[:, j]
    E = ksnr(e, M * π / 180)
    r = a * (1 - e * cos(E))
    pof = [cos(E) - e, sqrt(1 - e^2) * sin(E), 0] * a
    vof = [  - sin(E), sqrt(1 - e^2) * cos(E), 0] * (sqrt(mu * a) / r)
    rot = euler_rotd(Ω, i, ω)
    vr[:, j] = rot * pof
    vv[:, j] = rot * vof
end

"""
    xv_to_oe!(mu, j, positions, velocities, orbital_elements)

Convert cartesian vectors to orbital elements.
"""
function xv_to_oe!(mu::Float64, j::Int64,
                   rt::Array{Float64,2}, vt::Array{Float64,2},
                   oe::Array{Float64,2})
    tiny = 1e-10
    vr, vv = rt[:, j], vt[:, j]
    r = norm(vr)
    v = norm(vv)
    vh = cross(vr, vv)
    h = norm(vh)
    vn = [-vh[2], vh[1], 0]
    n = norm(vn)
    ve = cross(vv, vh) / mu - vr / r
    a = 1 / (2 / r - v^2 / mu)
    e = norm(ve)
    i = acos(vh[3] / h)
    if e > tiny
        if n > tiny
            Ω = acos(vn[1] / n)
            if vn[2] < 0
                Ω = 2π - Ω
            end
            ω = acos(dot(ve, vn) / (e * n))
            if ve[3] < 0
                ω = 2π - ω
            end
        else
            i = Ω = .0
            ω = atan2(ve[2], ve[3])
            if cross(ve, vr)[3] < 0
                ω = 2π - ω
            end
        end
        nu = acos(dot(ve, vr)/(e * r))
        if dot(vr, vv) < 0
            nu = 2π - nu
        end
    else
        #not implemented
    end
    E = 2 * atan(sqrt((1 - e) / (1 + e)) * tan(nu / 2))
    M = E - e * sin(E)
    if M < 0
        M += 2π
    end
    i, Ω, ω, M = [i, Ω, ω, M] * 180 / π
    oe[:, j] = [a, e, i, Ω % 360, ω % 360, M % 360]
end

"""
    grav_a!(no_of_bodies, masses, positions, acceleration)

Compute gravitational acceleration.
!!! note
    The `a` should be multiplied with the appropriate gravitational constant.
"""
function grav_a!(n::Int64, m::Array{Float64,1},
                 x::Array{Float64,2}, a::Array{Float64,2})
    fill!(a, 0)
    @inbounds for i = 1:(n - 1), j = (i + 1):n
        R = x[:, i] - x[:, j]
        Rd3 = R / norm(R)^3
        a[:, i] -= m[j] * Rd3
        a[:, j] += m[i] * Rd3
    end
end

"""
    mksnr(dt, α, n, σ0, r0, eps, max_iterations)

Solve modified Kepler's equation with NR method.
!!! note
    Argument `α` is the inverse of semimajor axis.
"""
function mksnr(dt::Float64, α::Float64, n::Float64, s::Float64, r0::Float64,
               eps::Float64 = 1e-16, nm::Int64 = 16)
    r0a, s0a, ndt = 1 - r0 * α, s * sqrt(α), n * dt
    x = π
    d = (x - r0a * sin(x) - s0a * (cos(x) - 1) - ndt) /
            (1 - r0a * cos(x) + s0a * sin(x))
    while abs(d) > eps && nm > 0
        x -= d
        d = (x - r0a * sin(x) - s0a * (cos(x) - 1) - ndt) /
                (1 - r0a * cos(x) + s0a * sin(x))
        nm -= 1
    end
    return x
end

"""
    gauss_step!(dt, mu, no_of_bodies, positions, velocities)

Keplerian drift using Gauss' f, g functions.
"""
function gauss_step!(dt::Float64, mu::Float64, np::Int64,
                     rt::Array{Float64,2}, vt::Array{Float64,2})
    @inbounds @simd for i = 1:np
        vr0, vv0 = rt[:, i], vt[:, i]
        r0, v0 = norm(vr0), norm(vv0)
        α = 2 / r0 - v0^2 / mu
        a, n = 1 / α, sqrt(mu * α^3)
        s = dot(vr0, vv0) / sqrt(mu)
        E = mksnr(dt, α, n, s, r0)
        cose, sine = cos(E), sin(E)
        r1 = a + (r0 - a) * cose + sqrt(a) * s * sine
        a1ce = a * (cose - 1)
        ivr0, ivr1 = 1 / r0, 1 / r1
        ft = 1 + a1ce * ivr0
        gt = dt + (sine - E) / n
        fd = - a^2 * n * sine * ivr1 * ivr0
        #gd = 1 + a1ce * ivr1
        gd = (1 + gt * fd) / ft
        rt[:, i] = ft * vr0 + gt * vv0
        vt[:, i] = fd * vr0 + gd * vv0
    end
end

"""
    eam_dh(no_of_bodies, masses, positions, velocities)

Compute energy and angular momentum, with democratic heliocentric vectors.
"""
function eam_dh(np::Int64, m::Array{Float64,1},
                x::Array{Float64,2}, v::Array{Float64,2})
    et, am = norm(sum(m' .* v, dims = 2))^2 / 2, zeros(3)
    @inbounds @simd for i = 1:np
        et += m[i] * norm(v[:, i])^2 / 2
        et -= m[i] / norm(x[:, i])
        am += m[i] * cross(x[:, i], v[:, i])
    end
    @inbounds for i = 1:(np - 1), j = (i + 1):np
        et -= m[i] * m[j] / norm(x[:, i] - x[:, j])
    end
    return et, norm(am)
end

"""
    plaint_dhi(iostream, total_iterations, dt, step_iterations,
               mu, masses, orbital_elements)

Perform planetary integration using democratic heliocentric integrator.
"""
function plaint_dhi(of::Array{IOStream, 1},
                    kmax::Float64, dt::Float64, kpri::Float64, mu::Float64,
                    mp::Array{Float64,1}, oe::Array{Float64,2})
    mt = 1 / (sum(mp) + 1)
    np = size(mp, 1)
    an = Array{Float64}(undef, 3, np)
    x = similar(an)
    v = similar(an)
    @inbounds @simd for j = 1:np
        oe_to_xv!(mu, j, oe, x, v)
    end
    v .-= mt * sum(mp' .* v, dims = 2)                  #to_dh
    el0 = eam_dh(np, mp, x, v)
    println("Initial energy:           ", el0[1], "\n",
            "Initial angular momentum: ", el0[2], "\n")
    d2 = dt / 2
    @inbounds for k = 1:kmax
        try
            x .+= d2 * sum(mp' .* v, dims = 2)
            gauss_step!(d2, mu, np, x, v)
            grav_a!(np, mp, x, an)
            v .+= dt * an
            gauss_step!(d2, mu, np, x, v)
            x .+= d2 * sum(mp' .* v, dims = 2)
            if k % kpri == 0
                eln = eam_dh(np, mp, x, v)
                va = v .+ mt * sum(mp' .* v, dims = 2)  #to_ac
                writedlm(of[end], [collect(eln .- el0); collect(@. 1 - eln / el0)]')
                @simd for j = 1:np
                    xv_to_oe!(mu, j, x, va, oe)
                    writedlm(of[j], oe[:, j]')
                end
            end
        catch
            println("Error: Exiting on ", k, "\n")
            break
        end
    end
    close.(of)
end

"""
    data_get(file)

Get planetary system's data and integration options.
"""
function data_get(file::String)
    data = YAML.load_file(file)
    tmax::Float64 = data["options"]["tmax"]
    step::Float64 = data["options"]["step"]
    ever::Float64 = data["options"]["ever"]
    unit = "solar"
    try
        unit = data["system"]["units"]
    catch
    end
    ms = data["system"]["star_mass"]
    datp = data["system"]["planets"]
    sort!(datp, by = x -> x["orbital_elements"][1])
    np = size(datp, 1)
    name = Array{String}(undef, np)
    mp = Array{Float64}(undef, np)
    oe = Array{Float64}(undef, 6, np)
    @inbounds @simd for i = 1:np
        name[i] = data["system"]["star_name"] * " " * datp[i]["id"]
        mp[i] = datp[i]["mass"]
        oe[:, i] = datp[i]["orbital_elements"]
    end
    return tmax, step, ever, ms, name, mp, oe, unit
end

"""
    data_prep(tmax, step, ever,
              star_mass, file_names, masses, orbital_elements, units)

Prepare data and options to feed to the integrator.
"""
function data_prep(tmax::Float64, step::Float64, ever::Float64,
                   ms::Float64, name::Array{String, 1},
                   mp::Array{Float64, 1}, oe::Array{Float64,2},
                   unit::String = "solar")
    if unit == "solar"
        mp = mp * 9.547919e-4 / ms
    else #natural
        mp = mp / ms
    end
    oe[1, :] /= oe[1, 1]
    tp = 2π
    dt = step * tp
    kmax = div(tmax * tp, dt)
    kpri = div(ever * tp, dt)
    of = map(f -> open("out/$f.txt", "w"), name)
    push!(of, open("out/eamerr.txt", "w"))
    println("Total simulation time is: ", tmax, " yr\n",
            "Time of every step to be: ", step, " yr\n",
            "Total number of steps is: ", kmax, " steps\n",
            "Output will be every:     ", ever, " yr\n",
            "or differently every:     ", kpri, " steps\n")
    return of, kmax, dt, kpri, 1.0, mp, oe
end

"""
    usage()

Return an exemplary session.
"""
function usage()
    println("## Example session:\n",
            "#set file containing data\n",
            "file = \"data/TRAPPIST-1.yml\";\n",
            "#get data from file\n",
            "tmax, step, ever, ms, name, mp, oe = data_get(file);\n",
            "#or enter manually\n",
            "dgr = (tmax, step, ever, ms, name, mp, oe);\n",
            "#prepare data\n",
            "of, kmax, dt, kpri, mu, mp, oe = data_prep(dgr...);\n",
            "#integrate, available: plaint_dhi, plaint_s6i\n",
            "plaint_dhi(of, kmax, dt, kpri, mu, mp, oe);\n",
            "#benchmark integration of a particular system\n",
            "@time plaint_dhi(data_prep(data_get(file)...)...);")
end

## Session
# Run as script or, interactively in a Julia session.

if ! isinteractive()
    @time plaint_dhi(data_prep(data_get(ARGS[1])...)...)
else
    usage()
end


perf_table = DataFrame()

perf_table.TestDescription = [
    "Rollout, 1 ft (s)",
    "60 mph (s)",
    "100 mph (s)",
    "150 mph (s)",
    "Rolling start, 5-60 mph (s)",
    "Top gear, 30-50 mph (s)",
    "Top gear, 50-70 mph (s)",
    "¼-mile time (s)",
    "¼-mile speed (mph)"
]

perf_table.CarAndDriver = [
    0.2,
    2.4,
    6.0,
    15.2,
    2.9,
    1.1,
    1.6,
    10.5,
    130
]

u0 = [0.0, 0, 0.0];
tspan=(0, 20.0);

condition(u,t,integrator) = u[1] - 0.3048; #150 mph in m/s
affect!(integrator) = terminate!(integrator);
cb = ContinuousCallback(condition,affect!);
rollout = solve(ODEProblem(taycan_state, u0, tspan),Rodas4(autodiff=false), saveat=1e-3, reltol=1e-9, callback=cb);
ro_t = rollout.t[end]

condition(u,t,integrator) = u[2] - 67.056; #150 mph in m/s
affect!(integrator) = terminate!(integrator);
cb = ContinuousCallback(condition,affect!);
zeroTo150 = solve(ODEProblem(taycan_state, u0, tspan),Rodas4(autodiff=false), saveat=1e-3, reltol=1e-9, callback=cb);
zeroTo150.t[end]

z60_t = zeroTo60.t[end] - ro_t
t_150 = zeroTo150.t
x_150 = [zeroTo150.u[x][1] for x in 1:length(zeroTo150.t)]
v_150 = [zeroTo150.u[x][2] for x in 1:length(zeroTo150.t)]

z100_t=t_150[findfirst((x)->x>=44.704, v_150)] - ro_t
z150_t=zeroTo150.t[end]-ro_t

z5_t=t_150[findfirst((x)->x>=2.2352, v_150)]
five60_t = zeroTo60.t[end]-z5_t

z30_t=t_150[findfirst((x)->x>=13.4112, v_150)]
z50_t=t_150[findfirst((x)->x>=22.352, v_150)]
z70_t=t_150[findfirst((x)->x>=31.2928, v_150)]

thirty50_t = z50_t-z30_t
fifty70_t = z70_t - z50_t

q_mile_t = t_150[findfirst((x)->x>=402.336, x_150)] - ro_t
q_mile_v = v_150[findfirst((x)->x>=402.336, x_150)]*2.23694

perf_table.SimulationResults = [
    ro_t,
    z60_t,
    z100_t,
    z150_t,
    five60_t,
    thirty50_t,
    fifty70_t,
    q_mile_t,
    q_mile_v
]

perf_table.AbsoluteError = perf_table.SimulationResults- perf_table.CarAndDriver
perf_table.RelativeError = (perf_table.AbsoluteError./perf_table.CarAndDriver)*100


# create recursion coefficient of logistic density numerically
# and compare against analytic solution

using PolyChaos, Test
import LinearAlgebra: norm

function myquad(N::Int64,pos::Bool)
    a,b = rm_laguerre(N+1)
    n,w = golubwelsch(a,b)
    w = w ./ ((1 .+ exp.(-n)).^2)
    pos ? (n,w) : (-n,w)
end

ε, tol = 1e-10, 1e-6
N = 1:3:15
quads = [ n->myquad(n,true); n->myquad(n,false) ]

@testset "Logistic density" begin
    for n in N, disc in [stieltjes, lanczos]
        a,b = mcdiscretization(n,quads;discretization=disc,Nmax=300,gaussquad=true,ε=ε)
        aa,bb = rm_logistic(n)
        @test isapprox(norm(b-bb,Inf),0.;atol=tol)
    end
end

function hrhermite(t)
    exp(-t^2)
end

AB=[[0. 3.];[3. 6.];[6. 9.];[9. Inf]];
quads = [ n->quadgp(hrhermite,AB[i,1],AB[i,2],n;quadrature=fejer) for i=1:size(AB,1)]
α, β = mcdiscretization(40,quads,Nmax=400,gaussquad=false,discretization=lanczos,ε=1e-7)


# Completed: 69/69

import Base: acos, acosd, acosh, acot, acotd, acoth, acsc, acscd, acsch, asec, asecd, asech, asin, asind, asinh, atan, atand, atanh, cbrt, clamp, cos, cosc, cosd, cosh, cospi, cot, cotd, coth, csc, cscd, csch, deg2rad, exp, exp10, exp2, expm1, exponent, hypot, ldexp, log, log10, log1p, log2, max, min, mod2pi, rad2deg, rem2pi, sec, secd, sech, significand, sin, sinc, sind, sinh, sinpi, sqrt, tan, tand, tanh

# julia> for n in names(Base.Math); println(n); end
# @evalpoly(z,c...): Evaluate the polynomial \sum_k c[k] z^{k-1}.
# Math: package name
# ^: Handled in base.jl
@primitive acos(x),dy (dy.*(-1 ./ sqrt.(1 .- abs2.(x))))
@primitive acosd(x),dy (dy.*(-(180 / pi) ./ sqrt.(1 .- abs2.(x))))
@primitive acosh(x),dy (dy.*(1 ./ sqrt.(abs2.(x) .- 1)))
@primitive acot(x),dy (dy.*(-1 ./ (1 .+ abs2.(x))))
@primitive acotd(x),dy (dy.*(-(180 / pi) ./ (1 .+ abs2.(x))))
@primitive acoth(x),dy (dy.*(1 ./ (1 .- abs2.(x))))
@primitive acsc(x),dy (dy.*(-1 ./ sqrt.((abs2.(x) .* (x .- 1)) .* (x .+ 1))))
@primitive acscd(x),dy (dy.*(-(180 / pi) ./ sqrt.((abs2.(x) .* (x .- 1)) .* (x .+ 1))))
@primitive acsch(x),dy (dy.*(-1 ./ sqrt.(abs2.(abs2.(x)) + abs2.(x))))
@primitive asec(x),dy (dy.*(1 ./ sqrt.(abs2.(abs2.(x)) - abs2.(x))))
@primitive asecd(x),dy (dy.*((180 / pi) ./ sqrt.(abs2.(abs2.(x)) - abs2.(x))))
@primitive asech(x),dy (dy.*(-1 ./ sqrt.(abs2.(x) - abs2.(abs2.(x)))))
@primitive asin(x),dy (dy.*(1 ./ sqrt.(1 .- abs2.(x))))
@primitive asind(x),dy (dy.*((180 / pi) ./ sqrt.(1 .- abs2.(x))))
@primitive asinh(x),dy (dy.*(1 ./ sqrt.(1 .+ abs2.(x))))
@primitive atan(x),dy (dy.*(1 ./ (1 .+ abs2.(x)))); @primitive atan(x1,x2),dy,y  unbroadcast(x1,dy.*(x2./(abs2.(x1)+abs2.(x2))))  unbroadcast(x2,dy.*(-x1./(abs2.(x1)+abs2.(x2))))
@primitive atand(x),dy (dy.*((180 / pi) ./ (1 .+ abs2.(x))))
@primitive atanh(x),dy (dy.*(1 ./ (1 .- abs2.(x))))
@primitive cbrt(x),dy,y (dy.*(1 ./ (3 .* abs2.(y))))
@primitive clamp(x,lo,hi),dy,y  unbroadcast(x,dy.*(lo .<= x .<= hi))
# clamp! overwrites.
@primitive cos(x),dy (dy.*(-(sin.(x))))
@primitive cosc(x),dy,y (dy.*((-2 .* y) ./ x - sinc.(x) .* (pi ^ 2)))
@primitive cosd(x),dy (dy.*((-(sind.(x)) .* pi) / 180))
@primitive cosh(x),dy (dy.*(sinh.(x)))
@primitive cospi(x),dy (dy.*(-(sinpi.(x)) .* pi))
@primitive cot(x),dy (dy.*(-(abs2.(csc.(x)))))
@primitive cotd(x),dy (dy.*((-(abs2.(cscd.(x))) .* pi) / 180))
@primitive coth(x),dy (dy.*(-(abs2.(csch.(x)))))
@primitive csc(x),dy,y (dy.*(-1 .* y .* cot.(x)))
@primitive cscd(x),dy,y (dy.*(((-1 .* y .* cotd.(x)) .* pi) / 180))
@primitive csch(x),dy,y (dy.*(-1 .* y .* coth.(x)))
@primitive deg2rad(x),dy (dy.*(pi / 180))
@primitive exp(x),dy,y (dy.*(y))
@primitive exp10(x),dy,y (dy.*(y .* log(10)))
@primitive exp2(x),dy,y (dy.*(y .* log(2)))
@primitive expm1(x),dy,y (dy.*(1 .+ y))
@zerograd exponent(x)
# frexp # returns 2 values
@primitive hypot(x1,x2),dy,y  unbroadcast(x1,dy.*(x1./y))  unbroadcast(x2,dy.*(x2./y))
@primitive ldexp(x,n),dy  (dy.*exp2.(n))
@primitive log(x),dy (dy.*(1 ./ x)); @primitive log(x1,x2),dy,y  unbroadcast(x1,dy.*(-log.(x2)./(x1.*abs2.(log.(x1)))))  unbroadcast(x2,dy.*(1 ./ (x2.*log.(x1))))
@primitive log10(x),dy (dy.*(1 ./ (log(10) .* x)))
@primitive log1p(x),dy (dy.*(1 ./ (1 .+ x)))
@primitive log2(x),dy (dy.*(1 ./ (log(2) .* x)))
@primitive max(x1,x2),dy,y  unbroadcast(x1,dy.*(y.==x1))  unbroadcast(x2,dy.*(y.==x2))
@primitive min(x1,x2),dy,y  unbroadcast(x1,dy.*(y.==x1))  unbroadcast(x2,dy.*(y.==x2))
# minmax # returns 2 values
@primitive mod2pi(x),dy dy
# modf # returns 2 values
@primitive rad2deg(x),dy (dy.*(180 / pi))
@primitive rem2pi(x,r),dy  dy
@primitive sec(x),dy,y (dy.*(y .* tan.(x)))
@primitive secd(x),dy,y (dy.*(((y .* tand.(x)) .* pi) / 180))
@primitive sech(x),dy,y (dy.*(-1 .* y .* tanh.(x)))
@primitive significand(x),dy (dy.*exp2.(-exponent.(x)))
@primitive sin(x),dy (dy.*(cos.(x)))
@primitive sinc(x),dy (dy.*(cosc.(x)))
# sincos # returns 2 values
@primitive sind(x),dy (dy.*((cosd.(x) .* pi) / 180))
@primitive sinh(x),dy (dy.*(cosh.(x)))
@primitive sinpi(x),dy (dy.*(cospi.(x) .* pi))
@primitive sqrt(x),dy,y (dy.*(1 ./ (2 .* y)))
@primitive tan(x),dy,y (dy.*(1 .+ abs2.(y)))
@primitive tand(x),dy,y (dy.*(((1 .+ abs2.(y)) .* pi) / 180))
@primitive tanh(x),dy,y (dy.*(1 .- abs2.(y)))


#= -------------------------------------------------------------------------
# @file problema2_PID_loop.jl
#
# @date 12/09/17 19:43:46
# @author Martin Noblia
# @email mnoblia@disroot.org
#
# @brief
#
# @detail
#
#  Licence:
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
---------------------------------------------------------------------------=#
# imports
using PyPlot
#------------------------------------------------------------------------------
#                            constants
#------------------------------------------------------------------------------
const  A      =  2.0
const  R      =  2.390057361376673e-5
const  cₑ     =  4.184e6
const  V      =  0.5
const  θ_ref  =  333.0
const  θₐ     =  293.0
const  θ₀     =  288.0
const  τ      = R * cₑ * V
const  G_inv  = 0.4 * 836800.0
const  K_p    = 41840.0

q(t) = K_p * (θ_ref - heat_panel_winter(t, A, θ₀, θₐ, τ, G_inv))
function heat_panel_winter(t, A, θ₀, θₐ, τ, G)
   (q(t) * R * exp(-t/τ)) / (τ) + ((θₐ + R*A*G) * (1 - exp(-t / τ))) + θ₀ * exp(-t / τ)
end

t = 0:0.01:100
θ = heat_panel_winter.(t, A, θ₀, θₐ, τ, G_inv)


module SI

    TSSGP = exp(841.2869278)                                # 26 sporadic group Product
    HFP = exp(674.5210288)                                  # Happy Family product
    O_M = 80801742479451287588645990496171075700575436800000000 #2^46 · 3^20 · 5^9 · 7^6 · 11^2 · 13^3 · 17 · 19 · 23 · 29 · 31 · 41 · 47 · 59 · 71
    O_B = 4154781481226426191177580544000000                #2^41 · 3^13 · 5^6 · 7^2 · 11 · 13 · 17 · 19 · 23 · 31 · 47
    ΔνCs = 9192631770e0;                                    #  Δν(133Cs)hfs=9192631770e0 s−1
    λe = 2.4263102367e-12;                                  # Electron Compton wavelength in m
    ƛe = 3.861592046068738e-13;                             # ƛe ≡ ħ/m_ec Reduced Electron Compton wavelength
    ElementaryCharge = 1.602176634e-19;                     #  coulomb (C). or A⋅s
    c = 299792458e0;                                        # m⋅s−1
    SpeedOfLight = 2.99792458e8;                            # m / s 299 792 458 m/s
    GravitationalConstant = 6.67408e-11;                    # m^3 / kg^-1 s^-2 or m^3 kg^-1 s^-2
    SanchezGravitationalConstant = 6.675453818e-11;         # m^3 / kg^-1 s^-2 or m^3 kg^-1 s^-2
    h = 6.62607015e-34;                                     # kg⋅m2⋅s−1
    PlanckConstantH = 6.62607015e-34;                       # J . s
    PlanckConstantHbar = 1.05457162825e-34;                 # J . s
    AstronomicalUnit = 1.49597870691e11;                    # m
    LightYear = 9.46053620707e15;                           # m
    Parsec = 3.08567758135e16;                              # m
    Jansky = 10e-26;                                        # W.m^−2.Hz^−1  \text{ Jansky}=10^{-26}\text{ Watts meters}^{-2}\text{ Hertz}^{-1}
    GravAccel = 9.80665e0;                                  # m / s^2
    ElectronVolt = 1.602176487e-19;                         # kg m^2 / s^2
    MassElectron = 9.10938356e-31;                          # kg
    MassMuon = 1.83615267e-28;                              # kg
    MassProton = 1.672621898e-27;                           # kg
    MassNeutron = 1.674927471e-27;                          # kg
    MassTau = 3.16773502e-27;                               # kg
    MassHydrogen = 1.6737236e-27;                           # kg
    Rydberg = 2.17987196968e-18;                            # kg m^2 / s^2
    k = 1.380649e-23;                                       # kg⋅m2⋅K−1⋅s−2
    Boltzmann = 1.380649e-23;                               # J / K                  [kg cm^2 / K s^2]
    AvogadroNumberNA = 6.02214076e23;                       # mol−1  reciprocal mole (mol−1)
    LuminousEfficacyKcd = 683e0;                            # cd⋅sr⋅s3⋅kg−1⋅m−2
    Minute = 6e1;                                           # s
    Hour = 3.6e3;                                           # s
    Day = 8.64e4;                                           # s
    Week = 6.048e5;                                         # s
    Inch = 2.54e-2;                                         # m
    Foot = 3.048e-1;                                        # m
    Yard = 9.144e-1;                                        # m
    Mile = 1.609344e3;                                      # m
    NauticalMile = 1.852e3;                                 # m
    Fathom = 1.8288e0;                                      # m
    Mil = 2.54e-5;                                          # m
    Point = 3.52777777778e-4;                               # m
    Texpoint = 3.51459803515e-4;                            # m
    Micron = 1e-6;                                          # m
    Angstrom = 1e-10;                                       # m
    Hectare = 1e4;                                          # m^2
    Acre = 4.04685642241e3;                                 # m^2
    Barn = 1e-28;                                           # m^2
    Liter = 1e-3;                                           # m^3
    UsGallon = 3.78541178402e-3;                            # m^3
    Quart = 9.46352946004e-4;                               # m^3
    Pint = 4.73176473002e-4;                                # m^3
    Cup = 2.36588236501e-4;                                 # m^3
    FluidOunce = 2.95735295626e-5;                          # m^3
    Tablespoon = 1.47867647813e-5;                          # m^3
    Teaspoon = 4.92892159375e-6;                            # m^3
    CanadianGallon = 4.54609e-3;                            # m^3
    UkGallon = 4.546092e-3;                                 # m^3
    MilesPerHour = 4.4704e-1;                               # m / s
    KilometersPerHour = 2.77777777778e-1;                   # m / s
    Knot = 5.14444444444e-1;                                # m / s
    PoundMass = 4.5359237e-1;                               # kg
    OunceMass = 2.8349523125e-2;                            # kg
    Ton = 9.0718474e2;                                      # kg
    MetricTon = 1e3;                                        # kg
    UkTon = 1.0160469088e3;                                 # kg
    TroyOunce = 3.1103475e-2;                               # kg
    Carat = 2e-4;                                           # kg
    UnifiedAtomicMass = 1.660539040e-27;                    # kg per u, Da or AMU [See Dalton]
    GramForce = 9.80665e-3;                                 # kg m / s^2
    PoundForce = 4.44822161526e0;                           # kg m / s^2
    KilopoundForce = 4.44822161526e3;                       # kg m / s^2
    Poundal = 1.38255e-1;                                   # kg m / s^2
    Calorie = 4.1868e0;                                     # kg m^2 / s^2
    Btu = 1.05505585262e3;                                  # kg m^2 / s^2
    Therm = 1.05506e8;                                      # kg m^2 / s^2
    Horsepower = 7.457e2;                                   # kg m^2 / s^3
    Bar = 1e5;                                              # kg / m s^2
    StdAtmosphere = 1.01325e5;                              # kg / m s^2
    Torr = 1.33322368421e2;                                 # kg / m s^2
    MeterOfMercury = 1.33322368421e5;                       # kg / m s^2
    InchOfMercury = 3.38638815789e3;                        # kg / m s^2
    InchOfWater = 2.490889e2;                               # kg / m s^2
    Psi = 6.89475729317e3;                                  # kg / m s^2
    Poise = 1e-1;                                           # kg m^-1 s^-1
    Stokes = 1e-4;                                          # m^2 / s
    Stilb = 1e4;                                            # cd / m^2
    Lumen = 1e0;                                            # cd sr
    Lux = 1e0;                                              # cd sr / m^2
    Phot = 1e4;                                             # cd sr / m^2
    Footcandle = 1.076e1;                                   # cd sr / m^2
    Lambert = 1e4;                                          # cd sr / m^2
    Footlambert = 1.07639104e1;                             # cd sr / m^2
    Curie = 3.7e10;                                         # 1 / s
    Roentgen = 2.58e-4;                                     # A s / kg
    Rad = 1e-2;                                             # m^2 / s^2
    SolarMass = 1.98892e30;                                 # kg
    BohrRadius = 5.291772083e-11;                           # m
    Newton = 1e0;                                           # kg m / s^2
    Dyne = 1e-5;                                            # kg m / s^2
    Joule = 1e0;                                            # kg m^2 / s^2
    Erg = 1e-7;                                             # kg m^2 / s^2
    StefanBoltzmannConstant = 5.67040047374e-8;             # kg / K^4 s^3
    ThomsonCrossSection = 6.65245893699e-29;                # m^2
    BohrMagneton = 9.27400899e-24;                          # A m^2
    NuclearMagneton = 5.05078317e-27;                       # A m^2
    ElectronMagneticMoment = 9.28476362e-24;                # A m^2
    ProtonMagneticMoment = 1.410606633e-26;                 # A m^2
    Faraday = 9.64853429775e4;                              # A s / mol
    ElectronCharge = 1.602176487e-19;                       # A s
    VacuumPermittivity = 8.854187817e-12;                   # A^2 s^4 / kg m^3
    VacuumPermeability = 1.25663706144e-6;                  # kg m / a^2 s^2
    Debye = 3.33564095198e-30;                              # A s^2 / m^2
    Gauss = 1e-4;                                           # kg / A s^2
    ThomsonCrossSection = 6.65245893699e-29;                # m^2
    KotovCosmicPeriodicity = 9600.061e0;                    # s
    CmbTemperature = 2.7255;
    SanchezCmbTemperature = 2.725820831e0;                  # K
    UniverseSanchezMass = 8.7936e52;                        # kg
    SanchezElectrConstant = 137.035999139e0;                # dimensionless Sanchez Electric Constant
    α = 7.2973525664e-3;                                    # Fine Structure Constant
    ε0 = 8.854187817e-12;                                   # F m-1 electric constant
    μ0 = 12.566370614e-7;                                   # N A-2     magnetic constant
    γ = 0.5772156649015328606065120900824024310421593359e0; # Euler–Mascheroni constant
    e = 2.7182818284590450907955982984276488423343144531e0; # Euler
    a = 137.035999139e0;                                    # Sanchez Electric Constant
    a_G = 1.32613243113089e-69                              # a_G gravitational coupling constant
    #α^-1 = 137.035999139e0                                 # Inverse fine-structure Constant
end


include("../functions/Functions.jl")
include("../evaluate/Evaluate.jl")
include("../opt/SAGA.jl")
include("../opt/SAGA_PROX.jl")
include("../opt/sgd.jl")
include("../data/data.jl")
include("../plot/Plotting.jl")

using ProximalOperators
using Plotting

X, y = getFullData()

# init initial point x_0
x0 = zeros(size(X,2),1)

lambda = 0.1

iter = 1000
step =0.001

#f = f_logreg(X,y,x0,lambda)
#g = sub_g_logreg()

f = f_svm(X,y,x0,lambda)
g = sub_g_svm()

prox_operator = SqrNormL2(lambda)

@time (theta, strings, iter,val_array_saga) = pgd(X,y,x0,f,g,1e-8,lambda,iter,step,size(X,1),prox_operator)

println(evaluate(X,y,theta,"svm",false))

#@time (theta,state,val_array,stop_array) = sgd(X,y,x0,f,g,1e-8,lambda,iter,step,size(X,1))

plotDoubleArray(val_array_saga,val_array, "saga1.svg")

println(evaluate(X,y,theta,"svm",false))
# Lippschitz = largest eigenvalue of Hessian*Hessian'


export solve_modelTRDiagAbs
function solve_modelTRDiagAbs(H, g, nlp_stop, PData::PDataFact, δ::T) where {T}
    # Solve the TR subproblem once diagonalized into Δ using the norm |Δ|
    # Setup the subproblem to recover identity regularizing matrix
    # Should this be in pre-process so that PData.λ is indeed used as updated?

    ḡ = PData.g̃
    n_g = norm(ḡ)
    ϵ = sqrt(eps(T)) / T(100.0) # *n_g
    Γ2 = max(maximum(abs.(PData.Δ)), ϵ)
    Γ = sqrt(Γ2)
    # @show Γ
    # @show PData.Δ

    Δ̄ = Γ .\ PData.Δ ./ Γ
    Δ = PData.Δ
    if PData.success    # initialize λmin, take care of hard case and Newton's direction interior
        l_m, i = findmin(Δ̄)
        λ = max(-l_m, 0.0)
        λmin = max(ϵ, λ .+ ϵ * (1.0 .+ λ)) #λ + ϵ #max(λ*(1+ϵ), ϵ)

        d̄ = -(Δ .+ λmin .* Γ2) .\ ḡ
        normd_s = sqrt(d̄ ⋅ d̄)
        #if isnan(normd_s)  @bp end
        if normd_s < δ
            if λ != 0.0  # λ == 0 is Newton's direction, nothing to do...
                #println(" hard case")
                d̄[i] = 0.0
                d̄[i] = -sign(PData.g̃[i]) * sqrt(δ^2 - d̄ ⋅ d̄)
            end
        else
            #println(" easy case normd_s = $normd_s")
            d̄, λ = solve_diagTR(λmin, Δ, ḡ, δ, ϵ, M = Γ2)
        end
    else # hard case impossible, λ already > λ_min
        # Solve the subproblem (Δ + λ I) d_s = -g_s such that δ = ||d_s||
        #println(" retry for success case")
        d̄, λ = solve_diagTR(PData.λ, Δ, ḡ, δ, ϵ, M = Γ2)
    end

    PData.λ = λ

    # Transform back d̄ into d
    d̃ = d̄ #Γ .\ d̄
    d = AInv(PData, d̃)
    #try assert((PData.g̃ + 0.5*PData.Δ .* d̃)⋅d̃ <= 0.0)  catch  @bp  end
    #println("*******SolveModelDiagAbs:  PData.g̃⋅d̃ = $(PData.g̃⋅d̃), 0.5 d̃'PData.Δd̃ = $(0.5*(PData.Δ .* d̃)⋅d̃)")
    return d, λ
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _premat(equinox1::T, equinox2::T, FK4::Bool) where {T<:AbstractFloat}
    t = (equinox2 - equinox1) / 1000
    if FK4
        st = (equinox1 - 1900) / 1000
        #  Compute 3 rotation angles
        a = sec2rad(t*(23042.53 + st*(139.75 + 0.06*st)
                       + t*(30.23 - 0.27*st + 18.0*t)))
        b = sec2rad(t*t*(79.27 + 0.66*st + 0.32*t)) + a
        c = sec2rad(t*(20046.85 - st*(85.33 + 0.37*st)
                       + t*(-42.67 - 0.37*st -41.8*t)))
    else
        st = (equinox1 / 1000) - 2
        # Compute 3 rotation angles
        a = sec2rad(t *(23062.181 + st*(139.656 + 0.0139*st)
                        + t*(30.188 - 0.344*st + 17.998*t)))
        b = sec2rad(t*t*(79.280 + 0.410*st + 0.205*t)) + a
        c = sec2rad(t*(20043.109 - st*(85.33 + 0.217*st)
                       + t*(-42.665 - 0.217*st - 41.833*t)))
    end
    sa, ca = sincos(a)
    sb, cb = sincos(b)
    sc, cc = sincos(c)
    return @SMatrix [ca * cb * cc - sa * sb   -ca * sb - sa * cb * cc   -cb*sc;
                     sa * cb + ca * sb * cc    ca * cb - sa * sb * cc   -sb*sc;
                     ca * sc                  -sa * sc                   cc]
end

"""
    premat(equinox1, equinox2[, FK4=true]) -> precession_matrix

### Purpose ###

Return the precession matrix needed to go from `equinox1` to `equinox2`.

### Explanation ###

This matrix is used by `precess` and `baryvel` to precess astronomical
coordinates.

### Arguments ###

* `equinox1`: original equinox of coordinates.
* `equinox2`: equinox of precessed coordinates.
* `FK4` (optional boolean keyword): if this keyword is set to `true`, the FK4
  (B1950.0) system precession angles are used to compute the precession matrix.
  When it is `false`, the default, use FK5 (J2000.0) precession angles.

### Output ###

A 3×3 matrix, used to precess equatorial rectangular coordinates.

### Example ###

Return the precession matrix from 1950.0 to 1975.0 in the FK4 system

```jldoctest
julia> using AstroLib

julia> premat(1950,1975,FK4=true)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9} with indices SOneTo(3)×SOneTo(3):
 0.999981    -0.00558775  -0.00242909
 0.00558775   0.999984    -6.78691e-6
 0.00242909  -6.78633e-6   0.999997
```

### Method ###

FK4 constants from "Computational Spherical Astronomy" by Taff (1983),
p. 24. (FK4). FK5 constants from "Explanatory Supplement To The Astronomical
Almanac" 1992, page 104 Table 3.211.1
(https://archive.org/details/131123ExplanatorySupplementAstronomicalAlmanac).

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
premat(eq1::Real, eq2::Real; FK4::Bool=false) =
    _premat(promote(float(eq1), float(eq2))..., FK4)


using ApproxFunOrthogonalPolynomials, ApproxFunBase, SpecialFunctions, LinearAlgebra, Test

@testset "broadcast" begin
    @testset "Special functions" begin
        x = Fun()
        @test exp(x) ≈ exp.(x) atol=10eps()
        f = Fun(exp)
        @test exp.(f) ≈ exp(f)
        @test besselj.(1,f) ≈ besselj(1,f)

        @test atan.(f,1)(0.1) ≈ atan(f(0.1),1)
        @test atan.(f,f)(0.1) ≈ atan(f(0.1),f(0.1))
    end

    @testset  "triplet with mixed scalar-fun" begin
        f = Fun(exp)
        @test ((a,b,c) -> exp(a +b+c)).(f,2.0,f) ≈ exp(2f + 2.0)
    end

    @testset "Arrays" begin
        x = Fun()
        @test exp.( x .+ [1,2,3]) isa Vector
        @test exp.(x .+ [1,2]) ≈ [exp(x+1),exp(x+2)]

        @test maximum.( [x;x] .+ [x;x]) isa Fun
        @test maximum.( [x;x] .+ [x;x]) ≈ 2x

        @test exp.([x,x] .+ 2 .+ [x;x]) isa Vector
        @test exp.( [x,x] .+ 2 .+ [x;x]) ≈ [exp(2x+2),exp(2x+2)]

        # since array broadcasting takes presidence, the following
        # does not try to call a constructor:
        @test ApproxFunBase.tocanonical.(x,[0.1,0.2]) ≈ [0.1,0.2]
    end

    @testset "broadcast!" begin
        x = Fun()
        f = Fun()
        f .= exp.(x)
        @test ≈(exp(x),f;atol=10eps())

        f = Fun(Ultraspherical(1))
        f .= exp.(x)
        @test f(0.1) ≈ exp(0.1)
    end

    @testset "norm" begin
        f = Fun(x ->[exp(x),sin(x)])
        F = [Fun(exp),Fun(sin)]
        @test norm.(f) ≈ Fun(x ->norm([exp(x),sin(x)]))
        @test norm.(F) ≈ [norm(Fun(exp)),norm(Fun(sin))] # isa Vector{Float64}

        @test norm.(f,[1,2]) ≈ norm.(F,[1,2]) ≈ [norm(f[1],1),norm(f[2],2)] # isa Vector{Float64}
    end
end


@symbol_func function T_rho(cur_reactor::AbstractReactor, cur_rho::AbstractSymbol)
  cur_T_bar = cur_reactor.T_bar

  cur_nu_T = cur_reactor.nu_T

  cur_T_rho = 1 - cur_rho ^ 2

  cur_T_rho ^= cur_nu_T

  cur_T_rho *= cur_T_bar

  cur_T_rho *= 1 + cur_nu_T

  cur_T_rho
end


struct ScalarFunctions{T,F}
    coord::T
    func::F
end

### Radial Functions
struct HeunConfluentRadial
    η::Complex{Float64}; α::Complex{Float64}; ξ::Complex{Float64}; ζ::Complex{Float64}
    r₊::Float64; r₋::Float64
    coeffs::Array{Complex{Float64},1}
end

function (ψᵣ::HeunConfluentRadial)(r::Union{Float64,Complex{Float64}})
    η = ψᵣ.η;
    α = ψᵣ.α;
    ξ = ψᵣ.ξ;
    ζ = ψᵣ.ζ;
    r₊ = ψᵣ.r₊;
    r₋ = ψᵣ.r₋
    asymptoticpart = (r₊-r₋)^(α)*(im*(r-r₋))^(η-α)*(im*(r-r₊))^(ξ)*exp(ζ*r)
    x = (r-r₊)/(r-r₋)
    finalsum = 0
    for n in 1:length(ψᵣ.coeffs)
       finalsum += ψᵣ.coeffs[n]*x^(n-1)
    end
    #print(asymptoticpart)
    asymptoticpart*finalsum
end

function ComplexPlot(ψ::HeunConfluentRadial; ztopleft = 0.8 + 1.0im, zbottomright = 2.0 - 0.5im)
    img = portrait(ztopleft, zbottomright, Ψ;
        point_color = cs_d(; colormap=hsv_colors()))
    display(img)
    img
end

### Angular Functions
struct SpinWeightedSpherical
    s::Int64; l::Int64; m::Int64
end

function sterlings(n)
    return sqrt(2*pi*n)*(n/exp(1))^n
end

function SpinWeightedSphericalCalculation(z,s,l,m)
    ### Convention is that θ goes from 0 to pi, and hence cos(θ)
    ### goes from 1 to -1, which means θ/2 is in the first quadrant
    ### hence all cot terms are positive, which is why I take the
    ### positive root whereever needed.
    ### Pulled from
    ### https://en.wikipedia.org/wiki/Spin-weighted_spherical_harmonics
    #Define the overall factor
    #println((l+m > 20),"  ",(l+s > 20),"  ",(l-m > 20),"  ",(l-s > 20))
    #println(((l+m > 20) | (l+s > 20) | (l-m > 20) | (l-s > 20)))
    if ((l+m > 20) | (l+s > 20) | (l-m > 20) | (l-s > 20))
        term1 = (2*l+1)*(sterlings(l+m)/sterlings(l+s))*(sterlings(l-m)/sterlings(l-s))
    else
        term1 = (2*l+1)*(factorial(l+m)/factorial(l+s))*(factorial(l-m)/factorial(l-s))
    end
    #println((l=l,m=m,s=s), (fact = factorial(l+m),), term1)
    A = sqrt(term1/(4*pi))*(-1)^m

    #The sin term right outside the summation
    sinterm = ((1-z)/2)^l

    #The summation terms
    sumterms = 0
    for r = 0:(l-s)
        consts = binomial(l-s,r)*binomial(l+s,r+s-m)*(-1)^(l-r-s)
        cotterm = ((1+z)/(1-z))^((2*r+s-m)/2)
        sumterms += consts*cotterm
    end

    return A*sinterm*sumterms
end

function (Ψ::SpinWeightedSpherical)(z)
    s = Ψ.s; l = Ψ.l; m = Ψ.m;
    if (l >= min(abs(m),abs(s))) & (l >= abs(m))
        return SpinWeightedSphericalCalculation(z,s,l,m)
    end
    return Complex(0.0)
end

function (Ψ::SpinWeightedSpherical)(z,ϕ)
    Ψ(z)*exp(im*Ψ.m*ϕ)
end

struct SpinWeightedSpheroidal
    s::Int64; l::Int64; m::Int64
    Cllʼ::Array{Complex{Float64},1}
    lmin::Int64; lmax::Int64
end

function SpinWeightedSpheroidal(s,l,m,Cllʼ)
    lmins = max(abs(s),abs(m));
    lmax = length(Cllʼ) + lmins -1
    SpinWeightedSpheroidal(s,l,m,Cllʼ,lmins, lmax)
end

function SpinWeightedSpheroidalCalculation(z,s,l,m,Cllʼ,lmin,lmax)
    N = lmax - lmin + 1
    val = Complex(0.0)
    for j = 1:N
        lʼ = j+lmin-1;
        #println(lʼ)
        val += Cllʼ[j]*SpinWeightedSphericalCalculation(z,s,lʼ,m)
    end
    val
end

function (Ψ::SpinWeightedSpheroidal)(z)
    s = Ψ.s; l = Ψ.l; m = Ψ.m;
    lmin = Ψ.lmin; lmax = Ψ.lmax; Cllʼ = Ψ.Cllʼ;
    SpinWeightedSpheroidalCalculation(z,s,l,m,Cllʼ,lmin,lmax)
end

(Ψ::SpinWeightedSpheroidal)(z,ϕ) = Ψ(z)*exp(im*Ψ.m*ϕ)

function RadialCoefficients(D₀, D₁, D₂, D₃, D₄; N = 250)
    αₙ(n) = (n+1)*(n+D₀)
    βₙ(n) = -2*n^2 + (D₁+2)*n + D₃
    γₙ(n) = (n-1)*(n+D₂-2) + D₄

    an(n) = βₙ(n)/αₙ(n) ; bn(n) = γₙ(n)/αₙ(n);

    ##Set largest rN value
    u1,u2,u3,u4 = rNCoeffs(D₀,D₁,D₂,D₃,D₄)
    rN = 1 + u1*N^(-0.5) + u2*N^(-1) + u3*N^(1.5) + u4*N^(-2)

    ComputeSeriesFromab(an,bn; rN=rN, PreN=0)
end

function ComputeSeriesFromab(an::Function,bn::Function; N=250, rN = 0.0*im, PreN = 40)
    ##Initialize rn and fn vectors
    rₙ = zeros(Complex{Float64},N+1)
    fₙ = zeros(Complex{Float64},N+1)

    rold = rN
    ##Startup Pass for rₙ
    for n = (N+PreN):-1:(N+1)
        rnew = -bn(n)/(an(n) + rold)
        rold = rnew
    end
    rₙ[N+1] = rold;
    ##Pass for rn
    for n= N:-1:1
        rₙ[n] = -bn(n)/(an(n) + rₙ[n+1])
    end

    fₙ[1] = 1;
    ##Pass for fn
    for n = 2:(N+1)
        fₙ[n] = rₙ[n-1]*fₙ[n-1]
    end
    fₙ
end

### Combining the Mode radial and angular Information
struct QuasinormalModeFunction
    s::Int64; l::Int64; m::Int64; a::Float64
    ω::Complex{Float64}
    Alm::Complex{Float64}
    R::HeunConfluentRadial
    S::SpinWeightedSpheroidal
end

function qnmfunction(; s=-2,l=2,m=2,n=0,a=0.00)
    ω, Alm, Cllʼ = qnm(l=l,m=m,s=s,n=n, a=a)

    ((ζ,ξ,η),(p,α,γ,δ,σ),(D₀,D₁,D₂,D₃,D₄)) = ParameterTransformations(l,m,s,a,ω,Alm)
    r₊ = 1 + sqrt(1-a^2); r₋ = 1 - sqrt(1-a^2)

    ##Radial WaveFunction
    aₙ = RadialCoefficients(D₀, D₁, D₂, D₃, D₄)
    Ψᵣ = HeunConfluentRadial(η,α,ξ,ζ,r₊,r₋,aₙ)

    ##Angular WaveFunction
    Ψᵪ = SpinWeightedSpheroidal(s,l,m,Cllʼ)

    QuasinormalModeFunction(s,l,m,a,ω,Alm,Ψᵣ,Ψᵪ)
end

(Ψ::QuasinormalModeFunction)(r::Number) =  Ψ.R(r)
(Ψ::QuasinormalModeFunction)(r::Number, θ::Number) =  Ψ.R(r)*Ψ.S(cos(θ))
(Ψ::QuasinormalModeFunction)(r::Number, θ::Number, ϕ::Number) =  Ψ.R(r)*Ψ.S(cos(θ))*exp(im*Ψ.m*ϕ)
(Ψ::QuasinormalModeFunction)(r::Number, θ::Number, ϕ::Number, t::Number) =  Ψ.R(r)*Ψ.S(cos(θ))*exp(im*Ψ.m*ϕ)*exp(-im*Ψ.ω*t)

(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r,),Tuple{Number}}) = Ψ.R(x[:r])
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:θ,),Tuple{Number}}) = Ψ.S(cos(x[:θ]))
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:z,),Tuple{Number}}) = Ψ.S(x[:z])
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r, :θ),Tuple{Number,Number}}) = Ψ.R(x[:r])*Ψ.S(cos(x[:θ]))
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r, :z),Tuple{Number,Number}}) = Ψ.R(x[:r])*Ψ.S(x[:z])
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r, :θ, :ϕ),Tuple{Number,Number,Number}}) = Ψ.R(x[:r])*Ψ.S(cos(x[:θ]))*exp(im*Ψ.m*ϕ)
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r, :z, :ϕ),Tuple{Number,Number,Number}}) = Ψ.R(x[:r])*Ψ.S(x[:z])*exp(im*Ψ.m*ϕ)
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r, :θ, :ϕ, :t),Tuple{Number,Number,Number,Number}}) = Ψ.R(x[:r])*Ψ.S(cos(x[:θ]))*exp(im*Ψ.m*ϕ)
(Ψ::QuasinormalModeFunction)(x::NamedTuple{(:r, :z, :ϕ, :t),Tuple{Number,Number,Number,Number}}) = Ψ.R(x[:r])*Ψ.S(x[:z])*exp(im*Ψ.m*ϕ)

print("Here")


using PDMats
import Statistics: mean, var
using LinearAlgebra

"""
   `calc_periodogram_orig(times, y_obs, σ_ys; options...)`
Computes a periodogram for times, observations and uncertainties.

Inputs:
- times: obseration times
- y_obs: observed values at each time
- σ_ys:  measurement uncertainty at each time (may be either a Vector or positive definite matrix)
Optional parameters:

- period_min (2.0)
- period_max (4*span of times)
- num_periods (4000)

Output (as NamedTuple):
- period_grid: orbital periods searched
- periodogram: χ² for best fit given each orbital period
- period_best_fit: period searched resulting to minimum χ²
- coeff_best_fit: best-fit coefficients for period_best_fit
- phase_best_fit: orbital phase for period_best_fit at time = 0
- predict: predictions of best-fit model at each time
- rms: root mean square of predictions (of best-fit model) minus observations
"""
function calc_periodogram_orig  end

function calc_periodogram_orig(t::Vector, y_obs::Vector, σ_ys::Vector; kwargs...)
 	calc_periodogram_orig(t, y_obs, PDiagMat(σ_ys.^2); kwargs...)
end

function calc_periodogram_orig(t::Vector, y_obs::Vector, covar_mat::MT;
	    		period_min::Real = 2.0, period_max::Real = 4*(maximum(t)-minimum(t)), num_periods::Integer = 4000) where {
						T3<:Real, MT<:AbstractPDMat{T3} }
	period_grid =  1.0 ./ range(1.0/period_max, stop=1.0/period_min, length=num_periods)
	periodogram = map(p->-0.5*calc_χ²_general_linear_least_squares_orig(calc_design_matrix_circ(p,t),covar_mat,y_obs),period_grid)
	period_fit = period_grid[argmax(periodogram)]
	design_matrix_fit = calc_design_matrix_circ(period_fit,t)
	coeff_fit = fit_general_linear_least_squares_orig(design_matrix_fit,covar_mat,y_obs)
	phase_fit = atan(coeff_fit[1],coeff_fit[2])
	pred = design_matrix_fit * coeff_fit
	rms = sqrt(mean((y_obs.-pred).^2))
	return (;period_grid=period_grid, periodogram=periodogram, period_best_fit = period_fit, coeff_best_fit=coeff_fit, phase_best_fit=phase_fit, predict=pred, rms=rms )
end

"""
   `calc_periodogram(times, y_obs, σ_ys; options...)`
Computes a periodogram for times, observations and uncertainties.

Inputs:
- times: obseration times
- y_obs: observed values at each time
- σ_ys:  measurement uncertainty at each time (may be either a Vector or positive definite matrix)
Optional parameters:

- period_min (2.0)
- period_max (4*span of times)
- num_periods (4000)

Output (as NamedTuple):
- period_grid: orbital periods searched
- periodogram: χ² for best fit given each orbital period
- period_best_fit: period searched resulting to minimum χ²
- coeff_best_fit: best-fit coefficients for period_best_fit
- phase_best_fit: orbital phase for period_best_fit at time = 0
- predict: predictions of best-fit model at each time
- rms: root mean square of predictions (of best-fit model) minus observations
"""
function calc_periodogram end

function calc_periodogram(t::Vector, y_obs::Vector, σ_ys::Vector; kwargs...)
 	calc_periodogram(t, y_obs, PDiagMat(σ_ys.^2), kwargs...)
end

function calc_periodogram(t::Vector, y_obs::Vector, covar_mat::MT;
				period_min::Real = 2.0, period_max::Real = 4*(maximum(t)-minimum(t)), num_periods::Integer = 4000) where {
						MT<:AbstractPDMat }
	period_grid =  1.0 ./ range(1.0/period_max, stop=1.0/period_min, length=num_periods)
	design_matrix_workspace = Array{eltype(t)}(undef,length(t),3)
	calc_χ²_workspace = prealloc_workspace(design_matrix_workspace,covar_mat,y_obs)
	periodogram = map(p->-0.5*calc_χ²_general_linear_least_squares(calc_design_matrix_circ!(design_matrix_workspace,p,t),covar_mat,y_obs,workspace=calc_χ²_workspace),period_grid)
	period_fit = period_grid[argmax(periodogram)]
	design_matrix_fit = calc_design_matrix_circ!(design_matrix_workspace,period_fit,t)
	coeff_fit = fit_general_linear_least_squares(design_matrix_fit,covar_mat,y_obs)
	phase_fit = atan(coeff_fit[1],coeff_fit[2])
	pred = design_matrix_fit * coeff_fit
	rms = sqrt(mean((y_obs.-pred).^2))
	return (;period_grid=period_grid, periodogram=periodogram, period_best_fit = period_fit, coeff_best_fit=coeff_fit, phase_best_fit=phase_fit, predict=pred, rms=rms )
end

# Internal routines called by calc_periodogram[_orig]
function calc_design_matrix_circ!(result::AM, period::Real, times::AV) where { R1<:Real, AM<:AbstractMatrix{R1}, AV<:AbstractVector{R1} }
	n = length(times)
	@assert size(result,1) == n
	@assert 2<= size(result,2) <=3
	for i in 1:n
		( result[i,1], result[i,2] ) = sincos(2π/period .* times[i])
	end
	if size(result,2) == 3
		result[:,3] .= one(eltype(result))
	end
	return result
end

function calc_design_matrix_circ(period::Real, times::AV) where { R1<:Real, AV<:AbstractVector{R1} }
	n = length(times)
	dm = Array{promote_type(typeof(period),eltype(times))}(undef,n,3)
	calc_design_matrix_circ!(dm,period,times)
	return dm
end

function prealloc_workspace( design_mat::ADM, covar_mat::APD, obs::AA) where { ADM<:AbstractMatrix, APD<:AbstractPDMat, AA<:AbstractArray }
	(;	invA_X = Array{Float64,2}(undef,size(covar_mat,1),size(design_mat,2)),
		Xt_invA_X =  Array{Float64,2}(undef,size(design_mat,2),size(design_mat,2)),
		invA_y =  Array{Float64,1}(undef,length(obs)),
		X_invA_y =  Array{Float64,1}(undef,size(design_mat,2)),
		AB_hat = Array{Float64,1}(undef,size(design_mat,2) ),
		predict =  Array{Float64,1}(undef,length(obs))  )
end

function fit_general_linear_least_squares( design_mat::ADM, covar_mat::APD, obs::AA;
			workspace = prealloc_workspace(design_mat,covar_mat,obs)
	 			) where { ADM<:AbstractMatrix, APD<:PDiagMat, AA<:AbstractArray }
	#workspace.invA_X .= covar_mat \ design_mat
	workspace.invA_X .= covar_mat.inv_diag .* design_mat
	mul!(workspace.Xt_invA_X,design_mat', workspace.invA_X)
	#workspace.inv A_y .= covar_mat \ obs
	workspace.invA_y .= covar_mat.inv_diag .* obs
	mul!(workspace.X_invA_y,  design_mat', workspace.invA_y)
	#workspace.AB_hat .= workspace.Xt_invA_X \ workspace.X_invA_y
	lufac = lu!( workspace.Xt_invA_X, check=false)
	workspace.AB_hat .= lufac \ workspace.X_invA_y
end

function fit_general_linear_least_squares_orig( design_mat::ADM, covar_mat::APD, obs::AA ) where { ADM<:AbstractMatrix, APD<:AbstractPDMat, AA<:AbstractArray }
	Xt_inv_covar_X = Xt_invA_X(covar_mat,design_mat)
	X_inv_covar_y =  design_mat' * (covar_mat \ obs)
	AB_hat = Xt_inv_covar_X \ X_inv_covar_y                            # standard GLS
end

function predict_general_linear_least_squares( design_mat::ADM, covar_mat::APD, obs::AA;
			workspace = prealloc_workspace(design_mat,covar_mat,obs)
			 	) where { ADM<:AbstractMatrix, APD<:AbstractPDMat, AA<:AbstractArray }
	#param =
	fit_general_linear_least_squares(design_mat,covar_mat,obs,workspace=workspace)
	mul!(workspace.predict, design_mat, workspace.AB_hat)
	return workspace.predict
end

function predict_general_linear_least_squares_orig( design_mat::ADM, covar_mat::APD, obs::AA ) where { ADM<:AbstractMatrix, APD<:AbstractPDMat, AA<:AbstractArray }
	param = fit_general_linear_least_squares_orig(design_mat,covar_mat,obs)
	design_mat * param
end

function calc_χ²_general_linear_least_squares_orig( design_mat::ADM, covar_mat::APD, obs::AA ) where { ADM<:AbstractMatrix, APD<:AbstractPDMat, AA<:AbstractArray }
	pred = predict_general_linear_least_squares_orig(design_mat,covar_mat,obs)
	invquad(covar_mat,obs-pred)
end

function calc_χ²_general_linear_least_squares( design_mat::ADM, covar_mat::APD, obs::AA;
				workspace = prealloc_workspace(design_mat,covar_mat,obs)
				 ) where { ADM<:AbstractMatrix, APD<:AbstractPDMat, AA<:AbstractArray }
	#workspace.predict .=
	predict_general_linear_least_squares(design_mat,covar_mat,obs,workspace=workspace)
	workspace.predict .-= obs
	invquad(covar_mat,workspace.predict)
end


using ValidatedNumerics

import Base.mod2pi

function standard_map(X::IntervalBox, k = 1.0)
    p, θ = X

    p′ = p + k*sin(θ)
    @show p′
    θ′ = mod2pi( θ + p′ )
    p′ = mod2pi(p′)

    @show p′, θ′
    @show typeof(p′), typeof(θ′)

    IntervalBox(p′, θ′)
end

function IntervalBox{T}(X::Vector{Interval{T}}, Y::Vector{Interval{T}})
    vec([IntervalBox(x, y) for x in X, y in Y])
end

function iterate(f, n, x)
    for i in 1:n
        x = f(x)
    end
    x
end

import Base.mod
function mod(X::Interval, width::Real)

    @show X, width

    X /= width

    if diam(X) >= 1.

        return [Interval(0, 1) * width]
    end

    a = X.lo - floor(X.lo)
    b = X.hi - floor(X.hi)

    if a < b
        return [Interval(a, b)*width]

    end

    return [Interval(0, b)*width, Interval(a, 1)*width]

end

function mod(X::IntervalBox, width::Real)
    x, y = X

    xx = mod(x, width)
    yy = mod(y, width)

    vec([IntervalBox(x, y) for x in xx, y in yy])
end

mod2pi{T}(x::Interval{T}) = mod(x, ValidatedNumerics.two_pi(T))

mod2pi{T}(X::Vector{Interval{T}}) = vcat(map(mod2pi, X)...)


using ConvexHull, JuMP, Polyhedra, Test

@testset "Throws error in case of linearity" begin
    simplex = LiftedHRepresentation([0 1 0; 0 0 1; 1 -1 -1], BitSet([3]))
    @test_throws ErrorException double_description(simplex)
    linv = LiftedHRepresentation([1 1], BitSet([1]))
    @test_throws ErrorException double_description(linv)
end

include("polyhedron.jl")


# test module for prob_stats

module TestProbStats
    using Test

    # target modules
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/sensor_data/print_sensor_data.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/freq_dist/histogram/draw_histogram.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/freq_dist/histogram/histogram_mean.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/freq_dist/variance/calc_variance.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/freq_dist/std_dev/calc_std_dev.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_dist/prob_dist.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_model/gauss_dist_model/gauss_dist_model.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_model/gauss_prob_dist/gauss_prob_dist.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_model/gauss_prob_dist_pdf/gauss_prob_dist_pdf.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_model/gauss_prob_dist_cdf/gauss_prob_dist_cdf.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_model/bar_prob_dist_cdf/bar_prob_dist_cdf.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/prob_model/expected_value/calc_expected_value.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/histogram_600mm/draw_hist_600.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/time_series_600mm/draw_time_series_600.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/group_by_hour/group_by_hour.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/hist_grp_by_hour/hist_grp_by_hour.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/group_hour_heatmap/group_hour_heatmap.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/prob_sum_hour/prob_sum_hour.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/prob_sum_lidar/prob_sum_lidar.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/cond_z_t_bar/cond_z_t_bar.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/probs_bayes/probs_bayes.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/bayes_estimation_630/bayes_estimation_630.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/bayes_est_at_5/bayes_est_at_5.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/complex_dist/bayes_est_at_11/bayes_est_at_11.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/covariance/calc_covariance.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/contour_pdf/contour_pdf.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/contour_pdf/contour_pdf_plus20.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/contour_pdf/contour_pdf_200.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/multiple_gauss_dist/multiple_gauss_dist.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/eigen/calc_plot_eigen.jl"))
    include(joinpath(split(@__FILE__, "test")[1], "src/prob_stats/multi_dim_gauss_dist/cov_matrix/calc_cov_mat.jl"))

    function main()
        @testset "ProbStats" begin
            @testset "SensorData" begin
                @test_nowarn PrintSensorData.print_data_200()
                @test_nowarn PrintSensorData.print_data_600()
            end
            @testset "FreqDist" begin
                @test_nowarn DrawHistogram.main(true)
                @test_nowarn HistogramMean.main(true)
                @test_nowarn CalcVariance.main()
                @test_nowarn CalcStdDev.main()
            end
            @testset "ProbDist" begin
                @test_nowarn ProbDist.main(true)
            end
            @testset "ProbModel" begin
                @testset "GaussDistModel" begin
                    @test GaussDistModel.p(190) == 2.064932530633184e-5
                    @test GaussDistModel.p(210) == 0.08231272044708592
                    @test GaussDistModel.p(230) == 1.2364903952298874e-5
                    @test_nowarn GaussDistModel.main(true)
                end
                @testset "GaussProbDist" begin
                    @test GaussProbDist.prob(193) == 0.0002254354061377026
                    @test GaussProbDist.prob(210) == 0.08187587024811953
                    @test GaussProbDist.prob(229) == 3.113715670488921e-5
                    @test_nowarn GaussProbDist.main(true)
                end
                @testset "GaussProbDistPdf" begin
                    @test_nowarn GaussProbDistPdf.main(true)
                end
                @testset "GaussProbDistCdf" begin
                    @test_nowarn GaussProbDistCdf.main(true)
                end
                @testset "BarProbDistCdf" begin
                    @test_nowarn BarProbDistCdf.main(true)
                end
                @testset "CalcExpValue" begin
                    @test_nowarn CalcExpValue.main()
                end
            end
            @testset "ComplexDist" begin
                @testset "DrawHist600" begin
                    @test_nowarn DrawHist600.main(true)
                end
                @testset "DrawTimeSeries600" begin
                    @test_nowarn DrawTimeSeries600.main(true)
                end
                @testset "GroupByHour" begin
                    @test_nowarn GroupByHour.main(true)
                end
                @testset "HistGrpByHour" begin
                    @test_nowarn HistGrpByHour.main(true)
                end
                @testset "GrpHrHeatmap" begin
                    @test_nowarn GrpHrHeatmap.main(true)
                end
                @testset "ProbSumHour" begin
                    @test ProbSumHour.main(true) == 1.0
                end
                @testset "ProbSumLidar" begin
                    @test ProbSumLidar.main(true) == 1.0
                end
                @testset "CondZtBar" begin
                    @test_nowarn CondZtBar.main(true)
                end
                @testset "ProbsBayes" begin
                    @test_nowarn ProbsBayes.main()
                end
                @testset "BayesEst630" begin
                    @test_nowarn BayesEst630.main(true)
                end
                @testset "BayesEstAt5" begin
                    @test_nowarn BayesEstAt5.main(true)
                end
                @testset "BayesEstAt11" begin
                    @test_nowarn BayesEstAt11.main(true)
                end
            end
            @testset "MultiDimGuassDist" begin
                @testset "CalcCovariance" begin
                    @test_nowarn CalcCovariance.main()
                end
                @testset "MultiGaussDist" begin
                    @test_nowarn MultiGaussDist.main(true)
                end
                @testset "CalcPlotEigen" begin
                    @test_nowarn CalcPlotEigen.main(true)
                end
                @testset "CalcCovMat" begin
                    @test_nowarn CalcCovMat.main()
                end
            end
        end
    end
end


using Plots

include(string(pwd(),"/src/simulation.jl"))
include(string(pwd(),"/gali_bayesian.jl"))
include(string(pwd(),"/src/likelihood.jl"))

dados,choques = simulate_dsge(GAMMA_0,GAMMA_1,PSI,PI,5000)

plot(dados[:,1], label = "Quartely CPI")
plot!(dados[:,2], label = "GDP")
plot!(dados[:,3], label = "i")

plot(dados[:,4], label = "v")
plot!(dados[:,3], label = "i")

param = [alfa,bet,epsilon,theta,sig,1,phi,phi_pi,phi_y,rho_v]

ss, filtered,loglike = log_like_dsge(param,dados[:,2])

plot(dados[1:100,1])
plot!(filtered[1:100,1])

param1 = [alfa*2,bet,epsilon,theta,sig,1,phi,phi_pi,phi_y,rho_v]

a2,ll = log_like_dsge(param1,dados[:,2])

bett = 0.70:0.005:1.15

vals_bet = collect(bett)

for i in 1:length(bett)
    param1 = [alfa,bett[i],epsilon,theta,sig,1,phi,phi_pi,phi_y,rho_v]
    ss,filtered,vals_bet[i] = log_like_dsge(param1,dados[:,2])
    println(i)
end

plot(bett,vals_bet)

alff = 0.10:0.01:0.9

vals_alfa = collect(alff)

for i in 1:length(alff)
    param1 = [alff[i],bet,epsilon,theta,sig,1,phi,phi_pi,phi_y,rho_v]
    ss,filtered,vals_alfa[i] = log_like_dsge(param1,dados[:,2])
end

plot(alff,vals)

bett = 0.70:0.005:1.15
alff = 0.10:0.01:0.9

vals = zeros(length(alff),length(bett))

for i in 1:length(alff), j in 1:length(bett)
    param1 = [alff[i],bett[j],epsilon,theta,sig,1,phi,phi_pi,phi_y,rho_v]
    ss,filtered,vals[i,j] = log_like_dsge(param1,dados[:,2])
end

plotly()

plot(alff,bett,vals, st=:surface, camera = (-45,45))
xaxis!("alfa")
yaxis!("beta")

using JLD2

@save "surface.jld2" alff bett vals


# Create example data
x = LinRange(-1, 1, 40)
y = LinRange(-1, 1, 40)
z = LinRange(-1, 1, 40)
v = 1 .- (x.^2 .+ y'.^2 .+ reshape(z,1,1,:).^2).^0.5 - 0.25 .* rand(40, 40, 40)
# Draw the 3d volume data using an emission model
volume(v, algorithm="mip")


# Postcompose with scaling and sum

"""
  Postcompose(f::ProximableFunction, a::Real=1.0, b::Real=0.0)

Returns the function `g(x) = a*f(x) + b`.
"""

immutable Postcompose{T <: ProximableFunction, S <: Real} <: ProximableFunction
  f::T
  a::S
  b::S
  function Postcompose(f::T, a::S, b::S)
    if a <= 0.0
      error("parameter a must be positive")
    else
      new(f, a, b)
    end
  end
end

is_prox_accurate(f::Postcompose) = is_prox_accurate(f.f)

Postcompose{T <: ProximableFunction, S <: Real}(f::T, a::S=1.0, b::S=0.0) = Postcompose{T, S}(f, a, b)

function (g::Postcompose){T <: RealOrComplex}(x::AbstractArray{T})
  return g.a*g.f(x) + g.b
end

function prox!{T <: RealOrComplex}(g::Postcompose, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  v = prox!(g.f, x, y, g.a*gamma)
  return g.a*v + g.b
end

function prox_naive{T <: RealOrComplex}(g::Postcompose, x::AbstractArray{T}, gamma::Real=1.0)
  y, v = prox_naive(g.f, x, g.a*gamma)
  return y, g.a*v + g.b
end

fun_name(f::Postcompose) = string("Postcomposition of ", fun_name(f.f))
fun_dom(f::Postcompose) = fun_dom(f.f)
fun_expr(f::Postcompose) = "x ↦ a*f(x)+b"
fun_params(f::Postcompose) = string("f(x) = ", fun_expr(f.f), ", a = $(f.a), b = $(f.b)")


# This implements the variant when A is dense and m > n

using LinearAlgebra

struct IndGraphSkinny{T <: RealOrComplex} <: IndGraph
  m::Int
  n::Int
  A::Array{T, 2}
  AA::Array{T, 2}
  F::Cholesky{T, Array{T, 2}} #LL factorization
  tmp::Array{T, 1}
end

function IndGraphSkinny(A::Array{T, 2}) where {T <: RealOrComplex}
  m, n = size(A)
  AA = A' * A

  F = cholesky(AA + I)
  #normrows = vec(sqrt.(sum(abs2.(A), 2)))

  #The tmp vector assumes that difference between m and n is not drastic.
  # If someone will decide to solve it using m >> n, then tmp vector might be
  # considered to have only n position required to prox esitmation and indicator
  # calculation might be converted to less efficient.
  tmp = Array{T, 1}(undef, m)
  IndGraphSkinny(m, n, A, AA, F, tmp)
end

function (f::IndGraphSkinny)(x::AbstractArray{T}, y::AbstractArray{T}) where
    {T <: RealOrComplex}
  # the tolerance in the following line should be customizable
  mul!(f.tmp, f.A, x)
  f.tmp .-= y
  if norm(f.tmp, Inf) <= 1e-10
    return 0.0
  end
  return +Inf
end

function prox!(
    x::AbstractArray{T},
    y::AbstractArray{T},
    f::IndGraphSkinny,
    c::AbstractArray{T},
    d::AbstractArray{T},
    gamma=1.0
    ) where {T <: RealOrComplex}

  # x[:] = f.F \ (c + f.A' * d)
  mul!(x, adjoint(f.A), d)
  x .+= c
  ldiv!(f.F, x)

  mul!(y, f.A, x)
  return 0.0
end

fun_name(f::IndGraphSkinny) = "Indicator of an operator graph defined by dense full column rank matrix"
# fun_dom(f::IndGraph) = "Array{Real,1}, Array{Complex,1}"
# fun_expr(f::IndGraph) = "x,y ↦ 0 if Ax = y, +∞ otherwise"
# fun_params(f::IndGraph) =
#   string( "A = ", typeof(f.A), " of size ", size(f.A))

function prox_naive(
    f::IndGraphSkinny,
    c::AbstractArray{T},
    d::AbstractArray{T},
    gamma=1.0
    ) where {T <: RealOrComplex}

  x = f.F \ (c + f.A' * d)
  return x, f.A * x, 0.0
end

function (f::IndGraphSkinny)(xy::AbstractVector{T}) where
  {T <: RealOrComplex}
  x, y = splitinput(f, xy)
  return f(x, y)
end

(f::IndGraphSkinny)(xy::Tuple{AbstractVector{T}, AbstractVector{T}}) where
  {T <: RealOrComplex} = f(xy[1], xy[2])


#    Problem : GROUP A
#    *********
#    The extended Dennis and Schnabel problem, as defined by Li.
#
#    Original SIF Source:
#    G. Li,
#    "The secant/finite difference algorithm for solving sparse
#    nonlinear systems of equations",
#    SIAM Journal on Optimization, (to appear), 1990.
#
#    Implementation translated from Source:
#    http://eprints.tsu.ge/234/14/Tests%20collection-K-F.pdf
#
#    EDENSCH.SIF classification  OUR2-AN-V-0
#
#    Number of variables (at least 2)
#
# Daniel Henderson, 08/2021

f = (x) -> begin
    fx = 0.0
    for i in firstindex(x):lastindex(x)-1
        item1 = x[i] - 2
        item2 = x[i]x[i+1] - 2x[i+1]
        item3 = x[i+1] + 1
        fx += 16 + item1^4 + item2^2 + item3^2
    end
    return fx
end

g! = (g, x) -> begin
    for i in firstindex(x):lastindex(x)-1
        item1 = x[i] - 2
        item2 = x[i]*x[i+1] - 2x[i+1]
        item3 = x[i+1] + 1
        g[i] += 4item1^3 + 2x[i+1]item2
        g[i+1] += 2item2*(x[i] - 2) + 2item3
    end
    return g
end

fg! = (g, x) -> begin
    fx = 0.0
    for i in firstindex(x):lastindex(x)-1
        item1 = x[i] - 2
        item2 = x[i]*x[i+1] - 2x[i+1]
        item3 = x[i+1] + 1
        fx += 16 + item1^4 + item2^2 + item3^2
        g[i] += 4item1^3 + 2x[i+1]item2
        g[i+1] += 2item2*(x[i] - 2) + 2item3
    end
    return fx, g
end

init = (n::Int=2000) -> begin

    return n, 8.0*ones(n)
end

@warn "TODO: Debug EDENSCH"
# TestSet["EDENSCH"] = UncProgram("EDENSCH", f, g!, fg!, init)


# Discrete univariate

convert(::Type{Binomial}, d::Bernoulli) = Binomial(1, d.p)

# Continuous univariate

convert(::Type{Gamma}, d::Exponential) = Gamma(1.0, d.θ)
convert(::Type{Gamma}, d::Erlang) = Gamma(d.α, d.θ)


# # Hypercube

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Extension/hypercube.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Extension/hypercube.ipynb)
# **Contributed by**: Benoît Legat

# Given a Sum-of-Squares constraint on an algebraic set:
# ```math
# g_1(x) = , \ldots, g_m(x) = 0 \Rightarrow p(x) \ge 0.
# ```
# We can either use the certificate:
# ```math
# p(x) = s(x) + \lambda_1(x) g_1(x) + \cdots + \lambda_m(x) g_m(x), s_0(x) \text{ is SOS},
# ```
# or
# ```math
# p(x) \equiv s(x) \pmod{\la g_1(x), \ldots, g_m(x) \ra}, s_0(x) \text{ is SOS}.
# ```
# the second one leads to a *simpler* SDP but needs to compute a *Gr\"obner* basis:
#  * SemialgebraicSets implements Buchberger's algorithm.
#  * The `@set` macro recognizes variable fix, e.g., `x = 1`
#     and provides shortcut.
#  * If you know a \alert{better} way to take modulo,
#    better create your \alert{own} type of algebraic set!

# We illustrate this in this example.

using DynamicPolynomials
@polyvar x[1:3]
p = sum(x)^2
using SumOfSquares
S = algebraicset([xi^2 - 1 for xi in x])

# We will now search for the minimum of `x` over `S` using Sum of Squares Programming.
# We first need to pick an SDP solver, see [here](https://jump.dev/JuMP.jl/v0.21.6/installation/#Supported-solvers) for a list of the available choices.

import CSDP
solver = optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true)

function min_algebraic(S)
    model = SOSModel(solver)
    @variable(model, α)
    @objective(model, Max, α)
    @constraint(model, c, p >= α, domain = S)
    optimize!(model)
    @show termination_status(model)
    @show objective_value(model)
end

min_algebraic(S)

# Note that the minimum is in fact `1`.
# Indeed, since each variables is odd (it is either `-1` or `1`)
# and there is an odd number of variables, their sum is odd.
# Therefore it cannot be zero!

# We can see that the Gröbner basis of `S` was computed

@show S.I.gröbnerbasis
S.I.algo

# The Gröbner basis is simple to compute in this case as the vector
# of `xi^2 - 1` is already a Gröbner basis.
# However, we still need to divide polynomials by the Gröbner basis
# which can be simplified in this case.

const MP = MultivariatePolynomials
const SS = SemialgebraicSets
struct HypercubeIdeal{V} <: SS.AbstractPolynomialIdeal
    variables::Vector{V}
end
struct HypercubeSet{V} <: SS.AbstractAlgebraicSet
    ideal::HypercubeIdeal{V}
end
MP.variables(set::HypercubeSet) = MP.variables(set.ideal)
MP.variables(ideal::HypercubeIdeal) = ideal.variables
MP.changecoefficienttype(set::HypercubeSet, ::Type) = set
SS.ideal(set::HypercubeSet) = set.ideal
function Base.rem(p, set::HypercubeIdeal)
    return MP.polynomial(map(MP.terms(p)) do term
        mono = MP.monomial(term)
        new_mono = one(mono)
        for (var, exp) in powers(mono)
            if var in set.variables
                exp = rem(exp, 2)
            end
            new_mono *= var^exp
        end
        MP.coefficient(term) * new_mono
    end)
end

H = HypercubeSet(HypercubeIdeal(x))

min_algebraic(H)

# Let's now try to find the correct lower bound:

function min_algebraic_rational(S, d)
    model = SOSModel(solver)
    @variable(model, q, SOSPoly(MP.monomials(x, 0:d)))
    deno = q + 1
    @constraint(model, c, deno * p >= deno, domain = S)
    optimize!(model)
    @show termination_status(model)
end

# With `d = 0`, it's the same as previously

min_algebraic_rational(H, 0)

# But with `d = 1`, we can find the correct lower bound

min_algebraic_rational(H, 1)


@testset "reducedMisfit" begin
	using jInv.Utils

(m,n) = (40,5)
M     = rand(m,n)
M     = (M.>.3)*1
R     = StrainRecon.sdiag(rand(m))
atrue = randomw(n)
Data  = M*atrue

param = getStrainReconParam(Data,n,R=R)

print("\t--test gradient of reduced misfit--")
function testGradRedMisfit(x,v=[])
    Jc,res,wres,M,dJ,H = reducedMisfit(param,x,doDerivative=true)
	if !(isempty(v))
        return Jc, dot(dJ,v)
	else
        return Jc
	end
end

atest = rand(n); atest /= sum(atest)
chkDer, = checkDerivative(testGradRedMisfit,atest)
@test chkDer

println("\t--test Hessian of reduced misfit--")
function testHessianRedMisfit(x,v=[])
    Jc,res,wres,M,dJ,H = reducedMisfit(param,x,doDerivative=true)
 	if !(isempty(v))
        return vec(dJ), H*v
	else
        return vec(dJ)
	end
end
chkDer, = checkDerivative(testHessianRedMisfit,atest)
@test chkDer
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

_rad2sec{T<:AbstractFloat}(rad::T) = rad2deg(rad)*3600.0

"""
    rad2sec(rad) -> seconds

### Purpose ###

Convert from radians to seconds.

### Argument ###

* `rad`: number of radians.  It can be either a scalar or a vector.

### Output ###

The number of seconds corresponding to `rad`.  If `rad` is an array, an array of
the same length is returned.

### Example ###

``` julia
rad2sec(1)
# => 206264.80624709636
```

### Notes ###

Use `sec2rad` to convert seconds to radians.
"""
rad2sec(rad::Real) = _rad2sec(float(rad))

@vectorize_1arg Real rad2sec


#Figure 5.9, 5.10
using Plots, SparseArrays, JLD2, LinearAlgebra, MultiscaleGraphSignalTransforms, PyCall
py"""
import numpy
xs = numpy.load('../data/texture_mask.npy')
"""
xs = PyArray(py"xs"o)

matrix_gt = reshape(xs,(512,512))[1:4:512,1:4:512]
heatmap(matrix_gt,ratio=1, yaxis =:flip, axis = false, color = :grays, colorbar = false)
max1 = maximum(matrix_gt[:])
min1 = minimum(matrix_gt[:])
matrix_gt = (matrix_gt[:] .- min1)./(max1 - min1)

JLD2.@load "../data/handcut_images.jld2"
heatmap(tmp["block5"],ratio=1, yaxis =:flip, axis = false, color = :grays, colorbar = false)
matrix = tmp["block5"]

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

sigma = 0.001
m, n = size(matrix)

############### Get the neighbors for affinity matrix computing
r = 5; # specify radius of neighbors
l = 2*r + 1

temp_x, temp_y = fill(1,l^2), fill(1,l^2)
temp_xy = CartesianIndices((1:l,1:l))
for i = 1:l^2
    temp_x[i], temp_y[i] = temp_xy[i][1], temp_xy[i][2]
end
temp_ind = ((temp_x .- (r + 1)).^2 + (temp_y .- (r + 1)).^2).<= r^2
neighbor_x = temp_x[temp_ind] .- (r + 1)
neighbor_y = temp_y[temp_ind] .- (r + 1)
# for any index(x,y), (x + neighbor_x, y - neigbor_y) are the neighbors to calculate affinity

################ Create affinity matrix
W = fill(0., (m*n,m*n))
for i = 1:m*n
    cur = CartesianIndices((m,n))[i]
    for j = 1:length(neighbor_x)
        if 1 <= cur[1] + neighbor_x[j] <= m && 1 <= cur[2] + neighbor_y[j] <= n
            tempd = LinearIndices((m,n))[cur[1] + neighbor_x[j], cur[2] + neighbor_y[j]]
            W[i,tempd] = exp(-(matrix_gt[i] - matrix_gt[tempd])^2/sigma)
        end
    end
end
W = sparse(W)

############### Preprocess to get G and GP
G = GraphSig(W, f = reshape(matrix,(length(matrix_gt),1)))
GP = partition_tree_fiedler(G)
dmatrix = ghwt_analysis!(G, GP=GP)

#####################
############# Haar
BS_haar = bs_haar(GP)
dvec_haar = dmatrix2dvec(dmatrix, GP, BS_haar)

############# Walsh
BS_walsh = bs_walsh(GP)
dvec_walsh = dmatrix2dvec(dmatrix, GP, BS_walsh)

############# GHWT_c2f
dvec_c2f, BS_c2f = ghwt_c2f_bestbasis(dmatrix, GP)

############# GHWT_f2c
dvec_f2c, BS_f2c = ghwt_f2c_bestbasis(dmatrix, GP)

############# eGHWT
dvec_eghwt, BS_eghwt = ghwt_tf_bestbasis(dmatrix, GP)

function top_vectors_plot(dvec::Array{Float64, 2}, BS::BasisSpec, GP::GraphPart; clims::Tuple{Float64,Float64} = (-0.01,0.01))
    sorted_ind = sortperm(abs.(dvec[:]), rev = true);
    plot(9,layout=(3,3),framestyle=:none, legend=false)
    for i=1:9
        dvecT = fill(0., size(dvec))
        #dvecT[sorted_ind[i]] = dvec_ghwt[sorted_ind[i]]
        dvecT[sorted_ind[i]] = 1
        f = ghwt_synthesis(dvecT, GP, BS)
        #print((maximum(f), minimum(f)))
        heatmap!(reshape(f, size(matrix)), subplot=i, ratio=1, yaxis=:flip, axis=false, color = :grays, clims = clims)
    end
    current()
end

#Figure 5.10b
top_vectors_plot(dvec_eghwt, BS_eghwt, GP)

function approx_error2(DVEC::Array{Array{Float64,1},1})
    plot(xaxis = "Fraction of Coefficients Retained", yaxis = "Relative Approximation Error")
    frac = 0.3
    T = ["eGHWT Haar basis","eGHWT Walsh basis","GHWT_c2f", "GHWT_f2c", "eGHWT best basis"]
    L = [(:dashdot,:orange),(:dashdot,:blue),(:solid, :red),(:solid, :green),(:solid, :black)]
    for i = 1:5
        dvec = DVEC[i]
        N = length(dvec)
        dvec_norm = norm(dvec,2)
        dvec_sort = sort(dvec.^2) # the smallest first
        er = sqrt.(reverse(cumsum(dvec_sort)))/dvec_norm # this is the relative L^2 error of the whole thing, i.e., its length is N
        p = Int64(floor(frac*N)) + 1 # upper limit
        plot!(frac*(0:(p-1))/(p-1), er[1:p], yaxis=:log, xlims = (0.,frac), label = T[i], line = L[i])
    end
end

#Figure 5.10a
approx_error2([dvec_haar[:], dvec_walsh[:], dvec_c2f[:], dvec_f2c[:], dvec_eghwt[:]])


```
oilrate(args...)  -->Plot

Recipe to Plot the Oil, Liquid Production.
Optionally you can make decline curve analysis, estimate a Forecast to either specific time or economic limit

Exponential Decline curve equation, Np estimation, Time to Economic Limit
`` { Q }_{ t }={ Q }_{ i }{ e }^{ -{ D }_{ i }t }``
``Np=\int _{ { t }_{ i } }^{ { t }_{ f } }{ { Q }_{ i }{ e }^{ -{ D }_{ i }t }dt }``
``Np=\frac { { -Q }_{ i } }{ { { D }_{ i } } } \left[ { e }^{ { -D }_{ i }{ t }_{ f } }-{ e }^{ { -D }_{ i }{ t }_{ i } } \right] ``
``{ t }_{ L }=\frac { -Ln\left( \frac { { Q }_{ L } }{ { Q }_{ i } }  \right)  }{ D } ``

The next table show the list of variables allowed:

|PropertyName|Args|Default|Input|Description
|---|---|---|---|---|
|Date|Mandatory|--|Array{Date,1}| Array of dates|
|OilRate|Mandatory|--|Array{Number,1}| Array of Rates|
|LiquidRate|Optional|--|Array{Number,1}| Array of Rates|
|Dlim|Optional|Dlim=false|Dlim=[Date(y,m,d)]|Initial Date to plot|
|mrange|Optional|mrange=false|mrange=number| Time interval in months to set xticks|
|dtf|Optional|dft="u yy"|dtf=DateFormat|Date format of xticks
|Liquid|Optional|Liquid=false|Liquid=Bool|If liquid array provided, choose if plot Liquid Rate|
|Dec|Optional|Dec=false|Dec=Bool|Set decline curve analysis to desired interval |
|Drange|Optional|Drange=false|Drange=[Date(y,m,d) Date(y,m,d)]|Decline curve range to fit|
|Forecast|Optional|Forecast=false|Forecast=Bool|Set to plot forecast estimation|
|Mforecast|Optional|Mforecast=12|Mforcast=Number|Set number of months to forecast <br> from last date of Declination range|
|EconLimit|Optional|EconLimit=false|EconLimit=Number|Set the Economic Limit to estimate time to forecast
|Ad|Optional|Ad=true|Ad=Bool|Set if apply Anormaly detection to decline curve|
```
@userplot OilRate

@recipe function f(h::OilRate; Dlim=false,Monthly=true, mrange=6, dft="u yy",
                                Liquid=false,
                                Dec=false, Drange=false,
                                Forecast=false, Mforecast=12,
                                Ad=true, EconLimit=false)
                  legend --> false
                  ylabel := "Rate [bbl/d]"

    if length(h.args)==2
        x, y = h.args
    elseif length(h.args)==3
        x, y, z = h.args
    end

    if Dec==true
        if Drange==false
            DaysDeclination=x
            RateDeclination=y
        else
            DaysDeclination=x[(x.>=Drange[1]) .& (x.<=Drange[2])]
            RateDeclination=y[(x.>=Drange[1]) .& (x.<=Drange[2])]
        end

        (S, DecTime, DecLine)=Declination(DaysDeclination,RateDeclination, AD=Ad)

        if Forecast==true
            if EconLimit!=false
                MforecastEL=round(Int,((log(EconLimit/S.Qi)/S.Di)/30.42))
                ForecastTime=DaysDeclination[end]:Dates.Month(1):DaysDeclination[1]+Dates.Month(MforecastEL)

            else
                ForecastTime=DaysDeclination[end]:Dates.Month(1):DaysDeclination[end]+Dates.Month(Mforecast)
            end
            ForecastLine=DeclinationForecast(S,ForecastTime)
            Np=(ForecastLine[1]-ForecastLine[end])/(-S.Di*1000)
            f=ForecastTime[end]+Dates.Month(1)
        else
            f=x[end]+Dates.Month(1)
        end
    else
        f=x[end]+Dates.Month(1)
    end

####
if Dlim==false
    if Monthly==true
      r = x[1]:Dates.Month(mrange):f
  else
     r = x[1]:Dates.Day(mrange):f
   end
    d = Dates.format.(r,dft)
else
    if Monthly==true
        r = Dlim[1]:Dates.Month(mrange):f
    else
        r= Dlim[1]:Dates.Day(mrange):f

    end
d = Dates.format.(r,dft)
end
xticks := (r,d)
if Dlim!=false
xlim := (Dates.value(Dlim[1]),Dates.value(Dlim[2]))
end
###
#=
    if Dlim==false
        r = x[1]:Dates.Month(mrange):f
        d = Dates.format.(r,dft)
    else
        r = Dlim[1]:Dates.Month(mrange):f
        d = Dates.format.(r,dft)
        xlim := (Dates.value(Dlim[1]),Dates.value(f))
    end
        xticks := (r,d)
=#
    #Oil rate Series
    @series begin
        seriestype := :path
        linewidth --> 3
        seriescolor --> :green

        label := "Oil Rate [bbl/d]"
        x, y
        end

    if Liquid==true
    #Liquid rate Series
    @series begin
        seriestype := :path
        linewidth --> 1
        seriescolor --> RGB(0.5, 0.5, 0.5)
        subplot := 1
        label --> "Liquid Rate [bbl/d]"
        x, z
        end
    end

    if Dec==true
        @series begin
            seriestype := :path
            linewidth --> 3
            seriescolor --> :red
            lab := "Di = $(round(S.Di*365/12,digits=2)) Monthly // $(round(S.Di*365,digits=2)) Annually"
            DecTime, DecLine
        end
    end

        if Forecast==true

        @series begin
            seriestype := :path
            linestyle --> :dash
            linewidth --> 1
            seriescolor --> :red
                lab := "$(Dates.format(ForecastTime[1],dft))--> $(Dates.format(ForecastTime[end],dft))// Np= $(round(Np,digits=1)) Mbbl))"
                ForecastTime, ForecastLine
        end

        end
end


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ f88b15de-cec6-4bc8-9b68-2a407b5aeded
begin
    using Pkg: Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(; name="Manifolds", version="0.7"),
        Pkg.PackageSpec(; name="Manopt", version="0.3"),
        Pkg.PackageSpec(; name="FiniteDiff", version="2.8"),
        Pkg.PackageSpec(; name="ReverseDiff", version="1.9"),
        # Pkg.PackageSpec(; name="PlutoUI"),
    ])
end

# ╔═╡ 856f336c-e232-4f1f-b1ac-759b4558acd1
using Manifolds, Manopt, Random, LinearAlgebra

# ╔═╡ b0769dfa-28cf-440e-9ba2-1ef488f171a9
using FiniteDiff, ReverseDiff

# ╔═╡ 0213d26a-18ac-11ec-03fd-ada5992bcea8
md"""
# Using (Euclidean) AD in Manopt.jl
"""

# ╔═╡ f3bc91ee-5871-4cba-ac89-190deb71ad0f
md"""
Since [Manifolds.jl](https://juliamanifolds.github.io/Manifolds.jl/latest/) 0.7 the support of automatic differentiation support has been extended.

This tutorial explains how to use Euclidean tools to derive a gradient for a real-valued function ``F\colon \mathcal M → ℝ``. We will consider two methods: an intrinsic variant and a variant employing the embedding. These gradients can then be used within any gradient based optimisation algorithm in [Manopt.jl](https://manoptjl.org).

While by default we use [FiniteDifferences.jl](https://juliadiff.org/FiniteDifferences.jl/latest/), you can also use [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl), [ForwardDiff.jl](https://juliadiff.org/ForwardDiff.jl/stable/), [ReverseDiff.jl](https://juliadiff.org/ReverseDiff.jl/), or  [Zygote.jl](https://fluxml.ai/Zygote.jl/).
"""

# ╔═╡ d9be6c2f-65fd-4685-9005-da22bf985e28
md"""
In this Notebook we will take a look at a few possibilities to approximate or derive the gradient of a function ``f:\mathcal M \to ℝ`` on a Riemannian manifold, without computing it yourself. There is mainly two different philosophies:

1. Working _instrinsically_, i.e. stay on the manifold and in the tangent spaces. Here, we will consider approximating the gradient by forward differences.

2. Working in an embedding – there we can use all tools from functions on Euclidean spaces – finite differences or automatic differenciation – and then compute the corresponding Riemannian gradient from there.

Let's first load all packages we need.
"""

# ╔═╡ 18d7459f-eed6-489b-a096-ac77ccd781af
md"""
## 1. (Intrinsic) Forward Differences

A first idea is to generalise (multivariate) finite differences to Riemannian manifolds. Let ``X_1,\ldots,X_d ∈ T_p\mathcal M`` denote an orthonormal basis of the tangent space ``T_p\mathcal M`` at the point ``p∈\mathcal M`` on the Riemannian manifold.

We can generalise the notion of a directional derivative, i.e. for the “direction” ``Y∈T_p\mathcal M`` let ``c\colon [-ε,ε]``, ``ε>0``, be a curve with ``c(0) = p``, ``\dot c(0) = Y`` and we obtain

```math
	Df(p)[Y] = \frac{\mathrm{d}}{\mathrm{d}t} f(c(t)) = \lim_{h \to 0} \frac{1}{h}(f(\exp_p(hY))-f(p))
```

We can approximate ``Df(p)[X]`` by a finite difference scheme for an ``h>0`` as

```math
DF(p)[Y] ≈ G_h(Y) := \frac{1}{h}(f(\exp_p(hY))-f(p))
```

Furthermore the gradient ``\operatorname{grad}f`` is the Riesz representer of the differential, ie.

```math
	Df(p)[Y] = g_p(\operatorname{grad}f(p), Y),\qquad \text{ for all } Y ∈ T_p\mathcal M
```

and since it is a tangent vector, we can write it in terms of a basis as

```math
	\operatorname{grad}f(p) = \sum_{i=1}^{d} g_p(\operatorname{grad}f(p),X_i)X_i
	= \sum_{i=1}^{d} Df(p)[X_i]X_i
```

and perform the approximation from above to obtain
```math
	\operatorname{grad}f(p) ≈ \sum_{i=1}^{d} G_h(X_i)X_i
```
for some suitable step size ``h``.This comes at the cost of ``d+1`` function evaluations and ``d`` exponential maps.
"""

# ╔═╡ a3df142e-94df-48d2-be08-d1f1f3854c76
md"""
This is the first variant we can use. An advantage is, that it is _intrinsic_ in the sense that it does not require any embedding of the manifold.
"""

# ╔═╡ 9a030ac6-1f44-4fa6-8bc9-1c0278e97fe2
md""" ### An Example: The Rayleigh Quotient

The Rayleigh quotient is concerned with finding Eigenvalues (and Eigenvectors) of a symmetric matrix $A\in ℝ^{(n+1)×(n+1)}$. The optimisation problem reads

```math
F\colon ℝ^{n+1} \to ℝ,\quad F(\mathbf x) = \frac{\mathbf x^\mathrm{T}A\mathbf x}{\mathbf x^\mathrm{T}\mathbf x}
```

Minimizing this function yields the smallest eigenvalue ``\lambda_1`` as a value and the corresponding minimizer ``\mathbf x^*`` is a corresponding eigenvector.

Since the length of an eigenvector is irrelevant, there is an ambiguity in the cost function. It can be better phrased on the sphere ``𝕊^n`` of unit vectors in ``\mathbb R^{n+1}``, i.e.

```math
\operatorname*{arg\,min}_{p \in 𝕊^n} f(p) = \operatorname*{arg\,min}_{p \in 𝕊^n} p^\mathrm{T}Ap
```

We can compute the Riemannian gradient exactly as

```math
\operatorname{grad} f(p) = 2(Ap - pp^\mathrm{T}Ap)
```

so we can compare it to the approximation by finite differences.
"""

# ╔═╡ 19747159-d383-4547-9315-0ed2494904a6
begin
    Random.seed!(42)
    n = 200
    A = randn(n + 1, n + 1)
    A = Symmetric(A)
    M = Sphere(n)
    nothing
end

# ╔═╡ 41c204dd-6e4e-4a70-8f06-209a469e0680
f1(p) = p' * A'p

# ╔═╡ 2e33de5e-ffaa-422a-91d9-61f588ed1211
gradf1(p) = 2 * (A * p - p * p' * A * p)

# ╔═╡ bbd9a010-1981-45b3-bf7d-c04bcd2c2128
md"""Manifolds provides a finite difference scheme in Tangent spaces, that you can introduce to use an existing framework (if the wrapper is implemented) form Euclidean space. Here we use `FiniteDiff.jl`."""

# ╔═╡ 08456b40-74ec-4319-93e7-130b5cf70ac3
r_backend = Manifolds.TangentDiffBackend(Manifolds.FiniteDiffBackend())

# ╔═╡ 12327b62-7e79-4381-b6a7-f85b08a8251b
gradf1_FD(p) = Manifolds.gradient(M, f1, p, r_backend)

# ╔═╡ 07f9a630-e53d-45ea-b109-3d4de190723d
begin
    p = zeros(n + 1)
    p[1] = 1.0
    X1 = gradf1(p)
    X2 = gradf1_FD(p)
    norm(M, p, X1 - X2)
end

# ╔═╡ 8e5f677d-dafa-49b9-b678-3f129be31dcf
md"We obtain quite a good approximation of the gradient."

# ╔═╡ 77769eab-54dd-41dc-8125-0382e5ef0bf1
md"""
## 2. Conversion of an Euclidean Gradient in the Embedding to a Riemannian Gradient of an (not necessarily isometrically) embedded Manifold

Let ``\tilde f\colon\mathbb R^m \to \mathbb R`` be a function in the embedding of an ``n``-dimensional manifold ``\mathcal M \subset \mathbb R^m`` and ``f\colon \mathcal M \to \mathbb R`` denote the restriction of ``\tilde f`` to the manifold ``\mathcal M``.

Since we can use the push forward of the embedding to also embed the tangent space ``T_p\mathcal M``, ``p\in \mathcal M``, we can similarly obtain the differential ``Df(p)\colon T_p\mathcal M \to \mathbb R`` by restricting the differential ``D\tilde f(p)`` to the tangent space.

If both ``T_p\mathcal M`` and ``T_p\mathcal R^m`` have the same inner product, or in other words the manifold is isometrically embedded in ``R^m`` (like for example the sphere ``\mathbb S^n\subset\mathbb R^{m+1}`` then this restriction of the differential directly translates to a projection of the gradient, i.e.

```math
\operatorname{grad}f(p) = \operatorname{Proj}_{T_p\mathcal M}(\operatorname{grad} \tilde f(p))
```

More generally we might have to take a change of the metric into account, i.e.

```math
\langle  \operatorname{Proj}_{T_p\mathcal M}(\operatorname{grad} \tilde f(p)), X \rangle
= Df(p)[X] = g_p(\operatorname{grad}f(p), X)
```

or in words: we have to change the Riesz representer of the (restricted/projected) differential of ``f`` (``\tilde f``) to the one with respect to the Riemannian metric. This is done using [`change_representer`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/metric.html#Manifolds.change_representer-Tuple{AbstractManifold,%20AbstractMetric,%20Any,%20Any}).
"""

# ╔═╡ 57cda07f-e432-46af-b771-5e5a3067feac
md"""
### A continued Example
We continue with the Rayleigh Quotient from before, now just starting with the defintion of the Euclidean case in the embedding, the function ``F``.
"""

# ╔═╡ c3f3aeba-2849-4715-94e2-0c44613a2ce9
F(x) = x' * A * x / (x' * x);

# ╔═╡ 786fce04-53ef-448d-9657-31208b35fb7e
md"The cost function is the same by restriction"

# ╔═╡ c1341fef-adec-4574-a642-a1a8a9c1fee5
f2(M, p) = F(p);

# ╔═╡ 0818a62f-1bef-44f7-a33f-1ab0054e853c
md"The gradient is now computed combining our gradient scheme with ReverseDiff."

# ╔═╡ 89cd6b4b-f9ef-47ac-afd3-cf9aacf43256
function grad_f2_AD(M, p)
    return Manifolds.gradient(
        M, F, p, RiemannianProjectionBackend(Manifolds.ReverseDiffBackend())
    )
end

# ╔═╡ 7c5a8a17-6f63-4587-a94a-6936bdd3cec6
X3 = grad_f2_AD(M, p)

# ╔═╡ b3e7f57f-d87a-47c5-b8ad-48b6d205fa73
norm(M, p, X1 - X3)

# ╔═╡ 893db402-283f-4e3e-8bf7-c6f22e485efb
md"""
### An Example for a nonisometrically embedded Manifold

on the manifold ``\mathcal P(3)`` of symmetric positive definite matrices.
"""

# ╔═╡ 8494a0d6-dbf2-4eb0-a555-f00e446fbe38
md"""
The following function computes (half) the distance squared (with respect to the linear affine metric) on the manifold ``\mathcal P(3)`` to the identity, i.e. $I_3$. denoting the unit matrix we consider the function

```math
	G(q) = \frac{1}{2}d^2_{\mathcal P(3)}(q,I_3) = \lVert \operatorname{Log}(q) \rVert_F^2,
```
where $\operatorname{Log}$ denotes the matrix logarithm and ``\lVert \cdot \rVert_F`` is the Frobenius norm.
This can be computed for symmetric positive definite matrices by summing the squares of the ``\log``arithms of the eigenvalues of ``q`` and divide by two:
"""

# ╔═╡ c93eb2da-89df-4751-b086-62be604d41e6
G(q) = sum(log.(eigvals(Symmetric(q))) .^ 2) / 2

# ╔═╡ e2bf6f55-7235-4d75-8bee-a325434e32ad
md"""
We can also interpret this as a function on the space of matrices and apply the Euclidean finite differences machinery; in this way we can easily derive the Euclidean gradient. But when computing the Riemannian gradient, we have to change the representer (see again [`change_representer`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/metric.html#Manifolds.change_representer-Tuple{AbstractManifold,%20AbstractMetric,%20Any,%20Any})) after projecting onto the tangent space ``T_p\mathcal P(n)`` at ``p``.

Let's first define a point and the manifold ``N=\mathcal P(3)``.
"""

# ╔═╡ 153378ca-703d-4a84-bc63-22347399a160
rotM(α) = [1.0 0.0 0.0; 0.0 cos(α) sin(α); 0.0 -sin(α) cos(α)]

# ╔═╡ 699f0177-2c5b-434b-9eca-b6fc573e497f
q = rotM(π / 6) * [1.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 3.0] * transpose(rotM(π / 6))

# ╔═╡ 9b52d05c-4cba-4922-85a4-1a2a3c74823e
N = SymmetricPositiveDefinite(3)

# ╔═╡ cb3fe7aa-1262-48f2-9ebc-3e959c72a33e
is_point(N, q)

# ╔═╡ 13712c64-48fd-4f2a-9ee4-1949e51d316f
md"""We could first just compute the gradient using `FiniteDiff.jl`, but this yields the Euclidean gradient:"""

# ╔═╡ 64beb3dd-9507-4792-be02-ae1405704690
FiniteDiff.finite_difference_gradient(G, q)

# ╔═╡ 2be4f9e8-0331-44ac-839f-7bb71d9edef9
md"""Instead, we use the [`RiemannianProjectedBackend`](https://juliamanifolds.github.io/Manifolds.jl/latest/features/differentiation.html#Manifolds.RiemannianProjectionBackend) of `Manifolds.jl`, which in this case internally uses `FiniteDiff.jl` to compute a Euclidean gradient but then uses the conversion explained above to derive the Riemannian gradient.

We define this here again as a function `grad_G_FD` that could be used in the `Manopt.jl` framework within a gradient based optimisation.
"""

# ╔═╡ 6f1d748f-27ce-496b-8561-f16972da50cc
function grad_G_FD(N, q)
    return Manifolds.gradient(
        N, G, q, RiemannianProjectionBackend(Manifolds.FiniteDiffBackend())
    )
end

# ╔═╡ 7dd656ea-08de-4172-8a92-87ad2228ce69
G1 = grad_G_FD(N, q)

# ╔═╡ 219573d2-283f-456c-a5c3-fadd734fc157
md"""
Now, we can agaon compare this to the (known) solution of the gradient, namely the gradient of (a half) the distance suqared, i.e. ``G(q) = \frac{1}{2}d^2_{\mathcal P(3)}(q,I_3)`` is given by ``\operatorname{grad} G(q) = -\operatorname{log}_q I_3``, where ``\operatorname{log}`` is the [logarithmic map](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/symmetricpositivedefinite.html#Base.log-Tuple{SymmetricPositiveDefinite,%20Vararg{Any,%20N}%20where%20N}) on the manifold.
"""

# ╔═╡ e28a2752-877c-4ab4-a253-8d26fa9a73c2
G2 = -log(N, q, Matrix{Float64}(I, 3, 3))

# ╔═╡ 25c65878-1be6-4fec-b65e-9c1741320a41
md"""Both terms agree up to ``1.2×10^{-10}``:"""

# ╔═╡ 9a66d4f3-508d-4285-9a93-df1323575202
norm(G1 - G2)

# ╔═╡ c07fb3d0-d12f-44d7-bcab-7a0d39e6af8d
isapprox(M, q, G1, G2; atol=2 * 1e-10)

# ╔═╡ f47c70b6-ca05-498f-9e10-58c3839ca427
md""" In this case we can not use `ReverseDiff.jl`, since it can not handle the `eigvals!` function that is called internally."""

# ╔═╡ 32d8d025-3993-4d31-9eea-3463e0af1c12
md"""
## Summary

This tutorial illustrates how to use tools from Euclidean spaces, finite differences or automatic differentiation, to compute gradients on Riemannian manifolds. The scheme allows to use _any_ differentiation framework within the embedding to derive a Riemannian gradient.
"""

# ╔═╡ Cell order:
# ╟─0213d26a-18ac-11ec-03fd-ada5992bcea8
# ╟─f3bc91ee-5871-4cba-ac89-190deb71ad0f
# ╟─d9be6c2f-65fd-4685-9005-da22bf985e28
# ╟─f88b15de-cec6-4bc8-9b68-2a407b5aeded
# ╠═856f336c-e232-4f1f-b1ac-759b4558acd1
# ╠═b0769dfa-28cf-440e-9ba2-1ef488f171a9
# ╟─18d7459f-eed6-489b-a096-ac77ccd781af
# ╟─a3df142e-94df-48d2-be08-d1f1f3854c76
# ╟─9a030ac6-1f44-4fa6-8bc9-1c0278e97fe2
# ╠═19747159-d383-4547-9315-0ed2494904a6
# ╠═41c204dd-6e4e-4a70-8f06-209a469e0680
# ╠═2e33de5e-ffaa-422a-91d9-61f588ed1211
# ╟─bbd9a010-1981-45b3-bf7d-c04bcd2c2128
# ╠═08456b40-74ec-4319-93e7-130b5cf70ac3
# ╠═12327b62-7e79-4381-b6a7-f85b08a8251b
# ╠═07f9a630-e53d-45ea-b109-3d4de190723d
# ╟─8e5f677d-dafa-49b9-b678-3f129be31dcf
# ╟─77769eab-54dd-41dc-8125-0382e5ef0bf1
# ╟─57cda07f-e432-46af-b771-5e5a3067feac
# ╠═c3f3aeba-2849-4715-94e2-0c44613a2ce9
# ╟─786fce04-53ef-448d-9657-31208b35fb7e
# ╠═c1341fef-adec-4574-a642-a1a8a9c1fee5
# ╟─0818a62f-1bef-44f7-a33f-1ab0054e853c
# ╠═89cd6b4b-f9ef-47ac-afd3-cf9aacf43256
# ╠═7c5a8a17-6f63-4587-a94a-6936bdd3cec6
# ╠═b3e7f57f-d87a-47c5-b8ad-48b6d205fa73
# ╟─893db402-283f-4e3e-8bf7-c6f22e485efb
# ╟─8494a0d6-dbf2-4eb0-a555-f00e446fbe38
# ╠═c93eb2da-89df-4751-b086-62be604d41e6
# ╟─e2bf6f55-7235-4d75-8bee-a325434e32ad
# ╠═153378ca-703d-4a84-bc63-22347399a160
# ╠═699f0177-2c5b-434b-9eca-b6fc573e497f
# ╠═9b52d05c-4cba-4922-85a4-1a2a3c74823e
# ╠═cb3fe7aa-1262-48f2-9ebc-3e959c72a33e
# ╟─13712c64-48fd-4f2a-9ee4-1949e51d316f
# ╠═64beb3dd-9507-4792-be02-ae1405704690
# ╟─2be4f9e8-0331-44ac-839f-7bb71d9edef9
# ╠═6f1d748f-27ce-496b-8561-f16972da50cc
# ╠═7dd656ea-08de-4172-8a92-87ad2228ce69
# ╟─219573d2-283f-456c-a5c3-fadd734fc157
# ╠═e28a2752-877c-4ab4-a253-8d26fa9a73c2
# ╟─25c65878-1be6-4fec-b65e-9c1741320a41
# ╠═9a66d4f3-508d-4285-9a93-df1323575202
# ╠═c07fb3d0-d12f-44d7-bcab-7a0d39e6af8d
# ╟─f47c70b6-ca05-498f-9e10-58c3839ca427
# ╟─32d8d025-3993-4d31-9eea-3463e0af1c12


"""
    compute_calcium_scores(image, spacing, mask, min_vol=nothing, max_vol=nothing)

Calculates Agatston score and volume given an input mask

Ported from
(https://github.com/sgmvanvelzen/calcium-scoring/blob/f63629f322ec2a22b3e7cb4a7eb6e494eba09fc5/src/calciumscoring/scores.py#L9)
"""
function compute_calcium_scores(image, spacing, mask, min_vol=nothing, max_vol=nothing)
	binary_mask = mask
	voxel_volume = spacing[1] * spacing[2] * spacing[3]

	agatston_score = 0
	calcium_volume = 0

	# Find individual lesions (in 3D) so that we can discard too small or too large lesions
	lesion_map = Images.label_components(binary_mask)
	n_lesions = length(unique(lesion_map))

	for lesion = 1:(n_lesions - 1)
		lesion_mask = map(x -> x == lesion, lesion_map)

		# Ignore too small or too large lesions
		lesion_volume = count(lesion_mask) * voxel_volume
		if ((min_vol !== nothing) && (lesion_volume < min_vol))
			continue
		end
		if ((max_vol !== nothing) && (lesion_volume > max_vol))
			continue
		end

		calcium_volume += lesion_volume

		# Calculate Agatston score for this lesion
		slices = sort(unique(lesion_mask)) .+ 1
		for z = 1:slices[end]
			fragment_mask = lesion_mask[z,:,:]
			n_pixels = count(fragment_mask)
			maximum_intensity = maximum(image[z,:,:][fragment_mask])
            if maximum_intensity < 200
                coefficient = 1
			elseif maximum_intensity < 300
				coefficient = 2
			elseif maximum_intensity < 400
				coefficient = 3
			else
				coefficient = 4
			end
			agatston_score += coefficient * n_pixels
		end
	end
	agatston_score *= spacing[1] / 3.0 * spacing[2] * spacing[3]
	return agatston_score, calcium_volume
end