import Base.rand

export Ellipsoid,
       shape_matrix

"""
    Ellipsoid{N<:AbstractFloat, VN<:AbstractVector{N},
              MN<:AbstractMatrix{N}} <: AbstractCentrallySymmetric{N}

Type that represents an ellipsoid.

It is defined as the set

```math
E = \\left\\{ x ∈ \\mathbb{R}^n : (x-c)^T Q^{-1} (x-c) ≤ 1 \\right\\},
```
where ``c \\in \\mathbb{R}^n`` is its *center* and ``Q \\in \\mathbb{R}^{n×n}``
its *shape matrix*, which should be a positive definite matrix.
An ellipsoid can also be characterized as the image of a Euclidean ball by an
invertible linear transformation. It is the higher-dimensional generalization
of an ellipse.

### Fields

- `center`       -- center of the ellipsoid
- `shape_matrix` -- real positive definite matrix, i.e. it is equal to its transpose
                    and ``x^\\mathrm{T}Qx > 0`` for all nonzero ``x``

## Notes

By default, the inner constructor checks that the given shape matrix is positive
definite. Use the flag `check_posdef=false` to disable this check.

### Examples

An ellipsoid can be created passing its center and shape matrix (which should be
positive definite).

For instance, we create a two-dimensional ellipsoid with center `[1, 1]`:

```jldoctest ellipsoid_constructor
julia> using LinearAlgebra

julia> E = Ellipsoid(ones(2), Diagonal([2.0, 0.5]))
Ellipsoid{Float64,Array{Float64,1},Diagonal{Float64,Array{Float64,1}}}([1.0, 1.0], [2.0 0.0; 0.0 0.5])
```

If the center is not specified, it is assumed that it is the origin. For instance,
a three-dimensional ellipsoid centered at the origin and the shape matrix being
the identity can be created with:

```jldoctest ellipsoid_constructor
julia> E = Ellipsoid(Matrix(1.0I, 3, 3))
Ellipsoid{Float64,Array{Float64,1},Array{Float64,2}}([0.0, 0.0, 0.0], [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0])

julia> dim(E)
3
```
The center and shape matrix of the ellipsoid can be retrieved with the functions
`center` and `shape_matrix` respectively:

```jldoctest ellipsoid_constructor
julia> center(E)
3-element Array{Float64,1}:
 0.0
 0.0
 0.0

julia> shape_matrix(E)
3×3 Array{Float64,2}:
 1.0  0.0  0.0
 0.0  1.0  0.0
 0.0  0.0  1.0
```

The function `an_element` returns some element of the ellipsoid:

```jldoctest ellipsoid_constructor
julia> an_element(E)
3-element Array{Float64,1}:
 0.0
 0.0
 0.0

julia> an_element(E) ∈ E
true
```

We can evaluate its support vector in a given direction, say `[1, 1, 1]`:

```jldoctest ellipsoid_constructor
julia> σ(ones(3), E)
3-element Array{Float64,1}:
 0.5773502691896258
 0.5773502691896258
 0.5773502691896258
```
"""
struct Ellipsoid{N<:AbstractFloat, VN<:AbstractVector{N},
                 MN<:AbstractMatrix{N}} <: AbstractCentrallySymmetric{N}
    center::VN
    shape_matrix::MN

    # default constructor with dimension check
    function Ellipsoid(c::VN, Q::MN;
                       check_posdef::Bool=true) where {N<:AbstractFloat,
                                                       VN<:AbstractVector{N},
                                                       MN<:AbstractMatrix{N}}

        @assert length(c) == checksquare(Q) "the length of the center and the size "
            "of the shape matrix do not match; they are $(length(c)) and $(size(Q)) respectively"

        if check_posdef
            isposdef(Q) || throw(ArgumentError("an ellipsoid's shape matrix " *
                                               "must be positive definite"))
        end
        return new{N, VN, MN}(c, Q)
    end
end

isoperationtype(::Type{<:Ellipsoid}) = false
isconvextype(::Type{<:Ellipsoid}) = true

# convenience constructor for ellipsoid centered in the origin
function Ellipsoid(Q::AbstractMatrix{N}; check_posdef::Bool=true) where {N<:AbstractFloat}
    # TODO: use similar vector type for the center, see #2032
    return Ellipsoid(zeros(N, size(Q, 1)), Q; check_posdef=check_posdef)
end

# --- AbstractCentrallySymmetric interface functions ---

"""
    center(E::Ellipsoid{N}) where {N<:AbstractFloat}

Return the center of the ellipsoid.

### Input

- `E` -- ellipsoid

### Output

The center of the ellipsoid.
"""
function center(E::Ellipsoid{N}) where {N<:AbstractFloat}
    return E.center
end

"""
    shape_matrix(E::Ellipsoid{N}) where {N<:AbstractFloat}

Return the shape matrix of the ellipsoid.

### Input

- `E` -- ellipsoid

### Output

The shape matrix of the ellipsoid.
"""
function shape_matrix(E::Ellipsoid{N}) where {N<:AbstractFloat}
    return E.shape_matrix
end

# --- LazySet interface functions ---

"""
    σ(d::AbstractVector{N}, E::Ellipsoid{N}) where {N<:AbstractFloat}

Return the support vector of an ellipsoid in a given direction.

### Input

- `d` -- direction
- `E` -- ellipsoid

### Output

Support vector in the given direction.

### Algorithm

Let ``E`` be an ellipsoid of center ``c`` and shape matrix ``Q = BB^\\mathrm{T}``.
Its support vector along direction ``d`` can be deduced from that of the unit
Euclidean ball ``\\mathcal{B}_2`` using the algebraic relations for the support
vector,

```math
σ_{B\\mathcal{B}_2 ⊕ c}(d) = c + Bσ_{\\mathcal{B}_2} (B^\\mathrm{T} d)
= c + \\dfrac{Qd}{\\sqrt{d^\\mathrm{T}Q d}}.
```
"""
function σ(d::AbstractVector{N}, E::Ellipsoid{N}) where {N<:AbstractFloat}
    if norm(d, 2) == zero(N)
        return an_element(E)
    end
    α = sqrt(dot(d, E.shape_matrix * d))
    return E.center .+ E.shape_matrix * d ./ α
end

"""
    ρ(d::AbstractVector{N}, E::Ellipsoid{N}) where {N<:AbstractFloat}

Return the support function of an ellipsoid in a given direction.

### Input

- `d` -- direction
- `E` -- ellipsoid

### Output

The support function of the ellipsoid in the given direction.

### Algorithm

The support value is ``cᵀ d + ‖Bᵀ d‖₂`` where ``c`` is the center and
``Q = B Bᵀ`` is the shape matrix of `E`.
"""
function ρ(d::AbstractVector{N}, E::Ellipsoid{N}) where {N<:AbstractFloat}
    return dot(center(E), d) + sqrt(inner(d, E.shape_matrix, d))
end

"""
    ∈(x::AbstractVector{N}, E::Ellipsoid{N}) where {N<:AbstractFloat}

Check whether a given point is contained in an ellipsoid.

### Input

- `x` -- point/vector
- `E` -- ellipsoid

### Output

`true` iff `x ∈ E`.

### Algorithm

The point ``x`` belongs to the ellipsoid of center ``c`` and shape matrix ``Q``
if and only if

```math
(x-c)^\\mathrm{T} Q^{-1} (x-c) ≤ 1.
```
"""
function ∈(x::AbstractVector{N}, E::Ellipsoid{N}) where {N<:AbstractFloat}
    @assert length(x) == dim(E)
    w, Q = x-E.center, E.shape_matrix
    return dot(w, Q \ w) ≤ 1
end

"""
    rand(::Type{Ellipsoid}; [N]::Type{<:AbstractFloat}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing)

Create a random ellipsoid.

### Input

- `Ellipsoid` -- type for dispatch
- `N`         -- (optional, default: `Float64`) numeric type
- `dim`       -- (optional, default: 2) dimension
- `rng`       -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`      -- (optional, default: `nothing`) seed for reseeding

### Output

A random ellipsoid.

### Algorithm

The center is a normally distributed vector with entries of mean 0 and standard
deviation 1.

The idea for the shape matrix comes from
[here](https://math.stackexchange.com/a/358092).
The matrix is symmetric positive definite, but also diagonally dominant.

```math
Q =  \\frac{1}{2}(S + S^T) + nI,
```
where ``n`` = `dim` (defaults to 2), and ``S`` is a ``n \\times n`` random
matrix whose coefficients are uniformly distributed in the interval ``[-1, 1]``.
"""
function rand(::Type{Ellipsoid};
              N::Type{<:AbstractFloat}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing)
    rng = reseed(rng, seed)
    center = randn(rng, N, dim)
    # random entries in [-1, 1]
    # this needs a bit of code because 'rand' only samples from [0, 1]
    shape_matrix = Matrix{N}(undef, dim, dim)
    for j in 1:dim
        for i in 1:dim
            entry = rand(N)
            if rand(Bool)
                entry = -entry
            end
            shape_matrix[i, j] = entry
        end
    end
    # make diagonally dominant
    shape_matrix = N(0.5) * (shape_matrix + shape_matrix') +
                   Matrix{N}(dim*I, dim, dim)
    return Ellipsoid(center, shape_matrix)
end

"""
    translate(E::Ellipsoid{N}, v::AbstractVector{N}; share::Bool=false
             ) where {N<:AbstractFloat}

Translate (i.e., shift) an ellipsoid by a given vector.

### Input

- `E`     -- ellipsoid
- `v`     -- translation vector
- `share` -- (optional, default: `false`) flag for sharing unmodified parts of
             the original set representation

### Output

A translated ellipsoid.

### Notes

The shape matrix is shared with the original ellipsoid if `share == true`.

### Algorithm

We add the vector to the center of the ellipsoid.
"""
function translate(E::Ellipsoid{N}, v::AbstractVector{N}; share::Bool=false
                  ) where {N<:AbstractFloat}
    @assert length(v) == dim(E) "cannot translate a $(dim(E))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    c = center(E) + v
    shape_matrix = share ? E.shape_matrix : copy(E.shape_matrix)
    return Ellipsoid(c, shape_matrix)
end


#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5

# Short description of this file: Simple Example Neural PDE

using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Plots
#GalacticOptim
# 3D PDE
@parameters x y t θ
@variables u(..)
@derivatives Dxx''~x
@derivatives Dyy''~y
@derivatives Dt'~t

# 3D PDE
eq  = Dt(u(x,y,t,θ)) ~ Dxx(u(x,y,t,θ)) + Dyy(u(x,y,t,θ))
# Initial and boundary conditions
bcs = [u(x,y,0,θ) ~ exp(x+y)*cos(x+y),
       u(0,y,t,θ) ~ exp(y)*cos(y+4t),
       u(2,y,t,θ) ~ exp(2+y)*cos(2+y+4t) ,
       u(x,0,t,θ) ~ exp(x)*cos(x+4t),
       u(x,2,t,θ) ~ exp(x+2)*cos(x+2+4t)]
# Space and time domains
domains = [x ∈ IntervalDomain(0.0,2.0),
           y ∈ IntervalDomain(0.0,2.0),
           t ∈ IntervalDomain(0.0,2.0)]

# Discretization
dx = 0.25; dy= 0.25; dt = 0.25
# Neural network
chain = FastChain(FastDense(3,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))

discretization = NeuralPDE.PhysicsInformedNN([dx,dy,dt],
                                             chain,
                                             strategy = NeuralPDE.StochasticTraining(include_frac=0.9))

pde_system = PDESystem(eq,bcs,domains,[x,y,t],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

res = GalacticOptim.solve(prob, ADAM(0.1), progress = false; cb = cb, maxiters=3000)
phi = discretization.phi

ts,xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains]

analytic_sol_func(t,x) = [exp(-t)*sin(pi*x), exp(-t)*cos(pi*x), (1+pi^2)*exp(-t)]
u_real  = [[analytic_sol_func(t,x)[i] for t in ts for x in xs] for i in 1:3]

u_predict  = [[phi([t,x],res.minimizer)[i] for t in ts for x in xs] for i in 1:3]
diff_u = [abs.(u_real[i] .- u_predict[i] ) for i in 1:3]

for i in 1:3
    p1 = plot(xs, ts, u_real[i], st=:surface,title = "u$i, analytic");
    p2 = plot(xs, ts, u_predict[i], st=:surface,title = "predict");
    p3 = plot(xs, ts, diff_u[i],linetype=:contourf,title = "error");
    plot(p1,p2,p3)
    savefig("sol_u$i")
end


"""
Module for statistics for common factor models (or approximate factor models), by Colin T Bowers \n
\n
This module implements the following statistics for use with common factor models: \n
    - numfactor: Estimating the number of factors using the methods from Bai, Ng (2002) Determining the Number
    of Factors in Approximate Factor Models
    - testobsfactor: Testing whether an observed factor spans the latent common factor space using the methods
    from Bai, Ng (2006) Evaluating Latent and Observed Factors in Macroeconomics and Finance
\n
This package has an MIT license. Please see associated LICENSE.md file for more detail.
"""
module CommonFactorModelStats

using Distributions

include("est_num_factor.jl")
include("test_obs_fac.jl")

# package code goes here

end # module


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# Hock-Schittkowski Nonlinear Test Suite  -  HS118
# This file is JuMP implementation of the model described in
#  W. Hock, K. Schittkowski, Test Examples for Nonlinear Programming
#  Codes, Lecture Notes in Economics and Mathematical Systems,
#  Springer, No, 187, 1981
# More information, including original model description, at
# http://www.ai7.uni-bayreuth.de/downloads.htm
#
# This problem has a quadratic objective with linear constraints.
#############################################################################

using JuMP
using Base.Test

let

m = Model()

L = zeros(15)
L[1] =  8.0
L[2] = 43.0
L[3] =  3.0

U = zeros(15)
U[1] = 21.0
U[2] = 57.0
U[3] = 16.0
for k in 1:4
    U[3*k+1] =  90.0
    U[3*k+2] = 120.0
    U[3*k+3] =  60.0
end

@variable(m, L[i] <= x[i=1:15] <= U[i])

@NLobjective(m, Min,
    sum{2.3     * x[3*k+1]   +
        0.0001  * x[3*k+1]^2 +
        1.7     * x[3*k+2]   +
        0.0001  * x[3*k+2]^2 +
        2.2     * x[3*k+3] +
        0.00015 * x[3*k+3]^2, k=0:4})

# setobjective(m, :Min,
#   sum([(2.3     * x[3*k+1]   +
#         0.0001  * x[3*k+1]^2 +
#         1.7     * x[3*k+2]   +
#         0.0001  * x[3*k+2]^2 +
#         2.2     * x[3*k+3]^2 +
#         0.00015 * x[3*k+3]^2) for k in 0:4 ]))

# constr1
for j in 1:4
    @constraint(m, x[3*j+1] - x[3*j-2] + 7 <= 13)
    @constraint(m, x[3*j+1] - x[3*j-2] + 7 >=  0)
end

# constr2
for j in 1:4
    @constraint(m, x[3*j+2] - x[3*j-1] + 7 <= 14)
    @constraint(m, x[3*j+2] - x[3*j-1] + 7 >=  0)
end

# constr3
for j in 1:4
    @constraint(m, x[3*j+3] - x[3*j  ] + 7 <= 13)
    @constraint(m, x[3*j+3] - x[3*j  ] + 7 >=  0)
end

@constraint(m, x[1] + x[2] + x[3]    >= 60)
@constraint(m, x[4] + x[5] + x[6]    >= 50)
@constraint(m, x[7] + x[8] + x[9]    >= 70)
@constraint(m, x[10] + x[11] + x[12] >= 85)
@constraint(m, x[13] + x[14] + x[15] >= 100)

# Initial solution (could also use 'start' keyword in @variable)
setvalue(x[1], 20.0)
setvalue(x[2], 55.0)
setvalue(x[3], 15.0)
setvalue(x[4], 20.0)
setvalue(x[5], 60.0)
setvalue(x[6], 20.0)
setvalue(x[7], 20.0)
setvalue(x[8], 60.0)
setvalue(x[9], 20.0)
setvalue(x[10], 20.0)
setvalue(x[11], 60.0)
setvalue(x[12], 20.0)
setvalue(x[13], 20.0)
setvalue(x[14], 60.0)
setvalue(x[15], 20.0)

solve(m)

#println(getvalue(x))

@test_approx_eq_eps getvalue(x[1]) 8.0  1e-5
@test_approx_eq_eps getvalue(x[2]) 49.0 1e-5
@test_approx_eq_eps getvalue(x[3]) 3.0  1e-5
@test_approx_eq_eps getvalue(x[4]) 1.0  1e-5
@test_approx_eq_eps getobjectivevalue(m) 664.82045 1e-5

end


@testset "PertBeta" begin
    @testset "basic" begin
        dist = PertBeta(4, 10, 20)
        @test minimum(dist) == 4
        @test maximum(dist) == 20
        @test mode(dist) == 10
        @test mean(dist) == (4 + 4*10 + 20)/6
    end

    @testset "$dist" for dist in (PertBeta(4, 10, 20), PertBeta(0.0,0.5,1.0), PertBeta(1.0, 1.01, 10.0), )
        samples = rand(dist, 1_000_000)
        @test minimum(dist) < minimum(samples)
        @test maximum(samples) < maximum(dist)
        @test mean(samples) ≈ mean(dist) atol=0.01
        @test std(samples) ≈ std(dist) atol=0.01
    end

    @testset "Constructor" begin
        @test_throws Exception PertBeta(0, 10, 5)  # mode must be less than max
        @test_throws Exception PertBeta(5, 1, 7)  # mode must be more than min
    end
end


# Air Constants
const GAMMA_AIR = 1.4  # Adiabatic index or ratio of specific heats (dry air at 20º C)
const R_AIR = 287.05287  # Specific gas constant for dry air (J/(Kg·K))

# Air at sea level conditions h=0 (m)
const RHO0 = 1.225  # Density at sea level (kg/m3)
const P0 = 101325.0  # Pressure at sea level (Pa)
const T0 = 288.15  # Temperature at sea level (K)
const A0 = 340.293990543  # Sound speed at sea level (m/s)

# EARTH CONSTANTS
const GRAVITY_ACCEL = 9.80665  # Gravity of Ethe Earth (m/s^2)
# Standard Gravitational Parameter
# product of the gravitational constant G and the mass M of the body
# TODO: review values with
# Rogers, R. M. (2007). Applied mathematics in integrated navigation systems.
# American Institute of Aeronautics and Astronautics. (Page 76, table 4.1)
const ROT_VELOCITY = 7.292115e-5  # Angular velocity ω_i/e (rad/s)
const STD_GRAV_PARAMETER = 3.986004418e14  # (m³/s²)
const EARTH_MASS = 5.9722e24  # Mass of the Earth (kg)
const GRAV_CONSTANT = 6.67384e11  # Gravitational constant (N·m²/kg²)
const EARTH_MEAN_RADIUS = 6371000  # Mean radius of the Earth (m)

# Ellipsoids
struct Ellipsoid
    a :: Real  # semi-major axis
    finv :: Real  # inverse of flattening
    f :: Real  # flattening
    b :: Real  # semi-minor axis
    e2 :: Real  # eccentricity squared
    ϵ2 :: Real  # second eccentricity squared
end

Ellipsoid(a, finv) = begin
    f = 1.0 / finv
    b = a * (1.0 - f)
    e2 = 1.0 - (1.0 - f) ^ 2.0
    ϵ2 = a*a / (b*b) - 1.0
    Ellipsoid(a, finv, f, b, e2, ϵ2)
end

# Rogers, R. M. (2007). Applied mathematics in integrated navigation systems.
# American Institute of Aeronautics and Astronautics. (Page 76, table 4.1)
Clarke1866           = Ellipsoid(6378206.4  , 294.9786982)
Clarke1880           = Ellipsoid(6378249.145, 294.465)
International        = Ellipsoid(6378388.0  , 297.0)
Bessel               = Ellipsoid(6377397.155, 299.1528128)
Everest              = Ellipsoid(6377276.345, 300.8017)
ModifiedEverest      = Ellipsoid(6377304.063, 300.8017)
AustralianNational   = Ellipsoid(6378160.0  , 298.25)
SouthAmerican1969    = Ellipsoid(6378160.0  , 298.25)
Airy                 = Ellipsoid(6377564.396, 299.3249646)
ModifiedAiry         = Ellipsoid(6377340.189, 299.3249646)
Hough                = Ellipsoid(6378270.0  , 297.0)
Fischer1960SouthAsia = Ellipsoid(6378155.0  , 298.3)
Fischer1960Mercury   = Ellipsoid(6378166.0  , 298.3)
Fischer1968          = Ellipsoid(6378150.0  , 298.3)
WGS60                = Ellipsoid(6378165.0  , 298.3)
WGS66                = Ellipsoid(6378145.0  , 298.25)
WGS72                = Ellipsoid(6378135.0  , 298.26)
WGS84                = Ellipsoid(6378137.0  , 298.257223563)

# CONVERSIONS
const LB2KG = 0.453592  # Pounds (lb) to kilograms (kg)
const KG2LB = 2.20462  # Kilograms (kg) to pounds (lb)
const FT2M = 0.3048  # Feet (ft) to meters (m)
const M2FT = 3.28084  # Meters (m) to feet (ft)
const IN2M = 0.0254  # Inches (in) to meters (m)
const M2IN = 39.37007874  # meters (m) to inches(in)
const SLUG2KG = 14.5939  # Slug to kilograms (kg)
const SLUGFT2_2_KGM2 = 1.35581795  # Slug*feet^2 to kilograms*meters^2 (kg*m^2)
const SLUGFT3_2_KGM3 = SLUG2KG / FT2M^3  # Slug*feet^3 to kilograms*meters^3 (kg*m^3)
const KT2MS = 0.514444  # knots (kt) to meters/second (m/s)
const MS2KT = 1.94384  # meters/second (m/s) to knots (kt)
const DEG2RAD = pi / 180  # degrees to radians
const RAD2DEG = 180 / pi  # radians to degrees
const HP2WAT = 745.7  # horse power to Watts
const WAT2HP = 0.00134102  # Wat to horse power
const LBF2N = 4.44822  # pounds force (lbf) to N (Newton)
const PA2PSF = 0.02088543  # Pascal to Pounds per square foot
const PSF2PA = 47.880258  # Pounds per square foot to Pascal
const RANK2KEL = 0.555556  # Rankine to Kelvin
const KEL2RANK = 1.8  # Kelvin to Rankine


using Test, DelayEmbeddings

println("\nTesting utils...")

@testset "utils" begin

    A = orthonormal(50, 50)
    @test size(A) == (50, 50)
    @test A isa Matrix

    B = orthonormal(3, 3)
    @test B isa SMatrix
end


module SmithNormalForm

# import Base: show, getindex
using LinearAlgebra
using SparseArrays

export snf, smith

include("bezout.jl")
include("snf.jl")

struct Smith{P,Q<:AbstractMatrix} <: Factorization{P}
    S::Q
    Sinv::Q
    T::Q
    Tinv::Q
    SNF::AbstractVector{P}
    Smith{P,Q}(S::AbstractMatrix{P}, Sinv::AbstractMatrix{P},
               T::AbstractMatrix{P}, Tinv::AbstractMatrix{P},
               A::AbstractVector{P}) where {P,Q} = new(S, Sinv, T, Tinv, A)
end
Smith(S::AbstractMatrix{P}, T::AbstractMatrix{P}, A::AbstractVector{P}) where {P} =
    Smith{P,typeof(A)}(S, similar(S, 0, 0), T, similar(T, 0, 0), A)

function smith(X::AbstractMatrix{P}; inverse=true, debug=false) where {P}
    S, T, A, Sinv, Tinv = snf(X, inverse=inverse, debug=debug)
    return Smith{P, typeof(X)}(S, Sinv, T, Tinv, diag(A))
end

"""Retrive SNF matrix from the factorization"""
function LinearAlgebra.diagm(F::Smith{P,Q}) where {P,Q}
    base = issparse(F.SNF) ? spzeros(P, size(F.S,1), size(F.T,1)) : zeros(P, size(F.S,1), size(F.T,1))
    for i in 1:length(F.SNF)
        base[i,i] = F.SNF[i]
    end
    return base
end

function Base.show(io::IO, F::Smith)
    println(io, "Smith normal form:")
    show(io, diagm(F))
end

end # module


# Cliffords

@test Clifford(1/sqrt(2) * [1 1; 1 -1]) == H
@test Clifford(eye(2)) == RI
@test Clifford(eye(4)) == kron(RI,RI)
@test Clifford([1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0]) == CNOT
@test Clifford(diagm(0 => [1,1,1,-1])) == CZ
@test Clifford(exp(-im*pi/4*[1 0; 0 -1])) == S

function eq_upto_phase(A, B)
    idx = findfirst(x -> x != 0, A)
    rel_phase = (B[idx] == 0) ? 1.0 : (A[idx] / B[idx])
    return rel_phase * A ≈ B
end
@test eq_upto_phase(complex(RI), eye(2))
@test eq_upto_phase(complex(H), 1/sqrt(2) * [1 1; 1 -1])
@test eq_upto_phase(complex(CNOT), [1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0])

# Cliffords * Paulis
@test H * Id == Id
@test H * X == Z
@test H * Z == X
@test H * Y == -Y

@test kron(RI, H) * II == II
@test kron(RI, H) * IX == IZ
@test kron(RI, H) * IZ == IX
@test kron(RI, H) * IY == -IY

@test CNOT * II == II
@test CNOT * IX == IX
@test CNOT * XI == XX
@test CNOT * YY == -XZ

# Cliffords \ Paulis
@test H \ Id == Id
@test H \ X == Z
@test H \ Z == X
@test H \ Y == -Y

@test kron(RI, H) \ II == II
@test kron(RI, H) \ IX == IZ
@test kron(RI, H) \ IZ == IX
@test kron(RI, H) \ IY == -IY

@test CNOT \ II == II
@test CNOT \ IX == IX
@test CNOT \ XI == XX
@test CNOT \ YY == -XZ

# Cliffords * Cliffords

@test H * S * H == Clifford(Dict(+X=>+X,+Z=>-Y),Dict(+X=>+X,+Z=>+Y))
@test H * S * H == Clifford(exp(-im*pi/4*[0 1; 1 0]))
@test RY * RX * RY == SelfInverseClifford(Dict(+X=>+X,+Z=>-Z))
@test kron(RI, H) * CNOT * kron(RI, H) == SelfInverseClifford(Dict(+IZ=>+IZ,+ZI=>+ZI,+XI=>+XZ,+IX=>+ZX))
@test kron(RI, H) * CNOT * kron(RI, H) == CZ
@test S * S == RZ

# 3 CNOTs is a SWAP
CNOT21 = Cliffords.expand(CNOT, [2,1], 2)
@test CNOT * CNOT21 * CNOT == SWAP

# iSWAP
iSWAP = CZ*kron(S,S)*SWAP
@test iSWAP == Clifford([1 0 0 0; 0 0 im 0; 0 im 0 0; 0 0 0 1])

# Cliffords \ Cliffords
@test H \ H == H * H
@test CNOT \ CNOT == CNOT * CNOT
@test RY * RX * RY == RY \ (RX * RY)
@test S \ S == RI
@test S \ RZ == S

# Local Cliffords
@test localclifford(1) == RI

# Cliffords -> Pauli
@test Clifford(Id) == RI
@test Clifford(X)  == RX
@test Clifford(Y)  == RY
@test Clifford(Z)  == RZ
@test Clifford(XX) == kron(RX,RX)
@test Clifford(XI) == kron(RX,RI)

# Clifford and Matrix interaction
@test typeof(promote(H, eye(2))) == Tuple{Array{Float64,2}, Array{Float64,2}}
@test typeof(promote(H, eye(ComplexF64, 2))) == Tuple{Array{ComplexF64,2}, Array{ComplexF64,2}}


module Outliers

export tukey_fences

"""
    tukey_fences(data; k=1.5)

Compute the values beyond which elements of `data` are considered anomalous by
to Tukey (1977; John W, Exploratory Data Analysis, Addison-Wesley, ISBN
0-201-07616-0, OCLC 3058187).  Larger values of `k` consider fewer elements
to be anomalous.
"""
function tukey_fences(data; k=1.5)
    q1,q3 = quantile(data,[0.25,0.75])
    iqr = q3-q1
    fence = k*iqr
    q1-fence, q3+fence
end

end # module


using ContinuumArrays, LinearAlgebra, FastTransforms, QuasiArrays, Test
import ContinuumArrays: Basis, Weight, Map, LazyQuasiArrayStyle, TransformFactorization, ExpansionLayout

"""
This is a simple implementation of Chebyshev for testing. Use ClassicalOrthogonalPolynomials
for the real implementation.
"""
struct Chebyshev <: Basis{Float64}
    n::Int
end

struct ChebyshevWeight <: Weight{Float64} end

Base.:(==)(::Chebyshev, ::Chebyshev) = true
Base.:(==)(::ChebyshevWeight, ::ChebyshevWeight) = true
Base.axes(T::Chebyshev) = (Inclusion(-1..1), Base.OneTo(T.n))
ContinuumArrays.grid(T::Chebyshev) = chebyshevpoints(Float64, T.n, Val(1))
Base.axes(T::ChebyshevWeight) = (Inclusion(-1..1),)

Base.getindex(::Chebyshev, x::Float64, n::Int) = cos((n-1)*acos(x))
Base.getindex(::ChebyshevWeight, x::Float64) = 1/sqrt(1-x^2)
Base.getindex(w::ChebyshevWeight, ::Inclusion) = w # TODO: make automatic

LinearAlgebra.factorize(L::Chebyshev) =
    TransformFactorization(grid(L), plan_chebyshevtransform(Array{Float64}(undef, size(L,2))))
LinearAlgebra.factorize(L::Chebyshev, n) =
    TransformFactorization(grid(L), plan_chebyshevtransform(Array{Float64}(undef, size(L,2),n),1))

# This is wrong but just for tests
QuasiArrays.layout_broadcasted(::Tuple{ExpansionLayout,Any}, ::typeof(*), a::ApplyQuasiVector{<:Any,typeof(*),<:Tuple{Chebyshev,Any}}, b::Chebyshev) = b * Matrix(I, 5, 5)

struct QuadraticMap{T} <: Map{T} end
struct InvQuadraticMap{T} <: Map{T} end

QuadraticMap() = QuadraticMap{Float64}()
InvQuadraticMap() = InvQuadraticMap{Float64}()

Base.getindex(::QuadraticMap, r::Number) = 2r^2-1
Base.axes(::QuadraticMap{T}) where T = (Inclusion(0..1),)
Base.axes(::InvQuadraticMap{T}) where T = (Inclusion(-1..1),)
Base.getindex(d::InvQuadraticMap, x::Number) = sqrt((x+1)/2)
ContinuumArrays.invmap(::QuadraticMap{T}) where T = InvQuadraticMap{T}()
ContinuumArrays.invmap(::InvQuadraticMap{T}) where T = QuadraticMap{T}()

@testset "Chebyshev" begin
    T = Chebyshev(5)
    w = ChebyshevWeight()
    wT = w .* T
    x = axes(T,1)

    @testset "basics" begin
        F = factorize(T)
        g = grid(F)
        @test T \ exp.(x) == F \ exp.(x) == F \ exp.(g) == chebyshevtransform(exp.(g), Val(1))
        @test all(checkpoints(T) .∈ Ref(axes(T,1)))

        @test T \ [exp.(x) cos.(x)] ≈ [F \ exp.(x) F \ cos.(x)]
    end

    @testset "Weighted" begin
        @test MemoryLayout(w) isa WeightLayout
        @test MemoryLayout(w[Inclusion(0..1)]) isa WeightLayout

        @test wT == wT

        wT2 = w .* T[:,2:4]
        wT3 = wT[:,2:4]
        @test MemoryLayout(wT) isa WeightedBasisLayout
        @test MemoryLayout(wT2) isa WeightedBasisLayout
        @test MemoryLayout(wT3) isa SubWeightedBasisLayout
        @test grid(wT) == grid(wT2) == grid(wT3) == grid(T)

        @test ContinuumArrays.unweighted(wT) ≡ T
        @test ContinuumArrays.unweighted(wT2) ≡ T[:,2:4]
        @test ContinuumArrays.unweighted(wT3) ≡ T[:,2:4]

        @test ContinuumArrays.weight(wT) ≡ ContinuumArrays.weight(wT2) ≡ ContinuumArrays.weight(wT3) ≡ w

        @test wT \ @.(exp(x) / sqrt(1-x^2)) ≈ T \ exp.(x)

        @test wT \ w ≈ [1; zeros(4)]

        @test (x .* wT)[0.1,:] ≈ 0.1 * wT[0.1,:]

        a = wT / wT \ @.(exp(x) / sqrt(1-x^2))
        @test wT \ (a .* T) == I # fake multiplication matrix
    end
    @testset "Mapped" begin
        y = affine(0..1, x)

        @test summary(T[y,:]) == "Chebyshev affine mapped to 0..1"
        @test MemoryLayout(wT[y,:]) isa MappedWeightedBasisLayout
        @test MemoryLayout(w[y] .* T[y,:]) isa MappedWeightedBasisLayout
        @test wT[y,:][[0.1,0.2],1:5] == (w[y] .* T[y,:])[[0.1,0.2],1:5] == (w .* T[:,1:5])[y,:][[0.1,0.2],:]
        @test MemoryLayout(wT[y,1:3]) isa MappedWeightedBasisLayout
        @test wT[y,1:3][[0.1,0.2],1:2] == wT[y[[0.1,0.2]],1:2]

        @testset "QuadraticMap" begin
            m = QuadraticMap()
            mi = InvQuadraticMap()
            @test 0.1 ∈ m
            @test -0.1 ∈ m
            @test 2 ∉ m
            @test 0.1 ∈ mi
            @test -0.1 ∉ mi

            @test m[findfirst(isequal(0.1), m)] ≈ 0.1
            @test m[findlast(isequal(0.1), m)] ≈ 0.1
            @test m[findall(isequal(0.1), m)] ≈ [0.1]

            @test m[Inclusion(0..1)] ≡ m
            @test_throws BoundsError m[Inclusion(-1..1)]
            T = Chebyshev(5)
            M = T[m,:]
            @test MemoryLayout(M) isa MappedBasisLayout
            @test MemoryLayout(M[:,1:3]) isa MappedBasisLayout
            @test M[0.1,:] ≈ T[2*0.1^2-1,:]
            x = axes(M,1)
            @test x == Inclusion(0..1)
            @test M \ exp.(x) ≈ T \ exp.(sqrt.((axes(T,1) .+ 1)/2))
        end
    end

    @testset "Broadcasted" begin
        T = Chebyshev(5)
        F = factorize(T)
        x = axes(T,1)
        a = 1 .+ x .+ x.^2
        # The following are wrong, just testing dispatch
        @test T \ (a .* T) == I
        @test T \ (a .* (T * (T \ a))) isa Vector
        f = exp.(x) .* a # another broadcast layout
        @test T \ f == F \ f

        ã = T * (T \ a)
        @test T \ (ã .* ã) ≈ [1.5,1,0.5,0,0]
    end
end


export HurwitzMatrix,
    HurwitzDeterminants,
    StabilityMatrix

struct HurwitzDeterminants{T}
    M::MatElem{T}
end

Base.convert(::Type{AbstractMatrix},M::HurwitzDeterminants) = Array(M.M)
Base.show(io::IO, M::HurwitzDeterminants) = print(io, M.M)
Base.show(io::IO, ::MIME"text/plain", z::HurwitzDeterminants{T}) where {T} =
    print(io, "HurwitzDeterminants{$T} iterator:\n", z)

# function Base.iterate(M::HurwitzDeterminants)
#     state = 1
#     return Base.iterate(M, state)
# end

function Base.iterate(M::HurwitzDeterminants, state=1)
    if state > size(M.M, 1)
        return nothing
    else
        return Nemo.det(M.M[1:state,1:state]), state + 1
    end
end

# https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-iteration-1
Base.length(M::HurwitzDeterminants) = size(M.M, 1)
Base.eltype(::Type{HurwitzDeterminants{T}}) where {T} = T

# Input: polynomial p in one variable
function HurwitzMatrix(R::AbstractAlgebra.Ring, p)
    n = length(p) - 1
    M = zero_matrix(R, n, n)
    for i in 1:n
        for j in 1:n
            if (n - 2 * i + j) in 0:n
                M[i,j] = coeff(p, n - 2 * i + j)
            end
        end
    end
    return M
end

function evenodd_coeffs(p)
    l = length(p); ## degree +1
    ## For the zero polynomial, 'odd' would be a 0-element array
    if l < 1
        even = [0]
        odd  = [0]
    else
        even = [coeff(p, i) for i in 0:2:l]
        odd  = [coeff(p, i) for i in 1:2:l]
    end
    return even, odd
end

function evenodd_polys(p)
    even, odd = evenodd_coeffs(p)
    return parent(p)(evalxsqrt(even)), parent(p)(evalxsqrt(odd))
end

function evalxsqrt(poly)
    n = size(poly, 1)
    R = parent(poly[1])
    out = fill(R(0), 2 * n - 1)
    out[1] = poly[1]
    for i in 2:n
        out[2 * i - 1] = poly[i]
    end
    return out
end

function HurwitzMatrix1(R::AbstractAlgebra.Ring, p)
    sylvester_matrix(evenodd_polys(p)...)
end

# ## k's are parameters
# R, (k1, k2) = PolynomialRing(Nemo.QQ, ["k1", "k2"]);
# RQ = FractionField(R)

# ## The ring for the polynomial system
# S, (x, y, z, w, t) = PolynomialRing(RQ, ["x", "y", "z", "w", "t"])
# SQ = FractionField(S)

# ## The ring for the characteristic polynomial.
# T, l = PolynomialRing(SQ, "A")

# P = T([x * y * k1, k1 * k2^2 + 2, k1, //(k1, k2 + k1 * k2) * x])
# Q = T([x, y, z, w, t])
# p, q = evenodd_polys(Q)

# sylvester_matrix(evenodd_polys(p)...)
# -x * resultant_subresultant(p, q) == det(HurwitzMatrix(S, Q))

# # add_row(a::MatrixElem, s::RingElement, i::Int, j::Int, cols = 1:ncols(a))

# R, (k1, k2, k3) = PolynomialRing(Nemo.QQ, ["k1", "k2", "k3"]);
# RQ = FractionField(R)

# ## The ring for the polynomial system
# S, (x, y, z) = PolynomialRing(RQ, ["x", "y", "z"])

# ## The ring for the characteristic polynomial.
# T, l = PolynomialRing(S, "A")

# chp = T([k1 * k2 * y + k2 * k3 * z + k1 * k3, k2 * y + k2 * z + k1 + k3, S(1)])

# p, q = evenodd_polys(chp)
# resultant_subresultant(p,q)
# M = sylvester_matrix(p, q)

# det(M) == resultant_subresultant(p, q)

# det(M) == det(HurwitzMatrix(S, chp))

function ISRO!(M, j)
    c = ncols(M)
    p = M[j,j]
    for i in (j + 1):c
        if !(iszero(M[i,j]))
            add_row!(M, M[i,j] // p, j, i, c)
        end
    end
end

function ISRO!(M)
    return j -> ISRO!(M, j)
end

function HurwitzDeterminants(R::AbstractAlgebra.Ring, p)
    return HurwitzDeterminants(HurwitzMatrix(R, p))
end

# mindeg = size(N,1)-rank(matrix(R,N))
function StabilityMatrix(R::AbstractAlgebra.Ring, J, mindeg::Integer=size(J, 1) - rank(J))
    P, x = PolynomialRing(R, "x")
    p = divexact(charpoly(P, J), x^mindeg)
    H = []
    iter = HurwitzDeterminants(R, p)
    print("==============================================\n")
    print("==============================================\n")
    print("$(length(iter)) expressions will be studied\n")
    for (i, D) in enumerate(iter)
        print("\n-- $i th determinant --\n")
        signs = unique(sign.(coeffs(D)))
        print("$(signs)\n")
        if -1 in signs
            H = [H; D]
        end
    end
    print("==============================================\n")
    print("==============================================\n")
    return H
end

function StabilityMatrix(io, R::AbstractAlgebra.Ring, J, mindeg::Integer=size(J, 1) - rank(J))
    P, x = PolynomialRing(R, "x")
    p = divexact(charpoly(P, J), x^mindeg)
    H = []
    iter = HurwitzDeterminants(R, p)
    write(io, "==============================================\n")
    write(io, "==============================================\n")
    write(io, "$(length(iter)) expressions will be studied\n")
    for (i, D) in enumerate(iter)
        write(io, "\n-- $i th determinant --\n")
        # signs = unique(sign.(coeffs(D)))
        signs = unique(sign.(reshape(hcat, coeffs.(coeffs(D)))))
        write(io, "$(signs)\n")
        if -1 in signs
            push!(H, D)
        end
    end
    write(io, "==============================================\n")
    write(io, "==============================================\n")
    return H
end


module Gap

using JuMP
using Cbc

include("Computation.jl")
using .Computation

include("Format.jl")
using .Format

include("FC.jl")
using .FC

include("Half.jl")
using .Half

export gap, vgap

# Determines upper bound alpha with the Gap Method
# Author: Jason Liu
function gap(m::Int64, s::Int64; output::Int64 = 2)
    output > 0 && printHeader(center("GAP METHOD"))

    if m <= 0 || s <= 0
        if output > 0
            printf("Gap Method does not apply", line = true)
            printEnd()
        end
        return 1 // 1
    elseif m < s
        alpha = m // s * gap(s, m, output = 0)
        if output > 1
            printfT(
                "Duality Theorem",
                "muffins($m,$s) = $m/$s · muffins($s,$m)",
                "Bounding muffins($s,$m) instead",
            )
            gap(s, m, output = output)
        elseif output > 0
            printfT("Gap Method", "Upper bound α of muffins($m,$s) is $(formatFrac(alpha))")
            printEnd()
        end
        return alpha
    elseif m % s == 0
        if output > 0
            printfT("Gap Method", "Since $m % $s = 0, muffins($m,$s) = 1")
            printEnd()
        end
        return 1 // 1
    end

    # Compute alpha candidates
    alphas = []
    upper = min(fc(m, s, output = 0), half(m, s, output = 0))
    for b = 1:m
        for n = Int64(ceil(b * s / 3)):Int64(floor(b * s * upper))
            append!(alphas, n // (b * s))
        end
    end
    append!(alphas, 1 // 3)
    alphas = sort(unique(alphas))

    # Test alpha candidates with binary search
    potAlphas = []
    (left, right) = (1, length(alphas))
    while left <= right
        mid = Int64(floor((left + right) / 2))
        append!(potAlphas, alphas[mid])
        if vgap(m, s, alphas[mid], output = 0)
            right = mid - 1
        else
            left = mid + 1
        end
    end

    for alpha in reverse(potAlphas)
        if vgap(m, s, alpha, output = 0)
            vgap(m, s, alpha, output = output)
            return alpha
        end
    end

    if output > 0
        printf("All α candidates were inconclusive, Gap Method inconclusive", line = true)
        printEnd()
    end
    1 // 1
end

# Helper function for gap -- verifies whether gap(m, s, alpha) is conclusive
function vgap(m::Int64, s::Int64, alpha::Rational{Int64}; output::Int64 = 2)
    if !(m > s > 0) || m % s == 0
        output > 0 && printf("VGap does not apply", line = true)
        false
    elseif alpha < 1 / 3
        output > 0 && printfT("No piece size < 1/3", "For m ≥ s, α must be ≥ 1/3")
        false
    elseif alpha > 1
        output > 0 && printfT("No piece size > 1", "α must be ≤ 1")
        false
    else
        (V, W, sV, sW) = sv(m, s)
        numV = (V)sV
        numW = (W)sW

        # Initialize Gap Method proof
        if output > 1
            # Define and format variables for proof
            alphaF = formatFrac(alpha)
            size = formatFrac(m // s)
            a = formatFrac(m // (V + 1)s)
            b = max(1 - m // (W - 1)s, 0)
            b1 = formatFrac(1 - b)
            b = formatFrac(b)

            # Establish assumptions and premises
            printHeader("CLAIM")
            printfT(
                "Claim",
                "For all ($m, $s) procedures where $s students each receive $size muffins, the smallest piece size is bounded above by α = $alphaF",
                "Proven by contradicting the assumption",
            )

            printHeader("ASSUMPTION")
            printfT("Assumption", "Assume muffins($m,$s) > α ≥ 1/3")
            printfT(
                "Property of Buddies",
                "Since muffins($m,$s) > 1/3, each muffin must be cut into 2 shs, totaling $(2m) shs",
                "",
                "Since each muffin is cut into 2 shs that are buddies, each sh of size x must imply the existence of a sh of size 1-x",
            )

            # Begin casework
            printHeader("CASEWORK")

            printfT(
                "Case 1",
                "Alice has ≥ $(V+1) shs",
                "One of them is ≤ ($size × 1/$(V+1)) = $a",
                "",
                "Contradicts assumption if α ≥ $a",
            )

            if W - 1 <= 1 && m / s > 1
                printfT(
                    "Case 2",
                    "Bob cannot have ≤ $(W-1) shs since $size > 1, so this case is impossible",
                )
            else
                printfT(
                    "Case 2",
                    "Bob has ≤ $(W-1) shs",
                    "One of them is ≥ ($size × 1/$(W-1)) = $b1",
                    "Its buddy is ≤ (1 - $b1) = $b",
                    "",
                    "Contradicts assumption if α ≥ $b",
                )
            end
        end

        # Check if V-Conjecture applies
        if m // s * 1 // (V + 1) > alpha || 1 - m // s * 1 // (W - 1) > alpha
            output > 0 && printf("V-Conjecture does not apply, VGap failed", line = true)
            return false
        end

        ((_, x), (y, _)) = findend(m, s, alpha, V)
        xOriginal = m // s - alpha * (V - 1)
        yOriginal = m // s - (1 - alpha) * (V - 2)

        # Check if FindEnd works
        if x != xOriginal && y != yOriginal
            output > 0 && printf("FindEnd inconclusive, VGap failed", line = true)
            return false
        end

        # Define and format variables for proof
        cd = lcm(s, denominator(x), denominator(y), 2)
        alphaF = formatFrac(alpha, cd)
        alpha1 = formatFrac(1 - alpha, cd)
        size = formatFrac(m // s)
        a = formatFrac(m // s - x, cd)
        b = formatFrac(m // s - y, cd)
        xF = formatFrac(x, cd)
        x1 = formatFrac(1 - x, cd)
        yF = formatFrac(y, cd)
        y1 = formatFrac(1 - y, cd)
        diff = abs(numV - numW)

        # Continue proof
        if output > 1
            # Continue casework
            printfT(
                "Note",
                "The remaining cases deal with everyone having either $W or $V shs, so:",
                "",
                "$(W)·s_$W + $(V)·s_$V = $(2m)  (total shs)",
                "s_$W + s_$V = $s  (total students)",
                "where s_N = # of students with N shs",
                "",
                "The solution to the system is:",
                "s_$W = $sW, s_$V = $sV",
                "So there are $numW $W-shs and $numV $V-shs",
            )

            if x == xOriginal
                printfT(
                    "Case 3",
                    "Alice has a $V-sh ≥ $xF",
                    "Her other $(V-1) $V-shs sum to ≤ ($size - $xF) = $a",
                    "One of them is ≤ ($a × 1/$(V-1)) = $alphaF",
                    "",
                    "Contradicts assumption if α ≥ $alphaF",
                )
            else
                printfT("Case 3", "FindEnd did not produce a conclusive bound for $V-shs")
            end

            if y == yOriginal
                printfT(
                    "Case 4",
                    "Bob has a $W-sh ≤ $yF",
                    "His other $(W-1) $W-shs sum to ≥ ($size - $yF) = $b",
                    "One of them is ≥ ($b × 1/$(W-1)) = $alpha1",
                    "Its buddy is ≤ (1 - $alpha1) = $alphaF",
                    "",
                    "Contradicts assumption if α ≥ $alphaF",
                )
            else
                printfT("Case 4", "FindEnd did not produce a conclusive bound for $W-shs")
            end

            printHeader("CASE 5: INTERVAL ANALYSIS")
            println()
            printf("The following intervals capture the negation of the previous cases:")
            if alpha < x <= y < 1 - alpha
                println(
                    "\n",
                    interval(
                        ["(", alphaF],
                        [")[", xF],
                        ["](", yF],
                        [")", alpha1],
                        labels = ["$numV $V-shs", "0", "$numW $W-shs"],
                    ),
                )
                printLine()
            else
                if xF != alphaF
                    println(
                        "\n",
                        interval(
                            ["(", alphaF],
                            [")[", xF],
                            [")", alpha1],
                            labels = ["$numV $V-shs", "0 $V-shs"],
                        ),
                    )
                else
                    println("\n", "No negation interval for $V-shs")
                end

                if yF != alpha1
                    println(
                        "\n",
                        interval(
                            ["(", alphaF],
                            ["](", yF],
                            [")", alpha1],
                            labels = ["0 $W-shs", "$numW $W-shs"],
                        ),
                    )
                else
                    println("\n", "No negation interval for $W-shs")
                end

                printfT(
                    "Case 5",
                    "The Gap Method is inconclusive on these intervals, VGap failed",
                )
            end
        end

        # Fail if interval inconclusive
        if x > y || x == alpha || y == 1 - alpha
            output == 1 &&
                printf("Could not generate disjoint intervals, VGap failed", line = true)
            return false
        end

        # Define and format variables for proof
        (vMin, vMax, sMin, sMax) = (V, W, sV, sW)
        (i, j, k, l) = (xF, yF, y1, x1)
        (rngMin, rngMax, rngS, rngL) = ((alphaF, xF), (x1, alpha1), (yF, y1), (x1, alpha1))
        if numV > numW
            (vMin, vMax, sMin, sMax) = (W, V, sW, sV)
            (i, j, k, l) = (y1, x1, xF, yF)
            (rngMin, rngMax, rngS, rngL) =
                ((yF, alpha1), (alphaF, y1), (alphaF, y1), (x1, xF))
        end
        numMin = (vMin)sMin

        # Conclude interval case with case by case analysis
        if output > 1
            printfT(
                "Property of Buddies",
                "Because 0 shs lie in [$xF,$yF], 0 shs must also lie in [$y1,$x1]",
                "",
                "Similary, the $numMin $vMin-shs in ($(rngMin[1]),$(rngMin[2])) implies $numMin $vMax-shs in ($(rngMax[1]),$(rngMax[2]))",
            )
            println()
            printf("The following intervals capture the previous statements:")
            println(
                "\n",
                interval(
                    ["(", alphaF],
                    [")[", i],
                    ["](", j],
                    [")[", k],
                    ["](", l],
                    [")", alpha1],
                    labels = [
                        "$numMin $V-shs",
                        "0",
                        "$diff $vMax-shs",
                        "0",
                        "$numMin $W-shs",
                    ],
                ),
            )
            printLine()

            printfT(
                "Note",
                "Let the $vMax-shs that lie in ($(rngS[1]),$(rngS[2])) be small $vMax-shs",
                "Let the $vMax-shs that lie in ($(rngL[1]),$(rngL[2])) be large $vMax-shs",
            )
        end

        # Determine possible small-sh, large-sh combinations
        posCombo = []
        for k = 0:vMax
            upperB = (vMax - k) * toFrac(rngL[2]) + (k) * toFrac(rngS[2])
            lowerB = (vMax - k) * toFrac(rngL[1]) + (k) * toFrac(rngS[1])
            if upperB > m // s > lowerB
                append!(posCombo, k)
            end
        end

        # Diverge from Interval Method
        if length(posCombo) > 0
            mid = formatFrac(1 // 2, cd)
            if output > 1
                combos =
                    ["$k small $vMax-shs and $(vMax-k) large $vMax-shs" for k in posCombo]
                printfT(
                    "Case 5",
                    "All combinations of small and large $vMax-shs are impossible except for:",
                    "",
                    combos...,
                    "",
                    "because the range of muffin amounts for the others does not include $size",
                )

                printHeader("CASE 5: MIDPOINT ANALYSIS")
                printfT(
                    "Property of Buddies",
                    "The # of $vMax-shs in ($j, $mid) is equal to the # of $vMax-shs in ($mid, $k)",
                )
            end

            bounds = [[toFrac(j), 1 // 2], [1 // 2, toFrac(k)], [toFrac(l), toFrac(alpha1)]]
            smallBounds = [toFrac(j), toFrac(k)]
            solid = [toFrac(l), toFrac(alpha1)]
            intervals = [["(", j], ["|", mid], [")[", k], ["](", l], [")", alpha1]]
            labels = ["z $W-shs", "z $W-shs", "0", "$numV $W-shs"]
            intervalDefs = [
                "A = ($j,$mid)",
                "B = ($mid,$k) (|A| = |B| = z)",
                "C = ($l,$alpha1) (|C| = $numV)",
            ]
            if numW < numV
                bounds = [[alpha, toFrac(i)], [toFrac(j), 1 // 2], [1 // 2, toFrac(k)]]
                smallBounds = [alpha, toFrac(i)]
                solid = [alpha, toFrac(i)]
                intervals = [["(", alphaF], [")[", i], ["](", j], ["|", mid], [")", k]]
                labels = ["$numW $V-shs", "0", "z $V-shs", "z $V-shs"]
                intervalDefs = [
                    "A = ($alphaF, $i) (|A| = $numW)",
                    "B = ($j, $mid)",
                    "C = ($mid, $k) (|B| = |C| = z)",
                ]
            end

            if output > 1
                println()
                printf("The following intervals capture what we know about the $V-shs:")
                println("\n", interval(intervals..., labels = labels))
                printLine()
                printfT("Note", "We define the following intervals:", intervalDefs...)
            end

            # Determine possible interval combinations
            posInt = intCombo(m, s, vMax, bounds, smallBounds, posCombo)

            if length(posInt) > 0
                if output > 1
                    printfT(
                        "Note",
                        "Let the notation (a, b, c, ...) denote a combination of $vMax-shs where there are:",
                        "a shs in interval A",
                        "b shs in interval B",
                        "c shs in interval C",
                        "...",
                        "",
                        "A (a, b, c, ...)-student denotes a student with that combination of $vMax-shs",
                    )
                    printfT(
                        "Midpoint Analysis I",
                        "Given the possible small and large $vMax-sh combinations,",
                        "all combinations of $vMax-shs from intervals A, B, and C are impossible except for:",
                        "",
                        posInt...,
                        "",
                        "because the range of muffin amounts for the others does not include $size",
                    )
                end

                # Check if posInt is too large
                if length(posInt) > 5
                    output > 1 && printf(
                        "The # of possible interval combinations is too large, skipping Midpoint Analysis",
                        line = true,
                    )
                    solution = [0, 0]
                else
                    # Set up and solve system of equations
                    iter = 1:length(posInt)
                    (a, b, c) =
                        (getindex.(posInt, 1), getindex.(posInt, 2), getindex.(posInt, 3))
                    (A, B, C) = ("A", "B", "C")
                    if numW < numV
                        (a, b, c) = (b, c, a)
                        (A, B, C) = (B, C, A)
                    end

                    equations = [
                        join(["$(a[i])·x_$i" for i in iter], " + ") *
                        " = " *
                        join(["$(b[i])·x_$i" for i in iter], " + "),
                        join(["$(c[i])·x_$i" for i in iter], " + ") * " = |$C| = $numMin",
                        join(["x_$i" for i in iter], " + ") * " = s_$vMax = $sMax",
                        "",
                        ["where x_$i is the # of $(posInt[i])-students" for i in iter]...,
                    ]

                    # Test all possible non-negative integer solutions with system
                    model = Model(Cbc.Optimizer)
                    set_optimizer_attribute(model, "logLevel", 0)
                    set_optimizer_attribute(model, "cuts", "off")

                    @variable(model, 0 <= x[i = 1:length(iter)], Int)

                    @constraint(model, sum(x .* a) == sum(x .* b))
                    @constraint(model, sum(x .* c) == numMin)
                    @constraint(model, sum(x) == sMax)

                    @objective(model, Max, x[1])
                    optimize!(model)
                    if termination_status(model) == MOI.OPTIMAL
                        x = convert.(Int64, round.(value.(x)))
                        solution = [
                            "One solution is " *
                            join(["x_$i = $(x[i])" for i in iter], ", "),
                            "The system has a non-negative integer solution, so Case 5 is still possible",
                        ]
                    else
                        solution =
                            ["The system has no non-negative integer solutions, so Case 5 is impossible"]
                    end
                end

                if length(solution) > 1
                    if output > 1
                        printf("Midpoint Analysis inconclusive, continuing to Gap Analysis", line=true)
                        printHeader("CASE 5: GAP ANALYSIS")
                    end

                    upChr = join('A':'Z')
                    (loChr, loOrd) = (join('z':-1:'a'), 1)
                    while true
                        # Find gaps
                        gaps = []
                        equations = []
                        for i = 1:length(bounds)
                            ints = []
                            for tup in posInt
                                if tup[i] > 0
                                    upper =
                                        m // s - sum(getindex.(bounds, 1) .* tup) +
                                        bounds[i][1]
                                    lower =
                                        m // s - sum(getindex.(bounds, 2) .* tup) +
                                        bounds[i][2]
                                    disjoint = true
                                    merged = []
                                    for j = 1:length(ints)
                                        if upper <= ints[j][1] || ints[j][2] <= lower
                                            continue
                                        else
                                            disjoint = false
                                            ints[j][1] = min(lower, ints[j][1])
                                            ints[j][2] = max(upper, ints[j][2])
                                            append!(merged, j)
                                        end
                                    end
                                    if length(merged) > 1
                                        append!(
                                            ints,
                                            [[
                                                minimum(ints[j][1] for j in merged),
                                                maximum(ints[j][2] for j in merged),
                                            ]],
                                        )
                                        deleteat!(ints, merged)
                                    end
                                    if disjoint
                                        append!(ints, [[lower, upper]])
                                    end
                                    append!(
                                        equations,
                                        ["For $tup-student: all $vMax-shs in interval $(upChr[i]) must lie in ($(formatFrac(lower, cd)), $(formatFrac(upper, cd)))"],
                                    )
                                end
                            end
                            sort!(ints)
                            append!(equations, [""])
                            for j = 1:length(ints)-1
                                (a, b, c, d) = (
                                    ints[j][2],
                                    ints[j+1][1],
                                    1 - ints[j+1][1],
                                    1 - ints[j][2],
                                )
                                append!(gaps, [[a, b], [c, d]])
                                append!(
                                    equations,
                                    [
                                        "Gap found: [$(formatFrac(a, cd)), $(formatFrac(b, cd))]",
                                        "Its buddy: [$(formatFrac(c, cd)), $(formatFrac(d, cd))]",
                                    ],
                                )
                            end
                            if length(ints) < 2
                                append!(
                                    equations,
                                    ["No gaps found in interval $(upChr[i])"],
                                )
                            end
                            append!(equations, [""])
                        end

                        # Break if no gaps found
                        if length(gaps) == 0
                            break
                        end

                        # Sort and merge gaps
                        sort!(gaps)
                        mergedGaps = []
                        i = 1
                        while i < length(gaps)
                            if gaps[i][2] < gaps[i+1][1]
                                append!(mergedGaps, [gaps[i]])
                            else
                                append!(mergedGaps, [[gaps[i][1], gaps[i+1][2]]])
                                i += 1
                            end
                            i += 1
                        end
                        if i == length(gaps)
                            append!(mergedGaps, [gaps[i]])
                        end
                        gaps = mergedGaps

                        if output > 1
                            printfT("New Gaps", equations...)
                        end

                        # Insert gaps
                        for gap in gaps
                            (l, u) = (gap[1], gap[2])
                            i = 1
                            while i <= length(bounds)
                                (bl, bu) = (bounds[i][1], bounds[i][2])
                                if bl < l <= u < bu
                                    # Update bounds
                                    insert!(bounds, i + 1, [u, bu])
                                    bounds[i][2] = l
                                    # Update intervals
                                    for j = 2:length(intervals)
                                        if toFrac(intervals[j-1][2]) == bl &&
                                           toFrac(intervals[j][2]) == bu
                                            insert!(intervals, j, ["](", formatFrac(u, cd)])
                                            insert!(intervals, j, [")[", formatFrac(l, cd)])
                                            break
                                        end
                                    end
                                end
                                i += 1
                            end
                        end

                        # Fail if gap finding excessive
                        if length(bounds) > 26
                            output > 0 &&
                                printf("Excessive gap finding, VGap failed", line = true)
                            return false
                        end

                        # Update posInt
                        posInt = intCombo(m, s, vMax, bounds, smallBounds, posCombo)
                    end

                    if length(bounds) == 3
                        output > 0 && printf(
                            "No further gaps could be generated, VGap failed",
                            line = true,
                        )
                        return false
                    end

                    # Update labels and intervalDefs
                    labels = repeat(["0"], length(intervals) - 1)
                    intervalDefs = repeat([""], length(bounds))
                    equivs = []
                    solidInt = 0
                    visited = []
                    iter = 1:length(intervals)-1
                    for i = 1:length(bounds)
                        if !(i in visited)
                            int = bounds[i]
                            chr = upChr[i]
                            for k in iter
                                if toFrac(intervals[k][2]) == int[1] &&
                                   toFrac(intervals[k+1][2]) == int[2]
                                    (lower, upper) = (intervals[k][2], intervals[k+1][2])
                                    if int == solid
                                        labels[k] = "$numMin $vMax-shs"
                                        intervalDefs[i] = "$chr = ($lower, $upper) (|$chr| = $numMin)"
                                        solidInt = i
                                    else
                                        chrL = loChr[loOrd]
                                        labels[k] = "$chrL $vMax-shs"
                                        intervalDefs[i] = "$chr = ($lower, $upper)"
                                        equiv = [i]
                                        for j = i:length(bounds)
                                            intOth = bounds[j]
                                            if intOth == [1 - int[2], 1 - int[1]] || (
                                                vMax == 2 &&
                                                intOth ==
                                                [m // s - int[2], m // s - int[1]]
                                            )
                                                for l in iter
                                                    if toFrac(intervals[l][2]) ==
                                                       intOth[1] &&
                                                       toFrac(intervals[l+1][2]) ==
                                                       intOth[2]
                                                        (lowerOth, upperOth) = (
                                                            intervals[l][2],
                                                            intervals[l+1][2],
                                                        )
                                                        labels[l] = "$chrL $vMax-shs"
                                                        intervalDefs[j] = "$(upChr[j]) = ($lowerOth, $upperOth)"
                                                        append!(equivs, [(i, j)])
                                                        append!(equiv, j)
                                                        append!(visited, j)
                                                    end
                                                end
                                            end
                                        end
                                        intervalDefs[equiv[length(
                                            equiv,
                                        )]] *= " ($(join(["|$(upChr[j])|" for j in equiv], " = ")) = $chrL)"
                                        loOrd += 1
                                    end
                                    break
                                end
                            end
                        end
                    end

                    if output > 1
                        println()
                        printf("After finding gaps and buddying, this is the interval of the $vMax-shs:")
                        println("\n", interval(intervals..., labels = labels))
                        printLine()
                        printfT(
                            "Note",
                            "We define the following intervals:",
                            intervalDefs...,
                        )
                    end

                    if length(posInt) > 0
                        if output > 1
                            printfT(
                                "Gap Analysis I",
                                "Given the possible small and large $vMax-sh combinations,",
                                "all combinations of $vMax-shs from the defined intervals are impossible except for:",
                                "",
                                posInt...,
                                "",
                                "because the range of muffin amounts for the others does not include $size",
                            )
                        end

                        # Set up and solve system of equations
                        iter = 1:length(posInt)
                        consts = [getindex.(posInt, i) for i = 1:length(bounds)]
                        l = solidInt

                        equations = []
                        for (j, k) in equivs
                            append!(
                                equations,
                                [
                                    join(["$(consts[j][i])·x_$i" for i in iter], " + ") *
                                    " = " *
                                    join(["$(consts[k][i])·x_$i" for i in iter], " + "),
                                ],
                            )
                        end
                        if l > 0
                            append!(
                                equations,
                                [
                                    join(["$(consts[l][i])·x_$i" for i in iter], " + ") * " = |$(upChr[l])| = $numMin",
                                ],
                            )
                        end
                        append!(
                            equations,
                            [
                                join(["x_$i" for i in iter], " + ") * " = s_$vMax = $sMax",
                                "",
                            ],
                        )

                        append!(
                            equations,
                            [
                                "where x_$i is the # of ($(join(posInt[i], ",")))-students"
                                for i in iter
                            ],
                        )

                        # Test all possible non-negative integer solutions with system
                        model = Model(Cbc.Optimizer)
                        set_optimizer_attribute(model, "logLevel", 0)
                        set_optimizer_attribute(model, "cuts", "off")

                        @variable(model, 0 <= x[i = 1:length(iter)], Int)

                        for (j, k) in equivs
                            @constraint(model, sum(x .* consts[j]) == sum(x .* consts[k]))
                        end
                        if l > 0
                            @constraint(model, sum(x .* consts[l]) == numMin)
                        end
                        @constraint(model, sum(x) == sMax)

                        @objective(model, Max, x[1])
                        optimize!(model)
                        if termination_status(model) == MOI.OPTIMAL
                            x = convert.(Int64, round.(value.(x)))
                            solution = [
                                "One solution is " *
                                join(["x_$i = $(x[i])" for i in iter], ", "),
                                "The system has a non-negative integer solution, so Case 5 is still possible",
                            ]
                        else
                            solution =
                                ["The system has no non-negative integer solutions, so Case 5 is impossible"]
                        end

                        if output > 1
                            printfT(
                                "Gap Analysis II",
                                "Given the new interval equalities, we can set up the following system of equations:",
                                "",
                                equations...,
                                "",
                                solution...,
                            )
                        end

                        if length(solution) > 1
                            output == 1 && printf(
                                "Could not generate conclusive system of equations, VGap failed",
                                line = true,
                            )
                            return false
                        end
                    else
                        if output > 1
                            printfT(
                                "Gap Analysis",
                                "Given the possible small and large $vMax-sh combinations,",
                                "all combinations of $vMax-shs from the intervals are impossible since their total amount can not be $size",
                            )
                        end
                    end

                else
                    if output > 1
                        printfT(
                            "Midpoint Analysis II",
                            "Since |$A| = |$B|, we can set up the following system of equations:",
                            "",
                            equations...,
                            "",
                            solution...,
                        )
                    end
                end
            else
                if output > 1
                    printfT(
                        "Midpoint Analysis",
                        "Given the possible small and large $vMax-sh combinations,",
                        "all combinations of $vMax-shs from intervals A, B, and C are impossible since their total amount can not be $size",
                    )
                end
            end

        else
            if output > 1
                printfT(
                    "Case 5",
                    "All combinations of small and large $vMax-shs are impossible since their total amount can not be $size",
                )
            end
        end

        # Conclude proof
        alpha = formatFrac(alpha)
        if output > 1
            # Conclude with alpha's value
            printHeader("CONCLUSION")
            printfT(
                "Conclusion",
                "All possible cases contradict the assumption iff. α ≥ $alpha",
                "muffins($m,$s) ≤ α, ∀ α ≥ $alpha",
                "",
                "muffins($m,$s) ≤ $alpha",
            )
            printEnd()
        elseif output > 0
            printfT("Gap Method", "Upper bound α of muffins($m,$s) is $alpha")
            printEnd()
        end

        true
    end
end

end


FitFunction(s::Symbol) = FitFunction(Float64, Val(s))
FitFunction(T::Type{<:AbstractFloat}, s::Symbol) = FitFunction(T, Val(s))

function FitFunction(T::Type{<:AbstractFloat}, s::Val{:Gauss})
    ff = FitFunction{T}( Gauss, 1, 3 )
    set_parameter_names!(ff, ["A", "σ", "μ"])
    set_initial_parameters!(ff, [1, 1, 0])
    set_fitranges!(ff, ([-1, 1],))
    return ff
end
function FitFunction(T::Type{<:AbstractFloat}, s::Val{:PDF_Gauss})
    ff = FitFunction{T}( PDF_Gauss, 1, 2 )
    set_parameter_names!(ff, ["σ", "μ"])
    set_initial_parameters!(ff, [1, 0])
    set_fitranges!(ff, ([-1, 1],))
    return ff
end

function FitFunction(T::Type{<:AbstractFloat}, s::Val{:GaussPlusConstantBackground})
    ff = FitFunction{T}( Gauss_plus_const_background, 1, 4 )
    set_parameter_names!(ff, ["A", "σ", "μ", "offset"])
    set_initial_parameters!(ff, [1, 1, 0, 1])
    set_fitranges!(ff, ([-1, 1],))
    return ff
end
function FitFunction(T::Type{<:AbstractFloat}, s::Val{:GaussPlusLinearBackground})
    ff = FitFunction{T}( Gauss_plus_linear_background, 1, 5 )
    set_parameter_names!(ff, ["A", "σ", "μ", "offset", "lin. slope"])
    set_initial_parameters!(ff, [1, 1, 0, 1, 0])
    set_fitranges!(ff, ([-1, 1],))
    return ff
end

function Gauss(x::T, A::T, σ::T, μ::T)::T where {T <: AbstractFloat}
    r::T = A * exp( -(x - μ)^2 / (2 * σ^2)) / sqrt(2π * σ^2)
    if σ <= 0 || A < 0 return T(Inf) end
    if r < 0 r = 0 end
    return r
end
function Gauss(x::Vector{T}, A::T, σ::T, μ::T)::Vector{T} where {T <: AbstractFloat}
    return Gauss.( x, A, σ, μ)
end
Gauss(x::Union{T, Vector{T}}, p::Vector{T}) where {T <: AbstractFloat} = Gauss(x, p... )

function PDF_Gauss(x::T, σ::T, μ::T)::T where {T <: AbstractFloat}
    r::T = exp( -(x - μ)^2 / (2 * σ^2)) / sqrt(2π * σ^2)
    if σ <= 0 return T(Inf) end
    if r < 0 r = 0 end
    return r
end
function PDF_Gauss(x::Vector{T}, σ::T, μ::T)::Vector{T} where {T <: AbstractFloat}
    return PDF_Gauss.( x, σ, μ)
end
PDF_Gauss(x::Union{T, Vector{T}}, p::Vector{T}) where {T <: AbstractFloat} = PDF_Gauss(x, p... )

function Gauss_plus_linear_background(x::T, A::T, σ::T, μ::T, offset::T, lin_slope::T)::T where {T <: AbstractFloat}
    r::T = A * exp( -(x - μ)^2 / (2 * σ^2)) / sqrt(2π * σ^2) + lin_slope * (x - μ) + offset
    if σ <= 0 || A < 0 || offset < 0 return T(Inf) end
    if r < 0 r = 0 end
    return r
end
function Gauss_plus_linear_background(x::Vector{T}, A::T, σ::T, μ::T, offset::T, lin_slope::T)::Vector{T} where {T <: AbstractFloat}
    return Gauss_plus_linear_background.( x, A, σ, μ, offset, lin_slope)
end
Gauss_plus_linear_background(x::Union{T, Vector{T}}, p::Vector{T}) where {T <: AbstractFloat} = Gauss_plus_linear_background(x, p... )

function Gauss_plus_const_background(x::T, A::T, σ::T, μ::T, offset::T)::T where {T <: AbstractFloat}
    r::T = A * exp( -(x - μ)^2 / (2 * σ^2)) / sqrt(2π * σ^2) + offset
    if σ <= 0 || A < 0 || offset < 0 return T(Inf) end
    if r < 0 r = 0 end
    return r
end
function Gauss_plus_const_background(x::Vector{T}, A::T, σ::T, μ::T, offset::T)::Vector{T} where {T <: AbstractFloat}
    return Gauss_plus_const_background.( x, A, σ, μ, offset)
    if σ <= 0 || A < 0 return T(Inf) end
end
Gauss_plus_const_background(x::Union{T, Vector{T}}, p::Vector{T}) where {T <: AbstractFloat} = Gauss_plus_const_background(x, p... )


using ApproxFun, SingularIntegralEquations , Plots
gr()

k = 50.
ω = 2π
d = (1,-1)
d = d[1]/hypot(d[1],d[2]),d[2]/hypot(d[1],d[2])
ui = (x,y) -> exp(im*k*(d[1]*x+d[2]*y))

# The Helmholtz Green's function, split into singular and nonsingular pieces.
g1 = (x,y) -> -besselj0(k*abs(y-x))/2
g2 = (x,y) -> x == y ? -(log(k/2)+γ)/2/π + im/4 : im/4*hankelh1(0,k*abs(y-x)) - g1(x,y).*logabs(y-x)/π

Γ = Curve(Fun(x->x+im*x^3))

#Γ = Segment(0.,1.0+im)

sp = Space(Γ)
cwsp = CauchyWeight(sp⊗sp,0)
uiΓ,⨍ = Fun(t->ui(real(t),imag(t)),sp),DefiniteLineIntegral(Γ)

@time G = GreensFun(g1,cwsp;method=:Cholesky) + GreensFun(g2,sp⊗sp;method=:Cholesky)

@time ∂u∂n = ⨍[G]\uiΓ

N=100
x = linspace(-3,3,N);y = linspace(-2,2,N)

us = (x,y) -> -logkernel(g1,∂u∂n,complex(x,y))-linesum(g2,∂u∂n,complex(x,y))

@time uvals = ui.(x,y') .+ us.(x,y')

contourf(x,y,real(uvals'))
    plot!(Γ)

println("The length of ∂u∂n is: ",ncoefficients(∂u∂n))
us = (x,y) -> -logkernel(g1,∂u∂n,complex(x,y))-linesum(g2,∂u∂n,complex(x,y))

t=0.1;us(t,t^3)-ui(t,t^3)
t=0.2;us(t,t^3)-ui(t,t^3)

t=0.1;us(t,t)+ui(t,t)
t=0.2;us(t,t)+ui(t,t)

Γ = Curve(Fun(x->x+im*x^3)) ∪ Curve(Fun(x->x-1+im*(x^4+4)))

plot(Γ)

sp = Space(Γ)
cwsp = CauchyWeight(sp⊗sp,0)

uiΓ,⨍ = Fun(t->ui(real(t),imag(t)),sp),DefiniteLineIntegral(Γ)

@time G = GreensFun(g1,cwsp;method=:Cholesky) + GreensFun(g2,sp⊗sp;method=:Cholesky)

@time ∂u∂n = ⨍[G]\uiΓ
println("The length of ∂u∂n is: ",ncoefficients(∂u∂n))
us = (x,y) -> -logkernel(g1,∂u∂n,complex(x,y))-linesum(g2,∂u∂n,complex(x,y))

N=300
x = linspace(-4,3,N);y = linspace(-2,6,N)

us = (x,y) -> -logkernel(g1,∂u∂n,complex(x,y))-linesum(g2,∂u∂n,complex(x,y))

@time uvals = ui.(x,y') .+ us.(x,y')

writecsv("/Users/solver/Desktop/uvals.csv",uvals)

contourf(x,y,real(uvals');legend=false,size=(1000,600))
    plot!(Γ;color=:black)
png("scat.png")

3

@time us(0.1,0.2)

writecsv("dudncoeffs.csv",∂u∂n.coefficients)

# on 0.5
dsp = domainspace(DefiniteLineIntegral(Γ))
cfs = eval.(parse.(vec(readcsv("/Users/solver/Desktop/dudncoeffs.csv"))))

∂u∂n = Fun(dsp,chop!(cfs,1E-5))

us = (x,y) -> -logkernel(g1,∂u∂n,complex(x,y))-linesum(g2,∂u∂n,complex(x,y))

us(0.1,0.2)

N=300
x = linspace(-4,3,N);y = linspace(-2,6,N)
xx,yy = x.+0.*y',0.*x.+y'

@time uvals = ui.(xx,yy) .+ us.(xx,yy)

using Plots
pyplot()

plot(Γ;color="black")

umax = maximum(abs,uvals)
contourf(x,y,real(uvals)')#,cmap="cubehelix")
    plot!(Γ;color="black")

xlabel("\$x\$");ylabel("\$y\$");colorbar(ax=gca(),shrink=0.685)#515#2/3)
savefig("ScatteringNeumannPDEplot.png";dpi=600,bbox_inches="tight")

function makegif(x,y,u,L;plotfunction=Main.PyPlot.contourf,seconds=1,cmap="seismic",vert=1)
    tm=string(time_ns())
    dr = pwd()*"/"*tm*"mov"
    mkdir(dr)

    umax = maxabs(u)
    fps = 24
    MLen = seconds*fps

    tt = linspace(-1.,1.,100)

    for k=1:MLen
        t = 2π/ω*(k-1)/fps
        Main.PyPlot.clf()
        Main.PyPlot.axes(aspect="equal")
        Main.PyPlot.plot(tt,tt.^3;color=:black)
        Main.PyPlot.plot(tt-1,tt.^4.+4;color=:black)
        plotfunction(x,y,real(u*exp(-im*ω*t))';vmin=-umax*vert,vmax=umax*vert,cmap="seismic")
        xlabel!("\$x\$");ylabel!("\$y\$")
        Main.PyPlot.savefig(dr * "/" * lpad(k,max(4,ceil(Int,log10(MLen))),0) * ".png";dpi=150,bbox_inches="tight")
    end
    # Requires: brew install imagemagick
    run(`convert -delay 6 -loop 0 $dr/*.png $dr/scattering.gif`)
    run(`open $dr/scattering.gif`)
end

makegif(x,y,uvals,100;seconds=1,cmap="seismic",vert=0.5)

fps = 24
seconds=1
MLen = seconds*fps

k=MLen
u=uvals;
x=xx;y=yy;
k=24
    t = 2π/ω*(k-1)/fps
    real(u*exp(-im*ω*t))
    Main.PyPlot.contourf(x,y,real(u*exp(-im*ω*t))')
    plot!(Γ;color=:black,legend=false)
xlabel!("\$x\$");ylabel!("\$y\$")

tm=string(time_ns())
dr = pwd()*"/"*tm*"mov"
mkdir(dr)

Main.PyPlot.savefig(dr * "/" * lpad(k,max(4,ceil(Int,log10(MLen))),0) * ".png";dpi=150,bbox_inches="tight")


module TestPseudoInverse
using LinearAlgebra
using LazyInverses
using Test

function scalar_tests(A, M, Inv)
    elty = eltype(A)
    x = randn(elty)
    @testset "scalar tests with eltype $elty" begin
        @test Matrix(Inv * x) ≈ M * x
        @test Matrix(x * Inv) ≈ x * M
        @test Matrix(Inv / x) ≈ M / x
        @test Matrix(x \ Inv) ≈ x \ M
        # @test Matrix(Inv \ x) ≈ A * x
        # @test Matrix(x / Inv) ≈ x * A
    end
end

@testset "pseudoinverse" begin
    n, m = 3, 2
    element_types = (Float32, Float64, ComplexF32, ComplexF64)
    for elty in element_types
        A = randn(elty, n, m)
        a = randn(elty, n)
        b = randn(elty, m)
        B = randn(elty, m, n)

        # pseudoinverse
        LInv = pseudoinverse(A)
        @test LInv isa PseudoInverse{<:elty}
        ML = Matrix(LInv)
        @test ML ≈ pinv(A)

        @test Matrix(adjoint(LInv)) ≈ adjoint(ML)
        @test Matrix(transpose(LInv)) ≈ transpose(ML)

        # scalar operations
        scalar_tests(A, ML, LInv)

        # vector operations
        @test LInv * a ≈ ML * a
        @test b' * LInv ≈ b' * ML

        # matrix operations
        @test LInv * A ≈ I(m)
        @test ML * A ≈ I(m)
        @test A * LInv ≈ A * ML
        @test B' * LInv ≈ B' * ML
        @test LInv * B' ≈ ML * B'
        @test B * LInv' ≈ B * ML'

        # factorize
        @test factorize(LInv) ≡ LInv

        # smart constructor
        A = randn(elty, n, n)
        @test pinverse(A) isa Inverse

        x = rand(elty, (1, 1))
        @test pinverse(x) isa elty
        @test pinverse(x) == inv(x[1]) # smart pseudo-constructor returns scalar on 1x1 matrix
        @test pinverse(Inverse(x)) == x[1]

        # scalar
        x = rand(elty)
        PInv = PseudoInverse(x)
        @test PInv * x ≈ 1
        @test x * PInv ≈ 1
        @test pinverse(x) ≈ 1/x
        @test AbstractMatrix(PInv) ≈ fill(1/x, (1, 1))
    end
end

end # module


"""
This package facilitates working with probability distributions by means of Monte-Carlo methods, in a way that allows for propagation of probability distributions through functions. This is useful for, e.g.,  nonlinear [uncertainty propagation](https://en.wikipedia.org/wiki/Propagation_of_uncertainty). A variable or parameter might be associated with uncertainty if it is measured or otherwise estimated from data. We provide two core types to represent probability distributions: `Particles` and `StaticParticles`, both `<: Real`. (The name "Particles" comes from the [particle-filtering](https://en.wikipedia.org/wiki/Particle_filter) literature.) These types all form a Monte-Carlo approximation of the distribution of a floating point number, i.e., the distribution is represented by samples/particles. **Correlated quantities** are handled as well, see [multivariate particles](https://baggepinnen.github.io/MonteCarloMeasurements.jl/stable/#Multivariate-particles-1) below.

A number of type `Particles` behaves just as any other `Number` while partaking in calculations. After a calculation, an approximation to the **complete distribution** of the output is captured and represented by the output particles. `mean`, `std` etc. can be extracted from the particles using the corresponding functions. `Particles` also interact with [Distributions.jl](https://github.com/JuliaStats/Distributions.jl), so that you can call, e.g., `Normal(p)` and get back a `Normal` type from distributions or `fit(Gamma, p)` to get a `Gamma`distribution. Particles can also be iterated, asked for `maximum/minimum`, `quantile` etc. If particles are plotted with `plot(p)`, a histogram is displayed. This requires Plots.jl. A kernel-density estimate can be obtained by `density(p)` is StatsPlots.jl is loaded.

## Quick start
```julia
julia> using MonteCarloMeasurements, Plots

julia> a = π ± 0.1 # Construct Gaussian uncertain parameters using ± (\\pm)
Particles{Float64,2000}
 3.14159 ± 0.1

julia> b = 2 ∓ 0.1 # ∓ (\\mp) creates StaticParticles (with StaticArrays)
StaticParticles{Float64,100}
 2.0 ± 0.0999

julia> std(a)      # Ask about statistical properties
0.09999231528930486

julia> sin(a)      # Use them like any real number
Particles{Float64,2000}
 1.2168e-16 ± 0.0995

julia> plot(a)     # Plot them

julia> b = sin.(1:0.1:5) .± 0.1; # Create multivariate uncertain numbers

julia> plot(b)                   # Vectors of particles can be plotted

julia> using Distributions

julia> c = Particles(500, Poisson(3.)) # Create uncertain numbers distributed according to a given distribution
Particles{Int64,500}
 2.882 ± 1.7
```

For further help, see the [documentation](https://baggepinnen.github.io/MonteCarloMeasurements.jl/stable), the [examples folder](https://github.com/baggepinnen/MonteCarloMeasurements.jl/tree/master/examples) or the [arXiv paper](https://arxiv.org/abs/2001.07625).
"""
module MonteCarloMeasurements
using LinearAlgebra, Statistics, Random, StaticArrays, RecipesBase, MacroTools, SLEEFPirates
using Distributed: pmap
import Base: add_sum

using Distributions, StatsBase, Requires

const DEFAULT_NUM_PARTICLES = 2000
const DEFAULT_STATIC_NUM_PARTICLES = 100

function pmean end

"""
The function used to reduce particles to a number for comparison. Defaults to `mean`. Change using `unsafe_comparisons`.
"""
const COMPARISON_FUNCTION = Ref{Function}(pmean)
const COMPARISON_MODE = Ref(:safe)

"""
    unsafe_comparisons(onoff=true; verbose=true)
Toggle the use of a comparison function without warning. By default `mean` is used to reduce particles to a floating point number for comparisons. This function can be changed, example: `set_comparison_function(median)`

    unsafe_comparisons(mode=:reduction; verbose=true)
One can also specify a comparison mode, `mode` can take the values `:safe, :montecarlo, :reduction`. `:safe` is the same as calling `unsafe_comparisons(false)` and `:reduction` corresponds to `true`.
If
"""
function unsafe_comparisons(mode=true; verbose=true)
    mode == false && (mode = :safe)
    mode == true && (mode = :reduction)
    COMPARISON_MODE[] = mode
    if mode != :safe && verbose
        if mode === :reduction
            @info "Unsafe comparisons using the function `$(COMPARISON_FUNCTION[])` has been enabled globally. Use `@unsafe` to enable in a local expression only or `unsafe_comparisons(false)` to turn off unsafe comparisons"
        elseif mode === :montecarlo
            @info "Comparisons using the monte carlo has been enabled globally. Call `unsafe_comparisons(false)` to turn off unsafe comparisons"
        end
    end
    mode ∉ (:safe, :montecarlo, :reduction) && error("Got unsupported comparison model")
end
"""
    set_comparison_function(f)

Change the Function used to reduce particles to a number for comparison operators
Toggle the use of a comparison Function without warning using the Function `unsafe_comparisons`.
"""
function set_comparison_function(f)
    if f in (mean, median, maximum, minimum)
        @warn "This comparison function ($(f)) is probably not the right choice, consider if you want the particle version (p$(f)) instead."
    end
    COMPARISON_FUNCTION[] = f
end

"""
    @unsafe expression
Activates unsafe comparisons for the provided expression only. The expression is surrounded by a try/catch block to robustly restore unsafe comparisons in case of exception.
"""
macro unsafe(ex)
    ex2 = if @capture(ex, assigned_vars__ = y_)
        if length(assigned_vars) == 1
            esc(assigned_vars[1])
        else
            esc.(assigned_vars[1].args)
        end
    else
        :(res)
    end
    quote
        previous_state = COMPARISON_MODE[]
        unsafe_comparisons(true, verbose=false)
        local res
        try
            res = ($(esc(ex)))
        finally
            unsafe_comparisons(previous_state, verbose=false)
        end
        $ex2 = res
    end
end

export ±, ∓, .., ⊠, ⊞, AbstractParticles,Particles,StaticParticles, MvParticles, sigmapoints, transform_moments, ≲,≳, systematic_sample, ess, outer_product, meanstd, meanvar, register_primitive, register_primitive_multi, register_primitive_single, ℝⁿ2ℝⁿ_function, ℝⁿ2ℂⁿ_function, ℂ2ℂ_function, ℂ2ℂ_function!, resample!, bootstrap, sqrt!, exp!, sin!, cos!, wasserstein, with_nominal, nominal, nparticles, particleeltype
# Plot exports
export errorbarplot, mcplot, ribbonplot

# Statistics reexport
export mean, std, cov, var, quantile, median
export pmean, pstd, pcov, pcor, pvar, pquantile, pmedian, pmiddle, piterate, pextrema, pminimum, pmaximum
# Distributions reexport
export Normal, MvNormal, Cauchy, Beta, Exponential, Gamma, Laplace, Uniform, fit, logpdf

export unsafe_comparisons, @unsafe, set_comparison_function

export bymap, bypmap, @bymap, @bypmap, @prob, Workspace, with_workspace, has_particles, mean_object

include("types.jl")
include("register_primitive.jl")
include("sampling.jl")
include("particles.jl")
include("distances.jl")
include("complex.jl")
include("sigmapoints.jl")
include("resampling.jl")
include("bymap.jl")
include("deconstruct.jl")
include("diff.jl")
include("plotting.jl")
include("optimize.jl")
include("sleefpirates.jl")
include("nominal.jl")

# This is defined here so that @bymap is loaded
LinearAlgebra.norm2(p::AbstractArray{<:AbstractParticles}) = bymap(LinearAlgebra.norm2,p)
Base.:\(x::AbstractVecOrMat{<:AbstractParticles}, y::AbstractVecOrMat{<:AbstractParticles}) = bymap(\, x, y)
Base.:\(x::Diagonal{<:AbstractParticles}, y::Vector{<:AbstractParticles}) = bymap(\, x, y) # required for ambiguity

function __init__()
    @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("forwarddiff.jl")
    @require Measurements="eff96d63-e80a-5855-80a2-b1b0885c5ab7" include("measurements.jl")
    @require Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" include("unitful.jl")
end

end


# Board struct
struct Board
  n :: Int
  lines :: Vector{Line}
  cells :: Matrix{Int}
end

# Creates board that can hold the given lines.
function Board(lines::Vector{Line})
  n = maximum(map(l -> max(l.startPt[1],l.startPt[2],l.endPt[1],l.endPt[2]),lines))
  return Board(n,lines,zeros(Int,(n,n)))
end

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

# Line struct
struct Line
  startPt :: CartesianIndex{2}
  endPt   :: CartesianIndex{2}
  points  :: Vector{CartesianIndex{2}}
end

# Creates line from the two endpoints.
function Line(p1::CartesianIndex{2}, p2::CartesianIndex{2})
  slope = (p2-p1)
  npts = max(abs(slope[1]), abs(slope[2]))
  slope = CartesianIndex(slope[1] ÷ npts,slope[2] ÷ npts)
  pts = [p1 + i*slope for i = 0:npts]
  return Line(p1, p2, pts)
end

# Check line properties
isHorizontal(l::Line) :: Bool = (l.startPt[1] == l.endPt[1])
isVertical(l::Line)   :: Bool = (l.startPt[2] == l.endPt[2])
isFlat(l::Line)       :: Bool = isHorizontal(l) || isVertical(l)

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

function parsePoint(s::String) :: CartesianIndex{2}
  coords = map(x -> parse(Int,x),split(s, ',', keepempty=false))
  return CartesianIndex(coords[1], coords[2])
end

function parseLine(s::String) :: Line
  pts = split(s, "->", keepempty=false)
  p1 = parsePoint(String(pts[1]))
  p2 = parsePoint(String(pts[2]))
  return Line(p1,p2)
end

function readInput() :: Vector{Line}
  f = open("./inputs/day5.txt")

  lines :: Vector{Line} = []
  for line in eachline(f)
    push!(lines, parseLine(line))
  end
  return lines
end

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

# Draw a single line into a matrix
function add!(line::Line, M::Matrix{Int})
  for p in line.points
    M[p] += 1
  end
end

# Draw all lines in a board
function drawLines!(b::Board)
  map(l -> add!(l,b.cells), b.lines)
end

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

lines = readInput()

# Problem 1
lines1 = deepcopy(filter(isFlat, lines)) # Only consider flat lines
board1 = Board(lines1)
drawLines!(board1)
res = count( map( x -> (x > 1) ? true : false, board1.cells))

# Problem 2
lines2 = deepcopy(lines)
board2 = Board(lines2)
drawLines!(board2)
res = count( map( x -> (x > 1) ? true : false, board2.cells))


# Common univariate distributions.
distset = [
    Normal(1.23, 3.45),
    Exponential(2.12),
    Gamma(4.3),
    Beta(2.123),
    LogNormal(4.5),
    DiscreteUniform(1, 10),
    Binomial(101, 0.45),
    Categorical([0.2, 0.35, 0.15, 0.3])
]

for dist in distset
    μ = mean(dist)
    σ = std(dist)
    # No convenience call.
    E = expectation(dist)
    @test E(x -> x) ≈ μ
    @test E(x -> x^2) - μ^2 ≈ σ^2
    @test E(x -> ((x - μ)/σ)^3) + 1. ≈ skewness(dist) + 1. # To avoid comparisons to 0.0 exactly.
    # Convenience call.
    @test expectation(x -> x, dist) ≈ μ
    @test expectation(x -> x^2, dist) - μ^2 ≈ σ^2
    @test expectation(x -> ((x - μ)/σ)^3, dist) + 1. ≈ skewness(dist) + 1.
    println(dist)
end

# Linear operator behavior.
distset = [
    DiscreteUniform(1., 10.),
    Normal(1.45),
    Beta(2.1)
]

for dist in distset
    E = expectation(dist)
    h(x) = 2*x
    z = nodes(E)
    @test E * h.(z) ≈ E(x -> 2*x) # Right-multiplying.
    @test weights(2E) ≈ 2*weights(E) # Left-multiplying.
    @test 3E*z ≈ (3E) * z ≈ 3*(E * z) # Linearity.
end

# Error handling.
distset = [ # Noncompact dists
    Beta(Inf),
    LogNormal(Inf, Inf),
    Normal(0, Inf),
    Uniform(-Inf, Inf),
    Poisson(3)
]

for dist in distset
    @test_throws ArgumentError expectation(dist)
end

# Other errors.
E = expectation(DiscreteUniform(1, 10))
@test_throws Exception E(x -> dot(x, ones(7))) # Non-applicable functions.
@test_throws MethodError (x -> 2*x).(nodes(E)) * E # Non-commutativity.


using LTVModels, LowLevelParticleFilters
using LTVModels, Plots, StatPlots, LinearAlgebra, Distributions
T = 1000
eye(n) = Matrix{Float64}(I,n,n)
σ_state_drift = 0.0001
σ_param_drift = 0.0001
σ_noise = 1std(x)
A,B,x,u,n,m,N = LTVModels.testdata(T=T, σ_state_drift=σ_state_drift, σ_param_drift=σ_param_drift, seed=1)

R1          = σ_param_drift^2*eye(n^2+n*m)
R2          = σ_state_drift^2*eye(n)
Re          = σ_noise^2*eye(n)

P0          = 1000000000eye(n^2+n*m)

xn = x + σ_noise*randn(size(x))
@show √(mean(abs2, xn-x))
xt,ut,y = xn[:,1:end-1],u[:,1:end-1],x[:,2:end]

vv(x) = [x[:,i] for i = 1:size(x,2)]

function iterated_kalman(xtrue,xi,u,R1,R2,P0; iters=10, kwargs...)
    T = size(xi,2)
    x = deepcopy(xi)
    pe,xe,me,ll,llx = zeros(iters), zeros(iters+1), zeros(iters), zeros(iters), zeros(iters)
    xe[1] = √(mean(abs2, xtrue-x))
    local model
    for i = 1:iters
        model = KalmanModel(x,u,R1,Re,P0; kwargs...)
        kf = KalmanFilter(model.At, model.Bt, I, 0, R2, Re, MvNormal(xi[:,1], Re))
        # xv,R,llxi = smooth(kf, vv(u), vv(x))
        _,xv,R,Rt,llxi = forward_trajectory(kf, vv(u), vv(x))
        # aₘ = 1/(1+iters/10+i)^(0.9)
        aₘ = ((2+iters/10)^(1))/(1+iters/10+i)^(1)
        # aₘ = 1/i
        x .+= aₘ*(hcat(xv...) - x)
        # xv = hcat(xv...)
        pe[i] = √(mean(abs2, y-predict(model,xt,ut)))
        me[i] = √(mean(abs2, A-model.At)) + √(mean(abs2, B-model.Bt))
        xe[i+1] = √(mean(abs2, xtrue-x))
        ll[i] = model.ll
        llx[i] = llxi
    end
    x,model,pe,xe,me,ll,llx
end

xh, model,pe,xe,me,ll,llx = iterated_kalman(x,xn,u,R1,R2,P0, extend=true, iters=3)

# plot(x', c=:blue, layout=3)
# plot!(xn', c=:red)
# plot!(xh', c=:green)

plot((x-xn)', c=:blue, layout=4, lab="Original noise", subplot=1:3)
plot!((x-xh)', c=:green, lab="Estimation error", subplot=1:3)
plot!(pe, subplot=4, lab="Prediction error")
plot!(me, subplot=4, lab="Model error")
# hline!([σ_noise], subplot=4, lab="sigma_e", l=(:dash,:black))
plot!(xe[1:end], subplot=4, lab="State error", yscale=:log10, xscale=:log10, ylims=(NaN,NaN), legend=:topright, legendfontsize=6) |> display

# plot(xe, lab="State error", yscale=:log10, xscale=:log10, xlabel="Iteration")
# plot!(me, lab="Model error")
# # savetikz("/local/home/fredrikb/phdthesis/ltvpaper/figs/iterated_kalman.tex")
#
# plot([ll[2:end] llx[2:end]],xscale=:log10, layout=2)
# # plot!(-llx)
plot(flatten(A), c=:blue, lab="")
plot!(flatten(model.At), c=:red, lab="") |> display

# serialize("/tmp/wtf", (x,xn,xh))

function dlyap(A,B)
    S = randn(size(A))
    for i = 1:100
        S = A*S*A' + B*B'
    end
    S
end

A = A[:,:,1]
B = B[:,:,1]
P = dlyap(A,B)
K = P/(Re+P)
C = K*Re*K'
Ce = cov((x-xh)')

P + A*Re*A'


export trouve_intervalle
function trouve_intervalle(h :: AbstractNLPModel,
                           t₀ :: Float64,
                           inc0 :: Float64;
                           verbose :: Bool = false)

        iter = 1
        h₀ = obj(h, [t₀])[1]
        g₀ = grad(h, [t₀])[1]
        sd = -sign(g₀)
        inc = inc0
        t₁ = t₀ + sd * inc

        h₁ = obj(h, [t₁])[1]
        g₁ = grad(h, [t₁])[1]

        verbose &&
            @printf("iter t₀        g₀        h₀         t₁        ")
        verbose &&
            @printf("g₀        h₀\n")
        verbose &&
            @printf("%4d %7.2e %7.2e  %7.2e  %7.2e  %7.2e  %7.2e \n",
                iter,  t₀, g₀, h₀, t₁, g₀, h₀)

        while (g₁ * sd < 0.0) && (h₁ < h₀) && (iter < 30)
          inc = inc * 4
          t₀ = t₁
          h₀ = h₁
          g₀ = g₁
          t₁ = t₀ + sd * inc
          h₁ = obj(h, [t₁])[1]
          g₁ = grad(h, [t₁])[1]
          verbose && @printf("%4d %7.2e %7.2e  %7.2e  %7.2e  %7.2e  %7.2e \n",
                              iter,  t₀, g₀, h₀, t₁, g₀, h₀)
          iter = iter + 1
        end

        while (g₁ * sd < 0.0) & (iter < 30)
          tₘ = (t₁ + t₀) / 2
          hₘ = obj(h, [tₘ])[1]
          gₘ = grad(h, [tₘ])[1]
          if gₘ * sd>0
            t₁ = tₘ
            g₁ = gₘ
            h₁ = hₘ
          else
            if hₘ < h₀
              t₀ = tₘ
              h₀ = hₘ
              g₀ = gₘ
            else
              t₁ = tₘ
              h₁ = hₘ
              g₁ = gₘ
            end
          end
          verbose &&
            @printf("%4d %7.2e %7.2e  %7.2e  %7.2e  %7.2e  %7.2e \n",
                    iter,  t₀, g₀, h₀, t₁, g₀, h₀)
          iter = iter + 1
        end

        a = min(t₀, t₁)
        b = max(t₀, t₁)

        return (a, b)

end


#!/usr/bin/env julia

using SparseArrays
using Base.Iterators
using ArgParse

function printml(arg)
    show(IOContext(stdout), "text/plain", arg)
    println()
end

function make_rebin_matrix_1d(sizein::T, ntogroup::T) where T
    sizeout = sizein÷ntogroup
    @assert sizein%ntogroup==0

    Kblock = sparse(ones(Int64, (ntogroup,1)))
    # println("Kblock: $(Array(Kblock))")

    blockdiag(repeated(Kblock, sizeout)...)
end

function make_rebin_matrices_2d(shapein::Tuple{T,T}, ntogroup::Tuple{T,T}) where T
    Kleft  = make_rebin_matrix_1d(shapein[1], ntogroup[1])'
    Kright = make_rebin_matrix_1d(shapein[2], ntogroup[2])

    Kleft, Kright
end

function check1d(; verbose::Bool)
    sizein=12::Int64
    ngroup=4::Int64

    K = make_rebin_matrix_1d(sizein, ngroup)
    @time K = make_rebin_matrix_1d(sizein, ngroup)
    if verbose
        println("Size in: $sizein")
        println("N to group: $ngroup")
        println("Shape K: $(size(K))")
        println("K: $(Array(K))")
        println()
    end

    arrin1  = ones(Float64, (1,sizein))

    arrout1 = arrin1*K

    arrin2 = ones(Float64, (3,sizein))
    arrin2[2,:].*=2
    arrin2[3,:].*=3

    arrout2 = arrin2*K

    if verbose
        println("Shape1 in: $(size(arrin1))")
        println("Shape1 out: $(size(arrout1))")
        println("Arr1 in: $arrin1")
        println("Arr1 out: $arrout1")
        println()
        println("Shape2 in: $(size(arrin2))")
        println("Shape2 out: $(size(arrout2))")
        println("Arr2 in: $arrin2")
        println("Arr2 out: $arrout2")
        println()
    end

end

function check2d(; verbose::Bool)
    verbose = Bool(verbose)
    shapein=(6, 4)
    ngroup=(3, 2)

    Kleft, Kright = make_rebin_matrices_2d(shapein, ngroup)
    @time Kleft, Kright = make_rebin_matrices_2d(shapein, ngroup)

    if verbose
        println("N to group: $ngroup")
        println("Shape K left: $(size(Kleft))")
        println("Shape in: $shapein")
        println("Shape K right: $(size(Kright))")
        println("K left:")
        println(Kleft)
        println("K right:")
        printml(Array(Kright))
        println()
    end

    arrin1 = reshape(Array(range(0, length=prod(shapein))), reverse(shapein))'

    arrout1 = Kleft*arrin1*Kright

    arrin2 = ones(Float64, shapein)
    arrin2[begin:3, begin:2] .= 1.0
    arrin2[begin:3, 3:end]   .= 2.0
    arrin2[4:end, begin:2]   .= 3.0
    arrin2[4:end, 3:end]     .= 4.0

    arrout2 = Kleft*arrin2*Kright

    if verbose
        println("Shape1 in: $(size(arrin1))")
        println("Shape1 out: $(size(arrout1))")
        println("Arr1 in:")
        printml(arrin1)
        println("Arr1 out:")
        printml(arrout1)
        println()
        println("Shape2 in: $(size(arrin2))")
        println("Shape2 out: $(size(arrout2))")
        println("Arr2 in:")
        printml(arrin2)
        println("Arr2 out:")
        printml(arrout2)
    end
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--verbose", "-v"
            help = "print data"
            action = :store_true
    end

    return parse_args(s, as_symbols=true)
end

function main()
    opts = parse_commandline()
    check1d(; opts...)
    check2d(; opts...)
end

main()