module ParameterDistributions

## Usings
using Distributions
using Statistics
using Random

#import (to add definitions)
import StatsBase: mean, var, cov, sample
import Base: size, length, ndims
## Exports

#types
export ParameterDistributionType
export ConstraintType

#objects
export Parameterized, Samples, VectorOfParameterized
export ParameterDistribution
export Constraint

#functions
export get_name, get_distribution, ndims, get_dimensions, get_all_constraints, get_n_samples
export sample
export no_constraint, bounded_below, bounded_above, bounded
export transform_constrained_to_unconstrained, transform_unconstrained_to_constrained
export get_logpdf, batch

export combine_distributions
## Objects
# for the Distribution
abstract type ParameterDistributionType end

"""
    Parameterized <: ParameterDistributionType

A distribution constructed from a parametrized formula (e.g Julia Distributions.jl)
"""
struct Parameterized <: ParameterDistributionType
    distribution::Distribution
end

"""
    Samples{FT <: Real} <: ParameterDistributionType

A distribution comprised of only samples, stored as columns of parameters.
"""
struct Samples{FT <: Real} <: ParameterDistributionType
    distribution_samples::AbstractMatrix{FT} #parameters are columns
    Samples(distribution_samples::AbstractMatrix{FT}; params_are_columns = true) where {FT <: Real} =
        params_are_columns ? new{FT}(distribution_samples) : new{FT}(permutedims(distribution_samples, (2, 1)))
    #Distinguish 1 sample of an ND parameter or N samples of 1D parameter, and store as 2D array
    Samples(distribution_samples::AbstractVector{FT}; params_are_columns = true) where {FT <: Real} =
        params_are_columns ? new{FT}(reshape(distribution_samples, 1, :)) : new{FT}(reshape(distribution_samples, :, 1))
end

"""
    VectorOfParameterized <: ParameterDistributionType

A distribution built from an array of Parametrized distributions.
A utility to help stacking of distributions where a multivariate equivalent doesn't exist
"""
struct VectorOfParameterized{DT <: Distribution} <: ParameterDistributionType
    distribution::AbstractVector{DT}
end

# For the transforms
abstract type ConstraintType end

"""
    Constraint <: ConstraintType

Contains two functions to map between constrained and unconstrained spaces.
"""
struct Constraint <: ConstraintType
    constrained_to_unconstrained::Function
    unconstrained_to_constrained::Function
end

"""
    no_constraint()

Constructs a Constraint with no constraints, enforced by maps x -> x and x -> x.
"""
function no_constraint()
    c_to_u = (x -> x)
    u_to_c = (x -> x)
    return Constraint(c_to_u, u_to_c)
end

"""
    bounded_below(lower_bound::FT) where {FT <: Real}

Constructs a Constraint with provided lower bound, enforced by maps `x -> log(x - lower_bound)`
and `x -> exp(x) + lower_bound`.
"""
function bounded_below(lower_bound::FT) where {FT <: Real}
    c_to_u = (x -> log(x - lower_bound))
    u_to_c = (x -> exp(x) + lower_bound)
    return Constraint(c_to_u, u_to_c)
end

"""
    bounded_above(upper_bound::FT) where {FT <: Real}

Constructs a Constraint with provided upper bound, enforced by maps `x -> log(upper_bound - x)`
and `x -> upper_bound - exp(x)`.
"""
function bounded_above(upper_bound::FT) where {FT <: Real}
    c_to_u = (x -> log(upper_bound - x))
    u_to_c = (x -> upper_bound - exp(x))
    return Constraint(c_to_u, u_to_c)
end

"""
    bounded(lower_bound::FT, upper_bound::FT) where {FT <: Real}

Constructs a Constraint with provided upper and lower bounds, enforced by maps
`x -> log((x - lower_bound) / (upper_bound - x))`
and `x -> (upper_bound * exp(x) + lower_bound) / (exp(x) + 1)`.

"""
function bounded(lower_bound::FT, upper_bound::FT) where {FT <: Real}
    if (upper_bound <= lower_bound)
        throw(DomainError("upper bound must be greater than lower bound"))
    end
    c_to_u = (x -> log((x - lower_bound) / (upper_bound - x)))
    u_to_c = (x -> (upper_bound * exp(x) + lower_bound) / (exp(x) + 1))
    return Constraint(c_to_u, u_to_c)
end

#extending Base.length
"""
    length(c<:ConstraintType)

A constraint has length 1.
"""
length(c::CType) where {CType <: ConstraintType} = length([c])

#extending Base.size
"""
    size(c<:ConstraintType)

A constraint has size 1.
"""
size(c::CType) where {CType <: ConstraintType} = size([c])

"""
    ndims(d<:ParametrizedDistributionType)

The number of dimensions of the parameter space
"""
ndims(d::Parameterized) = length(d.distribution)

ndims(d::Samples) = size(d.distribution_samples, 1)

ndims(d::VectorOfParameterized) = sum(length.(d.distribution))
"""
    n_samples(d<:Samples)

The number of samples in the array.
"""
n_samples(d::Samples) = size(d.distribution_samples)[2]

n_samples(d::Parameterized) = "Distribution stored in Parameterized form, draw samples using `sample` function"

n_samples(d::VectorOfParameterized) = "Distribution stored in Parameterized form, draw samples using `sample` function"

"""
    ParameterDistribution

Structure to hold a parameter distribution, always stored as an array internally.
"""
struct ParameterDistribution{PDType <: ParameterDistributionType, CType <: ConstraintType, ST <: AbstractString}
    distribution::AbstractVector{PDType}
    constraint::AbstractVector{CType}
    name::AbstractVector{ST}
end

"""
    ParameterDistribution(param_dist_dict::Union{Dict,AbstractVector})

Constructor taking in a Dict or array of Dicts. Each dict must contain the key-val pairs:
- `"distribution"` - a distribution of `ParameterDistributionType`
- `"constraint"` - constraint(s) given as a `ConstraintType` or array of `ConstraintType`s with length equal to the dims of the distribution
- `"name"` - a name of the distribution as a String.
"""
function ParameterDistribution(param_dist_dict::Union{Dict, AbstractVector})

    #check type
    if !(isa(param_dist_dict, Dict) || eltype(param_dist_dict) <: Dict)
        throw(ArgumentError("input argument must be a Dict, or <:AbstractVector{Dict}"))
    end

    # make copy as array
    param_dist_dict_array = !isa(param_dist_dict, AbstractVector) ? [param_dist_dict] : param_dist_dict
    # perform checks on the individual distributions
    for pdd in param_dist_dict_array
        # check all keys are present
        if !all(["distribution", "name", "constraint"] .∈ [collect(keys(pdd))])
            throw(
                ArgumentError("input dictionaries must contain the keys: \"distribution\", \"name\", \"constraint\" "),
            )
        end

        distribution = pdd["distribution"]
        name = pdd["name"]
        constraint = pdd["constraint"]

        # check key types
        if !isa(distribution, ParameterDistributionType)
            throw(
                ArgumentError(
                    "Value of \"distribution\" must be a valid ParameterDistribution object: Parameterized or Samples",
                ),
            )
        end
        if !isa(constraint, ConstraintType)
            if !isa(constraint, AbstractVector) #it's not a vector either
                throw(
                    ArgumentError(
                        "Value of \"constraint\" must be a ConstraintType, or <:AbstractVector(ConstraintType)",
                    ),
                )
            elseif !(eltype(constraint) <: ConstraintType) #it is a vector, but not of constraint
                throw(
                    ArgumentError(
                        "Value of \"constraint\" must be a ConstraintType, or <:AbstractVector(ConstraintType)",
                    ),
                )
            end
        end
        if !isa(name, String)
            throw(ArgumentError("Value of \"name\" must be a String"))
        end

        # 1 constraint per dimension check
        constraint_array = isa(constraint, ConstraintType) ? [constraint] : constraint
        n_parameters = ndims(distribution)

        if !(n_parameters == length(constraint_array))
            throw(
                DimensionMismatch(
                    "There must be one constraint dimension in a parameter distribution, use no_constraint() type if no constraint is required",
                ),
            )
        end
    end

    # flatten the structure
    distribution = getindex.(param_dist_dict_array, "distribution")
    flat_constraint = cat(getindex.(param_dist_dict_array, "constraint")..., dims = 1)
    name = getindex.(param_dist_dict_array, "name")

    # build the object
    return ParameterDistribution(distribution, flat_constraint, name)

end

"""
    function ParameterDistribution(distribution::ParameterDistributionType,
                                   constraint::Union{ConstraintType,AbstractVector{ConstraintType}},
                                   name::AbstractString)

constructor of a ParameterDistribution from a single `distribution`, (array of) `constraint`, `name`.
these can used to build another ParameterDistribution
"""
function ParameterDistribution(
    distribution::ParameterDistributionType,
    constraint::Union{ConstraintType, AbstractVector},
    name::AbstractString,
)

    if !(typeof(constraint) <: ConstraintType || eltype(constraint) <: ConstraintType) # if it is a vector, but not of constraint
        throw(ArgumentError("`constraint` must be a ConstraintType, or Vector of ConstraintType's "))
    end
    # 1 constraint per dimension check
    constraint_vec = isa(constraint, ConstraintType) ? [constraint] : constraint
    n_parameters = ndims(distribution)

    if !(n_parameters == length(constraint_vec))
        throw(
            DimensionMismatch(
                "There must be one constraint dimension in a parameter distribution, use no_constraint() type if no constraint is required",
            ),
        )
    end

    # flatten the structure
    distribution_vec = [distribution]
    flat_constraint_vec = cat(constraint_vec..., dims = 1)
    name_vec = [name]

    # build the object
    return ParameterDistribution(distribution_vec, flat_constraint_vec, name_vec)

end

"""
    combine_distributions(pd_vec::AbstractVector{ParameterDistribution}...; dims = 1)

concatenate from vector of single ParameterDistributions, dims keyword is unused
"""
function combine_distributions(pd_vec::AbstractVector{PD}; dims = 1) where {PD <: ParameterDistribution}
    # flatten the structure
    distribution = cat([d.distribution for d in pd_vec]..., dims = 1)
    constraint = cat([d.constraint for d in pd_vec]..., dims = 1)
    name = cat([d.name for d in pd_vec]..., dims = 1)
    return ParameterDistribution(distribution, constraint, name)
end

## Functions

"""
    get_name(pd::ParameterDistribution)

Returns a list of ParameterDistribution names.
"""
get_name(pd::ParameterDistribution) = pd.name

"""
    get_dimensions(pd::ParameterDistribution)

The number of dimensions of the parameter space.
"""
function get_dimensions(pd::ParameterDistribution)
    return [ndims(d) for d in pd.distribution]
end
function get_dimensions(d::VectorOfParameterized)
    return [length(dd) for dd in d.distribution]
end

function ndims(pd::ParameterDistribution)
    return sum(get_dimensions(pd))
end

"""
    get_n_samples(pd::ParameterDistribution)

The number of samples in a Samples distribution
"""
function get_n_samples(pd::ParameterDistribution)
    return Dict{String, Any}(pd.name[i] => n_samples(d) for (i, d) in enumerate(pd.distribution))
end

"""
    get_all_constraints(pd::ParameterDistribution)

Returns the (flattened) array of constraints of the parameter distribution.
"""
get_all_constraints(pd::ParameterDistribution) = pd.constraint

"""
    batch(pd::ParameterDistribution)

Returns a list of contiguous `[collect(1:i), collect(i+1:j),... ]` used to split parameter arrays by distribution dimensions.
"""
function batch(pd::Union{ParameterDistribution, VectorOfParameterized})
    #chunk xarray to give to the different distributions.
    d_dim = get_dimensions(pd) #e.g [4,1,2]
    d_dim_tmp = Array{Int64}(undef, size(d_dim)[1] + 1)
    d_dim_tmp[1] = 0
    for i in 2:(size(d_dim)[1] + 1)
        d_dim_tmp[i] = sum(d_dim[1:(i - 1)]) # e.g [0,4,5,7]
    end

    return [collect((d_dim_tmp[i] + 1):d_dim_tmp[i + 1]) for i in 1:size(d_dim)[1]] # e.g [1:4, 5:5, 6:7]
end

"""
    get_distribution(pd::ParameterDistribution)

Returns a `Dict` of `ParameterDistribution` distributions, with the parameter names
as dictionary keys. For parameters represented by `Samples`, the samples are returned
as a 2D (`parameter_dimension x n_samples`) array.
"""
function get_distribution(pd::ParameterDistribution)
    return Dict{String, Any}(pd.name[i] => get_distribution(d) for (i, d) in enumerate(pd.distribution))
end

get_distribution(d::Samples) = d.distribution_samples

get_distribution(d::Parameterized) = d.distribution

get_distribution(d::VectorOfParameterized) = d.distribution

"""
    sample([rng], pd::ParameterDistribution, [n_draws])

Draws `n_draws` samples from the parameter distributions `pd`. Returns an array, with
parameters as columns. `rng` is optional and defaults to `Random.GLOBAL_RNG`. `n_draws` is
optional and defaults to 1.
"""
function sample(rng::AbstractRNG, pd::ParameterDistribution, n_draws::IT) where {IT <: Integer}
    return cat([sample(rng, d, n_draws) for d in pd.distribution]..., dims = 1)
end

# define methods that dispatch to the above with Random.GLOBAL_RNG as a default value for rng
sample(pd::ParameterDistribution, n_draws::IT) where {IT <: Integer} = sample(Random.GLOBAL_RNG, pd, n_draws)
sample(rng::AbstractRNG, pd::ParameterDistribution) = sample(rng, pd, 1)
sample(pd::ParameterDistribution) = sample(Random.GLOBAL_RNG, pd, 1)

"""
    sample([rng], d::Samples, [n_draws])

Draws `n_draws` samples from the parameter distributions `d`. Returns an array, with
parameters as columns. `rng` is optional and defaults to `Random.GLOBAL_RNG`. `n_draws` is
optional and defaults to 1.
"""
function sample(rng::AbstractRNG, d::Samples, n_draws::IT) where {IT <: Integer}
    n_stored_samples = n_samples(d)
    samples_idx = sample(rng, collect(1:n_stored_samples), n_draws)
    if ndims(d) == 1
        return reshape(d.distribution_samples[:, samples_idx], :, n_draws) #columns are parameters
    else
        return d.distribution_samples[:, samples_idx]
    end
end

# define methods that dispatch to the above with Random.GLOBAL_RNG as a default value for rng
sample(d::Samples, n_draws::IT) where {IT <: Integer} = sample(Random.GLOBAL_RNG, d, n_draws)
sample(rng::AbstractRNG, d::Samples) = sample(rng, d, 1)
sample(d::Samples) = sample(Random.GLOBAL_RNG, d, 1)

"""
    sample([rng], d::Parameterized, [n_draws])

Draws `n_draws` samples from the parameter distributions `d`. Returns an array, with
parameters as columns. `rng` is optional and defaults to `Random.GLOBAL_RNG`. `n_draws` is
optional and defaults to 1.
"""
function sample(rng::AbstractRNG, d::Parameterized, n_draws::IT) where {IT <: Integer}
    if ndims(d) == 1
        return reshape(rand(rng, d.distribution, n_draws), :, n_draws) #columns are parameters
    else
        return rand(rng, d.distribution, n_draws)
    end
end

# define methods that dispatch to the above with Random.GLOBAL_RNG as a default value for rng
sample(d::Parameterized, n_draws::IT) where {IT <: Integer} = sample(Random.GLOBAL_RNG, d, n_draws)
sample(rng::AbstractRNG, d::Parameterized) = sample(rng, d, 1)
sample(d::Parameterized) = sample(Random.GLOBAL_RNG, d, 1)

"""
    sample([rng], d::VectorOfParameterized, [n_draws])

Draws `n_draws` samples from the parameter distributions `d`. Returns an array, with
parameters as columns. `rng` is optional and defaults to `Random.GLOBAL_RNG`. `n_draws` is
optional and defaults to 1.
"""
function sample(rng::AbstractRNG, d::VectorOfParameterized, n_draws::IT) where {IT <: Integer}
    samples = zeros(ndims(d), n_draws)
    batches = batch(d)
    dimensions = get_dimensions(d)
    for (i, dd) in enumerate(d.distribution)
        samples[batches[i], :] = rand(rng, dd, n_draws) #columns are parameters
    end
    return samples
end

sample(d::VectorOfParameterized, n_draws::IT) where {IT <: Integer} = sample(Random.GLOBAL_RNG, d, n_draws)
sample(rng::AbstractRNG, d::VectorOfParameterized) = sample(rng, d, 1)
sample(d::VectorOfParameterized) = sample(Random.GLOBAL_RNG, d, 1)

"""
    logpdf(pd::ParameterDistribution, xarray::Array{<:Real,1})

Obtains the independent logpdfs of the parameter distributions at `xarray`
(non-Samples Distributions only), and returns their sum.
"""
get_logpdf(d::Parameterized, xarray::AbstractVector{FT}) where {FT <: Real} = logpdf.(d.distribution, xarray)

function get_logpdf(d::VectorOfParameterized, xarray::AbstractVector{FT}) where {FT <: Real}
    # get the index of xarray chunks to give to the different distributions.
    batches = batch(d)
    dimensions = get_dimensions(d)
    lpdfsum = 0.0
    # perform the logpdf of each of the distributions, and returns their sum
    for (i, dd) in enumerate(d.distribution)
        if dimensions[i] == 1
            lpdfsum += logpdf.(dd, xarray[batches[i]])[1]
        else
            lpdfsum += logpdf(dd, xarray[batches[i]])
        end
    end
    return lpdfsum
end

function get_logpdf(pd::ParameterDistribution, xarray::AbstractVector{FT}) where {FT <: Real}
    #first check we don't have sampled distribution
    for d in pd.distribution
        if typeof(d) <: Samples
            throw(
                ErrorException(
                    "Cannot compute get_logpdf of Samples distribution. Consider using a Parameterized type for your prior.",
                ),
            )
        end
    end
    #assert xarray correct dim/length
    if size(xarray)[1] != ndims(pd)
        throw(DimensionMismatch("xarray must have dimension equal to the parameter space"))
    end

    # get the index of xarray chunks to give to the different distributions.
    batches = batch(pd)

    # perform the logpdf of each of the distributions, and returns their sum
    return sum(cat([get_logpdf(d, xarray[batches[i]]) for (i, d) in enumerate(pd.distribution)]..., dims = 1))
end

#extending StatsBase cov,var
"""
    var(pd::ParameterDistribution)
Returns a flattened variance of the distributions
"""
var(d::Parameterized) = var(d.distribution)
var(d::Samples) = var(d.distribution_samples, dims = 2)
function var(d::VectorOfParameterized)
    d_dims = get_dimensions(d)
    block_var = Array{Any}(undef, size(d_dims)[1])

    for (i, dimension) in enumerate(d_dims)
        block_var[i] = var(d.distribution[i])
    end
    return cat(block_var..., dims = 1)
end

function var(pd::ParameterDistribution)
    d_dims = get_dimensions(pd)
    block_var = Array{Any}(undef, size(d_dims)[1])

    for (i, dimension) in enumerate(d_dims)
        block_var[i] = var(pd.distribution[i])
    end
    return cat(block_var..., dims = 1) #build the flattened vector
end

"""
    cov(pd::ParameterDistribution)

Returns a dense blocked (co)variance of the distributions.
"""
cov(d::Parameterized) = cov(d.distribution)
cov(d::Samples) = cov(d.distribution_samples, dims = 2) #parameters are columns
function cov(d::VectorOfParameterized)
    d_dims = get_dimensions(d)

    # create each block (co)variance
    block_cov = Array{Any}(undef, size(d_dims)[1])
    for (i, dimension) in enumerate(d_dims)
        if dimension == 1
            block_cov[i] = var(d.distribution[i])
        else
            block_cov[i] = cov(d.distribution[i])
        end
    end

    return cat(block_cov..., dims = (1, 2)) #build the block diagonal (dense) matrix

end

function cov(pd::ParameterDistribution)
    d_dims = get_dimensions(pd)

    # create each block (co)variance
    block_cov = Array{Any}(undef, size(d_dims)[1])
    for (i, dimension) in enumerate(d_dims)
        if dimension == 1
            block_cov[i] = var(pd.distribution[i])
        else
            block_cov[i] = cov(pd.distribution[i])
        end
    end

    return cat(block_cov..., dims = (1, 2)) #build the block diagonal (dense) matrix

end

#extending mean
"""
    mean(pd::ParameterDistribution)

Returns a concatenated mean of the parameter distributions.
"""
mean(d::Parameterized) = mean(d.distribution)
mean(d::Samples) = mean(d.distribution_samples, dims = 2)
function mean(d::VectorOfParameterized)
    return cat([mean(dd) for dd in d.distribution]..., dims = 1)
end
function mean(pd::ParameterDistribution)
    return cat([mean(d) for d in pd.distribution]..., dims = 1)
end

#apply transforms

"""
    transform_constrained_to_unconstrained(pd::ParameterDistribution, xarray::Array{<:Real,1})

Apply the transformation to map (possibly constrained) parameters `xarray` into the unconstrained space.
"""
function transform_constrained_to_unconstrained(
    pd::ParameterDistribution,
    xarray::AbstractVector{FT},
) where {FT <: Real}
    return cat([c.constrained_to_unconstrained(xarray[i]) for (i, c) in enumerate(pd.constraint)]..., dims = 1)
end

"""
    transform_constrained_to_unconstrained(pd::ParameterDistribution, xarray::Array{<:Real,2})

Apply the transformation to map (possibly constrained) parameter samples `xarray` into the unconstrained space.
Here, `xarray` contains parameters as columns and samples as rows.
"""
function transform_constrained_to_unconstrained(
    pd::ParameterDistribution,
    xarray::AbstractMatrix{FT},
) where {FT <: Real}
    return Array(hcat([c.constrained_to_unconstrained.(xarray[i, :]) for (i, c) in enumerate(pd.constraint)]...)')
end

"""
    transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{<:Real,1})

Apply the transformation to map parameters `xarray` from the unconstrained space into (possibly constrained) space.
"""
function transform_unconstrained_to_constrained(
    pd::ParameterDistribution,
    xarray::AbstractVector{FT},
) where {FT <: Real}
    return cat([c.unconstrained_to_constrained(xarray[i]) for (i, c) in enumerate(pd.constraint)]..., dims = 1)
end

"""
    transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{<:Real,2})

Apply the transformation to map parameter samples `xarray` from the unconstrained space into (possibly constrained) space.
Here, `xarray` contains parameters as columns and samples as rows.
"""
function transform_unconstrained_to_constrained(
    pd::ParameterDistribution,
    xarray::AbstractMatrix{FT},
) where {FT <: Real}
    return Array(hcat([c.unconstrained_to_constrained.(xarray[i, :]) for (i, c) in enumerate(pd.constraint)]...)')
end

"""
    transform_unconstrained_to_constrained(pd::ParameterDistribution, xarray::Array{Array{<:Real,2},1})

Apply the transformation to map parameter sample ensembles `xarray` from the unconstrained space into (possibly constrained) space.
Here, `xarray` is an iterable of parameters sample ensembles for different EKP iterations.
"""
function transform_unconstrained_to_constrained(
    pd::ParameterDistribution,
    xarray, # ::Iterable{AbstractMatrix{FT}},
) where {FT <: Real}
    transf_xarray = []
    for elem in xarray
        push!(transf_xarray, transform_unconstrained_to_constrained(pd, elem))
    end
    return transf_xarray
end

end # of module


function test_adjtrans()
  @testset ExtendedTestSet "Adjoint/Transpose/Conjugate" begin
    A = rand(5, 3) + im * rand(5, 3)
    opA = LinearOperator(A)

    aopA = AdjointLinearOperator(opA)
    copA = ConjugateLinearOperator(opA)
    topA = TransposeLinearOperator(opA)

    for (foo, fop) in [(adjoint, aopA), (conj, copA), (transpose, topA)]
      @test foo(opA) === fop
      @test Matrix(fop) == foo(A)
      @test foo(fop) === opA

      @test Matrix(-fop) == foo(-A)
      @test Matrix((2 + 3im) * fop) == (2 + 3im) * foo(A)
      @test Matrix(fop * (2 + 3im)) == foo(A) * (2 + 3im)
    end

    @test adjoint(topA) === copA
    @test adjoint(copA) === topA
    @test conj(aopA) === topA
    @test conj(topA) === aopA
    @test transpose(copA) === aopA
    @test transpose(aopA) === copA

    v = rand(5) + im * rand(5)
    @test aopA * v == adjoint(A) * v
    @test topA * v == transpose(A) * v

    v = rand(3) + im * rand(3)
    @test copA * v == conj(A) * v
  end
end

test_adjtrans()


adj = [0 1 0 1;
       1 0 1 0;
       0 1 0 1;
       1 0 1 0]
deg = [2 0 0 0;
       0 2 0 0;
       0 0 2 0;
       0 0 0 2]
lap = [2 -1 0 -1;
       -1 2 -1 0;
       0 -1 2 -1;
       -1 0 -1 2]
norm_lap = [1. -.5 0. -.5;
           -.5 1. -.5 0.;
           0. -.5 1. -.5;
           -.5 0. -.5 1.]
scaled_lap =   [0 -0.5 0 -0.5;
                -0.5 0 -0.5 -0;
                0 -0.5 0 -0.5;
                -0.5 0 -0.5 0]

@testset "linalg" begin
    fg = FeaturedGraph(adj)
    @test fg.matrix_type == :adjm
    for T in [Int8, Int16, Int32, Int64, Int128]
        @test GraphSignals.adjacency_matrix(adj, T) == T.(adj)
        @test GraphSignals.adjacency_matrix(fg, T) == T.(adj)
        @test GraphSignals.degrees(fg; dir=:both) == [2, 2, 2, 2]
        dm = GraphSignals.degree_matrix(fg, T; dir=:out)
        @test dm == T.(deg)
        @test GraphSignals.degree_matrix(adj, T; dir=:in) == dm
        @test GraphSignals.degree_matrix(adj, T; dir=:both) == dm
        @test GraphSignals.laplacian_matrix(fg, T) == T.(lap)
    end

    fg = FeaturedGraph(Float64.(adj))
    @test fg.matrix_type == :adjm
    for T in [Float16, Float32, Float64]
        dm = GraphSignals.degree_matrix(fg, T; dir=:out)
        @test dm == T.(deg)
        @test GraphSignals.degree_matrix(adj, T; dir=:in) == dm
        @test GraphSignals.degree_matrix(adj, T; dir=:both) == dm
        @test GraphSignals.laplacian_matrix(fg, T) == T.(lap)
        @test GraphSignals.normalized_laplacian(fg, T) ≈ T.(norm_lap)
        @test GraphSignals.scaled_laplacian(fg, T) ≈ T.(scaled_lap)

        fg_ = GraphSignals.laplacian_matrix!(deepcopy(fg), T)
        @test fg_.graph == T.(lap)
        @test fg_.matrix_type == :laplacian
        fg_ = GraphSignals.normalized_laplacian!(deepcopy(fg), T)
        @test fg_.graph ≈ T.(norm_lap)
        @test fg_.matrix_type == :normalized
        fg_ = GraphSignals.scaled_laplacian!(deepcopy(fg), T)
        @test fg_.graph ≈ T.(scaled_lap)
        @test fg_.matrix_type == :scaled
    end
end


export antiplane_viscosity, mantle_viscosity
@doc raw"""
    antiplane_viscosity(ε::Union{PyObject, Array{Float64}}, σ::Union{PyObject, Array{Float64}},
    μ::Union{PyObject, Float64}, η::Union{PyObject, Float64}, Δt::Float64)

Calculates the stress at time $t_{n+1}$ given the strain at $t_{n+1}$ and stress at $t_{n}$. The governing equation is
```math
\dot\sigma + \frac{\mu}{\eta}\sigma = 2\mu \dot\epsilon
```
The discretization form is
```math
\sigma^{n+1} = \frac{1}{\frac{1}{\Delta t}+\frac{\mu}{\eta}}(2\mu\dot\epsilon^{n+1} + \frac{\sigma^n}{\Delta t})
```
"""
function antiplane_viscosity(dε::Union{PyObject, Array{Float64}}, σ::Union{PyObject, Array{Float64}},
        μ::Union{PyObject, Float64, Array{Float64}}, η::Union{PyObject, Float64, Array{Float64}}, Δt::Union{PyObject,Float64})
    dε, σ, μ, η = convert_to_tensor([dε, σ, μ,η], [Float64, Float64, Float64, Float64])
    if length(size(μ))==0
        η/(η/Δt + μ)*(2μ*dε+σ/Δt)
    elseif length(size(μ))==1
        repeat(η/(η/Δt + μ),1,2).*(2repeat(μ, 1, 2).*dε+σ/Δt)
    end
end

@doc raw"""
    mantle_viscosity(u::Union{Array{Float64}, PyObject},
        T::Union{Array{Float64}, PyObject}, m::Int64, n::Int64, h::Float64;
        σ_yield::Union{Float64, PyObject} = 300e6,
        ω::Union{Float64, PyObject},
        η_min::Union{Float64, PyObject} = 1e18,
        η_max::Union{Float64, PyObject} = 1e23,
        E::Union{Float64, PyObject} = 9.0,
        C::Union{Float64, PyObject} = 1000., N::Union{Float64, PyObject} = 2.)

```math
\eta = \eta_{\min} + \min\left( \frac{\sigma_{\text{yield}}}{2\sqrt{\epsilon_{II}}}, \omega\min(\eta_{\max}, \eta) \right)
```
with
```math
\epsilon_{II} = \frac{1}{2} \epsilon(u)\qquad \eta = C e^{E(0.5-T)} (\epsilon_{II})^{(1-n)/2n}
```

Here $\epsilon_{II}$ is the second invariant of the strain rate tensor, $C > 0$ is a viscosity pre-factor, $E > 0$ is the non-dimensional activation energy,
$n > 0$ is the nonlinear exponent, $η_\min$, $η_\max$ act as minimum and maximum bounds for the effective viscosity, and $σ_{\text{yield}} > 0$ is the yield
stress. $w\in (0, 1]$ is the weakening factor, which is used to incorporate phenomenological aspects that
cannot be represented in a purely viscous flow model, such as processes which govern mega-thrust faults along the subduction interface, or
partial melting near a mid-ocean ridge.

The viscosity of the mantle is governed by the high-temperature creep of silicates, for which laboratory experiments show that the creep
strength is temperature-, pressure-, compositional- and stress-dependent.

The output is a length $4mn$ vector.

!!! info
    See __Towards adjoint-based inversion of time-dependent mantle convection with nonlinear viscosity__ for details.
"""
function mantle_viscosity(u::Union{Array{Float64}, PyObject},
     T::Union{Array{Float64}, PyObject}, m::Int64, n::Int64, h::Float64;
     σ_yield::Union{Float64, PyObject} = 300e6,
     ω::Union{Float64, PyObject} = 1.0,
     η_min::Union{Float64, PyObject} = 1e-2,
     η_max::Union{Float64, PyObject} = 1e3,
     E::Union{Float64, PyObject} = 9.0,
     C::Union{Float64, PyObject} = 1000., N::Union{Float64, PyObject} = 2.)
    u, T, σ_yield, ω, η_min, η_max, E, C, N = convert_to_tensor(
        [u, T, σ_yield, ω, η_min, η_max, E, C, N], [Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64]
    )
    ε = eval_strain_on_gauss_pts(u, m, n, h)
    εII = 0.5 * sum(ε .* ε, dims = 2)
    T = reshape(repeat(T, 1, 4), (-1,))
    η = C * exp(E*(0.5 - T))*εII^((1-N)/2N)
    η_min + min(σ_yield/2/sqrt(εII), ω*min(η_max, η))
end


function PHI(x::Array{T,1} where T<:Real, c::Array{T,1} where T<:Real, rinv::Real)

    dist2 = pairwise(SqEuclidean(), reshape(x, 1, length(x)), reshape(c, 1, length(c)), dims=2)

    Φ = exp.(-0.5*dist2*rinv*rinv)

    @assert(size(Φ,1) == length(x)) # N

    @assert(size(Φ,2) == length(c)) # M

    return Φ

end

function centresongrid(x; numcentres=100)

    collect(LinRange(minimum(x), maximum(x), numcentres))

end


# Binary heap

struct MutableBinaryHeapNode{T}
    value::T
    handle::Int
end

#################################################
#
#   core implementation
#
#################################################

function _heap_bubble_up!(ord::Ordering,
    nodes::Vector{MutableBinaryHeapNode{T}}, nodemap::Vector{Int}, nd_id::Int) where T

    @inbounds nd = nodes[nd_id]
    v::T = nd.value

    swapped = true  # whether swap happens at last step
    i = nd_id

    while swapped && i > 1  # nd is not root
        p = i >> 1
        @inbounds nd_p = nodes[p]

        if Base.lt(ord, v, nd_p.value)
            # move parent downward
            @inbounds nodes[i] = nd_p
            @inbounds nodemap[nd_p.handle] = i
            i = p
        else
            swapped = false
        end
    end

    if i != nd_id
        nodes[i] = nd
        nodemap[nd.handle] = i
    end
end

function _heap_bubble_down!(ord::Ordering,
    nodes::Vector{MutableBinaryHeapNode{T}}, nodemap::Vector{Int}, nd_id::Int) where T

    @inbounds nd = nodes[nd_id]
    v::T = nd.value

    n = length(nodes)
    last_parent = n >> 1

    swapped = true
    i = nd_id

    while swapped && i <= last_parent
        il = i << 1

        if il < n   # contains both left and right children
            ir = il + 1

            # determine the better child
            @inbounds nd_l = nodes[il]
            @inbounds nd_r = nodes[ir]

            if Base.lt(ord, nd_r.value, nd_l.value)
                # consider right child
                if Base.lt(ord, nd_r.value, v)
                    @inbounds nodes[i] = nd_r
                    @inbounds nodemap[nd_r.handle] = i
                    i = ir
                else
                    swapped = false
                end
            else
                # consider left child
                if Base.lt(ord, nd_l.value, v)
                    @inbounds nodes[i] = nd_l
                    @inbounds nodemap[nd_l.handle] = i
                    i = il
                else
                    swapped = false
                end
            end

        else  # contains only left child
            nd_l = nodes[il]
            if Base.lt(ord, nd_l.value, v)
                @inbounds nodes[i] = nd_l
                @inbounds nodemap[nd_l.handle] = i
                i = il
            else
                swapped = false
            end
        end
    end

    if i != nd_id
        @inbounds nodes[i] = nd
        @inbounds nodemap[nd.handle] = i
    end
end

function _binary_heap_pop!(ord::Ordering,
    nodes::Vector{MutableBinaryHeapNode{T}}, nodemap::Vector{Int}, nd_id::Int=1) where T

    # extract node
    rt = nodes[nd_id]
    v = rt.value
    @inbounds nodemap[rt.handle] = 0

    if length(nodes) == 1
        # clear
        empty!(nodes)
    else
        # move the last node to the position of the removed node
        @inbounds nodes[nd_id] = new_rt = nodes[end]
        pop!(nodes)
        @inbounds nodemap[new_rt.handle] = nd_id

        if length(nodes) > 1
            if Base.lt(ord, new_rt.value, v)
                _heap_bubble_up!(ord, nodes, nodemap, nd_id)
            else
                _heap_bubble_down!(ord, nodes, nodemap, nd_id)
            end
        end
    end
    return v
end

function _make_mutable_binary_heap(ord::Ordering, ty::Type{T}, values) where T
    # make a static binary index tree from a list of values

    n = length(values)
    nodes = Vector{MutableBinaryHeapNode{T}}(undef, n)
    nodemap = Vector{Int}(undef, n)

    i::Int = 0
    for v in values
        i += 1
        @inbounds nodes[i] = MutableBinaryHeapNode{T}(v, i)
        @inbounds nodemap[i] = i
    end

    for i = 1 : n
        _heap_bubble_up!(ord, nodes, nodemap, i)
    end
    return nodes, nodemap
end

#################################################
#
#   Binary Heap type and constructors
#
#################################################

mutable struct MutableBinaryHeap{T, O <: Base.Ordering} <: AbstractMutableHeap{T, Int}
    ordering::O
    nodes::Vector{MutableBinaryHeapNode{T}}
    node_map::Vector{Int}

    function MutableBinaryHeap{T, O}() where {T, O}
        ordering = O()
        nodes = Vector{MutableBinaryHeapNode{T}}()
        node_map = Vector{Int}()
        new{T, O}(ordering, nodes, node_map)
    end

    function MutableBinaryHeap{T, O}(xs::AbstractVector{T}) where {T, O}
        ordering = O()
        nodes, node_map = _make_mutable_binary_heap(ordering, T, xs)
        new{T, O}(ordering, nodes, node_map)
    end
end

const MutableBinaryMinHeap{T} = MutableBinaryHeap{T, Base.ForwardOrdering}
const MutableBinaryMaxHeap{T} = MutableBinaryHeap{T, Base.ReverseOrdering}

MutableBinaryMinHeap(xs::AbstractVector{T}) where T = MutableBinaryMinHeap{T}(xs)
MutableBinaryMaxHeap(xs::AbstractVector{T}) where T = MutableBinaryMaxHeap{T}(xs)

function Base.show(io::IO, h::MutableBinaryHeap)
    print(io, "MutableBinaryHeap(")
    nodes = h.nodes
    n = length(nodes)
    if n > 0
        print(io, string(nodes[1].value))
        for i = 2 : n
            print(io, ", $(nodes[i].value)")
        end
    end
    print(io, ")")
end

#################################################
#
#   interfaces
#
#################################################

Base.length(h::MutableBinaryHeap) = length(h.nodes)

Base.isempty(h::MutableBinaryHeap) = isempty(h.nodes)

function Base.push!(h::MutableBinaryHeap{T}, v) where T
    nodes = h.nodes
    nodemap = h.node_map
    i = length(nodemap) + 1
    nd_id = length(nodes) + 1
    push!(nodes, MutableBinaryHeapNode(convert(T, v), i))
    push!(nodemap, nd_id)
    _heap_bubble_up!(h.ordering, nodes, nodemap, nd_id)
    return i
end

function Base.sizehint!(h::MutableBinaryHeap, s::Integer)
    sizehint!(h.nodes, s)
    sizehint!(h.node_map, s)
    return h
end

@inline Base.first(h::MutableBinaryHeap) = h.nodes[1].value

"""
    top_with_handle(h::MutableBinaryHeap)

Returns the element at the top of the heap `h` and its handle.
"""
function top_with_handle(h::MutableBinaryHeap)
    el = h.nodes[1]
    return el.value, el.handle
end

Base.pop!(h::MutableBinaryHeap{T}) where {T} = _binary_heap_pop!(h.ordering, h.nodes, h.node_map)

"""
    update!{T}(h::MutableBinaryHeap{T}, i::Int, v::T)

Replace the element at index `i` in heap `h` with `v`.
This is equivalent to `h[i]=v`.
"""
function update!(h::MutableBinaryHeap{T}, i::Int, v) where T
    nodes = h.nodes
    nodemap = h.node_map
    ordering = h.ordering

    nd_id = nodemap[i]
    v0 = nodes[nd_id].value
    x = convert(T, v)
    nodes[nd_id] = MutableBinaryHeapNode(x, i)
    if Base.lt(ordering, x, v0)
        _heap_bubble_up!(ordering, nodes, nodemap, nd_id)
    else
        _heap_bubble_down!(ordering, nodes, nodemap, nd_id)
    end
end

"""
    delete!{T}(h::MutableBinaryHeap{T}, i::Int)

Deletes the element with handle `i` from heap `h` .
"""
function Base.delete!(h::MutableBinaryHeap{T}, i::Int) where T
     nd_id = h.node_map[i]
    _binary_heap_pop!(h.ordering, h.nodes, h.node_map, nd_id)
    return h
end

Base.setindex!(h::MutableBinaryHeap, v, i::Int) = update!(h, i, v)
Base.getindex(h::MutableBinaryHeap, i::Int) = h.nodes[h.node_map[i]].value


function boardtest{Lib<:PolyhedraLibrary}(lib::Lib)
  A1 = -eye(Int, 9) # x >= 0
  b1 = zeros(Int, 9)
  A2 = eye(Int, 9) # x <= 1
  b2 = ones(Int, 9)
  A3 = zeros(Int, 9, 9)
  b3 = 3 * ones(Int, 9)
  i = 1
  for a = 1:3
    for b = (a+1):3
      for c = 1:3
        for d = (c+1):3
          ac = a + (c-1) * 3
          ad = a + (d-1) * 3
          bc = b + (c-1) * 3
          bd = b + (d-1) * 3
          A3[i, ac] = 1
          A3[i, ad] = 1
          A3[i, bc] = 1
          A3[i, bd] = 1
          i += 1
        end
      end
    end
  end
  A = [A1; A2; A3]
  b = [b1; b2; b3]
  ine = SimpleHRepresentation(A, b)
  poly = polyhedron(ine, lib)
  @test !isempty(poly)
  ext  = SimpleVRepresentation(getgenerators(poly))
  target = ones(Int, 9) * (3 // 4)
  ok = false
  for i = 1:size(ext.V, 1)
    # In julia v0.4 [i,:] returns a row matrix and in v0.5 it is
    # a 1D vector hence the use of vec
    if vec(ext.V[i,:]) == target
      ok = true
    end
  end
  @test ok

  cutA = ones(Int, 1, 9)
  cutb = 6
  Acut = [cutA; A]
  bcut = [cutb; b]
  inecut = SimpleHRepresentation(Acut, bcut)
  polycut = polyhedron(inecut, lib)
  @test !isempty(polycut)
  #(isredundant, certificate) = isredundantinequality(polycut, 1)
  #@test !isredundant
  #@test certificate == target
  @test !isredundantinequality(polycut, 1)
  @test IntSet([]) == getredundantinequalities(polycut)
 #(issredundant, scertificate) = isstronglyredundantinequality(polycut, 1)
 #@test !issredundant
 #@test scertificate == target
 #@test IntSet([]) == getstronglyredundantinequalities(polycut)
end


@testset "Result" begin
    @testset "Result+PathResult" begin
        R = solve(equations(heart()), save_all_paths=true, show_progress=false)
        @test R isa Result
        @test natinfinity(R) ≤ 572
        @test nfinite(R) == 4
        @test length(collect(R)) == 576
        @test finite(R) isa Vector{<:PathResult}

        @test length(finite(R, onlynonsingular=false)) == 4
        @test length(finite(R, onlynonsingular=true)) == 4
        @test length(finite(R, onlysingular=true)) == 0
        @test 572 - length(failed(R)) == natinfinity(R)
        @test length(real(R, tol=1e-6)) == 2
        @test nreal(R, tol=1e-6) == 2
        @test length(atinfinity(R)) ≤ 572
        @test length(results(R, onlyreal=true, realtol=1e-8)) == 2
        @test length(results(R, onlynonsingular=true, singulartol=1e9)) == 4
        @test length(finite(results(R, onlyreal=true))) == 2
        @test nresults(R, onlynonsingular=true, singulartol=1e9) == 4
        @test length(results(R, onlyfinite=false)) == 576
        @test nresults(R, onlyfinite=false) == 576
        @test nnonsingular(R) == 4
        @test length(nonsingular(R)) == 4
        @test length(singular(R)) == 0
        @test length(singular(real(R))) == 0

        @test_nowarn mapresults(solution, R)
        @test_nowarn mapresults(start_solution, R)
        @test_nowarn mapresults(residual, R)
        @test_nowarn mapresults(issingular, R)
        @test count(mapresults(isreal, R)) == 2
        # test fallback
        @test count(results(isreal, R)) == 2

        @test length(solutions(R)) == 4
        @test solutions(R) isa Vector{Vector{ComplexF64}}
        @test realsolutions(R, realtol=1e-8) isa Vector{Vector{Float64}}
        @test length(realsolutions(R, realtol=1e-8)) == 2

        @test_nowarn string(R)
        @test_nowarn string(R[end])
        @test_nowarn show(IOContext(devnull, :compact=>true), R)
        @test_nowarn show(IOContext(devnull, :compact=>true), R[end])

        test_treeviews(R)

        @polyvar x y z
        R = solve([(x-3)^3,(y-2)], affine_tracking=true)
        @test R isa Result
        test_treeviews(R)
        @test nnonsingular(R) == 0
        @test nsingular(R) == 3
        @test_nowarn sprint(show, R)

        R = solve([(x-3)^3,(y-2)], affine_tracking=true, system_scaling=nothing)
        @test nnonsingular(R) == 0
        @test nsingular(R) == 3
        @test_nowarn sprint(show, R)

        @polyvar x y z
        R = solve([(x-3z),(y-2z)])
        @test R isa Result{<:ProjectiveVectors.PVector}
        @test_nowarn string(R)
        test_treeviews(R)
        @test nnonsingular(R) == 1

        @polyvar x y z
        R = solve([(x-3z)^3,(y-2z)])
        @test R isa Result{<:ProjectiveVectors.PVector}
        test_treeviews(R)
        @test nnonsingular(R) == 0
        @test nsingular(R) == 3

        @polyvar x y z
        V = x^2 + y^2 - z^2
        L = randn(1,3) * [x, y, z]
        R = solve([V; L])
        @test_nowarn sprint(show, R)
    end

    @testset "multiplicities" begin
        @polyvar x y z
        f = (x-1*y)*(x-1.01*y)
        g = x - y - z
        S = solve([f,g], threading=false)
        @test length(multiplicities(S, tol=1e-1)) == 1
        @test length(multiplicities(S, tol=1e-5)) == 0
        @test length(multiplicities(S.pathresults, tol=1e-1)) == 1
        @test length(multiplicities(S.pathresults, tol=1e-5)) == 0

        @polyvar x z
        f = (x-1)*(x-1.01)
        g = x - 1 - z
        S = solve([f,g], threading=false)
        @test length(multiplicities(S, tol=1e-1)) == 1
        @test length(multiplicities(S, tol=1e-5)) == 0
        @test length(multiplicities(S.pathresults, tol=1e-1)) == 1
        @test length(multiplicities(S.pathresults, tol=1e-5)) == 0
        @test norm(solution(S[1]) - solution(S[2])) < 1e-1
    end

    @testset "uniquesolutions" begin
        @polyvar x
        f = (x-3)^3*(x-2)
        R = solve([f])
        @test length(uniquesolutions(R)) == 2
        @test uniquesolutions(R) isa Vector{Vector{Complex{Float64}}}
        @test length(uniquesolutions(R, multiplicities=true)) == 2
        a, b = uniquesolutions(R, multiplicities=true)
        @test (a[2] == 3 && b[2] == 1) || (b[2] == 3 && a[2] == 1)

        @polyvar x y
        f = (x-3y)^3*(x-2y)
        R = solve([f])
        @test length(uniquesolutions(R)) == 2
        @test uniquesolutions(R) isa Vector{ProjectiveVectors.PVector{Complex{Float64}, 1}}
        @test length(uniquesolutions(R, multiplicities=true)) == 2
        a, b = uniquesolutions(R, multiplicities=true)
        @test (a[2] == 3 && b[2] == 1) || (b[2] == 3 && a[2] == 1)
    end
end


struct KernelRegression{R <: Union{Vector, Matrix}, T <: Real}
    x::R
    y::Vector{T}
    k::Function
    h::Float64
end

# Want a version for x as a vector (this) and X as a matrix (doesn't exist yet)
function kernelregression(x::Vector{R}, y::Vector{T}; kernel::Symbol = :epanechnikov, bandwidth::Real = NaN, getbandwidth::Function = kNN.bandwidth) where {R <: Any, T <: Real}
    if isnan(bandwidth)
        h_x = getbandwidth(x)
        h_y = getbandwidth(y)
        h = sqrt(h_x * h_y)
    else
        h = bandwidth
    end
    k = kernels[kernel]
    return KernelRegression(x, y, k, h)
end

function StatsBase.predict(model::KernelRegression, x::Real)
    y, n = 0.0, length(model.x)
    h, k = model.h, model.k
    normalizer = 0.0
    for i in 1:n
        # TODO: Get definitions and usage of bandwidth right
        # w_i = (1 / h) * k((x - model.x[i]) / h)
        w_i = h * k.(h * (x - model.x[i]))
        y += model.y[i] * w_i
        normalizer += w_i
    end
    return y / normalizer
end

function StatsBase.predict!(ys::Vector, model::KernelRegression, xs::AbstractVector{T}) where T <: Real
    n = length(xs)
    # @assert length(ys) == n
    for i in 1:n
        ys[i] = predict(model, xs[i])
    end
    return
end

function StatsBase.predict(model::KernelRegression, xs::AbstractVector{T}) where T <: Real
    ys = Array{T}(undef, length(xs))
    predict!(ys, model, xs)
    return ys
end


module DiscreteExteriorCalculus

using StaticArrays: SVector, SMatrix
using LinearAlgebra: Symmetric
using UniqueVectors: UniqueVector

export Point
"""
    Point{N}(coords::SVector{N, Float64}) where N
    Point(coords::AbstractVector{<:Real})
    Point(coords::Vararg{<:Real})
    Point(b::Barycentric)
    Point(b::SimpleBarycentric)

A point in `N` dimensional space.
"""
struct Point{N}
    coords::SVector{N, Float64}
end

export Simplex
"""
    Simplex{N, K}(points::SVector{K, Point{N}}) where {N, K}
    Simplex(points::AbstractVector{Point{N}}) where N
    Simplex(points::Vararg{Point{N}}) where N
    Simplex(c::Cell)
    Simplex(s::SimpleSimplex)

A simplex of dimension `K-1` embedded in `N` dimensional space.
"""
struct Simplex{N, K}
    points::SVector{K, Point{N}}
    function Simplex{N, K}(points::SVector{K, Point{N}}) where {N, K}
        @assert 1 <= K <= N+1
        return new{N, K}(points)
    end
end

export SimpleSimplex
"""
    SimpleSimplex{N}(Vector{Point{N}}) where N
    SimpleSimplex(points::AbstractVector{Point{N}}) where N
    SimpleSimplex(s::Simplex)
    SimpleSimplex(s::Vector{SimpleBarycentric{N}}) where N

Like a `simplex` but not parametrized by dimension.
"""
struct SimpleSimplex{N}
    points::Vector{Point{N}}
end

export Metric
"""
    Metric{N, T<:SMatrix{N, N, Float64}}(mat::Symmetric{Float64, T})
    Metric(mat::T) where {N, T<:SMatrix{N, N, Float64}}
    Metric(mat::AbstractMatrix{<:Real})

A symmetric but not necessarily positive semi-definite metric.
"""
struct Metric{N, T<:SMatrix{N, N, Float64}}
    mat::Symmetric{Float64, T}
end

export Wedge
"""
    Wedge{N, K}(vectors::SVector{K, SVector{N, Float64}})
    Wedge(vectors::Vector{SVector{N, Float64}}) where N

The wedge product of `K` vectors in `N` dimensional space.
"""
struct Wedge{N, K}
    vectors::SVector{K, SVector{N, Float64}}
end

export Barycentric
"""
    Barycentric{N, K}(s::Simplex{N, K}, coords::SVector{K, Float64}) where {N, K}
    Barycentric(s::Simplex{N, K}, coords::AbstractVector{<:Real}) where {N, K}
    Barycentric(s::Simplex, coords::Vararg{<:Real})
    Barycentric(b::SimpleBarycentric)

A representation of a point with respect to a simplex using barycentric coordinates. If the
simplex `s` has vertices `Aᵢ` and `x=coords` then the point represented by
`b = Barycentric(s, coords)` is `Point(b) = ΣAᵢxᵢ`.
"""
struct Barycentric{N, K}
    simplex::Simplex{N, K}
    coords::SVector{K, Float64}
    function Barycentric{N, K}(s::Simplex{N, K}, coords::SVector{K, Float64}) where {N, K}
        @assert sum(coords) ≈ 1
        if sum(coords) != 1
            coords = SVector{K, Float64}([coords[1:end-1]..., 1 - sum(coords[1:end-1])])
        end
        return new{N, K}(s, coords)
    end
end

export SimpleBarycentric
"""
    SimpleBarycentric{N}(s::SimpleSimplex{N}, coords::Vector{Float64}) where N
    SimpleBarycentric(s::SimpleSimplex{N}, coords::AbstractVector{<:Real}) where N
    SimpleBarycentric(b::Barycentric)

Like a `Barycentric` but not parametrized by dimension.
"""
struct SimpleBarycentric{N}
    simplex::SimpleSimplex{N}
    coords::Vector{Float64}
    function SimpleBarycentric{N}(simplex::SimpleSimplex{N}, coords::Vector{Float64}) where N
        @assert length(simplex.points) == length(coords)
        @assert sum(coords) ≈ 1
        if sum(coords) != 1
            coords = [coords[1:end-1]..., 1 - sum(coords[1:end-1])]
        end
        return new{N}(simplex, coords)
    end
end

export Cell
"""
    Cell{N}(children::Vector{Cell{N}}, parents::Dict{Cell{N}, Bool},
        points::Vector{Point{N}}, K::Int) where N
    Cell(points::AbstractVector{Point{N}}, K::Int) where N
    Cell(s::Simplex{N, K}) where {N, K}

A polytope of dimension `K-1` embedded in `N` dimensional space.
"""
struct Cell{N}
    children::Vector{Cell{N}}
    parents::Dict{Cell{N}, Bool}
    points::Vector{Point{N}}
    K::Int
    function Cell{N}(children::Vector{Cell{N}}, parents::Dict{Cell{N}, Bool},
                     points::Vector{Point{N}}, K::Int) where N
        @assert 1 <= K <= N+1
        for c in children
            @assert c.K == K-1
        end
        for c in keys(parents)
            @assert c.K == K+1
        end
        return new{N}(children, parents, points, K)
    end
end

export CellComplex
"""
    CellComplex{N,K}(cells::SVector{K, UniqueVector{Cell{N}}}) where {N, K}
    CellComplex{N,K}() where {N, K}
    CellComplex{N,K}(cells::AbstractVector{Cell{N}}) where {N, K}
    CellComplex(simplices::AbstractVector{Simplex{N, K}})

A cell complex with cells of dimension `0` through `K-1` embedded in `N` dimensional space.
"""
struct CellComplex{N, K}
    cells::SVector{K, UniqueVector{Cell{N}}}
    function CellComplex{N,K}(cells::SVector{K, UniqueVector{Cell{N}}}) where {N, K}
        @assert 1 <= K <= N+1
        for k in 1:K
            for c in cells[k]
                @assert c.K == k
            end
        end
        return new{N,K}(cells)
    end
end

SignedSimpleSimplex{N} = Tuple{SimpleSimplex{N}, Bool}

export TriangulatedComplex
"""
    TriangulatedComplex{N, K}(complex::CellComplex{N, K},
        simplices::Dict{Cell{N}, Vector{Tuple{SimpleSimplex{N}, Bool}})
    TriangulatedComplex{N,K}() where {N, K}
    TriangulatedComplex(simplices::AbstractVector{Simplex{N, K}}) where {N, K}

A cell complex along with a decomposition of each cell into simplices. The decomposition
comes with signs which are used to compute a signed volume for each cell.
"""
struct TriangulatedComplex{N, K}
    complex::CellComplex{N, K}
    simplices::Dict{Cell{N}, Vector{SignedSimpleSimplex{N}}}
end

export Mesh
"""
    Mesh{N,K}(primal::TriangulatedComplex{N,K}, dual::TriangulatedComplex{N,K})
    Mesh(tcomp::TriangulatedComplex{N}, center::Function) where N

A pair of TriangulatedComplexes representing a primal and dual pair. The second constructor
takes a primal TriangulatedComplex and a function `center` that takes a Simplex{N, K} to a
Barycentric{N, K} (e.g. the circumcenter or centroid) and returns the corresponding dual
TriangulatedComplex.
"""
struct Mesh{N,K}
    primal::TriangulatedComplex{N,K}
    dual::TriangulatedComplex{N,K}
end

include("utils.jl")
include("metrics.jl")
include("simplices.jl")
include("cell_complex.jl")
include("components_and_boundary.jl")
include("orientation.jl")
include("mesh.jl")
include("operators.jl")
include("constraints.jl")
end


using Kalman
using Base.Test

s = UnscentedState([1.0, 0.0],4*eye(2),0.1,2.0,0.0)
su,wm,wc = sigma(s)

@test s.x == [1.0,0.0]
@test s.p == [4.0 0; 0 4.0]
@test_approx_eq su[:,5] [1.0,-sqrt(0.08)]

xt = zeros(s.x)
pt = zeros(s.p)

for i in 1:size(su,2)
    xt += wm[i] * su[:,i]
    pt += wc[i] * (su[:,i]-s.x)*(su[:,i]-s.x)'
end

@test_approx_eq xt s.x
@test_approx_eq pt s.p

dt = 0.01

function fu(x)
    x1 = zeros(x)
    x1[1] = x[1] + dt*x[2]
    x1[2] = x[2] - dt*x[1]
    x1
end

fm = AdditiveUnscentedModel(fu,1e-10*eye(2))

hu(x) = [x[1]]

zm = AdditiveUnscentedObservationModel(hu,[0.1]')

kf = AdditiveUnscentedKalmanFilter(s,fm,zm)


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra

@testset "GUN (native)" begin
    nep = nep_gallery("nlevp_native_gun")
    n = size(nep, 1)
    tol = 1e-11
    ref_eigenvalue = 22345.116783765 + 0.644998598im # from NLEIGS

    @bench @testset "Running algorithm" begin
        λ1,v1 = quasinewton(nep, λ = 150^2+1im, v = ones(n), logger = displaylevel, tol = tol, maxit = 500)

        v1 = v1 / norm(v1)

        @test norm(compute_Mlincomb(nep, λ1, v1)) < sqrt(tol)
        @test norm(compute_Mder(nep, λ1) * v1) < sqrt(tol)

        @test norm(λ1 - ref_eigenvalue) < sqrt(tol)*100
    end

    @bench @testset "Compute derivatives" begin
        λ = 150^2+2im
        v = randn(n)
        z1 = compute_Mlincomb(nep, λ, v, [1.0], 1)

        # Compare with divided difference
        ee = 1e-4
        z2 = (compute_Mlincomb(nep, λ+ee, v) - compute_Mlincomb(nep, λ-ee, v)) / (2*ee)

        @test norm(z2-z1) < (ee^2)*1000
    end

    @testset "Errors thrown" begin
        @test_throws MethodError nep_gallery("nlevp_native_gun", 15)
        @test_throws MethodError nep_gallery("nlevp_native_gun", t=15)
    end
end


const HILLCOEFFICIENT = Param(0.1; bounds=(0.0, 100))
const MEDIAN_LOGLOGISTIC = Param(0.1; bounds=(0.0, 100))
const THRESHOLD_EXPOSURE = Param(0.0; bounds=(0.0, 1e9))

"""
    Mortality <: CellRule

Abstract super type for rules of survival effect
For best performance these should be chained with other
CellRule or following an NeighborhoodRule.
"""
abstract type Mortality{R,W} <: CellRule{R,W} end

"""
    LoglogisticMortality <: Mortality

    LoglogisticMortality(; median, hillcoefficient, timestep, [nsteps_type])
    LoglogisticMortality{R}(; median, hillcoefficient, timestep, [nsteps_type])
    LoglogisticMortality{R,W}(; median, hillcoefficient, timestep, [nsteps_type])

Loglogistic mortality based on median, ``α`` and hill coefficient ``β``.

Cumulative function of loglogistic:
```math
F(x; α, β) =  x^β/(α^β + x^β)
```
where ``α>0`` is the scale and ``β>0`` is the shape.

# Keywords

- `median`: Median of the loglogistic function
- `hillcoefficient`: Hill's coefficient, a measure of ultrasensitivity (i.e. how steep is the response curve).
    May be a `Number`, an `Aux` array or another `Grid`.
- `timestep`: Time step for the mortality rate, in a type compatible with the simulation `tspan`.

Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
`R` is a 2 Grids `NamedTuple` like `Tuple{:population,:exposure}` and `W` return only the first grid `:population`.
"""
struct LoglogisticMortality{R,W,MD,HC,TS,S} <: Mortality{R,W}
    median::MD
    hillcoefficient::HC
    timestep::TS
    nsteps::S
end
function LoglogisticMortality{R,W}(;
    median=MEDIAN_LOGLOGISTIC,
    hillcoefficient=HILLCOEFFICIENT,
    timestep,
    nsteps_type=Float64,
) where {R,W}
    LoglogisticMortality{R,W}(median, hillcoefficient, timestep, zero(nsteps_type))
end

modifyrule(rule::LoglogisticMortality, data) = precalc_nsteps(rule, data)

@inline function applyrule(data, rule::LoglogisticMortality, (N, X), I)
    N > zero(N) || return zero(N)
    β = get(data, rule.hillcoefficient, I...)
    α = get(data, rule.median, I...)
    # NOTE: 1 - x^β / (α^β + x^β) = α^β / (α^β + x^β)
    return N * α^β / (α^β + X^β)
end

"""
    ExponentialMortality <: Mortality

    ExponentialMortality(; rate, threshold, timestep, [nsteps_type])
    ExponentialMortality{R}(; rate, threshold, timestep, [nsteps_type])
    ExponentialMortality{R,W}(; rate, threshold, timestep, [nsteps_type])

Exponential mortality based on exposure threshold and mortality rate parameter, using exact solution.

Exponential mortality based on exposure grid ``X``, an exposure threshold parameter ``z``
    and a mortality rate ``r`` using exact solution between time ``t`` and ```t+1``:

```math
N_{t+1} = N_{t}e^{-r t(X-z)}
```

# Keywords

- `rate`: Mortality rate.
- `threshold`: Exposure threshold under which there is no effect.
- `timestep`: Time step for the growth rate, in a type compatible with the simulation `tspan`.
- `nsteps_type`: Specify the floating point type to use when `nsteps` is generated from the
    timestep, if it is required for type-stability or performance. The default is `Float64`.

`rate` and `threshold` can be a `Number`, an `Aux` array or another Grid`.

Pass grid `Symbol`s to `R` or both `R` and `W` type parameters to use to specific grids.
`R` is a 2 Grids `NamedTuple` like `Tuple{:population,:exposure}` and `W` return only the first grid `:population`.
"""
struct ExponentialMortality{R,W,MR,ZX,TS,S} <: GrowthRule{R,W}
    rate::MR
    threshold::ZX
    timestep::TS
    nsteps::S
end
function ExponentialMortality{R,W}(;
    rate=INTRINSICRATE_PARAM,
    threshold=THRESHOLD_EXPOSURE,
    timestep,
    nsteps_type=Float64,
) where {R,W}
    ExponentialMortality{R,W}(rate, threshold, timestep, zero(nsteps_type))
end

modifyrule(rule::ExponentialMortality, data) = precalc_nsteps(rule, data)

@inline function applyrule(data, rule::ExponentialMortality, (N,X), I)
    N > zero(N) || return zero(N)
    rt = get(data, rule.rate, I...) * rule.nsteps
    z = get(data, rule.threshold, I...)
    return N * exp(- rt * max(0, X-z))
end


function trimolecular_1path_diff!(du,u,p,t,dens,cons)
    # unpack rates and constants
    nᵣ,nᵣₓ,n₃ = u
    kₓ,k₋ₓ,k₃,k₋₃ = p
    k₁,k₋₁ = cons
    mᵣ,mₗ,mₓ,A = dens

    # model
    du[1] = dnᵣ  = A*k₁*mᵣ*mₗ - k₋₁*nᵣ - A*kₓ*mₓ*nᵣ + k₋ₓ*nᵣₓ
    du[2] = dnᵣₓ = A*kₓ*mₓ*nᵣ - k₋ₓ*nᵣₓ - k₃*nᵣₓ + k₋₃*n₃
    du[3] = dn₃  = k₃*nᵣₓ - k₋₃*n₃

end


export
    WignerFunction,
    WignerSurface,
    wigner

#=
    Wigner function by Laguerre Polynominal
=#

function calc_wigner(
    x_range::AbstractRange,
    p_range::AbstractRange,
    dim::Integer
)
    𝐰 = Array{ComplexF64,4}(undef, dim, dim, length(x_range), length(p_range))
    xs = collect(x_range)
    ps = collect(p_range)
    @sync for m in 1:dim
        Threads.@spawn for n in 1:dim
            𝐰[m, n, :, :] .=
                gaussian_function(xs, ps) .*
                coefficient_of_wave_function(m, n) .*
                z_to_power(m, n, xs, ps) .*
                laguerre(m, n, xs, ps)
        end
    end

    return 𝐰
end

function load_wigner(
    bin_path::String,
    x_range::AbstractRange,
    p_range::AbstractRange,
    dim::Integer
)
    return load_𝐰(bin_path, x_range, p_range, dim)
end

function create_wigner(
    x_range::AbstractRange,
    p_range::AbstractRange,
    dim::Integer
)
    bin_name = gen_wigner_bin_name(x_range, p_range, dim)
    hash = my_artifact_hash(bin_name, my_artifacts[])

    if !isnothing(hash) && my_artifact_exists(hash)
        return load_wigner(my_artifact_path(hash), x_range, p_range, dim)
    end

    𝐰 = calc_wigner(x_range, p_range, dim)
    save_𝐰(bin_name, 𝐰)

    return 𝐰
end

"""
    WignerFunction

`WignerFunction` to calculate Wigner function for a quantum state in gevin `x` and `p` range.

## Arguments

* `x_range`: First physical domain range in phase space.
* `p_range`: Second physical domain range in phase space.
* `dim`: Photon number truncation of the quantum state.

## Example

```julia-repl
julia> xs = -1:0.1:1; ps = -1:0.1:1; dim=35;

julia> wf = WignerFunction(xs, ps, dim=dim);

julia> wf(VacuumState(Matrix, dim=dim));
```
"""
mutable struct WignerFunction{T<:Integer, U<:AbstractRange}
    𝐰::Array{ComplexF64,4}
    x_range::U
    p_range::U
    dim::T

    function WignerFunction(
        x_range::U,
        p_range::U,
        dim::T
    ) where {T<:Integer, U<:AbstractRange}
        𝐰 = create_wigner(x_range, p_range, dim)

        return new{T, U}(𝐰, x_range, p_range, dim)
    end
end

function WignerFunction(x_range::AbstractRange, p_range::AbstractRange; dim)
    return WignerFunction(x_range, p_range, dim)
end

struct WignerSurface{T<:AbstractRange}
    x_range::T
    p_range::T
    𝐰_surface::Matrix{Float64}

    function WignerSurface(
        x_range::T,
        p_range::T,
        𝐰_surface::Matrix{Float64}
    ) where {T<:AbstractRange}
        return new{T}(x_range, p_range, 𝐰_surface)
    end
end

function (wf::WignerFunction)(ρ::AbstractMatrix{T}) where {T}
    𝐰_surface = Matrix{Float64}(undef, length(wf.x_range), length(wf.p_range))
    @sync for i in 1:length(wf.x_range)
        Threads.@spawn for j in 1:length(wf.p_range)
            𝐰_surface[i, j] = real(sum(ρ .* wf.𝐰[:, :, i, j]))
        end
    end

    return WignerSurface(wf.x_range, wf.p_range, 𝐰_surface)
end

# TODO: (wf::WignerFunction)(state::StateVector) = wf(ρ(state))

"""
    wigner(ρ::AbstractMatrix{T}, x_range::AbstractRange, p_range::AbstractRange)

Calculate the Wigner fuction of the quantum state in the gevin `x` and `p` range.

## Arguments

* `ρ`: Density matrix of a quantum state.
* `x_range`: First physical domain range in phase space.
* `p_range`: Second physical domain range in phase space.

## Example

```julia-repl
julia> xs = -1:0.1:1; ps = -1:0.1:1;

julia> wigner(VacuumState(Matrix, dim=35), xs, ps);
```
"""
function wigner(ρ::AbstractMatrix{T}, x_range::AbstractRange, p_range::AbstractRange) where {T}
    dim = size(ρ, 1)
    wf = WignerFunction(x_range, p_range, dim=dim)

    return wf(ρ)
end

# TODO: wigner(v::AbstractVector{T}, x_range::AbstractRange, p_range::AbstractRange) where {T}


const timestamp_raw, buses_raw = readlines("day13.txt")

const timestamp = parse(Int, timestamp_raw)

const buses_p1 = parse.(Int, [b for b in split(buses_raw, ',') if b != "x"])

function part1(buses::Array{Int}, timestamp::Int)
    remainders = mod.(-timestamp, buses)
    bus_min = argmin(remainders)
    return buses[bus_min] * remainders[bus_min]
end

println("Part 1: $(part1(buses_p1, timestamp))")

buses_p2 = [[parse(Int, b), -i + 1] for (i, b) in enumerate(split(buses_raw, ",")) if b != "x"]

function chineseremainder(Ns::Vector{BigInt}, remainders::Vector{BigInt})
    n = BigInt(prod(Ns))
    Ni = n .÷ Ns
    invs = invmod.(Ni, Ns) .* Ni
    s = sum(remainders .* invs) % n
    return s >= 0 ? s : s + n
end

function part2(buses)
    Ns = @. BigInt(first(buses))
    remainders = @. BigInt(last(buses))
    return chineseremainder(Ns, remainders)
end

println("Part 2: $(part2(buses_p2))")


function _svd_(A::AbstractSymMatrix)
    T = eltype(A)
    S = vtype(A)

    u_blocks = SortedDict{Sector{S, 2}, Matrix{T}}()
    s_blocks = SortedDict{Sector{S, 2}, Diagonal{Float64, Vector{Float64}}}()
    v_blocks = SortedDict{Sector{S, 2}, Matrix{T}}()

    midl = SortedDict{S, Int}()
    midr = SortedDict{S, Int}()
    fact = svd(zeros(T,1,1), full=false)
    for (sect,blk) in A.blocks
        c1, c2 = sect
        try
            fact = svd(blk, full=false)
        catch e
            @warn "error $e"
            fact = svd(blk, full=false, alg=LinearAlgebra.QRIteration())
        end
        u_blocks[Sector(c1, c1)] = fact.U
        s_blocks[Sector(c1, c2)] = Diagonal(fact.S)
        v_blocks[Sector(c2, c2)] = fact.Vt
        midl[c1] = length(fact.S)
        midr[c2] = length(fact.S)
    end

    Vl = VectorSpace{S}(midl, A.space[1].isdual)
    Vr = VectorSpace{S}(midr, A.space[2].isdual)

    u  = SymMatrix{S,T}(zero(S), (A.space[1], dual(Vl)), u_blocks)
    s  = SymDiagonal{S,Float64}(A.charge, (Vl, Vr), s_blocks)
    v = SymMatrix{S,T}(zero(S), (dual(Vr), A.space[2]), v_blocks)

    u, s, v
end

function _svd_(A::AbstractArray)
    u,s,v = svd(A, full=false)
    # note that iterating the component produces U,S,V instead of U,S,Vt!
    u, Diagonal(s), v'
end

"""
    fermionswapgate(l1, l2)

Fot the two given legs, `l1` and `l2` make the fermionic swap gate
which is `-1` if the two legs has charge odd and `+1` otherwise.

The output legs are l2_afterX, l1_afterX, l2, l1
"""
function fermionswapgate(l1::VectorSpace{S}, l2::VectorSpace{S}) where{S}

    space = (l2, l1, dual(l2), dual(l1))

    sects, sizes = _allsectorsandsizes(0, space)
    blocks = SortedDict{Sector{S, 4}, Array{Float64, 4}}()

    for index in eachindex(sects)
        c2, c1, c2X, c1X = sects[index]
        if c1 == c1X && c2 == c2X
            d2, d1 = sizes[index][1:2]
            s = isodd(c1) && isodd(c2) ? -1 : +1
            blocks[sects[index]] =  s * reshape(I(d1*d2), d2, d1, d2, d1)
        else
            blocks[sects[index]] = zeros(sizes[index])
        end
    end
    SymTensor(zero(S), space, blocks)
end

# function isrightisometry(A::SymTensor)
#     N = rank(A)

# end

function isrightisometry(A::SymMatrix)
    if eye(eltype(A), A.space[1]) ≈ contract(A, (1, -1), dual(A), (2, -1))
        return true
    else
        println(contract(A, (1, -1), dual(A), (2, -1)))
    end
    false
end

function isleftisometry(A::SymMatrix)
    eye(eltype(A), dual(A.space[2])) ≈ contract(dual(A), (-1, 1), A, (-1, 2))
end

isunitary(A::SymMatrix) = isrightisometry(A) && isleftisometry(A)


using Test
using PolynomialBasis
using Roots
using IntervalArithmetic
using Revise
using ImplicitDomainQuadrature

IDQ = ImplicitDomainQuadrature

function circle_distance_function(coords, center, radius)
    difference = (coords .- center).^2
    distance = radius .- mapslices(sum,difference,dims=1)'
    return distance
end

numqp = 4
polyorder = 2
poly = InterpolatingPolynomial(1,2,polyorder)

center = [3.,0.5]
radius = 1.5
coords = [0.0  0.0  0.0  1.5  1.5  1.5  3.0  3.0  3.0
          0.0  0.5  1.0  0.0  0.5  1.0  0.0  0.5  1.0]
coeffs = circle_distance_function(coords,center,radius)
update!(poly,coeffs)

interpgrad = IDQ.interpolating_gradient(poly)

grad(x) = gradient(poly,x)

box = IntervalBox(-1..1,2)

# quad1d = ReferenceQuadratureRule(numqp)
# pquad = area_quadrature(poly,+1,box,quad1d)


using LinearAlgebra, ForwardDiff, NLPModels

export Bisseccao, Newton_rc_bissec

function Bisseccao(g, a, b, max_bissec; λ = 0)
    ϵ = 1e-4
    status= :resolvido
    iter = 0
    while abs(g(λ)) > ϵ
        if g(a)*g(b)==0 && g(b)==0
            λ=b
        else
            λ=a
        end
        if g(a) * g(b) < 0
            while abs(b-a) > ϵ
                λ = (b + a) / 2
                if g(λ) * g(a) < 0
                    b = λ
                elseif g(λ) * g(b) < 0
                    a = λ
                end
            end
        end
        if g(a) * g(b) > 0
            if b < 1000
                b = b*5+1e-3
            elseif a < 8000
                a = a + 40
            else
                status= :bisseccao_falhou
                break
            end
        end

        if iter > max_bissec
            println("Maximum bissection.")
            break
        end
        iter +=1
    end
    return λ, status
end

function Newton_rc_bissec(
    nlp::AbstractNLPModel;
    η1::Float64 = 1e-2,
    η2::Float64 = 0.75,
    Δ::Float64 = 5.0,
    max_bissec = 1000,
    a::Float64 = 0.0,
    b::Float64 = 1.0,
    atol::Real = 1e-6,
    rtol::Real = 1e-6,
    max_eval::Int = 1000,
    max_iter::Int = 0,
    max_time::Float64 = 10.0
    )

    if !unconstrained(nlp)
        error("Problem is not unconstrained")
    end

    x = copy(nlp.meta.x0)
    n = nlp.meta.nvar

    f(x) = obj(nlp, x)
    ∇f(x) = grad(nlp, x)
    H(x) = hess(nlp, x)

    fx = f(x)
    ∇fx = ∇f(x)
    Hx = Matrix(H(x))
    g(λ) = norm(inv(Hx + λ * I) * ∇f(x)) - Δ
    d = - Hx \ ∇fx

    ϵ = atol + rtol * norm(∇fx)
    t₀ = time()

    iter = 0
    Δt = time() - t₀
    solved = norm(∇fx) < ϵ # First order stationary
    tired = neval_obj(nlp) ≥ max_eval > 0|| iter ≥ max_iter > 0 || Δt ≥ max_time > 0 # Excess time, iteration, evaluations

    status = :unknown

    @info log_header(
    [:iter, :fx, :ngx, :nf, :Δt],
    [Int, Float64, Float64, Int, Float64],
    hdr_override=Dict(:fx => "f(x)", :ngx => "‖∇f(x)‖", :nf => "#f")
    )

    @info log_row(
    Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
    )

    while !(solved || tired)
        if norm(d) < Δ
            d = - Hx \ ∇f(x)
        else
            λ, stat = Bisseccao(g, a, b, max_bissec)
            if stat == :bisseccao_falhou
                @warn("Bissection fail")
                break
            end
            d = - (Hx + λ * I) \ ∇f(x)
        end

        Ared = f(x) - f(x + d)
        Pred = f(x) - (f(x) + dot(d, ∇f(x)) + dot(d, Hx * d) /  2)
        ρ = Ared / Pred
        if ρ < η1
            Δ = Δ / 2
        elseif ρ < η2
            x = x + d
        else
            x = x + d
            Δ = 2Δ
        end

        fx = f(x)
        ∇fx = ∇f(x)
        Hx = Matrix(H(x))

        iter += 1
        Δt = time() - t₀
        solved = norm(∇fx) < ϵ # First order stationary
        tired = neval_obj(nlp) ≥ max_eval > 0|| iter ≥ max_iter > 0 || Δt ≥ max_time > 0 # Excess time, iteration, evaluations

        @info log_row(
        Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
        )
    end

    if solved
        status = :first_order
    elseif tired
        if neval_obj(nlp) ≥ max_eval > 0
            status = :max_eval
        elseif iter ≥ max_iter > 0
            status = :max_iter
        elseif Δt ≥ max_time > 0
            status = :max_time
        end
    end

    return GenericExecutionStats(
    status,
    nlp,
    solution=x,
    objective=f(x),
    dual_feas=norm(∇fx),
    elapsed_time=Δt,
    iter=iter
    )
end


const test = "3,4,3,1,2"

function simulate(fish::AbstractVector{<: Integer},
                  days::Integer;
                  fertility_cycle::Integer = 7,
                  spawn_age::Integer = 8
                 )::Integer
    ages = [count(==(age), fish) for age in 0:spawn_age]
    for day in 1:days
        spawning = ages[1]
        ages = [ages[2:end]; spawning]
        ages[fertility_cycle] += spawning
    end
    return sum(ages)
end

let v = parse.(Int, split(test, ','))
    @assert simulate(v, 80) == 5934
    @assert simulate(v, 256) == 26984457539
end

v = parse.(Int, split(readline("day06.in"), ','))
println("Q1: ", simulate(v, 80))
println("Q2: ", simulate(v, 256))


include("utils.jl")

@testset "MultinomialMatrices manifold" begin
    M = MultinomialMatrices(3, 2)
    @testset "MultinomialMatrices manifold Basics" begin
        @test ProbabilitySimplex(2)^3 === MultinomialMatrices(3, 3)
        @test ProbabilitySimplex(2)^(3,) === PowerManifold(ProbabilitySimplex(2), 3)
        @test ^(ProbabilitySimplex(2), 2) === MultinomialMatrices(3, 2)
        @test typeof(^(ProbabilitySimplex(2), 2)) == MultinomialMatrices{3,2,2}
        @test repr(M) == "MultinomialMatrices(3,2)"
        @test representation_size(M) == (3, 2)
        @test manifold_dimension(M) == 4
        @test Manifolds.get_iterator(ProbabilitySimplex(3)^4) === Base.OneTo(4)
        p = [2, 0, 0]
        p2 = [p p]
        @test !is_manifold_point(M, p)
        @test_throws DomainError is_manifold_point(M, p, true)
        @test !is_manifold_point(M, p2)
        @test_throws CompositeManifoldError is_manifold_point(M, p2, true)
        @test !is_tangent_vector(M, p2, 0.0)
        @test_throws DomainError is_tangent_vector(M, p2, [-1.0, 0.0, 0.0], true)
        @test !is_tangent_vector(M, p2, [-1.0, 0.0, 0.0])
        @test_throws DomainError is_tangent_vector(M, p, [-1.0, 0.0, 0.0], true)
        @test injectivity_radius(M) ≈ 0
        x = [0.5 0.4 0.1; 0.5 0.4 0.1]'
        y = [0.6 0.3 0.1; 0.4 0.5 0.1]'
        z = [0.3 0.6 0.1; 0.6 0.3 0.1]'
        test_manifold(
            M,
            [x, y, z],
            test_injectivity_radius = false,
            test_forward_diff = false,
            test_reverse_diff = false,
            test_vector_spaces = true,
            test_project_tangent = false,
            test_musical_isomorphisms = true,
            test_default_vector_transport = false,
            is_tangent_atol_multiplier = 5.0,
        )
    end
end


@testset "Isogenies" begin
  K = GF(7)
  E1 = EllipticCurve(K, [1, 2, 3, 4, 5])
  E2 = EllipticCurve(K, [1, 2, 3, 1, 1])
  phi = @inferred isogeny_from_kernel(E1, division_polynomial_univariate(E1,3)[1])
  @test @inferred domain(phi) == E1
  @test @inferred codomain(phi) == E2
  @test is_isomorphic(E1, codomain(phi))
  @test @inferred image(phi, E1([1,1])) == E2([6,2])

  K = GF(3)
  E1 = EllipticCurve(K, [1, 0, 1, 1, 0])
  E2 = EllipticCurve(K, [1, 0, 1, 0, 2])
  Kx, x = PolynomialRing(K);
  f = x+2
  phi = @inferred isogeny_from_kernel(E1, f)
  @test @inferred domain(phi) == E1
  @test @inferred codomain(phi) == E2
  @test @inferred isinfinite(image(phi, E1([1,2])))

  E1 = EllipticCurve([1, 2, 3, 4, 5])
  E2 = EllipticCurve([1, 2, 3, 979//16, 19067//64])
  phi = @inferred isogeny_from_kernel(E1, division_polynomial_univariate(E1,2)[1])
  @test @inferred domain(phi) == E1
  @test @inferred codomain(phi) == E2
  @test is_isomorphic(E1, codomain(phi))

  K= GF(2,4)
  a = gen(K)
  E1 = EllipticCurve(K,[a^2,1-a,1,0,a])
  E2 = EllipticCurve(K,[a^2,1-a,1,a^8,1])
  f = division_polynomial_univariate(E1,5)[1]
  phi = @inferred isogeny_from_kernel(E1, f)
  @test @inferred domain(phi) == E1
  @test @inferred codomain(phi) == E2
  @test @inferred image(phi, E1([a^2,a^2])) == E2([1,a^14])
end


using LinearAlgebraicRepresentation
using Plasm, LinearAlgebra
Lar = LinearAlgebraicRepresentation
using SparseArrays

include("./LARinput.jl")

# s = Lar.Struct(store)
# V,CV,FV,EV = Lar.struct2lar(s)
# V = Plasm.normalize3D(V) TODO:  solve MethodError bug
# Plasm.view(V,CV)

cop_EV = Lar.coboundary_0(EV::Lar.Cells);
cop_EW = convert(Lar.ChainOp, cop_EV);
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
W = convert(Lar.Points, V');

V, copEV, copFE, copCF = Lar.Arrangement.spatial_arrangement(
	W::Lar.Points, cop_EW::Lar.ChainOp, cop_FE::Lar.ChainOp)

cc = [copEV, copFE, copCF]
output = Lar.lar2obj(V::Lar.Points, cc::Lar.ChainComplex)

f = open("test/out3d.obj", "w")
print(f, output)
close(f)

V,EVs,FVs = Lar.obj2lar("test/out3d.obj")

# for k=1:length(FVs)
# 	Plasm.view(V, FVs[k])
# end

Plasm.view(V, FVs[1])


function w_partial_t(w_x, w_y, phi_x, phi_y, psi_x, psi_y, am)

    _a = -(phi_x.*w_y - phi_y.*w_x)
    _b = -(psi_x.*w_y - psi_y.*w_x)
    return _a + am*_b
end

function psi_partial_t(psi_x, psi_y, phi_x, phi_y)

    return -(phi_x.*psi_y - phi_y.*psi_x)
end

function forward_euler(x, x_t, K1, K2, dt, gd::GhostData)

    return (x + dt*(K1*x_t)) ./ (1 .+ dt*diag(K2))
end

function solve_pde(A, x0, b, gd::GhostData)

    An, _ = require_boundary_conditions(A, gd)
    bn = gd.Proj*b

    bn = transpose(An)*bn
    An = transpose(An)*An
    # x = jacobi_preconditioned_gmres(An, gd.Proj*x0, bn, 1, 1)
    x = gd.Proj*x0

    gauss_seidel!(x, An, bn; maxiter=10)
    gmres!(x, An, bn)
    return gd.R*transpose(gd.Proj)*x
end

function solve_diffusion(x, k, Dxx, Dyy, gd::GhostData)

    M = I - k*(Dxx + Dyy)
    return solve_pde(M, x, x, gd)
end

function solve_vorticity_eqn(phi, w, Dxx, Dyy, gd::GhostData)

    return solve_pde(Dxx + Dyy, phi, w, gd)
end

function leapfrog!(w, psi, phi, Dx, Dxx, Dy, Dyy, K1, K2, gd, dt, am, k_psi, k_w)

    # take derivatives
    psi_x = Dx*psi[:,2]
    psi_y = Dy*psi[:,2]
    phi_x = Dx*phi
    phi_y = Dy*phi
    w_x = Dx*w[:,2]
    w_y = Dy*w[:,2]

    w_t = w_partial_t(w_x, w_y, phi_x, phi_y, psi_x, psi_y, am)
    psi_t = psi_partial_t(psi_x, psi_y, phi_x, phi_y)

    # compute predictor step
    psi[:,3] = forward_euler(0.5*(psi[:,1] + psi[:,2]), psi_t, K1, K2, dt, gd)
    w[:,3] = forward_euler(0.5*(w[:,1] + w[:,2]), w_t, K1, K2, dt, gd)

    # compute diffusion
    psi[:,3] = solve_diffusion(psi[:,3], k_psi, Dxx, Dyy, gd)
    w[:,3] = solve_diffusion(w[:,3], k_w, Dxx, Dyy, gd)

    # solve vorticity equation
    phi[:] = solve_vorticity_eqn(phi, w[:,3], Dxx, Dyy, gd)

    # take derivatives again
    psi_x = Dx*psi[:,3]
    psi_y = Dy*psi[:,3]
    phi_x = Dx*phi
    phi_y = Dy*phi
    w_x = Dx*w[:,3]
    w_y = Dy*w[:,3]

    w_t = w_partial_t(w_x, w_y, phi_x, phi_y, psi_x, psi_y, am)
    psi_t = psi_partial_t(psi_x, psi_y, phi_x, phi_y)

    # compute corrector step
    psi[:,3] = forward_euler(psi[:,2], psi_t, K1, K2, dt, gd)
    w[:,3] = forward_euler(w[:,2], w_t, K1, K2, dt, gd)

    # compute diffusion
    psi[:,3] = solve_diffusion(psi[:,3], k_psi, Dxx, Dyy, gd)
    w[:,3] = solve_diffusion(w[:,3], k_w, Dxx, Dyy, gd)

    # solve vorticity equation
    phi[:] = solve_vorticity_eqn(phi, w[:,3], Dxx, Dyy, gd)

    # shift indices
    psi[:,1:2] = psi[:,2:3]
    w[:,1:2] = w[:,2:3]
end


using Pkg;
using CSV, Plots;
using Statistics, SparseArrays, LinearAlgebra;
cd(@__DIR__) #src

include("func_slp.jl")

mutable struct inputcom
  y        :: Array
  x        :: Array
  w        :: Array
  H_min    :: Int
  H_max    :: Int
  type     :: String
  r        :: Int8
  λ        :: Int8
end

df = CSV.read("data.csv");
df = convert(Array{Float16}, df[:, 2:4]);
T, k = size(df);

## Parameterization
P = 4; # Number of lags used in LP - controlled variable

# start LP at H_min=0 or 1 (H_min=1 if impose no contemporanous impact)
H_max = 20;
ind_response = 1; # Endogenous variable - response
ind_shock    = 3; # Endogenous variable related to the shock

ind_lp = "reg"; # Type of linear regression

ind_contempo = filter(x -> x ≠ ind_shock, ind_all);
#,(P+1):end
y  = df[:, ind_response]; # endogenous variable
x  = df[:, ind_shock]; # endoegnous variable related to the shock

# control variables (contemporaneous vars, lagged vars)
w  = [ df[:, ind_contempo]  matlag( df , P ) ];
w[.!isfinite.(w)] .= 0;
# Packaing everything as input
r = 2; #(r-1)=order of the limit polynomial (so r=2 implies the IR is shrunk towards a line )
λ = 100; # value of penalty

testing₀ = inputcom(y, x, w, H_min, H_max, "reg", 3,4);
testing₁ = inputcom(y, x, w, H_min, H_max, "smooth", r, λ);

lp₀ = smoothlocalprojection(testing₀);

lp₁ = smoothlocalprojection(testing₁);
lp₁.δ
plot(lp₀.IR)
plot!(lp₁.IR)


"""
    Gkm(xy, r, window)

Return nearest neighbour distance distribution based on Kaplam-Meier estimator.
"""
function Gkm(pp, r)
    xy = pp.data
    if length(xy) == 1
        return zeros(Float64, length(r))
    end
    svec = @. SVector(getx(xy), gety(xy))
    tree = KDTree(svec)
    # Take 2 nearest neighbors because, the first nearest neighbor is the point itself
    _, D = knn(tree, svec, 2, true)
    D = last.(D)
    B = [bdist(window(pp), x) for x in xy]
    ind = D .<= B
    T = min.(D, B)
    km(ind, T, r)
end


# Quadratic programming in MATLAB-like style
#
#   minimize x^2 + xy + y^2 + yz + z^2
#
#   s.t.    x + 2 y + 3 z >= 4
#           x +   y       >= 1
#

using Xpress

model = xpress_model(
	name = "qp_02",
	f = [0., 0., 0.],
	H = [2. 1. 0.; 1. 2. 1.; 0. 1. 2.],
	A = -[1. 2. 3.; 1. 1. 0.],
	b = -[4., 1.])

optimize(model)

println(Xpress.get_lp_status(model))
println(Xpress.getq(model))

println("sol = $(get_solution(model))")
println("obj = $(get_objval(model))")


include("../examples/2D\ Cauchy\ Distribution.jl")
include("../examples/Airy\ equation.jl")
if isdir(Pkg.dir("RandomMatrices")) && isdir(Pkg.dir("PyPlot"))
    include("../examples/GUE\ Sampling.jl")
end
include("../examples/Lanczos.jl")
include("../examples/Lee\ &\ Greengard\ equation.jl")
include("../examples/Newton iteration for Nonlinear BVPs.jl")


@doc raw"""
    ∂log(Z::Rotor)

Return the gradient of `log(Z)` with respect to the components of `Z`.

The result includes "off-shell" components of the gradient, meaning that even
though change of `Z` in a direction that changes its norm would not be allowed
for a `Rotor`, we measure the gradient in that direction anyway.  That is, the
elements of the returned vector of quaternions is
```math
\begin{aligned}
  \left[
    \frac{\partial} {\partial Z_w} \log(Z),
    \frac{\partial} {\partial Z_x} \log(Z),
    \frac{\partial} {\partial Z_y} \log(Z),
    \frac{\partial} {\partial Z_z} \log(Z)
  \right].
\end{aligned}
```

Note that, even though `log(::Rotor)` is a `QuatVec`, the derivative (and
therefore each element of the result) is a general `Quaternion`.

See also [`∂exp`](@ref) for a similar function, as well as [`log∂log`](@ref)
for a function to compute the value along with the gradient.

# Examples
```julia
julia> ∂log∂w, ∂log∂x, ∂log∂y, ∂log∂z = ∂log(randn(QuatVecF64));

```
"""
function ∂log(Z::Rotor)
    # log(Z::Rotor) = log(abs2) / 2 + f(Z) * Z.vec
    # f(Z) = atan(absvec / Z.w) / absvec = acos(Z.w) / absvec
    a2 = abs2vec(Z)
    a = sqrt(a2)
    if a < 2eps(typeof(a))
        return Quaternion.([Quaternion(one(a)), zero(a)*imx, zero(a)*imy, one(a)*imz])
    end
    f = acos(Z.w) / a
    fprime = (Z.w - f) / a2
    V = 1 + fprime * QuatVec(Z)
    [
        Quaternion(conj(Z)),
        Z.x * V + f * imx,
        Z.y * V + f * imy,
        Z.z * V + f * imz,
    ]
end

"""
    log∂log(Z::Rotor)

Return the value and gradient of `log(Z)` with respect to the components of `Z`.

See [`∂log`](@ref) for more explanation of the components of the gradient.

# Examples
```julia
julia> l, ∂l = log∂log(randn(RotorF64));

```
"""
function log∂log(Z::Rotor)
    # log(Z::Rotor) = log(abs2) / 2 + f(Z) * Z.vec
    # f(Z) = atan(absvec / Z.w) / absvec = acos(Z.w) / absvec
    a2 = abs2vec(Z)
    a = sqrt(a2)
    if a < 2eps(typeof(a))
        return (
            QuatVec(zero(a)),
            Quaternion.([Quaternion(one(a)), zero(a)*imx, zero(a)*imy, one(a)*imz])
        )
    end
    f = acos(Z.w) / a
    fprime = (Z.w - f) / a2
    V = 1 + fprime * QuatVec(Z)
    ∂log = [
        Quaternion(conj(Z)),
        Z.x * V + f * imx,
        Z.y * V + f * imy,
        Z.z * V + f * imz,
    ]
    (f * QuatVec(Z), ∂log)
end

@doc raw"""
    ∂exp(Z::QuatVec)

Return the gradient of `exp(Z)` with respect to the components of `Z`.

The result includes "off-shell" components of the gradient, meaning that even
though a scalar component of `Z` would not be allowed for a `QuatVec`, we
measure the gradient in that direction anyway.  That is, the first element
of the returned vector of quaternions is
```math
\begin{aligned}
  \left.\frac{\partial} {\partial Z_w} \exp(Z) \right|_{Z_w=0}.
\end{aligned}
```

Note that, even though `exp(::QuatVec)` is a `Rotor`, the derivative (and
therefore each element of the result) is a general `Quaternion`.

See also [`∂log`](@ref) for a similar function, as well as [`exp∂exp`](@ref)
for a function to compute the value along with the gradient.

# Examples
```julia
julia> ∂exp∂w, ∂exp∂x, ∂exp∂y, ∂exp∂z = ∂exp(randn(QuatVecF64));

```
"""
function ∂exp(Z::QuatVec)
    # exp(Z::Quaternion) = exp(Z.w) * cos(absvec) + g(Z) * Z.vec
    # g(Z) = exp(Z.w) * sin(absvec) / absvec
    a2 = abs2vec(Z)
    a = sqrt(a2)
    if a < 2eps(typeof(a))
        return Quaternion.([Quaternion(one(a)), one(a)*imx, one(a)*imy, one(a)*imz])
    end
    s, c = sincos(a)
    sinca = s / a
    g = sinca
    gprime = (c - g) / a2
    V = -sinca + gprime * Z
    [
        c + sinca * Z,
        Z.x * V + g * imx,
        Z.y * V + g * imy,
        Z.z * V + g * imz,
    ]
end

"""
    exp∂exp(Z::QuatVec)

Return the value and gradient of `exp(Z)` with respect to the components of `Z`.

See [`∂exp`](@ref) for more explanation of the components of the gradient.

# Examples
```julia
julia> e, ∂e = exp∂exp(randn(QuatVecF64));

```
"""
function exp∂exp(Z::QuatVec)
    # exp(Z::Quaternion) = exp(Z.w) * cos(absvec) + g(Z) * Z.vec
    # g(Z) = exp(Z.w) * sin(absvec) / absvec
    a2 = abs2vec(Z)
    a = sqrt(a2)
    if a < 2eps(typeof(a))
        return (
            Quaternion(one(a)),
            Quaternion.([Quaternion(one(a)), one(a)*imx, one(a)*imy, one(a)*imz])
        )
    end
    s, c = sincos(a)
    sinca = s / a
    g = sinca
    gprime = (c - g) / a2
    V = -sinca + gprime * Z
    ∂exp = [
        c + sinca * Z,
        Z.x * V + g * imx,
        Z.y * V + g * imy,
        Z.z * V + g * imz,
    ]
    (c+g*Z, ∂exp)
end

@doc raw"""
    slerp∂slerp(q₁, q₂, τ)

Return the value and gradient of `slerp`.

The gradient is with respect to each of the input arguments in turn, with each
quaternion regarded as a series of four arguments.  That is, a total of 10
quaternions will be returned:
```math
\begin{aligned}
  &\big[\\
    &\quad \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₁.w} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₁.x} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₁.y} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₁.z} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₂.w} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₂.x} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₂.y} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial q₂.z} \mathrm{slerp}(q₁, q₂, τ), \\
    &\quad \frac{\partial}{\partial \tau} \mathrm{slerp}(q₁, q₂, τ), \\
  &\big]
\end{aligned}
```
For convenience, this will be a 4-tuple with the `slerp` as the first element,
the first four components of the derivative, followed by the next four
components of the derivative, followed by the last component of the derivative.

See also [`slerp`](@ref) for just the value, and [`slerp∂slerp∂τ`](@ref) for
just the value and the derivative with respect to `τ`.

# Examples
```julia
julia> (q1, q2), τ = randn(RotorF64, 2), rand();

julia> s, ∂s∂q₁, ∂s∂q₂, ∂s∂τ = slerp∂slerp(q₁, q₂, τ);

```
"""
function slerp∂slerp(q₁::Rotor, q₂::Rotor, τ::Real)
    r = q₂ / q₁
    basis = typeof(r)[1, imx, imy, imz]
    l, ∂l = log∂log(r)
    e, ∂e = exp∂exp(τ*l)
    s = e * q₁
    (
        s,
        [
            sum(
                ∂e[b] * τ * ∂l[c][b] * (
                    q₂ * conj(basis[a]) #/ abs2(q₁)
                    - 2q₁[a] * (q₂/q₁) #/ abs2(q₁)
                )[c] * q₁
                for b in 1:4
                for c in 1:4
            ) + e * basis[a]
            for a in 1:4
        ],
        [
            sum(
                ∂e[b] * τ * ∂l[c][b] * (
                    basis[a] / q₁
                )[c] * q₁
                for b in 1:4
                for c in 1:4
            )
            for a in 1:4
        ],
        l * s
    )
end

"""
    slerp∂slerp∂τ(q₁, q₂, τ)

Return the value and time-derivative of `slerp`.

See also [`slerp∂slerp`](@ref), which returns the value and *all* of the
derivatives of `slerp`.

"""
function slerp∂slerp∂τ(q₁::Rotor, q₂::Rotor, τ::Real)
    l = log(q₂ / q₁)
    e = exp(τ*l)
    s = e * q₁
    ∂s∂t = l * s
    (s, ∂s∂t)
end

"""
    squad∂squad∂t(qᵢ, A, B, qᵢ₊₁, ta, tb, t)

Compute the value and time-derivative of [`squad`](@ref).

This is primarily an internal helper function, taking various parameters
computed within the `squad` function.  This will be used to compute the
derivative of `squad` when the angular velocity is also requested.  To actually
obtain the derivative, simply pass the relevant keyword to the `squad`
function.

"""
function squad∂squad∂t(qᵢ, A, B, qᵢ₊₁, ta, tb, t)
    # squad = slerp(slerp(qᵢ, qᵢ₊₁, τ), slerp(A, B, τ), 2τ*(1-τ))
    # τ = (t - ta) / (tb - ta)
    # ∂τ∂t = 1 / (tb - ta)
    # ∂(2τ*(1-τ))/∂t = ∂τ/∂t * (2 - 4τ)

    τ = (t - ta) / (tb - ta)
    ∂τ∂t = 1 / (tb - ta)

    X, ∂X∂τ = slerp∂slerp∂τ(qᵢ, qᵢ₊₁, τ)
    ∂X∂t = ∂τ∂t * ∂X∂τ

    Y, ∂Y∂τ = slerp∂slerp∂τ(A, B, τ)
    ∂Y∂t = ∂τ∂t * ∂Y∂τ

    s, ∂s∂X, ∂s∂Y, ∂s∂τparabola = slerp∂slerp(X, Y, 2τ*(1-τ))
    ∂s∂t = ∂τ∂t * (2 - 4τ) * ∂s∂τparabola

    (
        s,
        sum(∂X∂t[b] * ∂s∂X[b] for b in 1:4)
        + sum(∂Y∂t[b] * ∂s∂Y[b] for b in 1:4)
        + ∂s∂t
    )
end


module Mods

import Base: isequal, (==), (+), (-), (*), (inv), (/), (//), (^), hash, show
import Base: zero, one, rand, conj

export Mod, modulus, value, AbstractMod
export isequal, ==, +, -, *, is_invertible, inv, /, ^
export hash, CRT

abstract type AbstractMod <: Number end

"""
`Mod{m}(v)` creates a modular number in mod `m` with value `mod(v,m)`.
`Mod{m}()` is equivalent to `Mod{m}(0)`.
"""
struct Mod{N} <: AbstractMod
    val::Int
    function Mod(x::Int,N::Int)
        @assert N>1 "modulus must be at least 2"
        new{N::Int}(mod(x,N))
    end
end

Mod{N}(x::Int=0) where N = Mod(x,N)

function Mod{N}(x::T) where {N,T<:Integer}
    xx = mod(x,N)
    return Mod{N}(Int(xx))
end
Mod(x::T, N::Int) where T<:Integer = Mod{N}(x)

"""
`modulus(a::Mod)` returns the modulus of this `Mod` number.
```
julia> a = Mod{13}(11);

julia> modulus(a)
13
```
"""
modulus(a::Mod{N}) where N = N

"""
`value(a::Mod)` returns the value of this `Mod` number.
```
julia> a = Mod{13}(11);

julia> value(a)
11
```
"""
value(a::AbstractMod) = a.val

zero(::Mod{N}) where N = Mod{N}()
zero(::Type{Mod{N}}) where N = Mod{N}()

one(::Mod{N}) where N = Mod{N}(1)
one(::Type{Mod{N}}) where N = Mod{N}(1)

conj(x::Mod) = x   # so matrix transpose will work

function hash(x::AbstractMod, h::UInt64= UInt64(0))
    v = value(x)
    m = modulus(x)
    return hash(v,hash(m,h))
end

# Test for equality
isequal(x::Mod, y::Mod) = value(x)==value(y) && modulus(x)==modulus(y)

==(x::Mod,y::Mod) = isequal(x,y)

"""
`modcheck(x::Mod,y::Mod)` checks if `x` and `y` have the same modulus.
If so, return than modulus. If not, throw an error.
"""
function modcheck(x::Mod, y::Mod)::Int
    mx = modulus(x)
    my = modulus(y)
    if mx != my
        error("Cannot operate on two Mod objects with different moduli")
    end
    return mx
end

# Easy arithmetic
function +(x::Mod, y::Mod)::Mod
    m = modcheck(x,y)
    s,flag = Base.add_with_overflow(x.val,y.val)
    if !flag
        return Mod(x.val+y.val, m)
    end
    s = widen(x.val) + widen(y.val)    # add with added precision
    s = Int(mod(s,m))                  # reduce by modulus
    return Mod(s,m)
end

function -(x::Mod{M}) where M
    return Mod(-x.val, M)
end

-(x::Mod,y::Mod) = x + (-y)

function *(x::Mod, y::Mod)
    m = modcheck(x,y)
    p,flag = Base.mul_with_overflow(x.val,y.val)
    if !flag
        return Mod(x.val*y.val, m)
    end
    p = widemul(x.val, y.val)         # multipy with added precision
    p = Int(mod(p,m))                  # reduce by the modulus
    return Mod(p,m) # return with proper type
end

# Division stuff
"""
`is_invertible(x::Mod)` determines if `x` is invertible.
"""
function is_invertible(x::Mod{M})::Bool where M
    return gcd(value(x),M) == 1
end

"""
`inv(x::Mod)` gives the multiplicative inverse of `x`.
This may be abbreviated by `x'`.
"""
function inv(x::Mod{M}) where M
    (g, v, ignore) = gcdx(x.val, M)
    if g != 1
        error("$x is not invertible")
    end
    return Mod(v, M)
end

function /(x::Mod, y::Mod)
    modcheck(x,y)
    return x * inv(y)
end

(//)(x::Mod,y::Mod) = x/y

# Operations with Integers

(+)(x::Mod{M}, k::Integer) where M = Mod(k,M)+x
(+)(k::Integer, x::Mod) = x+k

(-)(x::Mod, k::Integer) = x + (-k)
(-)(k::Integer, x::Mod) = (-x) + k

(*)(x::Mod{M}, k::Integer) where M = Mod(k,M) * x
(*)(k::Integer, x::Mod) = x*k

(/)(x::Mod{M}, k::Integer) where M = x / Mod(k, M)
(/)(k::Integer, x::Mod{M}) where M = Mod(k, M) / x

(//)(x::Mod{M}, k::Integer) where M = x / Mod(k, M)
(//)(k::Integer, x::Mod{M}) where M = Mod(k, M) / x

# Operations with rational numbers

Mod{N}(k::Rational) where N = Mod{N}(numerator(k))/Mod{N}(denominator(k))

function +(x::Mod{N}, k::Rational) where N
    return x + Mod{N}(k)
end
(+)(k::Rational,x::Mod) = x+k

(-)(x::Mod,k::Rational) = x + (-k)
(-)(k::Rational,x::Mod) = k + (-x)

function (*)(x::Mod{N},k::Rational) where N
    return x * Mod{N}(k)
end
(*)(k::Rational,x::Mod) = x*k

function (/)(x::Mod,k::Rational)
    return x * (1/k)
end
(/)(k::Rational,x::Mod) = k * inv(x)

(//)(x::Mod,k::Rational) = x/k
(//)(k::Rational,x::Mod) = k/x

# Comparison with Integers

isequal(x::Mod{M}, k::Integer) where M = mod(k,M) == x.val
isequal(k::Integer, x::Mod) = isequal(x,k)
(==)(x::Mod, k::Integer) = isequal(x,k)
(==)(k::Integer, x::Mod) = isequal(x,k)

# Comparisons with Rationals
function isequal(x::Mod{N}, k::Rational) where N
    return x == Mod{N}(k)
end
isequal(k::Rational,x::Mod) = isequal(x,k)
(==)(x::Mod, k::Rational) = isequal(x,k)
(==)(k::Rational, x::Mod) = isequal(x,k)

# Random

rand(::Type{Mod{N}}) where N = Mod{N}(rand(Int))
rand(::Type{Mod{N}},dims::Integer...) where N = Mod{N}.(rand(Int,dims...))

# Chinese remainder theorem functions

# private helper function
function CRT_work(x::Mod{n}, y::Mod{m}) where {n,m}
    # n = x.mod
    # m = y.mod
    if gcd(n,m) != 1
        error("Moduli must be pairwise relatively prime")
    end

    a = x.val
    b = y.val

    k = inv(Mod(n,m)) * (b-a)

    z = a + k.val*n

    return Mod(z, n*m)
end

# public interface
"""
`CRT(m1,m2,...)`: Chinese Remainder Theorem
```
julia> CRT( Mod(4,11), Mod(8,14) )
Mods.Mod(92,154)

julia> 92%11
4

julia> 92%14
8
```
"""
function CRT(mtuple::Mod...)
    n = length(mtuple)
    if n == 0
        return 1
    end

    result::Mod = mtuple[1]

    for k=2:n
        result = CRT_work(result,mtuple[k])
    end

    return result
end

include("GaussMods.jl")

end # end of module Mods


using JFVM
using PyPlot
include("CRT.jl")

#Calculating error as a function of timesteps

# define the domain and constants
Lx = 1.0 # [m]
ϕ  = 0.4 # porosity
u  = 1.0/(3600*24) # m/s, eq. 1 m/day
D  = [1e-9, 1e-9, 1e-9] # m^2/s diffusivity
K_eq = 10.0
k_f = 100.0
k_b = k_f/K_eq
nu_A = 1.0 # stoichiometric coefficient of component A
nu_B = 1.0 # stoichiometric coefficient of component B
nu_C = 1.0 # stoichiometric coefficient of component C
a = 1.0 #order of reaction
b = 1.0 #order of reaction
c = 1.0 #order of reaction
N_comp = 3 # number of components

# Defining time-stepping parameters
tol = 1e-10 # convergence tolerance

# boundary conditions
C_bc = [1.0, 0.5, 0.0] # boundary concentration
C_0 = zeros(N_comp) # initial concentration

NX = [11, 51, 101, 201, 501, 1001]
err_array = zeros(5, length(NX))

#Calculating reference "true solution"
Nx_ref = 10001;
m_ref = createMesh1D(Nx_ref,Lx)
dt_ref = Lx/(u/ϕ)/Nx_ref
T_ref = Nx_ref/2*dt_ref
C_ref = CRT.CRT_GIA_eps(m_ref, Nx_ref, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt_ref, T_ref, tol)

for i=1:length(NX)
  Nx = NX[i]  # number of grids
  m = createMesh1D(Nx, Lx)  # mesh

  #Time-stepping parameters
  dt = Lx/(u/ϕ)/Nx
  T = Nx/2*dt

  #Calculating alternatives
  C_SNIA        = CRT.CRT_SNIA(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T)
  C_SIA         = CRT.CRT_SIA(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T, tol)
  C_GIA_eps     = CRT.CRT_GIA_eps(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T, tol)
  C_GIA_kin_eq  = CRT.CRT_GIA_kin_eq(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, a, b, c, C_bc, C_0, N_comp, dt, T, tol)
  C_GIA_kin_fb  = CRT.CRT_GIA_kin_fb(m, Nx, ϕ, u, D, k_f, k_b, nu_A, nu_B, nu_C, a, b, c, C_bc, C_0, N_comp, dt, T, tol)

  for j = 1:N_comp
    reference = internalCells(C_ref[j])
    err_array[1, i] += sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]-internalCells(C_SNIA[j])))/sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]))
    err_array[2, i] += sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]-internalCells(C_SIA[j])))/sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]))
    err_array[3, i] += sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]-internalCells(C_GIA_eps[j])))/sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]))
    err_array[4, i] += sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]-internalCells(C_GIA_kin_eq[j])))/sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]))
    err_array[5, i] += sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]-internalCells(C_GIA_kin_fb[j])))/sum(abs(reference[Int(1):Int((Nx_ref-1)/(Nx-1)):Int(Nx_ref)]))
  end
  println(i)
end

figure()
plot(NX, err_array[1,:])
plot(NX, err_array[2,:])
plot(NX, err_array[3,:])
plot(NX, err_array[4,:])
plot(NX, err_array[5,:])

grid("on")
legend(["SNIA", "SIA", "GIA ("L"$\epsilon$"")", "GIA ("L"$K_{eq}$"")", "GIA ("L"$k_fk_b$"")"])
xlabel(L"$N_x$")
ylabel(L"||$C_{ref}-C||_1/||C_{ref}||_1$")
savefig("error_grid_plot.png")

figure()
semilogy(NX, err_array[1,:])
semilogy(NX, err_array[2,:])
semilogy(NX, err_array[3,:])
semilogy(NX, err_array[4,:])
semilogy(NX, err_array[5,:])

grid("on")
legend(["SNIA", "SIA", "GIA ("L"$\epsilon$"")", "GIA ("L"$K_{eq}$"")", "GIA ("L"$k_fk_b$"")"])
xlabel(L"$N_x$")
ylabel(L"log(||$C_{ref}-C||_1/||C_{ref}||_1)$")
savefig("error_grid_semilogy.png")

#CRT_SIA title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (SIA)"

#CRT_SNIA title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (SNIA)"

#CRT_GIA_eps title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (FIM)"

#CRT_GIA_kin_eq title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}(K_{eq}C_A^aC_B^b-C_C^c)=0,\quad K_{eq}=$""$K_eq"" (FIM)"

#CRT_GIA_kin_fb title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}(k_{f}C_A^aC_B^b-k_{b}C_C^c)=0,\quad K_{eq}=$""$K_eq"" (FIM)"


using Test, LinearFractionalTransformations, LinearAlgebra

@testset "Basics" begin
    f = LFT(1, 0, 0, 1)  # identity
    @test f(1 + im) == 1 + im

    f = LFT(2, 0, 0, 1)  # f(x) = 2x
    @test f(-im) == -2im

    g = inv(f)
    @test g(im) == 0.5im

    g = f * f
    @test g(2 - im) == 4 * (2 - im)

    f = LFT(1, 2 + im, 3, 0, 1 - im, 5)
    @test f(1) == 2 + im
    @test f(3) == 0
    @test f(1 - im) == 5

    g = inv(f)
    @test g(0) == 3

    @test LFT() == g * f

    f = LFT(0, 1, 1, 0)
    @test f(Inf) == 0
    @test isinf(f(0))

    f = LFT(3 - im, 2 + im, 6)
    @test f(3 - im) == 0
    @test f(2 + im) == 1
    @test isinf(f(6))

end

# check out LFTQ
@testset "LFTQ" begin
    for k = 1:5
        Q = [1 0; 0 -1]
        while det(Q) < 0.5
            Q, R = qr(randn(3, 3))
        end
        F = LFTQ(Q)
        v = randn() + randn() * im

        w1 = stereo(Q * stereo(v))
        w2 = F(v)
        @test abs(w1 - w2) <= 1e-8

        z = stereo(Q * stereo(0))
        @test abs(z - F(0)) <= 1e-8

        z = stereo(Q * stereo(1))
        @test abs(z - F(1)) <= 1e-8

        z = stereo(Q * stereo(Inf))
        @test abs(z - F(Inf)) <= 1e-8

    end
end


using Bridge, Distributions, StaticArrays

PLOT = :winston
include("plot.jl")

tt = range(0, stop=1, length=1000)
G =  GammaProcess(10., 1.)
X = sample(tt, G)
plot(X, "-")
GB =  GammaBridge(1., 1., G)
XB = sample(tt, GB)
oplot(XB, "-")

VG =  VarianceGammaProcess(0., 1., 1.)
X = sample(tt, VG)
plot(X, "-")

#n = 1000
#simulate gamma process:
#tt = cumsum(rand(Gamma(0.01, 1/(0.01n)), n))
#plot(tt)

X = sample(range(0, stop=1, length=500),Bridge.GammaBridge(1,1, GammaProcess(10,10)))
X = sample(range(0, stop=1, length=500), GammaProcess(10,10))


# v0.6

using Primes

function mersennefactor(p::Int)::Int
    q = 2p + 1
    while true
        if log2(q) > p / 2
            return -1
        elseif q % 8 in (1, 7) && Primes.isprime(q) && powermod(2, p, q) == 1
            return q
        end
    q += 2p
    end
end

for i in filter(Primes.isprime, push!(collect(1:20), 929))
    mf = mersennefactor(i)
    if mf != -1 println("M$i = ", mf, " × ", (big(2) ^ i - 1) ÷ mf)
    else println("M$i is prime") end
end