######################################################################
# weierstrass.jl : curves given by weierstrass equations
######################################################################

export Weierstrass, ShortWeierstrass, SeparatedWeierstrass, EllipticCurve, from_j_invariant

export curvetype, a_invariants, b_invariants, c_invariants, discriminant, j_invariant, equation, projective_equation, divisionpolynomial, coordinate_ring

######################################################################
# Type definitions
######################################################################

"""
Concrete type for elliptic curves in long Weierstrass form over a ring.
"""
struct Weierstrass{T} <: AbstractWeierstrass{T}
    a1::T
    a2::T
    a3::T
    a4::T
    a6::T
end

"""
Concrete types for elliptic curves in short Weierstrass form over a ring.
"""
struct ShortWeierstrass{T} <: AbstractWeierstrass{T}
    a::T
    b::T
end

struct SeparatedWeierstrass{T} <: AbstractWeierstrass{T}
	a2::T
	a4::T
	a6::T
end

curvetype(E::SeparatedWeierstrass) = SeparatedWeierstrass

curvetype(E::ShortWeierstrass) = ShortWeierstrass

curvetype(E::Weierstrass) = Weierstrass

######################################################################
# Calling with six arguments
######################################################################

function ShortWeierstrass(a1::T, a2::T, a3::T, a4::T, a6::T) where T
	a1 == 0 && a2 == 0 && a3 == 0 || throw(ArgumentError("Tried to call ShortWeierstrass with nonzero coefficients"))
	return ShortWeierstrass(a4, a6)
end

function SeparatedWeierstrass(a1::T, a2::T, a3::T, a4::T, a6::T) where T
	a1 == 0 && a3 == 0 || throw(ArgumentError("Tried to call SeparatedWeierstrass with nonzero coefficients"))
	return SeparatedWeierstrass(a2, a4, a6)
end

######################################################################
# Calling with the 'EllipticCurve' constructor
######################################################################

function EllipticCurve(a4::T, a6::T) where T
	return ShortWeierstrass(a4, a6)
end

function EllipticCurve(a2::T, a4::T, a6::T) where T
	return SeparatedWeierstrass(a2, a4, a6)
end

function EllipticCurve(a1::T, a2::T, a3::T, a4::T, a6::T) where T
	return Weierstrass(a1, a2, a3, a4, a6)
end

function from_j_invariant(j::T) where T
    a = -j * (j - 1728)
    b = -a * (j - 1728)
    return ShortWeierstrass(3*a, 2*b)
end

######################################################################
# Basic methods
######################################################################

base_ring(E::Weierstrass) = Nemo.parent(E.a1)

base_ring(E::ShortWeierstrass) = Nemo.parent(E.a)

base_ring(E::SeparatedWeierstrass) = Nemo.parent(E.a2)

"""
Get a description of an elliptic curve in long Weierstrass form.
"""
function show(io::IO, E::Weierstrass)
    print(io, "Elliptic Curve in long Weierstrass form y² + $(E.a1) xy + $(E.a3) y = x³ + $(E.a2) x² + $(E.a4) x + $(E.a6)  over ")
    show(io, base_ring(E))
end

"""
Get a description of an elliptic curve in short Weierstrass form.
"""
function show(io::IO, E::ShortWeierstrass)
    print(io, "Elliptic Curve in short Weierstrass form y² = x³ + $(E.a) x + $(E.b)  over ")
    show(io, base_ring(E))
end

"""
Get a description of an elliptic curve in short Weierstrass form.
"""
function show(io::IO, E::SeparatedWeierstrass)
    print(io, "Elliptic Curve in separated Weierstrass form y² = x³ + $(E.a2) x² + $(E.a4) x + $(E.a6)  over ")
    show(io, base_ring(E))
end

######################################################################
# Invariants and equations
######################################################################

"""
Get the a-invariants (a1, a2, a3, a4, a6) of an elliptic curve in long Weierstrass form.
"""
a_invariants(E::Weierstrass) = (E.a1, E.a2, E.a3, E.a4, E.a6)

"""
Get the a-invariants (a1, a2, a3, a4, a6) of an elliptic curve in short Weierstrass form.
"""
function a_invariants(E::ShortWeierstrass)
    zero = Nemo.zero(base_ring(E))
    return (zero, zero, zero, E.a, E.b)
end

"""
Get the a-invariants (a1, a2, a3, a4, a6) of an elliptic curve in separated Weierstrass form.
"""
function a_invariants(E::SeparatedWeierstrass)
    zero = Nemo.zero(base_ring(E))
    return (zero, E.a2, zero, E.a4, E.a6)
end

==(E1::AbstractWeierstrass, E2::AbstractWeierstrass) = a_invariants(E1) == a_invariants(E2)

"""
Get the b-invariants (b2, b4, b6, b8) of an elliptic curve in Weierstrass form.
"""
function b_invariants(E::AbstractWeierstrass)
    a1, a2, a3, a4, a6 = a_invariants(E)
    return (a1*a1 + 4*a2,
             a1*a3 + 2*a4,
             a3^2 + 4*a6,
             a1^2 * a6 + 4*a2*a6 - a1*a3*a4 + a2*a3^2 - a4^2)
end

"""
Get the c-invariants (c4, c6) of an elliptic curve in Weierstrass form.
"""
function c_invariants(E::AbstractWeierstrass)
    b2, b4, b6, b8 = b_invariants(E)
    return (b2^2 - 24*b4, -b2^3 + 36*b2*b4 - 216*b6)
end

"""
Get the discriminant of an elliptic curve in Weierstrass form.
"""
function discriminant(E::AbstractWeierstrass)
    b2, b4, b6, b8 = b_invariants(E)
    return -b2^2*b8 - 8*b4^3 - 27*b6^2 + 9*b2*b4*b6
end

"""
Check if a given instance of AbstractWeierstrass is indeed a nonsingular curve.
"""
isvalid(E::AbstractWeierstrass) = !iszero(discriminant(E))

"""
Get the j-invariant of an elliptic curve in Weierstrass form.
This may fail if the base ring is not a field.
"""
function j_invariant(E::AbstractWeierstrass)
    c4, _ = c_invariants(E)
    disc = discriminant(E)
	return c4^3 // disc
end

function equation(E::AbstractWeierstrass) #use nested polynomial rings for efficient reduction
	K = base_ring(E)
	RX, X = PolynomialRing(K, "X")
	RXY, Y = PolynomialRing(RX, "Y")
	a1, a2, a3, a4, a6 = a_invariants(E)
	return Y^2 + a1 * X * Y + a3 * Y - X^3 - a2 * X^2 - a4 * X - a6
end

function equation(E::AbstractWeierstrass, X, Y)
	a1, a2, a3, a4, a6 = a_invariants(E)
	return Y^2 + a1 * X * Y + a3 * Y - X^3 - a2 * X^2 - a4 * X - a6
end

function projective_equation(E::AbstractWeierstrass)
	K = base_ring(E)
	R, (X, Y, Z) = PolynomialRing(K, ["X", "Y", "Z"])
	a1, a2, a3, a4, a6 = a_invariants(E)
	return Z * Y^2 + a1 * X * Y * Z + a3 * Y * Z^2 - X^3 - a2 * X^2 * Z - a4 * X * Z^2 - a6 * Z^3
end

######################################################################
# Division polynomials for Weierstrass curves
######################################################################

"""
Internal function to compute the 2-division polynomial of an elliptic curve in Weierstrass form.
"""
_2divpoly(E::AbstractWeierstrass{T}, x::PolyElem{T}, D::Dict{Int, PolyElem{T}}) where T = _divpoly(E, -1, x, D)

"""
Internal function to compute the mth division polynomial of an elliptic curve in Weierstrass form, using the well-known induction formulas. If the index is even, *the usual factor 2*y is dropped*.
"""
function _divpoly(E::AbstractWeierstrass{T}, m::Integer, x::PolyElem{T}, D::Dict{Int, PolyElem{T}}) where T
	b2, b4, b6, b8 = b_invariants(E)
	R = parent(x)
	if m in keys(D)
		return D[m]
	elseif m == -2
		D[-2] = _divpoly(E, - 1, x, D)^2
		return D[-2]
	elseif m == -1
		D[-1] = 4 * x^3 + b2 * x^2 + 2 * b4  *x + b6
		return D[-1]
	elseif m <= 0
		throw(ArgumentError("Non-positive integers in _divpoly must be either -1 or -2, got $m"))
	elseif m == 1
		D[1] = R(1)
		return D[1]
	elseif m == 2
		D[2] = R(1)
		return D[2]
	elseif m == 3
		D[3] = 3 * x^4 + b2 * x^3 + 3 * b4 * x^2 + 3 * b6 * x + b8
		return D[3]
	elseif m == 4
		pm2 = _divpoly(E, -2, x, D)
		p3 = _divpoly(E, 3, x, D)
		D[4] = - pm2 + (6*x^2 + b2*x + b4) * p3
		return D[4]
	elseif isodd(m)
		k = div(m - 1, 2)
		m1 = _divpoly(E, k - 1, x, D)
		eq = _divpoly(E, k, x, D)
		p1 = _divpoly(E, k + 1, x, D)
		p2 = _divpoly(E, k + 2, x, D)
		if iseven(k)
			D[m] = _divpoly(E, -2, x, D) * p2 * eq^3 - m1 * p1^3
			return D[m]
		else #k is odd
			D[m] = p2 * eq^3 - _divpoly(E, -2, x, D) * m1 * p1^3
			return D[m]
		end
	else  #m is even
		k = div(m, 2)
		m2 = _divpoly(E, k - 2, x, D)
		m1 = _divpoly(E, k - 1, x, D)
		eq = _divpoly(E, k, x, D)
		p1 = _divpoly(E, k + 1, x, D)
		p2 = _divpoly(E, k + 2, x, D)
		D[m] = eq * (p2 * m1^2 - m2 * p1^2)
		return D[m]
	end
end

"""
Compute the division polynomial of an elliptic curve in Weierstrass form.

If m is even, you may specify an additional parameter two_torsion:
* 0 simply ignores the points of order two
* 1 returns the exact bivariate division polynomial
* 2 (default) adds extra roots at points of two-torsion in order to get a univariate polynomial.
"""

function divisionpolynomial(E::AbstractWeierstrass{T}, m::Nemo.Integer, two_torsion::Integer = 2) where T
	m <= -3 && throw(ArgumentError("m must be positive or -1 or -2"))

	if m<=0
		K = base_ring(E)
		R, x = Nemo.PolynomialRing(K, "x")
		D = Dict{Int, Nemo.PolyElem{T}}()
		poly = _divpoly(E, m, x, D)
		return poly
	elseif m%2 == 0
		if two_torsion == 1
			K = base_ring(E)
			a1, _, a3, _, _ = a_invariants(E)
			R, (x, y) = Nemo.PolynomialRing(K, ["x", "y"])
			D = Dict{Int, Nemo.PolyElem{T}}()
			poly = _divpoly(E, m, x, D)
			return (2 * y + a1 * x + a3) * poly
		elseif two_torsion == 2
			K = base_ring(E)
			R, x = Nemo.PolynomialRing(K, "x")
			D = Dict{Int, Nemo.PolyElem{T}}()
			poly = _divpoly(E, m, x, D)
			pm1 = _divpoly(E, -1, x, D)
			return pm1 * poly
		else #two_torsion == 0
			K = base_ring(E)
			R, x = Nemo.PolynomialRing(K, "x")
			D = Dict{Int, Nemo.PolyElem{T}}()
			poly = _divpoly(E, m, x, D)
			return poly
		end
	else # m is odd
		K = base_ring(E)
		R, x = Nemo.PolynomialRing(K, "x")
		D = Dict{Int, Nemo.PolyElem{T}}()
		poly = _divpoly(E, m, x, D)
		return poly
	end
end

######################################################################
# Coordinate ring
######################################################################

function coordinate_ring(E::AbstractWeierstrass)
	Eq = equation(E)
	Y = gen(parent(Eq))
	X = gen(base_ring(Eq))
	RXY = parent(Eq)
	return ResidueRing(RXY, Eq), X, Y
end


function demoBurglarDictPot()
    # This demo is the same as demoBurglar.jl, except that here we use the ability to define non-contiguously indexed potentials and variables by demonstrating the use of the Dictoinary of PotArray technique.

    burglar,earthquake,alarm,radio=10,2,30,5 # Variable order and indexing is arbitary (thianks to using DictPot and doesn't need to be contiguous
    yes, no= 1,2 # define states, starting from 1. These do need to be contiguous, starting from 1, since currently the probability tables are defined by arrays of contiguous states.

    # This defines the meaning of the variables and their states.
    # It is not necessary for the computation, but is useful for
    # printing the results in a meaningful form:
    DictVariable=Dict{Int,DiscreteVariable}()
    DictVariable[burglar]=DiscreteVariable("burglar",["yes","no"])
    DictVariable[earthquake]=DiscreteVariable("earthquake",["yes","no"])
    DictVariable[alarm]=DiscreteVariable("alarm",["yes","no"])
    DictVariable[radio]=DiscreteVariable("radio",["yes","no"])

    # Define the discrete probability tables:
    DictPot=Dict{Integer,PotArray}()
    DictPot[burglar]=PotArray(burglar,[0.01 0.99])
    DictPot[earthquake]=PotArray(earthquake,[0.000001 1-0.000001])

    table=zeros(2,2)
    table[yes,yes]=1
    table[no,yes]=0
    table[yes,no]=0
    table[no,no]=1
    DictPot[radio]=PotArray([radio earthquake],table)

    table=zeros(2,2,2)
    table[yes,yes,yes]=0.9999
    table[yes,yes,no]=0.99
    table[yes,no,yes]=0.99
    table[yes,no,no]=0.0001
    table[no,:,:]=1-table[yes,:,:]
    DictPot[alarm]=PotArray([alarm burglar earthquake],table)

    jointpot=prod(DictPot)

    println("p(burglar|alarm=yes):")
    show(condpot(setpot(jointpot,alarm,yes),burglar),DictVariable)

    println("p(burglar|alarm=yes,radio=yes):")
    show(condpot(setpot(jointpot,[alarm radio],[yes yes]),burglar),DictVariable)

end


abstract type AbstractGraphProblem <: AbstractProblem end

import Base: length

using LightGraphs, SimpleWeightedGraphs

struct MST <: AbstractGraphProblem
    g::Any
    optimum::Real
    function MST(g; optimum::Real = -mst_value(g))
        new(g, optimum)
    end
end

function fitness(x::BitArray, problem::MST)
    penalty = 100000000
    g = problem.g

    subg = graph_bitstring(g, x)

    cc = connected_components(subg)
    if length(cc) != 1
        return -1 * (penalty + length(cc))
    end
    weights = sum(subg.weights) / 2 #Becuase subg.weights counts each edge twice.

    return -weights
end

function optimum(problem::MST)
    return problem.optimum
end

length(problem::MST) = nv(problem.g)

function TG(n)
    p = floor(Int, n / 4)
    q = floor(Int, n / 2)
    a = n^2
    g = SimpleWeightedGraph(complete_graph(q), 1)
    add_vertices!(g, 2 * p)
    for j = q+1:2:q+2p
        add_edge!(g, j, j - 1, 2 * a)
        add_edge!(g, j, j + 1, 2 * a)
        add_edge!(g, j + 1, j - 1, 3 * a)
    end
    g
end

#
function ER(n::Integer)
    g = erdos_renyi(n, 2 * Base.Math.log(n) / n)
    egs = collect(edges(g))
    k = SimpleWeightedGraph(n)
    for edge in egs
        add_edge!(k, src(edge), dst(edge), rand(1:n^2))
    end
    if n - 1 != length(prim_mst(k))
        return ER(n)
    end
    k
end

function Kn(n::Integer)
    g = complete_graph(n)
    egs = collect(edges(g))
    k = SimpleWeightedGraph(n)
    for edge in egs
        add_edge!(k, src(edge), dst(edge), rand(1:n^2))
    end
    if n - 1 != length(prim_mst(k))
        return ER(n)
    end
    k
end

#
#
# using GraphPlot
# g = ER(100)
# g.weights
#
#

#
#
struct MaximumMatching <: AbstractGraphProblem
    g::Any
    optimum::Real
    function MaximumMatching(g, optimum::Real)
        new(g, optimum)
    end
end

function fitness(x::BitArray, problem::MaximumMatching)
    h = copy(problem.g)

    for item in findall(row -> row == 0, x)
        edge = collect(edges(problem.g))[item]
        rem_edge!(h, edge)
    end

    degreeh = degree(h)
    dh = degreeh
    penalty = (ne(problem.g) + 1) * sum([max(0, i - 1) for i in dh])
    ne(h) - penalty
end

function optimum(problem::MaximumMatching)
    return problem.optimum
end

length(problem::MaximumMatching) = ne(problem.g)

function GHL(h, l)

    g = SimpleGraph(h * l)
    function nodeindex(L, H, l, h)
        return (h - 1) * L + l
    end

    for i = 1:l
        for j = 1:h
            if i != l
                if i % 2 == 1
                    add_edge!(g, nodeindex(l, h, i, j), nodeindex(l, h, i + 1, j))
                else
                    for k = 1:h
                        add_edge!(g, nodeindex(l, h, i, j), nodeindex(l, h, i + 1, k))
                    end
                end
            end
        end
    end
    g
end


export ExpWeight

immutable ExpWeight{S,FF,DD} <: WeightSpace{S,RealBasis,DD,1}
    exponent::FF
    space::S
end

function ExpWeight(q,sp)
    @assert domain(q) == domain(sp)
    ExpWeight{typeof(sp),typeof(q),typeof(domain(q))}(q,sp)
end

weight(sp::ExpWeight,x) = exp(sp.exponent(x))
spacescompatible(a::ExpWeight,b::ExpWeight) = a.exponent ≈ b.exponent && spacescompatible(a.space,b.space)
setdomain(sp::ExpWeight,d::Domain) = ExpWeight(setdomain(sp.exponent,d),setdomain(sp.space,d))

for (RULE,TYP) in ((:conversion_rule,:conversion_type),(:maxspace_rule,:maxspace))
    @eval function $RULE(a::ExpWeight,b::ExpWeight)
        if a.exponent ≈ b.exponent
            sp=$TYP(a.space,b.space)
            if !isa(sp,NoSpace)
                return ExpWeight(a.exponent,sp)
            end
        end
        NoSpace()
    end
end

function Derivative(wsp::ExpWeight)
    q=wsp.exponent
    sp=wsp.space
    D=Derivative(sp) + q'
    DerivativeWrapper(SpaceOperator(D,wsp,ExpWeight(q,rangespace(D))),1)
end

function Derivative(S::ExpWeight,k::Integer)
    if k==1
        Derivative(S)
    else
        D=Derivative(S)
        DerivativeWrapper(TimesOperator(Derivative(rangespace(D),k-1),D),k)
    end
end

function Multiplication(f::Fun,sp::ExpWeight)
    M=Multiplication(f,sp.space)
    MultiplicationWrapper(f,SpaceOperator(M,sp,ExpWeight(sp.exponent,rangespace(M))))
end

function Conversion(a::ExpWeight,b::ExpWeight)
    @assert a.exponent ≈ b.exponent
    ConversionWrapper(SpaceOperator(Conversion(a.space,b.space),a,b))
end


"""
## loc_to_glob!

This subroutine transforms the local end reactions and
moments into the global system (3-d).

### Function
```julia
loc_to_glob!(loc, glob, gamma, coord)
```

### Arguments
```julia
* loc::Vector{Float64}       : Local force and momemts (Updated)
* glob::Vector{Float64}      : Globale forces and moments
* gamma::Float64             : Element orientation angle (3D)
* coord::Matrix{Float64}     : Nodal coordinates
```
"""
function loc_to_glob!(loc::Vector{Float64}, glob::Vector{Float64},
  gamma::Float64, coord::LinearAlgebra.Adjoint{Float64,Array{Float64,2}})
  #
  # This subroutine transforms the local end reactions and
  # moments into the global system (3-d).
  #
  # Updates glob[:]
  #
  t = zeros(12, 12)
  r0 = zeros(3, 3)

  x1 = coord[1,1]
  y1 = coord[1,2]
  z1 = coord[1,3]
  x2 = coord[2,1]
  y2 = coord[2,2]
  z2 = coord[2,3]
  xl = x2 - x1
  yl = y2 - y1
  zl = z2 - z1
  ell = sqrt(xl*xl+yl*yl+zl*zl)
  gamrad = gamma*pi/180.0
  cg = cos(gamrad)
  sg = sin(gamrad)
  den = ell*sqrt(xl*xl+zl*zl)
  if den !== 0.0
   r0[1,1] = xl/ell
   r0[2,1] = yl/ell
   r0[3,1] = zl/ell
   r0[1,2] = (-xl*yl*cg-ell*zl*sg)/den
   r0[2,2] = den*cg/(ell*ell)
   r0[3,2] = (-yl*zl*cg+ell*xl*sg)/den
   r0[1,3] = (xl*yl*sg-ell*zl*cg)/den
   r0[2,3] = -den*sg/(ell*ell)
   r0[3,3] = (yl*zl*sg+ell*xl*cg)/den
  else
   r0[1,1] = 0.0
   r0[3,1] = 0.0
   r0[2,2] = 0.0
   r0[2,3] = 0.0
   r0[2,1] = 1.0
   r0[1,2] = -cg
   r0[3,3] = cg
   r0[3,2] = sg
   r0[1,3] = sg
  end
  for i in 1:3
   for j in 1:3
     x = r0[i,j]
     for k in 0:3:9
       t[i+k, j+k] = x
     end
   end
  end
  for i in 1:12
   sum = 0.0
   for j in 1:12
     sum = sum+t[i,j]*loc[j]
   end
   glob[i] = sum
  end
  glob
end


using FFTW, LinearAlgebra, Plots
gr()

include("src/dataset.jl")
include("src/micmac.jl")

const epsilons = [10.0^(-i) for i in 0:6]
const xmin     = 0
const xmax     = 2π
const T        = 2π
const size_x   = [64]
const size_tau = [64]
const Tfinal   = 0.25
const nb_dt    = 5 # different values of dt

tabdt  = zeros(Float64, nb_dt)
taberr = zeros(Float64, (length(epsilons), nb_dt))

etime = @elapsed for nx in size_x, ntau in size_tau

    println(" nx                 : $nx ")
    println(" ntau               : $ntau ")

    for (kk, epsilon) in enumerate(epsilons)

        print(" $epsilon: ")

        data = CosSin(xmin, xmax, nx, epsilon, T, Tfinal)

        for hh in 0:nb_dt-1

            println("dt = $(2.0^(-hh)) ")

            dtmicmac = 2.0^(-hh) * data.Tfinal / 16

            solver = MicMac(data, ntau)

            u, v = solve(solver, dtmicmac)

            tabdt[hh+1] = dtmicmac

            err = compute_error(u,v,data)
            println(" error = $err ")
            taberr[kk,hh+1] = err

        end

        println()

    end

end

println("Elapsed time :", etime)

p = plot(layout=(1,2))
xlabel!(p[1,1],"dt")
ylabel!(p[1,1],"error")

for j in 1:size(taberr)[1]
    plot!(p[1,1],
          tabdt, taberr[j, :],
          markershape=:circle,
          xaxis=:log,
          yaxis=:log,
          label="epsilon=$(epsilons[j])")
end

for j in 1:size(taberr)[2]
    plot!(p[1,2],
          epsilons, taberr[:, j],
          markershape=:circle,
          xaxis=:log,
          yaxis=:log,
          label="dt=$(tabdt[j])")
end

xlabel!(p[1,2],"epsilon")

p


using Test
using SerialDependence
using DelimitedFiles

#testing the functions by reproducing the results from C. Weiss's book "An Introduction to Discrete-Valued Time Series".
#I could not find the raw data so I copied the values provided in the book (the raw time-series should be ~1300 elements long), explaining slight numerical differences in results.
pewee = Int64.(readdlm("pewee.txt", ',')[1,:])
@test round(cramer_coefficient(pewee, Int64(4))[1], digits = 2) == 0.46
@test round(cohen_coefficient(pewee, Int64(4))[1], digits = 2) == 0.55
# this one was not in the book, but the plot correctly reproduces the results.
@test round(theils_u(pewee,[1])[1], digits = 2) == 0.55
@test round(H(pewee), digits = 2) == 1.49


const BigFloats = Union{BigFloat,Complex{BigFloat}}

function fft(x::AbstractVector{F}) where F<:Fun
    n,T = length(x),mapreduce(eltype,promote_type,x)
    if ispow2(n) return fft_pow2(x) end
    ks = range(zero(real(T)),stop=n-one(real(T)),length=n)
    Wks = exp(-im*convert(T,π)*ks.^2/n)
    xq,wq = x.*Wks,conj([exp(-im*convert(T,π)*n);reverse(Wks);Wks[2:end]])
    return Wks.*conv(xq,wq)[n+1:2n]
end

ifft(x::AbstractVector{F}) where {F<:Fun} = conj(fft(conj(x)))/length(x)
function ifft!(x::AbstractVector{F}) where F<:Fun
    y = conj(fft(conj(x)))/length(x)
    x[:] = y
    return x
end

function conv(u::StridedVector{F}, v::StridedVector) where F<:Fun
    nu,nv = length(u),length(v)
    n = nu + nv - 1
    np2 = nextpow2(n)
    pad!(u,np2),pad!(v,np2)
    y = ifft_pow2(fft_pow2(u).*fft_pow2(v))
    #TODO This would not handle Dual/ComplexDual numbers correctly
    T = promote_type(mapreduce(eltype,promote_type,u),mapreduce(eltype,promote_type,v))
    y = T<:Real ? real(y[1:n]) : y[1:n]
end

######################################################################
# TO BE DEPRECATED FOR V0.4 UPGRADE
######################################################################

# plan_fft for BigFloats (covers Laurent svfft)

plan_fft(x::Vector{F}) where {F<:Fun} = fft
plan_ifft(x::Vector{F}) where {F<:Fun} = ifft
plan_ifft!(x::Vector{F}) where {F<:Fun} = ifft

# Chebyshev transforms and plans for BigFloats
# no plan exists at the moment, so we make a dummy plan
plan_chebyshevtransform!(x::Vector{T}, kind) where {T<:BigFloats} =
    error("In-place variant not implemented for BigFloat")

plan_ichebyshevtransform!(x::Vector{T}, kind) where {T<:BigFloats} =
    error("In-place variant not implemented for BigFloat")

plan_chebyshevtransform(x::AbstractVector{T}, ::Val{kind}) where {T<:BigFloats,kind} =
    ChebyshevTransformPlan{T,kind,false,Nothing}()
plan_ichebyshevtransform(x::AbstractVector{T}, ::Val{kind}) where {T<:BigFloats,kind} =
    IChebyshevTransformPlan{T,kind,false,Nothing}()

#following Chebfun's @Chebtech1/vals2coeffs.m and @Chebtech2/vals2coeffs.m
function *(P::ChebyshevTransformPlan{T,1,false,Nothing}, x::AbstractVector{T}) where T<:BigFloats
    n = length(x)
    if n == 1
        x
    else
        w = [2exp(im*convert(T,π)*k/2n) for k=0:n-1]
        ret = w.*ifft([x;reverse(x)])[1:n]
        ret = T<:Real ? real(ret) : ret
        ret[1] /= 2
        ret
    end
end

function *(P::ChebyshevTransformPlan{T,2,false}, x::AbstractVector{T}) where T<:BigFloats
    n = length(x)
    if n == 1
        x
    else
        ret = ifft([x;x[end:-1:2]])[1:n]
        ret = T<:Real ? real(ret) : ret
        ret[2:n-1] *= 2
        ret
    end
end

#following Chebfun's @Chebtech1/vals2coeffs.m and @Chebtech2/vals2coeffs.m
function *(P::IChebyshevTransformPlan{T,1,false,Nothing}, x::AbstractVector{T}) where T<:BigFloats
    n = length(x)
    if n == 1
        x
    else
        w = [exp(-im*convert(T,π)*k/2n)/2 for k=0:2n-1]
        w[1] *= 2;w[n+1] *= 0;w[n+2:end] *= -1
        ret = fft(w.*[x;one(T);x[end:-1:2]])
        ret = T<:Real ? real(ret) : ret
    end
end
function *(P::IChebyshevTransformPlan{T,2,false}, x::AbstractVector{T}) where T<:BigFloats
    n = length(x)
    if n == 1
        x
    else
        ##TODO: make thread safe
        x[1] *= 2;x[end] *= 2
        ret = chebyshevtransform(x, Val(2))
        x[1] /=2;x[end] /=2
        ret[1] *= 2;ret[end] *= 2
        ret *= .5*(n-1)
        ret
    end
end

# Fourier space plans for BigFloat

plan_transform(sp::Fourier{D,R},x::AbstractVector{T}) where {T<:BigFloat,D,R} =
    TransformPlan{T,typeof(sp),false,Nothing}(sp,nothing)
plan_itransform(sp::Fourier{D,R},x::AbstractVector{T}) where {T<:BigFloat,D,R} =
    ITransformPlan{T,typeof(sp),false,Nothing}(sp,nothing)

function *(::TransformPlan{T,Fourier{D,R},false},x::AbstractVector{T}) where {T<:BigFloat,D,R}
    l = length(x); n = div(l+1,2)
    v = fft(x)
    rmul!(v,convert(T,2)/l)
    v[1] /= 2
    mod(l,2) == 1 ? ApproxFunBase.interlace(real(v[1:n]),-imag(v[2:n])) :
      [ApproxFunBase.interlace(real(v[1:n]),-imag(v[2:n]));-real(v[n+1])/2]
end

function *(::ITransformPlan{T,Fourier{D,R},false},x::AbstractVector{T}) where {T<:BigFloat,D,R}
    l = length(x); n = div(l+1,2)
    # v = complex([x[1:n];x[n-1:-1:2]],[0;-x[2n-2:-1:n+1];0;x[n+1:2n-2]])
    v = mod(l,2) == 1 ?
        complex([2x[1];x[3:2:end];x[end:-2:3]],[0;-x[2:2:end];x[end-1:-2:2]]) :
        complex([2x[1];x[3:2:end-1];-2x[end];x[end-1:-2:3]],[0;-x[2:2:end];x[end-2:-2:2]])
    rmul!(v,convert(T,l)/2)
     real(ifft(v))
end

# SinSpace plans for BigFloat

plan_transform(sp::SinSpace{D,R},x::AbstractVector{T}) where {T<:BigFloat,D,R} =
    TransformPlan{T,typeof(sp),false,Nothing}(sp,nothing)
plan_itransform(sp::SinSpace{D,R},x::AbstractVector{T}) where {T<:BigFloat,D,R} =
    ITransformPlan{T,typeof(sp),false,Nothing}(sp,nothing)

function *(::TransformPlan{T,SinSpace{D,R},false},x::AbstractVector{T}) where {T<:BigFloat,D,R}
    v=imag(fft([0;-x;0;reverse(x)]))[2:length(x)+1]
    rmul!(v,convert(T,1)/(length(x)+1))
    v
end

*(::ITransformPlan{T,SinSpace{D,R},false},x::AbstractVector{T}) where {T<:BigFloat,D,R} =
    imag(fft([0;-x;0;reverse(x)]))[2:length(x)+1]/2

# Fourier space & SinSpace plans for Complex{BigFloat}

for SP in (:Fourier,:SinSpace), (pl,TransPlan) in ((:plan_transform,:TransformPlan),
                                                    (:plan_itransform,:ITransformPlan))
    @eval begin
        $pl(sp::$SP{D,R},x::AbstractVector{T}) where {T<:Complex{BigFloat},D,R} =
                $TransPlan(sp,$pl(sp,Array{T}(undef, length(x))),Val{false})
        *(P::$TransPlan{T,$SP{D,R},false},x::Vector{T}) where {T<:Complex{BigFloat},D,R} =
            complex(P.plan*real(x),P.plan*imag(x))
    end
end


using ConeProgramDiff

using MathOptInterface
using SCS
using Random
using LinearAlgebra

const MOI = MathOptInterface

function check_kkt(A, b, c, cone_prod, x, y, s; tol=1e-8)
    try
        @assert isapprox(A*x + s, b,atol=tol)
        @assert isapprox(A'*y + c, zeros(length(c)),atol=tol)
        @assert isapprox(s, project_onto_cone(s, cone_prod),atol=tol)
        @assert isapprox(y, project_onto_cone(y, [MOI.dual_set(c) for c in cone_prod]),atol=tol)
        @assert isapprox(s'*y, 0, atol=tol)
    catch e
        error("KKT conditions failed:
               \nprimal constraint: $(norm(A*x + s - b))
               \ndual constraint: $(norm(A'*y +c))
               \nK error: $(norm(s - project_onto_cone(s, cone_prod)))
               \nK* error: $(norm(y - project_onto_cone(y, [MOI.dual_set(c) for c in cone_prod])))
               \nCS cond: $(norm(s'*y))"
        )
    end
end

function check_adjoint(tol)
    Random.seed!(0)
    dims = (15, 10)
    A, b, c, cone_prod = l1_minimization_program(dims; solve=false)
    x, y, s, pf, pb = solve_and_diff(A, b, c, cone_prod; eps=1e-10, max_iters=10000)
    check_kkt(A, b, c, cone_prod, x, y, s, tol=1e-8)

    # Check adjoint
    pstar = dot(x, c)
    dA, db, dc = pb(c)
    del = 1e-6
    x_pert, _, _, _, _ = solve_and_diff(A+del*dA, b+del*db, c+del*dc, cone_prod;
                                                  eps=1e-10, max_iters=10000)
    pstar_pert = dot(x_pert, c)
    df = pstar_pert - pstar
    dp = del * (sum(dA .* dA) + db'*db + dc'*dc)
    if !isapprox(df, dp, atol=tol)
        error("Adjoint mismatch: $(abs(df-dp))\n\tdf = $df\n\tdp = $dp")
    end
    return true
end

function check_derivative()
    Random.seed!(0)
    dims = (15, 10)
    A, b, c, cone_prod = l1_minimization_program(dims; solve=false)
    x, y, s, pf, pb = solve_and_diff(A, b, c, cone_prod; eps=1e-10, max_iters=10000)
    check_kkt(A, b, c, cone_prod, x, y, s, tol=1e-6)

    # Check derivative
    del = 1e-6
    dA, db, dc = del*randn(size(A)), del*randn(size(b)), del*randn(size(c))
    dx, dy, ds = pf(dA, db, dc)
    x_pert, y_pert, s_pert, _, _ = solve_and_diff(A+dA, b+db, c+dc, cone_prod;
                                                  eps=1e-10, max_iters=10000)

    # TODO: tols are high here -- should investigate
    @assert isapprox(x_pert - x, dx, atol=1e-4)
    @assert isapprox(y_pert - y, dy, atol=1e-4)
    @assert isapprox(s_pert - s, ds, atol=2e-4)
    return true
end

@testset "solve_and_diff Tests" begin
    @test check_adjoint(2e-5)
    @test check_derivative()
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    intersecttype(l1, l2)

Compute the intersection type of two lines `l1` and `l2`.

The intersection type can be one of three types:

1. intersect at one point
2. overlap at more than one point
3. do not overlap nor intersect
"""
function intersecttype(l1::Line{3,T}, l2::Line{3,T}) where {T}
  a, b = l1(0), l1(1)
  c, d = l2(0), l2(1)

  if measure(Tetrahedron(a, b, c, d)) > 0
    NoIntersection()
  elseif isapprox(abs((b - a) × (c - d)), zero(T), atol=atol(T)^2)
    OverlappingLines(l1)
  else
    CrossingLines(intersectpoint(l1, l2))
  end
end

function intersecttype(l1::Line{2,T}, l2::Line{2,T}) where {T}
  a, b = l1(0), l1(1)
  c, d = l2(0), l2(1)

  if !isapprox(abs((b - a) × (c - d)), zero(T), atol=atol(T)^2)
    CrossingLines(intersectpoint(l1, l2))
  elseif isapprox(measure(Triangle(a, b, c)), zero(T), atol=atol(T)^2)
    OverlappingLines(l1)
  else
    NoIntersection()
  end
end

# compute the intersection of two lines assuming that it is a point
function intersectpoint(l1::Line{2}, l2::Line{2})
  a, b = l1(0), l1(1)
  c, d = l2(0), l2(1)

  v1  = a - b
  v2  = c - d
  v12 = a - c

  # the intersection point lies in between a and b at a fraction s
  # (https://en.wikipedia.org/wiki/Line-line_intersection#Formulas)
  s = (v12[1] * v2[2] - v12[2] * v2[1]) / (v1[1] * v2[2] - v1[2] * v2[1])

  a - s*v1
end


function colmeans(mat::Array{T,2})::Array{Float64,1} where {T<:Number}
    _, p = size(mat)
    means = Array{Float64,1}(undef, p)
    for i = 1:p
        means[i] = Statistics.mean(mat[:, i])
    end
    return means
end


# Text provided under a Creative Commons Attribution license, CC-BY.  All code is made available under the FSF-approved BSD-3 license.  (c) Lorena A. Barba, Gilbert F. Forsyth 2017. Thanks to NSF for support via CAREER award #1149784.
# [@LorenaABarba](https://twitter.com/LorenaABarba)
# Translation to Julia by [@miguelraz_](https://twitter.com/miguelraz_)

## 12 steps to Navier–Stokes
# ======
# ***

# You should have completed Steps [1](./01_Step_1.jl) and [2](./02_Step_2.jl) before continuing.
# This Jupyter notebook continues the presentation of the **12 steps to Navier–Stokes**,
# the practical module taught in the interactive CFD class of [Prof. Lorena Barba](http://lorenabarba.com).

## Step 3: Diffusion Equation in 1-D
# -----
# ***

# The one-dimensional diffusion equation is:

# $$\frac{\partial u}{\partial t}= \nu \frac{\partial^2 u}{\partial x^2}$$

# The first thing you should notice is that —unlike the previous two simple equations we have studied— this equation has a second-order derivative. We first need to learn what to do with it!

# ### Discretizing $\frac{\partial ^2 u}{\partial x^2}$

# The second-order derivative can be represented geometrically as the line tangent to the curve given by the first derivative.  We will discretize the second-order derivative with a Central Difference scheme: a combination of Forward Difference and Backward Difference of the first derivative.  Consider the Taylor expansion of $u_{i+1}$ and $u_{i-1}$ around $u_i$:

# $u_{i+1} = u_i + \Delta x \frac{\partial u}{\partial x}\bigg|_i + \frac{\Delta x^2}{2} \frac{\partial ^2 u}{\partial x^2}\bigg|_i + \frac{\Delta x^3}{3!} \frac{\partial ^3 u}{\partial x^3}\bigg|_i + O(\Delta x^4)$

# $u_{i-1} = u_i - \Delta x \frac{\partial u}{\partial x}\bigg|_i + \frac{\Delta x^2}{2} \frac{\partial ^2 u}{\partial x^2}\bigg|_i - \frac{\Delta x^3}{3!} \frac{\partial ^3 u}{\partial x^3}\bigg|_i + O(\Delta x^4)$

# If we add these two expansions, you can see that the odd-numbered derivative terms will cancel each other out.  If we neglect any terms of $O(\Delta x^4)$ or higher (and really, those are very small), then we can rearrange the sum of these two expansions to solve for our second-derivative.

# $u_{i+1} + u_{i-1} = 2u_i+\Delta x^2 \frac{\partial ^2 u}{\partial x^2}\bigg|_i + O(\Delta x^4)$

# Then rearrange to solve for $\frac{\partial ^2 u}{\partial x^2}\bigg|_i$ and the result is:

# $$\frac{\partial ^2 u}{\partial x^2}=\frac{u_{i+1}-2u_{i}+u_{i-1}}{\Delta x^2} + O(\Delta x^2)$$

# ### Back to Step 3

# We can now write the discretized version of the diffusion equation in 1D:

# $$\frac{u_{i}^{n+1}-u_{i}^{n}}{\Delta t}=\nu\frac{u_{i+1}^{n}-2u_{i}^{n}+u_{i-1}^{n}}{\Delta x^2}$$

# As before, we notice that once we have an initial condition, the only unknown is $u_{i}^{n+1}$, so we re-arrange the equation solving for our unknown:

# $$u_{i}^{n+1}=u_{i}^{n}+\frac{\nu\Delta t}{\Delta x^2}(u_{i+1}^{n}-2u_{i}^{n}+u_{i-1}^{n})$$

# The above discrete equation allows us to write a program to advance a solution in time. But we need an initial condition. Let's continue using our favorite: the hat function. So, at $t=0$, $u=2$ in the interval $0.5\le x\le 1$ and $u=1$ everywhere else. We are ready to number-crunch!

nₓ = 41 # n\_<TAB>x
Δx = 2 / (nₓ - 1)
nₜ = 20    #nₜ is the number of timesteps we want to calculate
nᵤ = 0.3
σ = 0.2
Δt = σ * Δx^2 / nᵤ  #Δt is the amount of time each timestep covers (delta t)

u = ones(nₓ)      #as before, we initialize u with every value equal to 1.
indexes = Int.((.5/Δx):(1/Δx))
u[indexes] .= 2 #then set u = 2 between 0.5 and 1 as per our I.C.s

uₙ = ones(nₓ) #initialize our placeholder array un, to hold the time-stepped solution

u = numpy.ones(nx)      #a numpy array with nx elements all equal to 1.
u[int(.5 / dx):int(1 / dx + 1)] = 2  #setting u = 2 between 0.5 and 1 as per our I.C.s

for n in nₜ  #iterate through time
    uₙ = copy(u) ##copy the existing values of u into uₙ
    for i in 1:nₓ  ##now we'll iterate through the u array

     ###This is the line from Step 1, copied exactly.  Edit it for our new equation.
     ###then uncomment it and run the cell to evaluate Step 2
           ###u[i] = uₙ[i] - c * dt / dx * (uₙ[i] - uₙ[i-1])
    u[i] = uₙ[i] + nu * Δt / Δx^2 * (uₙ[i+1] - 2 * uₙ[i] + uₙ[i-1])
    end
end

plot(range(start = 0, stop = 2, length = nₓ), u);

## Learn More

# For a careful walk-through of the discretization of the diffusion equation with finite differences (and all steps from 1 to 4), watch **Video Lesson 4** by Prof. Barba on YouTube.

# TODO
# YouTubeVideo('y2WaK7_iMRI')


################################################################################
#                                 MgO:LiNbO₃                                   #
################################################################################
export MgO_LiNbO₃

"""
These functions create a symbolic representation (using ModelingToolkit) of the
Sellmeier Equation model for the temperature and wavelength dependence
of 5% MgO:LiNbO3's (congruent LiNbO3 (CLN) not stoichiometric LiNbO3 (SLN))
ordinary and extraordinary indices of refraction.
Equation form is based on "Temperature and
wavelength dependent refractive index equations for MgO-doped congruent
and stoichiometric LiNbO3" by Gayer et al., Applied Physics B 91,
343?348 (2008)
This model is then exported to other functions that use it and its
derivatives to return index, group index and GVD values as a function
of temperature and wavelength.
Variable units are lm in [um] and T in [deg C]
"""
function n²_MgO_LiNbO₃_sym(λ, T; a₁, a₂, a₃, a₄, a₅, a₆, b₁, b₂, b₃, b₄, T₀)
    f = (T - T₀) * (T + T₀ + 2*273.16)  # so-called 'temperature dependent parameter'
    λ² = λ^2
    a₁ + b₁*f + (a₂ + b₂*f) / (λ² - (a₃ + b₃*f)^2) + (a₄ + b₄*f) / (λ² - a₅^2) - a₆*λ²
end

pₑ = (
    a₁ = 5.756,
    a₂ = 0.0983,
    a₃ = 0.202,
    a₄ = 189.32,
    a₅ = 12.52,
    a₆ = 1.32e-2,
    b₁ = 2.86e-6,
    b₂ = 4.7e-8,
    b₃ = 6.113e-8,
    b₄ = 1.516e-4,
    T₀ = 24.5,      # reference temperature in [Deg C]
)
pₒ = (
    a₁ = 5.653,
    a₂ = 0.1185,
    a₃ = 0.2091,
    a₄ = 89.61,
    a₅ = 10.85,
    a₆ = 1.97e-2,
    b₁ = 7.941e-7,
    b₂ = 3.134e-8,
    b₃ = -4.641e-9,
    b₄ = -2.188e-6,
    T₀ = 24.5,      # reference temperature in [Deg C]
)

pᵪ₂ = (
	d₃₃ =   20.3,    #   pm/V
	d₃₁ =   -4.1,    #   pm/V
	d₂₂ =   2.1,     #   pm/V
	λs  =  [1.313, 1.313, 1.313/2.0]
)

function make_MgO_LiNbO₃(;pₒ=pₒ,pₑ=pₑ,pᵪ₂=pᵪ₂)
	@variables λ, T, λs[1:3]
	nₒ² = n²_MgO_LiNbO₃_sym(λ, T; pₒ...)
	nₑ² = n²_MgO_LiNbO₃_sym(λ, T; pₑ...)
	ε 	= diagm([nₒ², nₒ², nₑ²])
	d₃₃, d₃₁, d₂₂, λᵣs = pᵪ₂
	χ⁽²⁾ᵣ = cat(
		[ 	0.0	 	-d₂₂ 	d₃₁			#	xxx, xxy and xxz
		 	-d₂₂	0.0 	0.0			#	xyx, xyy and xyz
			d₃₁	 	0.0		0.0		],	#	xzx, xzy and xzz
		[ 	-d₂₂	0.0 	0.0			#	yxx, yxy and yxz
			0.0	 	d₂₂ 	d₃₁			#	yyx, yyy and yyz
			0.0	 	d₃₁		0.0		],	#	yzx, yzy and yzz
		[ 	d₃₁	 	0.0 	0.0			#	zxx, zxy and zxz
			0.0	 	d₃₁ 	0.0			#	zyx, zyy and zyz
			0.0	 	0.0 	d₃₃		],	#	zzx, zzy and zzz
		 dims = 3
	)
	nₒ = sqrt(nₒ²)
	ngₒ = ng_model(nₒ,λ)
	gvdₒ = gvd_model(nₒ,λ)
	nₑ = sqrt(nₑ²)
	ngₑ = ng_model(nₑ,λ)
	gvdₑ = gvd_model(nₑ,λ)
	models = Dict([
		:nₒ		=>	nₒ,
		:ngₒ	=>	ngₒ,
		:gvdₒ	=>	gvdₒ,
		:nₑ		=>	nₑ,
		:ngₑ	=>	ngₑ,
		:gvdₑ	=>	gvdₑ,
		:ng		=>	diagm([ngₒ, ngₒ, ngₑ]),
		:gvd	=>	diagm([gvdₒ, gvdₒ, gvdₑ]),
		:ε 		=> 	ε,
		:χ⁽²⁾	=>	SArray{Tuple{3,3,3}}(Δₘ(λs,ε, λᵣs, χ⁽²⁾ᵣ)),
	])
	defaults =	Dict([
		:λ		=>		0.8,	# μm
		:T		=>		24.5,	# °C
		:λs₁	=>		1.064,	# μm
		:λs₂	=>		1.064,	# μm
		:λs₃	=>		0.532,	# μm

	])
	Material(
		models,
		defaults,
		:MgO_LiNbO₃,
		colorant"seagreen2", #RGB{N0f8}(0.22,0.596,0.149),	#
		)
end

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

MgO_LiNbO₃ = make_MgO_LiNbO₃()


using SpecialFunctions: lgamma
import SpecialFunctions: lbeta, beta

# Call back functions for `lbeta` and `beta`
# Those in `SpecialFunctions.jl` are implemented for `(::Number, ::Number)` only,
# for which `KnetArray` cannot broadcast through.
# NOTE: `_lbeta` and `_beta` are implemented for explictly test purpose.
# TODO: implement this broadcasting in Knet.jl
_lbeta(α, β) = lgamma.(α) + lgamma.(β) - lgamma.(α + β)
_beta(α, β) = exp.(_lbeta(α, β))
lbeta(α, β) = _lbeta(α, β)
beta(α, β) = _beta(α, β)


abstract type SphericalHarmonicPlan{T} end

function *(P::SphericalHarmonicPlan, X::AbstractMatrix)
    mul!(zero(X), P, X)
end

if VERSION < v"0.7-"
    function \(P::SphericalHarmonicPlan, X::AbstractMatrix)
        At_mul_B!(zero(X), P, X)
    end
else
    function \(P::SphericalHarmonicPlan, X::AbstractMatrix)
        mul!(zero(X), transpose(P), X)
    end
end

include("sphfunctions.jl")
include("slowplan.jl")
include("Butterfly.jl")
include("fastplan.jl")
include("thinplan.jl")
include("synthesisanalysis.jl")
include("vectorfield.jl")

function plan_sph2fourier(A::AbstractMatrix; opts...)
    M, N = size(A)
    if M ≤ 1023
        SlowSphericalHarmonicPlan(A)
    else
        ThinSphericalHarmonicPlan(A; opts...)
    end
end

sph2fourier(A::AbstractMatrix; opts...) = plan_sph2fourier(A; opts...)*A
fourier2sph(A::AbstractMatrix; opts...) = plan_sph2fourier(A; opts...)\A

"""
Computes the bivariate Fourier series given by the spherical harmonic expansion:

```math
{\\rm SHT} : \\sum_{\\ell=0}^n\\sum_{m=-\\ell}^{\\ell} f_{\\ell}^m Y_{\\ell}^m(\\theta,\\varphi) \\to \\sum_{\\ell=0}^n\\sum_{m=-n}^{n} g_{\\ell}^m \\left\\{\\begin{array}{ccc}\\cos\\ell\\theta & {\\rm for} & m{\\rm~even}\\\\ \\sin(\\ell+1)\\theta & {\\rm for} & m{\\rm~odd}\\end{array}\\right\\}\\times \\sqrt{\\frac{2-\\delta_{m,0}}{2\\pi}} \\left\\{\\begin{array}{ccc} \\cos m\\varphi & {\\rm for} & m \\ge 0,\\\\ \\sin(-m\\varphi) & {\\rm for} & m < 0.\\end{array}\\right.
```

The spherical harmonic expansion coefficients are organized as follows:

```math
F = \\begin{pmatrix}
f_0^0 & f_1^{-1} & f_1^1 & f_2^{-2} & f_2^2 & \\cdots & f_n^{-n} & f_n^n\\\\
f_1^0 & f_2^{-1} & f_2^1 & f_3^{-2} & f_3^2 & \\cdots & 0 & 0\\\\
\\vdots & \\vdots & \\vdots &  \\vdots &  \\vdots & \\ddots & \\vdots & \\vdots\\\\
f_{n-2}^0 & f_{n-1}^{-1} & f_{n-1}^1 & f_n^{-2} & f_n^2 &  & \\vdots & \\vdots\\\\
f_{n-1}^0 & f_n^{-1} & f_n^1 & 0 & 0 & \\cdots & 0 & 0\\\\
f_n^0 & 0 & 0 & 0 & 0 & \\cdots & 0 & 0\\\\
\\end{pmatrix},
```

and the Fourier coefficients are organized similarly:

```math
G = \\begin{pmatrix}
g_0^0 & g_0^{-1} & g_0^1 & \\cdots & g_0^{-n} & g_0^n\\\\
g_1^0 & g_1^{-1} & g_1^1 & \\cdots & g_1^{-n} & g_1^n\\\\
\\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots\\\\
g_{n-1}^0 & g_{n-1}^{-1} & g_{n-1}^1& \\cdots & g_{n-1}^{-n} & g_{n-1}^n\\\\
g_n^0 & 0 & 0 & \\cdots & g_n^{-n} & g_n^n\\\\
\\end{pmatrix}.
```
"""
sph2fourier(::AbstractMatrix; opts...)

"""
Computes the spherical harmonic expansion given by the bivariate Fourier series:

```math
{\\rm iSHT} : \\sum_{\\ell=0}^n\\sum_{m=-n}^{n} g_{\\ell}^m \\left\\{\\begin{array}{ccc}\\cos\\ell\\theta & {\\rm for} & m{\\rm~even}\\\\ \\sin(\\ell+1)\\theta & {\\rm for} & m{\\rm~odd}\\end{array}\\right\\}\\times \\sqrt{\\frac{2-\\delta_{m,0}}{2\\pi}} \\left\\{\\begin{array}{ccc} \\cos m\\varphi & {\\rm for} & m \\ge 0,\\\\ \\sin(-m\\varphi) & {\\rm for} & m < 0,\\end{array}\\right. \\to \\sum_{\\ell=0}^n\\sum_{m=-\\ell}^{\\ell} f_{\\ell}^m Y_{\\ell}^m(\\theta,\\varphi).
```

The spherical harmonic expansion coefficients are organized as follows:

```math
F = \\begin{pmatrix}
f_0^0 & f_1^{-1} & f_1^1 & f_2^{-2} & f_2^2 & \\cdots & f_n^{-n} & f_n^n\\\\
f_1^0 & f_2^{-1} & f_2^1 & f_3^{-2} & f_3^2 & \\cdots & 0 & 0\\\\
\\vdots & \\vdots & \\vdots &  \\vdots &  \\vdots & \\ddots & \\vdots & \\vdots\\\\
f_{n-2}^0 & f_{n-1}^{-1} & f_{n-1}^1 & f_n^{-2} & f_n^2 &  & \\vdots & \\vdots\\\\
f_{n-1}^0 & f_n^{-1} & f_n^1 & 0 & 0 & \\cdots & 0 & 0\\\\
f_n^0 & 0 & 0 & 0 & 0 & \\cdots & 0 & 0\\\\
\\end{pmatrix},
```

and the Fourier coefficients are organized similarly:

```math
G = \\begin{pmatrix}
g_0^0 & g_0^{-1} & g_0^1 & \\cdots & g_0^{-n} & g_0^n\\\\
g_1^0 & g_1^{-1} & g_1^1 & \\cdots & g_1^{-n} & g_1^n\\\\
\\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots\\\\
g_{n-1}^0 & g_{n-1}^{-1} & g_{n-1}^1& \\cdots & g_{n-1}^{-n} & g_{n-1}^n\\\\
g_n^0 & 0 & 0 & \\cdots & g_n^{-n} & g_n^n\\\\
\\end{pmatrix}.
```
"""
fourier2sph(::AbstractMatrix; opts...)

"""
Pre-computes the spherical harmonic transform.
"""
plan_sph2fourier(::AbstractMatrix; opts...)


rosenbrock2d(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
RANDOM_SEED = 809230
rng = MersenneTwister(RANDOM_SEED)

@testset "SimulatedAnnealing" begin

    # * Single-run tests for user-defined functions
    @test begin
        dimension = 2
        ground_truth = 0.0
        val = SimulatedAnnealing(
            rosenbrock2d, -5.0, 5.0, dimension, rng;
            low_temp=5000
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end

    # * Single-run tests for benchmark functions
    @test begin
        dimension = 3
        ground_truth = 0.0
        val = SimulatedAnnealing(
            Rosenbrock(), -5.0, 5.0, dimension, rng;
            low_temp=20000
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end

    # * Tests for random seeds
    @test begin
        dimension = 3
        ground_truth = 0.0
        val = SimulatedAnnealing(
            Rosenbrock(), -5.0, 5.0, dimension, rng;
            low_temp=20000
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end
end

@testset "GeneralizedSimulatedAnnealing" begin
    # * Single-run tests for user-defined functions
    @test begin
        dimension = 2
        ground_truth = 0.0
        val = GeneralSimulatedAnnealing(
            rosenbrock2d, -5.0, 5.0, dimension, rng;
            low_temp=35000
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end

    # * Single-run tests for different Tsallis' parameters
    @test begin
        dimension = 2
        ground_truth = 0.0
        val = GeneralSimulatedAnnealing(
            rosenbrock2d,
            -5.0,
            5.0,
            dimension,
            rng;
            low_temp=30000,
            qv=2.5,
            qa=1.0000001
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end

    # * Single-run tests for benchmark functions
    @test begin
        dimension = 3
        ground_truth = 0.0
        val = GeneralSimulatedAnnealing(
            Ackley(),
            -3.0,
            3.0,
            dimension,
            rng;
            low_temp=10000
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end

    # * Test for random seeds
    @test begin
        dimension = 3
        ground_truth = 0.0
        val = GeneralSimulatedAnnealing(
            Ackley(),
            -3.0,
            3.0,
            dimension,
            rng;
            low_temp=10000
        )
        @show val
        isapprox(val.min, ground_truth, atol=1e-2)
    end
end


using ITensors
using ITensorUnicodePlots
using PastaQ
using Zygote

function ising(n; h)
  H = OpSum()
  for j in 1:(n - 1)
    H -= "Z", j, "Z", j + 1
  end
  for j in 1:n
    H += h, "X", j
  end
  return H
end

n = 30
i = [Index(2, "S=1/2") for j in 1:n];

Zp = MPS(i, "Zp");
Zm = MPS(i, "Zm");

maxlinkdim(Zp) == 1 # product state

ψ = (Zp + Zm)/√2;
maxlinkdim(ψ) == 2 # entangled state

inner(Zp, Zp) ≈ 1
inner(Zp, ψ) ≈ 1/√2

j = n ÷ 2
Xj = op("X", i[j])
XjZp = apply(Xj, Zp);
st = [k == j ? "Zm" : "Zp" for k in 1:n];
XjZp = MPS(i, st);

maxlinkdim(XjZp) == 1
inner(Zp, XjZp) == 0

h = 0.5 # field
H = MPO(ising(n; h), i);
maxlinkdim(H) == 3

Zp = MPS(i, "Zp");
inner(Zp', H, Zp) ≈ -(n-1)

# Optimize
# kwargs allow for approximate arithmatic
function minimize(f, ∂f, x; nsteps, γ, kwargs...)
  for n in 1:nsteps
  x = -(x, γ * ∂f(x); kwargs...)
  println("Step = $(n), f(x) = $(f(x))")
  end
  return x
end

# Optimize energy
function E(ψ)
  ψHψ = inner(ψ', H, ψ)
  ψψ = inner(ψ, ψ)
  return ψHψ / ψψ
end
∂E(ψ) = gradient(E, ψ)[1]

ψ0 = MPS(i, "Zp");
ψ = minimize(E, ∂E, ψ0; nsteps=50, γ=0.1, cutoff=1e-5);
ψ /= norm(ψ);

Edmrg, ψdmrg = dmrg(H, ψ0; nsweeps=10, cutoff=1e-5);

inner(ψ, ψdmrg)

maxlinkdim(ψ0) == 1
maxlinkdim(ψ) == 3
maxlinkdim(ψdmrg) == 3
E(ψ0) # ≈ -29
E(ψ) # ≈ -31.0317917
E(ψdmrg) # ≈ -31.0356110

Ry_layer(θ, i) = [op("Ry", i[j]; θ=θ[j]) for j in 1:n]
CX_layer(i) = [op("CX", i[j], i[j+1]) for j in 1:2:(n-1)]

function U(θ, i; nlayers)
  n = length(i)
  Uθ = ITensor[]
  for l in 1:(nlayers-1)
    θl = θ[(1:n) .+ l*n]
    Uθ = [Uθ; Ry_layer(θl, i)]
    Uθ = [Uθ; CX_layer(i)]
  end
  return Uθ
end

# Optimize energy
ψ0 = MPS(i, "Zp");
nlayers = 6
function E(θ)
  ψθ = apply(U(θ, i; nlayers),ψ0)
  return inner(ψθ',H,ψθ)
end
∂E(θ) = gradient(E, θ)[1]

θ0 = zeros(nlayers * n);
θ = minimize(E, ∂E, θ0; nsteps=20, γ=0.1);

ψθ = apply(U(θ, i; nlayers), ψ0);
maxlinkdim(ψ0) == 1
maxlinkdim(ψθ) == 3
E(θ0) # ≈ -29
E(θ) # ≈ -31.017062

# Optimize fidelity
ψ0 = MPS(i, "Zp");
nlayers = 6
function F(θ)
  ψθ = apply(U(θ, i; nlayers),ψ0)
  return -abs2(inner(ψ,ψθ))
end
∂F(θ) = gradient(F, θ)[1]

θ0 = zeros(nlayers * n);
θ = minimize(F, ∂F, θ0; nsteps=20, γ=0.1);

maxlinkdim(ψ0) == 1
maxlinkdim(ψθ) == 2
F(θ0) # ≈ -0.556066
F(θ) # ≈ -0.989336


using AlgebraPDF
using DelimitedFiles

using Plots
using LaTeXStrings
theme(:wong2, grid=false, frame=:box,
    guidefontvalign=:top, guidefonthalign=:right,
    foreground_color_legend = nothing,
    xlim=(:auto,:auto), ylim=(:auto,:auto),
    xlab=L"m(\Xi_c^+K^-)\,\,[\mathrm{GeV}]", ylab=L"\mathrm{Candidates}")
#
#
const xth = 2.96
const fitrange = (2.960, 3.180)

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

data = readdlm(joinpath("plots", "data.txt"))[:,1]
Nev = size(data,1)

# amplitude for the signal
Φ2(x) = sqrt(x-xth)
Γ(x,m,Γ₀) = Γ₀*Φ2(x)/Φ2(m)
breitwigner(x,m,Γ₀) = 1/(m^2-x^2-1im*m*Γ(x,m,Γ₀)) # complex function

# create
bw = FunctionWithParameters((x;p)->breitwigner(x,p.m,p.Γ), (m=3.05,Γ=0.02))
# check
bw(3.02) # works on a number
bw([3.05,3.065]) # works on a vector
pars(bw) # gives namedtuple
bw(3.02; p=(m=3.05,Γ=0.02)) # works parameters
bw(3.02, [3.05, 0.02]) #

# operators like `abs2`, `log` can be applied
bwsq = abs2(bw) # real function
plot(bwsq, fitrange...) # can plot, range is needed

# making PDF from the function just by normalizing
bwsq_norm = Normalized(bwsq, fitrange) # needs limits (ranges, support)
plot(bwsq_norm) # can plot

# fit would work already
testfit_summary = fit(bwsq_norm, data)
# same as
# fit(NegativeLogLikelihood(bwsq_norm, data))

# define a type SimpleBW and method `func` for dispatch
struct SimpleBW{P} <: AbstractFunctionWithParameters
    p::P
end
import AlgebraPDF:func
function func(bw::SimpleBW, x::NumberOrTuple; p=pars(bw))
    m,Γ = (getproperty(p,s) for s in keys(bw.p))
    breitwigner(x, m, Γ)
end

# create four instances of the same PDF with different parameters
signalamps = []
push!(signalamps, SimpleBW(( m1=3.00, Γ1=6.5e-3 )))
push!(signalamps, SimpleBW(( m2=3.05, Γ2=2.3e-3 )))
push!(signalamps, SimpleBW(( m3=3.06, Γ3=4.0e-3 )))
push!(signalamps, SimpleBW(( m4=3.09, Γ4=9.9e-3 )))
@time plot(); plot!.(abs2.(signalamps), fitrange...); plot!()

# |A|^2 * phase_space
phasespace = FunctionWithParameters((x;p)->Φ2(x), ∅)
signalpdfs = [Normalized(abs2(A)*phasespace, fitrange) for A in signalamps]
@time plot(); plot!.(signalpdfs, fill=0, α=0.3); plot!()

#  - BW S-wave
@makefuntype SimpleBWg(x;p) =
    1/(p.m0^2 - x^2 - 1im*p.g^2*Φ2(x))
#
signalpdf0 = Normalized(
        abs2(SimpleBWg((m0=2.95, g=0.01)))*phasespace,
        fitrange)
plot!(signalpdf0, fill=0, α=0.3)
#
backgrpdf = Normalized(phasespace, fitrange)
plot(backgrpdf)

#########################################################################
# full model
model0 =
    signalpdf0 * (f0=0.05Nev,) +
    signalpdfs[1] * (f1=0.25Nev,) + signalpdfs[2] * (f2=0.15Nev,) +
    signalpdfs[3] * (f3=0.15Nev,) + signalpdfs[4] * (f4=0.55Nev,) +
    backgrpdf * Ext(fb=0.3Nev,) # Ext - to be able to fix

model1 = fixpar(model0, :fb, 17.2)

@time fit_summary = fit(
    Extended(NegativeLogLikelihood(model1, data)))
#
let bins = range(fitrange..., length=80)
    Ns = scaletobinneddata(bins)
    stephist(data; bins, lc=:black, lab="Data")
    for i in 1:5
        plot!(fit_summary.best_model[i], Ns, 300,
            fillto=0.0, α=0.4, lab="Signal$(i)")
    end
    plot!(fit_summary.best_model, Ns, 300, lc=:red, lab="Total fit")
    #
    m = fit_summary.measurements
    for (i,(k,v)) in enumerate(zip(keys(m), m))
        annotate!(xth+0.01,25-1.2i,text("$k=$(string(v))",:left, 6))
    end
    plot!()
end
savefig(joinpath("plots", "sixcompfit.png"))


#######################################################################################
# Simplest Loop method
# Makes the simplest possible method for teaching and performance testing
#######################################################################################
struct LoopRK4 <: DiffEqBase.AbstractODEAlgorithm end
export LoopRK4

# Out-of-place
# No caching, good for static arrays, bad for arrays
@muladd function DiffEqBase.__solve(prob::ODEProblem{uType,tType,false},
                          alg::LoopRK4;
                          dt = error("dt is required for this algorithm"),
                          save_everystep = true,
                          save_start = true,
                          adaptive = false,
                          dense = false,
                          save_end = true,
                          kwargs...) where {uType,tType}
  @assert !adaptive
  @assert !dense
  u0 = prob.u0
  tspan = prob.tspan
  f = prob.f
  p = prob.p
  t = tspan[1]
  tf = prob.tspan[2]
  ts = tspan[1]:dt:tspan[2]

  if save_everystep && save_start
      us = Vector{typeof(u0)}(undef,length(ts))
      us[1] = u0
  elseif save_everystep
      us = Vector{typeof(u0)}(undef,length(ts)-1)
  elseif save_start
      us = Vector{typeof(u0)}(undef,2)
      us[1] = u0
  else
      us = Vector{typeof(u0)}(undef,1) # for interface compatibility
  end

  u = u0
  half = convert(eltype(u0),1//2)
  sixth = convert(eltype(u0),1//6)

  for i in 2:length(ts)
      uprev = u; t = ts[i]
      k1 = f(u,p,t)
      tmp = uprev+dt*half*k1
      k2 = f(tmp, p, t+half*dt)
      tmp = uprev+dt*half*k2
      k3 = f(tmp, p, t+half*dt)
      tmp = uprev+dt*k3
      k4 = f(tmp, p, t+dt)
      u = uprev+dt*sixth*(k1+2k2+2k3+k4)
      save_everystep && (us[i] = u)
  end

  !save_everystep && save_end && (us[end] = u)

  sol = DiffEqBase.build_solution(prob,alg,ts,us,
                                  k = nothing, destats = nothing,
                                  calculate_error = false)
  DiffEqBase.has_analytic(prob.f) && DiffEqBase.calculate_solution_errors!(sol;timeseries_errors=true,dense_errors=false)
  sol
end

# In-place
# Good for mutable objects like arrays
# Use DiffEqBase.@.. for simd ivdep
@muladd function DiffEqBase.solve(prob::ODEProblem{uType,tType,true},
                          alg::LoopRK4;
                          dt = error("dt is required for this algorithm"),
                          save_everystep = true,
                          save_start = true,
                          adaptive = false,
                          dense = false,
                          save_end = true,
                          kwargs...) where {uType,tType}
  @assert !adaptive
  @assert !dense
  u0 = prob.u0
  tspan = prob.tspan
  f = prob.f
  p = prob.p
  t = tspan[1]
  tf = prob.tspan[2]
  ts = tspan[1]:dt:tspan[2]

  if save_everystep && save_start
      us = Vector{typeof(u0)}(undef,length(ts))
      us[1] = u0
  elseif save_everystep
      us = Vector{typeof(u0)}(undef,length(ts)-1)
  elseif save_start
      us = Vector{typeof(u0)}(undef,2)
      us[1] = u0
  else
      us = Vector{typeof(u0)}(undef,1) # for interface compatibility
  end

  u = copy(u0); uprev = copy(u0)
  k1 = zero(u0); k2 = zero(u0); k3 = zero(u0); k4 = zero(u0)
  half = convert(eltype(u0),1//2)
  sixth = convert(eltype(u0),1//6)

  for i in 2:length(ts)
      uprev .= u; t = ts[i]
      f(k1,u,p,t)
      DiffEqBase.@.. u = uprev+dt*half*k1
      f(k2, u, p, t+half*dt)
      DiffEqBase.@.. u = uprev+dt*half*k2
      f(k3, u, p, t+half*dt)
      DiffEqBase.@.. u = uprev+dt*k3
      f(k4, u, p, t+dt)
      DiffEqBase.@.. u = uprev+dt*sixth*(k1+2k2+2k3+k4)
      save_everystep && (us[i] = copy(u))
  end

  !save_everystep && save_end && (us[end] = u)

  sol = DiffEqBase.build_solution(prob,alg,ts,us,
                                  k = nothing, destats = nothing,
                                  calculate_error = false)
  DiffEqBase.has_analytic(prob.f) && DiffEqBase.calculate_solution_errors!(sol;timeseries_errors=true,dense_errors=false)
  sol
end


using Test
include("simpson_integrate.jl")

testfunc(x) = 4.0x^2 + 2.0x - 1.0
stepsize = 1e-2
xrange = 0:stepsize:1
funcvals = testfunc.(xrange)
integral = simpson_integrate(funcvals,stepsize)
exactintegral = 4.0/3
@test integral ≈ exactintegral


"""
#
# PEJROOT calculates (multiple) roots of polynomial using an iterative
# method developed by Z. Zeng. For details, you can request the paper
#
#     "Computing Multiple roots of polynomials: pejorative condition and accurate
#     computation", by Zhonggang Zeng
#
# This code is dated Dec. 4, 2002. Please sent bug/failure report to
#
#             Prof. Zhonggang Zeng
#             Department of Mathematics
#             Northeastern Illinois University
#             Chicago, IL 60625
#
#             email: zzeng@neiu.edu
#
# This code is freely released for research exchange only. The author
# is not responsible for any demage caused by using this code.
#
# Calling syntax: The simplest way to call is
#
#     pejroot(f,y,m)
#
# where the input items are
#          f --- (row vector) the target polynomial
#          y --- (row vector) the initial iterate (of the roots)
#          m --- (row vector) the multiplicity structure
#
# For more advanced usage:
#
#	 [z, e, c] = pejroot(f, y, l; noi, tol, style, prtsty)
#
# The output
#          z --- (matrix) distinct roots (1st column) and corresponding
#                         multiplicities (2nd column)
#          e --- (scalar) the backward error of the roots
#          c --- (scalar) the pejorative condition number
#
#
#  Additional input parameters
#
#        noi --- (integer)    the number of iteration allowed
#                               (default: 10)
#        tol --- (numeric)    the error tolerance
#                               (default: 1.0d-8)
#      style --- (integer)    backward error minimization style:
#                               1 : overall (absolute)
#                               2 : coefficientwise (relative,default)
#     prtsty --- (integer)    intermediate results showing style
#                               0: minimal (default)
#                               1: plus intermediate backward error and
#                                     root corrections
#                               2: plus intemediate root approximations
#  Example:
#
#  f = poly([ones(1,20) 2*ones(1,10) 3*ones(1,5)]); # construct the test
#                					# polynomial with multiple roots
#  y = [0.95,2.05,2.95];           # prepare initial iterate
#  m = [20, 10, 5];                # prepare the multiplicity structure
#  pejroot(f,y,m)                  # running the program
# ans =
#
#   3.00000000000000   5.00000000000000
#   2.00000000000000  10.00000000000000
#   1.00000000000000  20.00000000000000
#
"""
function pejroot(f::AbstractVector{T}, z::AbstractVector{S}, mult; noi::Int = 10, tol::U = 1e-8, style::Int = 2, prtsty::Int = 0) where {T<:Number,S<:Number,U<:AbstractFloat}

    E = one(U)
    Z = zero(U)

    m = length(mult)	# number of variables (different roots)
    n = sum(mult)		# number of equations (degree of polynomial f)

    length(f) == n+1 || throw(ArgumentError("degree of polynomial $(length(f)-1) does not match multiplicies $(n)"))
    f[1] != Z || throw(ArgumentError("polynomial has not full degree compared to its length"))

    f = f / f[1]		# make the polynomial monic

    # sort initial values to enhance accuracy
    # It is interesting, sorting really improves accuracy
    jj = sortperm(abs.(z), rev=false)	# greatest first
    y0 = z[jj]
    ll = mult[jj]
    CT, zs = zerostructure(y0, ll)
    if !(CT<:T)
        f = CT.(f)
    end

    job = 0						# initialize job success return value
    y = copy(y0)				# pass the initial iterate
    h = copy(f[2:n+1])	        # make the RHS
    delta = zeros(U, 1, noi)	# space for sizes of the correction
    bcker = zeros(U, 1, noi)	# space for backward errors

    w = ones(T, n)				# set weight
    if style == 2
       for j = 1:n
           afj = abs(f[j+1])
           if afj >= E
               w[j] = E / afj
           end
       end
    end

    if prtsty >= 1
       println("  step   bckerr        delta")
    end

    Df = zeros(T, n ,m)			# open the space for Df
    bkerr = Z
    pjcnd = Z
    kk = 0

    for k = 1:noi
        kk = k

        if prtsty >= 2
            println(y)
        end

        # evaluate the coefficient operator and its Jacobian
        pr = products(y, ll)
        b = polyG(pr) - h
        Df = polyDG(y, ll, pr, zs)

        if style == 2   # scale
            b = w .* b
            lmul!(Diagonal(w), Df)
         end

         d = realtocoeff(Df \ b, S, zs) # least squares

         delta[k] = norm(d, 2)
         bcker[k] = norm(b, 2)

         if prtsty >= 1
             @printf(" %2.0f    %10.2e    %10.2e\n", k, bcker[k], delta[k])
         end

         if delta[k] < tol
            job = 1
         end     							# convergence criterion 1
         if k > 1
             if delta[k] >= delta[k-1] && bcker[k] >= bcker[k-1]
                 if job == 1
                     bkerr = bcker[k-1]
                     # accept result of previous iteration
                     break
                 end
             elseif delta[k] < delta[k-1]	# criterion 2
                 if delta[k]^2 / (delta[k-1] - delta[k]) < tol
                     job = 1
                 end
            end
        end

         yy = y
         step = 1.0
         stepmin = 1/16
         while step >= stepmin
             y = yy - d * step
             pr = products(y, ll)
             b = (polyG(pr) - h)
             if style == 2
                 b = w .* b
             end
             bks = norm(b, 2)
             if isnan(bks)
                 stepmin *= 0.5
             end
             if bks < bcker[k]
                 break
             end
             step *= 0.5
         end
         bkerr = bcker[k]
      end # k-loop

      if job == 1
          s = svdvals(Df)
          pjcnd = 1 / s[m]  # get pej. cond. number
          jj = sortperm(ll)
          y = y[jj]
          ll = ll[jj]   # sort by multiplicities
      end

      PolyZeros(y, ll), bkerr, pjcnd, job, kk
end


# Copyright 2021 The FirstOrderLp Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

using DelimitedFiles

@testset "trust_region_utils" begin
  @testset "solve_bound_constrained_trust_region unbounded" begin
    for solve_approximately in [true false]
      # min -x
      # || x || <= 5.0
      variable_lower_bounds = [-Inf]
      variable_upper_bounds = [Inf]
      current_point = [0.0]
      current_gradient = [-1.0]
      target_radius = 5.0
      norm_weights = [1.0]
      result = FirstOrderLp.solve_bound_constrained_trust_region(
        current_point,
        current_gradient,
        variable_lower_bounds,
        variable_upper_bounds,
        norm_weights,
        target_radius,
        solve_approximately,
      )
      @test result.value == -5.0
      @test result.solution == [5.0]

      # min x + y
      # 2 * x^2 + y^2 <= 6.0
      # [x*, y*] = [-1.0, -2.0]
      variable_lower_bounds = [-Inf, -Inf]
      variable_upper_bounds = [Inf, Inf]
      current_point = [0.0, 0.0]
      current_gradient = [1.0, 1.0]
      norm_weights = [2.0, 1.0]
      target_radius = sqrt(6.0)
      result = FirstOrderLp.solve_bound_constrained_trust_region(
        current_point,
        current_gradient,
        variable_lower_bounds,
        variable_upper_bounds,
        norm_weights,
        target_radius,
        solve_approximately,
      )
      @test result.solution ≈ [-1.0, -2.0] atol = 1.0e-8
      @test result.value ≈ -3.0 atol = 1.0e-8
    end
  end

  @testset "solve_bound_constrained_trust_region" begin
    # min -x
    # || x || <= 5.0
    variable_lower_bounds = [-Inf]
    variable_upper_bounds = [Inf]
    current_point = [0.0]
    current_gradient = [-1.0]
    target_radius = 5.0
    norm_weights = [1.0]
    result = FirstOrderLp.solve_bound_constrained_trust_region(
      current_point,
      current_gradient,
      variable_lower_bounds,
      variable_upper_bounds,
      norm_weights,
      target_radius,
      false,
    )
    @test result.solution == [5.0]

    # min -x
    # || x || <= 5.0
    # x <= 0.0
    variable_lower_bounds = [-Inf]
    variable_upper_bounds = [0.0]
    current_point = [0.0]
    current_gradient = [-1.0]
    result = FirstOrderLp.solve_bound_constrained_trust_region(
      current_point,
      current_gradient,
      variable_lower_bounds,
      variable_upper_bounds,
      norm_weights,
      target_radius,
      false,
    )
    @test result.solution == [0.0]

    # min -x
    # || x || <= 5.0
    # x <= 2.0
    variable_lower_bounds = [-Inf]
    variable_upper_bounds = [2.0]
    result = FirstOrderLp.solve_bound_constrained_trust_region(
      current_point,
      current_gradient,
      variable_lower_bounds,
      variable_upper_bounds,
      norm_weights,
      target_radius,
      false,
    )
    @test result.solution == [2.0]

    # min -2 * x - y
    # || [x; y] || <= 5.0
    # x <= 3.0
    # [x*, y*] = [3.0, 4.0]
    variable_lower_bounds = [-Inf, -Inf]
    variable_upper_bounds = [3.0, Inf]
    current_point = [0.0, 0.0]
    current_gradient = [-2.0, -1.0]
    norm_weights = [1.0, 1.0]
    result = FirstOrderLp.solve_bound_constrained_trust_region(
      current_point,
      current_gradient,
      variable_lower_bounds,
      variable_upper_bounds,
      norm_weights,
      target_radius,
      false,
    )
    @test result.solution ≈ [3.0, 4.0] atol = 1.0e-8

    # min -x
    # || [x; y] || <= 5.0
    # x <= 2.0
    # [x*, y*] = [2.0, 0.0]
    variable_lower_bounds = [-Inf, -Inf]
    variable_upper_bounds = [2.0, Inf]
    current_point = [0.0, 0.0]
    current_gradient = [-1.0, 0.0]
    result = FirstOrderLp.solve_bound_constrained_trust_region(
      current_point,
      current_gradient,
      variable_lower_bounds,
      variable_upper_bounds,
      norm_weights,
      target_radius,
      false,
    )
    @test result.solution == [2.0, 0.0]

    # Test norm_weights:
    # min -4.0 * x - 3.0 * y
    # 4.0^2 * x^2 + 3.0^2 * y^2 <= 2.0
    # [x*, y*] = [1.0 / 4.0, 1.0 / 3.0]
    # to see this substitute a = 4.0 x, b = 3.0 y, yielding:
    # min -a - b
    # a^2 + b^2 <= 2.0
    # [a*, b*] = [1.0, 1.0]
    variable_lower_bounds = [-Inf, -Inf]
    variable_upper_bounds = [Inf, Inf]
    current_point = [0.0, 0.0]
    current_gradient = [-4.0, -3.0]
    norm_weights = [4.0^2, 3.0^2]
    target_radius = sqrt(2.0)
    result = FirstOrderLp.solve_bound_constrained_trust_region(
      current_point,
      current_gradient,
      variable_lower_bounds,
      variable_upper_bounds,
      norm_weights,
      target_radius,
      false,
    )
    @test FirstOrderLp.weighted_norm(result.solution, norm_weights) ≈ sqrt(2.0) atol =
      1.0e-8
    @test result.solution ≈ [1.0 / 4.0, 1.0 / 3.0] atol = 1.0e-8

    # Hundred dimensional test
    # min -sum(x[i])
    # || x || <= r
    # x[i] <= i
    #
    # x*[i] = min(i, m)
    # Set r = || x* ||.
    n = 100
    variable_lower_bounds = zeros(n)
    variable_upper_bounds = 1.0 * collect(1:n)
    current_point = zeros(n)
    current_gradient = -ones(n)
    acceptable_relative_error = 1e-6
    norm_weights = ones(length(variable_lower_bounds))
    for m in [10.0, 50.0]
      target_radius = sqrt(sum([min(i, m)^2 for i in 1:n]))
      result = FirstOrderLp.solve_bound_constrained_trust_region(
        current_point,
        current_gradient,
        variable_lower_bounds,
        variable_upper_bounds,
        norm_weights,
        target_radius,
        false,
      )
      @test result.solution ≈ [min(i, m) for i in 1:n] atol = 1.0e-8
    end
  end

  @testset "bound_primal_and_dual_objective" begin
    lp = example_lp()
    primal_norm_weights = ones(4)
    dual_norm_weights = ones(3)
    # Test cases with zero or close to zero duality gap
    for norm in [FirstOrderLp.MAX_NORM, FirstOrderLp.EUCLIDEAN_NORM]
      primal_solution = [1.0, 0.0, 6.0, 2.0]
      dual_solution = [0.5, 4.0, 0.0]
      distance_to_optimality = 10.0
      result = FirstOrderLp.bound_optimal_objective(
        lp,
        primal_solution,
        dual_solution,
        primal_norm_weights,
        dual_norm_weights,
        distance_to_optimality,
        FirstOrderLp.MAX_NORM;
        solve_approximately = false,
      )

      @test result.lower_bound_value == -1.0
      @test result.upper_bound_value == -1.0

      primal_solution = [1.0, 0.0, 5.99999, 2.0]
      dual_solution = [0.50001, 4.0, 0.0]
      result = FirstOrderLp.bound_optimal_objective(
        lp,
        primal_solution,
        dual_solution,
        primal_norm_weights,
        dual_norm_weights,
        distance_to_optimality,
        FirstOrderLp.MAX_NORM;
        solve_approximately = false,
      )
      @test -1.01 < result.lower_bound_value < -1.0
      @test -1.0 < result.upper_bound_value < -0.99
    end

    primal_solution = [1.0, 0.0, 6.0, 1.0]
    dual_solution = [0.0, 4.0, 0.0]
    distance_to_optimality = 2.0
    result = FirstOrderLp.bound_optimal_objective(
      lp,
      primal_solution,
      dual_solution,
      primal_norm_weights,
      dual_norm_weights,
      distance_to_optimality,
      FirstOrderLp.MAX_NORM;
      solve_approximately = false,
    )
    @test result.lower_bound_value == -4.0
    @test result.upper_bound_value == 2.0
    @test result.lower_bound_value ==
          FirstOrderLp.corrected_dual_obj(lp, primal_solution, dual_solution)

    primal_solution = [3.0, 0.0, 6.0, 0.0]
    dual_solution = [0.0, 4.0, 0.0]
    distance_to_optimality = 5.0
    result = FirstOrderLp.bound_optimal_objective(
      lp,
      primal_solution,
      dual_solution,
      primal_norm_weights,
      dual_norm_weights,
      distance_to_optimality,
      FirstOrderLp.EUCLIDEAN_NORM;
      solve_approximately = false,
    )
    @test result.lower_bound_value == -4.0
    @test result.lagrangian_value == -1.0
    @test distance_to_optimality^2 ==
          norm(result.primal_solution - primal_solution, 2)^2 +
          norm(result.dual_solution - dual_solution, 2)^2

    @test result.upper_bound_value == 7.0

    primal_solution = [1.0, 1.0, 4.0, 1.0]
    dual_solution = [0.0, 0.0, 0.0]
    distance_to_optimality = 10.0
    result = FirstOrderLp.bound_optimal_objective(
      example_lp(),
      primal_solution,
      dual_solution,
      primal_norm_weights,
      dual_norm_weights,
      distance_to_optimality,
      FirstOrderLp.MAX_NORM;
      solve_approximately = false,
    )
    corrected_dual_obj =
      FirstOrderLp.corrected_dual_obj(lp, primal_solution, dual_solution)
    @test result.lower_bound_value ==
          FirstOrderLp.corrected_dual_obj(lp, primal_solution, dual_solution)

    primal_solution = [0.5, 0.5, 0.5, 1.0, 1.0, 1.0] # Solution in the interior.
    dual_solution = zeros(3)
    primal_norm_weights = ones(6)
    dual_norm_weights = ones(3)

    result = FirstOrderLp.bound_optimal_objective(
      example_cc_star_lp(),
      primal_solution,
      dual_solution,
      primal_norm_weights,
      dual_norm_weights,
      distance_to_optimality,
      FirstOrderLp.MAX_NORM;
      solve_approximately = false,
    )
    @test result.lagrangian_value == result.upper_bound_value
    @test result.lower_bound_value < result.lagrangian_value
  end
end