using RCall, JLD

R"
library(gridExtra)
library(cowplot)
library(tidyverse)
library(latex2exp)
library(viridis)
"

DF =load("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_6_6.jld")
grid_vec= DF["grid_vec"]
betas = first.(grid_vec)
sigmas = last.(grid_vec)

R"
range01 <- function(x){(x-min(x))/(max(x)-min(x))}
#df_r$beta_scaled= range01(df_r$beta_log)
#df_r$color = rgb(df_r$beta_scaled, df_r$sigma,0.5,1)

#betas = exp(seq(log(1), log(200), length.out=10))
#sigmas = seq(0.05, 0.99, length.out = 10)

betas = $betas
sigmas = $sigmas

library(tidyverse)

to_plot = expand_grid(betas, sigmas) %>%
  mutate(col = rgb(red = range01(log(betas)), green = sigmas, blue = 0.5, alpha = 1))

pt = to_plot %>%
  ggplot(aes(x = betas, y = sigmas)) +
  geom_point(aes(colour = col),size=4) + scale_x_log10() + annotation_logticks(base = 10)+ scale_color_identity()
 pt = pt +  geom_line(data =to_plot, aes(x=betas, y = sigmas, group=betas),linetype = 'longdash', alpha = 0.5,size=0.3)+
   geom_line(data =to_plot, aes(x=betas, y = sigmas, group=sigmas),linetype = 'dotted', alpha = 0.5,size=0.3)+ theme_classic()+
   theme(plot.title = element_text(hjust = 0.5,size = 15), axis.text = element_text(size=15), axis.title=element_text(size=15),legend.position='none')+
   ylab(TeX(sprintf('$\\sigma$')))+ xlab(TeX(sprintf('$\\tau$')))"

 R"
 pdf('/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/figures/figure_6_6.pdf')
 plot(pt)
 dev.off()"


using Plots
using Random

αs = [0.0, 0.9, 0.98]
n = 100
p = plot()

function getpricepath(α, n)
    x = zeros(n + 1)
    x[1] = 0.0
    for t in 1:n
        x[t+1] = α * x[t] + rand()
    end
    x
end

pricepaths = [getpricepath(α, n) for α in αs]

for (idx, pricepath) in enumerate(pricepaths)
    plot!(p, pricepath, title="Random price paths", label="alpha = $(αs[idx])", xlabel="Timestep", ylabel="Price")
end

p


using DesignStructureMatrix
using Test

@testset "Clustering" begin
 DSM =[ 0 1 0 0 1 1 0;
    0 0 0 1 0 0 1;
    0 1 0 1 0 0 1;
    0 1 1 0 1 0 1;
    0 0 0 1 0 1 0;
    1 0 0 0 1 0 0;
    0 1 1 1 0 0 0];

    label = ["A","B","C","D","E","F","G"];

    A,labelA = Clustering(DSM,label)

        @test A[1,1] == 0
        @test A[7,7] == 0

end

@testset "Sequencing" begin

         DSM = [0 0 0 0 0 0 0 0 0 0 0 0 0;
                0 0 0 1 0 0 0 0 0 0 0 0 0;
                1 0 0 1 1 0 0 1 0 1 0 0 1;
                1 0 0 0 0 0 0 0 0 0 0 0 0;
                1 1 1 1 0 0 1 1 0 1 1 0 1;
                0 1 0 0 0 0 0 0 1 0 0 0 0;
                0 0 1 1 1 0 0 1 0 1 0 0 1;
                0 1 0 1 0 0 0 0 0 1 0 0 0;
                0 0 1 1 0 0 0 0 0 1 0 0 1;
                0 1 0 1 0 0 0 1 0 0 0 0 0;
                0 0 1 1 1 0 0 1 0 1 0 0 1;
                0 0 0 0 0 0 0 0 0 0 0 0 1;
                0 1 1 1 1 0 1 1 1 1 1 1 0];

        label = ["1", "2", "3", "4", "5", "6", "7", "8",
                "9", "10", "11", "12", "13"];

        A,labelA = Sequencing(DSM,label)

    @test labelA == ["1","4","2","8","10","3","5","7","9","11","12","13","6"]

end

@testset "Plot" begin

        using Luxor

        DSM = [0 0 0 0 0 0 0 0 0 0 0 0 0;
               0 0 0 1 0 0 0 0 0 0 0 0 0;
               1 0 0 1 1 0 0 1 0 1 0 0 1;
               1 0 0 0 0 0 0 0 0 0 0 0 0;
               1 1 1 1 0 0 1 1 0 1 1 0 1;
               0 1 0 0 0 0 0 0 1 0 0 0 0;
               0 0 1 1 1 0 0 1 0 1 0 0 1;
               0 1 0 1 0 0 0 0 0 1 0 0 0;
               0 0 1 1 0 0 0 0 0 1 0 0 1;
               0 1 0 1 0 0 0 1 0 0 0 0 0;
               0 0 1 1 1 0 0 1 0 1 0 0 1;
               0 0 0 0 0 0 0 0 0 0 0 0 1;
               0 1 1 1 1 0 1 1 1 1 1 1 0];

       label = ["1", "2", "3", "4", "5", "6", "7", "8",
               "9", "10", "11", "12", "13"];

       p = plotDSM(DSM,label)

   @test  typeof(p) == Luxor.Drawing

end


"""
    embedding(small::AbstractMesh, big::AbstractMesh) -> S::AbstractSparseArray
"""
function embedding(small::AbstractMesh, big::AbstractMesh)

    @assert dimension(small) == 1
    @assert dimension(big)   == 1

    T = coordtype(small)
    tol = sqrt(eps(T))

    # store the edges of the big mesh in an octree
    ctrs_big = [cartesian(center(chart(big,v))) for v in cells(big)]
    tree_big = Octree(ctrs_big, zeros(size(ctrs_big)))

    rows = 1:numcells(small)
    cols = zeros(numcells(small))
    sgns = zeros(numcells(small))

    cells_big = collect(cells(big))
    for (i,c) in enumerate(cells(small))
        # V = vertices(small,c)
        V = chart(small, c).vertices
        ch = chart(small,c)
        x = cartesian(center(ch))
        pred(c,s) = fitsinbox(x, 0.0, c, s+tol)
        found = false
        for box in boxes(tree_big, pred)
            for j in box
                c′ = cells_big[j]
                # V′ = vertices(big,c′)
                V′ = chart(big,c′).vertices
                ch′ = chart(big, c′)
                if norm(V[1]-V′[1]) + norm(V[2]-V′[2]) < tol
                    cols[i] = j
                    sgns[i] = +1
                    found = true
                    break
                elseif norm(V[1]-V′[2]) + norm(V[2]-V′[1]) < tol
                    cols[i] = j
                    sgns[i] = -1
                    found = true
                    break
                end
            end
            found && break
        end
    end

    @assert !any(cols .== 0)
    @assert !any(sgns .== 0)
    sparse(rows, cols, sgns, numcells(small), numcells(big))
end


using ProjectiveDictionaryPairLearning
using MAT, Base.Test

###### Accelerating the normalization with testing
exampleMat = [1.0 2 3; 4 5 6; 7 8 9; 3 4 5]

normcol_equal(matin) = matin ./ repmat(sqrt(sum(matin .^ 2, 1) + eps()), size(matin, 1), 1)

tic()
exampleResult = normcol_equal(exampleMat)
normcol_equalTime = toq()

# tic()
# normalize!(exampleMat, 2 ,1)
# normalizeTime = toq()

# @test_approx_eq(exampleMat, exampleResult)
# @test normalizeTime < normcol_equalTime

# println("normcol_equalTime : $normcol_equalTime s")
# println("normalizeTime     : $normalizeTime s")

###### Accelerating the Class-specific reconstruction error calculation
# load TtData, DictMat, EncoderMat, DictSize
# prepared using
    # matwrite("class-specific.mat", @compat Dict(
    #     "TtData" => TtData,
    #     "DictMat" => DictMat,
    #     "EncoderMat" => EncoderMat,
    #     "DictSize" => DictSize
    # ))
matFileName = "class-specific.mat"
data = matread(joinpath(dirname(@__FILE__), matFileName))
TtData, DictMat, EncoderMat, DictSize = data["TtData"], data["DictMat"], data["EncoderMat"], data["DictSize"]

# prepare
ClassNum    = length(DictMat)
PredictCoef = EncoderMat * TtData

ErrorMATLAB = zeros(ClassNum, size(TtData, 2))
ErrorSubtract = zeros(ClassNum, size(TtData, 2))

tic()
# Class-specific reconstruction error calculation
for i = 1:ClassNum
    ErrorMATLAB[i, :]= sum((DictMat[i] * PredictCoef[(i - 1) * DictSize + 1 : i * DictSize, :]- TtData) .^ 2, 1)
end
ErrorMATLABTime = toq()

tic()
# Class-specific reconstruction error calculation
for i=1:ClassNum
    @inbounds tempMat = DictMat[i] * PredictCoef[(i-1)*DictSize+1:i*DictSize, :]
    tempMat -= TtData
    @inbounds ErrorSubtract[i,:] = sum(abs2(tempMat), 1)
end
ErrorSubtractTime = toq()

@test_approx_eq(ErrorMATLAB, ErrorSubtract)
@test ErrorSubtractTime < ErrorMATLABTime

println("Finished performance tests!")


using Echidna
using Test

# TEST ZDT1
x = [.1 for i in 1:30]
@show obj = ZDT1(x)
@test obj[1] == 0.1
@test obj[2] == 1.464110105645933

nobjs = 2
nvars = 30
zdt1_problem = Problem(
    2,
    [false for i in 1:nobjs],
    nvars,
    [MOGA_Real(0.0, 1.0) for i in 1:30],
    ZDT1
)
sol = Solution(
    zdt1_problem,
    [.1 for i in 1:nvars],
    false,
    [0.0, 0.0],
    0.0,
    0
)
# @show sol.objectives = ZDT1(sol)
@show evaluate!(sol)
@test sol.objectives[1] == 0.1
@test sol.objectives[2] == 1.464110105645933
@test sol.evaluated

#### TEST WITH TOURNAMENT
pop_size = 100
pop = [random_candidate(zdt1_problem) for i in 1:pop_size]
for i in 1:pop_size
    evaluate!(pop[i])
end


#=
Julia implementation of the SPIKE distance, from the description of the Scholarpedia
http://www.scholarpedia.org/article/SPIKE-distance and the paper
Kreuz et al. (2013), [Monitoring spike train synchrony](https://doi.org/10.1152/jn.00873.2012)

Implementation by Geoge Datseris.
=#
export SPIKE_distance_profile, SPIKE_distance

"""
    SPIKE_distance_profile(y1, y2) -> t, S(t)
Calculate the SPIKE distance profile function ``S(t)``, Eq. (19) of [1], given input
spike trains `y1, y2`. Return `t, S(t)`.

The keywords `t0 = min(y1[1], y2[1]) - 1` and
`tf = max(y1[end], y2[end]) + 1` add auxilary spikes to both trains
at these times, as demanded by the algorithm. (Auxilary spikes are only added
in case `y1, y2` do not start or end with the given `t0, tf`)

[1] : Kreuz et al. (2013), [Monitoring spike train synchrony](https://doi.org/10.1152/jn.00873.2012)
"""
function SPIKE_distance_profile(y1, y2;
        t0 = min(y1[1], y2[1]) - 1,
        tf = max(y1[end], y2[end]) + 1,
    )
    @assert issorted(y1) && issorted(y2)
    @assert y1 == unique(y1) && y2 == unique(y2)
    y1 = copy(y1); y2 = copy(y2)
    # Add auxilary spikes:
    y1[1] != t0 && pushfirst!(y1, t0)
    y2[1] != t0 && pushfirst!(y2, t0)
    y1[end] != tf && push!(y1, tf)
    y2[end] != tf && push!(y2, tf)
    tvec, indices = _initialize_SPIKE_data(y1, y2)
    # S are the values of S(t), for t ∈ tvec.
    # Its values are only stored exactly before and exactly on a spike
    # (as only these values are need for the full S(t)) function)
    S⁻ = zeros(Float64, length(tvec))
    S⁺ = zeros(Float64, length(tvec))
    i1, i2 = 1, 1 # indices of the previous spike for each train
    @inbounds for k in 2:length(tvec)-1
        t = tvec[k]
        # First compute S given at the time of the next spike, but before
        # being exactly ON the spike
        S⁻[k] = _compute_S(t, y1, y2, i1, i2)
        # Then find out which indices to increment (where is current spike from)
        if indices[k] == 3 # both trains share current spike
            i1 += 1; i2 += 1
        elseif indices[k] == 1 # current spike from train 1
            i1 += 1
        else # current spike from train 2
            i2 += 1
        end
        # Then compute S exactly ON the current spike
        S⁺[k] = _compute_S(t, y1, y2, i1, i2)
    end
    tvec = sort!(vcat(tvec, tvec)) # duplicate time vector
    S = zeros(length(tvec))
    S[1:2:end] .= S⁻
    S[2:2:end] .= S⁺
    return tvec, S
end

"""
Core computatation for  ``S(t)`` at time ``t``. Whether we are exactly
before or exactly on top of a spike is taken care of by `i1, i2`.
"""
function _compute_S(t, y1, y2, i1, i2)
    # Because of the way `_initialize_SPIKE_data` works,
    # this is guaranteed to be correct from the data preperation:
    tP1 = y1[i1]; tP2 = y2[i2]
    tF1 = y1[i1+1]; tF2 = y2[i2+1]
    xISI1 = tF1 - tP1;  xISI2 = tF2 - tP2
    xP1 = t - tP1; xP2 = t - tP2
    xF1 = tF1 - t; xF2 = tF2 - t
    ΔtP1 = minimum_Δt(tP1, y2, i2); ΔtF1 = minimum_Δt(tF1, y2, i2);
    ΔtP2 = minimum_Δt(tP2, y1, i1); ΔtF2 = minimum_Δt(tF2, y1, i1)
    S1 = (ΔtP1*xF1 + ΔtF1*xP1)/xISI1; S2 = (ΔtP2*xF2 + ΔtF2*xP2)/xISI2
    S = 2(S1*xISI2 + S2*xISI1)/(xISI1 + xISI2)^2
    return S
end

"""
Create smart containers for the time vector (that has ALL spikes sorted) as
well as the indices of where each spike comes from. This allows for exceptionally
simplified source code for the SPIKE distance profile ``S(t)``.
"""
function _initialize_SPIKE_data(y1::AbstractVector{T1}, y2::AbstractVector{T2}) where {T1, T2}
    T = promote_type(T1, T2)
    tvec = Vector{T}()
    indices = Vector{Int8}()
    i1 = i2 = 1 # indices of the spike trains
    while i1 ≤ length(y1) && i2 ≤ length(y2)
        v1, v2 = y1[i1], y2[i2]
        if v1 == v2
            push!(tvec, v1); push!(indices, 3)
            i1 += 1; i2 += 1;
        elseif v1 < v2
            push!(tvec, v1); push!(indices, 1)
            i1 += 1
        else
            push!(tvec, v2); push!(indices, 2)
            i2 += 1
        end
    end
    @assert tvec == sort!(unique(vcat(y1, y2)))
    return tvec, indices
end

"""
    minimum_Δt(t, y, i = 1) -> Δt
Return the minimum distance `Δt` between time point `t` and the spikes in `y`.
Start searching from index `i` of `y` and efficiently scan indices starting from
`i` and incrementally increasing and decreasing.
"""
function minimum_Δt(t::T, y, i = 1) where {T}
    dleft, dright = typemax(t), typemax(t)
    for j in i:-1:1 # scan from i until 1
        z = abs(t - y[j])
        z > dleft ? break : (dleft = z)
    end
    for j in i+1:1:length(y) # scan from i+1 until the end
        z = abs(t - y[j])
        z > dright ? break : (dright = z)
    end
    return min(dright, dleft)
end

"""
    SPIKE_distance(y1, y2) -> D_S
Calculate the SPIKE distance ``D_S``, which is the average of
[`SPIKE_distance_profile`](@ref).
"""
function SPIKE_distance(y1, y2; kwargs...)
    t, S = SPIKE_distance_profile(y1, y2; kwargs...)
    return trapezoid_integral(t, S)/(t[end] - t[1])
end

function trapezoid_integral(t, S::AbstractVector{T}) where {T}
    I = zero(T)
    for i in 2:length(t)
        @inbounds I += (S[i] + S[i-1])*(t[i] - t[i-1])
    end
    return T(0.5) * I
end


# NOTE: assemble weakly-singular kernels using quadgk. Should be slow, but this is
# useful to verify that integrals converge as they should.

function assemble_gk(iop;compress=Matrix)
    X,Y       = target_surface(iop), source_surface(iop)
    @timeit_debug "assemble dense part" begin
        out        = compress(iop)
    end
    @timeit_debug "compute near interaction list" begin
        dict_near = near_interaction_list(dofs(X),Y;atol=0.1)
    end
    @timeit_debug "compute singular correction" begin
        correction = singular_weights_gk(iop,dict_near)
    end
    axpy!(1,correction,out) # out <-- out + correction
end

function singular_weights_gk(iop::IntegralOperator,dict_near)
    X,Y       = target_surface(iop), source_surface(iop)
    T         = eltype(iop)
    Is    = Int[]
    Js    = Int[]
    Vs    = T[]
    # loop over all integration elements by types
    for E in keys(Y)
        iter      = ElementIterator(Y,E)
        qreg      = Y.etype2qrule[E]
        lag_basis = lagrange_basis(qnodes(qreg))
        @timeit_debug "singular weights" begin
            _singular_weights_gk!(Is,Js,Vs,iop,iter,dict_near,lag_basis,qreg)
        end
    end
    Sp = sparse(Is,Js,Vs,size(iop)...)
    return Sp
end

function _singular_weights_gk!(Is,Js,Vs,iop,iter,dict_near,lag_basis,qreg)
    yi = qnodes(qreg)
    K = kernel(iop)
    X = target_surface(iop)
    Y = source_surface(iop)
    E = eltype(iter)
    list_near = dict_near[E]
    # Over the lines below we loop over all element, find which target nodes
    # need to be regularized, and compute the regularized weights on the source
    # points for that target point.
    for (n,el) in enumerate(iter)               # loop over elements
        jglob = Y.elt2dof[E][:,n]
        for (i,jloc) in list_near[n]            # loop over near targets
            ys = yi[jloc]
            xdof   = dofs(X)[i]
            for (m,l) in enumerate(lag_basis)   # loop over lagrange basis
                val,_ = quadgk(0,ys[1],1;atol=1e-14) do ŷ
                    y    = el(ŷ)
                    jac  = jacobian(el,ŷ)
                    μ    = integration_measure(jac)
                    ydof = NystromDOF(y,-1.0,jac,-1,-1)
                    K(xdof,ydof)*μ*l(ŷ)
                end
                # push data to sparse matrix
                push!(Is,i)
                j = jglob[m]
                push!(Js,j)
                push!(Vs,val-iop[i,j])
            end
        end
    end
    return Is,Js,Vs
end


module DifferentialGeometry

using Match
using LinearAlgebra, Polynomials, Statistics
using MiniQhull

include("io.jl")
using .DataIO: read_obj, read_ply

export Mesh
export mesh

export Surface
export basis, pullback

export Manifold
export manifold, makefield

numdims(x) = sum(1 for d in 1:ndims(x) if size(x,d) > 1)

# ------------------------------------------------------------------------
# mesh type

"""
    struct Mesh{T <: Real}
        r   :: Array{T, 2}
        vₙ  :: Array{T, 2}
        tri :: Array{Int, 2}
    end

Store a 2 dimensional triangular mesh, embedded into arbitrary dimensions.
`r` and `vₙ` denote vertex positions and normals respectively.
`tri` denotes (non-oriented) triangular faces.
"""
struct Mesh{T <: Real}
    r   :: Array{T, 2}     # vertex positions
    vₙ  :: Array{T, 2}     # vertex normals
    tri :: Array{Int, 2}   # triangular faces
end

"""
    mesh(io::IO, type::Symbol=:obj)

Load a `Mesh` object from stream `io` formatted with `type`.
As of now, only `.obj` files are supported.
"""
function mesh(io::IO, type::Symbol=:obj)
    if type != :obj
        error("not implemented")
    end

    data = read_obj(io)

    return Mesh{eltype(data.verts[1])}(data.verts, data.norms, data.faces)
end

# ------------------------------------------------------------------------
# surface type

# internal functions
"""
    ellipse(r)

Fit an ellipse to 2D point cloud `r`.
`r` is assumed to be sized ``N \\times 2``.
"""
function ellipse(r)
    # constrained least squares (constraint: 4AC-B² = 1)
    D = hcat(r[:,1].*r[:,1], r[:,1].*r[:,2], r[:,2].*r[:,2], r[:,1], r[:,2], ones(size(r, 1),1))
    S = D'*D

    C      = zeros(6,6)
    C[1,3] = C[3,1] = 2
    C[2,2] = -1

    λ, ν = eigen(inv(S) * C)
    n    = argmax(λ)
    p    = real(ν[:,n])

    # equations taken from wikipedia
    A, B, C, D, E, F = p[1], p[2]/2, p[3], p[4]/2, p[5]/2, p[6]
    num = B^2 - A*C

    x₀ = (C*D-B*E)/num
    y₀ = (A*E-B*D)/num
    ϕ  = 0.5*atan(2*B/(A-C))

    up    = 2*(A*E^2+C*D^2+F*B^2-2*B*D*E-A*C*F)
    down1 = (B^2-A*C)*( (C-A)*√(1+4*B^2/((A-C)*(A-C)))-(C+A))
    down2 = (B^2-A*C)*( (A-C)*√(1+4*B^2/((A-C)*(A-C)))-(C+A))
    a, b  = √(abs(up/down1)), √(abs(up/down2))

    return [x₀, y₀, a, b, ϕ]
end

"""
    struct Surface{T <: Real}
        Θ  :: Array{T, 2}
        x  :: Array{T}
        L  :: T
        Λ  :: Array{Polynomial{T}}
        ∂Λ :: Array{Polynomial{T}}
    end

Store a representation of a 2D surface embedded into higher dimensional Euclidean space.
Fits the surface by:
  1. Partition the x axis so that each resultant interval has, on average, 20 points.
  2. Fit points within each partition to a 2D ellipse.
  3. Fit a polynomial function to each elliptical parameter over all partitions.
  4. Use polynomials to estimate tangent vectors.
`Θ` denotes the elliptical parameters fit per partition.
`x` denotes the input data.
`L` denotes the length along the x axis.
`Λ` denotes the polynomials for each elliptical parameter.
`∂Λ` denotes the derivative of polynomials for each elliptical parameter.
"""
struct Surface{T <: Real}
    Θ  :: Array{T, 2}
    x  :: Array{T}
    L  :: T
    Λ  :: Array{Polynomial{T}}
    ∂Λ :: Array{Polynomial{T}}

    # constructors
    Surface(Θ, x, L, Λ, ∂Λ) = new{eltype(x)}(Θ, x, L, Λ, ∂Λ)

    function Surface(x; order=12, debug=false)
        nbin = size(x,1)÷20
        xbin = LinRange(minimum(x[:,1]), maximum(x[:,1]), nbin+1)

        select(x, index) = hcat([x[index,d] for d in 2:size(x,2)]...)

        T  = eltype(x)
        x₀ = T[]
        Θ  = Array{T}[]
        for (lb, ub) in zip(xbin[1:end-1],xbin[2:end])
            ι = (x[:,1] .>= lb) .& (x[:,1] .<= ub)
            if sum(ι) < 6
                continue
            end

            push!(x₀, mean(x[ι,1]))
            push!(Θ,  ellipse(select(x, ι)))
        end

        Θ = hcat(Θ...)' # [x₀, y₀, a, b, ϕ]
        Λ = Array{Polynomial}(undef, size(Θ,2))
        L = (maximum(x₀) - minimum(x₀))./2
        for p in 1:size(Θ,2)
            M    = hcat([(x₀./L).^n for n in 0:order]...)
            α    = M \ Θ[:,p]
            Λ[p] = Polynomial(α)
            if debug
                χ = range(minimum(x₀), maximum(x₀), length=1000)
                Plots.scatter(x₀, Θ[:,p])
                Plots.plot!(χ, Λ[p].(χ./L), ylims=(minimum(Θ[:,p]),maximum(Θ[:,p])))
                gui()
                readline()
            end
        end

        ∂Λ = [derivative(Λ[i]) for i in 1:length(Λ)]

        new{T}(Θ, x₀, L, Λ, ∂Λ)
    end
end

"""
    function pullback(s::Surface, r)

Transform Euclidean point cloud `r` into a 2D cylindrical projection defined by `s`.
Cylindrical coordinates are [x,φ].
Assumes all points within `r` are distributed over the surface.
Assumes `r` is sized ``N \\times 3``.
"""
function pullback(s::Surface, r)
    z₀ = r[:,1] ./ s.L
    x₀ = s.Λ[1].(z₀)
    y₀ = s.Λ[2].(z₀)
    a  = s.Λ[3].(z₀)
    b  = s.Λ[4].(z₀)
    ϕ  = s.Λ[5].(z₀)

    v = Array{eltype(r),2}(undef, length(a), 2)
    for (i, p) in enumerate(ϕ)
        v[i,:] = [cos(p) +sin(p); -sin(p) cos(p)]*(r[i,2:3] .- [x₀[i], y₀[i]])
    end
    θ = mod2pi.(atan.(v[:,2]./b, v[:,1]./a) .+ π .- 1.65) .- π

    v[:,1] = z₀
    v[:,2] = θ

    return v
end

"""
    function basis(s::Surface, r)

Compute the tangent vectors ``\\hat{\\bm{e}}_\\phi, \\hat{\\bm{e}}_x`` associated to each point within `r`.
Assumes all points within `r` are distributed over the surface.
Assumes `r` is sized ``N \\times 3``.
"""
function basis(s::Surface, r; surface=false, normalized=true)
    x  = surface ? r : pullback(s, r)

    # function
    x₀ = s.Λ[1].(x[:,1])
    y₀ = s.Λ[2].(x[:,1])
    a  = s.Λ[3].(x[:,1])
    b  = s.Λ[4].(x[:,1])
    ϕ  = s.Λ[5].(x[:,1])

    # derivatives
    ∂x₀ = s.∂Λ[1].(x[:,1])./s.L
    ∂y₀ = s.∂Λ[2].(x[:,1])./s.L
    ∂a  = s.∂Λ[3].(x[:,1])./s.L
    ∂b  = s.∂Λ[4].(x[:,1])./s.L
    ∂ϕ  = s.∂Λ[5].(x[:,1])./s.L

    eθ = hcat(zeros(length(a)),
              -a.*cos.(ϕ).*sin.(x[:,2]) .- b.*sin.(ϕ).*cos.(x[:,2]) .+ x₀,
              -a.*sin.(ϕ).*sin.(x[:,2]) .+ b.*cos.(ϕ).*cos.(x[:,2]) .+ y₀)

    ez = hcat(ones(length(a)),
              +(∂a.*cos.(ϕ).-a.*sin.(ϕ).*∂ϕ).*cos.(x[:,2]) .- (∂b.*sin.(ϕ).+b.*cos.(ϕ).*∂ϕ).*sin.(x[:,2]) .+ ∂x₀,
              +(∂a.*sin.(ϕ).+a.*cos.(ϕ).*∂ϕ).*cos.(x[:,2]) .+ (∂b.*cos.(ϕ).-b.*sin.(ϕ).*∂ϕ).*sin.(x[:,2]) .+ ∂y₀)

    if normalized
        ez = ez ./ sqrt.(ez⋅ez)
        eθ = eθ ./ sqrt.(eθ⋅eθ)
    end

    return ez, eθ
end

# ------------------------------------------------------------------------
# manifold type

"""
    struct Manifold{T <: Real}
        mesh :: Mesh{T}
        surf :: Surface{T}
    end

Store the representation of a differential geometry object.
`mesh` is the empirical point cloud.
`surf` is the estimated differentiable surface.
"""
struct Manifold{T <: Real}
    mesh :: Mesh{T}
    surf :: Surface{T}
end

"""
    function pullback(ℳ::Manifold)

Returns the 2D cylindrical projection of the mesh, as estimated by the internal surface.
"""
function pullback(ℳ::Manifold)
    x    = pullback(ℳ.surf, ℳ.mesh.r)
    maxₗ = maximum(hcat(abs.(x[ℳ.mesh.tri[:,1],2]-x[ℳ.mesh.tri[:,2],2]),
                        abs.(x[ℳ.mesh.tri[:,2],2]-x[ℳ.mesh.tri[:,3],2]),
                        abs.(x[ℳ.mesh.tri[:,3],2]-x[ℳ.mesh.tri[:,1],2])),
                   dims=2)
    return x, ℳ.mesh.tri[getindex.(findall(maxₗ .> π/2),1), :]
end

pullback(ℳ::Manifold, r) = pullback(ℳ.surf, r)

"""
    function order(tri, r)

Reorder the triangulation so that all labels are counterclockwise.
"""
function order(tri, r)
    r21 = r[tri[:,2],:] .- r[tri[:,1],:]
    r32 = r[tri[:,3],:] .- r[tri[:,2],:]

    det = r21[:,1].*r32[:,2] - r21[:,2].*r32[:,1]
    idx = findall(det .< 0)

    tmp = tri[idx, 2]
    tri[idx, 2] .= tri[idx,3]
    tri[idx, 3] .= tmp

    return tri
end

"""
    function triangulation(r₀, rᵢ)

Return the triangle containing each point given in `rᵢ`.
The triangulation used is delaunay triangulation defined by point cloud `r₀`.
"""
function triangulation(r₀, rᵢ)
    tri = order(delaunay(Array{eltype(r₀),2}(r₀'))', r₀)

    # compute containing faces
    t1  = r₀[tri[:,2],:] - r₀[tri[:,1],:]
    t2  = r₀[tri[:,3],:] - r₀[tri[:,2],:]
    t3  = r₀[tri[:,1],:] - r₀[tri[:,3],:]

    t1 /= sqrt.(t1⋅t1)
    t2 /= sqrt.(t2⋅t2)
    t3 /= sqrt.(t3⋅t3)

    t1  = t1*[0 +1; -1 0]
    t2  = t2*[0 +1; -1 0]
    t3  = t3*[0 +1; -1 0]

    e1  = .5*(r₀[tri[:,2],:] + r₀[tri[:,1],:])
    e2  = .5*(r₀[tri[:,3],:] + r₀[tri[:,2],:])
    e3  = .5*(r₀[tri[:,1],:] + r₀[tri[:,3],:])

    θ1  = t1*rᵢ' .- sum(t1.*e1,dims=2)
    θ2  = t2*rᵢ' .- sum(t2.*e2,dims=2)
    θ3  = t3*rᵢ' .- sum(t3.*e3,dims=2)

    τ(i)  = findfirst((θ1[:,i] .> 0) .& (θ2[:,i] .> 0) .& (θ3[:,i] .> 0))
    faces = [τ(i) for i in 1:size(rᵢ,1)]
    return tri, faces
end

function continuum(tri, ϕ, r)
    function fit(ϕ, r)
        M = hcat(r[:,1], r[:,2], ones(size(r,1)))
        ρ = M \ ϕ

        f(x) = ρ[1]*x[1] + ρ[2]*x[2] + ρ[3]
        return f
    end

    return [fit(ϕ[tri[t,:]], r[tri[t,:],:]) for t in 1:size(tri,1)]
end

"""
    function interpolate(M, ϕ, r)

Interpolate the tensor field `ϕ`, defined at points `r` onto the vertices of manifold `M`.
Interpolation is computed by finding the containing triangle within the mesh of `M` for each point `r`.
Linear interpolation is performed per face.
"""
function interpolate(M, ϕ, r)
    x, _ = pullback(M)
    y    = pullback(M.surf, r)

    tri, face = triangulation(y, x)
    dist = sum( (x[:,d] .- y[:,d]').^2 for d in 1:size(x,2) )

    @match numdims(ϕ) begin
        1 => begin
            field  = continuum(tri, ϕ, y)
            result = Array{eltype(ϕ)}(undef, size(x,1))
            for i in 1:size(x,1)
                if isnothing(face[i])
                    result[i] = ϕ[argmin(dist[i,:])]
                else
                    result[i] = field[face[i]](x[i,:])
                end
            end

            return result
        end
        2 => begin
            ψ = Array{eltype(ϕ)}(undef, size(x,1), 3)
            for d in 1:3
                field  = continuum(tri, ϕ[:,d], y)
                for i in 1:size(x,1)
                    if isnothing(face[i])
                        ψ[i,d] = ϕ[argmin(dist[i,:]), d]
                    else
                        ψ[i,d] = [field[face[i]](x[i,:]) for i in 1:size(x,1)]
                    end
                end
            end
            return ψ
        end
        _ => return error("unsupported field dimension = $(size(ϕ))")
    end
end

"""
    rescale(x)

Rescale array `x` to run between ``[0,1]``.
"""
function rescale(x)
    min = minimum(x)
    max = maximum(x)
    return (x.-min)./(max.-min)
end

"""
    scalar(M::Manifold, ϕ, field::Symbol)

Return a scalar field `ϕ` interpolated onto either the embedded space of manifold `M` or the pullback.
If `field` is :ℝ² the pullback is computed.
If `field` is :ℝ³ the embedding space is returned.
"""
function scalar(M::Manifold, ϕ, field::Symbol)
    @match field begin
        :ℝ² => begin
            q, tri = pullback(ℳ)
            return (
                x1  = q[:,1],
                x2  = q[:,2],
                tri = tri',
                ϕ   = ϕ,
            )
        end
        :ℝ³ => begin
            return (
                x1  = M.mesh.r[:,1],
                x2  = M.mesh.r[:,2],
                x3  = M.mesh.r[:,3],
                tri = M.mesh.tri',
                ϕ   = ϕ,
            )
        end
        _   => return error("unsupported field $(field)")
    end
end

"""
    scalar(M::Manifold, ϕ, field::Symbol)

Return a vector field field `ϕ` interpolated onto either the embedded space of manifold `M` or the pullback.
If `field` is :ℝ² the pullback is computed. Only the tangent space component is kept.
If `field` is :ℝ³ the embedding space is returned.
"""
function vector(ℳ::Manifold, ϕ, field::Symbol)
    @match field begin
        :ℝ² => begin
                q, tri = pullback(ℳ)
                ez, eθ = basis(ℳ.surf, ℳ.mesh.r)
                ϕz, ϕθ = sum(ϕ.*ez,dims=2), sum(ϕ.*eθ,dims=2)

                return (
                    x1  = q[:,1],
                    x2  = q[:,2],
                    tri = tri',
                    v1  = ϕz,
                    v2  = ϕθ,
                )
            end
        :ℝ³ => begin
                return (
                    x1  = M.mesh.r[:,1],
                    x2  = M.mesh.r[:,2],
                    x3  = M.mesh.r[:,3],
                    tri = M.mesh.tri',
                    v1  = ϕ[:,1],
                    v2  = ϕ[:,2],
                    v3  = ϕ[:,3],
                )
            end
        _   => return error("unsupported field $field")
    end
end

"""
    makefield(ℳ::Manifold, ϕ::AbstractArray; field::Symbol = :ℝ³)

Return an arbitrary tensor field defined over _either_ the embedding space or the cylindrical pullback of manifold `M`.
Higher level function than either [`scalar`](@ref) or [`vector`](@ref).
"""
function makefield(ℳ::Manifold, ϕ::AbstractArray; field::Symbol = :ℝ³)
    @match numdims(ϕ) begin
        1                         => return scalar(ℳ, ϕ, field)
        2, if size(ϕ,2) == 3 end  => return vector(ℳ, ϕ, field)
        _                         => return error("unsupported field dimension = $(size(ϕ))")
    end
end

function manifold(io::IO, type::Symbol=:obj)
    m = mesh(io, type)
    s = Surface(m.r)

    return Manifold(m, s)
end

end


# Conversion of all complex variables to real ones in the Poly structures

"""
  realPart, imagPart = cplx2real(expo::Exponent)

  Convert a complex Exponent in complex variables into `realPart` and
  `imagPart` polynomials of twice as many variables, real and imag parts of
  `expo` variables. Done recursively with the `cplx2real_rec` function.
"""
function cplx2real(expo::Exponent)
  vars, inds = collect(keys(expo.expo)), collect(values(expo.expo))
  return cplx2real_rec(vars, inds, Polynomial()+1, Polynomial()+0, length(expo)+1, Degree(0,0))
end

"""
realPart, imagPart = cplx2real_rec(vars::Array{Variable}, degs::Array{Degree}, realPart::Polynomial, imagPart::Polynomial, cur_ind::Int, cur_deg::Degree)

Transform recursively the complex exponent represented by the `vars` and `degs`
arrays into its real and imag parts, functions of its imag and real part
variables.
`cur_ind` decreases to 0, `cur_deg` decreases to Degree(0,0) for each step of
cur_ind. Terminaison case is reached at 0, Degree(0,0).
Initial arrays `vars` and `degs` are read only.

### Arguments
- vars::Array{Variable}
- degs::Array{Degree}
- realPart::Polynomial
- imagPart::Polynomial
- cur_ind::Int
- cur_deg::Degree
"""
function cplx2real_rec(vars::Array{Variable}, degs::Array{Degree}, realPart::Polynomial, imagPart::Polynomial, cur_ind::Int, cur_deg::Degree)
  ## Final case:
  if cur_ind == 1 && cur_deg == Degree(0,0)
    return (realPart, imagPart)
  ## One less variable to deal with:
  elseif cur_deg == Degree(0,0)
    return cplx2real_rec(vars, degs, realPart, imagPart, cur_ind-1, degs[cur_ind-1])
  ## Recursion rule, decrease the current variable exponent until it reaches Degree(0,0):
  else
    var = vars[cur_ind]
    if iscomplex(var)
      var_real, var_imag = varname_cplx2real(var.name)
      var_R, var_I = Variable(var_real, Real), Variable(var_imag, Real)
      if cur_deg.explvar > 0
        return cplx2real_rec(vars, degs, var_R * realPart - var_I * imagPart, var_R * imagPart + var_I * realPart, cur_ind, Degree(cur_deg.explvar-1, cur_deg.conjvar))
      elseif cur_deg.conjvar > 0
        cur_deg.explvar == 0 || warn("cur_deg.explvar should be 0 (and not $(cur_deg.explvar)), set to this value")
        return cplx2real_rec(vars, degs, var_R * realPart + var_I * imagPart, var_R * imagPart - var_I * realPart, cur_ind, Degree(0, cur_deg.conjvar-1))
      end
    elseif isbool(var)
      return cplx2real_rec(vars, degs, var*realPart, var*imagPart, cur_ind, Degree(0,0))
    else
      return cplx2real_rec(vars, degs, Exponent(Dict(var=>cur_deg))*realPart, Exponent(Dict(var=>cur_deg))*imagPart, cur_ind, Degree(0,0))
    end
  end
end

"""
  realPart, imagPart = cplx2real(pol::Polynomial)

  Convert a complex polynomial in complex variables into `realPart` and
  `imagPart` polynomials of twice as many variables, real and imag parts of
  `pol` variables.
"""
function cplx2real(pol::Polynomial)
  realPart = Polynomial()
  imagPart = Polynomial()

  for (expo, λ) in pol
    realexpo, imagexpo = cplx2real(expo)

    realPart += realexpo*real(λ) - imagexpo*imag(λ)
    imagPart += imagexpo*real(λ) + realexpo*imag(λ)
  end
  return (realPart, imagPart)
end

function cplx2real(pt_C::Point)
  pt = Point()
  for (var, val) in pt_C
    if var.kind <: Complex
      var_real, var_imag = varname_cplx2real(var.name)
      if real(val) != 0
        pt[Variable(var_real, Real)] = real(val)
      end
      if imag(val) != 0
        pt[Variable(var_imag, Real)] = imag(val)
      end
    else
      if val != 0
        pt[var] = real(val)
      end
    end
  end
  return pt
end

function real2cplx(pt::Point)
    ptC = Point()
    for (var, val) in pt
        if ismatch(r"(_Re|_Im)", var.name)
            if ismatch(r"_Re", var.name)
                varname = var.name[1:(end-3)]
                ptC[Variable(varname, Complex)] = val + im*pt[Variable(varname*"_Im", Real)]
            end
        else
            ptC[var] = val
        end
    end
    return ptC
end


"""
Run existing CellML model and save output membrane potential Vm
"""

using CellMLToolkit, DifferentialEquations, OrdinaryDiffEq, Sundials
using PyPlot
PyPlot.pygui(true)

function LoadAndRun(model_name::String; show_params=false, show_states=false)
    """
    Load and run model with model name `model_name`
    """
    # path to cellml models
    cellml_path = ".//data//CellML_models//$model_name"

    # CellModel factory function creates an ODESystem
    ml = CellMLToolkit.CellModel(cellml_path; dependency=false)

    # model parameters
    if show_params
        params = list_params(ml)
        @info params
    end

    # model states
    states = list_states(ml)
    if show_states
        @info states
    end

    # find index of membrane voltage in list of state variables
    idx = findfirst([occursin("V", string(s)) for s in states])
    @info "Index of voltage variable:", idx

    # generate and solve ODEProblem
    tspan = (0., 0.4)
    ts = collect(0:1e-5:tspan[2])
    prob = CellMLToolkit.ODEProblem(ml, tspan);
    sol = solve(prob, TRBDF2(), saveat=ts)
    # extract membrane voltage
    Vm = sol[idx, :]

    # plot
    f, ax = plt.subplots(figsize=(10, 5))
    ax.set_ylabel("Voltage (mV)")
    ax.set_xlabel("Time (ms)")
    ax.plot(sol.x, Vm, lw=2)
    f.tight_layout()
    PyPlot.show()
end

model_name = "severi_fantini_charawi_difrancesco_2012.cellml"
LoadAndRun(model_name; show_params=true, show_states=true)


using Test

using BEAST
using CompScienceMeshes

V = [
    point(0,0,0),
    point(1,0,0),
    point(0,1,0)]
F = [index(1,2,3)]

m1 = Mesh(V,F)
m2 = CompScienceMeshes.rotate(m1, 0.5pi*point(1,0,0))
m3 = CompScienceMeshes.rotate(m1, 1.0pi*point(1,0,0))

m = weld(m1,m2,m3)

@test numcells(skeleton(m,0)) == 5
@test numcells(m) == 3

Nd = BEAST.nedelec(m)
@test numfunctions(Nd) == 7

function interior(mesh::Mesh, edges=skeleton(mesh,1))
    @assert dimension(mesh) == 2
    @assert vertices(mesh) === vertices(edges)

    C = connectivity(edges, mesh)
    @assert size(C) == (numcells(mesh), numcells(edges))

    nn = vec(sum(abs.(C), dims=1))
    T = CompScienceMeshes.celltype(edges)
    interior_edges = Vector{T}()
    for (i,edge) in pairs(cells(edges))
        nn[i] > 1 && push!(interior_edges, edge)
    end
    Mesh(vertices(mesh), interior_edges)
end

interior_edges = interior(m)
@test numcells(interior_edges) == 1
Nd = BEAST.nedelec(m, interior_edges)
@test numfunctions(Nd) == 1


# TODO: check trees are on the same leave set

"""
    RF(tree1::T, tree2::T)::Int64 where T <:GeneralNode

Calculate the Robinson-Foulds distance between the two trees.
In its current form the function assumes the trees have identical leave sets.

Returns result of algorithm as integer.

* `tree1` : tree used to determine RF distance.

* `tree2` : tree used to determine RF distance.
"""
function RF(tree1::T, tree2::T)::Int64 where T <: GeneralNode
    bt3 = get_bipartitions(tree1)
    bt4 = get_bipartitions(tree2)
    length(bt3)+length(bt4) - 2* length(intersect(bt3, bt4))
end

"""
    get_bipartitions(tree::T)::Vector{Tuple} where T <:GeneralNode

Get a vector of all bipartions of `tree`.

Returns a vector containing Tuples of sets representing the bipartions.
"""
function get_bipartitions(tree::T)::Vector{Tuple} where T <:GeneralNode
    po_vect= post_order(tree)[1:end-1]
    bt = Vector{Tuple}(undef, length(po_vect))
    Base.Threads.@threads for ind in eachindex(po_vect)
        elem = po_vect[ind]::T
        inset = Set{Int64}()
        outset = Set{Int64}()
        @simd for elem2 in po_vect
            startswith(elem2.binary, elem.binary) ? push!(inset, elem2.num) : push!(outset, elem2.num)
        end # for
        @inbounds bt[ind] = (inset, outset)
    end # for
    bt
end

"""
    BHV_bounds(tree1::T, tree2::T)::Tuple{Float64, Float64} where T <:GeneralNode

This function calculates the lower and upper bounds of the geodesic in the
Billera-Holmes-Vogtman space.

Returns tuple of floats.
"""
function BHV_bounds(tree1::T, tree2::T)::Tuple{Float64, Float64} where T <:GeneralNode
    res_upper_1 = Threads.Atomic{Float64}(0.0)
    res_upper_2 = Threads.Atomic{Float64}(0.0)
    res_upper_3 = Threads.Atomic{Float64}(0.0)
    po::Vector{T} = post_order(tree1)
    Base.Threads.@threads for node in po[1:end-1]
        nom = find_num(tree2, node.num)
        if node.mother.num == nom.mother.num
            Threads.atomic_add!(res_upper_3, (node.inc_length-nom.inc_length)^2)
        else
            Threads.atomic_add!(res_upper_2 , nom.inc_length^2)
            Threads.atomic_add!(res_upper_1, node.inc_length^2)
        end # if

    end # for
    res_low::Float64 = res_upper_1[]+res_upper_2[]+res_upper_3[]
    res_high::Float64 = (sqrt(res_upper_2[])+sqrt(res_upper_1[]))^2+res_upper_3[]

    sqrt(res_low), sqrt(res_high)
end


# Most code below is taken from https://github.com/richardkwo/InvariantCausal.jl,
# specifically SEM.jl with minor modifications to update for Julia 1.0
using LinearAlgebra # For identity matrix, I
using Distributions

abstract type SCM end

struct LinAddSCM <: SCM
   p ::Int64 # Number of variables
   B ::Matrix{Float64} # Weighted Adjacency Matrix of coefficients
   variances ::Vector{Float64} # Variances for each variables error term
   function LinAddSCM(B, variances)
       @assert length(variances) == size(B,1) == size(B,2)
       @assert all(variances .> 0)
       p = length(variances)
       new(p, B, variances)
   end
end

struct NonLinAddSCM <: SCM
   p ::Int64 # Number of variables
   A ::Matrix{Float64} # Adjacency Matrix (signed so weights can only be {-1,0,1})
   fs::Vector{Function} # Vector of Non-linear Functions for each variable
   # variances ::Vector{Float64} # Variances for each variables error term; NOTE: will reintroduce after the first sim
   function NonLinAddSCM(A, fs)#, variances)
       @assert size(A,1) == size(A,2) == length(fs)# == length(variances)
       # @assert all(variances .> 0)
       # p = length(variances)
       p = size(A,1)
       new(p, A, fs)#, variances)
   end
end

function Base.show(io::IO, scm::LinAddSCM)
    println(io, "LinAdd SCM with $(scm.p) variables.")
    println(io, "σ² = $(scm.variances).")
end

"""
    simulate(scm; d::UnivariateDistribution)
    simulate(scm, [n]; d::UnivariateDistribution)
    simulate(scm, [do_variables, do_values], [shift, additive], [n]; d::UnivariateDistribution)

Simulate from a Linear Additive (LinAdd) scm `scm`. `n` is the sample size.
`d` is a keyword argument that specifies which distribution to sample from for the error terms.
do-interventions can be performed by specifying vectors of `do_variables` and `do_values`.
shift-interventions can be performed by setting `shift=true` with `shift_values = do_variables`
and `shift_values = do_values`.
shift-interventions can be either additive by letting `additive=true` or multiplicative `additive=false`.
"""
function simulate(scm::LinAddSCM, d::UnivariateDistribution)
    ϵ = rand(d, scm.p) .* sqrt.(scm.variances) # Scale variances to our chosen variances
    return (I - scm.B) \ ϵ
end

function simulate(scm::LinAddSCM, n::Int64; d::UnivariateDistribution=Normal(0,1))
    return vcat(map(i -> simulate(scm, d), 1:n)'...)
end

function simulate(scm::LinAddSCM, d::UnivariateDistribution,
                  interventions::Tuple{Vector{Int64}, Vector{Float64}},
                  shift, additive)
    do_variables, do_values = interventions
    @assert length(do_variables) == length(do_values)
    p = scm.p
    ϵ = rand(d, p) .* sqrt.(scm.variances)
    if !shift # If shift intervention is false, perform an atomic intervention
        ϵ[do_variables] .= do_values
        B = copy(scm.B)
        B[do_variables, :] .= 0
    else
        shift_variables, shift_values = do_variables, do_values
        if additive
            ϵ[shift_variables] .+= shift_values
        else
            ϵ[shift_variables] .*= shift_values
        end
        B = copy(scm.B)
    end
    return (I - B) \ ϵ
end

function simulate(scm::LinAddSCM, interventions::Tuple{Vector{Int64}, Vector{Float64}}, n::Int64;
                  d::UnivariateDistribution=Normal(0,1),
                  shift=false, additive=true)
    return vcat(map(i -> simulate(scm, d, interventions, shift, additive), 1:n)'...)
end

# simulate observationa data, but allowing to give the full vector of distributions for every noises separately in the SCM
function simulate(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, n::Int64)
    # check
    @assert length(noise_dist) == scm.p

    # interventions
    jointdist = product_distribution(noise_dist)
    B = copy(scm.B)

    # simulate
    ϵ = rand(jointdist, n) .* sqrt.(scm.variances)

    return ((I - B) \ ϵ)'
end

function simulate(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, n,
                  noise=nothing, sum_or_prod=:sum, top_order=collect(1:6))
    # check
    @assert length(noise_dist) == scm.p
    if sum_or_prod == :sum
        g = xs -> dropdims(sum(xs, dims=2), dims=2)
    else
        g = xs -> dropdims(prod(xs, dims=2), dims=2)
    end

    # simulate
    if isnothing(noise)
        jointdist = product_distribution(noise_dist)
        ϵ = rand(jointdist, n) #.* sqrt.(scm.variances) # p × n; will reintroduce this term after the first sim we do
    else
        ϵ = noise
    end
    A = scm.A
    fs = scm.fs
    x = zeros(n, p)
    for j in top_order
        inds = findall(!iszero, A[:, j])
        if length(inds) == 0
            x[:,j] .= ϵ[j, :]
        else
            signs = A[inds,j]
            x[:, j] .= g(fs[j].(x[:, inds] .* signs')) .+ ϵ[j, :]
        end
    end
    return x
end

"""
        stochastic_inter_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64)

Simulate from an SCM with stochastic interventions, that is do(X:=N), with N indep. of the rest, for a set list of variables X.

* `scm`         : a linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `inter_vect`  : tuple of indexes of variables to intervene on, and their corresponding interventional distr. (~N)
* `n`           : nbr of samples
"""
function stochastic_inter_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                                Vector{<:UnivariateDistribution}}, n::Int64)
    # check
    inter_variables, inter_distr = inter_vect
    @assert length(noise_dist) == scm.p >= length(inter_variables) == length(inter_distr)

    # interventions
    dist_vect = copy(noise_dist)
    dist_vect[inter_variables] = inter_distr
    jointdist = product_distribution(dist_vect)
    B = copy(scm.B)
    B[inter_variables,:] .= 0

    # simulate
    ϵ = rand(jointdist, n) .* sqrt.(scm.variances)

    return ((I - B) \ ϵ)'
end

"""
        stochastic_inter_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64, fs::Vector{Function})

Simulate from a non-linear SCM with stochastic interventions, that is do(X:=N), with N indep. of the rest, for a set list of variables X.

* `scm`         : a non-linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `inter_vect`  : tuple of indexes of variables to intervene on, and their corresponding interventional distr. (~N)
* `n`           : nbr of samples
* `fs`          : vector of non-linear functions used for the graph
"""
function stochastic_inter_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect, n::Int64, fs::Vector{Function})
    # interventions
    dist_vect = copy(noise_dist)
    inter_variables, inter_distr = inter_vect
    dist_vect[inter_variables] = inter_distr
    A = copy(scm.A)
    A[:,inter_variables] .= 0
    new_scm = NonLinAddSCM(A, fs)
    return simulate(new_scm, dist_vect, n)
end

"""
        mech_changes_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64)

Simulate from an SCM with mechanism changes, by either rescaling (i.e. X = β*X' + A*ϵ, with A random) or shifting the noise (i.e. X = β*X' + C+ϵ, with C random) or changing the structural equation's coefficients (i.e. X = β'*X' + ϵ) - or both coeff change + shift/rescale noise.

* `scm`         : a linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `noise_change`: tuple of indexes of variables for scale/shift noise change, and corresponding distr. (~A/C)
* `coeff_change`: tuple of indexes of variables for coeff. changes, and corresponding new coeffs (β')
* `n`           : nbr of samples
* `additive`    : whether we shift (=true) or rescale (=false) the intervened noise
"""
function mech_changes_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, noise_change::Tuple{Vector{Int64},
                                Vector{<:UnivariateDistribution}}, coeff_change::Tuple{Vector{Int64}, Array{Float64,2}}, n::Int64;
                                additive = true)
    # check
    noise_change_variables, noise_change_distr::Array{UnivariateDistribution{Continuous},1} = noise_change
    coeff_variables, coeff_mat = coeff_change
    @assert size(coeff_mat)[1] == length(coeff_variables) <= size(coeff_mat)[2] == length(noise_dist) == scm.p >=                      length(noise_change_variables) == length(noise_change_distr)

    # interventions
    joint_change_distr = product_distribution(noise_change_distr)
    B = copy(scm.B)
    B[coeff_variables,:] = coeff_mat

    # simulate
    jointdist = product_distribution(noise_dist)
    ϵ = rand(jointdist, n) .* sqrt.(scm.variances)
    noise_mat = (additive ? zeros(scm.p, n) : ones(scm.p, n))
    noise_mat[noise_change_variables,:] = rand(joint_change_distr, n)

    return ((I - B) \ (additive ? ϵ .+ noise_mat : ϵ .* noise_mat))'

end

"""
        mech_changes_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64)

Simulate from an SCM with mechanism changes, by either rescaling (i.e. X = f(X) + A*ϵ, with A random) or shifting the noise
(i.e. X = f(X) + C+ϵ, with C random) or changing the structural equation's coefficients (i.e. X = β'*X' + ϵ) - or both coeff change + shift/rescale noise.

* `scm`         : a non-linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `noise_change`: tuple of indexes of variables for scale/shift noise change, and corresponding distr. (~A/C)
* `n`           : nbr of samples
* `additive`    : whether we shift (=true) or rescale (=false) the intervened noise
"""
function mech_changes_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, noise_change::Tuple{Vector{Int64},
                                Vector{<:UnivariateDistribution}}, n::Int64;
                                additive = true)
    # check
    noise_change_variables, noise_change_distr::Array{UnivariateDistribution{Continuous},1} = noise_change
    @assert length(noise_dist) == scm.p >= length(noise_change_variables) == length(noise_change_distr)

    # interventions
    joint_change_distr = product_distribution(noise_change_distr)

    # simulate
    jointdist = product_distribution(noise_dist)
    ϵ = rand(jointdist, n) #.* sqrt.(scm.variances) # NOTE: pretend it doesnt have variances for now
    noise_mat = (additive ? zeros(scm.p, n) : ones(scm.p, n))
    noise_mat[noise_change_variables,:] = rand(joint_change_distr, n)
    # although it seems like noise_dist is getting used here, it is isn't and the noise distribution is being passed in.
    return simulate(scm, noise_dist, n, (additive ? ϵ .+ noise_mat : ϵ .* noise_mat))
end

# get direct causes for i
function causes(scm::LinAddSCM, i::Int64)
    @assert 1 <= i <= scm.p
    return (1:scm.p)[scm.B[i, :].!=0]
end

# get markov blannket for i
function markov_blanket(scm::LinAddSCM, i::Int64)
    blanket = causes(scm, i)
    children = (1:scm.p)[scm.B[:,i].!=0]
    blanket = vcat(blanket, children)
    for j in children
        blanket = vcat(blanket,(1:scm.p)[scm.B[j, :].!=0])
    end
    return unique(blanket[blanket .!= i])
end

# function returning the structural equation coeff. for variable i
function causalcoeff(scm::LinAddSCM, i::Int64)
    @assert 1 <= i <= scm.p
    return (scm.B[i, 1:scm.p .!= i])
end

"""
    fixed_linear_SCM(p, k; [lb=-2, ub=2, var_min=0.5, var_max=2])

Generate a fixed-graph acyclic SCM with `p` variables with `edges` connecting them, and random coefficients.
* Edges should be of form [(child, parent), ...]
* `lb`, `ub`: coeff  ~ unif[`lb`, `ub`] with random sign
* `var_min`, `var_max`: var of error ~ unif[`var.min`, `var.max`]
"""
function fixed_linear_SCM(edges::Array{Tuple{Int64, Int64},1}, p::Int64; lb=-2, ub=2, var_min=0.5^2, var_max=2^2)
    B = zeros(p, p)
    for e in edges
        B[e...] = (rand() * (ub - lb) + lb) * sign(randn())
    end
    variances = rand(p) * (var_max - var_min) .+ var_min
    return LinAddSCM(B, variances)
end

"""
    random_linear_SCM(p, k; [lb=-2, ub=2, var_min=0.5, var_max=2])

Generate a random-graph acyclic SCM with `p` variables and `k` average degree, and random coefficients.
* `lb`, `ub`: coeff  ~ unif[`lb`, `ub`] with random sign
* `var_min`, `var_max`: var of error ~ unif[`var.min`, `var.max`]
"""
function random_linear_SCM(p::Int64, k::Int64; lb=-2, ub=2, var_min=0.5^2, var_max=2^2)
    B = zeros(p, p)
    B[rand(p, p) .< k / (p-1)] .= 1
    B[UpperTriangular(B).!=0] .= 0
    m = sum(B.==1)
    B[B.==1] .= (rand(m) * (ub - lb) .+ lb) .* sign.(randn(m))
    err_var = rand(p) * (var_max - var_min) .+ var_min
    _order = sample(1:p, p, replace=false)
    B = B[_order, _order]
    return LinAddSCM(B, err_var)
end

"""
    random_noise_intervened_SCM(scm::LinAddSCM, [p_intervened=2, noise_multiplier_min=0.5, noise_multiplier_max=2., avoid=[],
                                                   prob_coeff_unchanged=2/3, lb=-2, ub=2])

Produce a new SCM based on original SCM by changing coefficients and noise variances.
* `p_intervened`: randomly choose `p_intervened` variables to intervene; will avoid those specified in `avoid`
* [`noise_multiplier_min`, `noise_multiplier_max`]: interval that noise multiplier is uniformly sampled from
* `prob.coeff.unchanged`: probability that coefficient is not changed
* `[lb, ub]`: if to change, coefficient is drawn uniformly from this interval with random sign

Return: `scm_new`, `intervened_variables`
"""
function random_noise_intervened_SCM(scm::LinAddSCM;
                                     p_intervened=2, noise_multiplier_min=0.5, noise_multiplier_max=2., avoid=[],
                                     prob_coeff_unchanged=2/3, lb=-2, ub=2)
    B = copy(scm.B)
    p = scm.p
    err_var = copy(scm.variances)
    vars = sample(setdiff(collect(1:p), avoid), p_intervened, replace=false)
    for i in vars
        noise_multiplier = rand() * (noise_multiplier_max - noise_multiplier_min) + noise_multiplier_min
        err_var[i] = err_var[i] * noise_multiplier
        if rand() > prob_coeff_unchanged
            _J = (1:p)[B[i, :].!=0]
            B[i, _J] .= rand(length(_J)) * (ub - lb) .+ lb
        end
    end
    return LinAddSCM(B, err_var), vars
end

"""
        high_dim_SCM(p_anterior::Int64, p_posterior::Int64, k::Int64, parents_nbr::Int64 ;
                            lb=-2, ub=2, var_min=0.5, var_max=2, shuffle = false)

Generate a linear SCM, by allowing to fix the number of the direct causes (i.e. parents) of the target.
Mainly used for high-dimensional simulations. The variables are returned in their causal order.

* `p_anterior`          : nbr of non-desc variables to the target (which is the (`p_anterior`+1)th variable)
* `p_posterior`         : nbr of variables non-asc variables to the target
* `k`                   : averaged degree of every other nodes
* `parents_nbr`         : nbr of direct causes of the target
* `lb`, `ub`            : coeff  ~ unif[`lb`, `ub`] with random sign
* `var_min`, `var_max`  : var of error ~ unif[`var.min`, `var.max`]
"""
function high_dim_SCM(p_anterior::Int64, p_posterior::Int64, k::Int64, parents_nbr::Int64 ;
                      lb=-2, ub=2, var_min=0.5^2, var_max=2^2)
    # check
    @assert parents_nbr <= p_anterior

    # vars
    p = p_anterior + p_posterior + 1
    target_idx = p_anterior + 1

    # adjacency matrix
    B = zeros(p, p)
    draws = rand(p, p)
    draws[UpperTriangular(draws).> 0] .= 0
    B[draws .> 1 - min(k,p-1) / (p-1)] .= 1

    # target parents
    B[target_idx,sortperm(draws[target_idx,:])] = vcat(zeros(p - parents_nbr), ones(parents_nbr))
    m = sum(B.==1)

    # generate the coeff
    B[B.==1] .= (rand(m) * (ub - lb) .+ lb) .* rand([-1,1],m)

    # generate the variances
    err_var = rand(p) * (var_max - var_min) .+ var_min

    # Normalize X*β so that var(X*β) ∝ var(ϵ)
    B_inv = inv(I - B)
    for i in 2:p
        b²ᵢ = B_inv[i, :].^2
        vᵢ = sqrt(dot(b²ᵢ, err_var))
        B[i,:] ./= vᵢ
    end

    return LinAddSCM(B, err_var)
end


import Base.max
function max(pot::ArrayPotential,variables;MaxOver=true,ReturnState=false,Ind2Sub=false)
    #MAX Maximise a multi-dimensional array over a set of dimenions
    # maxval=max(x,variables)
    # find the values and states that maximize the multidimensional array x
    # over the dimensions in maxover
    #
    if !MaxOver
        variables = setdiff(pot.variables,variables)
    end

    newvars=setdiff(pot.variables,variables)
    max_variables=memberinds(variables,pot.variables)

    maxval,maxind=findmax(pot.content,max_variables) ## PROBABLY SHOULD CHANGE TO THE MAX FUNCTION -- NEED TO ENSURE THIS CAN WORK FOR ANY POTENTIAL

    #maxval=squeeze(maxval,find(memberinds(size(maxval),1)))
    maxval=squeeze(maxval,tuple(find(memberinds(size(maxval),1))...))

    outpot=deepcopy(pot)
    outpot.variables=newvars; outpot.content=maxval

    if !ReturnState
        return outpot
    end

    if !Ind2Sub
        return outpot, maxind
    end

    s=ind2sub(size(pot.content),maxind[:]) # return in more readable form
    maxstate=zeros(Int64,length(s[1]),length(s))
    for i=1:length(pot.variables)
        maxstate[:,i]=s[i]
    end

    return outpot, maxstate

end

##################################################################
# PotArray type functions:

export numstates
function numstates{T<:ArrayPotential}(A::T)

    if isavector(A.content)
        n=prod(size(A.content));
    else
        n=size(A.content);
    end
    return n
end

export standardise
function standardise{T<:ArrayPotential}(A::T)
    # If the ArrayPotential is a vector, make this a column vector, otherwise leave unchanged
    if isavector(A.content)
        A.content=vec(A.content)
        end
    return A
end


using Flux
using Flux: flatten
using ExplainableAI: flatten_model, has_output_softmax, check_output_softmax, activation
using ExplainableAI: stabilize_denom, batch_dim_view, drop_batch_index
using Random

pseudorand(dims...) = rand(MersenneTwister(123), Float32, dims...)

# Test `activation`
@test activation(Dense(5, 2, gelu)) == gelu
@test activation(Conv((5, 5), 3 => 2, softplus)) == softplus
@test activation(BatchNorm(5, selu)) == selu
@test isnothing(activation(flatten))

# flatten_model
@test flatten_model(Chain(Chain(Chain(abs)), sqrt, Chain(relu))) == Chain(abs, sqrt, relu)
@test flatten_model(Chain(abs, sqrt, relu)) == Chain(abs, sqrt, relu)
@test_deprecated flatten_chain(Chain(identity))

# has_output_softmax
@test has_output_softmax(Chain(abs, sqrt, relu, softmax)) == true
@test has_output_softmax(Chain(abs, sqrt, relu, tanh)) == false
@test has_output_softmax(Chain(Chain(abs), sqrt, Chain(Chain(softmax)))) == true
@test has_output_softmax(Chain(Chain(abs), Chain(Chain(softmax)), sqrt)) == false
@test has_output_softmax(Chain(Dense(5, 5, softmax), Dense(5, 5, softmax))) == true
@test has_output_softmax(Chain(Dense(5, 5, softmax), Dense(5, 5, relu))) == false
@test has_output_softmax(Chain(Dense(5, 5, softmax), Chain(Dense(5, 5, softmax)))) == true
@test has_output_softmax(Chain(Dense(5, 5, softmax), Chain(Dense(5, 5, relu)))) == false

# check_output_softmax
@test_throws ArgumentError check_output_softmax(Chain(abs, sqrt, relu, softmax))

# strip_softmax
d_softmax = Dense(2, 2, softmax; init=pseudorand)
d_softmax2 = Dense(2, 2, softmax; init=pseudorand)
d_relu = Dense(2, 2, relu; init=pseudorand)
d_identity = Dense(2, 2; init=pseudorand)
# flatten to remove softmax
m = strip_softmax(Chain(Chain(abs), sqrt, Chain(Chain(softmax))))
@test m == Chain(abs, sqrt)
m1 = strip_softmax(Chain(d_relu, Chain(d_softmax)))
m2 = Chain(d_relu, d_identity)
x = rand(Float32, 2, 10)
@test typeof(m1) == typeof(m2)
@test m1(x) == m2(x)
# don't do anything if there is no softmax at the end
@test strip_softmax(Chain(Chain(abs), Chain(Chain(softmax)), sqrt)) ==
    Chain(Chain(abs), Chain(Chain(softmax)), sqrt)
@test strip_softmax(Chain(d_softmax, Chain(d_relu))) == Chain(d_softmax, Chain(d_relu))

# stabilize_denom
A = [1.0 0.0 1.0e-25; -1.0 -0.0 -1.0e-25]
S = @inferred stabilize_denom(A, 1e-3)
@test S ≈ [1.001 1e-3 1e-3; -1.001 1e-3 -1e-3]
S = @inferred stabilize_denom(Float32.(A), 1e-2)
@test S ≈ [1.01 1.0f-2 1.0f-2; -1.01 1.0f-2 -1.0f-2]

# batch_dim_view
A = [1 2; 3 4]
V = @inferred batch_dim_view(A)
@test size(V) == (2, 2, 1)
@test V[:, :, 1] == A

# drop_batch_index
I1 = CartesianIndex(5, 3, 2)
I2 = @inferred drop_batch_index(I1)
@test I2 == CartesianIndex(5, 3)
I1 = CartesianIndex(5, 3, 2, 6)
I2 = @inferred drop_batch_index(I1)
@test I2 == CartesianIndex(5, 3, 2)


module image_transform
export image_to_vector

using Images

function image_to_vector(filepath::String)
    image = load(filepath)
    inverted_bw = broadcast(x -> 1.0 - x, Gray.(image)) # black on white to white on black
    not_black = broadcast(!=, inverted_bw, Gray{N0f8}(0.0)) # pixels which are not black
    bounding_box = (
        up = findall(any(not_black, dims=2)[:, 1])[begin],
        down = findall(any(not_black, dims=2)[:, 1])[end],
        left = findall(any(not_black, dims=1)[1, :])[begin],
        right = findall(any(not_black, dims=1)[1, :])[end]
    )
    canvas_dims = (28, 28)
    canvas = zeros((canvas_dims))

    bounding_box_dims = (bounding_box.down - bounding_box.up, bounding_box.right - bounding_box.left)
    rescaled_height, rescaled_width = round.(Int, (bounding_box_dims) .* (20 / maximum(bounding_box_dims)))

    # image within bounding box is scaled (maintaining aspect ratio) to fit in 20x20 and pasted in the center of the canvas
    canvas[
        canvas_dims[1]÷2 - rescaled_height÷2 : canvas_dims[1]÷2 + (rescaled_height - rescaled_height÷2) - 1,
        canvas_dims[2]÷2 - rescaled_width÷2 : canvas_dims[2]÷2 + (rescaled_width - rescaled_width÷2) - 1,
        ] = imresize(
        inverted_bw[bounding_box.up:bounding_box.down, bounding_box.left:bounding_box.right]
            |> dilate
            |> dilate
            |> dilate,
        (rescaled_height, rescaled_width)
    )

    return reduce(vcat, Float64.(canvas)')
end
end


struct Cglsr
    B::Matrix{Float64}
    g::Vector{Float64}
    xmeans::Vector{Float64}
    ymeans::Vector{Float64}
    F::Union{Array{Float64}, Nothing}
end

"""
    cglsr(X, y; nlv, reorth = true, filt = false)
Conjugate gradient algorithm for the normal equations (CGLS) (Björck 1996).
* `X` : X-data.
* `y` : Univariate Y-data.
* `nlv` : Nb. CG iterations.
* `reorth` : If `true`, a Gram-Schmidt reorthogonalization of the normal equation
    residual vectors is done.
* `filt` : If `true`, the CG filter factors are computed (output `F`).

CGLS algorithm "7.4.1" Bjorck 1996, p.289

The code computing the re-orthogonalization (Hansen 1998) and filter factors (Vogel 1987, Hansen 1998)
is a transcription (with few adaptations) of the Matlab function `cgls`
(Saunders et al. https://web.stanford.edu/group/SOL/software/cgls/; Hansen 2008).

`X` and `y` are internally centered. The model is computed with an intercept.

## References

Björck, A., 1996. Numerical Methods for Least Squares Problems, Other Titles in Applied Mathematics.
Society for Industrial and Applied Mathematics. https://doi.org/10.1137/1.9781611971484

Hansen, P.C., 1998. Rank-Deficient and Discrete Ill-Posed Problems, Mathematical Modeling and Computation.
Society for Industrial and Applied Mathematics. https://doi.org/10.1137/1.9780898719697

Hansen, P.C., 2008. Regularization Tools version 4.0 for Matlab 7.3.
Numer Algor 46, 189–194. https://doi.org/10.1007/s11075-007-9136-9

Manne R. Analysis of two partial-least-squares algorithms for multivariate calibration. Chemometrics Intell.
Lab. Syst. 1987; 2: 187–197.

Phatak A, De Hoog F. Exploiting the connection between
PLS, Lanczos methods and conjugate gradients: alternative proofs of some properties of PLS. J. Chemometrics
2002; 16: 361–367.

Vogel, C. R.,  "Solving ill-conditioned linear systems using the conjugate gradient method",
Report, Dept. of Mathematical Sciences, Montana State University, 1987.
"""
function cglsr(X, y; nlv, reorth = true, filt = false)
    cglsr!(copy(X), copy(y); nlv = nlv, reorth = reorth, filt = filt)
end
function cglsr!(X, y; nlv, reorth = true, filt = false)
    X = ensure_mat(X)
    y = vec(y)
    n = size(X, 1)
    p = size(X, 2)
    xmeans = colmeans(X)
    ymeans = colmeans(y)
    center!(X, xmeans)
    center!(y, ymeans)
    # Pre-allocation and initialization
    B = similar(X, p, nlv)
    b = zeros(p)
    r = y       # r = y - X * b, with b = 0
    s = X' * r
    zp = copy(s)
    q = similar(X, n)
    g = dot(s, s)
    gnew = similar(X, nlv)
    if(reorth)
        A = similar(X, p, nlv + 1)
        A[:, 1] .= s ./ sqrt(g)
    end
    F = nothing
    if(filt)
        eig = svd(X).S.^2
        if(n < p)
            eig = [eig ; zeros(p - n)]
        end
        F = similar(X, p, nlv)
        Fd = similar(X, p)
    end
    fudge = 1e-4
    # End
    @inbounds for j in 1:nlv
        mul!(q, X, zp)
        alpha = g / dot(q, q)
        b .+= alpha * zp
        r .-= alpha * q
        mul!(s, X', r)
        # Reorthogonalize s to previous s-vectors
        if reorth
            @inbounds for i in 1:j
                v = vcol(A, i)
                s .-= dot(v, s) * v
            end
            A[:, j + 1] .= s ./ sqrt(dot(s, s))
        end
        # End
        gnew[j] = dot(s, s)
        beta = gnew[j] / g
        g = copy(gnew[j])
        zp .= s .+ beta * zp
        B[:, j] .= b
        # Filter factors
        # fudge threshold is used to prevent filter factors from exploding
        if filt
            if j == 1
                F[:, 1] .= alpha * eig
                Fd .= eig .- eig .* F[:, 1] .+ beta * eig
            else
                F[:, j] .= F[:, j - 1] .+ alpha * Fd
                Fd .= eig - eig .* F[:, j] .+ beta * Fd
            end
            if j > 2
                u = (abs.(F[:, j - 1] .- 1) .< fudge) .* (abs.(F[:, j - 2] .- 1) .< fudge)
                u = findall(u)
                if length(u) > 0
                    F[u, j] .= ones(length(u))
                end
            end
        end
        # End
    end
    Cglsr(B, gnew, xmeans, ymeans, F)
end

"""
    coef(object::Cglsr)
Compute the b-coefficients of a fitted model.
* `object` : The fitted model.
* `nlv` : Nb. iterations to consider. If nothing, it is the maximum nb. iterations.
"""
function coef(object::Cglsr; nlv = nothing)
    a = size(object.B, 2)
    isnothing(nlv) ? nlv = a : nlv = min(nlv, a)
    B = object.B[:, nlv:nlv]
    int = object.ymeans' .- object.xmeans' * B
    (B = B, int = int)
end

"""
    predict(object::Cglsr, X; nlv = nothing)
Compute Y-predictions from a fitted model.
* `object` : The maximal fitted model.
* `X` : X-data for which predictions are computed.
* `nlv` : Nb. iterations, or collection of nb. iterations, to consider.
If nothing, it is the maximum nb. iterations.
"""
function predict(object::Cglsr, X; nlv = nothing)
    a = size(object.B, 2)
    isnothing(nlv) ? nlv = a : nlv = (max(minimum(nlv), 0):min(maximum(nlv), a))
    le_nlv = length(nlv)
    pred = list(le_nlv, Matrix{Float64})
    @inbounds for i = 1:le_nlv
        z = coef(object; nlv = nlv[i])
        pred[i] = z.int .+ X * z.B
    end
    le_nlv == 1 ? pred = pred[1] : nothing
    (pred = pred,)
end


#=##############################################################################
# DESCRIPTION
    Windcraft vehicle kinematic maneuver definition.

# AUTHORSHIP
  * Author    : Judd Mehr and Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Apr 2020
  * License   : MIT
=###############################################################################

function generate_maneuver_windcraft_kinematic(nrevs;
                                                 disp_plot=false,
                                                 includetail=true,
                                                 includewing=true,
                                                 includecontrols=true,
                                                 includerotors=true,
                                                 optargs...)

    omegamean = pi/180 * mean(omegastar.(range(0, 1, length=361); optargs...))

    ############################################################################
    # AIRCRAFT VELOCITY
    ############################################################################
    """
     Receives a nondimensional time between 0 and 1, and returns the
     vector of velocity of the vehicle at that instant.
    """
    function Vvehicle(t)

        theta = thetastar_periodic(t, nrevs; optargs...)*pi/180
        omega = omegastar_periodic(t, nrevs; optargs...)*pi/180
        scaling = omega/omegamean               # Scales velocity to a maximum value of 1

        Vcomp = [0.0, cos(theta), -sin(theta)]  # Counter-clockwise rotation
        return scaling*Vcomp
    end

    ############################################################################
    # AIRCRAFT ANGLES
    ############################################################################
    """
        Receives a nondimensional time between 0 and 1, and returns the angle
        (in degrees) of the vehicle.
        Returns: (angle_aircraft_x, angle_aircraft_y, angle_aircraft_z)
    """
    function anglevehicle(t)

        angle = [0.0, 0.0, thetastar_periodic(abs(t), nrevs; optargs...)]

        if t < 0.0
            return -angle
        else
            return angle
        end

    end

    ############################################################################
    # ANGLES OF CONTROL SURFACES
    ############################################################################

    """
        Receives a nondimensional time between 0 and 1, and returns the angles
        (in degrees) of the tilting system relative to the aircraft axis.
    """
    function angle_default(t)
        return [0, 0, 0]
    end

    if includetail == true
        angle = Tuple(angle_default for i in 1:6)  # Angle of each tilting system

    elseif includewing == true && includecontrols == true
        angle = Tuple(angle_default for i in 1:4)

    else
        angle = ()
    end

    ############################################################################
    # ROTORS RPM
    ############################################################################
    """
        Receives a nondimensional time between 0 and 1, and returns the RPM
        non-dimensionalized by the hover RPM.
    """
    function RPM_default(t)
        tsr = 4.0
        R = 1.0

        v = Vvehicle(t)
        vmag = sqrt(v[1]^2 + v[2]^2 + v[3]^2)
        omega = tsr*vmag/R
        rpm = omega*pi/30

        return rpm
    end

    if includerotors == true
        RPM = (RPM_default,)        # RPM of each rotor system
    else
        RPM = ()
    end

    ############################################################################
    # MANEUVER OBJECT
    ############################################################################
    maneuver = uns.KinematicManeuver(angle, RPM, Vvehicle, anglevehicle)

    if disp_plot
        uns.plot(maneuver)
    end

    return maneuver

end

"""
angle theta output in degrees
"""
function thetastar(tstar; theta0=0.0, thetan=360.0, omegan=[4.0/9.0;20.0/27.0;4.0/9.0], tn=[0.0;0.5;1.0])
    #find index of max tn less than tstar and set to index variable m.
    if tstar==tn[end]
        m = length(tn)
    elseif tstar==tn[1]
        m = 1
    else

        m = maximum(find(x->x<tstar,tn))
    end
    # m = searchsortednearest(tn,tstar)
    #loop from 1 to m and sum areas.
    areas = 0.0
    den = 0.0
    for i = 2:length(tn)
        if i <= m
            areas += (tn[i]-tn[i-1])*(omegan[i]+omegan[i-1])/2.0
        end
        den += (tn[i]-tn[i-1])*(omegan[i]+omegan[i-1])/2.0
    end

    #add integral of final partial area up to tstar
    if m < length(tn) #if we haven't already integrated over everything
        integral = ( tstar - tn[m] ) * ( omegan[m] + (omegan[m+1]-omegan[m])/2.0 * (tstar-tn[m])/(tn[m+1]-tn[m]) )
    else
        integral = 0.0
    end

    #solve for scale factor, which
    sf = thetan / den
    #add theta0 (constant from integration)
    theta = sf * (areas + integral) + theta0

    return theta
end

function omegastar(tstar; h=1e-8, optargs...)
    return (thetastar(tstar+h; optargs...) - thetastar(tstar; optargs...))/h
end

"""
Transform time of thetastar and make it periodic to extend the angle output
outside the range [0, 360]
"""
function thetastar_periodic(t, nrevs; optargs...)
    tstar = (t*nrevs)%1    # Convert general time to time of one revolution
    return 360*floor(t*nrevs) + thetastar(tstar; optargs...)
end

function omegastar_periodic(tstar, nrevs; h=1e-8, optargs...)
    return (thetastar_periodic(tstar+h, nrevs; optargs...)
                - thetastar_periodic(tstar, nrevs; optargs...))/h
end


# -------------------------------------------------
#
# An explicit tensor constructor of the coefficients
# for \prod_i f_i (x_i) given the coefficients of
# each f(x_i).
#
# -------------------------------------------------

# This is a useful way to reduce a difficult O(N log^(d-1) N)
# space of integrations down to O(N d)
# for specifying coefficients of initial conditions

# Efficiency criticality: MEDIUM

# Accuracy criticality: LOW
# Most of the accuracy dependent on functions
# called from other scripts

function tensor_construct(D::Int, k::Int, n::Int, coeff_array::Array{Dict{CartesianIndex{1},
    Array{Array{T, 1}, 1}}, 1}; scheme = "sparse") where T <: Real
    tensor_construct(Val(D), k, n, coeff_array, Val(Symbol(scheme)))
end
function tensor_construct(::Val{D}, k::Int, n::Int, coeff_array::Array{Dict{CartesianIndex{1},
    Array{Array{T, 1}, 1}}, 1}, scheme::Val{Scheme}) where {D, T <: Real, Scheme}

    coeffs = Dict{CartesianIndex{D}, Array{Array{T,D},D}}()
    ls::NTuple{D, Int}    = ntuple(i-> (n+1), D)
    modes::NTuple{D, Int} = ntuple(i-> k, D)

    for level in CartesianIndices(ls)
        cutoff(scheme, level, n) && continue

        cells::NTuple{D, Int} = ntuple(i -> 1<<max(0, level[i]-2), D)
        level_coeffs = Array{Array{T,D}}(undef, cells)
        for cell in CartesianIndices(cells)
            cell_coeffs = Array{T}(undef, modes)
            for mode in CartesianIndices(modes)
                val = one(T)
                for d in 1:D
                    coeff = coeff_array[d]
                    tup::NTuple{1, Int} = (level[d],) # This is somehow still slow
                    l = CartesianIndex{1}(tup);
                    c = cell[d]
                    m = mode[d]
                    val *= coeff[l][c][m]
                end
                cell_coeffs[mode] = val
            end
            level_coeffs[cell] = cell_coeffs
        end
        coeffs[level] = level_coeffs
    end
    return coeffs
end

function tensor_construct(D::Int, k::Int, n::Int, vcoeff_array::Array{Array{T, 1}};
							scheme = "sparse") where T <: Real

	coeff_array = [V2D(1, k, n, vcoeff_array[i]) for i in 1:length(vcoeff_array)]
	return D2V(D, k, n, tensor_construct(D, k, n, coeff_array); scheme=scheme)

end


# Time series analysis of 2020 worldwide coronavirus outbreak
#
# data based on John Hopkins CSSE repository
# https://github.com/CSSEGISandData/COVID-19

module C

using  Dates, DSP, PyPlot, DifferentialEquations, Polynomials

include("read_data.jl")

function smooth_log(x)
# make log-transform and smooth data
    lx = log10.(x)
    b = (1/2)*ones(2)     # flat (rectangular) window -> moving average
    lx = filtfilt(b,lx)
    10 .^ lx, lx
end

function local_fit(x,y,di)
# local slope of graph y(x) from regression in running window
# window width: 2*di+1
  len = length(x)
  slope = zeros(len)
  for i=1:len
    i1 = max(1,i-di)
    i2 = min(len,i+di)
    p = polyfit(x[i1:i2],y[i1:i2],1)
    slope[i] = coeffs(p)[2]
  end
  slope, log(10.) .* slope
end

# doubling time when slope is taken from log10-transformed data
doubling_time(slope) = log(2.)/log(10.)/slope

function nlys()

   today = " March 19"

   t_de,I_de,R_de,D_de = covid_ts("Germany",       Date(2020,2,23))
   t_it,I_it,R_it,D_it = covid_ts("Italy",         Date(2020,2,20))
   t_es,I_es,R_es,D_es = covid_ts("Spain",         Date(2020,2,25))
   t_fr,I_fr,R_fr,D_fr = covid_ts("France",        Date(2020,2,25))
   t_uk,I_uk,R_uk,D_uk = covid_ts("United Kingdom",Date(2020,2,27))
   t_at,I_at,R_at,D_at = covid_ts("Austria",       Date(2020,2,29))
   t_nl,I_nl,R_nl,D_nl = covid_ts("Netherlands",   Date(2020,2,23))
   t_ch,I_ch,R_ch,D_ch = covid_ts("Switzerland",   Date(2020,2,25))
   t_cn,I_cn,R_cn,D_cn = covid_ts("China",         Date(2020,1,22))
   t_jp,I_jp,R_jp,D_jp = covid_ts("Japan",         Date(2020,1,24))
   t_kr,I_kr,R_kr,D_kr = covid_ts("Korea, South",  Date(2020,2,17))
   t_us,I_us,R_us,D_us = covid_ts("US",            Date(2020,2,22))
   t_ir,I_ir,R_ir,D_ir = covid_ts("Iran",          Date(2020,2,19))
   t_il,I_il,R_il,D_il = covid_ts("Israel",        Date(2020,2,27))
   t_no,I_no,R_no,D_no = covid_ts("Norway",        Date(2020,2,28))
   t_ca,I_ca,R_ca,D_ca = covid_ts("Canada",        Date(2020,2,27))

  I_de, lI_de = smooth_log(I_de)
  I_it, lI_it = smooth_log(I_it)
  I_fr, lI_fr = smooth_log(I_fr)
  I_es, lI_es = smooth_log(I_es)
  I_jp, lI_jp = smooth_log(I_jp)
  I_cn, lI_cn = smooth_log(I_cn)
  I_kr, lI_kr = smooth_log(I_kr)
  I_ir, lI_ir = smooth_log(I_ir)
  I_us, lI_us = smooth_log(I_us)
  I_ch, lI_ch = smooth_log(I_ch)
  I_il, lI_il = smooth_log(I_il)
  I_at, lI_at = smooth_log(I_at)
  I_uk, lI_uk = smooth_log(I_uk)
  I_nl, lI_nl = smooth_log(I_nl)
  I_no, lI_no = smooth_log(I_no)
  I_ca, lI_ca = smooth_log(I_ca)

  di = 2 # length of running window: 2*di+1
  slope_de,growth_de = local_fit(t_de,lI_de,di)
  slope_it,growth_it = local_fit(t_it,lI_it,di)
  slope_es,growth_es = local_fit(t_es,lI_es,di)
  slope_fr,growth_fr = local_fit(t_fr,lI_fr,di)
  slope_jp,growth_jp = local_fit(t_jp,lI_jp,di)
  slope_cn,growth_cn = local_fit(t_cn,lI_cn,di)
  slope_kr,growth_kr = local_fit(t_kr,lI_kr,di)
  slope_ir,growth_ir = local_fit(t_ir,lI_ir,di)
  slope_us,growth_us = local_fit(t_us,lI_us,di)
  slope_ch,growth_ch = local_fit(t_ch,lI_ch,di)
  slope_at,growth_at = local_fit(t_at,lI_at,di)
  slope_il,growth_il = local_fit(t_il,lI_il,di)
  slope_uk,growth_uk = local_fit(t_uk,lI_uk,di)
  slope_nl,growth_nl = local_fit(t_nl,lI_nl,di)
  slope_no,growth_no = local_fit(t_no,lI_no,di)
  slope_ca,growth_ca = local_fit(t_ca,lI_ca,di)

  p_tot_de1 = polyfit(t_de[4:12],lI_de[4:12],1)  # fit full slope Germany
  coef = coeffs(p_tot_de1)[2]
  println("slope Germany ",coef, " T12 ", doubling_time(coef))
  lI_fit_de1= polyval(p_tot_de1, t_de)

  p_tot_de2 = polyfit(t_de[12:end],lI_de[12:end],1)  # fit full slope Germany
  coef = coeffs(p_tot_de2)[2]
  println("slope Germany ",coef, " T12 ", doubling_time(coef))
  lI_fit_de2= polyval(p_tot_de2, t_de)

  p_tot_it1 = polyfit(t_it[4:9],lI_it[4:9],1)  # fit slope 1st phase Italy
  coef = coeffs(p_tot_it1)[2]
  println("slope Italy Phase 1 ",coef, " T12 ", doubling_time(coef))
  lI_fit_it1 = polyval(p_tot_it1, t_it)

  p_tot_it = polyfit(t_it[10:end],lI_it[10:end],1)  # fit slope 2nd phase Italy
  coef = coeffs(p_tot_it)[2]
  println("slope Italy Phase 2 ",coef, " T12 ", doubling_time(coef))
  lI_fit_it = polyval(p_tot_it, t_it)

  p_tot_es1 = polyfit(t_es[6:18],lI_es[6:18],1)  # fit full slope Spain
  lI_fit_es1= polyval(p_tot_es1, t_es)
  p_tot_es2 = polyfit(t_es[19:24],lI_es[19:24],1)  # fit full slope Spain
  lI_fit_es2= polyval(p_tot_es2, t_es)

  p_tot_fr1 = polyfit(t_fr[3:14],lI_fr[3:14],1)  # fit full slope France
  lI_fit_fr1= polyval(p_tot_fr1, t_fr)
  p_tot_fr2 = polyfit(t_fr[14:end],lI_fr[14:end],1)  # fit full slope France
  lI_fit_fr2= polyval(p_tot_fr2, t_fr)

  p_tot_cn = polyfit(t_cn,lI_cn,1)  # fit full slope China
  lI_fit_cn= polyval(p_tot_cn, t_cn)

  p_tot_kr = polyfit(t_kr,lI_kr,1)  # fit full slope South Korean
  lI_fit_kr= polyval(p_tot_kr, t_kr)

  p_tot_jp = polyfit(t_jp,lI_jp,1)  # fit full slope Japan
  lI_fit_jp= polyval(p_tot_jp, t_jp)

  p_tot_ir = polyfit(t_ir,lI_ir,1)  # fit full slope Iran
  lI_fit_ir= polyval(p_tot_ir, t_ir)

  p_tot_us = polyfit(t_us,lI_us,1)  # fit full slope USA
  lI_fit_us= polyval(p_tot_us, t_us)

  p_tot_ch = polyfit(t_ch[11:22],lI_ch[11:22],1)  # fit full slope Switzerland
  lI_fit_ch= polyval(p_tot_ch, t_ch)

  p_tot_at = polyfit(t_at,lI_at,1)  # fit full slope Austria
  lI_fit_at= polyval(p_tot_at, t_at)

  p_tot_il = polyfit(t_il,lI_il,1)  # fit full slope Israel
  lI_fit_il= polyval(p_tot_il, t_il)

  p_tot_uk = polyfit(t_uk,lI_uk,1)  # fit full slope UK
  lI_fit_uk= polyval(p_tot_uk, t_uk)

  p_tot_no = polyfit(t_no,lI_no,1)  # fit full slope Norway
  lI_fit_no= polyval(p_tot_no, t_no)

  p_tot_nl = polyfit(t_nl[15:23],lI_nl[15:23],1)  # fit full slope Netherlands
  lI_fit_nl= polyval(p_tot_nl, t_nl)

  p_tot_ca = polyfit(t_ca,lI_ca,1)  # fit full slope Norway
  lI_fit_ca= polyval(p_tot_ca, t_ca)

  #majorformatter = matplotlib.dates.DateFormatter("%d.%m.%Y")
  #majorlocator = matplotlib.dates.DayLocator(interval=1)

# ********************************************
  fig = figure("Outbreak Europe",figsize=(12,8))
  clf()
  subplot(3,4,1)         # Germany
      plot(t_de,I_de,"r*")
      ylabel("active cases",fontsize=12)

      title("Germany,"* today)
  ax=subplot(3,4,5)
     # p1 = plot_date(Date(2020,2,17),y,linestyle="-",marker="None",label="test")
      semilogy(t_de,I_de,"r*")
      semilogy(t_de,10 .^ lI_fit_de1,"k")
      semilogy(t_de,10 .^ lI_fit_de2,"k")
      ax.yaxis.grid()
      ylim(5,5e4)
      ylabel("active cases",labelpad=20,fontsize=12)
   ax = subplot(3,4,9)
      plot(t_de,growth_de,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      #ax.yaxis.grid(which="major", color="r", linestyle="-", linewidth=2)
      ax.yaxis.grid()
     # annotate("slope 0.107",xy=[20,2],
     # xycoords="axes fraction"
     # )
      #ax.xaxis.set_major_formatter(majorformatter)
      #fig.autofmt_xdate(bottom=0.2,rotation=30,ha="right")
      xlabel("time (days)",fontsize=12)
      ylabel("growth rate",labelpad=20,fontsize=12)

  subplot(3,4,2)          # Italy
      plot(t_it,I_it,"r*")
      #ylabel("infected Italy")
      title("Italy,"*today)
  ax=subplot(3,4,6)
      semilogy(t_it,I_it,"r*")
      semilogy(t_it,10 .^ lI_fit_it1,"k")
      semilogy(t_it,10 .^ lI_fit_it,"k")
      ylim(5,5e4)
      ax.yaxis.grid()
  ax=subplot(3,4,10)
      plot(t_it,growth_it,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("slope")

  subplot(3,4,3)              # Spain
      plot(t_es,I_es,"r*")
      title("Spain,"*today)
  ax=subplot(3,4,7)
      semilogy(t_es,I_es,"r*")
      semilogy(t_es,10 .^ lI_fit_es1,"k")
      semilogy(t_es,10 .^ lI_fit_es2,"k")
      ax.yaxis.grid()
      ylim(5,5e4)
   ax= subplot(3,4,11)
      plot(t_es,growth_es,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)

  subplot(3,4,4)              # France
      plot(t_fr,I_fr,"r*")
      title("France,"*today)
  ax=subplot(3,4,8)
      semilogy(t_fr,I_fr,"r*")
      semilogy(t_fr,10 .^ lI_fit_fr1,"k")
      semilogy(t_fr,10 .^ lI_fit_fr2,"k")
      ax.yaxis.grid()
      ylim(5,5e4)
  ax = subplot(3,4,12)
      plot(t_fr,growth_fr,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  PyPlot.tight_layout()
  savefig("../figures/outbreak_europe.png")

# ********************************************
  fig=figure("Outbreak Asia",figsize=(12,8))
  clf()
  subplot(3,4,1)         # China
      plot(t_cn,I_cn,"r*")
      ylabel("active cases",fontsize=12)
      title("China,"*today)
  ax=subplot(3,4,5)
      semilogy(t_cn,I_cn,"r*")
      #plot(t_cn,lI_fit_cn,"k")
      ylim(1e2,1e5)
      ax.yaxis.grid()
      ylabel("active cases", labelpad=20,fontsize=12)
  ax = subplot(3,4,9)
      plot(t_cn,growth_cn,"r*-")
      ylim(-0.1,0.5)
      #yticks([-0.05,0,0.05,0.1,0.15,0.2])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      ylabel("growth",fontsize=12,labelpad=10)

  subplot(3,4,2)          # Japan
      plot(t_jp,I_jp,"r*")
      #ylabel("infected Italy")
      title("Japan,"*today)
  ax=subplot(3,4,6)
      semilogy(t_jp,I_jp,"r*")
      semilogy(t_jp,10 .^ lI_fit_jp,"k")
      ax.yaxis.grid()
      ylim(5,2e4)
      #ylabel("log10 infected Italy")
  ax=subplot(3,4,10)
      plot(t_jp,growth_jp,"r*-")
      ylim(-0.1,0.5)
      #yticks([-0.05,0,0.05,0.1,0.15,0.2])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  subplot(3,4,3)              # South Korean
      plot(t_kr,I_kr,"r*")
      #ylabel("infected Spain")
      title("South Korea,"*today)
  ax=subplot(3,4,7)
      semilogy(t_kr,I_kr,"r*")
      ax.yaxis.grid()
      ylim(5,2e4)
      #plot(t_kr,lI_fit_kr,"k")
      #ylabel("log10 infected Spain")
  ax= subplot(3,4,11)
      plot(t_kr,growth_kr,"r*-")
      ylim(-0.1,0.5)
      #yticks([-0.05,0,0.05,0.1,0.15,0.2])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  subplot(3,4,4)              # Iran
      plot(t_ir,I_ir,"r*")
      #ylabel("infected")
      title("Iran,"*today)
  ax=subplot(3,4,8)
      semilogy(t_ir,I_ir,"r*")
      semilogy(t_ir,10 .^ lI_fit_ir,"k")
      ax.yaxis.grid()
      ylim(5,2e4)
      #ylabel("log10 infected")
  ax = subplot(3,4,12)
      plot(t_ir,growth_ir,"r*-")
      ylim(-0.1,0.5)
      #yticks([-0.05,0,0.05,0.1,0.15,0.2])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")
  PyPlot.tight_layout()
  savefig("../figures/outbreak_Asia.png")

# ********************************************
  fig=figure("Outbreak World",figsize=(12,8))
  clf()
  subplot(3,4,1)         # USA
      plot(t_us,I_us,"r*")
      ylabel("active cases",fontsize=12)
      title("USA,"*today)
  ax=subplot(3,4,5)
      semilogy(t_us,I_us,"r*")
      semilogy(t_us,10 .^ lI_fit_us,"k")
      ax.yaxis.grid()
      ylim(5,2e4)
      ylabel("active cases",fontsize=12,labelpad=12)
  ax = subplot(3,4,9)
      plot(t_us,growth_us,"r*-")
      ax.yaxis.grid()
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      xlabel("time (days)",fontsize=12)
      ylabel("growth",fontsize=12,labelpad=12)

  subplot(3,4,2)              # Canada
      plot(t_ca,I_ca,"r*")
      #ylabel("infected")
      title("Canada,"* today)
  ax=subplot(3,4,6)
      semilogy(t_ca,I_ca,"r*")
      semilogy(t_ca,10 .^ lI_fit_ca,"k")
      ax.yaxis.grid()
      ylim(5,2e4)
      #ylabel("log10 infected")
  ax = subplot(3,4,10)
      plot(t_ca,growth_ca,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  subplot(3,4,3)              # UK
      plot(t_uk,I_uk,"r*")
      #ylabel("infected Spain")
      title("UK,"*today)
  ax=subplot(3,4,7)
      semilogy(t_uk,I_uk,"r*")
      semilogy(t_uk,10 .^ lI_fit_uk,"k")
      ax.yaxis.grid()
      ylim(5,2e4)
      #ylabel("log10 infected Spain")
  ax= subplot(3,4,11)
      plot(t_uk,growth_uk,"r*-")
      ax.yaxis.grid()
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  subplot(3,4,4)         # Israel
      plot(t_il,I_il,"r*")
      title("Israel,"*today)
  ax=subplot(3,4,8)
      semilogy(t_il,I_il,"r*")
      semilogy(t_il,10 .^ lI_fit_il,"k")
      ax.yaxis.grid()
      ylim(1,1e3)
  ax = subplot(3,4,12)
      plot(t_il,growth_il,"r*-")
      ax.yaxis.grid()
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      xlabel("time (days)",fontsize=12)

  PyPlot.tight_layout()
  savefig("../figures/outbreak_World.png")

# ********************************************
  fig=figure("Outbreak Small Europe",figsize=(12,8))
  clf()
  subplot(3,4,1)         # Netherlands
      plot(t_nl,I_nl,"r*")
      ylabel("active cases",fontsize=12)

      title("Netherlands,"* today)
  ax=subplot(3,4,5)
      semilogy(t_nl,I_nl,"r*")
      semilogy(t_nl,10 .^ lI_fit_nl,"k")
      ax.yaxis.grid()
      ylim(5,1e4)
      ylabel("active cases",labelpad=16,fontsize=12)
   ax = subplot(3,4,9)
      plot(t_nl,growth_nl,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      #ax.yaxis.grid(which="major", color="r", linestyle="-", linewidth=2)
      ax.yaxis.grid()
     # annotate("growth 0.107",xy=[20,2],
     # xycoords="axes fraction"
     # )
      #ax.xaxis.set_major_formatter(majorformatter)
      #fig.autofmt_xdate(bottom=0.2,rotation=30,ha="right")
      xlabel("time (days)",fontsize=12)
      ylabel("growth",fontsize=12,labelpad=20)

  subplot(3,4,2)          # Norway
      plot(t_no,I_no,"r*")
      #ylabel("infected Italy")
      title("Norway,"*today)
  ax=subplot(3,4,6)
      semilogy(t_no,I_no,"r*")
      semilogy(t_no,10 .^ lI_fit_no,"k")
      ax.yaxis.grid()
      ylim(5,1e4)
  ax=subplot(3,4,10)
      plot(t_no,growth_no,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  subplot(3,4,3)              # Switzerland
      plot(t_ch,I_ch,"r*")
      #ylabel("infected Spain")
      title("Switzerland,"* today)
  ax=subplot(3,4,7)
      semilogy(t_ch,I_ch,"r*")
      semilogy(t_ch,10 .^ lI_fit_ch,"k")
      ax.yaxis.grid()
      ylim(5,1e4)
      #ylabel("log10 infected Spain")
  ax= subplot(3,4,11)
      plot(t_ch,growth_ch,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")

  subplot(3,4,4)              # Austria
      plot(t_at,I_at,"r*")
      #ylabel("infected")
      title("Austria,"* today)
  ax=subplot(3,4,8)
      semilogy(t_at,I_at,"r*")
      semilogy(t_at,10 .^ lI_fit_at,"k")
      ax.yaxis.grid()
      ylim(5,1e4)
      #ylabel("log10 infected")
  ax = subplot(3,4,12)
      plot(t_at,growth_at,"r*-")
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2,0.25])
      ax.yaxis.grid()
      xlabel("time (days)",fontsize=12)
      #ylabel("growth")
  PyPlot.tight_layout()
  savefig("../figures/outbreak_europe1.png")

# ********************************************

   lt_it = log10.(t_it)
   p_tot_it = polyfit(lt_it,lI_it,1)  # fit full slope Germany
   println("loglog It ", coeffs(p_tot_it)[2])
   I_fit_it = 10 .^ polyval(p_tot_it, lt_it)

   lt_de = log10.(t_de)
   p_tot_de = polyfit(lt_de,lI_de,1)  # fit full slope Germany
   println("loglog It ", coeffs(p_tot_de)[2])
   I_fit_de = 10 .^ polyval(p_tot_de, lt_de)

   figure("ts")
   clf()
   subplot(2,1,1)
   loglog(t_it,I_it,"*-")
   loglog(t_it,I_fit_it,"-")

   subplot(2,1,2)
   loglog(t_de,I_de,"*-")
   loglog(t_de,I_fit_de,"-")
   #semilogy(t_es,I_es,"r*-")

end

# **********************************************************
# **********************************************************
# **********************************************************

function nlys_ItDe()
  today = " March 18"
  t_it,I_it,R_it,D_it = covid_ts("Italy",Date(2020,2,21))
  I_it, lI_it = smooth_log(I_it)
  di = 2
  slope_it,growth_it = local_fit(t_it,lI_it,di)

  p_tot_it1 = polyfit(t_it[3:8],lI_it[3:8],1)  # fit slope 1st phase Italy
  coef = coeffs(p_tot_it1)[2]
  println("slope Italy Phase 1 ",coef, " T12 ", doubling_time(coef))
  lI_fit_it1 = polyval(p_tot_it1, t_it)

  p_tot_it2 = polyfit(t_it[9:22],lI_it[9:22],1)  # fit slope 2nd phase Italy
  coef = coeffs(p_tot_it2)[2]
  println("slope Italy Phase 2 ",coef, " T12 ", doubling_time(coef))
  lI_fit_it2 = polyval(p_tot_it2, t_it)

  p_tot_it3 = polyfit(t_it[23:end],lI_it[23:end],1)  # fit slope 2nd phase Italy
  coef = coeffs(p_tot_it3)[2]
  println("slope Italy Phase 3 ",coef, " T12 ", doubling_time(coef))
  lI_fit_it3 = polyval(p_tot_it3, t_it)

  # ********************
  t_de,I_de,R_de,D_de = covid_ts("Germany",  Date(2020,2,24))
  t_de = t_de .+ 3 # bring to the same time axis
  I_de, lI_de = smooth_log(I_de)
  di = 2
  slope_de, growth_de = local_fit(t_de,lI_de,di)

  p_tot_de1 = polyfit(t_de[4:11],lI_de[4:11],1)  # fit full slope Germany
  coef = coeffs(p_tot_de1)[2]
  println("slope Germany ",coef, " T12 ", doubling_time(coef))
  lI_fit_de1= polyval(p_tot_de1, t_de)

  p_tot_de2 = polyfit(t_de[12:end],lI_de[12:end],1)  # fit full slope Germany
  coef = coeffs(p_tot_de2)[2]
  println("slope Germany ",coef, " T12 ", doubling_time(coef))
  lI_fit_de2= polyval(p_tot_de2, t_de)

   figure("Germany_Italy", figsize=(12,6))
   clf()
   subplot(1,2,1)
      semilogy(t_it[1:12],10 .^ lI_fit_it1[1:12],"k")
      semilogy(t_it[5:25],10 .^ lI_fit_it2[5:25],"k")
      semilogy(t_it[19:28],10 .^ lI_fit_it3[19:28],"k")
      semilogy(t_it,I_it,"b*", label="Italy", markersize=8)
      semilogy(t_de[1:15],10 .^ lI_fit_de1[1:15],"k")
      semilogy(t_de[8:end],10 .^ lI_fit_de2[8:end],"k")
      semilogy(t_de,I_de,"r*",label="Germany",markersize=8)
      #plot([10,10],[8,5e4],"grey",linewidth=1)
      #ylim(0.5,5)
      ylim(8,5e4)
      ylabel("active cases",fontsize=14)
      xlabel("time (Feb 21 - March 19)", fontsize=14)
      legend(loc="upper left",fancybox="true")
   ax = subplot(1,2,2)
      plot(t_it,growth_it,"b*-",label="Italy",markersize=8)
      plot(t_de,growth_de,"r*-",label="Germany",markersize=8)
      #plot([10,10],[0,0.6],"grey",linewidth=1)
      #ylim(0,0.22)
      #yticks([0,0.05,0.1,0.15,0.2])
      ylim(0,0.6)
      #yticks([0,0.05,0.1,0.15,0.2])
      ax.yaxis.grid()
      #xlabel("time (days)",fontsize=14)
      xlabel("time (Feb 21 - March 18)", fontsize=14)
      ylabel("growth rate",fontsize=14)
      legend(loc="upper right",fancybox="true")
  PyPlot.tight_layout()
  savefig("../figures/ItalyGermany_March19.png")

end

#country = "Italy"; date = Date(2020,2,20)

#nlys_country(country, date)
nlys_ItDe()

nlys()

end   # module