using Test
using Random
using QuantileSummaries

@testset "String" begin
    # string
    qb = QuantileStat{String}()

    for i = 1:1000
        fit!(qb, randstring(10))
    end

    summary = value(qb)
    @test true #passed without erroring
end

@testset "Correctness" begin

    eps = 0.01
    qb = QuantileStat(eps)

    n = 10_000
    r = randn(n)
    fit!(qb, r)
    sort!(r)

    qs = summarize(qb)

    # Test if the right percentile is identified (within eps)
    for percentile = eps:eps/2:1

        valueforpercentile = r[ Int(floor(percentile * n)) ]
        predictedpercentile = QuantileSummaries.percentile(qs, valueforpercentile)
        @test abs(predictedpercentile-percentile) < eps

    end

    for i = 1:1000
        percentile = rand()
        idx = 1 + Int(floor(percentile * n))
        valueforpercentile = r[idx]
        predictedpercentile = QuantileSummaries.percentile(qs, valueforpercentile)
        @test abs(predictedpercentile-percentile) < eps
    end

    # check if computed value is within eps of expected
    # Percentiles & values are the same in this setup
    eps = 0.01
    qb = QuantileStat(eps)

    r = collect(0:0.0001:1)
    shuffle!(r)
    fit!(qb, r)
    qs = summarize(qb)

    for p in 0:0.001:1
        value = qvalue(qs, p)
        @test( abs(p-value) < eps )
    end
end

@testset "merge!" begin
    eps = 0.01
    qb1 = QuantileStat(eps)
    qb2 = QuantileStat(eps)

    n = 10_000
    r = randn(n)
    for value in r
        if rand() < 0.1
            fit!(qb1, value)
        else
            fit!(qb2, value)
        end
    end

    merge!(qb1, qb2)
    sort!(r)

    qs = summarize(qb1)
    for percentile = eps:eps/2:1
        valueforpercentile = r[ Int(floor(percentile * n)) ]
        predictedpercentile = QuantileSummaries.percentile(qs, valueforpercentile)
        @test abs(predictedpercentile-percentile) < eps
    end

end


#lowerBoundsSymmetric

#computes the maximum over the geometric grid
#sep_1(a, b, l, u) solves  min_{t in [l, u]}  at + 2bt H(t)
#sep_2(a, l, u) solves min_{t in [l, u]} a H(t)
#user must ensure sep_1, sep_2 and h are consistent.
#H(t) = 1/2t int_[-t, t] h(1+s)ds
function _rev_p_lb_symmetric(S, h, sep_1, sep_2, H, H_un, pj, numCuts, print_trace, TOL)
	m = Model(solver=GurobiSolver(OutputFlag=false))
	@variable(m, theta)
	@variable(m, lam)
	@objective(m, Min, theta)

	#add endpoints to get things started
	#t = 0
	ind = pj <= 1 ? pj : 0
	@constraint(m, lam * H(0) >= ind - theta)

	#t = abs(pj -1)
	tstar = abs(pj -1)
	@constraint(m, lam * H(tstar) >= ind - theta)
	@constraint(m, (2theta - pj) * tstar + lam * H_un(tstar) >= pj * (1 - pj))

	#t = S-1
	@constraint(m, (2theta - pj) * (S - 1) + lam * H_un(S - 1) >= pj * (1 - pj))

	for iter = 1:numCuts
		status = solve(m)
		feasible = true

		print_trace && println("Iteration: \t", iter)
		#First constraint [0, abs(pj -1)]
		rhs =  ind - getvalue(theta)
		val, tstar = sep_2(getvalue(lam), 0, abs(pj -1))
		print_trace && println("Iter: $iter Val $val RHS $rhs  : tstar $tstar")
		if val < rhs - TOL
			@constraint(m, lam * H(tstar) >= ind - theta)
			feasible = false
		end

		#2nd constraint [abs(pj-1) S-1]
		rhs = pj * (1 - pj)
		val, tstar = sep_1(2 * getvalue(theta) - pj, getvalue(lam), abs(pj-1), S-1)

		print_trace && println("Iter: $iter Val $val RHS $rhs : tstar $tstar")
		if val < rhs - TOL
			cnst = @constraint(m, (2theta - pj) * tstar + lam * H_un(tstar) >= rhs)
			feasible = false
		end

		if feasible
			if status != :Optimal
				println("Problem likely infeasible.  Dual Status:\t", status)
				return -1.0
			else #correctly solved
				return getobjectivevalue(m)
			end
		end
	end
	##Exhausted iteration count
	throw("Max Iteration Count Reached:  h(1): $(h(1))")
	return -1.
end

function vopp_lb_symmetric_MAD(mu, S, M, D; delta = S/100, numCuts=100, print_trace=false, TOL=1e-6, pj=-1, safe_fail=false)
	#standardize
	c = mu * (1 - M)
	Sc = vopp.comp_Sc(S, M)

	#round down Sc if it's too large and throw otherwise
	Sc = min(Sc, 2)
	@assert Sc > 1 "Scale $S is too small for any distribution"

	if D/M > .25
		println("Deviation $(D/M) > .25, problem likely infeasible")
		if safe_fail
			println("Safe_fail set to true. Aborting Calc")
			return -1.0
		end
	end

	#create the functions/separators closed-form
	h(t) = M/2 * abs(t-1) - D
	H(t) = M * t / 4 - D
	H_un(t) = 2t * H(t)

	#sep_1(a, b, l, u) solves  min_{t in [l, u]}  at + b * H_un(t)
	#min_{t in [l, u]} (a - 2bD) t + bM/2 * t^2
	function sep_1(a, b, l, u)
		f(t) = (a - 2b * D) * t + b * M/2 * t^2
		valstar, tstar = f(l), l
		if f(u) < valstar
			valstar, tstar = f(u), u
		end
		if b > 0
			tmid = (2b * D - a) / b / M
			valmid = Inf
			if l <= tmid <= u
				valmid = f(tmid)
			end
			if valmid < valstar
				valstar, tstar = valmid, tmid
			end
		end
		valstar, tstar
	end

	#solves min_{t in [l, u]} a H(t)
	function sep_2(a, l, u)
		if a * H(u) < a * H(l)
			return a * H(u), u
		end
		return a * H(l), l
	end

	#search over the geometric price ladder
	if pj > 0
		ps = [pj]
	else
		ps = geom_price_ladder(Sc, delta)
	end
	rstar = -Inf
	pstar = 0.
	rj = 0.

	for pj in ps
		#kick i tto the workhorse
		rj = _rev_p_lb_symmetric(Sc, h, sep_1, sep_2, H, H_un, pj, numCuts, print_trace, TOL)
		print_trace && println("\n Price Pj:\t", pj, "\t", "Rj:\t", rj, "\n")

		if rj > rstar
			rstar = rj
			pstar = pj
		end
	end

	return 1/rstar, pstar
end

function vopp_lb_symmetric_CV(mu, S, M, C; delta = S/100, numCuts=100, print_trace=false, TOL=1e-6, pj=-1, safe_fail=false)
	#standardize
	c = mu * (1 - M)
	Sc = vopp.comp_Sc(S, M)

	#round down Sc if it's too large and throw otherwise
	Sc = min(Sc, 2)
	@assert Sc > 1 "Scale $S is too small for any distribution"

	max_cv = max_std_cv_unimodal_guess(1, Sc, M, true)
	if  C / M > max_cv
		println("CV $(C/M) > $max_cv, so problem likely infeasible")
		if safe_fail
			println("Safe_fail set to true. Aborting Calc")
			return -1.0
		end
	end

	#create the functions/separators closed-form
	h(t) = M^2 * (t-1)^2 - C^2
	H(t) = M^2 / 3 * t^2 - C^2
	H_un(t) = 2t * H(t)

	#sep_1(a, b, l, u) solves
	#min_{t in [l, u]}  at + b * H_un(t)
	function sep_1(a, b, l, u)
		f(t) = a * t + 2 * b * M^2 / 3 * t^3 - 2b * t * C^2
		valstar, tstar = f(l), l
		if f(u) < valstar
			valstar, tstar = f(u), u
		end
		if (2b * C^2 - a) / (2b * M^2) > 0
			tmid = sqrt((2b * C^2 - a) / (2b * M^2))
			valmid = Inf
			if l <= tmid <= u
				valmid = f(tmid)
			end
			if valmid < valstar
				valstar, tstar = valmid, tmid
			end
		end
		valstar, tstar
	end

	#solves min_{t in [l, u]} a H(t)
	function sep_2(a, l, u)
		f(t) = a * H(t)
		tstar, valstar = l, f(l)
		if f(u) < valstar
			tstar, valstar = u, f(u)
		end
		if l <= 0. <= u
			if f(0) < valstar
				tstar, valstar = 0., f(0.)
			end
		end
		return valstar, tstar
	end

	#search over the geometric price ladder
	if pj > 0
		ps = [pj]
	else
		ps = geom_price_ladder(Sc, delta)
	end
	rstar = -Inf
	pstar = 0.
	rj = 0.

	for pj in ps
		#kick i tto the workhorse
		rj = _rev_p_lb_symmetric(Sc, h, sep_1, sep_2, H, H_un, pj, numCuts, print_trace, TOL)
		print_trace && println("\n Price Pj:\t", pj, "\t", "Rj:\t", rj, "\n")

		if rj > rstar
			rstar = rj
			pstar = pj
		end
	end

	return 1/rstar, pstar
end


##
x = 5
##
typeof(x)

##
"""
`add5`
======

This adds the value of 5 to the input variable.

Returns: the same type as the input variable + the integer 5.

Example
-------
        add5(5.0)
        typeof(add5(5.0)) # returns Float64
        add5(5)
        typeof(add5(5)) # returns Int64
"""
function add5(x)
    return x+5
end
@show typeof(add5(5.0))
@show add5(5.0)

##
@show typeof(add5(5))
@show add5(5)
##

## Types of Matrices and Vectors
# a vector in Julia looks like:

x = [5,6,7]
@show typeof(x)

##
x = [5.0,6,7]
@show typeof(x)

##
A = [5 6 7]
##
A = [5.0 6 7; 1 2 3; 4 5 6]
##
B = A .+ 5

##
C = B' # this is the transpose
@show typeof(C)
##
D = copy(C)
##
E = D+B

##


using Printf

function makesofn(n::Integer)
    buf = Vector{typeof(n)}(0)
    i = 0
    return function sofn(item)
        i += 1
        if i ≤ n
            push!(buf, item)
        else
            j = rand(1:i)
            if j ≤ n buf[j] = item end
        end
        return buf
    end
end

nhist = zeros(Int, 10)
for _ in 1:10^5
    kas = makesofn(3)
    for j in 0:8 kas(j) end
    for k in kas(9) nhist[k+1] += 1 end
end

println("Simulating sof3(0:9) 100000 times:")
for (i, c) in enumerate(nhist)
    @printf("%5d → %5d\n", i-1, c)
end


using Polynomials, SimplePartitions
import Base.length, Base.setindex!, Base.getindex
export chrome_poly, reset_cpm, size_cpm

const CPM_pair = Tuple{SimpleGraph,Poly{Int}}
const CPM_dict = Dict{UInt64,Vector{CPM_pair}}

# This is a device to record previously computed chromatic polynomials
mutable struct ChromePolyMemo
    D::CPM_dict
    function ChromePolyMemo()
        new( CPM_dict() )
    end
end

_CPM = ChromePolyMemo()

"""
`reset_cpm()` clears the datastructure of previously computed
chromatic polynomials.
"""
function reset_cpm()
    global _CPM = ChromePolyMemo()
    return true
end

"""
`size_cpm()` reports the number of graphs held in the datastructure of
previously computed chromatic polynomials.
"""
function size_cpm()
    return length(_CPM)
end

function show(io::IO, CPM::ChromePolyMemo)
    print(io,"ChromePolyMemo with $(length(CPM)) graphs")
end

function length(CPM::ChromePolyMemo)
    result = 0
    for ds in keys(CPM.D)
        result += length(CPM.D[ds])
    end
    return result
end

"""
`remember!(CPM,G,P)` save the polynomial `P` associated with `G` in
this data structure (if it wasn't in there already).

Short form: `CPM[G] = P`.
"""
function remember!(CPM::ChromePolyMemo, G::SimpleGraph, P::Poly{Int})
    index::Int128 = uhash(G)  # get the signature for this graph

    # have we seen this graph's uhash before?
    # if no, create a new entry in the Memo
    if !haskey(CPM.D,index)
        CPM.D[index] = [(G,P)]
        return
    end

    # Otherwise, look through the values associated with this uhash to
    # see if we've seen this graph (or one isomorphic to it) before.

    # Recall those graphs with matching uhash
    for (H,Q) in CPM.D[index]
        if is_iso(G,H) # if isomorphic, nothing to add
            return
        end
    end

    # otherwise, add this polynomial to the end of the list
    push!(CPM.D[index], (G,P))
end

setindex!(CPM,P,G) = remember!(CPM,G,P)

"""
`recall(CPM,G)` look up in CPM to see if we have already computed the
chromatic polynomial of this graph (or something isomorphic to it).

Short form: `P = CPM[G]`.
"""
function recall(CPM::ChromePolyMemo, G::SimpleGraph)
    ds = uhash(G)
    SG = CPM.D[ds]  # This may throw an error if not found. That's good.

    for (H,Q) in SG
        if is_iso(G,H)
            return Q
        end
    end

    # Not found, so it's an error
    error("Graph not found")
end

getindex(CPM,G) = recall(CPM,G)

"""
`chrome_poly(G)` computes the chromatic polynomial of the graph `G`.
If `G` is a directe graph, this returns the chromatic polynomial of
`G`'s underlying simple graph.

This function builds a datastructure to prevent computing the
chromatic polynomial of the same graph twice. To do this, it uses
frequent isomorphism checks.

See `size_cpm` and `reset_cpm`.
"""
function chrome_poly(G::SimpleGraph) #, CPM::ChromePolyMemo = _CPM)
  if cache_check(G,:chrome_poly)
    return G.cache[:chrome_poly]
  end

    n::Int = NV(G)
    m::Int = NE(G)

    # Special case: no vertices
    if n==0
        return Poly([1])
    end

    # Special case: no edges
    if m==0
        return Poly([zeros(Int,n);1])
    end

    # Special case: complete graph
    if 2m == n*(n-1)
        return poly(0:n-1)
    end

    # Special case: disconnected graph
    comps = parts(components(G))
    if length(comps) > 1
        result = Poly([1])
        for A in comps
            H = induce(G,A)
            result *= chrome_poly(H)
        end
        return result
    end

    # Special case: Tree
    if m == n-1
        return Poly([0,1]) * Poly([-1,1])^(n-1)
    end

    # See if we've computed this chromatic polynomial of this graph before
    try P = _CPM[G]
        return P
    catch
    end

    # Resort to inclusion/exclusion

    # Find a vertex u of minimum degree
    min_d = minimum(deg(G))
    smalls = filter(v -> deg(G,v) == min_d, G.V)
    u = first(smalls)

    # And choose any neighbor
    v = first(G[u])

    # Special case to speed up handling leaves
    if min_d==1
        GG = deepcopy(G)
        delete!(GG,u)
        p1 = chrome_poly(GG)
        P = p1 * Poly([-1,1])
        _CPM[G] = P
        return P
    end

    # p1 is chrome_poly of G-e
    GG = deepcopy(G)
    delete!(GG,u,v)
    p1 = chrome_poly(GG)

    # p2 is chrome_poly of G/e
    GG = deepcopy(G)
    contract!(GG,u,v)
    p2 = chrome_poly(GG)

    P = p1 - p2

    _CPM[G] = P # save in case we see this graph again
    G.cache[:chrome_poly] = P  # and save in graph's cache too
    return P
end

chrome_poly(G::SimpleDigraph) = chrome_poly(simplify(G))


#  Unit "LogLike.jl" of the Diagonalization.jl package for Julia language
#
#  MIT License
#  Copyright (c) 2019, 2020
#  Marco Congedo, CNRS, Grenoble, France:
#  https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#  This unit implements two Dinh-Tuan Pham's algorithms based on the
#  log-likelyhood (Kullback-Leibler divergence) criterion.
#  D.-T. Pham (2000) Joint approximate diagonalization of positive definite
#  matrices, SIAM Journal on Matrix Analysis and Applications, 22(4), 1136–1152.
#  They are adapted in Julia from code freely made available from the author.
#
#  These algorithms handles the AJD diagonalization procedure, corresponding
#  to the case m=1, k>1 according to the taxonomy adopted in this package.

# function to get the weights from argment `w`
function _logLikeWeights(w, 𝐂, type)
	if 	   w isa Function
		   w=[w(C) for C∈𝐂]
	elseif w===○
		   w=ones(type, length(𝐂))
	end
	return w, sum(w)
end

#  PRIMITIVE LogLike algorithm:
#  It takes as input a n·nk matrix holding k horizontally stacked n·n real or
#  complex matrices, such as C=[C_1...C_k].
#  It finds a non-singular matrix B such that the
#  congruences B'*C_κ*B are as diagonal as possible for all κ=1:k.
#  `tol` is the convergence to be attained.
#  `maxiter` is the maximum number of iterations allowed.
#  if `verbose`=true, the convergence attained at each iteration and other
#  information will be printed.
#
#  NB: Pham's algorithm proceeds by transforming pairs of vectors of B.
#  A sweep goes over all (n*(n+1))/2 ij pairs, i>j. Thus it can be optimized
#  by multi-threading the optimization of the pairs as it is done for
#  algorithms based on Givens rotations (e.g., round-Robin tournament scheme).
#
#  RETURN: B, the number of iterations and the convergence attained (a 3-tuple)
function logLike(C::Matrix{T};
	     		 tol 	 = 0.,
				 maxiter = 1000,
				 verbose = false) where T<:Union{Real, Complex}

	@inline function phamSweep!()
		decr = 0.
		@inbounds for i = 2:n, j = 1:i-1
			c1 = C[i, i:n:nk]
			c2 = C[j, j:n:nk]
			g₁₂ = mean(C[i, j:n:nk]./c1)		# this is g_{ij}
			g₂₁ = mean(C[i, j:n:nk]./c2)		# conjugate of g_{ji}
			𝜔₂₁ = mean(c1./c2)
			𝜔₁₂ = mean(c2./c1)
			𝜔 = √(𝜔₁₂*𝜔₂₁)
			𝜏 = √(𝜔₂₁/𝜔₁₂)
			𝜏₁ = (𝜏*g₁₂ + g₂₁)/(𝜔 + 1.)
			if T<:Real 𝜔=max(𝜔 - 1., e) end
			𝜏₂ = (𝜏*g₁₂ - g₂₁)/𝜔 		#max(𝜔 - 1., e)	# in case 𝜔 = 1
			h₁₂ = 𝜏₁ + 𝜏₂				# this is twice h_{ij}
			h₂₁ = conj((𝜏₁ - 𝜏₂)/𝜏)	# this is twice h_{ji}
			decr += k*(g₁₂*conj(h₁₂) + g₂₁*h₂₁)/2.

			𝜏 = 1. + 0.5im*imag(h₁₂*h₂₁)	# = 1 + (h₁₂*h₂₁ - conj(h₁₂*h₂₁))/4
			𝜏 = 𝜏 + √(𝜏^2 - h₁₂*h₂₁)
			Γ = [1 conj(-h₂₁/𝜏); conj(-h₁₂/𝜏) 1]
			C[[i, j], :] = Γ'*C[[i, j], :]		# new i, j rows of C
			ijInd = vcat(collect(i:n:nk), collect(j:n:nk))
			C[:, ijInd] = reshape(reshape(C[:, ijInd], n*k, 2)*Γ, n, k*2) # new i,j columns of C
			B[:, [i, j]] = B[:, [i, j]]*Γ # update the columns of B
		end
		return decr
	end # phamSweep

	(n, nk) = size(C)
	k, e = nk÷n, T(eps(real(T)))
	B=Matrix{T}(I, n, n) # initialize AJD algorithm

    iter, conv = _iterate!("LogLike", phamSweep!, maxiter, T, tol, verbose)

    return B, iter, conv
end

#  ADVANCED LogLike algorithm:
#  It takes as input a vector of k real symmetric or complex Hermitian
#  matrices 𝐂 and finds a non-singular matrix B such that the
#  congruences B'*𝐂_κ*B are as diagonal as possible for all κ=1:k.
#  It handles both real and complex data input.
#
#  NB: For the moment being the weights are not supported. The `w` argument
#  is left in for syntax homogeneity with other AJD procedures.
#
#  if `whitening` = true is passed, the Jeffrey mean of the matrices in 𝐂 is
#  computed (using the PosDefManifold.jl package) and the matrices in 𝐂
#  are pre-transformed using the whitening matrix of the mean.
#  Dimensionality reduction can be obtained at this stage using optional
#  arguments `eVar` and `eVarMeth` (see documentation of the AJD constructors).
#
#  if sort=true (default) the column vectors of the B matrix are normalized
#  to unit norm and permuted so as to sort in descending order the mean over
#  κ=1:k of the diagonal elements of B'*𝐂_κ*B.
#  Note that if `whitening` is true the output B will not have unit norm columns
#  as it is multiplied by the whitener after being scaled and sorted
#  and before being returned.
#
#  if  `whitening` = false (default), a matrix can be provided with the `init`
#  argument in order to initialize B. In this case the actual AJD
#  will be given by init*B, where B is the output of the algorithms.
#
#  `tol` is the convergence to be attained.
#
#  `maxiter` is the maximum number of iterations allowed.
#
#  if `verbose`=true, the convergence attained at each iteration and other
#  information will be printed.
#
#  return: B, its pseudo-inverse, the mean diagonal elements of B'*mean(𝐂)*B,
#          the number of iterations and the convergence attained
function logLike(𝐂::Union{Vector{Hermitian}, Vector{Symmetric}};
				 w			:: Union{Tw, Function} = ○,
				 preWhite	:: Bool = false,
				 sort      	:: Bool = true,
				 init		:: Union{Matrix, Nothing} = ○,
				 tol     	:: Real = 0.,
				 maxiter 	:: Int  = 1000,
				 verbose 	:: Bool = false,
			  eVar 	   :: TeVaro = ○,
			  eVarMeth :: Function = searchsortedfirst)

	# pre-whiten or initialize and stack matrices horizontally
	W, C = _preWhiteOrInit(𝐂, preWhite, Jeffrey, eVar, eVarMeth, init, :stacked)

	(n, nk) = size(C)

	B, iter, conv = logLike(C; tol=tol, maxiter=maxiter, verbose=verbose)

	# scale and permute the vectors of B
    D=Diagonal([mean(C[i, i:n:nk]) for i=1:n])
    λ = sort ? _permute!(_scale!(B, D, n)...) : diag(D)

	return preWhite ? (W.F*B, pinv(B)*W.iF, λ, iter, conv) :
                      (B, pinv(B), λ, iter, conv)
end

######################################################
#  ADVANCED 'c-style' version of the LogLike algorithm
#  It supports only real data.
#  It uses the same API as 'logLike' function.
#  It supports weights:
#  `w` is an optional vector of k non-negative weights for each matrix in 𝐂.
#  Pham's criterion being invariant by scaling, the weights act on the cost
#  function, not as a weights for the entries of the matrices.
#  Notice that in this algorithm the weights can be zero,
#  amounting to ignoring the corresponding matrices.
#  By default all weights are equal to one.
#  A function can be passed as the `w` argument, in which case the kth weight
#  is found as the output of the function applied to the kth matrix in 𝐂.
#  A good choice in general is the `nonD` function declared in tools.jl unit.
#
#  All other arguments are in 'logLike' function
function logLikeR(𝐂::Union{Vector{Hermitian}, Vector{Symmetric}};
				  w 		:: Union{Tw, Function} = ○,
				  preWhite  :: Bool = false,
				  sort      :: Bool = true,
				  init 	 	:: Union{Matrix, Nothing} = ○,
				  tol     	:: Real = 0.,
				  maxiter 	:: Int  = 1000,
				  verbose 	:: Bool = false,
			eVar 	  :: TeVaro = ○,
			eVarMeth  :: Function = searchsortedfirst)

	@inline function phamSweepR!()
	   det, decr, i, ic  = 1., 0., 1, n
	   @inbounds while i<n
			j, jc = 0, 0
			while j<i
				ii, jj, ij = i + ic, j + jc, i + jc
				q1, p2, p, q = 0., 0., 0., 0.
				for κ=1:k
					if w[κ]>e
						wκ, t1, t2, t = w[κ], 𝐜[κ][ii+1], 𝐜[κ][jj+1], 𝐜[κ][ij+1]
						p += wκ*t/t1; q += wκ*t/t2; q1 += wκ*t1/t2; p2 += wκ*t2/t1
					end
				end

				# find rotation
				q1 /= ∑w; p2 /= ∑w; p  /= ∑w; q  /= ∑w; β = 1. - p2*q1
				if (q1 ≤ p2)
					α = p2*q - p
					abs(α)-β<e ? begin β = -1.; γ = p/p2; end : γ = - (p*β + α)/p2 # p1 = 1
					decr += ∑w*(p^2 - α^2/β)/p2
				else
					γ  = p*q1 - q		# p1 = 1
					abs(γ)-β<e ? begin β = -1.; α = q/q1; end : α = -(q*β + γ)/q1	# q2 = 1
					decr += ∑w*(q^2 - γ^2/β)/q1
				end
				t = 2/(β - sqrt(β^2 - 4*α*γ))
				B₁₂, B₂₁ = γ*t, α*t

				for κ=1:k
					if w[κ]>e
					   	ii, jj = i, j
						while ii<ij
							𝜏 = 𝐜[κ][ii+1]
							𝐜[κ][ii+1] = 𝐜[κ][ii+1] + B₁₂*𝐜[κ][jj+1]
							𝐜[κ][jj+1] = 𝐜[κ][jj+1] + B₂₁*𝜏 # at exit ii = ij = i + jc
							ii += n
							jj += n
						end
						𝜏 = 𝐜[κ][i+ic+1]
						𝐜[κ][i+ic+1] += (B₁₂*(2*𝐜[κ][ij+1] + B₁₂*𝐜[κ][jj+1]))
						𝐜[κ][jj+1] += B₂₁*𝐜[κ][ij+1]
						𝐜[κ][ij+1] += B₂₁*𝜏	# element of index j,i */
						while ii<ic
							𝜏 = 𝐜[κ][ii+1]
							𝐜[κ][ii+1] += B₁₂*𝐜[κ][jj+1]
							𝐜[κ][jj+1] += B₂₁*𝜏
							ii += n
							jj += 1
						end

						jj += 1
						ii += 1
						while jj<(jc+n)
							𝜏 = 𝐜[κ][ii+1]
							𝐜[κ][ii+1] += B₁₂*𝐜[κ][jj+1]
							𝐜[κ][jj+1] += B₂₁*𝜏
							jj += 1
							ii += 1
						end
					end
				end

				@inbounds for r=1:n # rotate B
					𝜏 = B[i+1, r]
					B[i+1, r] += B₁₂*B[j+1, r]
					B[j+1, r] += B₂₁*𝜏
				end

				det *= 1. - B₁₂*B₂₁ # compute determinant
				j  += 1
				jc +=  n
			end # while j
			i += 1
			ic += n
	   end # while i

		#= useless computing the logdet crit as we only need its decrement
		ld += (2*∑w*log(det));
		𝜏 = 0.
		for κ=1:k # OK
			if w[κ]>e
				det = 1.
				ii = 1
				while ii≤n^2 # OK
					 det *= 𝐜[κ][ii]
					 ii += (n+1)
				end
				𝜏 += w[κ]*log(det)
			end
		end
		# return 𝜏 - ld
		=#
		return decr
	end # phamSweepR

	n, k, type =size(𝐂[1], 1), length(𝐂), eltype(𝐂[1])
	e = eps(type)*100

	w, ∑w = _logLikeWeights(w, 𝐂, type) # weights and sum of weights

	# pre-whiten, initialize and write matrices in vectorized form
	if preWhite
		W=whitening(PosDefManifold.mean(Jeffrey, 𝐂); eVar=eVar, eVarMeth=eVarMeth)
		𝐜=[(W.F'*C*W.F)[:] for C∈𝐂]
		n=size(W.F, 2) # subspace dimension
	else
		# initialization only if preWhite is false
		init≠nothing ? 𝐜=[(init'*C*init)[:] for C∈𝐂] : 𝐜=[C[:] for C∈𝐂]
	end

	B=Matrix{type}(I, n, n)
    iter, conv = _iterate!("LogLikeR", phamSweepR!, maxiter, type, tol, verbose)
	B = Matrix(B')

	# sort the vectors of solver
	M=mean(𝐂)
	D=Diagonal([PosDefManifold.quadraticForm(B[:, i], M) for i=1:n])
	# scale and permute the vectors of B
    λ = sort ? _permute!(_scale!(B, D, n)...) : diag(D)

	return preWhite ? (W.F*B, pinv(B)*W.iF, λ, iter, conv) :
                      (B, pinv(B), λ, iter, conv)
end


# by lazarusA # HIDE
# using GLMakie # HIDE
using CairoMakie
CairoMakie.activate!() # HIDE
let
    testField(x,y) = Point(-x, 2y) # x'(t) = -x, y'(t) = 2y
    x =  -2:0.1:4
    y =  -2:0.1:2
    u2(x,y) = -x
    v2(x,y) = 2y
    z = [log10(sqrt(u2(x,y)^2 + v2(x,y)^2)) for x in x, y in y]

    fig = Figure(resolution = (700, 450), fontsize = 18, font = "sans")
    ax = fig[1, 1] = Axis(fig, xlabel = "x", ylabel = "y")
    fs = heatmap!(ax, x, y, z, colormap = Reverse(:plasma))
    streamplot!(ax, testField, x, y, colormap = Reverse(:plasma),
        gridsize= (32,32), arrow_size = 10)
    #limits!(ax, -2, 4, -2, 2)
    cbar = fig[1,2] = Colorbar(fig, fs, label = "log10[(u²+v²)½]",
        labelsize = 14, ticklabelsize=14)
    cbar.width = 20
    colgap!(fig.layout, 5)
    fig
    save(joinpath(@__DIR__, "output", "streamplotAndHeatmap.png"), fig, px_per_unit = 2) # HIDE
end
using Pkg # HIDE
Pkg.status("CairoMakie") # HIDE


module TestTransformer

using Test, MLJBase, MLJModels
using Tables, CategoricalArrays, Random

#### STATIC TRANSFORMER ####

t  = StaticTransformer(f=log)
f, = fit(t, 1, nothing)
@test transform(t, f, 5) ≈ log(5)

infos = info_dict(t)
@test infos[:input_scitype]  == MLJBase.Table(MLJBase.Scientific)
@test infos[:output_scitype] == MLJBase.Table(Scientific)

#### FEATURE SELECTOR ####

N = 100
X = (Zn   = rand(N),
     Crim = rand(N),
     x3   = categorical(rand("YN", N)),
     x4   = categorical(rand("YN", N)))

namesX   = Tables.schema(X).names |> collect
selector = FeatureSelector()
f,       = fit(selector, 1, X)

@test f == namesX

Xt = transform(selector, f, selectrows(X, 1:2))

@test Set(Tables.schema(Xt).names) == Set(namesX)
@test length(Xt.Zn) == 2

selector = FeatureSelector(features=[:Zn, :Crim])
f,       = fit(selector, 1, X)

@test transform(selector, f, selectrows(X, 1:2)) ==
        selectcols(selectrows(X, 1:2), [:Zn, :Crim])

infos = info_dict(selector)
@test infos[:input_scitype]  == MLJBase.Table(Scientific)
@test infos[:output_scitype] == MLJBase.Table(Scientific)

#  To be added with FeatureSelectorRule X = (n1=["a", "b", "a"], n2=["g", "g", "g"], n3=[7, 8, 9],
#               n4 =UInt8[3,5,10],  o1=[4.5, 3.6, 4.0], )
# MLJBase.schema(X)
# Xc = coerce(X,  :n1=>Multiclass, :n2=>Multiclass)

# t = Discretizer(features=[:o1, :n3, :n2, :n1])
# @test Xt.features == [:o1, :n3, :n2, :n1]
# @test Xt.is_ordinal == [true, false, false, false]
# @test Xt.A == [512 1 1 1; 1 2 1 2; 256 3 1 1]

#### UNIVARIATE DISCRETIZATION ####

# TODO: move this test to MLJBase:
# test helper function:
v = collect("qwertyuiopasdfghjklzxcvbnm1")
X = reshape(v, (3, 9))
Xcat = categorical(X)
Acat = Xcat[:, 1:4] # cat vector with unseen levels
element = Acat[1]
@test transform(element, X) == Xcat
@test transform(element, X[5]) == Xcat[5]

v = randn(10000)
t = UnivariateDiscretizer(n_classes=100);
result, = fit(t, 1, v)
w = transform(t, result, v)
bad_values = filter(v - MLJBase.inverse_transform(t, result, w)) do x
    abs(x) > 0.05
end
@test length(bad_values)/length(v) < 0.06

# scalars:
@test transform(t, result, v[42]) == w[42]
r =  inverse_transform(t, result, w)[43]
@test inverse_transform(t, result, w[43]) ≈ r

# test of permitted abuses of argument:
@test inverse_transform(t, result, get(w[43])) ≈ r
@test inverse_transform(t, result, map(get, w)) ≈
    inverse_transform(t, result, w)

# all transformed vectors should have an identical pool (determined in
# call to fit):
v2 = v[1:3]
w2 = transform(t, result, v2)
@test levels(w2) == levels(w)

#### UNIVARIATE STANDARDIZER ####

stand = UnivariateStandardizer()
f,    = fit(stand, 1, [0, 2, 4])

@test round.(Int, transform(stand, f, [0,4,8])) == [-1.0,1.0,3.0]
@test round.(Int, inverse_transform(stand, f, [-1, 1, 3])) == [0, 4, 8]

infos = info_dict(stand)

#### STANDARDIZER ####

N = 5
X = (OverallQual  = rand(UInt8, N),
     GrLivArea    = rand(N),
     Neighborhood = categorical(rand("abc", N)),
     x1stFlrSF    = rand(N),
     TotalBsmtSF  = rand(N))

# introduce a field of type `Char`:
x1 = categorical(map(Char, (X.OverallQual |> collect)))

# introduce field of Int type:
x4 = [round(Int, x) for x in X.x1stFlrSF]

X = (x1=x1, x2=X[2], x3=X[3], x4=x4, x5=X[5])

stand = Standardizer()
f,    = fit(stand, 1, X)
Xnew  = transform(stand, f, X)

@test Xnew[1] == X[1]
@test std(Xnew[2]) ≈ 1.0
@test Xnew[3] == X[3]
@test Xnew[4] == X[4]
@test std(Xnew[5]) ≈ 1.0

stand.features = [:x1, :x5]
f,   = fit(stand, 1, X)
Xnew = transform(stand, f, X)
f,   = fit(stand, 1, X)

@test issubset(Set(keys(f)), Set(Tables.schema(X).names[[5,]]))

Xt = transform(stand, f, X)

@test Xnew[1] == X[1]
@test Xnew[2] == X[2]
@test Xnew[3] == X[3]
@test Xnew[4] == X[4]
@test std(Xnew[5]) ≈ 1.0

infos = info_dict(stand)

@test infos[:name] == "Standardizer"
@test infos[:input_scitype] == MLJBase.Table(Scientific)
@test infos[:output_scitype] == MLJBase.Table(Scientific)

#### UNIVARIATE BOX COX TRANSFORMER ####

# create skewed non-negative vector with a zero value:
Random.seed!(1551)
v = abs.(randn(1000))
v = v .- minimum(v)

t  = UnivariateBoxCoxTransformer(shift=true)
f, = fit(t, 2, v)

@test sum(abs.(v - inverse_transform(t, f, transform(t, f, v)))) <= 5000*eps()

infos = info_dict(t)

@test infos[:name] == "UnivariateBoxCoxTransformer"
@test infos[:input_scitype] == AbstractVector{MLJBase.Continuous}
@test infos[:output_scitype] == AbstractVector{MLJBase.Continuous}

#### ONE HOT ENCODER ####

X = (name   = categorical(["Ben", "John", "Mary", "John"], ordered=true),
     height = [1.85, 1.67, 1.5, 1.67],
     favourite_number = categorical([7, 5, 10, 5]),
     age    = [23, 23, 14, 23])

t  = OneHotEncoder()
f, = @test_logs((:info, r"Spawning 3"), (:info, r"Spawning 3"), fit(t, 1, X))

Xt = transform(t, f, X)

@test Xt.name__John == float.([false, true, false, true])
@test Xt.height == X.height
@test Xt.favourite_number__10 == float.([false, false, true, false])
@test Xt.age == X.age
@test schema(Xt).names == (:name__Ben, :name__John, :name__Mary,
                           :height, :favourite_number__5,
                           :favourite_number__7, :favourite_number__10,
                           :age)

# test that *entire* pool of categoricals is used in fit, including
# unseen levels:
f, = @test_logs((:info, r"Spawning 3"), (:info, r"Spawning 3"),
                      fit(t, 1, MLJBase.selectrows(X,1:2)))
Xtsmall = transform(t, f, X)
@test Xt == Xtsmall

# test that transform can be applied to subset of the data:
@test transform(t, f, MLJBase.selectcols(X, [:name, :age])) ==
    MLJBase.selectcols(transform(t, f, X),
                       [:name__Ben, :name__John, :name__Mary, :age])

# test exclusion of ordered factors:
t  = OneHotEncoder(ordered_factor=false)
f, = fit(t, 1, X)
Xt = transform(t, f, X)
@test :name in Tables.schema(Xt).names
@test :favourite_number__5 in Tables.schema(Xt).names

# test that one may not add new columns:
X = (name       = categorical(["Ben", "John", "Mary", "John"], ordered=true),
     height     = [1.85, 1.67, 1.5, 1.67],
     favourite_number = categorical([7, 5, 10, 5]),
     age        = [23, 23, 14, 23],
     gender     = categorical(['M', 'M', 'F', 'M']))
@test_throws Exception transform(t, f, X)

infos = info_dict(t)

@test infos[:name] == "OneHotEncoder"
@test infos[:input_scitype] == MLJBase.Table(Scientific)
@test infos[:output_scitype] == MLJBase.Table(Scientific)

#### FILL IMPUTER ####

X = (
    x = [missing,ones(10)...],
    y = [missing,ones(10)...],
    z = [missing,ones(10)...]
    )

imp = FillImputer()
f,  = fit(imp, 1, X)
Xt  = transform(imp, f, X)
@test all(.!ismissing.(Xt.x))
@test Xt.x isa Vector{Float64} # no missing
@test all(Xt.x .== 1.0)

imp = FillImputer(features=[:x,:y])
f,  = fit(imp, 1, X)
@test_throws ErrorException transform(imp, f, X) # in X there's :z which we haven't trained on

X = (x = categorical([missing, missing, missing, missing, "Old", "Young", "Middle", "Young",
                      "Old", "Young", "Middle", "Young"]), )

mode_ = mode(["Old", "Young", "Middle", "Young", "Old", "Young", "Middle", "Young"])

imp = FillImputer()
f,  = fit(imp, 1, X)
Xt  = transform(imp, f, X)
@test all(.!ismissing.(Xt.x))
@test all(Xt.x[ismissing.(X.x)] .== mode_)

X  = (x = [missing, missing, 1, 1, 1, 1, 1, 5], )
f, = fit(imp, 1, X)
Xt = transform(imp, f, X)
@test Xt.x == [1, 1, 1, 1, 1, 1, 1, 5]

X = (x = categorical([missing, missing, missing, missing, "Old", "Young", "Middle", "Young",
                      "Old", "Young", "Middle", "Young"]),
     y = [missing, ones(11)...],
     z = [missing, missing, 1,1,1,1,1,5,1,1,1,1],
     a = rand("abc", 12))

f, = fit(imp, 1, X)
Xt = transform(imp, f, X)

@test all(.!ismissing.(Xt.x))
@test all(.!ismissing.(Xt.y))
@test all(.!ismissing.(Xt.z))
@test all(.!ismissing.(Xt.a))

@test Xt.x[1] == mode_
@test Xt.y[1] == 1
@test Xt.z[1] == 1

end
true


#
#
#   These methods are for testing purposes only.
#   They are not optimized at all.
#
#

"""
Calculate the condition number from the given singular values
"""
function LinearAlgebra.cond(svs::AbstractVector)
    mi, ma = extrema(svs)
    ma / mi
end

"""
    calc_Bchain(B,N) -> (R, svs)

Calculate B^N as B*B*B*B*...*B from right to left.

Returns the result `R` and the singular values (column of `svs`) of all intermediate powers (rows of `svs`).
"""
function calc_Bchain(B, N)
    bs = size(B,1)
    R = Matrix{eltype(B)}(I, bs, bs)
    svs = zeros(Float64, N,bs)
    svc = 1
    for k in 1:N
        R = B * R
        svs[svc,:] = log.(svdvals(R))
        svc += 1
    end
    return (R, svs)
end

"""
Same as `calc_Bchain` but stabilizes the
matrix products by intermediate matrix decompositions.
"""
function calc_Bchain_svd(B, N; svdalg = svd)
    svs = zeros(Float64, N,size(B,1))
    F = svdalg(B)
    U, S, Vt = F.U, F.S, F.Vt
    svs[1,:] = log.(S)
    svc = 2

    for k in 2:N
        # multiply B from the left to USVt
        # and update USVt
        U = B * U
        U *= Diagonal(S)
        M = svdalg(U)
        U, S, Vtnew = M.U, M.S, M.Vt
        Vt = Vtnew * Vt

        # keep singular values
        svs[svc,:] = log.(S)
        svc += 1
    end

    return SVD(U, S, Vt), svs
end

"""
Same as `calc_Bchain` but stabilizes the
matrix products by intermediate matrix decompositions.
"""
function calc_Bchain_qr(B, N)
    svs = zeros(Float64, N,size(B,1))
    U, D, T = udt(B)
    svs[1,:] = log.(D)
    svc = 2

    for k in 2:N
        # multiply B from the left to UDT
        # and update UDT
        U = B * U
        U *= Diagonal(D)
        U, D, Tnew = udt(U)
        T = Tnew * T

        # keep singular values
        svs[svc,:] = log.(D)
        svc += 1
    end

    return UDT(U, D, T), svs
end

"""
Calculate fake "G(τ, 0)", i.e. [B^-N1 + B^N2]^-1
"""
function calc_tdgf_qr(B, N1, N2; inv_sum_method = inv_sum)
  if N1 != 0
    udtl = calc_Bchain_qr(inv(B), N1)[1]
  else
    udtl = udt!(Matrix(I*one(eltype(B)), size(B)))
  end

  if N2 != 0
    udtr = calc_Bchain_qr(B, N2)[1]
  else
    udtr = udt!(Matrix(I*one(eltype(B)), size(B)))
  end

  inv_sum_method(udtl, udtr)
end

"""
Calculate fake "G(τ, 0)", i.e. [B^-N1 + B^N2]^-1
"""
function calc_tdgf_svd(B, N1, N2; inv_sum_method = inv_sum, svdalg_inv = svd!, svdalg_chain = svd)
  # svdalg_nonmutating = Symbol(replace(string(nameof(svdalg)), "!" => ""))
  if N1 != 0
    svdl = calc_Bchain_svd(inv(B), N1; svdalg = svdalg_chain)[1]
  else
    svdl = svdalg(Matrix(I*one(eltype(B)), size(B)))
  end

  if N2 != 0
    svdr = calc_Bchain_svd(B, N2; svdalg = svdalg_chain)[1]
  else
    svdr = svdalg(Matrix(I*one(eltype(B)), size(B)))
  end

  inv_sum_method(svdl, svdr; svdalg = svdalg_inv)
end


"""
**SeisMWNI**

*Minimum Weighted Norm Interpolation of seismic records.*

**IN**

* d_in: input data that can have up to 5 dimensions
* dt=0.001 sampling rate along the time axis (in seconds)
* fmax=99999. maximum temporal frequency to process.
* padt=2 padding to use for the time axis
* padx=1 padding to use for the spatial axes
* Niter_internal=10 number of internal iterations for Conjugate Gradients
* Niter_external=3 number of external iterations for iterative reweighting

**OUT**

* d_out: interpolated data
"""
function SeisMWNI(in;dt=0.001,fmax=99999.,padt=2,padx=1,Niter_internal=10,Niter_external=3,mu=0)

    nt = size(in,1)
    nx1 = size(in,2)
    nx2 = size(in,3)
    nx3 = size(in,4)
    nx4 = size(in,5)
    d = zeros(Float32,nt,nx1,nx2,nx3,nx4)
    d[1:nt,1:nx1,1:nx2,1:nx3,1:nx4] = in
    nf = padt*nextpow2(nt)
    dw = 2.*pi/nf/dt
    nw = round(Int,nf/2) + 1
    iw_max = fmax*dt*nf < nw ? round(Int,floor(fmax*dt*nf)) : round(Int,floor(0.5/dt))
    nx1 > 1 ? nk1 = padx*nextpow2(nx1) : nk1 = 1
    nx2 > 1 ? nk2 = padx*nextpow2(nx2) : nk2 = 1
    nx3 > 1 ? nk3 = padx*nextpow2(nx3) : nk3 = 1
    nx4 > 1 ? nk4 = padx*nextpow2(nx4) : nk4 = 1
    nk = nk1*nk2*nk3*nk4
    # generate sampling operator from the padded data
    T = CalculateSampling(d)
    T = T[1:1,:,:,:,:]
    wm = ones(T)
    if (sum(T[:])/length(T[:]) < 0.05)
	println(sum(T[:])/length(T[:]))
	return in
    else
	T = Pad5D(T,1,nk1,nk2,nk3,nk4)
	T = T[1,:,:,:,:]
	operators = [WeightingOp,FFTOp,WeightingOp]
	@compat parameters = [Dict(:w=>T),Dict(:normalize=>true),
                              Dict(:w=>ones(T))]
	d = Pad5D(d,nf,nk1,nk2,nk3,nk4)
	D = fft(d,1)
	for iw=1:iw_max
	    x = D[iw,:,:,:,:]
	    y = copy(x)
	    Y,cost = IRLS(x,operators,parameters;Niter_external=Niter_external,
                          Niter_internal=Niter_internal,mu=mu)
	    y = FFTOp(Y,false)
	    D[iw,:,:,:,:] = y
	end
	# symmetries
	for iw=nw+1:nf
	    D[iw,:,:,:,:] = conj(D[nf-iw+2,:,:,:,:])
	end
	d = ifft(D,1)
	d = real(d[1:nt,1:nx1,1:nx2,1:nx3,1:nx4])
	return d
    end
end


struct LatentVariable
    user::Array{Float64,2}
    item::Array{Float64,2}
end

function svd(dataset::Persa.CFDatasetAbstract, features::Int)::LatentVariable
    matrix = Persa.getmatrix(dataset)
    (result, -) = svds(matrix, nsv = features)
    U = result.U
    V = result.Vt
    return LatentVariable(U, V')
end

Base.length(lv::LatentVariable) = maximum(size(lv))

user(lv::LatentVariable) = lv.user
item(lv::LatentVariable) = lv.item

user(lv::LatentVariable, id::Int) = user(lv)[id,:]
item(lv::LatentVariable, id::Int) = item(lv)[id,:]

Base.size(lv::LatentVariable) = (size(lv.user)[2], size(lv.item)[2])


using ImagePhaseCongruency, Images, PyPlot

img = circsine(offset = pi/4, ampexponent = -1.5);

imshow(img)

fimg = imfilter(img, KernelFactors.gaussian((2, 2)))

(grad_y, grad_x) = imgradients(fimg, KernelFactors.ando3)

imshow(grad_y)

orient = orientation(grad_x, grad_y)
mag = magnitude(grad_x, grad_y)

thinned, subpix = thin_edges_nonmaxsup(mag, orient, radius=1.35, theta=pi/180)

imshow(thinned)


"""
    simulate(μ, x)

Simulation of a dichotomous vector via the probit link.

# Arguments
- `μ::Vector{<:Real}`: real-valued vector of coefficients
- `x::Matrix{<:Real}`: real-valued vector or matrix of inputs

# Output
- `y::Vector{Bool}`: boolean vector of outcomes
- `p::Vector{float}`: vector of probabilities
- `z::Vector{float}`: vector of latent variables
"""

function simulate(β::Vector{T} where T<:Real, α::Real, X::Matrix{T} where T<:Real)
    N, J = size(X)
    @assert length(β) == J "μ is a $(length(β)) vector and X is a $Nx$J matrix."
    @assert α > 0 "σ must be a positive scalar."
    λ = lambda(β, α, X)
    Time = rand.(Weibull.(α, λ))
    Censor = rand.(Weibull.(α, λ))
    Indicator = Time .> Censor
    Time = ifelse.(Indicator, Censor, Time)
    return (Time, Indicator, λ, β, α)
end


hi = 61
car = carmichael(hi)

curp = tcnt = 0
print("Carmichael 3 (p×q×r) pseudoprimes, up to p = $hi:")
for j in sortperm(1:size(car)[2], by=x->(car[1,x], car[2,x], car[3,x]))
    p, q, r = car[:, j]
    c = prod(car[:, j])
    if p != curp
        curp = p
        @printf("\n\np = %d\n  ", p)
        tcnt = 0
    end
    if tcnt == 4
        print("\n  ")
        tcnt = 1
    else
        tcnt += 1
    end
    @printf("p× %d × %d = %d  ", q, r, c)
end
println("\n\n", size(car)[2], " results in total.")


"""
  TCG is a trust region sub-problem solver that approximately solves the quadratic
  programming problem subject to the solution being in a Euclidean Ball. It is
  appropriate for all types of Hessians. It solves the problem by using a conjugate
  gradient approach to solving the unconstrained quadratic optimization problem. If
  an interior solution exists it will be found according to the usual CG theory. If the
  problem is unbounded (negative curvature) or the solution is outside of the trust region
  we step to the boundary.

  TCG accepts all Hessians and is therefore well-suited for Newton's method for
  any vexity. It is relatively cheap for (very) large problems.

  The implementation follows Algorithm 7.5.1 on p. 205 of [ConnGouldTointBook] and
  the algorithm on p. 628 in [Steihaug1983]
"""
struct TCG <: TRSPSolver
end
summary(::TCG) = "Steihaug-Toint Truncated CG"

function (ms::TCG)(∇f, H, Δ::T, s, scheme, λ0=0; abstol=1e-10, maxiter=min(5,length(s)), κeasy=T(1)/10, κhard=T(2)/10) where T
    M = I
    # We only know that TCG has its properties if we start with s = 0
    s .= T(0)
    # g is the gradient of the quadratic model with the step α*p
    if all(iszero, ∇f)
        return tr_return(;λ=NaN, ∇f=∇f, H=H, s=s, interior=true, solved=true, hard_case=false, Δ=Δ)
    end
    g = copy(∇f)
    v = M\g
    p = -v
    for i = 1:maxiter
        # Check for negative curvature
        κ = dot(p, H*p)

        c = dot(s, M*s)
        b′ = dot(s, M*p) # b/2
        a = dot(p, M*p)

        if κ <= 0
          	# This branch just catches that sigma can be nan if a is exactly 0.
        	  if !iszero(a)
	              σ = (-b′ + sqrt(b′^2 + a*(Δ^2 - c)))/a
                @. s = s + σ*p
            end
            return tr_return(;λ=NaN, ∇f=∇f, H=H, s=s, interior=false, solved=true, hard_case=false, Δ=Δ)
        end

        α = dot(g, v)/κ

        if a*α^2 + α*2*b′ + c >= Δ^2
            σ = (-b′ + sqrt(b′^2 + a*(Δ^2 - c)))/a
            @. s = s + σ*p
            return tr_return(;λ=NaN, ∇f=∇f, H=H, s=s, interior=false, solved=true, hard_case=false, Δ=Δ)
        end
        @. s = s + α*p
        den = dot(g, v)
        g .= g .+ α.*(H*p)
        v .= M\g
        β = dot(g, v)/den
        @. p = -v + β*p
    end
    tr_return(;λ=NaN, ∇f=∇f, H=H, s=s, interior=true, solved=true, hard_case=false, Δ=Δ)
end


@testset "Group algebras" begin
  G = small_group(8, 4)
  A = AlgGrp(FlintQQ, G)

  @testset "Regular matrix algebra" begin
    B, BtoA = Hecke.regular_matrix_algebra(A)

    @test dim(B) == dim(A)

    for i = 1:dim(A)
      for j = 1:dim(A)
        @test BtoA(B[i])*BtoA(B[j]) == BtoA(B[i]*B[j])
        @test (BtoA\A[i])*(BtoA\A[j]) == BtoA\(A[i]*A[j])
      end
    end
  end

  @testset "Generators" begin
    g, full_basis, v = gens(A, Val{true})

    @test length(full_basis) == dim(A)

    M = zero_matrix(FlintQQ, dim(A), dim(A))
    for i = 1:dim(A)
      Hecke.elem_to_mat_row!(M, i, full_basis[i])
    end
    @test rank(M) == dim(A)

    for i = 1:dim(A)
      b = full_basis[i]
      a = one(A)
      for (j, k) in v[i]
        a *= g[j]^k
      end
      @test b == a
    end
  end

end


module DuffingVanDerPol

using DiffEqBase: ODEProblem
using Parameters: @with_kw, @unpack
using Setfield: @lens

using ...Bifurcations: BifurcationProblem

@with_kw struct DuffingVanDerPolParam{A, B, D}
    a::A = 1.0
    b::B = 0.0
    d::D = 1.0
end

function f(du, u, p, t)
    @unpack a, b, d = p
    du[1] = u[2]
    du[2] = - a * u[1] - b * u[1]^3 - d * (u[1]^2 - 1) * u[2]
    nothing
end

make_prob(
        p = DuffingVanDerPolParam();
        u0 = [0.0, 0.0],
        tspan = (0.0, 2π),
        ode = ODEProblem(f, u0, tspan, p),
        param_axis = (@lens _.d),
        t_domain = (-2.0, 8.0),
        kwargs...) =
    BifurcationProblem(ode, param_axis, t_domain;
                       kwargs...)

prob = make_prob()
ode = prob.p.de_prob
param_axis = prob.p.param_axis
t_domain = prob.t_domain

end  # module


Pkg.update()
Pkg.add("ScikitLearn")
Pkg.add("PyPlot")

using ScikitLearn
using PyPlot

@sk_import datasets: (make_circles, make_moons, make_blobs)
@sk_import cluster: (estimate_bandwidth, MeanShift, MiniBatchKMeans, AgglomerativeClustering, SpectralClustering)
@sk_import cluster: (DBSCAN, AffinityPropagation, Birch)
@sk_import preprocessing: StandardScaler
@sk_import neighbors: kneighbors_graph

srand(33)

# Generate datasets. We choose the size big enough to see the scalability
# of the algorithms, but not too big to avoid too long running times
n_samples = 1500
noisy_circles = make_circles(n_samples=n_samples, factor=.5, noise=.05)
noisy_moons = make_moons(n_samples=n_samples, noise=.05)
blobs = make_blobs(n_samples=n_samples, random_state=8)
no_structure = rand(n_samples, 2), nothing

colors0 = collect("bgrcmykbgrcmykbgrcmykbgrcmyk")
colors = vcat(fill(colors0, 20)...)

clustering_names = [
    "MiniBatchKMeans", "AffinityPropagation", "MeanShift",
    "SpectralClustering", "Ward", "AgglomerativeClustering",
    "DBSCAN", "Birch"];

figure(figsize=(length(clustering_names) * 2 + 3, 9.5))
subplots_adjust(left=.02, right=.98, bottom=.001, top=.96, wspace=.05,
                    hspace=.01)

plot_num = 1

datasets = [noisy_circles, noisy_moons, blobs, no_structure]

for (i_dataset, dataset) in enumerate(datasets)

    X, y = dataset
    # normalize dataset for easier parameter selection
    X = fit_transform!(StandardScaler(), X)

    # estimate bandwidth for mean shift
    bandwidth = estimate_bandwidth(X, quantile=0.3)

    # connectivity matrix for structured Ward
    connectivity = kneighbors_graph(X, n_neighbors=10,
                include_self=false)[:todense]()

    # PyCall does not support numpy sparse matrices
    # make connectivity symmetric
    connectivity = 0.5 * (connectivity + connectivity')

    # create clustering estimators
    ms = MeanShift(bandwidth=bandwidth, bin_seeding=true)
    two_means = MiniBatchKMeans(n_clusters=2)
    ward = AgglomerativeClustering(n_clusters=2, linkage="ward",
                                   connectivity=connectivity)
    spectral = SpectralClustering(n_clusters=2,
                                  eigen_solver="arpack",
                                  affinity="nearest_neighbors")
    dbscan = DBSCAN(eps=.2)
    affinity_propagation = AffinityPropagation(damping=.9, preference=-200)

    average_linkage = AgglomerativeClustering(
        linkage="average", affinity="cityblock", n_clusters=2,
        connectivity=connectivity)

    birch = Birch(n_clusters=2)
    clustering_algorithms = [
        two_means, affinity_propagation, ms, spectral, ward, average_linkage,
        dbscan, birch]

    for (name, algorithm) in zip(clustering_names, clustering_algorithms)
        fit!(algorithm, X)
        y_pred = nothing
        try
            y_pred = predict(algorithm, X)
        catch e
            if isa(e, KeyError)
                y_pred = map(Int, algorithm[:labels_])
                clamp!(y_pred, 0, 27) # not sure why some algorithms return -1
            else rethrow() end
        end
        subplot(4, length(clustering_algorithms), plot_num)
        if i_dataset == 1
            title(name, size=18)
        end

        for y_val in unique(y_pred)
            selected = y_pred.==y_val
            scatter(X[selected, 1], X[selected, 2], color=string(colors0[y_val+1]), s=10)
        end

        xlim(-2, 2)
        ylim(-2, 2)
        xticks(())
        yticks(())
        plot_num += 1
    end
end


# Generate a time series of stock prices by repeatedly solving the SPDE
# from Mastromatteo et al. (2014) to model the latent order book.
# using the previously generated mid price as the new initial mid price and
# boundary mid point during each iteration to generate a sequence of T prices,
# where T is the number of measured prices stored in the loaded
# ObjectiveFunction object.

mutable struct ReactionDiffusionPricePaths
    num_paths::Int64
    T::Int64
    p₀::Float64
    M::Int64
    β::Float64
    L::Float64
    D::Float64
    σ::Float64
    nu::Float64
    source_term::SourceTerm
    x::Array{Float64, 1}
    Δx::Float64
    Δt::Float64
end
function ReactionDiffusionPricePaths(num_paths, T, p₀, M, β,
    L, D, σ, nu, source_term)
    x₀ = p₀ - 0.5*L
    xₘ = p₀ + 0.5*L
    @assert x₀ >= 0
    # if D < 1.5 * σ
    #     @warn "D=$D is less than 1.5 times σ=$σ which runs the risk of unstable LOB ghost points"
    # end
    # Δx = L/M
    # if (β*Δx/D < log(1/3)/2) | (β*Δx/D > log(3)/2)
    #     @warn "The choice of β, D, L, M might result in jump probabilities > 0.75, thus causing large shifts in the LOB density"
    # end

    x = collect(Float64, range(x₀, xₘ, length=M+1))
    Δx = L/M
    Δt = (Δx^2) / (2.0*D)
    return ReactionDiffusionPricePaths(num_paths, T, p₀, M, β,
        L, D, σ, nu, source_term, x, Δx, Δt)
end

function (rdpp::ReactionDiffusionPricePaths)(seed::Int=-1)
    if seed == -1
            seeds = Int.(rand(MersenneTwister(), UInt32, rdpp.num_paths))
        else
            seeds = Int.(rand(MersenneTwister(seed), UInt32, rdpp.num_paths))
    end

    Δx = rdpp.L/rdpp.M
    Δt = (Δx^2) / (2.0*rdpp.D)
    time_steps = floor(Int ,rdpp.T / Δt)

    raw_price_paths = ones(Float64, time_steps + 1, rdpp.num_paths)
    raw_price_paths[1, :] .= rdpp.p₀
    sample_price_paths = ones(Float64, rdpp.T + 1, rdpp.num_paths)
    sample_price_paths[1, :] .= rdpp.p₀
    lob_densities = zeros(Float64, rdpp.M+1, time_steps + 1, rdpp.num_paths)
    P⁺s = ones(Float64, time_steps, rdpp.num_paths)
    P⁻s = ones(Float64, time_steps, rdpp.num_paths)

    for path in 1:rdpp.num_paths
         Random.seed!(seeds[path])
        lob_densities[:, :, path], raw_price_paths[:, path],
            sample_price_paths[:, path], P⁺s[:, path],
            P⁻s[:, path]  = dtrw_solver(rdpp)
    end

    return lob_densities, raw_price_paths, sample_price_paths, P⁺s, P⁻s
end

ReactionDiffusionPricePaths(dict)=ReactionDiffusionPricePaths(
    dict["num_paths"], dict["T"], dict["p₀"], dict["M"],
    dict["β"], dict["L"], dict["D"],
    dict["σ"], dict["nu"], SourceTerm(dict["λ"], dict["μ"]))

ReactionDiffusionPricePaths(;num_paths=1, T::Int64=100,
    p₀::Real=100.0, M::Int64=100, β::Real=1.0,
    L::Real=50.0, D::Real=4.0, σ::Real=1.0,
    nu::Real=0.0, λ::Real=1.0, μ::Real=0.5) =
    ReactionDiffusionPricePaths(num_paths, T, p₀,
    M, β, L, D, σ, nu, SourceTerm(λ, μ))


@testset "Truncated SVD" begin
    Q = diagm([2.0^(-j) for j in 1:100])
    tol = sum(2.0^(-j) for j in 27:100)

    Q_red = truncated_svd(Q, tol = tol)
    @test rank(Q_red) == 26
    @test norm(Matrix(Q_red) - diagm([j <= 26 ? 2.0^(-j) : 0 for j in 1:100])) < 1e-9

    Q_red = truncated_svd(Q, rmin = 30)
    @test rank(Q_red) == 30
    @test norm(Matrix(Q_red) - diagm([j <= 30 ? 2.0^(-j) : 0 for j in 1:100])) < 1e-9

    Q_red = truncated_svd(Q, rmax = 12)
    @test rank(Q_red) == 12
    @test norm(Matrix(Q_red) - diagm([j <= 12 ? 2.0^(-j) : 0 for j in 1:100])) < 1e-9

    Q_red = truncated_svd(Q, TSVD(), rmax = 12)
    @test rank(Q_red) == 12
    @test norm(Matrix(Q_red) - diagm([j <= 12 ? 2.0^(-j) : 0 for j in 1:100])) < 1e-9

    Q_red = truncated_svd(Q, 12)
    @test rank(Q_red) == 12
    @test norm(Matrix(Q_red) - diagm([j <= 12 ? 2.0^(-j) : 0 for j in 1:100])) < 1e-9

    Q_red = truncated_svd(Q, 12, TSVD())
    @test rank(Q_red) == 12
    @test norm(Matrix(Q_red) - diagm([j <= 12 ? 2.0^(-j) : 0 for j in 1:100])) < 1e-9
end

@testset "Rounding SVDLikeRepresentation" begin
    Q = qr(randn(100,100)).Q
    D = Diagonal([2.0^(-j) for j in 1:100])
    A = exp(Matrix(Q))*D*exp(Matrix(transpose(Q)))
    LRA = truncated_svd(A, tol = 0.01)

    # svd based rounding
    LRA_rounded = round(LRA + LRA)
    @test typeof(LRA_rounded) <: SVDLikeRepresentation
    @test rank(LRA_rounded) == rank(2*LRA)
    @test norm(Matrix(LRA_rounded) - Matrix(2*LRA)) <= 1e-8

    LRA_rounded = round(LRA .* LRA)
    @test rank(LRA_rounded) == rank(LRA.^2)
    @test norm(Matrix(LRA_rounded) - Matrix(LRA.^2)) <= 1e-8

    # tsvd based rounding
    LRA_tsvd_rounded = round(LRA .* LRA, TSVD(), rmax = 5)
    LRA_svd_rounded = round(LRA .* LRA, rmax = 5)
    @test typeof(LRA_tsvd_rounded) <: SVDLikeRepresentation
    @test norm(Matrix(LRA_tsvd_rounded) - Matrix(LRA_svd_rounded)) <= 1e-8
end

@testset "Rounding TwoFactorRepresentation" begin
    Q = qr(randn(100,100)).Q
    D = Diagonal([2.0^(-j) for j in 1:100])
    A = exp(Matrix(Q))*D*exp(Matrix(transpose(Q)))
    LRA = TwoFactorRepresentation(truncated_svd(A, tol = 0.01))

    LRA_rounded = round(LRA + LRA)
    @test typeof(LRA_rounded) <: TwoFactorRepresentation
    @test rank(LRA_rounded) == rank(2*LRA)
    @test norm(Matrix(LRA_rounded) - Matrix(2*LRA)) <= 1e-8

    LRA_rounded = round(LRA .* LRA)
    @test rank(LRA_rounded) == rank(LRA.^2)
    @test norm(Matrix(LRA_rounded) - Matrix(LRA.^2)) <= 1e-8

    # tsvd based rounding
    LRA_tsvd_rounded = round(LRA .* LRA, TSVD(), rmax = 5)
    LRA_svd_rounded = round(LRA .* LRA, rmax = 5)
    @test typeof(LRA_tsvd_rounded) <: TwoFactorRepresentation
    @test norm(Matrix(LRA_tsvd_rounded) - Matrix(LRA_svd_rounded)) <= 1e-8
end


function uhlig06()
#
#  test polynomial used by F. Uhlig
#
    p = [1.0;0;0;0;0;0;0;0;-1];
    z = [-1.00000000000000
         -0.70710678118655 + 0.70710678118655im
         -0.70710678118655 - 0.70710678118655im
          0 + 1.00000000000000im
          0 - 1.00000000000000im
          1.00000000000000
          0.70710678118655 + 0.70710678118655im
          0.70710678118655 - 0.70710678118655im];
    z = [z ones(8)]

    p, PolyZeros(z)
end


# This is the baseline distribution (i.e. a guassian process or iid uniform)
abstract type TimeseriesDistribution end

## Time series that are independent and identically distributed
struct IID{D} <: TimeseriesDistribution where {D}
    x::Array{Float64} # Sample locations
    distribution::D # Sampling distributions
end

# Number of points in the time series
N_pts(t::TimeseriesDistribution) = length(t.x)

# logpdf is just the summation over the individual values
Distributions.logpdf(iid::IID{D}, y) where {D} = sum([logpdf(iid.distribution, yi) for yi in y])

# Uniform distribution
function Base.rand(rng::AbstractRNG, iid::IID{Uniform{Float64}}, lb::Array{Float64} = -Inf*ones(N_pts(iid)), ub::Array{Float64} = Inf*ones(N_pts(iid)))
    samples = Array{Float64}(undef, length(iid.x))
    l, u = max.(lb, iid.distribution.a),  min.(ub, iid.distribution.b)
    for i=1:length(iid.x)
        samples[i] = l[i] == u[i] ? l[i] : rand(rng, Uniform(l[i], u[i]))
    end
    samples
end

# Gaussian Distribution
function Base.rand(rng::AbstractRNG, iid::IID{Normal{Float64}}, lb::Array{Float64} = -Inf*ones(N_pts(iid)), ub::Array{Float64} = Inf*ones(N_pts(iid)))
    samples = Array{Float64}(undef, length(iid.x))
    for i=1:length(iid.x)
        samples[i] = lb[i] == ub[i] ? lb[i] : rand(rng, Truncated(iid.distribution, lb[i], ub[i]))
    end
    samples
end

# Categorical
function Base.rand(rng::AbstractRNG, iid::IID{Categorical{Float64, Array{Float64,1}}}, feasible::Array{Array{Int64,1},1} = fill(collect(1:length(iid.distribution.p)), N_pts(iid)))
    samples = Array{Int64}(undef, length(iid.x))
    for i=1:length(iid.x)
        probs = zeros(length(iid.distribution.p))
        probs[feasible[i]] = iid.distribution.p[feasible[i]]
        d = Categorical(probs ./ sum(probs))
        samples[i] = rand(rng, d)
    end
    samples
end

# Bernoulli
function Base.rand(rng::AbstractRNG, iid::IID{Bernoulli{Float64}}, feasible::Array{Array{Int64,1},1} = fill([0,1], N_pts(iid)))
    samples = Array{Bool}(undef, length(iid.x))
    for i=1:length(iid.x)
        if length(feasible[i]) == 1
            samples[i] = Bool(feasible[i][1])
        else
            samples[i] = rand(rng, iid.distribution)
        end
    end
    samples
end

## Time series distributed according to a Gaussian Processs
squared_exp_kernel(;l=2, σ2 = 1) = (x,xp) -> σ2*exp(-(x-xp)^2/(2*l^2))

# Construct the mean vector
μ(X::Array{Float64}, m) = [m(x) for x in X]

# Construct the covariance matrix
K(X::Array{Float64}, X′::Array{Float64}, k) = [k(x,x′) for x in X, x′ in X′]

# Struct defining a Guassian Process
@with_kw struct GaussianProcess <: TimeseriesDistribution
    m::Function # Mean function (x_i) -> μ_i
    k::Function # Kernel function (x_i, x_j) -> k_ij
    x::Array{Float64} # Sample locations
    σ2::Float64 = 1e-6 # std deviation of noise on the samples
end

# Generate a random sample from the gaussian process
function Base.rand(rng::AbstractRNG, gp::GaussianProcess, lb::Array{Float64} = -Inf*ones(N_pts(gp)), ub::Array{Float64} = Inf*ones(N_pts(gp)))
    Xν, Xs, X, y = gp.x, gp.x, Float64[], Float64[]

    # Convert equality constraints to known points
    eq = lb .≈ ub
    l, u = copy(lb), copy(ub)
    X, y = Float64[X..., Xν[eq]...], Float64[y..., l[eq]...]
    l[eq] .= -Inf
    u[eq] .= Inf

    # Setup gram matrices
    N, Ns, Nν = length(X), length(Xs), length(Xν)
    Kxx = K(X, X, gp.k)
    Kxxν = K(X, Xν, gp.k)
    Kxxs = K(X, Xs, gp.k)
    Kxνxν = K(Xν, Xν, gp.k)
    Kxsxν = K(Xs, Xν, gp.k)
    Kxsxs = K(Xs, Xs, gp.k)
    Ixx = Matrix{Float64}(I, N, N)
    Ixsxs = Matrix{Float64}(I, Ns, Ns)
    Ixνxν = Matrix{Float64}(I, Nν, Nν)

    # Compute cholesky factors and important matrices
    L = cholesky(Kxx + gp.σ2*Ixx).L
    v1 = L \ Kxxν
    v2 = L \ Kxxs
    A1 = (L' \ v1)'
    A2 = (L' \ v2)'
    B1 = Kxνxν + gp.σ2*Ixνxν - v1' * v1
    B2 = Kxsxs + gp.σ2*Ixsxs - v2' * v2
    B3 = Kxsxν - v2' * v1

    L1 = cholesky(B1).L
    v3 = L1 \ B3'
    A = (L1' \ v3)'
    B = A2 - A * A1
    Σ = B2 - v3' * v3

    # Get mean vectors
    μo = μ(X, gp.m)
    μs = μ(Xs, gp.m)
    μν = μ(Xν, gp.m)

    # Sample from the posterior
    Q = cholesky(Σ).L
    Cm = mvrandn_μ(rng, μν + A1*(y - μo), l, u, B1, 1)
    Um = rand(rng, Normal(), (Ns, 1))

    dropdims(μs + B*(y .- μo) .+ (A*(-μν .+ Cm) + Q*Um), dims=2)
end

# Get the log probability of the sample point
function Distributions.logpdf(gp::GaussianProcess, y::Array{Float64})
    μ_vals = μ(gp.x, gp.m)
    Σ = K(gp.x, gp.x, gp.k) .+ gp.σ2*Matrix(I, N_pts(gp), N_pts(gp))
    logpdf(MvNormal(μ_vals, Σ), y)
end


#------------------------------
#   Solve for the equilibrium interest rate in the Hugget economy with a general diffusion process.
#
#  First, solve the HJB equation to get the A ("transition") matrix, then
#   solve the discretized KF equations. Then iterate over interest rates to
#   get the asset supply function. Update interest rate until convergence.
#
#   This code uses an implicit method for the solution: use
#   an 'upwind' scheme consisting of forwards and backwards difference
#   approximation to the derivative of the value function.
#
#   Assume: CRRA utility function u(c) = c^(1-s)/(1-s)
#   Note that r < rho
#
#   Written by Ben Moll et al (2015). Optimized for speed by SeHyoun Ahn. Translation into Julia by James Graham
#------------------------------

type Huggett_model
  gamma::Float64            # risk aversion coefficient
  rho::Float64              # discount rate.
  ou_var::Float64           # Ornstein Uhlenbeck variance
  zmean::Float64            # Mean of log-normal distriution for the Ornstein Uhlenbeck process
  ou_corr::Float64          # Ornstein Uhlenbeck correlation
  theta::Float64            # Ornstein Uhlenbeck param
  sig2::Float64            # Ornstein Uhlenbeck variance
  I::Int64                  # Number of grids for assets
  J::Int64                  # Number of grids for income
  Ir::Int64                 # Number of interest rate iterations
  amin::Float64             # Min assets = borrowing constraint
  amax::Float64             # Max assets
  zmin::Float64             # Min income
  zmax::Float64             # Max income
  rmin::Float64             # Min interest rate
  rmax::Float64             # Max interest rate
  agrid::Array{Float64}    # Grid space for assets
  zgrid::Array{Float64}    # Grid space for incoomes
  aa::Array{Float64}        # Matrix for asset grid space
  zz::Array{Float64}        # Matrix for income grid space
  da::Float64               # Asset grid step size
  dz::Float64               # Income grid step size
  dz2::Float64              # Income grid step size squared
  Delta::Float64            # Time step-size
  mu::Array{Float64}       # Drift from Ito's lemma
  s2::Array{Float64}       # Std Dev from Ito's lemma
  chi::Array{Float64}      # elements in the transition matrix
  nu::Array{Float64}       # elements in the transition matrix
  zeta::Array{Float64}     # elements in the transition matrix
  Cswitch::SparseMatrixCSC{Float64,Int64}  # Transition matrix
  v0::Array{Float64}       # Initial guess for value function
  r0::Float64                 # Intial guess for interest rate
  maxit::Int64              # Max iterations for solving HJB
  crit::Float64
  crit_S::Float64           # Tolerance for asset mkt clearing
end

# Create model function
function Huggett_model(;gamma::Float64=2.0,
  rho::Float64=0.05,
  ou_var::Float64=0.1,
  ou_corr::Float64=0.5,
  I::Int64=100,
  J::Int64=100,
  Ir::Int64=40,
  amin::Float64=-0.15,
  amax::Float64=2.0,
  rmin::Float64=-0.05,
  rmax::Float64=0.045,
  maxit::Int64=100,
  crit::Float64=1e-6,
  crit_S::Float64=1e-5,
  Delta::Float64=1000.0,
  r0::Float64=0.03)

  ## Ornstein-Uhlenbeck: dlog(z) = -theta*log(z)dt + sig*dW
  zmean   = exp(ou_var)
  theta   = -log(ou_corr)
  sig2    = theta*ou_var
  ## Ornstein-Uhlenbeck: dz = theta(zmean - z)dt + sig*dW
  #theta   = -log(ou_corr)
  #ou_var  = sig2/(theta)
  #zmean   = 1

  zmin    = zmean*0.65
  zmax    = zmean*1.25
  agrid   = reshape(collect(linspace(amin,amax,I)),I,1)
  zgrid   = reshape(collect(linspace(zmin,zmax,J)),J,1)
  aa      = repmat(agrid,1,J)
  zz      = repmat(zgrid',I,1)
  da      = (amax-amin)/(I-1)
  dz      = (zmax-zmin)/(J-1)
  dz2     = dz^2
  ## Ornstein-Uhlenbeck: in logs
  mu      = (-theta*log(zgrid) + sig2/2).*zgrid
  s2      = sig2.*zgrid.^2
  ## Ornstein-Uhlenbeck: in LEVELS
  #mu      = theta*(zmean - zgrid)
  #s2      = sig2.*ones(1,J)

  chi     = s2/(2*dz2)
  nu      = - s2/dz2 - mu/dz
  zeta    = mu/dz + s2/(2*dz2)

  v0      = (zz + r0.*aa).^(1-gamma)/(1-gamma)/rho

  #This will be the upperdiagonal of the C_switch
  updiag    = zeros(I,1)
  for j in 1:J-1
      updiag=vcat(updiag, repmat([zeta[j]],I,1))
  end

  #This will be the center diagonal of the B_switch
  centdiag = repmat([chi[1]+nu[1]],I,1)
  for j in 2:J-1
      centdiag = vcat(centdiag, repmat([nu[j]],I,1))
  end
  centdiag = vcat(centdiag, repmat([nu[J]+zeta[J]],I,1))

  #This will be the lower diagonal of the B_switch
  lowdiag = repmat([chi[2]],I,1)
  for j in 3:J
      lowdiag = vcat(lowdiag, repmat([chi[j]],I,1))
  end

  # Add upper, center, and lower diagonal into a sparse matrix
  Cswitch     = spdiagm(centdiag[:,1],0,I*J,I*J) + spdiagm(lowdiag[:,1],-I,I*J,I*J) + spdiagm(updiag[I+1:end,1],I,I*J,I*J)

  Huggett_model(gamma, rho, ou_var, zmean, ou_corr, theta, sig2, I, J, Ir, amin, amax, zmin, zmax, rmin, rmax, agrid, zgrid, aa, zz, da, dz, dz2, Delta, mu, s2, chi, nu, zeta, Cswitch, v0, r0, maxit, crit, crit_S)
end

# Creates an instance of the model
mod = Huggett_model()

# Define Utility funciton
function util(x; gamma::Float64=mod.gamma)
  u = x.^(1-gamma)/(1-gamma)
  return u
end

# Define derivative of Utility funciton
function utilprime(x; gamma::Float64=mod.gamma)
  uprime = x.^(-gamma)
  return uprime
end

# Define inverse of derivative of utility function
function utilprimeinv(x; gamma::Float64=mod.gamma)
  uprimeinv = x.^(-1/gamma)
  return uprimeinv
end

# Function to compute the KFE
function KFE(A::SparseMatrixCSC{Float64,Int64}; I::Int64=mod.I, J::Int64=mod.J, da::Float64=mod.da, dz::Float64=mod.dz)
  AT          = A'
  b           = zeros(I*J,1)
  i_fix       = 1    # fix one value so that matrix is not singular
  b[i_fix]    = 0.1
  row         = hcat(zeros(1,i_fix-1), 1, zeros(1,I*J-i_fix))
  AT[i_fix,:] = row
  #Solve linear system
  gg          = AT\b          # A'*gg = b --> gg = (A')^(-1)*b
  g_sum       = gg'*ones(I*J,1)*da*dz
  gg          = gg./g_sum     # Normalize so sum(gg)*da) = 1
  g           = reshape(gg,I,J)
  return g
end

# Initialization of value functions and consumption policy
Vaf         = zeros(mod.I,mod.J)
Vab         = zeros(mod.I,mod.J)
Vzf         = zeros(mod.I,mod.J)
Vzb         = zeros(mod.I,mod.J)
Vzz         = zeros(mod.I,mod.J)
c           = zeros(mod.I,mod.J)
r_r         = zeros(mod.Ir,1)
rmin_r      = zeros(mod.Ir,1)
rmax_r      = zeros(mod.Ir,1)
g_r         = zeros(mod.I,mod.J,mod.Ir)     # Stationary distribution
adot        = zeros(mod.I,mod.J,mod.Ir)    # Policy function
V_r         = zeros(mod.I,mod.J,mod.Ir)    # Value function
S           = zeros(mod.Ir,1)      # Aggregate asset supply
V           = zeros(mod.I,2)        # Value function
A           = sparse(zeros(mod.J*mod.I,mod.J*mod.I))   # Transition matrix
r           = mod.r0
dist        = zeros(mod.maxit,1)
r_end       = 0                   # Last iteration value

for ir in 1:mod.Ir
  r_r[ir]     = r
  rmin_r[ir]  = mod.rmin
  rmax_r[ir]  = mod.rmax

  if ir>1
    mod.v0 = V_r[:,:,ir-1]
  end

  v = mod.v0

  for n in 1:mod.maxit
    V               = v
    # forward difference
    Vaf[1:mod.I-1,:]    = (V[2:mod.I,:]-V[1:mod.I-1,:])/mod.da
    Vaf[mod.I,:]        = utilprime(mod.zgrid + r.*mod.amax)
    # backward difference
    Vab[1,:]            = utilprime(mod.zgrid + r.*mod.amin)  #state constraint boundary condition
    Vab[2:mod.I,:]      = (V[2:mod.I,:]-V[1:mod.I-1,:])/mod.da

    #consumption and savings with forward difference
    cf              = utilprimeinv(Vaf)
    sf             = mod.zz + r.*mod.aa - cf
    #consumption and savings with backward difference
    cb              = utilprimeinv(Vab)
    sb             = mod.zz + r.*mod.aa - cb
    #consumption and derivative of value function at steady state
    c0              = mod.zz + r.*mod.aa
    Va0             =utilprime(c0)

    # dV_upwind makes a choice of forward or backward differences based on the sign of the drift
    If              = sf .> 0    # positive drift --> forward difference
    Ib              = sb .< 0    # negative drift --> backward difference
    I0              = (1-If-Ib)   # at steady state
    #make sure backward difference is used at amax
    #Ib(I,:) = 1; If(I,:) = 0;
    #STATE CONSTRAINT at amin: USE BOUNDARY CONDITION UNLESS muf > 0:
    #already taken care of automatically

    Va_Upwind       = Vaf.*If + Vab.*Ib + Va0.*I0    #important to include third term
    c               = utilprimeinv(Va_Upwind)
    u               = util(c)

    #CONSTRUCT MATRIX
    X               = - min(sb,0)/mod.da
    Y               = - max(sf,0)/mod.da + min(sb,0)/mod.da
    Z               = max(sf,0)/mod.da

    updiag          = 0
    for j = 1:mod.J
        updiag      = vcat(updiag, Z[1:mod.I-1,j], 0)      # Note: Z_I,j = 0 for all j
    end
    updiag = updiag[2:end-1,1]

    centdiag = reshape(Y,mod.I*mod.J,1)

    lowdiag=X[2:mod.I,1]
    for j=2:mod.J
        lowdiag=vcat(lowdiag, 0, X[2:mod.I,j])          # Note: X_1,j = 0 for all j
    end

    B              = spdiagm(centdiag[:,1],0,mod.I*mod.J,mod.I*mod.J) + spdiagm(updiag[:,1],1,mod.I*mod.J,mod.I*mod.J) + spdiagm(lowdiag[:,1],-1,mod.I*mod.J,mod.I*mod.J)

    A              = B + mod.Cswitch
    InvMat         = (1/mod.Delta + mod.rho)*speye(mod.I*mod.J) - A

    u_stacked      = reshape(u,mod.I*mod.J,1)
    V_stacked      = reshape(V,mod.I*mod.J,1)

    b              = u_stacked + V_stacked/mod.Delta
    V_stacked      = InvMat\b  #SOLVE SYSTEM OF EQUATIONS
    V              = reshape(V_stacked,mod.I,mod.J)
    Vchange        = V - v
    v              = V

    dist[n] = maximum(maximum(abs(Vchange)))
    if dist[n]<mod.crit
      @printf("-----------------------\n")
      @printf("Interest rate = %f\n", r)
      @printf("Value Function Converged, Iteration = %i\n", n)
      break
    end
  end

  # Solve Kolmogorov Forward equation (Fokker-Plankck equation)
  g = KFE(A)
  g_r[:,:,ir]     = g
  adot[:,:,ir]    = mod.zz + r.*mod.aa - c
  V_r[:,:,ir]     = V

  S[ir]           = ((g'*mod.agrid*mod.da)'*ones(mod.I,1)*mod.dz)[1,1]

  #UPDATE INTEREST RATE
  r_end = ir

  if S[ir]>mod.crit_S
    @printf("Excess Supply at r = %0.2f\n", r)
    mod.rmax = r
    r = 0.5*(r+mod.rmin)
  elseif S[ir]<-mod.crit_S
    @printf("Excess Demand at r = %0.2f\n", r)
    mod.rmin = r
    r = 0.5*(r+mod.rmax)
  elseif abs(S[ir])<mod.crit_S
    @printf("Equilibrium Found, Interest rate = %0.2f\n", r)
    break
  end

end

## EXPORT DATA FOR GRAPHING IN PYTHON
using DataFrames
df = DataFrame(adot[:,:,r_end])
writetable("adot_diff.csv", df)

df = DataFrame(c)
writetable("cons_diff.csv", df)

df = DataFrame(g_r[:,:,r_end])
writetable("dist_diff.csv", df)

## GRAPHS
using Gadfly

# Distribution using a contour plot
plot(z=adot[:,:,r_end], x = mod.agrid, y = mod.zgrid, Geom.contour, Theme(line_width=1mm),Guide.title("Stationary distribution"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Income, z"),
Coord.cartesian(xmin=mod.amin, xmax=2), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))

# Savings Policy function
plot(layer(x=mod.agrid,y=adot[:,1,r_end],Geom.line,Theme(line_width=1.5mm)), layer(x=mod.agrid,y=adot[:,2,r_end],Geom.line, Theme(line_width=1.5mm,default_color=colorant"red")),
xintercept=[mod.amin], yintercept=[0], Geom.hline(color=colorant"black", size=0.5mm), Geom.vline(color=colorant"black", size=0.5mm),
Guide.title("Savings policy function"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Savings, si(a)"),
Coord.cartesian(xmin=mod.amin, xmax=mod.amax), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))

# Distributions
plot(layer(x=mod.agrid,y=g_r[:,50,r_end],Geom.line,Theme(line_width=1.5mm)), layer(x=mod.agrid,y=g_r[:,60,r_end],Geom.line, Theme(line_width=1.5mm,default_color=colorant"red")),
xintercept=[mod.amin], Geom.vline(color=colorant"black", size=0.5mm), yintercept=[0], Geom.hline(color=colorant"black", size=0.5mm),
Guide.title("Stationary distributions"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Density, gi(a)"),
Coord.cartesian(xmin=mod.amin, xmax=1), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))


using Test
using Complementarity

@testset "Complementarity Tests" begin

    include("../examples/mcp/transmcp.jl")
    include("transmcp_numeric.jl")

    @testset "Array{Variable,4} Test" begin
        m = MCPModel()
        @variable(m, x[1:2, 1:3, 1:4, 1:5] >=0 )
        @mapping(m, F[i=1:2, j=1:3, k=1:4, l=1:5], x[i,j,k,l])
        # for i=1:2, j=1:3, k=1:4, l=1:5
        #     @complementarity(m, F[i,j,k,l], x[i,j,k,l])
        # end
        @complementarity(m, F, x)
        status = solveMCP(m)
        @test status == :Solved
    end

    include("hhoeschle_mcp1.jl")

    include("runtests_mcp.jl")

    include("runtests_mcp_nlsolve.jl")

    include("runtests_mpec.jl")

    include("test_bounds.jl")
end


@testset "test_pid_design" begin
@test stabregionPID(s->exp(-sqrt(s)), exp10.(range(-3, stop=3, length=50)), kd=1)[2][1] ≈ -1.022356911142034

P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
gangoffourplot(ss(P),ss(1))
ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=true)
@test kp ≈ 0.82274734724854
@test ki ≈ -0.45472580052281375

C = pid(;kp = 1, ki = 1, kd=1)
@test C == tf(1) + tf(1,[1,0]) + tf([1,0],[1])
pidplots(C, :nyquist,:gof,:pz,:controller, grid=true) # Simply test that the functions runs and not errors
pidplots(C, :nyquist,:gof,:pz,:controller, grid=false)
leadlinkcurve()

stabregionPID(tf(1,[1,0]))
stabregionPID(s -> exp(-sqrt(s)))

P = tf(1,[1, 1, 1])
kp,ki,C = loopshapingPI(P,10; phasemargin = 30, doplot = false)
_,_,_,pm = margin(P*C)
@test pm[] > 30

P = tf(1,[1, 1])
piparams,C = placePI(P, 2, 0.7)
@test poles(feedback(P, C)) ≈ [-1.4 + √2.04im, -1.4 - √2.04im]
@test [piparams[:Kp], piparams[:Ti]] ≈ [9/5, 9/20]

end


export ambient_space, rank, gram_matrix, inner_product, involution,
       islocal_square, isequivalent, isrationally_equivalent, quadratic_space,
       hermitian_space, diagonal, invariants, hasse_invariant, witt_invariant, orthogonal_basis, fixed_field

################################################################################
#
#  Basic invariants
#
################################################################################

@doc Markdown.doc"""
    isquadratic(V::AbsSpace) -> Bool

Returns whether $V$ is quadratic.
"""
isquadratic(V::AbsSpace)

@doc Markdown.doc"""
    ishermitian(V::AbsSpace) -> Bool

Returns whether $V$ is hermitian.
"""
ishermitian(V::AbsSpace)

@doc Markdown.doc"""
    rank(V::AbsSpace) -> Int

Return the rank of the space `V`.
"""
rank(L::AbsSpace) = nrows(L.gram)

@doc Markdown.doc"""
    dim(V::AbsSpace) -> Int

Return the dimension of the space `V`.
"""
dim(V::AbsSpace) = rank(V)

@doc Markdown.doc"""
    gram_matrix(V::AbsSpace) -> MatElem

Return the Gram matrix of `V`.
"""
gram_matrix(V::AbsSpace) = V.gram

# Once we have quaternion spaces the following makes more sense

@doc Markdown.doc"""
    base_ring(V::AbsSpace) -> NumField

Return the base field of `V`.
"""
base_ring(V::AbsSpace) = _base_algebra(V)

@doc Markdown.doc"""
    fixed_field(V::AbsSpace) -> NumField

Return the fixed field of `V`.
"""
fixed_field(::AbsSpace)

@doc Markdown.doc"""
    involution(V::AbsSpace) -> NumField

Return the involution of `V`.
"""
involution(V::AbsSpace)

# Maybe cache this

@doc Markdown.doc"""
    isregular(V::AbsSpace) -> Bool

Return whether `V` is regular, that is, if the Gram matrix
has full rank.
"""
function isregular(V::AbsSpace)
  G = gram_matrix(V)
  return rank(G) == nrows(G)
end

################################################################################
#
#  Determinant and discriminant
#
################################################################################

function det(V::AbsSpace)
  d = det(gram_matrix(V))
  return fixed_field(V)(d)
end

@doc Markdown.doc"""
    det(V::AbsSpace) -> FieldElem

Returns the determinant of the space `V` as an element of the fixed field.
"""
det(::AbsSpace)

@doc Markdown.doc"""
    discriminant(V::AbsSpace) -> FieldElem

Returns the discriminant of the space `V` as an element of the fixed field.
"""
function discriminant(V::AbsSpace)
  d = det(V)
  n = mod(rank(V), 4)
  if n == 0 || n == 1
    return d
  else
    return -d
  end
end

function _discriminant(G)
  d = det(G)
  n = nrows(G)
  if mod(n, 4) == 0 || mod(n, 4) == 1
    return d
  else
    return -d
  end
end

################################################################################
#
#  Computing inner products
#
################################################################################

@doc Markdown.doc"""
    gram_matrix(V::AbsSpace, M::MatElem) -> MatElem

Returns the gram matrix of the rows of `M`.
"""
function gram_matrix(V::AbsSpace{T}, M::MatElem{S}) where {S, T}
  if S === elem_type(T)
    return M * gram_matrix(V) * transpose(_map(M, involution(V)))
  else
    Mc = change_base_ring(base_ring(V), M)
    return Mc * gram_matrix(V) * transpose(_map(Mc, involution(V)))
  end
end

@doc Markdown.doc"""
    gram_matrix(V::AbsSpace, S::Vector{Vector}) -> MatElem

Returns the gram matrix of the sequence `S`.
"""
function gram_matrix(V::AbsSpace{T}, S::Vector{Vector{U}}) where {T, U}
  m = zero_matrix(base_ring(V), length(S), rank(V))
  for i in 1:length(S)
    if length(S[i]) != rank(V)
      throw(error("Vectors must be of length $(rank(V))"))
    end
    for j in 1:rank(V)
      m[i, j] = S[i][j]
    end
  end
  return gram_matrix(V, m)
end

@doc Markdown.doc"""
    inner_product(V::AbsSpace, v::Vector, w::Vector) -> FieldElem

Returns the inner product of `v` and `w` with respect to `V`.
"""
inner_product(V::AbsSpace, v::Vector, w::Vector)

################################################################################
#
#  Diagonalization
#
################################################################################

@doc Markdown.doc"""
    orthogonal_basis(V::AbsSpace) -> MatElem

Returns a matrix `M`, such that the rows of `M` form an orthgonal basis of `V`.
"""
function orthogonal_basis(V::AbsSpace)
  _, B = _gram_schmidt(gram_matrix(V), involution(V))
  return B
end

@doc Markdown.doc"""
    diagonal(V::AbsSpace) -> Vector{FieldElem}

Returns a vector of elements $a_1,\dotsc,a_n$ such that `V` is isometric to
the diagonal space $\langle a_1,\dotsc,a_n \rangle$.

The elements will be contained in the fixed field of `V`.
"""
diagonal(V::AbsSpace)

################################################################################
#
#  Gram-Schmidt
#
################################################################################

# Clean this up
function _gram_schmidt(M::MatElem, a, nondeg = true)
  F = deepcopy(M)
  K = base_ring(F)
  n = nrows(F)
  S = identity_matrix(K, n)
  okk = isdiagonal(F)
  if !okk
    for i in 1:n
      if iszero(F[i,i])
        T = identity_matrix(K, n)
        ok = 0
        for j in (i + 1):n
          if !iszero(F[j, j])
            ok = j
            break
          end
        end
        #ok = findfirst(j -> !iszero(F[j, j]), (i + 1):n)
        if ok != 0 # ok !== nothing
          #j = ok + i # findfirst gives the index
          j = ok
          T[i,i] = 0
          T[j,j] = 0
          T[i,j] = 1
          T[j,i] = 1
        else
          ok = 0
          for j in (i + 1):n
            if !iszero(F[i, j])
              ok = j
              break
            end
          end
          if ok == 0
            if nondeg
              error("Matrix is not of full rank")
            end
          else
            j = ok
            T[i, j] = 1 // (2 * F[j, i])
          end
        end
        S = T * S
        F = T * F * transpose(_map(T, a))
      end
      T = identity_matrix(K, n)
      for j in (i + 1):n
        T[j, i] = divexact(-F[j, i], F[i, i])
      end
      F = T * F * transpose(_map(T, a))
      S = T * S
    end
    @assert isdiagonal(F)
  end
  return F, S
end

################################################################################
#
#  Local equivalence
#
################################################################################

@doc Markdown.doc"""
    isequivalent(L::AbsSpace, M::AbsSpace, p::Union{InfPlc, NfOrdIdl}) -> Bool

Returns whether `L` and `M` are equivalent over the completion at `p`.
"""
isequivalent(L::AbsSpace, M::AbsSpace, p)

################################################################################
#
#  Definitness
#
################################################################################

# Returns 0 if V is not definite
# Returns an element a != 0 such that a * canonical_basis of V has
# positive Gram matrix
function _isdefinite(V::AbsSpace)
  K = fixed_field(V)
  R = maximal_order(K)
  if !istotally_real(K) || (ishermitian(V) && !istotally_complex(K))
    return zero(K)
  end
  D = diagonal(V)
  signs_to_consider = Tuple{InfPlc, Int}[]
  for v in real_places(K)
    S = Int[sign(d, v) for d in D]
    if length(unique(S)) != 1
      return zero(K)
    else
      push!(signs_to_consider, (v, S[1]))
    end
  end
  if length(signs_to_consider) == 1
    return K(signs_to_consider[1][2])
  else
    return element_with_signs(K, signs_to_consider)
  end
end

function ispositive_definite(V::AbsSpace)
  E = base_ring(V)
  if isquadratic(V)
    K = E
  else
    K = base_field(E)
  end
  if !istotally_real(K)
    return false
  end
  D = diagonal(V)
  for d in D
    if !istotally_positive(d)
      return false
    end
  end
  return true
end

function isnegative_definite(V::AbsSpace)
  E = base_ring(V)
  if isquadratic(V)
    K = E
  else
    K = base_field(E)
  end

  if !istotally_real(K)
    return false
  end
  D = diagonal(V)
  for d in D
    if !istotally_positive(-d)
      return false
    end
  end
  return true
end

function isdefinite(V::AbsSpace)
  return ispositive_definite(V) || isnegative_definite(V)
end

################################################################################
#
#  Isotropic
#
################################################################################

isisotropic(V::AbsSpace, p::InfPlc) = _isisotropic(V, p)

function _isisotropic(V::AbsSpace, p::InfPlc)
  n = rank(V)
  d = det(V)
  E = base_ring(V)
  if d == 0
    return true
  elseif n <= 1
    return false
  elseif iscomplex(p)
    return true
  else
    D = diagonal(V)
    return length(unique!(Int[sign(d, p) for d in D])) == 2
  end
end

################################################################################
#
#  Helper functions (sort them later)
#
################################################################################

function image(f::NfRelToNfRelMor{T, T}, I::NfRelOrdIdl{T, S}) where {T, S}
  #f has to be an automorphism!!!!
  O = order(I)
  @assert ismaximal(O) # Otherwise the order might change
  K = nf(O)

  B = absolute_basis(I)

  if I.is_prime == 1
    lp = prime_decomposition(O, minimum(I))
    for (Q, e) in lp
      if I.splitting_type[2] == e
        if all(b -> f(b) in Q, B)
          return Q
        end
      end
    end
  end

  pb = pseudo_basis(I)
  pm = basis_pmatrix(I)

  m = zero(matrix(pm))

  c = coefficient_ideals(pm)

  for i in 1:length(pb)
    cc = coordinates(O(f(pb[i][1])))
    for j in 1:length(cc)
      m[i, j] = cc[j]
    end
  end

  J = ideal(O, pseudo_matrix(m, c))

  if isdefined(I, :minimum)
    J.minimum = I.minimum
  end

  J.has_norm = I.has_norm

  if isdefined(I, :norm)
    J.norm = I.norm
  end

  if isdefined(I, :is_prime)
    J.is_prime = I.is_prime
  end

  if isdefined(I, :splitting_type)
    J.splitting_type = I.splitting_type
  end

  return J
end

function image(f::NfRelToNfRelMor{T, T}, I::NfRelOrdFracIdl{T, S}) where {T, S}
  #S has to be an automorphism!!!!
  O = order(I)
  @assert ismaximal(O) # Otherwise the order might change
  K = nf(O)

  pb = pseudo_basis(I)

  z = sum(b * (f(a) * O) for (a, b) in pb)
  return z
end

# An element is locally a square if and only if the quadratic defect is 0, that is
# the valuation is inf.
# (see O'Meara, Introduction to quadratic forms, 3rd edition, p. 160)
function islocal_square(a, p)
  return quadratic_defect(a, p) isa PosInf
end

function _map(a::AbstractAlgebra.MatrixElem, f)
  z = similar(a)
  for i in 1:nrows(a)
    for j in 1:ncols(a)
      z[i, j] = f(a[i, j])
    end
  end
  return z
end

# I think I need a can_change_base_ring version

function element_with_signs(K, D::Dict{InfPlc, Int})
  return _element_with_signs(K, D)
end

function _element_with_signs(K, D)
  OK = maximal_order(K)
  G, mG = infinite_primes_map(OK, real_places(K), 1*OK)
  r = real_places(K)
  z = id(G)
  for (v, s) in D
    for i in 1:length(r)
      if s == 1
        ss = 0
      else
        ss = 1
      end

      if v == r[i]
        z = z + ss * G[i]
      end
    end
  end
  zz = elem_in_nf(mG(z))::elem_type(K)
  @assert all(u -> sign(zz, u[1]) == u[2], D)
  return zz
end

function element_with_signs(K, P::Vector{InfPlc}, S::Vector{Int})
  return _element_with_signs(K, zip(P, S))
end

function element_with_signs(K, A::Vector{Tuple{InfPlc, Int}})
  return _element_with_signs(K, A)
end

function prime_ideals_up_to(O::NfRelOrd, n::Union{Int, fmpz})
  p = 2
  z = ideal_type(O)[]
  while p <= n
    lp = prime_decomposition(base_ring(O), p)
    for q in lp
      if norm(q[1]) > n
        continue
      else
        lq = prime_decomposition(O, q[1])
        for Q in lq
          if absolute_norm(Q[1]) <= n
            push!(z, Q[1])
          end
        end
      end
    end
    p = next_prime(p)
  end
  return sort!(z, by = a -> absolute_norm(a))
end

################################################################################
#
#  Restriction of scalars
#
################################################################################

function restrict_scalars(V::AbsSpace, K::Field = FlintQQ, alpha = one(base_ring(V)))
  E = base_ring(V)
  n = rank(V)
  d = absolute_degree(E)
  B = absolute_basis(E)
  v = elem_type(E)[zero(E) for i in 1:n]
  w = elem_type(E)[zero(E) for i in 1:n]
  G = zero_matrix(FlintQQ, d * n, d * n)
  r = 1
  c = 1
  indices = Vector{Tuple{Int, Int}}(undef, n * d)
  for i in 1:n
    for bi in 1:length(B)
      v[i] = B[bi]
      c = 1
      for j in 1:n
        for bj in 1:length(B)
          w[j] = B[bj]
          G[r, c] = trace(alpha * inner_product(V, v, w), FlintQQ)
          w[j] = zero(E)
          c = c + 1
        end
      end
      v[i] = zero(E)
      indices[r] = (i, bi)
      r = r + 1
    end
  end

  VabstoV = function(v)
    @assert length(v) == d * n
    z = Vector{elem_type(E)}(undef, n)
    for i in 1:n
      z[i] = zero(E)
    end
    for k in 1:(d * n)
      (i, j) = indices[k]
      z[i] = z[i] + v[k] * B[j]
    end
    return z
  end

  VtoVabs = function(w)
    z = Vector{fmpq}(undef, d * n)
    k = 1
    for i in 1:n
      y = w[i]
      for j in 1:d
        z[k] = absolute_coeff(y, j - 1)
        k = k + 1
      end
    end
    return z
  end

  return quadratic_space(FlintQQ, G), VabstoV, VtoVabs
end

# Careful, starts at 0!
function absolute_coeff(z::nf_elem, i)
  return coeff(z, i)
end

function absolute_coeff(z::NumFieldElem, i)
  d = absolute_degree(base_field(parent(z)))
  rowindex = fld(i, d)
  colindex = (i % d)
  return absolute_coeff(coeff(z, rowindex), colindex)
end

function absolute_basis(K::NumField)
  k = base_field(K)
  kabs = absolute_basis(k)
  res = elem_type(K)[]
  for b in basis(K)
    for bb in kabs
      push!(res, bb * b)
    end
  end
  return res
end

function absolute_basis(K::NumField{fmpq})
  return basis(K)
end

#

istotally_real(::FlintRationalField) = true

istotally_positive(x::fmpq) = x > 0

# This function is really slow...
function denominator(M::fmpq_mat)
  d = one(FlintZZ)
  for i in 1:nrows(M)
    for j in 1:ncols(M)
      d = lcm!(d, d, denominator(M[i, j]))
    end
  end
  return d
end

function _weak_approximation_coprime(IP, S, M)
  R = order(M)
  A, _exp, _log = infinite_primes_map(R, IP, M)

  t = (1 + _exp(A([ S[j] == 1 ? 0 : -1 for j in 1:length(IP)])))
  @assert all(i -> sign(t, IP[i]) == S[i], 1:length(IP))
  return t
end


#=
LeapFrogger.jl
Jacob Child
Flow Lab Onboarding Project
April 23rd, 2022

Equations
V(vector) = (Gamma(vector) cross r(vector))/(2*pi*r^2)

Pseudocode- Using the reference image and the right hand rule, compute the velocities for each
vortex "point" individually. The bottom left point is 1, then go clockwise.
For Pt1- compute the effects of pts 2,3,4 on it, add them together for Pt1TotVeloc, then using a
small time step, compute its new location and repeat the process :||
All points have the same process
Things to watch- don't update the position of any points before all the calculations for
that time step are done.

Coordinate system:
(0; 0; 0) = [x; y; z] = inbetween and in line with pts 1 and 2
Pt 1 = btm left, the rest follow clockwise.
Righthand rule = out of page = positive, into page = negative.
=#

#needed packages
using LinearAlgebra
using Plots
#using DataFrame

#Inputs/Variables
d = 1 #distance between vortex ring top and bottom points, and between rings
G = 1   #strength of vorticity = Gamma, same for all rings/points
dt = .01 #time step size
x = collect(0:dt:480) #gives a time vector from 0-xseconds with .01 second step size
plot()  #to clear the plot

#p1 = zeros(lenght(x), 3)
#Point 1
Pt1Gamma = [0; 0; -G]
Pt1Location = [0; -d/2; 0]
x1 = []
y1 = []

#Point 2
Pt2Gamma = [0; 0; G]
Pt2Location = [0; d/2; 0]
x2 = []
y2 = []

#Point 3
Pt3Gamma = [0; 0; G]
Pt3Location = [d; d/2; 0]
x3 = []
y3 = []

#Point 4
Pt4Gamma = [0; 0; -G]
Pt4Location = [d; -d/2; 0]
x4 = []
y4 = []

#before putting everything in functions, practice for pt 1
#2 on 1
# R21 = Pt1Location - Pt2Location #make sure positive vs negative is correct!
# println(R21)
# MagnitudeR21 = sqrt(R21[1]^2 + R21[2]^2 + R21[3]^2) #the magnitude of the distance btwn the 2 pts
# V21Inf = cross(Pt2Gamma, R21) / (2*pi*MagnitudeR21) #The influence pt2 has on 1 for velocity
# println(V21Inf)
# dt = .01    #time step
# Pt1dp = V21Inf * dt
# Pt1Location = Pt1Location + Pt1dp
# println(Pt1Location)

function v_influence(Location1, Location2, Gamma)
    R = Location1 - Location2 #how to make sure this is right positive vs negative?
    MagnitudeR = sqrt(R[1]^2 + R[2]^2 + R[3]^2) #the magnitude of the distance btwn the 2 pts
    VInf = cross(Gamma, R) / (2*pi*MagnitudeR^2) #The influence pt2 has on 1 for velocity
    return VInf
end

function total_influence(VAInf,VBInf, VCInf)
    VTotInf = VAInf + VBInf +VCInf
    return VTotInf

end

function update_location(Location, VInfTot, dt)
    dp = VInfTot*dt
    NewLocation = Location + dp
    return NewLocation
end

###Looping Function

for i in x #i=1:length(x)

    #Point 1
    V21Inf = v_influence(Pt1Location, Pt2Location, Pt2Gamma)
    V31Inf = v_influence(Pt1Location, Pt3Location, Pt3Gamma)
    V41Inf = v_influence(Pt1Location, Pt4Location, Pt4Gamma)
    V1InfTot = total_influence(V21Inf, V31Inf, V41Inf)
    #println(V21Inf)

    #Point 2
    V12Inf = v_influence(Pt2Location, Pt1Location, Pt1Gamma)
    V32Inf = v_influence(Pt2Location, Pt3Location, Pt3Gamma)
    V42Inf = v_influence(Pt2Location, Pt4Location, Pt4Gamma)
    V2InfTot = total_influence(V12Inf, V32Inf, V42Inf)

    #Point 3
    V13Inf = v_influence(Pt3Location, Pt1Location, Pt1Gamma)
    V23Inf = v_influence(Pt3Location, Pt2Location, Pt2Gamma)
    V43Inf = v_influence(Pt3Location, Pt4Location, Pt4Gamma)
    V3InfTot = total_influence(V13Inf, V23Inf, V43Inf)

    #Point 4
    V14Inf = v_influence(Pt4Location, Pt1Location, Pt1Gamma)
    V24Inf = v_influence(Pt4Location, Pt2Location, Pt2Gamma)
    V34Inf = v_influence(Pt4Location, Pt3Location, Pt3Gamma)
    V4InfTot = total_influence(V14Inf, V24Inf, V34Inf)

    #Update Location of All points
    global Pt1Location = update_location(Pt1Location, V1InfTot, dt) #Make sure to update the positions last!
    # p1[i,:] = update_location(Pt1Location, V1InfTot, dt)
    global Pt2Location = update_location(Pt2Location, V2InfTot, dt)
    global Pt3Location = update_location(Pt3Location, V3InfTot, dt)
    global Pt4Location = update_location(Pt4Location, V4InfTot, dt)
    global x1 = push!(x1,Pt1Location[1]) #creates a series of x and y values for each point
    global y1 = push!(y1,Pt1Location[2])
    global x2 = push!(x2,Pt2Location[1])
    global y2 = push!(y2,Pt2Location[2])
    global x3 = push!(x3,Pt3Location[1])
    global y3 = push!(y3,Pt3Location[2])
    global x4 = push!(x4,Pt4Location[1])
    global y4 = push!(y4,Pt4Location[2])
    #x1,y1,_ = update_location(Pt1Location, V1InfTot, dt)
    # #debugging
    # A = V12Inf[1]   #x direction velocity influence of pt 3 on pt 2
    # #B = V12Inf[2]   #y direction velocity influence of pt 3 on pt2
    # display(plot!([A], legend = false, m = 3))

end
#plots
#Vortex Ring 1 (back)
display(plot!(x1,y1,  label = false, color = :blue, xlabel="Position (m)", ylabel = "Position (m)")) #plot only does vectors, not floats etc, number values are plot commands
display(plot!(x2,y2,  label = false, color = :blue, title = "Leap Frogging Vortex Ring Positions Over Time"))

#Vortex Ring 2 (front)
display(plot!(x3,y3,  label = false, color = :green))
display(plot!(x4,y4,  label = false, color = :green))

# plt = plot(xaxis="X distance (m)", yaxis="Y distance (m)")
# plot!(x1, y1)
# display(plt)

#println(Pt1Location, Pt2Location, Pt3Location, Pt4Location)


using ForwardDiff, NLopt

l = x->loss(x,full=false) #from uncertainty.jl

function loss_nlopt(x,∇)
    length(∇) > 0 ? ForwardDiff.gradient!(∇, l,x) : nothing
    l(x)
end

function gimbal_optimize(x₀)
    #x₀ = [0.01,0.029,0.16,0.06,0.009,0.01,7000.0]
    #x₀ = [0.011,0.0319,0.0319,0.066,0.0081,0.011,7000.0]
    opt = Opt(:LD_MMA, 7)
    opt.lower_bounds = x₀*0.8
    opt.upper_bounds = x₀*1.2
    opt.xtol_rel = 1e-3
    #opt.maxtime = 10 #time limit in seconds
    opt.min_objective = loss_nlopt
    (minf,minx,ret) = NLopt.optimize(opt, x₀)
end

export gimbal_optimize


type FiveConstantArms <: BanditEnvironment
end

function arms(env::FiveConstantArms, trial::Int64)
  [
   BernoulliArm(0.1),
   BernoulliArm(0.3),
   BernoulliArm(0.5),
   BernoulliArm(0.4),
   BernoulliArm(0.2)
   ]
end

function n_arms(env::FiveConstantArms)
  5
end


@plottest begin
    # LIF params
    τm = 100
    vreset = 0.0
    vth = 0.1
    R = 1.75

    # Input spike train params
    rate = 0.05
    T = 1000

    lif = Neuron(QueuedSynapse(Synapse.Delta()), LIF(τm, vreset, R), Threshold.Ideal(vth))
    input = ConstantRate(rate)
    spikes = excite!(lif, input, T)

    # callback to record voltages
    voltages = Float64[]
    record = let lif = lif
        () -> push!(voltages, getvoltage(lif))
    end

    # compile
    simulate!(lif, T; cb = record)
    reset!(lif)
    voltages = Float64[]
    excite!(lif, spikes)

    # simulate
    simulate!(lif, T; cb = record, dense = true)

    # plot dense voltage recording
    plot(collect(1:T), voltages, title = "LIF Membrane Potential Over Time", xlabel = "Time (sec)", ylabel = "Potential (V)", label = "Dense")

    # repeat with dense simulation
    reset!(lif)
    voltages = Float64[]
    excite!(lif, spikes)
    output = simulate!(lif, T; cb = record)

    # plot sparse voltage recording
    voltage_plot = plot!(spikes, voltages,
        title = "LIF Membrane Potential Over Time", xlabel = "Time (sec)", ylabel = "Potential (V)", label = "Sparse (default)")

    # plot raster plot
    raster_plot = rasterplot(spikes, output, label = ["Input", "Output"], title = "Raster Plot", xlabel = "Time (sec)")

    plot(raster_plot, voltage_plot, layout = grid(2, 1))
end joinpath(datadir, "Lif-Test.png") !isci


# Adaptation of some examples in MosekTools
# Refer to:
# https://github.com/JuliaOpt/MosekTools.jl/blob/master/test/jump_soc.jl
#
# MIT License
# Copyright (c) 2017, Ulf Worsøe, Mosek Aps

using KNITRO
using JuMP
using Test
using MathOptInterface
const MOI = MathOptInterface

model = Model(with_optimizer(KNITRO.Optimizer, outlev=0))
@variable(model, x)
@variable(model, y)
@variable(model, t >= 0.0)
@objective(model, Min, t)
@constraint(model, x + y >= 1.0)
@constraint(model, [t, x, y] in MOI.SecondOrderCone(3))

JuMP.optimize!(model)

@test JuMP.has_values(model)
@test JuMP.termination_status(model) == MOI.LOCALLY_SOLVED
@test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
@test JuMP.objective_value(model) ≈ sqrt(1/2) atol=1e-6
@test JuMP.value.([x,y,t]) ≈ [0.5, 0.5, sqrt(1/2)] atol=1e-3


#
#--------------------------------------#
export pcg
#--------------------------------------#
"""
 Preconditioned conjugate gradient

 We can't just use any packaged iterative solver for the PCG because
 we need to overwrite not just A*x, A'*x ops for Ax=b which is
 complicated enough since x,b  are multidimensional arrays (would have
 to allocate memory in converting and reconverting arrays to vecs), but
 also <u,v>, the dot product. because of the SEM construction, the dot
 product isnt just sum(u.*v) , but it's sum(u.*v.* SEM_mult)

"""
function pcg(b,opA
            ;opM=x->x           # preconditioner
            ,mult=ones(size(b)) # SEM mult
            ,ifv=false          # verbose flag
            ,tol=1e-8           # tolerance
            ,maxiter=length(b)) # maximum number of iterations

n = length(b);

x   = zero(b) # zero out initial condition
Ax  = zero(b)
ra  = b - Ax;
ha  = zero(b)
hp  = zero(b)
hpp = zero(b)
rp  = zero(b)
rpp = zero(b)
u   = zero(b)
k   = 0

while(norm(ra,Inf) > tol)
ha = opM * ra # preconditinoer
#println("PCG iter: ",k,", res: ",norm(ra,2));
if(k==maxiter) println("warning: res:",norm(ra,Inf)); return x; end;
k  += 1;
hpp = hp;
rpp = rp;
hp  = ha;
rp  = ra;
t   = sum(rp.*hp.*mult);
if(k==1)
    u = copy(hp);
else
    u = hp+(t/sum(rpp.*hpp.*mult))*u;
end
Au = opA * u # operator
a = t / sum(u.*Au.*mult);
x = x + a * u;
ra = rp - a * Au;
end

if(ifv) println("PCG iter: ",k,", res: ",norm(ra,2)); end

return x
end
#--------------------------------------#
export pcg!
#--------------------------------------#
function pcg!(x,b,opA
             ;opM=x->x           # preconditioner
             ,mult=ones(size(b)) # SEM mult
             ,ifv=false          # verbose flag
             ,tol=1e-8           # tolerance
             ,maxiter=length(b)) # maximum number of iterations

    x .= pcg(b,opA
            ;opM=opM
            ,mult=mult
            ,ifv=ifv
            ,tol=tol
            ,maxiter=maxiter)

return
end
#--------------------------------------#


using StatsBase
using Base.Test

##### weighted var & std

x = rand(10)
wv = weights(rand(10))
m = mean(x, wv)

@test var(x, wv)           ≈ sum(abs2.(x .- m), wv) ./ sum(wv)
@test var(x, wv; mean=0)   ≈ sum(abs2.(x), wv) ./ sum(wv)
@test var(x, wv; mean=1.0) ≈ sum(abs2.(x .- 1.0), wv) ./ sum(wv)

@test std(x, wv)           ≈ sqrt(var(x, wv))
@test std(x, wv; mean=0)   ≈ sqrt(var(x, wv; mean=0))
@test std(x, wv; mean=1.0) ≈ sqrt(var(x, wv; mean=1.0))

(m, v) = mean_and_var(x)
@test m == mean(x)
@test v == var(x)

(m, s) = mean_and_std(x)
@test m == mean(x)
@test s == std(x)

(m, v) = mean_and_var(x, wv)
@test m == mean(x, wv)
@test v == var(x, wv)

(m, s) = mean_and_std(x, wv)
@test m == mean(x, wv)
@test s == std(x, wv)

x = rand(5, 6)
w1 = rand(5)
w2 = rand(6)
wv1 = weights(w1)
wv2 = weights(w2)
m1 = mean(x, wv1, 1)
m2 = mean(x, wv2, 2)

@test var(x, wv1, 1; mean=0) ≈ sum(abs2.(x) .* w1, 1) ./ sum(wv1)
@test var(x, wv2, 2; mean=0) ≈ sum(abs2.(x) .* w2', 2) ./ sum(wv2)

@test var(x, wv1, 1; mean=m1) ≈ sum(abs2.(x .- m1) .* w1, 1) ./ sum(wv1)
@test var(x, wv2, 2; mean=m2) ≈ sum(abs2.(x .- m2) .* w2', 2) ./ sum(wv2)

@test var(x, wv1, 1) ≈ sum(abs2.(x .- m1) .* w1, 1) ./ sum(wv1)
@test var(x, wv2, 2) ≈ sum(abs2.(x .- m2) .* w2', 2) ./ sum(wv2)

@test std(x, wv1, 1)          ≈ sqrt.(var(x, wv1, 1))
@test std(x, wv2, 2)          ≈ sqrt.(var(x, wv2, 2))
@test std(x, wv1, 1; mean=0)  ≈ sqrt.(var(x, wv1, 1; mean=0))
@test std(x, wv2, 2; mean=0)  ≈ sqrt.(var(x, wv2, 2; mean=0))
@test std(x, wv1, 1; mean=m1) ≈ sqrt.(var(x, wv1, 1; mean=m1))
@test std(x, wv2, 2; mean=m2) ≈ sqrt.(var(x, wv2, 2; mean=m2))

for d in 1:2
    (m, v) = mean_and_var(x, d)
    @test m == mean(x, d)
    @test v == var(x, d)

    (m, s) = mean_and_std(x, d)
    @test m == mean(x, d)
    @test s == std(x, d)
end

(m, v) = mean_and_var(x, wv1, 1)
@test m == mean(x, wv1, 1)
@test v == var(x, wv1, 1)

(m, v) = mean_and_var(x, wv2, 2)
@test m == mean(x, wv2, 2)
@test v == var(x, wv2, 2)

(m, s) = mean_and_std(x, wv1, 1)
@test m == mean(x, wv1, 1)
@test s == std(x, wv1, 1)

(m, s) = mean_and_std(x, wv2, 2)
@test m == mean(x, wv2, 2)
@test s == std(x, wv2, 2)

##### skewness & kurtosis

wv = weights(ones(5) * 2.0)

@test skewness(1:5)             ≈  0.0
@test skewness([1, 2, 3, 4, 5]) ≈  0.0
@test skewness([1, 2, 2, 2, 5]) ≈  1.1731251294063556
@test skewness([1, 4, 4, 4, 5]) ≈ -1.1731251294063556

@test skewness([1, 2, 2, 2, 5], wv) ≈ 1.1731251294063556

@test kurtosis(1:5)             ≈ -1.3
@test kurtosis([1, 2, 3, 4, 5]) ≈ -1.3
@test kurtosis([1, 2, 3, 3, 2]) ≈ -1.1530612244897953

@test kurtosis([1, 2, 3, 4, 5], wv) ≈ -1.3

##### general moments

x = collect(2.0:8.0)
@test moment(x, 2) ≈ sum((x .- 5).^2) / length(x)
@test moment(x, 3) ≈ sum((x .- 5).^3) / length(x)
@test moment(x, 4) ≈ sum((x .- 5).^4) / length(x)
@test moment(x, 5) ≈ sum((x .- 5).^5) / length(x)

@test moment(x, 2, 4.0) ≈ sum((x .- 4).^2) / length(x)
@test moment(x, 3, 4.0) ≈ sum((x .- 4).^3) / length(x)
@test moment(x, 4, 4.0) ≈ sum((x .- 4).^4) / length(x)
@test moment(x, 5, 4.0) ≈ sum((x .- 4).^5) / length(x)

w = weights([1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0])
x2 = collect(2.0:6.0)
@test moment(x, 2, w) ≈ sum((x2 .- 4).^2) / 5
@test moment(x, 3, w) ≈ sum((x2 .- 4).^3) / 5
@test moment(x, 4, w) ≈ sum((x2 .- 4).^4) / 5
@test moment(x, 5, w) ≈ sum((x2 .- 4).^5) / 5