#=
shepplogan.jl
=#

export shepp_logan, ellipse_parameters
export SheppLogan
export SheppLoganEmis
export SheppLoganBrainWeb
export SheppLoganToft
export SouthPark

"""
    EllipsePhantomVersion
Parent type for different versions of ellipse phantoms:
* `SheppLogan` original CT version from Shepp&Logan paper
* `SheppLoganEmis` higher contrast version suitable for emission tomography
* `SheppLoganBrainWeb` integer index version based on brainweb
* `SheppLoganToft` higher contrast version from Toft, 1996
* `SouthPark` for fun
"""
abstract type EllipsePhantomVersion end

"""
    SheppLogan
Original version from :
Larry A Shepp, Benjamin F Logan,
"The Fourier reconstruction of a head section,"
IEEE Transactions on Nuclear Science, 21(3):21-42, June 1974.
[doi](http://doi.org/10.1109/TNS.1974.6499235)

Also in Kak and Slaney 1988 text, p. 255.
[doi](http://doi.org/10.1137/1.9780898719277)
"""
struct SheppLogan <: EllipsePhantomVersion end

"""
    SheppLoganToft
Toft, Peter Aundal & Sørensen, John Aasted
"The Radon transform-theory and implementation,"
Technical University of Denmark (DTU), 1996. Page 201.
https://files.openpdfs.org/1ra51GP6gJO.pdf
"""
struct SheppLoganToft <: EllipsePhantomVersion end

struct SheppLoganEmis <: EllipsePhantomVersion end
struct SheppLoganBrainWeb <: EllipsePhantomVersion end
struct SouthPark <: EllipsePhantomVersion end

"""
    values = shepp_logan_values(::EllipsePhantomVersion)
Return 10 Shepp-Logan ellipse amplitudes for various versions.
"""
shepp_logan_values(::SheppLogan) =
    [2.0, -0.98, -0.02, -0.02, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01]
shepp_logan_values(::SheppLoganEmis) =
    [1, 1, -2, 2, 3, 4, 5, 6, 1, 1]
shepp_logan_values(::SheppLoganBrainWeb) =
    [1, 0, 2, 3, 4, 5, 6, 7, 8, 9] # brainweb uses index 1-10
shepp_logan_values(::SheppLoganToft) =
    [1.0, -0.8, -0.2, -0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]

"""
    phantom = Ellipse(n × 6 AbstractMatrix)
Return vector of Ellipse objects, one for each row of input matrix.
"""
function Ellipse(params::AbstractMatrix{<:RealU})
    size(params,2) == 6 || throw("ellipses need 6 parameters")
    return [Ellipse(params[n,:]) for n in 1:size(params,1)]
end

function shepp_logan(case::EllipsePhantomVersion; kwargs...)
    return Ellipse(ellipse_parameters(case; kwargs...))
end

"""
    image = shepp_logan(M, [N,], case, options...)
Convenience method for generating `M×N` samples of Shepp-Logan phantoms.

In
* `M::Int` : horizontal size
* `N::Int` : vertical size, defaults to `M`
* `case::EllipsePhantomVersion = SheppLogan()`

# Options
* `oversample::Int = 3` (usually)
* `yflip::Bool = true` (reverse `y` samples for convenience.)
* `kwargs...` remaining options passed to `ellipse_parameters` for parameters.

# Out
* `image` : `M × N` matrix

The default here is 3× over-sampling along both axes (9 samples per pixel),
except for the `SheppLoganBrainWeb` phantom that consists of integer indices.
"""
function shepp_logan(
    M::Int, N::Int,
    case::EllipsePhantomVersion = SheppLogan() ;
    oversample::Int = case == SheppLoganBrainWeb() ? 1 : 3,
    yflip::Bool = true,
    kwargs...
)
    ob = shepp_logan(case ; kwargs...)
    x = LinRange(-0.5, 0.5, M)
    y = LinRange(-0.5, 0.5, N)
    if yflip
        y = reverse(y)
    end
    return oversample > 1 ?
        phantom(x, y, ob, oversample) :
        phantom(x, y, ob) # type unstable because of over-sampling
end

shepp_logan(M::Int, case::EllipsePhantomVersion = SheppLogan(); kwargs...) =
    shepp_logan(M, M, case ; kwargs...)

"""
    params = ellipse_parameters(case::Symbol; fovs::NTuple{2}, u::Tuple, disjoint::Bool)

By default the first four columns are unitless "fractions of field of view",
so columns 1,3 are scaled by `xfov` and columns 2,4 are scaled by `yfov`,
where `(xfov, yfov) = fovs`.
The optional 3-tuple `u` specifies scaling and/or units:
* columns 1-4 (center, radii) are scaled by `u[1]` (e.g., mm),
* column 5 (angle) is scaled by `u[2]` (e.g., `1` or `°`),
* column 6 (value) is scaled by `u[3]` (e.g., `1/cm`) for an attenuation map.
If `disjoint==true` then the middle ellipse positions are adjusted to avoid overlap.
"""
function ellipse_parameters(
    case::EllipsePhantomVersion = SheppLogan() ;
    fovs::NTuple{2,RealU} = (1,1),
    u::NTuple{3,Any} = (1,1,1), # unit scaling
    disjoint::Bool = false,
)
    (xfov, yfov) = fovs

    # The "Number" here is to enable units
    params = Number[ # original CT version
    0       0       0.92    0.69    90    2     # skull
    0       -0.0184 0.874   0.6624  90    -0.98 # brain
    0.22    0       0.31    0.11    72    -0.02 # right big
    -0.22   0       0.41    0.16   108    -0.02 # left big
    0       0.35    0.25    0.21    90    0.01 # top
    0       0.1     0.046   0.046    0    0.01 # middle high
    0       -0.1    0.046   0.046    0    0.01 # middle low
    -0.08   -0.605  0.046   0.023    0    0.01 # bottom left
    0       -0.605  0.023   0.023    0    0.01 # bottom center
    0.06    -0.605  0.046   0.023   90    0.01 # bottom right
    ]

    if disjoint # tweak so that middle ellipses are disjoint
        params[4,1] += -0.04
        params[5,2] += 0.07
    end

    params[:,[1,3]] .*= xfov/2
    params[:,[2,4]] .*= yfov/2
    params[:,5] .*= π/180
    params[:,6] = shepp_logan_values(case)

    return Number[params[:,1:4] * u[1] params[:,5] * u[2] params[:,6] * u[3]]
end

"""
    params = ellipse_parameters(SouthPark() ; fovs::NTuple{2,Number} = (100,100))
Ellipse parameters for "South Park" phantom.
"""
function ellipse_parameters(::SouthPark ; fovs::NTuple{2,RealU} = (100,100))
    params = Float64[
        0 0 85 115 0 100
        0 -60 30 20 0 -80 # mouth
        30 20 25 35 30 20 # eyes
        -30 20 25 35 -30 20
        35 25 7 7 0 -100 # pupils
        -15 25 7 7 0 -100
        0 75 60 15 0 -50  # hat
    ]
    params[:,[1,3]] .*= fovs[1] / 256
    params[:,[2,4]] .*= fovs[2] / 256
    params[:,5] .*= π/180
    return params
end


# This file was generated, do not modify it. # hide
hpn  = xx.Horsepower
Xnew = DataFrame(hp1=hpn, hp2=hpn.^2, hp3=hpn.^3)

yy1 = predict(lr1, Xnew)
yy2 = predict(lr2, Xnew)
yy3 = predict(lr3, Xnew)

figure(figsize=(8,6))
plot(X.Horsepower, y, ls="none", marker="o")
plot(xx.Horsepower, yy1, lw=3, label="Order 1")
plot(xx.Horsepower, yy2, lw=3, label="Order 2")
plot(xx.Horsepower, yy3, lw=3, label="Order 3")

legend(fontsize=14)

xlabel("Horsepower", fontsize=14)
xticks(50:50:250, fontsize=12)
yticks(10:10:50, fontsize=12)
ylabel("MPG", fontsize=14)

savefig(joinpath(@OUTPUT, "ISL-lab-5-g3.svg")) # hide


using DelimitedFiles

function ReadData(fname::AbstractString)
    dataStr = readdlm(fname, String)

    data = zeros(Int, length(dataStr), length(dataStr[1]))
    for i = 1:length(dataStr)
        data[i,:] = map(x -> parse(Int, x), split(dataStr[i],""))
    end

    return data
end

function FindMinima(data::Matrix)
    vOffsetU = circshift(data, (-1,0)) # data shifted up one
    hOffsetL = circshift(data, (0,-1)) # data shifted to the left
    vOffsetD = circshift(data, (1,0))  # data shifted up one
    hOffsetR = circshift(data, (0,1))  # data shifted to the left

    # set the edges to be 9 (the maximum possible value)
    vOffsetU[size(data)[1],:] = fill(9, size(data)[2])
    hOffsetL[:,size(data)[2]] = fill(9, size(data)[1])
    vOffsetD[1,:] = fill(9, size(data)[2])
    hOffsetR[:,1] = fill(9, size(data)[1])

    return intersect([findall(x -> x<0, data-offset) for offset in [vOffsetU, hOffsetL, vOffsetD, hOffsetR]]...)
end

function Part1(data::Matrix)
    minimaPositions = FindMinima(data)

    minima = map(x -> getindex(data, x), minimaPositions)

    return sum(minima)+length(minima)
end

function Part2(data::Matrix)
    vOffsetU = circshift(data, (-1,0)) # data shifted up one
    hOffsetL = circshift(data, (0,-1)) # data shifted to the left
    vOffsetD = circshift(data, (1,0))  # data shifted up one
    hOffsetR = circshift(data, (0,1))  # data shifted to the left

    # set the edges to be 9 (the maximum possible value)
    vOffsetU[size(data)[1],:] = fill(9, size(data)[2])
    hOffsetL[:,size(data)[2]] = fill(9, size(data)[1])
    vOffsetD[1,:] = fill(9, size(data)[2])
    hOffsetR[:,1] = fill(9, size(data)[1])

    minimaPositions = intersect([findall(x -> x<0,  data-offset) for offset in [vOffsetU, hOffsetL, vOffsetD, hOffsetR]]...)
    maximaPositions = intersect([findall(x -> x>=0, data-offset) for offset in [vOffsetU, hOffsetL, vOffsetD, hOffsetR]]...) # this will catch saddle points

    basinSizes = []

    for x in minimaPositions
        # start building basin
        currentBasin = [x]

        basinSize = 0

        Q = [x]
        explored = [x]
        while ! isempty(Q)
            v = pop!(Q)
            if getindex(data, v) == 9
                continue
            else
                basinSize += 1
            end
            i, j = v[1], v[2]
            for k=-1:2:1
                if CartesianIndex(i+k, j) ∉ explored && 1 <= i+k <= size(data)[1]
                    push!(explored, CartesianIndex(i+k, j))
                    push!(Q, CartesianIndex(i+k, j))
                end

                if CartesianIndex(i, j+k) ∉ explored && 1 <= j+k <= size(data)[2]
                    push!(explored, CartesianIndex(i, j+k))
                    push!(Q, CartesianIndex(i, j+k))
                end
            end
        end

        push!(basinSizes, basinSize)
    end

    sort!(basinSizes, rev=true)

    return basinSizes[1]*basinSizes[2]*basinSizes[3]

end

#data = ReadData("../data/test9.txt")
data = ReadData("../data/day9.txt")

println(Part2(data))


module ActuaryUtilities

using Dates
using ForwardDiff
using QuadGK
using Yields
import Roots
import StatsBase

include("financial_math.jl")
include("risk_measures.jl")
include("derivatives.jl")

# https://timholy.github.io/SnoopCompile.jl/stable/snoopi_deep_parcel/
if Base.VERSION >= v"1.4.2"
    include("precompile.jl")
    _precompile_()
end

"""
    Years_Between(d1::Date, d2::Date)

Compute the number of integer years between two dates, with the
first date typically before the second. Will return negative number if
first date is after the second. Use third argument to indicate if calendar
anniversary should count as a full year.

# Examples
```jldoctest
julia> d1 = Date(2018,09,30);

julia> d2 = Date(2019,09,30);

julia> d3 = Date(2019,10,01);

julia> years_between(d1,d3)
1
julia> years_between(d1,d2,false) # same month/day but `false` overlap
0
julia> years_between(d1,d2) # same month/day but `true` overlap
1
julia> years_between(d1,d2) # using default `true` overlap
1
```
"""
function years_between(d1::Date, d2::Date, overlap = true)
    iy, im, id = Dates.year(d1), Dates.month(d1), Dates.day(d1)
    vy, vm, vd = Dates.year(d2), Dates.month(d2), Dates.day(d2)
    dur = vy - iy
    if vm == im
        if overlap
            if vd >= id
                dur += 1
            end
        else
            if vd > id
                dur += 1
            end
        end
    elseif vm > im
        dur += 1
    end

    return dur - 1
end

"""
    duration(d1::Date, d2::Date)

Compute the duration given two dates, which is the number of years
since the first date. The interval `[0,1)` is defined as having
duration `1`. Can return negative durations if second argument is before the first.

```jldoctest
julia> issue_date  = Date(2018,9,30);

julia> duration(issue_date , Date(2019,9,30) )
2
julia> duration(issue_date , issue_date)
1
julia> duration(issue_date , Date(2018,10,1) )
1
julia> duration(issue_date , Date(2019,10,1) )
2
julia> duration(issue_date , Date(2018,6,30) )
0
julia> duration(Date(2018,9,30),Date(2017,6,30))
-1
```

"""
function duration(issue_date::Date, proj_date::Date)
    return years_between(issue_date, proj_date, true) + 1
end

"""
    accum_offset(x; op=*, init=1.0)

A shortcut for the common operation wherein a vector is scanned with an operation, but has an initial value and the resulting array is offset from the traditional accumulate.

This is a common pattern when calculating things like survivorship given a mortality vector and you want the first value of the resulting vector to be `1.0`, and the second value to be `1.0 * x[1]`, etc.

Two keyword arguments:

- `op` is the binary (two argument) operator you want to use, such as `*` or `+`
- `init` is the initial value in the returned array

# Examples

```julia=repl
julia> accum_offset([0.9, 0.8, 0.7])
3-element Array{Float64,1}:
 1.0
 0.9
 0.7200000000000001

julia> accum_offset(1:5) # the product of elements 1:n, with the default `1` as the first value
5-element Array{Int64,1}:
  1
  1
  2
  6
 24

julia> accum_offset(1:5,op=+)
5-element Array{Int64,1}:
  1
  2
  4
  7
 11

```

"""
function accum_offset(x; op=*,init=1.0)
    xnew = similar(x)
    xnew[1] = init
    for i in 2:length(x)
      xnew[i] = op(xnew[i-1],x[i-1])
    end
    return xnew
end

export years_between, duration,
    irr, internal_rate_of_return,
    pv, present_value, price, present_values,
    breakeven,moic,
    accum_offset,
    Macaulay,Modified,DV01,KeyRatePar,KeyRateZero,KeyRate,duration, convexity,
    VaR,ValueAtRisk,CTE,ConditionalTailExpectation,ExpectedShortfall,
    eurocall, europut

end # module


##### NewtonsMethod

export NewtonsMethod

"""
    NewtonsMethod(
            f!::F!,
            j!::J!,
            x_init::A,
        ) where {F! <: Function, J! <: Function, A <: AbstractArray}

A non-linear system of equations type.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct NewtonsMethod{FT, F!, J!, A, JA} <: AbstractNonlinearSolverMethod{FT}
    "Function to find the root of"
    f!::F!
    "Jacobian of `f!`"
    j!::J!
    "Initial guess"
    x_init::A
    "Storage"
    x1::A
    "Storage"
    J::JA
    "Storage"
    J⁻¹::JA
    "Storage"
    F::A
    function NewtonsMethod(
        f!::F!,
        j!::J!,
        x_init::A,
    ) where {F! <: Function, J! <: Function, A <: AbstractArray}
        x1 = similar(x_init)
        J = similar(x_init, (length(x_init), length(x_init)))
        J⁻¹ = similar(J)
        F = similar(x_init)
        JA = typeof(J)
        FT = eltype(x_init)
        return new{FT, F!, J!, A, JA}(f!, j!, x_init, x1, J, J⁻¹, F)
    end
end

method_args(m::NewtonsMethod) = (m.x_init, m.x1, m.f!, m.F, m.j!, m.J, m.J⁻¹)
function solve!(
    ::NewtonsMethod,
    x0::AT,
    x1::AT,
    f!::F!,
    F::FA,
    j!::J!,
    J::JA,
    J⁻¹::J⁻¹A,
    soltype::SolutionType,
    tol::AbstractTolerance{FT},
    maxiters::Int,
) where {FA, J⁻¹A, JA, F! <: Function, J! <: Function, AT, FT}

    x_history = init_history(soltype, AT)
    F_history = init_history(soltype, AT)
    if soltype isa VerboseSolution
        f!(F, x0)
        j!(J, x0)
        push_history!(x_history, x0, soltype)
        push_history!(F_history, F, soltype)
    end
    for i in 1:maxiters
        f!(F, x0)
        j!(J, x0)
        x1 .= x0 .- J \ F
        push_history!(x_history, x1, soltype)
        push_history!(F_history, F, soltype)
        if tol(x0, x1, F)
            return SolutionResults(
                soltype,
                x1,
                true,
                F,
                i,
                x_history,
                F_history,
            )
        end
        x0 = x1
    end
    return SolutionResults(
        soltype,
        x0,
        false,
        F,
        maxiters,
        x_history,
        F_history,
    )
end


const default_accuracy = 1.0e-12
const default_max_iterations = 1000
const near_int_threshold = 1.0e-6
const series_transition_threshold = 0.25

struct Diagnostics # this immutable type has no fields, so constructing it is essentially free
end

"""
    polylog(s, z)

Calculates the Polylogarithm function ``{Li}_s(z)`` defined by

``{Li}_s = \\sum_{n=1}^{\\infty} \\frac{z^n}{n^s}``

Uses double precision complex numbers (not arbitrary precision).
It's goal is an relative error bound 10^{-12}.

## Input Arguments
* ``s`` `::Complex`: the 'fractional' parameter
* ``z`` `::Complex`: the point at which to calculate it

It should also accept input arguments as `Real` or `Rational` or `Integer` but these aren't completely tested.

There are additional keywords, but these are currently intended for testing not use.

## Output Arguments
* ``Li_s(z)``: The result

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> polylog(0.35, 0.2)
0.23803890574407033
```
"""
function polylog(s::Number, z::Number;
                 level=1, # keep track of recursion
                 accuracy::Float64=default_accuracy,
                 min_iterations::Integer=0,
                 max_iterations::Integer=default_max_iterations)
    polylog(s, z, Diagnostics();
            level=level, accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)[1] # just output the result
end

# this is the main version, but outputs diagnosstics, which I guess most people won't want
"""
    polylog(s, z, Diagnostics())

Calculates the Polylogarithm function ``{Li}_s(z)`` defined by

``{Li}_s = \\sum_{n=1}^{\\infty} \\frac{z^n}{n^s}``

Uses double precision complex numbers (not arbitrary precision).
It's goal is an relative error bound 10^{-12}.

## Input Arguments
* ``s::`` `Complex`: the 'fractional' parameter
* ``z`` `::Complex`: the point at which to calculate it
* `::Diagnostics`: use this to indicate that the output should include extra information

It should also accept input arguments as `Real` or `Rational` or `Integer` but these aren't completely tested.

There are additional keywords, but these are currently intended for testing not use.

## Output Arguments
* ``Li_s(z)``: The result
* ``n``:       The number of elements used in each series
* `series`:    The series used to compute results (4 = reciprocal)

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> polylog(0.35, 0.2, Diagnostics() )
(0.23803890574407033, 17, 1)
```
"""
function polylog(s::Number, z::Number, ::Diagnostics;
                 level=1, # keep track of recursion
                 accuracy::Float64=default_accuracy,
                 min_iterations::Integer=0,
                 max_iterations::Integer=default_max_iterations)
    tau_threshold = 1.0e-3
    μ = log(convert(Complex{Float64}, z)) # input z could be an integer or anything
    t = abs(μ/twoπ)
    T = 0.512 # the duplication formula seems to work if T=0.5 in the small associated wedge, but wyh risk it?
    if abs(μ) < 1.0e-14
        # deal with this case separately or it causes pain in Series 2 and 3, which should be resolvable, but its just easier here
        # there are lots of other special cases to add here eventually
        if real(s) > 1
            return SpecialFunctions.zeta(s), 0, 0
        else
            return typeof(z)(Inf), 0 , 0
        end
    elseif abs(z) <= 0.5 && abs(z) < t
        return polylog_series_1(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    elseif t <= T && ( abs(round(real(s))-s) > tau_threshold || real(s)<= 0 )
        return polylog_series_2(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    elseif t <= T
        return polylog_series_3(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
#    elseif t <= 2.0
    else
        # println("  main level $level, z=$z,  abs(μ)/2π = ", abs(log(z))/twoπ  )
        return polylog_duplication(s, z; level=level, accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    # elseif abs(z) > 1
    #     return polylog_reciprocal(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    # else
    #     throw(DomainError(z, "bad z value $z with |z|=$(abs(z)) and |log(z)/(twoπ)|=$(abs(μ/twoπ))"))
    end
    # we could have a lot more special cases here, particularly for integer input
    # to make the code faster for these cases, but at the moment may use such values for testing
end

# old version without square root
    # if abs(z) <= 0.5 && abs(z) < abs(μ/twoπ)
    #     return polylog_series_1(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    # elseif abs(μ/twoπ) < series2_threshold && ( abs(round(real(s))-s) > tau_threshold || real(s)<= 0 )
    #     return polylog_series_2(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    # elseif abs(μ/twoπ) < series2_threshold
    #     return polylog_series_3(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    # elseif abs(z) > 1
    #     return polylog_reciprocal(s, z; accuracy=accuracy, min_iterations=min_iterations, max_iterations=max_iterations)
    # else
    #     throw(DomainError(z, "bad z value $z with |z|=$(abs(z)) and |log(z)/(twoπ)|=$(abs(μ/twoπ))"))
    # end

# function polylog(s::Number, z::Number, accuracy::Real=default_accuracy)
#     if z ≈ 1.0
#         if real(s) > 1
#             return zeta(s)
#         else
#             return Inf
#         end
#     elseif z ≈ -1.0
#         return -eta(s)
#     elseif s ≈ 0.0
#         return z ./ (1-z)
#     elseif abs(z) <= T
#         ifconfig
#         return polylog_direct(s, z, accuracy)
#     elseif abs(z) >= 1/T && isinteger(s) && real(s) < 0
#         # use reciprocal formula to calculate in terms of Li_n(1/z)
#         # but note for negative integer s, it collapses to something small
#         return -(-1.0)^s .*polylog_direct(s, 1/z, accuracy)
#     elseif  abs(z) >= 1/T
#         # use reciprocal formula to calculate in terms of Li_s(1/z)
#         twopi = 2π
#         z = convert(Complex{Float64}, z)
#         G = (twopi*im)^s * zeta( 1-s, 0.5 + log(-z)/(twopi*im) ) /  gamma(s)
#         F = complex(-1.0)^s * polylog_direct(s, 1/z, accuracy)

#         A = twopi*im*log(z)^(s-1)/(gamma(s))
#         if ( isreal(z) && real(z)>=1 )
#             Θ = 1
#         else
#             Θ = 0
#         end
#         # println("G = $G, F=$F, Θ=$Θ, A=$A")
#         return ( G - F - Θ*A )
#     else
#         # power series around mu=0, for z = e^mu
#         polylog_series_mu(s, z, accuracy)
#     end
# end

####################################################################
#### these are component functions and aren't exported at this point
#### note that for consistency they all have keywords arguments like "accuracy" but
#### these aren't intended for general use, just for testing (at the moment)

# # calculate using the relationship to the Hurwitz zeta function
# function polylog_zeta(s::Number, z::Number, accuracy=default_accuracy)
#     # compute using the Hurwitz-zeta function identity
#     #   N.B. this doesn't seem to work as well as you might think
#     x = im * (log(convert(Complex{Float64}, -z)) / twoπ)
#     ss = 1-s
#     ip = im^ss
#     return ( SpecialFunctions.gamma(ss)/twoπ^(ss) ) * (ip * SpecialFunctions.zeta(ss, 0.5+x) + conj(ip) * SpecialFunctions.zeta(ss, 0.5-x))
# end

# calculate using the duplication formula
function polylog_duplication(s::Number, z::Number;
                             level=0, # keep track of recursion
                             accuracy::Float64=default_accuracy,
                             min_iterations::Integer=0,
                             max_iterations::Integer=default_max_iterations)
    z = convert(Complex{Float64}, z)
    f = min( 0.5, 2.0^(1-real(s)) )
    # println("  dup level $level, z=$z,  abs(μ)/2π = ", abs(log(z))/twoπ  )
    (Li1, k1, series1) = polylog(s,  sqrt(z), Diagnostics();
                                 level=level+1, accuracy=f*accuracy,
                                 min_iterations=min_iterations, max_iterations=max_iterations)
    (Li2, k2, series2) = polylog(s, -sqrt(z), Diagnostics();
                                 level=level+1, accuracy=f*accuracy,
                                 min_iterations=min_iterations, max_iterations=max_iterations)
    if typeof(s) <: Real
        s = convert(Float64, s) # convert s into a double
    elseif typeof(s) <: Complex
        s = convert(Complex{Float64}, s) # convert s into doubles
    end
    return (2^(s-1) * ( Li1 + Li2 ), k1 + k2, 10+series1+series2)
end

# # calculate using the reciprocal formula
# function polylog_reciprocal(s::Number, z::Number;
#                             accuracy::Float64=default_accuracy,
#                             min_iterations::Integer=0,
#                             max_iterations::Integer=default_max_iterations)
#     # z = convert(Complex{Float64}, z)
#     if abs(z) <= 1
#         throw(DomainError(z, "only use this function for |z|>1, and pref |z| > 2"))
#     end
#     # if abs(z) < 2
#     #     warn("Slow convergence for  |z| < 2")
#     # end
#     if abs(s) < 0.1*accuracy
#         return (z/(1-z), 0) # use the identity
#     elseif real(s) < 0 &&  abs(imag(s)) < 0.1*accuracy && abs( round(s) - s ) < 0.1*accuracy
#         G = 0.0 # pole of the Gamma function
#         A = 0.0
#     else
#         # G = (twoπ*im)^s * SpecialFunctions.zeta( 1-s, 0.5 + log(complex(-z))/(twoπ*im) ) /  SpecialFunctions.gamma(s)
#         # A = twoπ*im*log(z)^(s-1) / SpecialFunctions.gamma(s)
#         tmp = exp( s*log(twoπ*im) - SpecialFunctions.loggamma(complex(s)) )  # (twoπ*im)^s /  SpecialFunctions.gamma(s)
#         G = tmp * SpecialFunctions.zeta( 1-s, 0.5 + log(complex(-z))/(twoπ*im) )
#         A = twoπ*im*log(z)^(s-1) / SpecialFunctions.gamma(s)
#     end
#     # accuracy of overall result depends on size of total, which includes these other parts
#     (Li, k, series) = polylog_series_1(s, 1/z; accuracy=0.1*accuracy,
#                                min_iterations=min_iterations, max_iterations=max_iterations, existing_total=G)
#     F = complex(-1.0)^s * Li
#     if ( imag(z) == 0 ) &&  ( real(z) >= 1 )
#         Θ = 1.0
#     else
#         Θ = 0.0
#     end
#     # println("G = $G, F=$F, Θ=$Θ, A=$A")
#     return ( G - F - Θ*A, k, 3+series )
# end

# calculate using direct definition
function polylog_series_1(s::Number, z::Number;
                          accuracy::Float64=default_accuracy,
                          min_iterations::Integer=0,
                          max_iterations::Integer=default_max_iterations,
                          existing_total::Number=0.0)
    # prolly should convert z to a double or complex-double
    if abs(z) > 1 || ( abs(z) ≈ 1  && real(s) <= 2)
        throw(DomainError(z))
    end
    if abs(z) > 1/2
        throw(DomainError(z, "Slow convergence for  |z| > 1/2"))
    end
    total = 0.0
    converged = false
    a = z
    k = 0
    if real(s) < 0
        min_iterations = ceil( real(s) / log(abs(z)) )
    end
    while k<=max_iterations && ~converged
        k = k+1
        total += a
        a *= z * ( k/(k+1.0) )^s
        # println("   total = $total")
        if k > min_iterations && abs(a)/abs(total) < 0.5*accuracy
            converged = true
        end
    end
    return (total, k, 1)
end

# calculate using power series around μ = log(z) = 0
# this should not be used near positive integer values of s, but we allow it here in order to test
function polylog_series_2(s::Number, z::Number;
                          accuracy::Float64=default_accuracy,
                          min_iterations::Integer=0,
                          max_iterations::Integer=default_max_iterations )
    μ = log(convert(Complex{Float64}, z)) # input z could be an integer or anything
    if typeof(s) <: Real
        s = convert(Float64, s) # convert s into a double
    elseif typeof(s) <: Complex
        s = convert(Complex{Float64}, s) # convert s into doubles
    end
    # println("μ = $μ")
    if abs(μ) > twoπ
        throw(DomainError(z, "we need |log(z)|< 2π for this series"))
    end
    # if real(s) > 0
        # min_iterations = ceil( real(s) ) + 1
    # else
        # min_iterations = ceil( -real(s) ) + 1
        # min_iterations = ceil( real(s) / log(abs( log(z)/twoπ ) ) )
    # end
    oneminuss = 1.0 - s
    total = SpecialFunctions.gamma(oneminuss) * (-μ)^(- oneminuss)
    # total = exp( SpecialFunctions.loggamma(oneminuss) - oneminuss*log(-μ) )
    converged = false
    tmp = 1
    k = 0
    a = Inf
    a_2 = Inf
    # A = abs( 2.0*twoπ^real(s) * exp(abs(imag(π*s)))  )
    # this doesn't work if z=1.0, and hence μ=0, even when that converges, but should already be delt with
    while k<=max_iterations && ~converged
        a_3 = a_2
        a_2 = a
        a = tmp * SpecialFunctions.zeta(s - k)
        total += a
        tmp *= μ/(k+1)
        if k > min_iterations &&
            abs(a)/abs(total) < 0.5*accuracy &&
            abs(a_2)/abs(total) < 0.5*accuracy &&
            abs(a_2) > abs(a)
            # abs( A * (k-real(s))^real(-s) * (μ/twoπ)^k )/abs(total) < accuracy
            # && abs( 2*twoπ^real(s) * (μ/twoπ)^k )/abs(total) < accuracy
            # the stopping rule should be implemented more efficiently as part of the calculation above
            converged = true
        end
        k = k + 1
    end
    # get correct value along the branch
    if isreal(z) && real(z)>=1
        # total -= 2*π*im*μ^(s-1)/SpecialFunctions.gamma(s)
        total -= exp( log(twoπ*im) + (s-1)*log(μ) - SpecialFunctions.loggamma(s) )
    end
    return (total, k, 2)
end

function c_closed(n::Integer, j::Integer,  ℒ::Number)
    d2 = SpecialFunctions.digamma(n+1) - ℒ
    if j==0
        return harmonic(n) - ℒ
    elseif j==1
        # Wood:+stieltjes(1) - d2^2/2 + π^2/6 + SpecialFunctions.polygamma(1,n+1)/2
        return -stieltjes(1) - d2^2/2 - π^2/6 + SpecialFunctions.polygamma(1,n+1)/2
    elseif j==2
        # Wood:stieltjes(2)   + d2^3/6 + d2*( π^2/6 + SpecialFunctions.polygamma(1,n+1)/2 ) + SpecialFunctions.polygamma(2,n+1)/6
        return stieltjes(2)/2 + d2^3/6 + d2*( π^2/6 - SpecialFunctions.polygamma(1,n+1)/2 ) + SpecialFunctions.polygamma(2,n+1)/6
   end
end

function Q_closed(n::Integer, τ::Number, ℒ::Number; n_terms::Integer=3)
    # τ is the distance from the pole s=n>0, ℒ = log(-μ) = log(-log( z ))
    max_n_terms = 3
    if n_terms < 1 || n_terms > max_n_terms
        throw(DomainError(n_terms))
    end
    return sum( c_closed.(n, 0:n_terms-1,  ℒ) .* τ.^(0:n_terms-1) )
end

function Q(n::Integer, τ::Number, ℒ::Number; n_terms::Integer=5) # Crandall,2012, p.35
    # τ is the distance from the pole s=n>0, ℒ = log(-μ) = log(-log( z ))
    if abs(τ) <= 1.0e-14
        # if really close to the integer, then ignore the extra terms
        return c_closed(n, 0,  ℒ)
    else
        max_n_terms = 7
        if n_terms < 1 || n_terms > max_n_terms
            throw(DomainError(n_terms))
        end
        if n_terms <= 3
            # use the direct method in this case
            return Q_closed(n, τ, ℒ; n_terms=n_terms)
        end
        return sum( c_crandall.(n, 0:n_terms-1,  ℒ) .* τ.^(0:n_terms-1) )
    end
end

function c_crandall(k::Integer, j::Integer,  ℒ) # Crandall,2012, p.35
    return (-1)^j*stieltjes(j)/SpecialFunctions.gamma(j+1) - b_crandall(k,j+1,ℒ)
end

function b_crandall(k::Integer, j::Integer,  ℒ) # Crandall,2012, p.36
    total = 0
    for q=0:j
        for t=0:j-q
            p = j-q-t
            a1 = ℒ^p / SpecialFunctions.gamma(p+1)
            a2 = (-1)^t * f_crandall(k,q) # Bailey and Borwein, 2015 correct Crandall (t+q - > t)
            a3 = g_crandall(t)/SpecialFunctions.gamma(t+1)
            total += a1 * a2 * a3
        end
    end
    return total
end

const gamma_t = [1.0,-0.5772156649015315,1.9781119906559432,-5.44487445648531,23.561474084025583,-117.83940826837748,715.0673625273184,-5019.848872629852,40243.62157333573,-362526.2891146549,3.627042412756892e6,-3.990708415143132e7,4.7894329176518273e8,-6.226641351546061e9,8.717563381070836e10]
function g_crandall(t::Integer) # Crandall,2012, p.17
    # t derivate of Gamma function at 1
    # see "gamma_derivatives.jl" for derivations of these numbers
    if t<0
        throw(DomainError(t))
    elseif t>14
        throw(DomainError(t, "only calculate the 1st 14 derivatives"))
    else
        return gamma_t[t+1]
    end
end

function f_crandall(k::Integer, q::Integer) # Crandall,2012, p.36
    # Crandall omits the k=0 case. but this is given in Bailey and Borwein and matches other text
    if k == 0 && q == 0
        return 1
    elseif k == 0
        return 0
    elseif q == 0
        return 1
    # elseif q == 1
    #     return -harmonic(k)
    # elseif q == 2
    #     return (harmonic(k)^2 + harmonic(k,2))/2
    else
        h = 0:q
        return sum( (-1).^h .* f_crandall.(k-1, q .- h)  ./ k.^h )
    end
end

# For the special case that s is near a postive integer n>0
# Calculate in a power series around z=1, and s=n
function polylog_series_3(s::Number, z::Number;
                          accuracy::Float64=default_accuracy,
                          min_iterations::Integer=0,
                          max_iterations::Integer=default_max_iterations,
                          n_terms::Integer=5 )
    μ = log(convert(Complex{Float64}, z))
    if abs(μ) > twoπ
        throw(DomainError(z, "does not converge for abs(ln(z)) > twoπ"))
    end
    if real(s)<=0.5
        throw(DomainError(s, "for this function s should be near a positive integer"))
    end
    # this series assumes s is near a positive integer
    n = Int(round(real(s)))
    τ = s - n
    # if real(s) > 0
    #     min_iterations = ceil( real(s) )
    # end
    ℒ = log(complex(-μ))  # '\u2112'
    # total = μ^(n-1)*Q(n-1, τ, ℒ; n_terms=n_terms)/SpecialFunctions.gamma(n)
    total = exp( (n-1)*log(μ) + log(Q(n-1, τ, ℒ; n_terms=n_terms))  - SpecialFunctions.loggamma(n) )
    converged = false
    a = Inf
    a_2 = Inf
    tmp = 1
    k = 0
    while k<=max_iterations && ~converged
        if n - k != 1
            a_2 = a
            a = tmp * SpecialFunctions.zeta(s - k)
            total += a
        end
        tmp *= μ/(k+1)
        if k > min_iterations &&
            abs(a)/abs(total) < 0.5*accuracy &&
            abs(a_2)/abs(total) < 0.5*accuracy &&
            abs(a_2) > abs(a)
            # abs( (μ/twoπ)^k )/abs(total) < 0.05*accuracy
            converged = true
        end
        k = k + 1
    end
    # get correct value along the branch
    if isreal(z) && real(z)>=1
        total -= 2*π*im*μ^(s-1)/SpecialFunctions.gamma(s)
    end
    return (total, k, 3)
end

# # asymptotic expansion
# function polylog_series_4(s::Number, z::Number;
#                           accuracy::Float64=default_accuracy,
#                           min_iterations::Integer=0,
#                           max_iterations::Integer=default_max_iterations,
#                           existing_total::Number=0.0)
#     x = log(-convert(Complex{Float64}, z))
#     total = 0.0
#     converged = false
#     k = 0
#     a = Inf
#     while k<=max_iterations && ~converged
#         old_a = a
#         a = (-1)^k * (1.0 - 2.0^(1-2*k)) * twoπ^(2*k) * bernoulli(2*k,0.0) * x^(s-2*k) / ( SpecialFunctions.gamma(2*k+1) * SpecialFunctions.gamma(s+1-2*k) )
#         if abs(a) < abs(old_a)
#             total += a
#         else
#             converged = true
#         end
#         k = k+1
#     end
#     return (total, k, 4)
# end

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

# # Special case s=integer for use in testing etc.
#     function polylog(n::Integer, z::Number;
#                       accuracy::Float64=default_accuracy, max_iterations::Integer=default_max_iterations )
#     L = Int(ceil(-log10(accuracy)*log2(10))) # revisit this limit
#     k = 0
#     if n>1
#         μ = log(convert(Complex{Float64}, z))
#         if abs(μ) > twoπ
#             throw(DomainError(z))
#         end
#         total = μ^(n-1)*(harmonic(n-1) - log(-μ))/SpecialFunctions.gamma(n)
#         tmp = 1
#         for m=0:L
#             if n - m != 1
#                 total += tmp * zeta(n - m)
#             end
#             tmp *= μ/(m+1)
#             if abs(tmp)/abs(total) < 1.0e-30
#                 break
#             end
#         end
#         if  isreal(z) && real(z)>=1
#             total -= 2*pi*im*μ^(s-1)/SpecialFunctions.gamma(n)
#         end
#     elseif n==1
#         total = -log(complex(1-z))
#     elseif n==0
#         total = z / (1-z)
#     elseif n==-1
#         total = z / (1-z)^2
#     elseif n<-1
#         # Crandall's 1.5 for s integer
#         total = factorial(-n) * (-μ)^(n-1)
#         tmp = 1
#         for k=0:L
#             # total -= μ^k * bernoulli(k-n+1, 0.0) / ( gamma(k+1)*(k-n+1) )
#             total -= tmp * bernoulli(k-n+1, 0.0) / (k-n+1)
#             tmp *= μ/(k+1)
#             if abs(tmp)/abs(total) < 1.0e-30
#                     break
#             end
#         end
#     else
#         error("Should not get this case")
#     end

#     if isreal(z) && real(z)>=1
#         total -= 2*pi*im*μ^(s-1)/gamma(s)
#     end
#     return (total,k)
# end


# ---
# title: Template Matching
# cover: assets/template_matching.png
# author: Michael Pusterhofer
# date: 2021-06-22
# ---

# This demo shows how to find objects in an image using template matching.

# The main idea is to check the similarity between a search target(template) and a subsection of the image. The subsection is usually the same size as the template. Therefore by using [`mapwindow`](@ref) we can assign a value for every subsection of the image.

# At first we import the following packages.
using ImageCore: Gray
using ImageMorphology: label_components, component_centroids
using ImageFiltering: mapwindow, Fill, imfilter, KernelFactors
using ImageDistances: sqeuclidean
using ImageContrastAdjustment: adjust_histogram, LinearStretching
using TestImages
using Plots: scatter!, plot

# `ImageCore` enables the generation of images, `ImageFiltering` provides the [`mapwindow`](@ref) function and `ImageFeatures` provides functions to label segments of an image. To calculate the similarity `squeclidean` is used. This defines a distance between two images as:
#
# ```math
# \textrm{distance} = \sum_i^{\textrm{number of pixels}} (\textrm{template}_i - \textrm{subsection}_i)^2
# ```
#
# `ImageContrastAdjustment` provides functions to adjust the histogram, which is useful when an image contains values bigger than 1 or smaller than 0.  The `Testimages` package will provide our image and plot is used to overlay a scatter plot onto the image.

# To start we first load our image.
img = testimage("moonsurface")

# Let's say we want to find the medium sized craters in the image based one of the impact areas. For this we generate a template from a subsection of the image and apply a small gaussian blur using [`imfilter`](@ref). The gaussian blur often helps when the search targets are not exactly the same.
template = img[12:22,20:30]
template = imfilter(template,KernelFactors.gaussian((0.5,0.5)))

# Now that we have an image and a template, the next step is to define how we measure the similarity between a section of the image and the template. This can be done in multiple way, but a sum of square distances should work quite well. The `ImageDistance` package provides an already optimized version called `sqeuclidean`, which can be used to define a function for [`mapwindow`](@ref).  Let's call it `SDIFF`.

function SDIFF(template)
  (subsection)->sqeuclidean(subsection, template)
end

# To actually generate our similarity map we use [`mapwindow`](@ref) in the following way.

res = mapwindow(SDIFF(template), img, size(template), border=Fill(1)) .|> Gray
rescaled_map = adjust_histogram(res, LinearStretching())

# If the subsection is located at the border of the image the image has to be extended and in our case we will fill all values outside the image with 1. One thing to keep in mind is that because all of the square differences will be summed up per subsection the resulting sum can be bigger than 1. This will be a problem if we just convert it to an image to check the values. To rescale the values to be between 0 and 1 we can use `adjust_histogram`.

# To find the best locations we have to look for small values on the similarity map. This can be done by comparing if the pixel is below a certain value. Let's chose a value of `0.05`.

threshold = rescaled_map .< 0.05
Gray.(threshold)

# Now we see small blobs at the locations which match our template and we can label the connected regions by `label_components`.  This will enumerate are connected regions and `component_centroids` can be used to get the centroid of each region.  `component_centroids` also return the centroid for the background region, which is at the first position and we will omit it.

centroids = component_centroids(label_components(threshold))[2:end]

# To check if it worked correctly we can overlay the centroids with the original image using the `Plots` package. As the images are stored using the first index for rows we have to reverse the order of the coordinates to match the order of the plotting library.
plot(Gray.(img), size=(512,512))
scatter!(reverse.(centroids), label="centroids", ms=10, alpha=0.5, c=:red, msw=3)

using Plots: savefig #src
savefig("assets/template_matching.png") #src


# Functions:
function fft_len(x::AbstractVector, len::Int)
    if len > length(x)
        return fft([x; zeros(len - length(x))])
    else
        return fft(x[1:len])
    end
end

function conv_same(x::AbstractVector, y::AbstractVector)
    # RoundUp is used to handle scenarios where length(y) is odd
    s = div(length(y), 2, RoundUp)
    e = length(x) + s - 1
    return conv(x, y)[s:e]
end

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

# Julia code for Example 10.5
using Plots

x = zeros(1000, 20)
t = range(-2, 2, length=1000)

for col in eachcol(x)
    col .= rand() * cos.(2π * t)
end

plot(t, x, linewidth=2, color=:gray80, legend=false)
plot!(t, 0.5 * cos.(2π * t), linewidth=4, color=:darkred)

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

# Julia code for Example 10.6
using Plots

x = zeros(1000, 20)
t = range(-2, 2, length=1000)

for col in eachcol(x)
    col .= cos.(2π*t .+ 2π*rand())
end

plot(t, x, linewidth=2, color=:gray80, legend=false)
hline!([0], linewidth=4, color=:darkred)

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

# Julia code for Example 10.7
using Plots

x = zeros(21, 20)
t = 0:20

for col in eachcol(x)
    col .= rand() .^ t
end

sticks(t, x, linewidth=2, shape=:circle, color=:gray80, legend=false)
sticks!(t, 1 ./ (t .+ 1), linewidth=2, shape=:circle, markersize=6, color=:darkred)

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

# Julia code for Example 10.11: Plot the time function
using Plots

x = zeros(1000, 20)
t = range(-2, 2, length=1000)

for col in eachcol(x)
    col .= rand() * cos.(2π * t)
end

plot(t, x, linewidth=2, color=:gray80, legend=false)
plot!(t, 0.5 * cos.(2π * t), linewidth=4, color=:darkred)
scatter!(zeros(20), x[501, :], markersize=4, color=:orange)
scatter!(0.5 * ones(20), x[626, :], markersize=4, color=:blue)

# Julia code for Example 10.11: Plot the autocorrelation function
t = range(-1, 1, length=1000)
R = 1/3 * cos.(2π * t) * cos.(2π * t')

heatmap(t, t, R, yticks= -1:0.25:1, xticks= -1:0.25:1, legend=false, yflip=true)

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

# Julia code for Example 10.12: Plot the time function
using ToeplitzMatrices
using Plots

x = zeros(1000, 20)
t = range(-2, 2, length=1000)

for col in eachcol(x)
    col .= cos.(2π*t .+ 2π*rand())
end

plot(t, x, linewidth=2, color=:gray80, legend=false)
plot!(t, 0 * cos.(2π * t), linewidth=4, color=:darkred)
scatter!(zeros(20), x[501, :], markersize=4, color=:orange)
scatter!(0.5 * ones(20), x[626, :], markersize=4, color=:blue)

# Julia code for Example 10.12: Plot the autocorrelation function
t = range(-1, 1, length=1000)
R = SymmetricToeplitz(0.5 * cos.(2π * t))

heatmap(t, t, R, yticks= -1:0.25:1, xticks= -1:0.25:1, legend=false, yflip=true)

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

# Julia code to demonstrate autocorrelation
using Plots
using DSP

# Figure 1
X₁ = randn(1000)
X₂ = randn(1000)

plot(X₁, linewidth=2, linestyle=:dot, color=:blue, legend=false)
plot!(X₂, linewidth=2, linestyle=:dot, color=:black)

# Figure 2
N  = 1000 # number of sample paths
T  = 1000 # number of time stamps
X  = randn(N, T)
xc = zeros(N, 2T - 1)

for i in 1:N
    xc[i, :] .= xcorr(X[i, :], X[i, :]) / T
end

plot(xc[1, :], linewidth=2, linestyle=:dot, color=:blue, label="correlation of sample 1")
plot!(xc[2, :], linewidth=2, linestyle=:dot, color=:black, label="correlation of sample 2")
####################################################################

# Julia code for Example 10.15
using Plots, LaTeXStrings
using Images
using DSP

# Figure 1
t = -10:0.001:10
L = length(t)
X = randn(L)
h = 10 * max.(1 .- abs.(t), 0) / 1000
Y = imfilter(X, h, "circular")

plot(t, X, linewidth=0.5, color=:gray60, label=L"X(t)")
hline!([0], linewidth=4, color=:gray20, label=L"\mu_X(t)")
plot!(t, Y, linewidth=2, color=:red, label=L"Y(t)")
hline!([0], linewidth=4, linestyle=:dot, color=:yellow, label=L"\mu_Y(t)")
xlims!((-10, 10))

# Figure 2
h₂ = conv(h, h)
Rₓ = zeros(40001)
Rₓ[20001] = 0.2

plot(-20:0.001:20, Rₓ, linewidth=2, color=:gray60, label=L"R_X(t)")
plot!(-20:0.001:20, h₂, linewidth=2, color=:red, label=L"R_Y(t)")
xlims!((-2, 2))
ylims!((-0.05, 0.2))

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

# Julia code to solve the Yule Walker Equation
using Plots, LaTeXStrings
using ToeplitzMatrices
using DelimitedFiles
using DSP

y = vec(readdlm("ch10_LPC_data.txt"))
K = 10
N = length(y)
y_corr = xcorr(y, y)
R = SymmetricToeplitz(y_corr[N:N+K-1])
lhs = y_corr[N+1:N+K]
h = R \ lhs

# Figure 1
plot(y, linewidth=4, color=:blue, label=L"Y[n]")

# Figure 2
plot(y_corr, linewidth=4, color=:black, label=L"R_Y[k]")

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

# Julia code to predict the samples
using ToeplitzMatrices
using DelimitedFiles
using Plots
using DSP

y = vec(readdlm("ch10_LPC_data_02.txt"))
K = 10
N = length(y)

y_corr = xcorr(y, y)
R = SymmetricToeplitz(y_corr[N:N+K-1])
lhs = y_corr[N+1:N+K]
h = R \ lhs

z = y[311:320]
ŷ = zeros(340)
ŷ[1:320] .= y

for t in 1:20
    pred = z' * h
    z = [z[2:10]; pred]
    ŷ[320+t] = pred
end

plot(ŷ, linewidth=3, color=:red, label="Prediction", legend=:bottomleft)
plot!(y, linewidth=4, linestyle=:dash, color=:gray60, label="Input")

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

# Julia code for Wiener filtering
using DSP, FFTW
using Plots

w = 0.05 * randn(320)
x = y + w

Ry = xcorr(y, y)
Rw = xcorr(w, w)
Sy = fft(Ry)
Sw = fft(Rw)
H = Sy ./ (Sy + Sw)

Ŷ = H .* fft_len(x, 639)
ŷ = real.(ifft(Ŷ))

# Figure 1
plot(x, linewidth=4, color=:gray, label="Noise Input X[n]", legend=:bottomleft)
plot!(ŷ[1:320], linewidth=2, color=:red, label="Wiener Filtered Yhat[n]")
plot!(y, linewidth=2, linestyle=:dot, color=:black, label="Ground Truth Y[n]")

# Figure 2
plot(Rw, linewidth=4, color=:blue, label="h[n]")
####################################################################

# Julia code to solve the Wiener deconvolution problem
using DelimitedFiles
using DSP, FFTW
using Plots

y = vec(readdlm("ch10_wiener_deblur_data.txt"))
g = ones(32) / 32
w = 0.02 * randn(320)
x = conv_same(y, g) + w

Ry = xcorr(y, y)
Rw = xcorr(w, w)
Sy = fft(Ry)
Sw = fft(Rw)
G  = fft_len(g, 639)

H = @. (conj(G) * Sy) / (abs(G)^2 * Sy + Sw)
Ŷ = H .* fft_len(x, 639)
ŷ = real.(ifft(Ŷ))

plot(x, linewidth=4, color=:gray, label="Noise Input X[n]")
plot!(16:320+15, ŷ[1:320], linewidth=2, color=:red, label="Wiener Filtered Yhat[n]")
plot!(y, linewidth=2, linestyle=:dot, color=:black, label="Ground Truth Y[n]")


# A development file for testing implementation of non-stationary rouwenhorst method (à la FGP2019)
# built on existing function for stationary processes from package QuantEcon

# Citations
# FGP2019: Fella, Gallipoli, Pan, 2019
# KV2010: Kaplan, Violante, 2010

using QuantEcon, Plots, Random, Statistics

function simulate_PTincome(N, T, σ_u, σ_ε::Real, ρ::Real, σ_0)
# Simulates persistent-transitory, non-stationary income process as per FGP2019
# Only gaussian persistent shock is implemented
# Inputs are standard deviations
# σ_0 is the standard deviation of η_0 (denoted simply as σ in FGP2019)
    Y = zeros(Float64, N, T)
    η_t1 = σ_0 .* randn(N)

    for t in 1:T
        η_t = ρ * η_t1 + σ_ε .* randn(N)
        Y[:,t] = exp.( η_t + σ_u .* randn(N) )
        η_t1 = copy(η_t)
    end
    return Y'
end

### Check simulation
σ_ε = sqrt(0.0161)    #variance is σ_η^2=0.0161 in FGP2019
σ_u = 0. #sqrt(0.063)   #variance is σ_u^2=0.063 in FGP2019
σ_η0 = 0.             # In FGP2019
ρ = 0.9            # ρ ∈ {0.95,0.98,1} in FGP2019 section 4.1.1
N = 10000
T = 20
Random.seed!(1234)
Y = simulate_PTincome(N, T, σ_u, σ_ε, ρ, σ_η0)
plot(1:T,Y[:,1])
plot(var(Y,dims=2))
histogram(Y[1,:], bins = 100)
histogram!(Y[end,:], bins = 100, alpha = 0.5)
histogram!(Y[end-1,:], bins = 100, alpha = 0.5)

### Test rouwenhorst_ns
σ_ε = sqrt(0.0161)    #variance is σ_η^2=0.0161 in FGP2019
σ_u = sqrt(0.063)   #variance is σ_u^2=0.063 in FGP2019
σ_η0 = 0.             # In FGP2019
ρ = 0.98            # ρ ∈ {0.95,0.98,1} in FGP2019 section 4.1.1
T = 15
test = rouwenhorst_ns(19, T, 0.9, σ_ε, σ_η0)
test[1].state_values
test[1].p
heatmap(test[15].p)

# Compare to stationary
test_stat = rouwenhorst(19, 0.9, σ_ε)
heatmap(test_stat.p)
maximum(test[15].p - test_stat.p)

# Difference due to initial variance? This is constrained for stationary case
test = rouwenhorst_ns(19, T, 0.9, σ_ε, σ_ε/sqrt(1-ρ^2))
maximum(test[15].p - test_stat.p)
# virtually identical at steady state
heatmap(test[5].p)

### Markov Chain simulation for non-stationary

# Check that my simulation function gives same results as QuantEcon's for
# stationary process
T = 25
N = 100000
ρ = 0.9
mc_qe = rouwenhorst(19, ρ, σ_ε)

# QuantEcon's MC simulation
# simulate(mc, T, init = 10) |> plot
Y_mc_qe = zeros(T+1,N)
[Y_mc_qe[:,i] = simulate(mc_qe, T+1, init = 10) for i in 1:N]
plot(0:T, var(Y_mc_qe, dims = 2), label="QE MC", xlabel="t", ylabel="Variance", legend=:bottomright)

# DGP simulation
Y = log.(simulate_PTincome(N, T, 0, σ_ε, ρ, 0.))
plot!(1:T, var(Y, dims = 2), label="DGP")

# My MC simulation with QE MC
mcs = fill(mc_qe,T)
Y_mc_qe1 = zeros(length(mcs)+1,N)
[Y_mc_qe1[:,i] = simulate_ns(mcs, 10) for i in 1:N]
plot!(0:T, var(Y_mc_qe1, dims = 2), label="MC")

# Compare to theoretical variance
expected_vars = function(σ, ρ, σ_y0, T)
    var_y = zeros(T)
    for t in 1:T
        var_y[t] = ρ^2 * σ_y0^2 + σ^2
        σ_y0 = sqrt(var_y[t])
    end
    return var_y
end
scatter!(1:T, expected_vars(σ_ε,ρ,0.,T), label = "Theory")

# My MC simulation with my MC
mcs_ns = rouwenhorst_ns(19, T, ρ, σ_ε, 0.)
Y_mc_ns = zeros(length(mcs_ns)+1,N)
[Y_mc_ns[:,i] = simulate_ns(mcs_ns, 10) for i in 1:N]
plot!(0:T, var(Y_mc_ns, dims = 2), label="MC all mine")

# Means
plot(mean(Y, dims=2))
plot!(mean(Y_mc_qe, dims = 2))
plot!(mean(Y_mc_ns, dims = 2))

# Check for ρ=1
ρ=1
T=30
scatter(1:T, expected_vars(σ_ε,ρ,0.,T), label = "Theory", legend=:bottomright)

# DGP simulation
Y = log.(simulate_PTincome(N, T, 0, σ_ε, ρ, 0.))
plot!(1:T, var(Y, dims = 2), label="DGP")

# My MC simulation with my MC
mcs_ns = rouwenhorst_ns(19, T, ρ, σ_ε, 0.)
Y_mc_ns = zeros(length(mcs_ns)+1,N)
[Y_mc_ns[:,i] = simulate_ns(mcs_ns, 10) for i in 1:N]
plot!(0:T, var(Y_mc_ns, dims = 2), label="MC all mine")

### Now let's try to discretize the transitory-persistent process!!!
σ_ε = sqrt(0.0161)    #variance is σ_η^2=0.0161 in FGP2019
σ_u = sqrt(0.063)   #variance is σ_u^2=0.063 in FGP2019
σ_η0 = 0.             # In FGP2019
ρ = 1.            # ρ ∈ {0.95,0.98,1} in FGP2019 section 4.1.1
N = 200000
T = 40
n = 11

# Simulate DGP
Random.seed!(1234)
Y = simulate_PTincome(N, T, σ_u, σ_ε, ρ, σ_η0)
plot(1:T, var(Y, dims = 2), label="DGP", legend=:bottomright)

# Markov chain for persistent component
mcs_pers = rouwenhorst_ns(n, T, ρ, σ_ε, σ_η0)
η = zeros(length(mcs_FGP)+1,N)
[η[:,i] = simulate_ns(mcs_pers, round(Int,n/2)) for i in 1:N]
# plot!(0:T, var(η, dims = 2), label="eta_FGP")

# Markov simulation with actual transitory component
u = σ_u*randn(size(η))
Y_m = exp.(η + u)
plot!(0:T, var(Y_m, dims = 2), label="Y_m_FGP")

# Markov simulation with discretized transitory component
# Transitory component
mc_trans = rouwenhorst(n, 0., σ_u)
u_mc = zeros(T+1,N)
[u_mc[:,i] = simulate(mc_trans, T+1, init = round(Int,n/2)) for i in 1:N]
# plot!(0:T, var(u_mc, dims = 2), label="u_mc", xlabel="t", ylabel="Variance", legend=:bottomright)

# Put together
Y_mm = exp.(η + u_mc)
plot!(0:T, var(Y_mm, dims = 2), label="Y_mm_FGP")
xlims!(1.,35.)
ylims!(0.,1.8)

Y_mm = Y_mm[2:end,:]
sim_error = (var(Y_mm, dims=2)-var(Y,dims=2))./var(Y,dims=2)*100
sim_error[end]
sim_error = (mean(Y_mm, dims=2)-mean(Y,dims=2))./mean(Y,dims=2)*100
sim_error[end]
# These compare quite well with FPG2019 Table 1

### Need to include initial variance in simulate_ns argument if non-zero!

### Then we calculate expectations!!!

# PROBLEMS
# NONE

stop

# Boneyard
# function simulate_AR1(N, T, σ_ε, ρ, σ_0)
# # η_t = ρ_t *η_{t-1} + ε_t, ε_t ~ N(0,σ_εt)
# # σ_0 is the standard deviation of η_0 (denoted simply as σ in FGP2019)
#     Y = zeros(Float64, N, T-1)
#     η_t1 = sqrt(σ_0) .* randn(N)
#
#     for t in 1:(T-1)
#         η_t = η_t1 + sqrt(σ_ε) .* randn(N)
#         η_t1 = copy(η_t)
#     end
#     return Y
# end

# function rouwenhorst(N::Integer, ρ::Real, σ::Real, μ::Real=0.0; stationary::Bool=true)
# # Note: μ is not implemented for non-stationary case
#     if stationary
#         return rouwenhorst(N, ρ, σ, μ)
#     else
#         return rouwenhorst_ns(N, T, repeat(ρ,T), repeat(σ,T))
# end

# Y[:,2]
# test = function()
#     pY = fill(plot(1:T-1,Y[:,1]),10)
#     for i in 1:10
#         pY[i] = plot(1:T-1,Y[:,i])
#     end
#     return pY
# end
# test()
# display(pY[2])
#
# plot()
# for i in 1:10
#     display(plot!(1:T+1,Y_mc_qe[:,i]))
# end
# plot()
# for i in 1:10
#     display(plot!(1:T+1,Y_mc_qe1[:,i]))
# end
#
#
# #######
# plot()
# for i in 100:110
#     plot!(Y[:,i], legend=:none) |>display
#     plot!(Y_m[:,i], line=:dash) |>display
#     # plot!(η[:,i]) |>display
# end

# @doc doc"""
# Simulate one sample path of the non-stationary vector of Markov chains `mcs`.
# The resulting vector has the state values of `mcs` as elements.
#
# ### Arguments
#
# - `mcs::Vector{MarkovChain}` : Vector of MarkovChains.
# - `;init::Int` : Index of initial state
#
# ### Returns
#
# - `X::Vector` : Vector containing the sample path, with length
#   length(mcs)+1 (includes initial state)
# """
# function simulate_ns(mcs::Vector, init::Int)
#     ind_0 = init
#     ind = zeros(Int, length(mcs))
#     X = zeros(Real, length(mcs))
#
#     for t in eachindex(mcs)
#         X[t] = simulate(mcs[t], 2, init = ind_0)[2]
#         ind[t] = findmin( abs.(mcs[t].state_values .- X[t]) )[2]
#         if t < length(mcs)
#             ind_0 = ind[t] #findmin( abs.(mcs[t].state_values .- X[t]) )[2]
#         end
#     end
#     ind = [init; ind]
#     X = [mcs[1].state_values[init]; X]
# end


using StatsPlots
using LinearAlgebra

"""
    barplot_THD(u, nodeID; h_min=1, h_max=0, yticks=0:0.1:1)

Barplot of harmonic voltage magnitudes at one node.
"""
function barplot_THD(u, nodeID; h_min=1, h_max=0, yticks=0:0.1:1)
    if h_max == 0
        h_max = maximum(keys(u))
    end
    h_max_idx = Int((h_max+1)/2)
    h_min_idx = Int((h_min+1)/2)
    x = sort(collect(keys(u)))[h_min_idx:h_max_idx]

    y = [u[h][nodeID, "v"] for h in x]

    THD = HarmonicPowerFlow.THD(u)
    bar(x, y,
        xticks = x,
        yticks = yticks,
        xlabel = "Harmonic Frequency",
        ylabel = "|V| [pu]",
        label = "Harmonic voltages at node 4",
        title = "THD = "*string(round(100*THD.THD_F[nodeID]; digits= 2))*"%")
end

"""
    heatmap_NE(nodes, settings; save=png)

Export heatmap indicating magnitude of Norton equivalent parameters tp `.png` or `.pdf`.
"""
function heatmap_NE(nodes, settings; save="png")
    harmonics = settings.harmonics
    NE = HarmonicPowerFlow.import_Norton_Equivalents(nodes, settings)
    l = @layout [a{0.66w}  b{1h}]
    for device in collect(keys(NE))
        device = string(device)
        if settings.coupled
            p1 = heatmap(harmonics, harmonics, abs.(NE[device][2]), ticks = harmonics, title = "Y_Norton Matrix ("*device*", coupled)", yflip=true)
            p2 = heatmap(transpose(abs.(NE[device][1])), aspect_ratio=1, yflip=true, title = "I_Norton ("*device*")", yticks = nothing, border = :none)
            p = plot(p1, p2, layout = l)
        else
            p1 = heatmap(harmonics, harmonics, diagm(vec(abs.(NE[device][2]))), ticks = harmonics, title = "Y_Norton Vector ("*device*", uncoupled)", yflip=true)
            p2 = heatmap(transpose(abs.(NE[device][1])), aspect_ratio=1, yflip=true, title = "I_Norton ("*device*")", yticks = nothing, border = :none)
            p = plot(p1, p2, layout = l)
        end
        if save == "png"
            savefig(p, "heatmap_"*device)
            println("Saved heatmap to file heatmap_"*device*".png")
        elseif save == "pdf"
            savefig(p, "heatmap_"*device*".pdf")
            println("Saved heatmap to file heatmap_"*device*".pdf")
        end
        # show(p) doesn't work, for now only returned if saved to file
    end
end


using Roots
using Compat.Test
import SpecialFunctions.erf

include("./test_find_zero.jl")
include("./test_fzero.jl")
include("./test_find_zeros.jl")
include("./test_newton.jl")
include("./test_simple.jl")
include("./RootTesting.jl")

#include("./test_composable.jl")

#run_benchmark_tests()

#include("./test_fzero3.jl")
#run_robustness_test()

#include("./test_derivative_free.jl")


using DiffEqFlux, DiffEqSensitivity, Flux, OrdinaryDiffEq, Zygote, Optim, NLopt, BlackBoxOptim, Test #using Plots
using MultistartOptimization, NLopt

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = (α - β*y)x
  du[2] = dy = (δ*x - γ)y
end
p = [2.2, 1.0, 2.0, 0.4]
u0 = [1.0,1.0]
prob = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)

# Reverse-mode

function predict_rd(p)
  Array(solve(prob,Tsit5(),p=p,saveat=0.1,reltol=1e-4))
end
function loss_rd(p)
  sum(abs2,x-1 for x in predict_rd(p))
end

loss_rd(p)

grads = Zygote.gradient(loss_rd, p)
@test !iszero(grads[1])

cb = function (p,l)
  display(l)
  false
end

# Display the ODE with the current parameter values.
loss1 = loss_rd(p)
pmin = DiffEqFlux.sciml_train(loss_rd, p, ADAM(0.1), cb = cb, maxiters = 100)
loss2 = loss_rd(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_rd, p, BFGS(initial_stepnorm = 0.01), cb = cb)
loss2 = loss_rd(pmin.minimizer)
@test 10loss2 < loss1
pmin = DiffEqFlux.sciml_train(loss_rd, p, Fminbox(BFGS(initial_stepnorm = 0.01)), lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_rd(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_rd, p, maxiters=100, lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_rd(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_rd, p, TikTak(100),
                              local_method = NLopt.LN_BOBYQA,
                              maxiters=100,
                              lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_rd(pmin.minimizer)
@test 10loss2 < loss1

#=
using QuadDIRECT
pmin = DiffEqFlux.sciml_train(loss_rd, p, QuadDIRECT(),
                              splits = (),
                              maxiters=100,
                              lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_rd(pmin.minimizer)
@test 10loss2 < loss1
=#

# Forward-mode, R^n -> R^m layer

p = [2.2, 1.0, 2.0, 0.4]
function predict_fd(p)
  vec(Array(solve(prob,Tsit5(),p=p,saveat=0.0:0.1:1.0,reltol=1e-4,sensealg=ForwardDiffSensitivity())))
end
loss_fd(p) = sum(abs2,x-1 for x in predict_fd(p))
loss_fd(p)

grads = Zygote.gradient(loss_fd, p)
@test !iszero(grads[1])

opt = ADAM(0.1)
cb = function (p,l)
  display(l)
  #display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
  false
end

# Display the ODE with the current parameter values.
loss1 = loss_fd(p)
pmin = DiffEqFlux.sciml_train(loss_fd, p, opt, cb = cb, maxiters = 100)
loss2 = loss_fd(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_fd, p, BFGS(initial_stepnorm = 0.01), cb = cb)
loss2 = loss_fd(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_fd, p, TikTak(100),
                              local_method = NLopt.LN_BOBYQA,
                              maxiters=100,
                              lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_fd(pmin.minimizer)
@test 10loss2 < loss1

#=
pmin = DiffEqFlux.sciml_train(loss_fd, p, QuadDIRECT(),
                              splits = (),
                              maxiters=100,
                              lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_fd(pmin.minimizer)
@test 10loss2 < loss1
=#

# Adjoint sensitivity
p = [2.2, 1.0, 2.0, 0.4]
ps = Flux.params(p)
function predict_adjoint(p)
    vec(Array(solve(prob,Tsit5(),u0=eltype(p).(prob.u0),p=p,saveat=0.1,reltol=1e-4)))
end
loss_reduction(sol) = sum(abs2,x-1 for x in vec(sol))
loss_adjoint(p) = loss_reduction(predict_adjoint(p))
loss_adjoint(p)

grads = Zygote.gradient(loss_adjoint, p)
@test !iszero(grads[1])

opt = ADAM(0.1)
cb = function (p,l)
  display(l)
  #display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
  false
end

# Display the ODE with the current parameter values.
loss1 = loss_adjoint(p)
pmin = DiffEqFlux.sciml_train(loss_adjoint, p, opt, cb = cb, maxiters = 100)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_adjoint, p, BFGS(initial_stepnorm = 0.01), cb = cb)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_adjoint, p, Fminbox(BFGS(initial_stepnorm = 0.01)), lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

opt = Opt(:LD_MMA, 4)
pmin = DiffEqFlux.sciml_train(loss_adjoint, p, opt)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_adjoint, p, TikTak(100),
                              local_method = NLopt.LN_BOBYQA,
                              maxiters=100,
                              lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

#=
pmin = DiffEqFlux.sciml_train(loss_adjoint, p, QuadDIRECT(),
                              splits = (),
                              maxiters=100,
                              lower_bounds = [0.0 for i in 1:4], upper_bounds = [5.0 for i in 1:4], cb = cb)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1
=#

function lotka_volterra2(u,p,t)
  x, y = u
  α, β, δ, γ = p
  dx = (α - β*y)x
  dy = (δ*x - γ)y
  [dx,dy]
end
p = [2.2, 1.0, 2.0, 0.4]
u0 = [1.0,1.0]
prob = ODEProblem{false}(lotka_volterra2,u0,(0.0,10.0),p)
function predict_adjoint(p)
    vec(Array(solve(prob,Tsit5(),p=p,saveat=0.1,reltol=1e-4)))
end
loss_reduction(sol) = sum(abs2,x-1 for x in vec(sol))
loss_adjoint(p) = loss_reduction(predict_adjoint(p))

pmin = DiffEqFlux.sciml_train(loss_adjoint, p, Newton())
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_adjoint, p, NewtonTrustRegion())
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1

pmin = DiffEqFlux.sciml_train(loss_adjoint, p, Optim.KrylovTrustRegion(), maxiters = 100)
loss2 = loss_adjoint(pmin.minimizer)
@test 10loss2 < loss1


"""
    Ysize(ℓₘᵢₙ, ℓₘₐₓ)

Compute total size of array of mode weights

Parameters
----------
ℓₘᵢₙ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ
ℓₘₐₓ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ

Returns
-------
i : int
    Total size of array of mode weights arranged as described below

See Also
--------
Yrange : Array of (ℓ, m) indices corresponding to this array
Yindex : Index of a particular element of the mode weight array

Notes
-----
This assumes that the modes are arranged (with fixed s value) as

    [
        Y(s, ℓ, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function Ysize(ℓₘᵢₙ, ℓₘₐₓ)
    # from sympy import symbols, summation, horner
    # from sympy.printing.pycode import pycode
    # ℓ,m,ℓₘᵢₙ,ℓₘₐₓ = symbols('ℓ,m,ℓₘᵢₙ,ℓₘₐₓ', integer=True)
    #
    # def nice(expr)
    #     return horner(expr.expand().simplify(), (m′ₘₐₓ, ℓₘᵢₙ, ℓₘₐₓ))
    #
    # Ysize = horner(
    #     summation(summation(1, (m, -ℓ, ℓ)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ)).expand().simplify(),
    #     (ℓₘᵢₙ, ℓₘₐₓ)
    # )
    return ℓₘₐₓ * (ℓₘₐₓ + 2) - ℓₘᵢₙ^2 + 1
end

"""
    deduce_limits(ysize, [ℓmin])

Deduce the value of `(ℓmin, ℓmax)` that produces ``Y`` arrays of the given size.

If `ℓmin` is not given, it is assumed to be 0.  If it is set to `nothing`, the
smallest possible value of `ℓmin` will be used.  However, note that this is not
a well-posed problem; multiple combinations of `(ℓmin, ℓmax)` can give rise to
``Y`` arrays of the same size.

See also [`Ysize`](@ref)
"""
function deduce_limits(ysize, ℓmin=0)
    if ℓmin === nothing
        maxℓmax = (ysize-1) ÷ 2
        ℓmin_range = 0:maxℓmax
    else
        ℓmin_range = [ℓmin]
    end
    for ℓmin in ℓmin_range
        ℓmax = round(Int, √(ysize + ℓmin^2) - 1, RoundDown)
        if ℓmax * (ℓmax + 2) - ℓmin^2 + 1 == ysize
            return (ℓmin, ℓmax)
        end
    end
    error("Input size $ysize does not correspond to a possible array of modes with ℓmin=$ℓmin")
end

"""
    Yrange(ℓₘᵢₙ, ℓₘₐₓ)

Create an array of (ℓ, m) indices as in Y array

Parameters
----------
ℓₘᵢₙ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ
ℓₘₐₓ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ

Returns
-------
i : int
    Total size of array of mode weights arranged as described below

See Also
--------
Ysize : Total size of array of mode weights
Yindex : Index of a particular element of the mode weight array

Notes
-----
This assumes that the modes are arranged (with fixed s value) as

    [
        Y(s, ℓ, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function Yrange(ℓₘᵢₙ, ℓₘₐₓ)
    r = zeros(typeof(ℓₘᵢₙ), (Ysize(ℓₘᵢₙ, ℓₘₐₓ), 2))
    i = 1
    for ℓ in ℓₘᵢₙ:ℓₘₐₓ
        for m in -ℓ:ℓ
            r[i, 1] = ℓ
            r[i, 2] = m
            i += 1
        end
    end
    return r
end

"""
    Yindex(ℓ, m, ℓₘᵢₙ=0)

Compute index into array of mode weights

Parameters
----------
ℓ : int
    Integer satisfying ℓₘᵢₙ <= ℓ <= ℓₘₐₓ
m : int
    Integer satisfying -ℓ <= m <= ℓ
ℓₘᵢₙ : int, optional
    Integer satisfying 0 <= ℓₘᵢₙ.  Defaults to 0.

Returns
-------
i : int
    Index of a particular element of the mode-weight array as described below

See Also
--------
Ysize : Total size of array of mode weights
Yrange : Array of (ℓ, m) indices corresponding to this array

Notes
-----
This assumes that the modes are arranged (with fixed s value) as

    [
        Y(s, ℓ, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
@inline function Yindex(ℓ, m, ℓₘᵢₙ=0)
    # from sympy import symbols, summation, horner
    # from sympy.printing.pycode import pycode
    # ℓ,m,m′,ℓₘᵢₙ, = symbols('ℓ,m,m′,ℓₘᵢₙ', integer=True)
    #
    # def nice(expr)
    #     return horner(expr.expand().simplify(), (ℓₘᵢₙ, ℓ, m))
    #
    # Yindex = horner(
    #     (Ysize.subs(ℓₘₐₓ, ℓ-1) + summation(1, (m′, -ℓ, m)) - 1).expand().simplify(),
    #     (ℓₘₐₓ, ℓ, m)
    # )
    return ℓ*(ℓ + 1) - ℓₘᵢₙ^2 + m + 1
end

"""
    theta_phi(Nθ, Nϕ, [T=Float64])

Construct (theta, phi) grid in `spinsfast` order.

Note that this order is different from the one assumed by this package;
use [`phi_theta`](@ref) for the opposite ordering.

"""
function theta_phi(nθ, nϕ, T=Float64)
    [[θ, ϕ][i] for θ in range(0, T(π), length=nθ), ϕ in range(0, 2*T(π), length=nϕ+1)[begin:end-1], i in 1:2]
end

"""
    phi_theta(Nϕ, Nθ, [T=Float64])

Construct (phi, theta) grid in order expected by this package.

See also [`theta_phi`](@ref) for the opposite ordering.

"""
function phi_theta(nϕ, nθ, T=Float64)
    [[ϕ, θ][i] for ϕ in range(0, 2*T(π), length=nϕ+1)[begin:end-1], θ in range(0, T(π), length=nθ), i in 1:2]
end


using Statistics
using PolynomialBasis
using ImplicitDomainQuadrature
using BenchmarkTools
# using Revise
using CutCellDG
include("../../../test/useful_routines.jl")
include("transformation-elasticity-solver.jl")

TES = TransformationElasticitySolver

polyorder = 2
nelmts = 15
penaltyfactor = 1e3

################################################################################
width = 1.0              # mm
interfacecenter = [0.5, 0.5]
interfaceradius = 0.4
outerradius = 1.2
################################################################################

################################################################################
K1, K2 = 247.0, 192.0    # GPa
mu1, mu2 = 126.0, 87.0   # GPa
lambda1 = lame_lambda(K1, mu1)
lambda2 = lame_lambda(K2, mu2)
stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)
rho1 = 3.93e3           # Kg/m^3
rho2 = 3.68e3           # Kg/m^3
V01 = 1.0 / rho1
V02 = 1.0 / rho2
theta0 = -0.067
ΔG0Jmol = -14351.0  # J/mol
molarmass = 0.147
ΔG0 = ΔG0Jmol / molarmass
transfstress =
    CutCellDG.plane_strain_transformation_stress(lambda1, mu1, theta0)
################################################################################

meshwidth = [width, width]
numqp = required_quadrature_order(polyorder) + 2

levelsetbasis = HermiteTensorProductBasis(2)
quad = tensor_product_quadrature(2, 4)
elasticitybasis = LagrangeTensorProductBasis(2, polyorder)

basispts = interpolation_points(elasticitybasis)
dim, numpts = size(interpolation_points(levelsetbasis))
cgmesh = CutCellDG.CGMesh([0.0, 0.0], meshwidth, [nelmts, nelmts], numpts)
dgmesh = CutCellDG.DGMesh([0.0, 0.0], meshwidth, [nelmts, nelmts], basispts)

@btime levelset = CutCellDG.LevelSet(
    x -> -circle_distance_function(x, interfacecenter, interfaceradius)[1],
    cgmesh,
    levelsetbasis,
    quad,
)

elementsize = CutCellDG.element_size(cgmesh)
minelmtsize = minimum(elementsize)
maxelmtsize = maximum(elementsize)
penalty = penaltyfactor / minelmtsize * 0.5 * (lambda1 + mu1 + lambda2 + mu2)
tol = minelmtsize^(polyorder + 1)
boundingradius = 1.5 * maxelmtsize

cutmesh = CutCellDG.CutMesh(dgmesh, levelset)

cellquads = CutCellDG.CellQuadratures(cutmesh,levelset,numqp)
@btime interfacequads = CutCellDG.InterfaceQuadratures(cutmesh, levelset, numqp)
@btime facequads = CutCellDG.FaceQuadratures(cutmesh, levelset, numqp)


function Damper(; name, b = 1.0)
    # 2 pin component
    @named oneport = OnePort()
    @unpack u, f = oneport
    ps = @parameters b=b
    eqs = [
        u ~ f / b
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

function Mass(; name, m = 1.0)
    @named oneport = OnePort()
    @unpack u, f = oneport
    ps = @parameters m=m
    D = Differential(t)
    eqs = [
        D(u) ~ f / m
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

function Spring(; name, k = 1.0)
    @named oneport = OnePort()
    @unpack u, f = oneport
    ps = @parameters k=k
    D = Differential(t)
    eqs = [
        D(f) ~ k * u
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end


using RungeKutta.Tableaus: get_radau_1_nodes, get_radau_1_weights, get_radau_1_coefficients,
                           get_radau_2_nodes, get_radau_2_weights, get_radau_2_coefficients

@testset "$(rpad("Radau Tableaus",80))" begin

    @test_throws ErrorException get_radau_1_nodes(1)
    @test_throws ErrorException get_radau_1_weights(1)
    @test_throws ErrorException get_radau_1_coefficients(1)

    @test_throws ErrorException get_radau_2_nodes(1)
    @test_throws ErrorException get_radau_2_weights(1)
    @test_throws ErrorException get_radau_2_coefficients(1)

    function _TableauRadauIA2(T=Float64)
        a = [[1//4  -1//4  ]
             [1//4   5//12 ]]
        b = [1//4, 3//4]
        c = [0,    2//3]

        Tableau{T}(:RadauIA2, 3, a, b, c)
    end

    function _TableauRadauIA3(T=Float64)
        a = [
                [ 1//9     (- 1 -    √6)/18    (- 1 +    √6)/18  ]
                [ 1//9     ( 88 +  7*√6)/360   ( 88 - 43*√6)/360 ]
                [ 1//9     ( 88 + 43*√6)/360   ( 88 -  7*√6)/360 ]
            ]
        b = [1/9,         (16+√6)/36,  (16-√6)/36 ]
        c = [0,           ( 6-√6)/10,  ( 6+√6)/10 ]

        Tableau{T}(:RadauIA3, 5, a, b, c)
    end

    function _TableauRadauIB2(T=Float64)
        a = [[1//8  -1//8  ]
             [7//24  3//8  ]]
        b = [1//4, 3//4]
        c = [0,    2//3]

        Tableau{T}(:RadauIB2, 3, a, b, c)
    end

    function _TableauRadauIB3(T=Float64)
        a = [
                [            1//18    (- 1 -     √6)/36    (- 1 +     √6)/36   ]
                [ ( 52 +  3*√6)/450   ( 16 +     √6)/72    (472 - 217*√6)/1800 ]
                [ ( 52 -  3*√6)/450   (472 + 217*√6)/1800  ( 16 -     √6)/72   ]
            ]
        b = [1/9,         (16+√6)/36,  (16-√6)/36 ]
        c = [0,           ( 6-√6)/10,  ( 6+√6)/10 ]

        Tableau{T}(:RadauIB3, 5, a, b, c)
    end

    function _TableauRadauIIA2(T=Float64)
        a = [[5//12  -1//12]
             [3//4    1//4 ]]
        b = [3//4, 1//4]
        c = [1//3, 1//1]

        Tableau{T}(:RadauIIA2, 3, a, b, c)
    end

    function _TableauRadauIIA3(T=Float64)
        a = [
                [ 11/45 -  7*√6/360    37/225-169*√6/1800   -2/225+√6/75 ]
                [ 37/225+169*√6/1800   11/45 +  7*√6/360    -2/225-√6/75 ]
                [  4/9  -    √6/36      4/9  +    √6/36      1/9         ]
            ]
        b = [4/9-√6/36,   4/9+√6/36,  1/9 ]
        c = [2/5-√6/10,   2/5+√6/10,  1   ]

        Tableau{T}(:RadauIIA3, 5, a, b, c)
    end

    function _TableauRadauIIB2(T=Float64)
        a = [[3//8   -1//24]
             [7//8    1//8 ]]
        b = [3//4, 1//4]
        c = [1//3, 1//1]

        Tableau{T}(:RadauIIB2, 3, a, b, c)
    end

    function _TableauRadauIIB3(T=Float64)
        a = [
                [( 16 -     √6)/72    (328 - 167*√6)/1800   (-2 + 3*√6)/450 ]
                [(328 + 167*√6)/1800  ( 16 +     √6)/72     (-2 - 3*√6)/450 ]
                [( 85 -  10*√6)/180   ( 85 +  10*√6)/180              1/18  ]
            ]
        b = [4/9-√6/36,   4/9+√6/36,  1/9 ]
        c = [2/5-√6/10,   2/5+√6/10,  1   ]

        Tableau{T}(:RadauIIB3, 5, a, b, c)
    end

    @test_throws ErrorException TableauRadauIA(1)
    @test_throws ErrorException TableauRadauIIA(1)

    @test TableauRadauIA(2)  ≈ _TableauRadauIA2()
    @test TableauRadauIA(3)  ≈ _TableauRadauIA3()

    @test TableauRadauIB(2)  ≈ _TableauRadauIB2()
    @test TableauRadauIB(3)  ≈ _TableauRadauIB3()

    @test TableauRadauIIA(2) ≈ _TableauRadauIIA2()
    @test TableauRadauIIA(3) ≈ _TableauRadauIIA3()

    @test TableauRadauIIB(2) ≈ _TableauRadauIIB2()
    @test TableauRadauIIB(3) ≈ _TableauRadauIIB3()

end


using EvolutionaryPlus
using Test

## Evolutionary Algorithm
println("=====Evolutionary Algorithm=====")
function f1(pop::DE.Population)
    for i in 1:size(pop.x, 1)
        pop.fit[i] = sum(pop.x[i, j]^2 for j in 1:size(pop.x, 2))
    end
end

DE.DE_evolution(f1, 50, 50, 1000)

## Covariance Matrix Adaptation Evolution Strategy
println("=====Covariance Matrix Adaptation Evolution Strategy=====")
function f2(p::CMA_ES.Population)
    for i in 1:size(p.x, 1)
        p.fit[i] = sum((-1 + 2 * p.x[i, j])^2 for j in 1:size(p.x, 2))
    end
end

CMA_ES.CMAES_evolution(f2, 50, 1000)

## Surrogate Model using Gaussian Processes
println("=====Surrogate Model using Gaussian Processes=====")
f(x) = sin(x)

tr_x = Array(0:0.5:10)
tr_y = f.(tr_x)

σ_f, l, L, α = GPs.model_training(tr_x, tr_y)

pr_x = Array(0:0.2:10)
pr_y = Array{Float64}(undef, length(pr_x), 2)
for i in 1:length(pr_x)
    pr_y[i, 1], pr_y[i, 2] = GPs.predict(σ_f, l, L, α, tr_x, pr_x[i])
end
te_y = f.(pr_x)

RMSE = sum((te_y[i] - pr_y[i, 1])^2 for i in length(te_y))^(1/2)

println("RMSE: ", RMSE)

#
# @testset "EvolutionaryPlus.jl" begin
#     # Write your own tests here.
#     f(x) = sin(x)
#
#     tr_x = Array(0:0.5:10)
#     tr_y = f.(tr_x)
#
#     σ_f, l, L, α = model_training(tr_x, tr_y)
#
#     pr_x = Array(0:0.2:10)
#     pr_y = Array{Float64}(undef, length(pr_x), 2)
#     for i in 1:length(pr_x)
#         pr_y[i, 1], pr_y[i, 2] = predict(σ_f, l, L, α, tr_x, pr_x[i])
#     end
#
#     @test abs(pr_y[1, 1] - tr_y[1]) < 1
# end
#


using AutoGrad
using Knet

include(Knet.dir("data", "housing.jl"))

predict(w,x) = w[1]*x .+ w[2]

loss(w,x,y) = mean(abs2, y - predict(w,x))

lossgradient = grad(loss)

function train(w, data; lr=0.1)
    for (x,y) in data
        dw = lossgradient(w,x,y)
        for i = 1:length(w) # 0 -> weights, 1 -> bias
            w[i] = w[i] - lr*dw[i]
        end
    end
    return w
end

x,y = housing()
w = Any[0.1 * randn(1,13), 0.0]
for i = 1:10
    train(w, [(x,y)])
    println(loss(w,x,y))
end


function linear_index(ind)
    j = round(Int, floor(-0.5 + 0.5*sqrt(1 + 8*(ind - 1))) + 2)
    i = round(Int, j .* (3 - j)/2 + ind -1)
    return i, j
end


# solve lyapunov equation (dH/dt) = B(t) H(t) + H(t) B(t)' - a(t), backward, conditional on Hend
function lyapunovpsdbackward_step(t, y, dt, P)
    B = Bridge.B(t - dt/2, P)
    ϕ = (I + 1/2*dt*B)\(I - 1/2*dt*B)
    ϕ *(y + 1/2*dt*a(t - dt, P))* ϕ' + 1/2*dt*Bridge.a(t, P)
end

lyapunovpsdbackward(t, P, Hend) = lyapunovpsdbackward!(samplepath(t, zero(Hend)), P, Hend)
function lyapunovpsdbackward!(H, P, Hend) # backward solve on grid tt, with final value Hend
    H.yy[end] = Hend
    t = H.tt
    for i in length(t)-1:-1:1
         H.yy[i] = lyapunovpsdbackward_step(t[i+1], H.yy[i+1], t[i+1] - t[i], P)
    end
    H
end


"F, the residual level Fct"
function residual( z::Array{T,1}, alpha::Array{T,1}, c::SLiSeFilters.Config ):T
    X = calcX(z, c);
    W = calcW(z, c);
    Z = calcZ(z, c);
    Y = calcY(z, c);
    theta = calctheta(z, c);

    XZa = (X - Z) * alpha
    WYa = (W - Y) * conj(alpha)

    return (2*real(alpha' * (XZa + WYa -2*theta)) .+ real(caseh_h( c )))[1]
end

"The gradient of F"
function Fgradient( z::Array{T,1}, alpha::Array{T,1}, c::SLiSeFilters.Config )
    X = calcX(z, c);
    Z = calcZ(z, c);
    theta = calctheta(z, c);
    Wc = calcW(conj(z), c);
    Yc = calcY(conj(z), c);

    da = 4*(alpha' * (X - Z) + transpose(alpha) * (Wc - Yc) - theta' );

    dYc = calcDeltaYc( z, c );
    dWc = calcDeltaWc( z, c );
    dX = calcDeltaX( z, c );
    dZ = calcDeltaZ( z, c );
    dTheta = calcDeltaTheta( z, c );

    dz = 4*(alpha' * (dX - dZ) + transpose(alpha) * (dWc - dYc) - dTheta') .* transpose(alpha); # elementwise multiplication

    # we conjugate as we want column vectors
    return dz',da'
end

# the approximation to the hessian required for levMar
function hessianApprox( z::Array{T,1}, a::Array{T,1}, c::SLiSeFilters.Config )

    ddX = calcddX( z, c );
    ddZ = calcddZ( z, c );

    X = calcX( z, c);
    Z = calcZ( z, c);

    dX = calcDeltaX( z, c );
    dZ = calcDeltaZ( z, c );

    H2 = 2*[  (diagm(a)' * ( ddX - ddZ) * diagm(a)) (diagm(a)' * (dX - dZ)')  ;
              ( ( dX -  dZ ) * diagm(a) )          (X - Z)                  ];

    dz, da = Fgradient( z, a, c );

    H = H2;
    d = (0.5.*[dz;da]);
    d = (d[:])

    return H, d;
end


# Adapted from:
# SOSDEMO7 --- Chebyshev polynomials
# Section 3.7 of SOSTOOLS User's Manual

@testset "SOSDEMO7 with $(factory.constructor)" for factory in sdp_factories
    isscs(factory) && continue
    ndeg = 8   # Degree of Chebyshev polynomial

    @polyvar x

    Z = monomials((x,), 0:ndeg-1)

    m = SOSModel(factory)

    @variable m γ
    @variable m p1 Poly(Z)

    p = p1 + γ * x^ndeg # the leading coeff of p is γ

    dom = @set x >= -1 && x <= 1
    @constraint(m, p <= 1, domain = dom)
    @constraint(m, p >= -1, domain = dom)

    @objective m Max γ

    JuMP.optimize!(m)

    @test JuMP.primal_status(m) == MOI.FeasiblePoint

    @test isapprox(JuMP.value(p), 128x^8 - 256x^6 + 160x^4 - 32x^2 + 1, ztol=1e-7, atol=1e-7)
    @test isapprox(JuMP.value(γ), 128)
end


using InteratomicPotentials
using AtomsBase
using UnitfulAtomic
using Unitful
using StaticArrays
using LinearAlgebra
using Statistics

# Load data
include("LJCluster/load_data.jl")
systems, energies, forces = load_data(;num_entries = 100);

# Split into training, testing
train_systems, train_energies, train_forces = systems[1:50], energies[1:50], forces[1:50];
test_systems, test_energies, test_forces = systems[51:end], energies[51:end], forces[51:end];

## Create RPI Basis (2body, 8 polynomial degree)
n_body = 2
max_deg = 8
r0 = 1.0
rcutoff = 4.0
wL = 1.0
csp = 1.0
rpi_params = RPIParams([:Ar], n_body, max_deg, wL, csp, r0, rcutoff)

## Calculate descriptors
train_descriptors = [evaluate_basis(sys, rpi_params) for sys in train_systems]
test_descriptors = [evaluate_basis(sys, rpi_params) for sys in test_systems]
mean_d = mean(descriptors)
cov_d = cov(descriptors)
inv_cov_d = inv(cov_d)
mah_dist_train = [ sqrt( 0.5*(td - mean_d)' * inv_cov_d * (td - mean_d) ) for td in train_descriptors ]
mah_dist_test = [ sqrt( 0.5*(td - mean_d)' * inv_cov_d * (td - mean_d) ) for td in test_descriptors ]

# Estimate β
β = hcat(train_descriptors...)' \ train_energies
rpi = RPI(β, rpi_params)

potential_energy(test_systems[1], rpi)


@testset "cplx2real - Exponent" begin
    z1 = MPC.Variable("z1", Complex)
    z2 = MPC.Variable("z2", Complex)
    x = MPC.Variable("x", Real)
    b = MPC.Variable("b", Bool)

    expo_cplx = Exponent(SortedDict(z1=>Degree(1,0),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))

    z1_re = MPC.Variable("z1_Re", Real)
    z1_im = MPC.Variable("z1_Im", Real)

    expo_real = Exponent(SortedDict(z1_re=>Degree(1,0),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))
    expo_imag = Exponent(SortedDict(z1_im=>Degree(1,0),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))

    preal, pimag = cplx2real(expo_cplx, 1.)
    @test preal == Polynomial(SortedDict{Exponent, Number}(expo_real=>1.0))
    @test pimag == Polynomial(SortedDict{Exponent, Number}(expo_imag=>1.0))

    expo_cplx = Exponent(SortedDict(z1=>Degree(0,1),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))

    preal, pimag = cplx2real(expo_cplx, 1.)
    @test preal == Polynomial(SortedDict{Exponent, Number}(expo_real=>1.0))
    @test pimag == Polynomial(SortedDict{Exponent, Number}(expo_imag=>-1.0))
end

@testset "cplx2real - Polynomial" begin
    z1 = MPC.Variable("z1", Complex)
    z2 = MPC.Variable("z2", Complex)
    x = MPC.Variable("x", Real)
    b = MPC.Variable("b", Bool)

    expo_cplx = Exponent(SortedDict(z1=>Degree(1,0),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))
    p = Polynomial(SortedDict{Exponent, Number}(expo_cplx=>3+5im,
                                                Exponent()=>5im))

    # Corresponding real quantities
    z1_re = MPC.Variable("z1_Re", Real)
    z1_im = MPC.Variable("z1_Im", Real)
    expo_real = Exponent(SortedDict(z1_re=>Degree(1,0),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))
    expo_imag = Exponent(SortedDict(z1_im=>Degree(1,0),
                                    x=>Degree(4,0),
                                    b=>Degree(1,0)))

    preal, pimag = cplx2real(p)
    @test preal == Polynomial(SortedDict{Exponent,Number}(expo_real=>3.0,
                                                          expo_imag=>-5.0))
    @test pimag == Polynomial(SortedDict{Exponent,Number}(Exponent()=>5,
                                                          expo_imag=>3.0,
                                                          expo_real=>5.0))

end

@testset "cplx2real - Point" begin
    z1 = MPC.Variable("z1", Complex)
    z2 = MPC.Variable("z2", Complex)
    x = MPC.Variable("x", Real)
    b = MPC.Variable("b", Bool)

    pt_c = Point(SortedDict(z1=>1+2im,
                            z2=>3im,
                            x=>3.5,
                            b=>1))

    z1_Re = MPC.Variable("z1_Re", Real)
    z1_Im = MPC.Variable("z1_Im", Real)
    z2_Im = MPC.Variable("z2_Im", Real)

    pt_re = Point(SortedDict(z1_Re=>1,
                             z1_Im=>2,
                             z2_Im=>3,
                             x=>3.5,
                             b=>1))

    @test cplx2real(pt_c) == pt_re

    @test real2cplx(pt_re) == pt_c
end


"""
         simulate(p    :: Drip;
                  T    :: Integer = 500,                    # Optional: length of simulation
                  burn :: Integer = 100,                    # Optional: length of initial burn (in addition to T)
                  N    :: Union{Int4,Nothing} = nothing,    # Optional: number of simulated agents (returns the average beliefs of a large set of agents by default)
                  seed :: Union{Integer,Nothing} = nothing  # Optional: seed number for fundamental shocks
                  ) -> Path
Returns a `Path` structure with a simulated path of the fundamental (`x`), beliefs (`x_hat`)
    and actions (`a`)
    **under the steady state information structure**. In particular, if `n` is the
    dimension of `x` and `m` is the dimension of `a`, then

* `x` has dimension `n*T` where `x(:,t)` is the simulated value of `x` at time `t`.
* `x_hat` has dimension `n*T*N` where `x_hat(:,i,t)` is the simulated value of `x_hat` of agent `i` at time `t`.
* `a` has dimension `m*T*N` where `a(:,i,t)` is the the simulated value of `x_hat` of agent `i` at time `t`.
"""
function simulate(p    :: Drip;
                  T    :: Integer = 500,                    # Optional: length of simulation
                  burn :: Integer = 100,                    # Optional: length of initial burn (in addition to T)
                  N    :: Union{Integer,Nothing} = nothing, # Optional: number of simulated agents (returns the average beliefs of a large set of agents by default)
                  seed :: Union{Integer,Nothing} = nothing  # Optional: seed number for fundamental shocks
                  )
    # create flag for if N is nothing and change N to 1
    average_belief = (N == nothing)
    if average_belief; N = 1; end
    # get dimensions
    (n,m)    = size(p.H)
    (_,k)    = size(p.Q)
    # set seed for fundamental shocks
    seed != nothing ? Random.seed!(seed) : Random.seed!()
    x_Shocks = randn(k,burn+T);

    Random.seed!()
    a_Shocks = average_belief ? zeros(m,burn+T,N) : randn(m,burn+T,N);
    x        = zeros(n,burn+T);
    x_hat    = zeros(n,burn+T,N);
    a        = zeros(m,burn+T,N);

    SqΣz = average_belief ? 0 : sqrt(p.ss.Σ_z);

    for t in 2:T+burn
        x[:,t] = p.A*x[:,t-1]+p.Q*x_Shocks[:,t];
        for i in 1:N
            x_hat[:,t,i] = p.A*x_hat[:,t-1,i]+(p.ss.K*p.ss.Y')*(x[:,t]-p.A*x_hat[:,t-1,i])+p.ss.K*SqΣz*a_Shocks[:,t,i];
            a[:,t,i]    .= p.H'*x_hat[:,t,i];
        end
    end
    x     = x[:,burn+1:T+burn];
    x_hat = average_belief ? x_hat[:,burn+1:T+burn] : x_hat[:,burn+1:T+burn,:];
    a = average_belief ? a[:,burn+1:T+burn] : a[:,burn+1:T+burn,:];
    return(Path(T,x,x_hat,a))
end


#

#
# Horner initially stores a list of coefficients of a polynomial, a evaluation point,
# and optionally a corresponding list of absolute errors of the coefficients.
# On demand it modifies the values in order to obtain the values of the selected derivatives
# at the given point and the corresponding error value.
#

mutable struct Horner{S<:Number,T<:Number,U<:AbstractFloat,V<:AbstractFloat}
    x::S							# evaluation spot
    coeff::AbstractVector{T}		# polynomial coefficients
    devx::U							# absolute error estimatin for x
    devcoeff::AbstractVector{V}		# absolute error estimatins for coefficients
    valid::Int						# evaluation status
end

# Create Horner table for evaluation at point x.
# Provide polynomial coefficients and optionally absolute errors for x and the coefficients.
function Horner(x::S, coeff::AbstractVector{T}, devx::U, devcoeff::AbstractVector{V}) where {S<:Number,T<:Number,U<:AbstractFloat,V<:AbstractFloat}

    if devcoeff == []
        devcoeff = abs(coeff) * eps(promote_type(T,Float64))
    else
        devcoeff = copy(devcoeff)
    end
    if devx < 0
        devx = abs(x) * eps(real(S))
    end
    if length(coeff) != length(devcoeff) throw(ArgumentError("devcoeff length wrong")) end
    R = promote_type(S,T)
	Horner(x, map(R,coeff), devx, devcoeff, -1)
end

# Evaluate Horner table up to and including m deviations.
# return the m'th deviation of the polynomial and the error estimation
# If m >= degree, all deviations ot the polynomial have been calculated at x.
# Further calls just return the cached results.
function call(h::Horner{S,T,U,V}, m::Int) where {S,T,U,V}
#function (h::Horner{S,T,U,V})(m::Int)
    if m < 0 throw(ArgumentError("index $m is not >= 0")) end
	n = length(h.coeff) - 1
    if m > h.valid
        if h.x != 0
            ep = eps(real(T))
            b, e = h.coeff, h.devcoeff
            x, tol = h.x, h.devx
            ax = abs(x)
            for k = max(h.valid+1,0):min(m, n-1)
                for j = 1:n-k
                    bjp = b[j+1] + b[j] * x
                    b[j+1] = bjp
                    e[j+1] = hypot(hypot(abs(e[j+1]), abs(b[j])*tol +ax*e[j]), abs(bjp)*ep)
                end
            end
        end
        h.valid = m
    end
    m <= n ? (h.coeff[n+1-m], h.devcoeff[n+1-m]) : (zero(T), zero(V))
end

#
# evaluate polynomial f and all derivatives at given approximate root-values z
#
function horner_analysis(z::AbstractVector{S}, f::AbstractVector{T}) where {S<:Number,T<:Number}
    thresh = 1000.0
    m = length(z)
    n = length(z)-1
    if m > length(f)-1
        throw(ArgumentError("polynomial of degree $n with $m zeros"))
    end
	h = Array{Horner}(m)
    for i = 1:m
        h[i] = Horner(z[i], f, eps(abs(z[i])) * 100, map(eps, f) * 100)
    end
	for i = 1:m
        println("root $i: $(z[i])")
        mult = 0
        search = true
		for k = 0:n
            hv, he = h[i](k)
            deviation = abs(hv / (he + eps(zero(real(S)))))
            @printf("%d %d %6.2g %8.8g %8.8g\n", i, k, deviation, abs(hv), he)
            if search && deviation < thresh
                mult = k + 1
            else
                search = false
            end
        end
        println("root $i: $(z[i]) estimated multiplicity: $mult")
    end
end