###################################################################################################
#
# FlowBalances.jl
#
# by Walt McNab (May 2020)
#
# (1) Approximate potentiometric surface and hydraulic conductivity distributions by interpolation with kriging
# (2) Estimate groundwater flow vector field by Darcy's law; optional particle tracking
# (3) Numerical integration across polygon drawn on flow field to compute fluxes
#
# Current version assumes datum for head values is equal to an elevation of zero.
#
###################################################################################################

using DelimitedFiles
using Statistics
using Polynomials
using QuadGK

# miscellaneous fixed parameters
const epsilon = 0.01                	# factor used to prevent binning problems with variogram construction
const delta = 0.1 						# distance offset used to estimate hydraulic gradient

### defined types ###

mutable struct Polygon
    x::Array{Float64, 1}                # apex locations
    y::Array{Float64, 1}
    m::Array{Float64, 1}                # slope of polygon side
    n::Array{Any, 1}                    # normal vector to polygon side
    L::Array{Float64, 1}                # length of polygon side
end

mutable struct TrackParams
	dStep::Float64 						# step size for particle tracking
	phi::Float64 						# porosity
end

mutable struct Sources              	# source/sink terms (i.e., wells)
	activeSources::Bool
    x::Array{Float64, 1}
    y::Array{Float64, 1}
    Q::Array{Float64, 1}
end

mutable struct Particle
	x::Array{Float64, 1} 				# particle locations, starting with initial
	y::Array{Float64, 1}
	vx::Array{Float64, 1} 				# particle velocities (for plotting)
	vy::Array{Float64, 1}
	t::Array{Float64, 1}
	fwd::Float64 						# movement direction: 1 = forward, -1 = backward
	active::Bool
end

mutable struct Krig 					# variogram parameters
	varMatrix::Array{Float64, 2}
	m::Float64
	b::Float64
end

mutable struct Points
    x::Array{Float64, 1}
    y::Array{Float64, 1}
    v::Array{Float64, 1} 				# head measurements, or log hydraulic conductivity
end

mutable struct Grid
	x0::Float64
	xf::Float64
	y0::Float64
	yf::Float64
    x::Array{Float64, 1}
    y::Array{Float64, 1}
end

### utility functions ###

function DistVector(xp::Float64, yp::Float64, xArray::Array{Float64, 1}, yArray::Array{Float64, 1})::Array{Float64, 1}
	# return a vector of distances
	d = Float64[]
    for i = 1:length(xArray)
		dist = sqrt((xp-xArray[i])^2 + (yp-yArray[i])^2)
        push!(d, dist)
    end
	return d
end

function dot(a::Array{Float64, 1}, b::Array{Float64, 1})::Float64
    # vector dot product
    return sum(a.*b)
end

### hydro modeling functions ###

function Grad(x::Float64, y::Float64, headData::Points, krigHead::Krig)
	# estimate hydraulic gradient by finite differences
	h = Interpolate(x, y, headData, krigHead)
	hx = Interpolate(x+delta, y, headData, krigHead)
	hy = Interpolate(x, y+delta, headData, krigHead)
	dhdx = (hx-h)/delta
	dhdy = (hy-h)/delta
	return dhdx, dhdy
end

function Darcy(x::Float64, y::Float64, headData::Points, krigHead::Krig, KData::Points, krigLogK::Krig)
	# local groundwater Darcy velocity
	K = 10.0^Interpolate(x, y, KData, krigLogK)
	dhdx, dhdy = Grad(x, y, headData, krigHead)
	qx = -K * dhdx
	qy = -K * dhdy
	return qx, qy
end

function Displace(x0::Float64, y0::Float64, vx::Float64, vy::Float64, d::Float64, fwd::Float64)
	# particle movement
	vm = sqrt(vx^2 + vy^2) 	# velocity magnitude
	dt = d/vm 				# time step to move "d"
	x = x0 + fwd*vx*dt 		# implied new position
	y = y0 + fwd*vy*dt
	return x, y, dt
end

function CheckEndState(part::Particle, trackParams::TrackParams, grid::Grid, wells::Sources)::Particle
	# note if particle has exited model domain or has approached too close to a well
	part.active = (part.x[end]>=grid.x0) & (part.x[end]<=grid.xf) & (part.y[end]>=grid.y0) & (part.y[end]<=grid.yf)
	if wells.activeSources == true
		dWells = DistVector(part.x[end], part.y[end], wells.x, wells.y)
		if minimum(dWells) <= trackParams.dStep
			part.active = false
		end
	end
	return part
end

function Track(part::Particle, trackParams::TrackParams, headData::Points, krigHead::Krig, KData::Points, krigLogK::Krig, wells::Sources, grid::Grid)::Particle
	# follow the progress of a particle (initial conditions already defined)
	while part.active == true
		# single-pass particle position estimator
		xf, yf, dt = Displace(part.x[end], part.y[end], part.vx[end], part.vy[end], trackParams.dStep, part.fwd)
		vxf, vyf = Darcy(xf, yf, headData, krigHead, KData, krigLogK)./trackParams.phi
		vx = 0.5*part.vx[end] + 0.5*vxf
		vy = 0.5*part.vy[end] + 0.5*vyf
		x, y, dt = Displace(part.x[end], part.y[end], vx, vy, trackParams.dStep, part.fwd)
		push!(part.x, x)
		push!(part.y, y)
		push!(part.vx, vx)
		push!(part.vy, vy)
		push!(part.t, part.t[end]+dt)
		part = CheckEndState(part, trackParams, grid, wells)        # check to see if exit condition is met
	end
	return part
end

function FluxNorm(x::Float64, headData::Points, krigHead::Krig, KData::Points, krigLogK::Krig, polygon::Polygon, indx::Int64)
    # vertically-integrated, normal groundwater flux at a point along a transect
    y = polygon.y[indx] + polygon.m[indx]*(x-polygon.x[indx])
    qx, qy = Darcy(x, y, headData, krigHead, KData, krigLogK).*Interpolate(x, y, headData, krigHead)
    return -dot([qx, qy], polygon.n[indx])
end

function SumPolyFluxes(polygon::Polygon, headData::Points, krigHead::Krig, KData::Points, krigLogK::Krig)::Float64
    # integrate groundwater fluxes across polygon/polyline sides
    QTot = 0.0
    for i = 1:length(polygon.x)-1
        x0 = minimum([polygon.x[i], polygon.x[i+1]])
        xf = maximum([polygon.x[i], polygon.x[i+1]])
        QTot += (polygon.L[i]/abs(xf-x0)) * quadgk(x -> FluxNorm(x, headData, krigHead, KData, krigLogK, polygon, i), x0, xf)[1]
    end
    return QTot
end

### kriging functions ###

function MakeKrig(data, setName)
	# parameters to inform kriging
    S = Squareform(data.x, data.y)      			# distance matrix
    m, b = Variogram(data, S, setName)       		# linear variogram
    varMatrix = (S.*m).+b
	return Krig(varMatrix, m, b)
end

function Squareform(xArray::Array{Float64, 1}, yArray::Array{Float64, 1})::Array{Float64, 2}
	# squareform distance matrix
	N = length(xArray)
	S = zeros(Float64, N, N)
    for i = 1:N
		S[i, :] = DistVector(xArray[i], yArray[i], xArray, yArray)
    end
	return S
end

function Variogram(points::Points, S::Array{Float64, 2}, setName::AbstractString)
    # create variogram; write variogram summary to file and return linear variogram model
    dist = Float64[]
    gamma = Float64[]
    lagGroup = []
    lagMeanDist = Float64[]
    xVar = Float64[]
    yVar = Float64[]
    N = length(points.x)        # number of lag bins
    # difference in value versus distance offset
    for i = 1:N
        for j = i+1:N
            push!(dist, S[i,j])
            push!(gamma, (points.v[i]-points.v[j])^2)
        end
    end
    # populate bins
    lagStepSize = maximum(dist)/N
    for i = 1:N
        push!(lagGroup, [])                             # set up list of lag sets
        push!(lagMeanDist, (i-0.5)*lagStepSize)         # mean distance offset for each of the lag groups
    end
    # assign lags to groups
    for i = 1:length(dist)
        lagIndex = floor(Int, (dist[i]-epsilon)/lagStepSize) + 1
        push!(lagGroup[lagIndex], gamma[i])
    end
    # calculate variogram
    for i = 1:N
        if length(lagGroup[i]) > 1     # binned lag group is large enough to calculate a variance
            push!(xVar, lagMeanDist[i])
            push!(yVar, mean(lagGroup[i])/2)
        end
    end
    WriteVariogram(xVar, yVar, setName)					# write to output file
    # linear regression (i.e. linear variogram model)
    r = fit(xVar, yVar, 1)
    return r[1], r[0]               # = return m, b
end

function Interpolate(x::Float64, y::Float64, points::Points, krig::Krig)::Float64
    # interpolate field value at (x, y)
    dist = DistVector(x, y, points.x, points.y)          # calculate covariance vector for distances between (x,y) and data points
    k = (dist.*krig.m).+krig.b
    wt = \(krig.varMatrix, k)                            # compute weights
    z = sum(wt.*points.v)
    return z
end

### I/O functions ###

function ReadParticles(trackParams::TrackParams, headData::Points, krigHead::Krig, KData::Points, krigLogK::Krig)::Array{Particle, 1}
	# initialize particles
	part = []
	data = readdlm("particleStarts.txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        x = Float64(data[1][i, 1])
        y = Float64(data[1][i, 2])
        fwd = Float64(data[1][i, 3])
		vx, vy = Darcy(x, y, headData, krigHead, KData, krigLogK)./trackParams.phi
		t = 0.0
		active = true
        push!(part, Particle([x], [y], [vx], [vy], [t], fwd, active))
    end
    println("Read particle starting positions and initialized.")
	return part
end

function ReadPolygon()::Polygon
    # clockwise list convention; repeat 1st point as last point
    x = Float64[]
    y = Float64[]
    m = Float64[]
    n = []
    L = Float64[]
    data = readdlm("polygon.txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        push!(x, Float64(data[1][i, 1]))
        push!(y, Float64(data[1][i, 2]))
    end
    for i = 1:length(x)-1
        dx = x[i+1] - x[i]
        dy = y[i+1] - y[i]
        push!(m, dy/dx)
        d = sqrt(dx^2 + dy^2)
        push!(L, d)
        push!(n, [-dy/d, dx/d])
    end
    println("Read and processed bounding polygon.")
    return Polygon(x, y, m, n, L)
end

function ReadPoints(dataSet::AbstractString)::Points
    # read point data set and return array of point structs
    x = Float64[]
    y = Float64[]
    v = Float64[]
    data = readdlm(dataSet * ".txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        push!(x, Float64(data[1][i, 1]))
        push!(y, Float64(data[1][i, 2]))
        push!(v, Float64(data[1][i, 3]))
    end
    println("Read " * dataSet * ".")
    return Points(x, y, v)
end

function ReadGrid()
    # set up grid for kriging and inverse model
    x = Float64[]
    y = Float64[]
    data = readdlm("grid.txt", '\t', header=true)
    x0 = Float64(data[1][1, 2])
    y0 = Float64(data[1][1, 3])
    xf = Float64(data[1][2, 2])
    yf = Float64(data[1][2, 3])
    nx = Int64(data[1][3, 2])
    ny = Int64(data[1][3, 3])
    dx = (xf-x0)/nx
    dy = (yf-y0)/ny
    for j = 1:ny
        for i = 1:nx
            push!(x, x0 + (i-0.5)*dx)
            push!(y, y0 + (j-0.5)*dy)
        end
    end
    println("Read grid constraints.")
    return Grid(x0, xf, y0, yf, x, y)
end

function ReadSources(dataSet::AbstractString)::Sources
    # read source/sink terms
    x = Float64[]
    y = Float64[]
    Q = Float64[]
	activeSources = true
    data = readdlm(dataSet * ".txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        push!(x, Float64(data[1][i, 1]))
        push!(y, Float64(data[1][i, 2]))
        push!(Q, Float64(data[1][i, 3]))
    end
    println("Read " * dataSet * ".")
    return Sources(activeSources, x, y, Q)
end

function ReadTrack()
	# miscellaneous particle tracking parameters
    data = readdlm("trackParams.txt", '\t', header=false)
    dStep = Float64(data[1, 2])
	phi = Float64(data[2, 2])
    println("Read tracking parameters.")
	return TrackParams(dStep, phi)
end

function WriteVariogram(xVar::Array{Float64, 1}, yVar::Array{Float64, 1}, setName::AbstractString)
	# variogram output file
    csvfile = open("variogram_" * setName * ".csv","w")
    line_out = "h" * "," * "gamma"
    println(csvfile, line_out)
    for i = 1:length(xVar)
        line_out = string(xVar[i]) * "," * string(yVar[i])
        println(csvfile, line_out)
    end
    close(csvfile)
end

function WriteGridOut(grid::Grid, head::Array{Float64, 1}, logK::Array{Float64, 1})
	# write grid output, with interpolated heads and log K estimates, to file
    csvfile = open("gridResults.csv","w")
    line_out = "x" * "," * "y" * "," * "head" * "," * "logK"
    println(csvfile, line_out)
    for i = 1:length(grid.x)
        line_out = string(grid.x[i]) * "," * string(grid.y[i]) * "," * string(head[i]) * "," * string(logK[i])
        println(csvfile, line_out)
    end
    close(csvfile)
	println("Wrote grid output file.")
end

function WriteParticles(part::Array{Particle, 1})
	# write grid output, with interpolated heads and log K estimates, to file
    csvfile = open("particleTracks.csv","w")
    line_out = "particle" * "," * "time" * "," * "x" * "," * "y" * "," * "vx" * "," * "vy" * "," * "v-magn"
    println(csvfile, line_out)
    for (i, p) in enumerate(part)
		for j = 1:length(p.x)
			vMagn = sqrt(p.vx[j]^2 + p.vy[j]^2)
			line_out = string(i) * "," * string(p.t[j]) * "," * string(p.x[j]) * "," * string(p.y[j]) * "," * string(p.vx[j]) * "," * string(p.vy[j]) * "," * string(vMagn)
			println(csvfile, line_out)
		end
    end
    close(csvfile)
	println("Wrote particle tracks output file.")
end

### main ###

function FlowBalances(gridFlag, trackFlag, polyFlag, wellsFlag)

    # problem setup
	dataSet = ["heads", "logKs"]
    grid = ReadGrid()                       	# read grid
	headData = ReadPoints(dataSet[1])          	# head measurements
	KData = ReadPoints(dataSet[2])          	# log K measurements
	if wellsFlag == true 						# well locations and pumping rates, if applicable
		wells = ReadSources("wells")
	else
		wells = Sources(false, Float64[], Float64[], Float64[]) 		# empty struct, with activeSources = false
	end

	# set up kriging models
	krigHead = MakeKrig(headData, dataSet[1])
	krigLogK = MakeKrig(KData, dataSet[2])

	# read polygon, if used
    if polyFlag == true
        polygon = ReadPolygon()                                                     # read input file
        QIntegrate = SumPolyFluxes(polygon, headData, krigHead, KData, krigLogK)    # integrate
        println("Flux across polygon/polyline = ", "\t", QIntegrate)
    end

	# conduct particle tracking, if implemented
	if trackFlag == true
		trackParams = ReadTrack()													# read parameters file
		part = ReadParticles(trackParams, headData, krigHead, KData, krigLogK)		# initialize particles
		for i = 1:length(part)														# track particles
			if part[i].active == true
				part[i] = Track(part[i], trackParams, headData, krigHead, KData, krigLogK, wells, grid)
			end
		end
		WriteParticles(part)							# write to output file
	end

	# write to output grid for plotting
    if gridFlag == true
        head = Float64[]
        logK = Float64[]
        println("Interpolating heads across grid ...")
        for i = 1:length(grid.x)
            push!(head, Interpolate(grid.x[i], grid.y[i], headData, krigHead))
        end
        println("Interpolating log K estimates across grid ...")
        for i = 1:length(grid.x)
            push!(logK, Interpolate(grid.x[i], grid.y[i], KData, krigLogK))
        end
        WriteGridOut(grid, head, logK)        # write to output file
	end

    println("Finished.")

end

### run model ###

gridFlag = false
trackFlag = false
polyFlag = true
wellsFlag = true
FlowBalances(gridFlag, trackFlag, polyFlag, wellsFlag)


include("bounds.jl")

# Compute the negative log prior:

function nlog_prior_noecc(x::Array{T,1}) where {T <: Real}
# We will place a joint prior on eccentricity vector
# such that each planet has an eccentricity which lies between
# 0 and emax2, with a gradual decrease from emax1 to emax2:
# Eccentricity priors:
emax1 = 0.2; emax2 = 0.3
# Mass priors:
lgm1 = -8.0; lgm2 = -7.0
lgm3 = -3.0; lgm4 = -2.0
# Period priors:
P1 = [1.49,2.40,4.03,6.08,9.18,12.33,18.75]; P2 = [1.50,2.41,4.04,6.09,9.19,12.34,18.76]
P3 = [1.52,2.43,4.06,6.11,9.22,12.36,18.78]; P4 = [1.53,2.44,4.07,6.12,9.23,12.37,18.79]
# t0 priors:
t01 = [7257.53,7258.57,7257.05,7257.81,7257.05,7257.70,7249.58]
t02 = [7257.54,7258.58,7257.06,7257.82,7257.06,7257.71,7249.59]
t03 = [7257.56,7258.60,7257.08,7257.84,7257.08,7257.73,7249.61]
t04 = [7257.57,7258.61,7257.09,7257.85,7257.09,7257.74,7249.62]
# Define -log(prior):
nlprior = zero(T)
nparam = size(x)[1]
n = div(nparam-2,5)+1
# Loop over planets:
for i=2:n
  # Place prior on eccentricities:
  ecc = sqrt(x[(i-2)*5+4]^2+x[(i-2)*5+5]^2)
  nlprior_tmp,dpdx = log_bounds_upper(ecc,emax1,emax2)
  nlprior -= nlprior_tmp

  # We will place broad lower and upper limits on log_10(mass):
  log10mass = log10(x[(i-2)*5+1])
  nlprior_tmp,dpdx = log_bounds_lower(log10mass,lgm1,lgm2)
  nlprior -= nlprior_tmp
  nlprior_tmp,dpdx = log_bounds_upper(log10mass,lgm3,lgm4)
  nlprior -= nlprior_tmp
  # Broad priors on period:
  P = x[(i-2)*5+2]
  nlprior_tmp,dpdx = log_bounds_lower(P,P1[i-1],P2[i-1])
  nlprior -= nlprior_tmp
  nlprior_tmp,dpdx = log_bounds_upper(P,P3[i-1],P4[i-1])
  nlprior -= nlprior_tmp
  # Broad priors on t0:
  t0 = x[(i-2)*5+3]
  nlprior_tmp,dpdx = log_bounds_lower(t0,t01[i-1],t02[i-1])
  nlprior -= nlprior_tmp
  nlprior_tmp,dpdx = log_bounds_upper(t0,t03[i-1],t04[i-1])
  nlprior -= nlprior_tmp
end
# Add in prior on ndof & lnV1:
ndof1 = 0.5; ndof2 = 1.0
ndof3 = 50.0; ndof4 = 100.0
ndof = exp(x[nparam-1])
nlprior_tmp,dpdx = log_bounds_lower(ndof,ndof1,ndof2)
nlprior -= nlprior_tmp
nlprior_tmp,dpdx = log_bounds_upper(ndof,ndof3,ndof4)
nlprior -= nlprior_tmp
lnV11 = -2.0; lnV12 = -1.0
lnV13 = 5.0; lnV14 = 10.0
lnV1 = log(x[nparam]/exp(0.5/ndof))
nlprior_tmp,dpdx = log_bounds_lower(lnV1,lnV11,lnV12)
nlprior -= nlprior_tmp
nlprior_tmp,dpdx = log_bounds_upper(lnV1,lnV13,lnV14)
nlprior -= nlprior_tmp
# That's it!
return nlprior::T
end

# Take gradient and hessian:
grad_nlog_prior_noecc= x -> ForwardDiff.gradient(nlog_prior_noecc,x)
hess_nlog_prior_noecc= x -> ForwardDiff.hessian(nlog_prior_noecc,x)


"""
    restrict(refspace, element1, element2)

Computes the restriction of a set of local shape functions on `element1` as linear
combinations of the set of local shape functions on `element2`. More precisely `restrict`
returns an `NxM` matrix `P` such that the `i`-th local shape ``g_i`` function on element2
can be written as:

``g_i = sum_{j=1}^{M} P_{ij} f_j``
"""
function restrict end

function restrict(sp::Space, submesh::CompScienceMeshes.AbstractMesh)

    fns = similar(sp.fns)
    S = eltype(eltype(fns))
    for (i,fn) in pairs(sp.fns)
        fns[i] = Vector{S}()
        for sh in fn
            cell = CompScienceMeshes.restrict(submesh, sh.cellid)
            cell == 0 && continue
            push!(fns[i], S(cell, sh.refid, sh.coeff))
        end
    end

    similar(sp, submesh, fns, sp.pos)
end


using LinearAlgebra
using JuMP

function test(n::Int)
    model = Model()
    a = rand(n, n)
    @variable(model, x[1:n, 1:n])
    b = rand(n, n, n)
    @variable(model, y[1:n, 1:n, 1:n])
    c = rand(n)
    @variable(model, z[1:n])

    #initialize
    @constraint(model, sum(c[i] * z[i] for i in 1:n) <= 0)

    #Vector
    elapsed = @elapsed @constraint(model, sum(c[i] * z[i] for i in 1:n) <= 1)
    println("Vector with sum(): $(elapsed)")

    elapsed = @elapsed @constraint(model, dot(c, z) <= 1)
    println("Vector with vecdot() : $(elapsed)")

    #2D Matrix
    elapsed = @elapsed @constraint(
        model,
        sum(a[i, j] * x[i, j] for i in 1:n, j in 1:n) <= 1
    )
    println("2D Matrix with sum(): $(elapsed)")

    elapsed = @elapsed @constraint(model, dot(a, x) <= 1)
    println("2D Matrix with bigvecdot(): $(elapsed)")

    #3D Matrix
    elapsed = @elapsed @constraint(
        model,
        sum(b[i, j, k] * y[i, j, k] for i in 1:n, j in 1:n, k in 1:n) <= 1
    )
    println("3D Matrix with sum(): $(elapsed)")

    elapsed = @elapsed @constraint(model, dot(b, y) <= 1)
    println("3D Matrix with vecdot(): $(elapsed)")
    return 0
end

for n in [10, 50, 100, 200, 300]
    println("n = $n")
    test(n)
end


# test_direct_discretization.jl

# direct discretization for elliptic and anisotropic problem in 2D/3D
function get_problem_2d_elliptic(n)
    A = NotSoSimpleMultigrid.MultigridMethod((n, m)->elliptic2d(f(n, m)), (n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of elliptic PDE, 2D" begin
    for n in [16 32 64 128 256 512]
        (A, b) = get_problem_2d_elliptic(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^2
        log(A,2)
    end
end

function get_problem_2d_anisotropic(n)
    A = NotSoSimpleMultigrid.MultigridMethod((n,m)->anisotropic2d(1e-8,n,m), (n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of anisotropic PDE, 2D" begin
    for n in [16 32 64 128 256 512]
        (A, b) = get_problem_2d_anisotropic(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^2
        log(A,2)
    end
end

function get_problem_2d_rotated_anisotropic(n)
    A = NotSoSimpleMultigrid.MultigridMethod((n,m)->anisotropic2d(1e-4,10*π/180,n,m), (n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of rotated anisotropic PDE, 2D" begin
    for n in [16 32 64 128 256 512]
        (A, b) = get_problem_2d_anisotropic(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^2
        log(A,2)
    end
end

function get_problem_3d_elliptic(n)
    A = NotSoSimpleMultigrid.MultigridMethod((n,m,k)->elliptic3d(f(n,m,k)), (n, n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of elliptic PDE, 3D" begin
    for n in [4 8 16 32 64]
        (A, b) = get_problem_3d_elliptic(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^3
        log(A,3)
    end
end

function get_problem_3d_anisotropic(n)
    A = NotSoSimpleMultigrid.MultigridMethod((n,m,k)->anisotropic3d(1e-8, 1, n, m, k), (n, n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of anisotropic PDE, 3D" begin
    for n in [4 8 16 32 64]
        (A, b) = get_problem_3d_anisotropic(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^3
        log(A,3)
    end
end


# Elementary vector spaces
#--------------------------
abstract ElementarySpace{F} <: VectorSpace
typealias IndexSpace ElementarySpace # index spaces for tensors
# Elementary finite-dimensional vector space over a field F that can be used as the index space corresponding to the indices
# of a tensor. Every elementary vector space should respond to the methods conj and dual, returning the complex conjugate space
# and the dual space respectively. The complex conjugate of the dual space is obtained as dual(conj(V)) === conj(dual(V)). These
# different spaces should be of the same type, so that a tensor can be defined as an element of a homogeneous tensor product
# of these spaces

Base.eltype{F}(V::ElementarySpace{F}) = F
Base.eltype{F}(::Type{ElementarySpace{F}}) = F
Base.eltype{S<:ElementarySpace}(::Type{S}) = eltype(super(S))

Base.isfinite(::ElementarySpace) = true

# The complex conjugate vector space of a real vector space is equal to itself
Base.conj{F<:Real}(V::ElementarySpace{F}) = V

abstract InnerProductSpace{F} <: ElementarySpace{F}
abstract HilbertSpace{F} <: InnerProductSpace{F}
# An inner product space, has the possibility to raise or lower indices of tensors using the metric

const ℝ=Real
const ℂ=Complex{Real}

abstract EuclideanSpace{F<:Union{ℝ,ℂ}} <: HilbertSpace{F}
# Elementary finite-dimensional space R^d or C^d with standard (Euclidean) inner product (i.e. orthonormal basis)

Base.conj(V::EuclideanSpace) = dual(V)
# The complex conjugate space of a Euclidean space is naturally isomorphic to its dual space

dual(V::EuclideanSpace{ℝ}) = V
# For a real Euclidean space, the dual space is naturally isomorphic to the vector space

# Functionality for extracting and iterating over elementary space
Base.length(V::ElementarySpace) = 1
Base.endof(V::ElementarySpace) = 1
Base.getindex(V::ElementarySpace, n::Integer) = (n == 1 ? V : throw(BoundsError()))

Base.start(V::ElementarySpace) = false
Base.next(V::ElementarySpace, state) = (V,true)
Base.done(V::ElementarySpace, state) = state

# CartesianSpace: standard real space R^N with euclidean structure (canonical orthonormal/cartesian basis)
include("spaces/cartesianspace.jl")

# ComplexSpace: standard complex space C^N with euclidean structure (canonical orthonormal basis)
include("spaces/complexspace.jl")

# GeneralSpace: a general elementary vector space (without inner product structure)
include("spaces/generalspace.jl")

# Graded elementary vector spaces
#---------------------------------
# Vector spaces which have a natural decomposition as a direct sum of several sectors

# Sector: hierarchy of types for labelling different sectors, e.g. irreps, quantum numbers, anyon types, ...
abstract Sector
# Abelian: sectors that have an abelian fusion structure
abstract Abelian <: Sector
*(s::Abelian)=s

# Explicit realizations
include("spaces/sectors.jl")

# UnitaryRepresentationSpace: Family of euclidean spaces that are graded corresponding to unitary representations
abstract UnitaryRepresentationSpace{G<:Sector} <: EuclideanSpace{ℂ}

sectortype{G}(V::UnitaryRepresentationSpace{G}) = G
sectortype{G}(::Type{UnitaryRepresentationSpace{G}}) = G
sectortype{S<:UnitaryRepresentationSpace}(::Type{S}) = sectortype(super(S))

#This is a useful shorthand for handling, e.g. the output of sectors()
Base.conj{N,S<:Sector}(nsec::NTuple{N,S}) = map(conj, nsec)

# AbelianSpace: general space that is graded by abelian sectors (i.e. one-dimensional representations)
include("spaces/abelianspace.jl")

# Auxiliary functionality:
#This returns the dimension of the degeneracy space specified by the irrep label assignment s.
_dim{G<:Sector,N}(spaces::NTuple{N,UnitaryRepresentationSpace{G}},s::NTuple{N,G})=(d=1;for n=1:N;d*=dim(spaces[n],s[n]);end;return d)

#This function returns all possible irrep label assignments for a space of indices
@generated function _sectors{G<:Sector,N}(spaces::NTuple{N,UnitaryRepresentationSpace{G}})
    quote
        numsectors=1
        @nexprs $N i->(s_i=collect(sectors(spaces[i]));numsectors*=length(s_i))
        sectorlist=Array(NTuple{N,G},numsectors)
        counter=0
        @nloops $N i d->1:length(s_d) begin
            counter+=1
            sectorlist[counter]=@ntuple $N k->s_k[i_k]
        end
        return sectorlist
    end
end

#This function finds the combinations of irrep labels that can habe nonzero value in an invariant tensor
@generated function _invariantsectors{G<:Abelian,N}(spaces::NTuple{N,AbelianSpace{G}})
    quote
        @nexprs $N i->(s_i=collect(sectors(spaces[i]))) #generates s_n = collect(sectors(spaces[i])) for n in 1:N
        sectorlist=Array(NTuple{N,G},0) #a sector is specified by N irrep labels (instances of G)
        @nloops $N i d->1:length(s_d) begin #generates N nested loops: for i_n in 1:length(s_n)
            sector=@ntuple $N k->s_k[i_k]   #creates an N-tuple: (s_1[i_1], s_2[i_2], ..., s_N[i_N])
            c=one(G)
            for n=1:N
                c*=sector[n] #fuses all the irrep labels in sector
            end
            if c==one(G) #if they fuse to the trivial irrep, we've found an allowed block: add to the list!
                push!(sectorlist,sector)
            end
        end
        return sectorlist
    end
end


# lanczos absorption

"""
    lanczos(A, u_init, k[, M])

Perform k steps of Lanczos method for the operator `A` and initial
vector `u_init`. The output are the vectors `α`, `β` and `U` with
`A * U ≈ U * SymTridiagonal(α, β[1:(end - 1)])`.
"""
function lanczos(A, u_init, k, M = I)
    α = zeros(k)
    β = zeros(k)
    u = copy(u_init)
    u_old = zero(u)
    v_old = zero(u)
    U = zeros(eltype(u), length(u), k + 1)
    U[:, 1] .= u

    v = M * u
    x = A * v
    α[1] = real(v' * x)
    x .-= α[1] .* u
    y = M * x
    β[1] = sqrt(x' * y)
    u_old .= u
    u .= x ./ β[1]
    U[:, 2] .= u
    v_old .= v
    v .= y ./ β[1]
    @showprogress 1 "Lanczos iteration ..." for j in 2:k
        x .= A * v .- β[j - 1] .* u_old
        α[j] = real(v' * x)
        x .-= α[j] .* u
        y .= M * x
        β[j] = sqrt(real(x' * y))
        u_old .= u
        u .= x ./ β[j]
        U[:, j + 1] .= u
        v_old .= v
        v .= y ./ β[j]
    end
    return α, β, U
end

"""
    lanczos_eig(H, d, N_iter)

Compute eigenvalues and eigenvectors of an approximation of `H` based
on `N_iter` steps of Lanczos method with initial vector `d`.
"""
function lanczos_eig(H, d, N_iter)
    α, β, U = lanczos(H, d, N_iter)
    return lanczos_eig(α, β, N_iter)
end

function lanczos_eig(α::AbstractVector, β, N_iter)
    T̂ = SymTridiagonal(vcat(α[1:N_iter], α[(N_iter - 1):-1:1]), vcat(β[1:N_iter], β[(N_iter - 2):-1:1]))
    F = eigen(T̂)
    ev_lanczos, ef_lanczos = F.values, F.vectors
    return ev_lanczos, ef_lanczos
end

"""
    lanczos_optical_absorption(prob, isdf, direction, N_iter, g, Erange)

Compute the optical absorption function of `prob` in direction
`a_mat[i]` using `N_iter` steps of a Lanczos method. The absorption
spectrum is computed on the points in `Erange` and the function `g`
is used for broadening.
"""
function lanczos_optical_absorption(prob::AbstractBSEProblem, isdf::ISDF, direction, N_iter, g, Erange)
    Ω0_vol = Ω0_volume(prob)
    N_k = size(prob)[3]

    H = setup_H(prob, isdf)
    d = optical_absorption_vector(prob, direction)

    return lanczos_optical_absorption(H, d, N_iter, g, Erange, norm(d)^2 * 8 * π^2 / (Ω0_vol * N_k))
end

function lanczos_optical_absorption(H, d, N_iter, g, Erange, scaling)
    α, β, U = lanczos(H, normalize(d), N_iter)
    return lanczos_optical_absorption(α, β, N_iter, g, Erange, scaling)
end

function lanczos_optical_absorption(α::AbstractVector, β, N_iter, g, Erange, scaling)
    ev_lanczos, ef_lanczos = lanczos_eig(α, β, N_iter)
    return lanczos_optical_absorption(ev_lanczos, abs2.(ef_lanczos[1, :]), g, Erange, scaling)
end

function lanczos_optical_absorption(ev::AbstractVector, weights::AbstractVector, g, Erange, scaling)
    optical_absorption = zeros(length(Erange))
    for j in 1:length(ev)
        if ev[j] >= 0
            optical_absorption .+= weights[j] .* (g.(Erange .- ev[j]) .- g.(Erange .+ ev[j]))
        end
    end
    optical_absorption .*= scaling
    return optical_absorption
end


export constraint2coarse

"""
Adapts constraint definitions for the original grid, to coarser grids.
Some of these 'rules' are empirical or user defined. Some constraints are not adapted as
they are often considered invariant under subsampling the grid for imaging applications.
"""
function constraint2coarse(constraint,comp_grid,coarsening_factor)

  nr_constraints = length(constraint)

  #bound constraints: same as on fine grid

  #rank: same as on fine grid (limit by minimum dimension)
  for i=1:nr_constraints
    if constraint[i].set_type == "rank"
      constraint[i].max = min(constraint[i].max,minimum(comp_grid.n))
    end
  end

  #cardinality in a transform-domain: same as on fine grid (limit by number of elements)
  for i=1:nr_constraints
    if constraint[i].set_type == "cardinality"
      constraint[i].max = min(constraint[i].max,prod(comp_grid.n))
    end
  end

  #define below what to do with norms (l1, l2 and nuclear)
  # will probably depend on constraints and interpolation type that is used
  # below we use simple heuristics, better methods may improve performance of the multilevel scheme

if length(comp_grid.n)==3 && comp_grid.n[3]>1 #use 3D

  # #for point-wise bound constraints in a transform-domain on a 3D grid if the transform-domain is a derivative operator:
  # for i=1:3
  #   if haskey(constraint,string("use_TD_bounds_",i)) && (constraint[string("use_TD_bounds_",i)]==true)
  #     if (constraint[string("TDB_operator_",i)] in ["D_x","D_y","D_z","DFT"])
  #       constraint[string("TD_LB_",i)]=constraint[string("TD_LB_",i)].*coarsening_factor
  #       constraint[string("TD_UB_",i)]=constraint[string("TD_UB_",i)].*coarsening_factor
  #     end
  #   end
  # end

  #for l1 norm in a transform-domain: on a 3D grid: ||.||_1(coarse)=||.||_1(fine)/(coarsening factor^3)
  for i=1:nr_constraints
    if constraint[i].set_type == "l1"
      constraint[i].max = constraint[i].max/(coarsening_factor^3)
    end
  end

  #for l2 norm in a transform-domain: on a 3D grid: ||.||_2(coarse)=||.||_2(fine)/sqrt(coarsening factor^3)
  for i=1:nr_constraints
    if constraint[i].set_type == "l2"
      constraint[i].max = constraint[i].max/(sqrt(coarsening_factor^3))
    end
  end

else #use 2D

  # #for point-wise bound constraints in a transform-domain on a 2D grid if the transform-domain is a derivative operator:
  # for i=1:3
  #   if haskey(constraint,string("use_TD_bounds_",i)) && (constraint[string("use_TD_bounds_",i)]==true)
  #     if (constraint[string("TDB_operator_",i)] in ["D_x","D_y","D_z","DFT"])
  #       constraint[string("TD_LB_",i)]=constraint[string("TD_LB_",i)].*coarsening_factor
  #       constraint[string("TD_UB_",i)]=constraint[string("TD_UB_",i)].*coarsening_factor
  #     end
  #   end
  # end

    #for l1 norm in a transform-domain: on a 2D grid: ||.||_1(coarse)=||.||_1(fine)/(coarsening factor^2)
    for i=1:nr_constraints
      if constraint[i].set_type == "l1"
        constraint[i].max = constraint[i].max/(coarsening_factor^2)
      end
    end

    #for l2 norm in a transform-domain: on a 2D grid: ||.||_2(coarse)=||.||_2(fine)/(coarsening factor)
    for i=1:nr_constraints
      if constraint[i].set_type == "l2"
      constraint[i].max = constraint[i].max/(coarsening_factor)
      end
    end

    #nuclear norm:
    for i=1:nr_constraints
      if constraint[i].set_type == "nuclear"
        constraint[i].max = constraint[i].max/2.7
      end
    end

end #END if 2D or 3D

  # #subspace constraint
  # if constraint[i].set_type == "subspace"
  #   error("still need to define how to map subspace to a coarser grid in: function constraint2coarse")
  #   #maybe just use coarsening of matrix to start simple
  # end

return constraint
end


### Types
abstract type ApproximationScheme end

struct RosenfeldFMT{T} <: ApproximationScheme
    A::T
    B::T
    G::T
end

struct PercusYevick{T} <: ApproximationScheme
    α::T
    β::T
    δ::T
end

struct VerletWeis{T} <: ApproximationScheme
    coreliquid::HardSpheres{T}
    subscheme::PercusYevick{T}
    α::T
end

struct SharmaSharma{S, T} <: ApproximationScheme
    coreliquid::HardSpheres{T}
    subscheme::S
end

struct NonLinearSharmaSharma{S, T} <: ApproximationScheme
    coreliquid::HardSpheres{T}
    subscheme::S
end

struct MSA{S, T} <: ApproximationScheme
    coreliquid::HardSpheres{T}
    subscheme::S
    α₀::T
    α₁::T
    α₂::T
    α₃::T
    β₀::T
    β₁::T
    β₂::T
    β₃::T
    a₁::T
    a₃::T
    b₁::T
    b₃::T
end

### Constructors
function RosenfeldFMT(liquid::HardDisks)
    η  = volume_fraction(liquid)
    η² = η * η
    η³ = η² * η
    η⁴ = η³ * η
    η₋⁻¹ = 1 / (1 - η)
    η₋⁻² = η₋⁻¹ * η₋⁻¹

    Z  = (1 + 0.128η² + 0.027η³ + 0.06η⁴) * η₋⁻²
    Z′ = ((0.256η + 0.081η² + 0.24η³) * η₋⁻²) + (2Z * η₋⁻¹)
    χ  = Z + η * Z′

    G = sqrt(Z′ / 2)
    A = (1 + (2η - 1) * χ + 2η * G) / η
    B = (-1 + (1 - η) * χ - 3η * G) / η

    return RosenfeldFMT{typeof(A)}(A, B, G)
end

function PercusYevick(liquid::HardSpheres)
    η = volume_fraction(liquid)
    η₀⁴ = @fastmath (1 - η)^4
    η₁² = (1 + 2η)^2
    η₂² = (2 +  η)^2

    α = -η₁² / η₀⁴
    β = 3η * η₂² / 2η₀⁴
    δ = -η * η₁² / 2η₀⁴

    return PercusYevick{typeof(α)}(α, β, δ)
end

function VerletWeis(liquid::HardSpheres)
    η  = volume_fraction(liquid)
    κ  = 1 - η / 16
    η′ = η * κ
    α  = ∛(κ)

    coreliquid = HardSpheres(η′)
    subscheme = PercusYevick(coreliquid)

    return VerletWeis{typeof(α)}(coreliquid, subscheme, α)
end

function SharmaSharma{S}(liquid::AttractiveHardSpheres{TT}) where
    {S<:ApproximationScheme, TT}

    η = volume_fraction(liquid)
    coreliquid = HardSpheres(η)
    subscheme = S(coreliquid)
    SS = typeof(subscheme)

    return SharmaSharma{SS, TT}(coreliquid, subscheme)
end

function NonLinearSharmaSharma{S}(liquid::AttractiveHardSpheres{TT}) where
    {S <: ApproximationScheme, TT}

    η = volume_fraction(liquid)
    coreliquid = HardSpheres(η)
    subscheme = S(coreliquid)
    SS = typeof(subscheme)

    return NonLinearSharmaSharma{SS, TT}(coreliquid, subscheme)
end

function MSA{S}(liquid::DipolarHardSpheres{TT}; tol = √(eps(TT))) where
    {S<:ApproximationScheme, TT}

    η = volume_fraction(liquid)
    T = temperature(liquid)
    coreliquid = HardSpheres(η)
    subscheme = S(coreliquid)
    SS = typeof(subscheme)
    κ = dhs_msa_parameter(liquid, tol)

    T⁻¹ = 1 / T
    ξ   = κ * η

    ξ₁² = (1 - 2ξ)^2
    ξ₁⁴ = ξ₁² * ξ₁²
    ξ₂⁴ = (1 +  ξ)^4
    ξ₃² = (1 + 4ξ)^2

    r₁ = ξ₃² / ξ₁⁴
    r₂ = ξ₁² / ξ₂⁴
    r₃ = 8 * (1 + ξ)^2 / ξ₁⁴
    r₄ = (2 - ξ)^2 / ξ₂⁴

    α₀′ = 2κ * (-r₁ + r₂)
    α₁′ = 3ξ * κ * ( r₃ + r₄)
    α₃′ = -ξ * κ * (2r₁ + r₂)
    β₁′ = 3ξ * κ / 8 * (2r₃ - r₄)
    β₃′ = -ξ * κ / 4 * (4r₁ - r₂)

    α₀ = α₀′ / 9  + α₁′ / 12  + α₃′ / 18
    α₁ = α₀′ / 90 + α₁′ / 108 + α₃′ / 144
    α₂ = α₀′ +  α₁′ +  α₃′
    α₃ = α₀′ + 2α₁′ + 4α₃′

    β₀ = -T⁻¹ / 9
    β₁ =  β₁′ / 270 + β₃′ / 360 + β₀ / 10
    β₂ =  β₁′ +  β₃′ - T⁻¹
    β₃ = 4β₁′ + 6β₃′ + β₂

    a₁ = 2 * (α₁′ + 8β₁′) / 3
    a₃ = 8 * (α₃′ + 4β₃′)
    b₁ = 2 * (α₁′ - 4β₁′) / 3
    b₃ = 8 * (α₃′ - 2β₃′)

    return MSA{SS, TT}(coreliquid, subscheme,
                       α₀, α₁, α₂, α₃, β₀, β₁, β₂, β₃, a₁, a₃, b₁, b₃)
end


export SO2,
        matrix,
        generator,
        transformation

using StaticArrays
using LinearAlgebra

mutable struct SO2{T<:AbstractFloat}
    DoF::Int8
    tangent::T
    generator::SMatrix{2, 2, T}

    function SO2(tangent::AbstractFloat)
        generator = SMatrix{2, 2, typeof(tangent)}(0, -1, 1, 0)
        new{typeof(tangent)}(Int8(1), tangent, generator)
    end
end

function generator(group::SO2)
    return group.generator
end

function matrix(group::SO2)
    return group.tangent * SMatrix{2, 2, typeof(group.tangent)}(0, -1, 1, 0)
end

function transformation(group::SO2)
    # println(typeof(group.tangent))
    return exp.(group.tangent * group.generator)
end


s = open("input.txt") do file
    readlines(file);
end

startTime = parse(Int, s[1]);

nums = split(s[2], ',');
offsets = findall(x->x!="x", nums) .- 1;
busNums = map(x->parse(Int,x), nums[offsets.+1]);

function FindNextMatchup(startTime, iter, offset, num)
    curTime = startTime;
    while (curTime+offset)%num != 0
        curTime += iter;
    end
    println((curTime, iter*num));
    return (curTime, iter*num);
end

function GetPatternTime(start, startIter, indOrder, offsets, nums)
    iter = startIter;
    time = start;
    for i = 1:length(indOrder)
        curInd = indOrder[i];
        time, iter = FindNextMatchup(time, iter, offsets[curInd], nums[curInd]);
    end

    return (time, iter);
end

function CheckTime(time, offsets, nums)
    println(time);
    times = offsets .+ time;
    mods = times.%nums;
    println(mods);

    return all(mods.==0);
end

startTime = 100000000000000;
while startTime % busNums[1] != 0
    startTime += 1
end

# The start time came from past brute force efforts,
# But you could just use the above
sortedIndices = sortperm(busNums[2:end]);
theTime, bigIter = GetPatternTime(141883788274894, busNums[1], sortedIndices, offsets[2:end], busNums[2:end]);
CheckTime(theTime, offsets, busNums)


include("../src/csd.jl")

using DataFrames
# using DelimitedFiles
using CSV

"""
Timing and backward stability testing of the CS Decomposition (special case)

First, we generate an (m+l)-by-(m+l) random matrix A
Then, we compute QR decomposition of A, A = QR
Afterwards, we choose Q1 = Q[1:m, 1:l]
Q2 = Q[m+1:m+l, 1:l], such that:

      Q1' * Q1 + Q2' * Q2 = I

Call csd function, return U, V, Z, cosine, sine.
Construct C and S.

Assume ||Q1' * Q1 + Q2' * Q2 - I|| ≈ eps,
compute:

      res_q1 = ||U * C * Z' - Q1||/(l * ||Q1|| * eps)

      res_q2 = ||V * S * Z' - Q2||/(l * ||Q2|| * eps)

      orthog_cs = ||C' * C + S' * S -I||/(l * eps)

      orthog_u = ||U' * U - I||/(m * eps)

      orthog_v = ||V' * V - I||/(p * eps)

      orthog_z = ||Z' * Z - I||/(l * eps)

argument
m: num of rows of Q1
l: num of cols of Q1, Q2 and num of rows of Q2

no return
"""

function test(m, l)
    if m > l
        throw(ArgumentError("m <= l is required"))
    else
        F = qr(randn(m+l, l))

        Q1 = @views Matrix(F.Q)[1:m, :]
        Q2 = @views Matrix(F.Q)[m+1:m+l, :]

        e = eps(Float64)
        U, V, Z, cosine, sine = csd2by1(Q1, Q2)
        if m == l
            C = Diagonal(cosine)
            S = Diagonal(sine)
        else
            C = [Diagonal(cosine[1:m]) zeros(m, l-m)]
            S = Diagonal([sine[1:m]; ones(l-m)])
        end
        res_q1 = opnorm(U * C * Z' - Q1, 1)/(l*opnorm(Q1, 1)*e)
        res_q2 = opnorm(V * S * Z' - Q2, 1)/(l*opnorm(Q2, 1)*e)
        orthog_cs = opnorm(C' * C + S' * S - I, 1)/(l*e)
        orthog_u = opnorm(U' * U - I, 1)/(m*e)
        orthog_v = opnorm(V' * V - I, 1)/(l*e)
        orthog_z = opnorm(Z' * Z - I, 1)/(l*e)
        return res_q1, res_q2, orthog_cs, orthog_u, orthog_v, orthog_z
    end
end

function test()
    df = DataFrame(m = Int64[], l = Int64[], res_q1 = Float64[],
    res_q2 = Float64[], orthog_cs = Float64[], orthog_u = Float64[],
    orthog_v = Float64[], orthog_z = Float64[])

    # case 1: m < l
    println("===============================================================")
    for i = 1:20
        a, b, c, d, e, f = test(40, 60)
        push!(df, [40 60 a b c d e f])
        println("---------------------------------------------------------------")
    end
    for i = 1:20
        a, b, c, d, e, f = test(200, 300)
        push!(df, [200 300 a b c d e f])
        println("---------------------------------------------------------------")
    end
    for i = 1:20
        a, b, c, d, e, f = test(600, 900)
        push!(df, [600 900 a b c d e f])
        println("---------------------------------------------------------------")
    end
    # May throw LAPACK Exception at line of svd(Q2)
    # for i = 1:20
    #     a, b, c, d, e, f = test(1000, 1500)
    #     push!(df, [1000 1500 a b c d e f])
    #     println("---------------------------------------------------------------")
    # end

    # case 2: m = l
    println("===============================================================")
    for i = 1:20
        a, b, c, d, e, f = test(50, 50)
        push!(df, [50 50 a b c d e f])
        println("---------------------------------------------------------------")
    end
    for i = 1:20
        a, b, c, d, e, f = test(250, 250)
        push!(df, [250 250 a b c d e f])
        println("---------------------------------------------------------------")
    end
    for i = 1:20
        a, b, c, d, e, f = test(750, 750)
        push!(df, [750 750 a b c d e f])
        println("---------------------------------------------------------------")
    end
    for i = 1:20
        a, b, c, d, e, f = test(1250, 1250)
        push!(df, [1250 1250 a b c d e f])
        println("---------------------------------------------------------------")
    end

    CSV.write("rand_matrix_csd.csv", df, append=true)
    return df
end

# Example 2 in Chapter 3 of Jenny's thesis
function test2()
    Q1 = [1/sqrt(7) 0 0 1/sqrt(3);
          1/sqrt(7) -2/sqrt(10) -1/2 -1/(2*sqrt(3));
          1/sqrt(7) -1/sqrt(10) 1/4  3/(4*sqrt(3))]
    Q2 = [1/sqrt(7) 0 3/4 -3/(4*sqrt(3));
          1/sqrt(7) 0 0 0;
          1/sqrt(7) 1/sqrt(10) -1/4 -3/(4*sqrt(3));
          1/sqrt(7) 2/sqrt(10) -1/4 1/(4*sqrt(3))]
    m = 3; l = 4;
    e = eps(Float64)

    return U, V, Z, cosine, sine = csd2by1(Q1, Q2)
end


using DelimitedFiles

function solve(input)
    raw = vcat(readdlm(input,Int16)...)
    p1 = process(copy(raw),aggregate)
    @show p1
    p2 = process(copy(raw),findRootValue)
    @show p2

end

function process(list, f)
    numChildren = popfirst!(list)
    numMeta = popfirst!(list)
    return f(ntimes(numChildren,list,f), splice!(list,1:numMeta))
end

function ntimes(n,list, f)
    result = Int16[]
    for i in 1:n
        push!(result,process(list,f))
    end
    return result
end

function aggregate(child,metadata)
    return sum(child) + sum(metadata)
end

function childValue((child,value),index)
    if length(child)>index
        value+=child[index+1]
    else
        value+= 0
    end
    return (child,value)
end

function findRootValue(child,metadata)
    pushfirst!(child,0)
    if length(child) == 1
        return sum(metadata)
    else
        return reduce(childValue,metadata,init=(child,0))[2]
    end
end

solve("input.txt")


"""
    create_saturated_ols_approximation(dd::DataFrame, y::Symbol, x_variables::Array{Symbol,1}, degree::Int; intercept::Bool = true,  bases::Dict{Symbol,Float64} = Dict{Symbol,Float64}(x_variables .=> repeat([0.0],length(x_variables))))
This creates MultivariateFunction from an OLS regression predicting some variable. You input a dataframe and specify what column in that dataframe
is to be predicted by inputting a symbol y. you also put in an array of what x_variables should be used in prediction. A saturated ols model is then calculated up to the specified degree which is returned as a MultivariateFunction.
"""
function create_saturated_ols_approximation(dd::DataFrame, y::Symbol, x_variables::Array{Symbol,1}, degree::Int; intercept::Bool = true,  bases::Dict{Symbol,Float64} = Dict{Symbol,Float64}(x_variables .=> repeat([0.0],length(x_variables))))
    model = Array{PE_Function,1}()
    if intercept
        append!(model, [PE_Function(1.0,Dict{Symbol,PE_Unit}())] )
    end
    if degree > 0
        number_of_variables = length(x_variables)
        linear_set = Array{PE_Function,1}(undef, number_of_variables)
        for i in 1:length(x_variables)
            linear_set[i] = PE_Function(1.0,Dict{Symbol,PE_Unit}(x_variables[i] => PE_Unit(0.0,bases[x_variables[i]],1) ))
        end
        higher_order_terms = Array{Array{PE_Function,1},1}(undef,degree)
        higher_order_terms[1] = linear_set
        for i in 2:degree
            degree_terms = Array{PE_Function,1}()
            for j in 1:number_of_variables
                append!(degree_terms, linear_set[j] .* hcat(higher_order_terms[i-1]))
            end
            higher_order_terms[i] = degree_terms
        end
        append!(model, vcat(higher_order_terms...))
    end
    sum_of_functions = Sum_Of_Functions(model) # We put it through here to remove duplicates.
    return create_ols_approximation(dd, y, sum_of_functions)
end

"""
    create_ols_approximation(dd::DataFrame, y::Symbol, model::MultivariateFunction; allowrankdeficient = true)
    create_ols_approximation(dd::DataFrame, y::Symbol, model::Sum_Of_Functions; allowrankdeficient = true)
    create_ols_approximation(dd::DataFrame, y::Symbol, model::Sum_Of_Piecewise_Functions; allowrankdeficient = true)
This creates MultivariationFunction from an OLS regression predicting some variable. You input a dataframe and specify what column in that dataframe
is to be predicted by inputting a symbol y. You also input the regression model. This is input as a Array{MultivariateFunction,1}.
Each function that is input will be multiplied by the ols coefficient and will return a new function with these coefficients
incorporated.
"""
function create_ols_approximation(dd::DataFrame, y::Symbol, model::Array; allowrankdeficient = true)
    X = hcat(evaluate.(model, Ref(dd))...)
    y = dd[y]
    reg = fit(LinearModel, X,y, allowrankdeficient)
    coefficients = reg.pp.beta0
    updated_model = Sum_Of_Piecewise_Functions(model .* coefficients)
    return updated_model, reg
end
function create_ols_approximation(dd::DataFrame, y::Symbol, model::Sum_Of_Functions)
    return create_ols_approximation(dd, y, model.functions_)
end
function create_ols_approximation(dd::DataFrame, y::Symbol, model::Sum_Of_Piecewise_Functions)
    return create_ols_approximation(dd, y, vcat( model.functions_, model.global_funcs_ ))
end
"""
    create_ols_approximation(y::Array{Float64,1}, x::Array{Float64,1}, degree::Int; intercept::Bool = true, dim_name::Symbol = default_symbol, base_x::Float64 = 0.0)
    create_ols_approximation(y::Array{Float64,1}, x::Array{Date,1}, degree::Int; intercept::Bool = true, dim_name::Symbol = default_symbol, base_date::Date = global_base_date)

This predicts a linear relationship between the y and x arrays and creates a MultivariateFunction containing the approximation function. The degree specifies how many higher
order terms of x should be used (for instance degree 2 implies x and x^2 are both used to predict y).
"""
function create_ols_approximation(y::Array{Float64,1}, x::Array{Float64,1}, degree::Int; intercept::Bool = true, dim_name::Symbol = default_symbol, base_x::Float64 = 0.0)
    dd = DataFrame()
    dd[dim_name] = x
    dd[:y]       = y
    base_dict = Dict{Symbol,Float64}(dim_name => base_x)
    return create_saturated_ols_approximation(dd, :y, [dim_name], degree; intercept = intercept, bases = base_dict)
end

function create_ols_approximation(y::Array{Float64,1}, x::Array{Date,1}, degree::Int; intercept::Bool = true, dim_name::Symbol = default_symbol, base_date::Date = global_base_date)
    return  create_ols_approximation(y, years_from_global_base.(x), degree; intercept = intercept, dim_name = dim_name, base_x = years_from_global_base(base_date))
end


using Flux, MLDataPattern
using Flux: onehotbatch
include("spectrograms.jl")

# Image size constants
img_n = 150
img_m = 100

# Definition of a dataset with training and testing partitions
struct DataSet
  X_train::Array{Float32, 2}
  Y_train::Flux.OneHotMatrix
  X_test::Array{Float32, 2}
  Y_test::Flux.OneHotMatrix
end

# Split the data
function create_data_splits(n=img_n, m=img_m)
  @info("Generating spectrograms of size $(n)x$(m)...")

  # Read labels from source CSV file
  file_to_class = gen_label_dict("train/train.csv")
  labels = simple_labels(file_to_class)

  # Create spectrograms with values between 0 and 1 (via sigmoid)
  spec_base = create_spectrogram_base(file_to_class, n, m)

  # Format dataset
  X = hcat(float.(reshape.(spec_base, :))...) |> gpu
  XY = vcat(X, labels')

  get_label(x) = convert(Int, x[end])

  @info("Generating train and test splits...")

  # Stratified split: 80% for training and 20% for testing
  # Encode Ys with one-hot encoding
  train_data, test_data = stratifiedobs(get_label, XY, 0.8)
  X_train = train_data[1:end-1,:]
  Y_train = onehotbatch(convert(Array{Int}, train_data[end,:]), 0:9)
  X_test = test_data[1:end-1,:]
  Y_test = onehotbatch(convert(Array{Int}, test_data[end,:]), 0:9)

  @info("Done")

  return DataSet(X_train, Y_train, X_test, Y_test)
end


export evaluate

"""
__Univariate__
```
evaluate(n::Int,x::Array{<:Real},a::AbstractVector{<:Real},b::AbstractVector{<:Real})
evaluate(n::Int,x::Real,a::AbstractVector{<:Real},b::AbstractVector{<:Real})
evaluate(n::Int,x::AbstractVector{<:Real},op::AbstractOrthoPoly)
evaluate(n::Int,x::Real,op::AbstractOrthoPoly)
```
Evaluate the `n`-th univariate basis polynomial at point(s) `x`
The function is multiply dispatched to facilitate its use with the composite type `AbstractOrthoPoly`

If several basis polynomials (stored in `ns`) are to be evaluated at points `x`, then call
```
evaluate(ns::AbstractVector{<:Int},x::AbstractVector{<:Real},op::AbstractOrthoPoly) = evaluate(ns,x,op.α,op.β)
evaluate(ns::AbstractVector{<:Int},x::Real,op::AbstractOrthoPoly) = evaluate(ns,[x],op)
```

If *all* basis polynomials are to be evaluated at points `x`, then call
```
evaluate(x::AbstractVector{<:Real},op::AbstractOrthoPoly) = evaluate(collect(0:op.deg),x,op)
evaluate(x::Real,op::AbstractOrthoPoly) = evaluate([x],op)
```
which returns an Array of dimensions `(length(x),op.deg+1)`.

!!! note
    - `n` is the degree of the univariate basis polynomial
    - `length(x) = N`, where `N` is the number of points
    - `(a,b)` are the recursion coefficients

__Multivariate__
```
evaluate(n::AbstractVector{<:Int},x::AbstractMatrix{<:Real},a::Vector{<:AbstractVector{<:Real}},b::Vector{<:AbstractVector{<:Real}})
evaluate(n::AbstractVector{<:Int},x::AbstractVector{<:Real},a::Vector{<:AbstractVector{<:Real}},b::Vector{<:AbstractVector{<:Real}})
evaluate(n::AbstractVector{<:Int},x::AbstractMatrix{<:Real},op::MultiOrthoPoly)
evaluate(n::AbstractVector{<:Int},x::AbstractVector{<:Real},op::MultiOrthoPoly)
```
Evaluate the n-th p-variate basis polynomial at point(s) x
The function is multiply dispatched to facilitate its use with the composite type `MultiOrthoPoly`

If several basis polynomials are to be evaluated at points `x`, then call
```
evaluate(ind::AbstractMatrix{<:Int},x::AbstractMatrix{<:Real},a::Vector{<:AbstractVector{<:Real}},b::Vector{<:AbstractVector{<:Real}})
evaluate(ind::AbstractMatrix{<:Int},x::AbstractMatrix{<:Real},op::MultiOrthoPoly)
```
where `ind` is a matrix of multi-indices.

If *all* basis polynomials are to be evaluated at points `x`, then call
```
evaluate(x::AbstractMatrix{<:Real},mop::MultiOrthoPoly) = evaluate(mop.ind,x,mop)
```
which returns an array of dimensions `(mop.dim,size(x,1))`.

!!! note
    - `n` is a multi-index
    - `length(n) == p`, i.e. a p-variate basis polynomial
    - `size(x) = (N,p)`, where `N` is the number of points
    - `size(a)==size(b)=p`.
"""
function evaluate(n::Int,x::AbstractArray{<:Real},a::AbstractVector{<:Real},b::AbstractVector{<:Real})
    @assert n >= 0 "Degree n has to be non-negative (currently n=$n)."
    # if length(a)==0 warn("Length of a is 0.") end
    @assert length(a) == length(b) "Inconsistent number of recurrence coefficients."
    @assert n <= length(a) "Specified degree is $n, but you only provided $(length(a)) coefficients."
    # recurrence relation for orthogonal polynomials
    nx = length(x)
    pminus, p = zeros(nx), ones(nx)
    if n==0
        if nx==1
            return first(p)
        else
            return p
        end
    end
    pplus = (x .- first(a)).*p .- first(b)*pminus
    for k in 2:n
        pminus = p
        p = pplus
        @inbounds pplus = (x .- a[k]).*p .- b[k]*pminus
    end
    nx == 1 ? first(pplus) : pplus
end
evaluate(n::Int,x::Real,a::AbstractVector{<:Real},b::AbstractVector{<:Real}) = evaluate(n,[x],a,b)
evaluate(n::Int,x::AbstractVector{<:Real},op::AbstractOrthoPoly) = evaluate(n,x,op.α,op.β)
evaluate(n::Int,x::Real,op::AbstractOrthoPoly) = evaluate(n,[x],op)

# univariate + several bases
function evaluate(ns,x::AbstractArray{<:Real},a::AbstractVector{<:Real},b::AbstractVector{<:Real})
    hcat(map(i->evaluate(i,x,a,b),ns)...)
end
evaluate(ns,x::Real,a::AbstractVector{<:Real},b::AbstractVector{<:Real}) = evaluate(ns,[x],a,b)

evaluate(ns,x::AbstractVector{<:Real},op::AbstractOrthoPoly) = evaluate(ns,x,op.α,op.β)
evaluate(ns,x::Real,op::AbstractOrthoPoly) = evaluate(ns,[x],op)
evaluate(x::AbstractVector{<:Real},op::AbstractOrthoPoly) = evaluate(collect(0:op.deg),x,op)
evaluate(x::Real,op::AbstractOrthoPoly) = evaluate([x],op)

# multivariate
function evaluate(n::AbstractVector{<:Int},x::AbstractMatrix{<:Real},a::AbstractVector{<:AbstractVector{<:Real}},b::AbstractVector{<:AbstractVector{<:Real}})
    @assert length(n) == size(x,2) "number of univariate bases (= $(length(n))) inconsistent with columns points x (= $(size(x,2)))"
    val = ones(Float64,size(x,1))
    for i in 1:length(n)
        @inbounds val = val.*evaluate(n[i],x[:,i],a[i],b[i])
    end
    return val
end
evaluate(n::AbstractVector{<:Int},x::AbstractVector{<:Real},a::AbstractVector{<:AbstractVector{<:Real}},b::AbstractVector{<:AbstractVector{<:Real}}) = evaluate(n,reshape(x,1,length(x)),a,b)
evaluate(n::AbstractVector{<:Int},x::AbstractMatrix{<:Real},op::MultiOrthoPoly) = evaluate(n,x,coeffs(op)...)
evaluate(n::AbstractVector{<:Int},x::AbstractVector{<:Real},op::MultiOrthoPoly) = evaluate(n,reshape(x,1,length(x)),op)

# using multi-index + multivariate
function evaluate(ind::AbstractMatrix{<:Int},x::AbstractMatrix{<:Real},a::AbstractVector{<:AbstractVector{<:Real}},b::AbstractVector{<:AbstractVector{<:Real}})
    vals = map(i->evaluate(ind[i,:],x,a,b),Base.OneTo(size(ind,1)))
    hcat(vals...) |> transpose |> Matrix
end

evaluate(ind::AbstractMatrix{<:Int},x::AbstractMatrix{<:Real},op::MultiOrthoPoly) = evaluate(ind,x,coeffs(op)...)
evaluate(x::AbstractMatrix{<:Real},mop::MultiOrthoPoly) = evaluate(mop.ind,x,mop)

evaluate(ind::AbstractMatrix{<:Int},x::AbstractVector{<:Real},a::AbstractVector{<:AbstractVector{<:Real}},b::AbstractVector{<:AbstractVector{<:Real}}) = evaluate(ind,reshape(x,1,length(x)),a,b)
evaluate(ind::AbstractMatrix{<:Int},x::AbstractVector{<:Real},op::MultiOrthoPoly) = evaluate(ind,reshape(x,1,length(x)),coeffs(op)...)
evaluate(x::AbstractVector{<:Real},mop::MultiOrthoPoly) = evaluate(mop.ind,reshape(x,1,length(x)),mop)


"""
    sample_vdpoisson(shape::Tuple{Int64},redFac;seed::Int64=1234,kargs...)

generates a Poisson disk sampling pattern for an Vector of size `shape` with a subsampling factor `redFac`.
The arguments are the same as in the 2d case with the exception that shape is of type `Tuple{Int64}`
"""
function sample_vdpoisson(shape::Tuple{Int64},redFac;seed::Int64=1234,kargs...)
  return sample_vdpoisson((shape[1],1),redFac,seed=seed;kargs...)
end

function sample_vdpoisson(shape::Tuple{Int64,Int64,Int64},redFac;seed::Int64=1234,kargs...)
  error("Not implemented")
end

"""
    sample_vdpoisson(shape::Tuple{Int64,Int64},redFac::Float64; seed::Int64=1234,kargs...)

generates a variable density Poisson disk sampling pattern for an Array of size `shape` with a subsampling factor `redFac`.

# Arguments
* `shape::NTuple{2,Int64}` - size of the Array to be sampled
* `redFac::Float64`        - subsampling factor
* (`seed=1234`)            - seed for the random number generator
"""
function sample_vdpoisson(shape::Tuple{Int64,Int64},redFac::Float64; seed::Int64=1234,densityFac::Float64=1/pi,kargs...)
  M,N = shape
  samples = collect(1:M*N)
  chosenSamples = []
  forbiddenSamples =[]

  # set number of samples and mean distance
  numSamples = floor(Int64,M*N/redFac)
  rmean = densityFac*M*N/numSamples

  # srand(seed)
  while length(chosenSamples)<numSamples
    # check if samples are available. Otherwise reduce rmean and recalculate samples
    if isempty(samples)
      rmean  = rmean/2.0
      forbiddenSamples = updateForbiddenSamples(M,N,chosenSamples,rmean)
      samples = setdiff(collect(1:M*N),forbiddenSamples)
      continue
    end

    # pick a sample from the available samples and add them to chosen samples
    chosen = samples[ rand(1:length(samples)) ]
    push!(chosenSamples,chosen)

    # calculate forbidden radius and points within
    fSphere = forbiddenSphere(M,N,chosen,rmean)
    forbiddenSamples = union(forbiddenSamples,fSphere)

    # update available samples
    samples = setdiff(samples,forbiddenSamples)
  end

  return chosenSamples
end

function rmin(kx::Float64, ky::Float64, ai::Float64, a0::Float64, rmean::Float64)
  m = rmean/(ai-3*a0)
  c = rmean*(ai-5*a0)/(2*(ai-3*a0))
  return m*sqrt(kx^2+ky^2)+c
end

function updateForbiddenSamples(M,N,samples,rmean)
  forbiddenSamples = []
  for i=1:length(samples)
    fSphere = forbiddenSphere(M,N,samples[i],rmean)
    forbiddenSamples = union(forbiddenSamples,fSphere)
  end
  return forbiddenSamples
end

function forbiddenSphere(M::Int64,N::Int64,chosen::Int64,rmean::Float64)
  x,y = Tuple(CartesianIndices((M,N))[chosen])
  kx = x-(M+1.)/2.0
  ky = y-(N+1.)/2.0
  rad_min = rmin(kx, ky, 0.5, 0.0, rmean)
  fsphere = []
  for xi = floor(Int64,x-M*rad_min):ceil(Int64,x+M*rad_min)
    # xi must be positive and smaller then M
    (xi<1) && continue
    (xi>M) && continue

    for yi =  floor(Int64,y-N*rad_min):ceil(Int64,y+N*rad_min)
      # yi must be positive and smaller then N
      (yi<1) && continue
      (yi>N) && continue

      if (xi-x)^2+(yi-y)^2 < rad_min^2
        idx = LinearIndices((M,N))[xi,yi]
        push!(fsphere,idx)
      end
    end
  end
  return fsphere
end

function generateMask(M,N,pat)
  mask = zeros(Int64,M,N)
  for i=1:length(pat)
    mask[pat[i]] = 1.0
  end
  return mask
end


using DeterminantalPointProcesses
using Test
using Random, LinearAlgebra, IterTools
using KernelFunctions
import Combinatorics: combinations

rng = MersenneTwister(42)

n = 5
k = 2
A = rand(rng, n, n)
L = Symmetric(A * A')
kernel = SqExponentialKernel()
X = [rand(2) for _ in 1:3]
dpp = DPP(L)

@testset "Ensure correct pdf and logpdf computation" begin
    @testset "Kernel API" begin
        K = kernelmatrix(kernel, X)
        @test DPP(K).L ≈ DPP(kernel, X).L
        @test DPP(K).L ≈ DPP(kernel, reduce(hcat, X); obsdim=2).L
    end
    @testset "For DPP" begin
        # compute the true distribution
        true_pdf = Float64[]
        true_logpdf = Float64[]
        for z in subsets(1:n)
            push!(true_pdf, pdf(dpp, z))
            push!(true_logpdf, logpdf(dpp, z))
        end
        @test sum(true_pdf) ≈ 1.0
        @test all(true_pdf .<= 1.0)
        @test all(true_logpdf .<= 0.0)
    end

    @testset "For k-DPP" begin
        # compute the true distribution
        true_pdf = Float64[]
        true_logpdf = Float64[]
        for z in combinations(1:n, k)
            push!(true_pdf, pdf(dpp(k), z))
            push!(true_logpdf, logpdf(dpp(k), z))
        end
        @test sum(true_pdf) ≈ 1.0
        @test all(true_pdf .<= 1.0)
        @test all(true_logpdf .<= 0.0)
    end
end

@testset "Ensure correct sampling from DPP" begin
    # compute the true distribution
    all_subsets = []
    true_pdf = Float64[]
    for (i, z) in enumerate(subsets(1:n))
        push!(true_pdf, pdf(dpp, z))
        push!(all_subsets, (z, i))
    end
    global subset_to_idx = Dict(all_subsets)

    @testset "Exact sampling" begin
        nb_samples = 1000
        samples = rand(dpp, nb_samples)

        # compute the empirical distribution
        empirical_pdf = zeros(Float64, length(true_pdf))
        for z in samples
            empirical_pdf[subset_to_idx[z]] += 1
        end
        empirical_pdf ./= nb_samples

        # ensure that the empirical pdf is close to the true pdf
        total_variation = maximum(abs.(true_pdf - empirical_pdf))
        @test total_variation ≈ 0.0 atol = 1e-1
    end

    @testset "MCMC sampling" begin
        nb_samples = 1000
        samples, state = randmcmc(dpp, nb_samples; return_final_state=true)

        # ensure that L_z_inv makes sense (i.e., noise did not accumulate)
        z, L_z_inv = state
        @test sum(L_z_inv[z, z] * dpp.L[z, z] - I) ≈ 0 atol=1e-10

        # compute the empirical distribution
        empirical_pdf = zeros(Float64, length(true_pdf))
        for z in samples
            empirical_pdf[subset_to_idx[z]] += 1
        end
        empirical_pdf ./= nb_samples

        # ensure that the empirical pdf is close to the true pdf
        total_variation = maximum(abs.(true_pdf - empirical_pdf))
        @test total_variation ≈ 0.0 atol = 1e-1
    end
end

@testset "Ensure correct sampling from k-DPP" begin
    # compute the true distribution
    all_k_subsets = []
    true_pdf = Float64[]
    for (i, z) in enumerate(combinations(1:n, k))
        push!(true_pdf, pdf(dpp, z))
        push!(all_k_subsets, (z, i))
    end
    true_pdf ./= sum(true_pdf)
    k_subset_to_idx = Dict(all_k_subsets)

    @testset "Exact sampling" begin
        nb_samples = 10000
        samples = rand(dpp(k), nb_samples)

        # ensure that samples are of proper cardinality
        @test all(map(length, samples) .== k)

        # compute the empirical distribution
        empirical_pdf = zeros(Float64, length(true_pdf))
        for z in samples
            empirical_pdf[k_subset_to_idx[z]] += 1
        end
        empirical_pdf ./= nb_samples

        # ensure that the empirical pdf is close to the true pdf
        total_variation = maximum(abs.(true_pdf - empirical_pdf))
        @test total_variation ≈ 0.0 atol = 1e-1
    end

    @testset "MCMC sampling" begin
        nb_samples = 1000
        samples, state = randmcmc(dpp(k), nb_samples; return_final_state=true)

        # ensure that L_z_inv makes sense (i.e., noise did not accumulate)
        z, L_z_inv = state
        @test sum(L_z_inv[z, z] * dpp.L[z, z] - I) ≈ 0 atol = 1e-1

        # ensure that samples are of proper cardinality
        @test all(map(length, samples) .== k)

        # compute the empirical distribution
        empirical_pdf = zeros(Float64, length(true_pdf))
        for z in samples
            empirical_pdf[k_subset_to_idx[z]] += 1
        end
        empirical_pdf ./= nb_samples

        # ensure that the empirical pdf is close to the true pdf
        total_variation = maximum(abs.(true_pdf - empirical_pdf))
        @test_broken total_variation ≈ 0.0
        atol = 1e-1
    end
end


function besselroots(nu::Float64, n::Integer)
#BESSELROOTS    The first N roots of the function J_v(x)

# DEVELOPERS NOTES:
#   V = 0 --> Full double precision for N <= 20 (Wolfram Alpha), and very
#     accurate approximations for N > 20 (McMahon's expansion)
#   -1 <= V <= 5 : V ~= 0 --> 12 decimal figures for the 6 first zeros
#     (Piessens's Chebyshev series approximations), and very accurate
#     approximations for the others (McMahon's expansion)
#   V > 5 --> moderately accurate for the 6 first zeros and good
#     approximations for the others (McMahon's expansion)

# This code was originally written by L. L. Peixoto in MATLAB.
# Later modified by A. Townsend to work in Julia

    if n < 0
        error("Input N must be a positive integer")
    end

    x = zeros(n)
    if n > 0 && nu == 0
        for k in 1:min(n,20)
            x[k] = J0_roots[k]
        end
        for k in min(n,20)+1:n
            x[k] = McMahon(nu, k)
        end
    elseif n>0 && nu >= -1 && nu <= 5
        correctFirstFew = Piessens( nu )
        for k in 1:min(n,6)
            x[k] = correctFirstFew[k]
        end
        for k in min(n,6)+1:n
            x[k] = McMahon(nu, k)
        end
    elseif nu > 5
        for k in 1:n
            x[k] = McMahon(nu, k)
        end
    end
    x
end

function McMahon(nu::Float64, k::Integer)
    # McMahon's expansion. This expansion gives very accurate approximation
    # for the sth zero (s >= 7) in the whole region V >=- 1, and moderate
    # approximation in other cases.
    mu = 4nu^2
    a1 = 1 / 8
    a3 = (7mu-31) / 384
    a5 = 4*(3779+mu*(-982+83mu)) / 61440 # Evaluate via Horner's method.
    a7 = 6*(-6277237+mu*(1585743+mu*(-153855+6949mu))) / 20643840
    a9 = 144*(2092163573+mu*(-512062548+mu*(48010494+mu*(-2479316+70197mu)))) / 11890851840
    a11 = 720 *(-8249725736393+mu*(1982611456181+mu*(-179289628602+mu*(8903961290 +
          mu*(-287149133 + 5592657mu))))) / 10463949619200
    a13 = 576 *(423748443625564327 + mu*(-100847472093088506 + mu*(8929489333108377 +
        mu*(-426353946885548+mu*(13172003634537+mu*(-291245357370 + 4148944183mu)))))) / 13059009124761600
    b = 0.25 * (2nu+4k-1)*pi
    # Evaluate using Horner's scheme:
    x = b - (mu-1)*( ((((((a13/b^2 + a11)/b^2 + a9)/b^2 + a7)/b^2 + a5)/b^2 + a3)/b^2 + a1)/b)
end

const J0_roots =
    [   2.4048255576957728
        5.5200781102863106
        8.6537279129110122
        11.791534439014281
        14.930917708487785
        18.071063967910922
        21.211636629879258
        24.352471530749302
        27.493479132040254
        30.634606468431975
        33.775820213573568
        36.917098353664044
        40.058425764628239
        43.199791713176730
        46.341188371661814
        49.482609897397817
        52.624051841114996
        55.765510755019979
        58.906983926080942
        62.048469190227170  ]

const Piessens_C = [
       2.883975316228  8.263194332307 11.493871452173 14.689036505931 17.866882871378 21.034784308088
       0.767665211539  4.209200330779  4.317988625384  4.387437455306  4.435717974422  4.471319438161
      -0.086538804759 -0.164644722483 -0.130667664397 -0.109469595763 -0.094492317231 -0.083234240394
       0.020433979038  0.039764618826  0.023009510531  0.015359574754  0.011070071951  0.008388073020
      -0.006103761347 -0.011799527177 -0.004987164201 -0.002655024938 -0.001598668225 -0.001042443435
       0.002046841322  0.003893555229  0.001204453026  0.000511852711  0.000257620149  0.000144611721
      -0.000734476579 -0.001369989689 -0.000310786051 -0.000105522473 -0.000044416219 -0.000021469973
       0.000275336751  0.000503054700  0.000083834770  0.000022761626  0.000008016197  0.000003337753
      -0.000106375704 -0.000190381770 -0.000023343325 -0.000005071979 -0.000001495224 -0.000000536428
       0.000042003336  0.000073681222  0.000006655551  0.000001158094  0.000000285903  0.000000088402
      -0.000016858623 -0.000029010830 -0.000001932603 -0.000000269480 -0.000000055734 -0.000000014856
       0.000006852440  0.000011579131  0.000000569367  0.000000063657  0.000000011033  0.000000002536
      -0.000002813300 -0.000004672877 -0.000000169722 -0.000000015222 -0.000000002212 -0.000000000438
       0.000001164419  0.000001903082  0.000000051084  0.000000003677  0.000000000448  0.000000000077
      -0.000000485189 -0.000000781030 -0.000000015501 -0.000000000896 -0.000000000092 -0.000000000014
       0.000000203309  0.000000322648  0.000000004736  0.000000000220  0.000000000019  0.000000000002
      -0.000000085602 -0.000000134047 -0.000000001456 -0.000000000054 -0.000000000004               0.
       0.000000036192  0.000000055969  0.000000000450  0.000000000013               0.               0.
      -0.000000015357 -0.000000023472 -0.000000000140 -0.000000000003               0.               0.
       0.000000006537  0.000000009882  0.000000000043  0.000000000001               0.               0.
      -0.000000002791 -0.000000004175 -0.000000000014               0.               0.               0.
       0.000000001194  0.000000001770  0.000000000004               0.               0.               0.
      -0.000000000512 -0.000000000752              0.               0.              0.               0.
       0.000000000220  0.000000000321               0.               0.               0.               0.
      -0.000000000095 -0.000000000137               0.               0.               0.               0.
       0.000000000041  0.000000000059               0.               0.               0.               0.
      -0.000000000018 -0.000000000025               0.               0.               0.               0.
       0.000000000008  0.000000000011               0.               0.               0.               0.
      -0.000000000003 -0.000000000005               0.               0.               0.               0.
       0.000000000001  0.000000000002               0.               0.               0.               0.]

function Piessens( nu::Float64 )
    # Piessens's Chebyshev series approximations (1984). Calculates the 6 first
    # zeros to at least 12 decimal figures in region -1 <= V <= 5:
    C = Piessens_C
    T = Array{Float64}(undef,size(C,1))
    pt = (nu-2)/3
    T[1], T[2] = 1., pt
    for k = 2:size(C,1)-1
        T[k+1] = 2*pt*T[k] - T[k-1]
    end
    y = C'*T
    y[1] *= sqrt(nu+1)                  # Scale the first root.
    y
end


## UC Model
function UC(Ω, ρ, RES; printing = 0)
    mUC = Model(() -> Gurobi.Optimizer(GUROBI_ENV))
    set_optimizer_attribute(mUC, "OutputFlag", printing)

        ### Variables
        # 1st stage
        @variables(mUC, begin
            v[G,T], (Bin) # Unit state
            0.0 ≤ c_su[G,T] # start-up cost
            0.0 ≤ c_rc[G,T] # reserve capacity cost
            0.0 ≤ obj1
            0. ≤ p_exp[g in G,T] ≤ GenData.Pmax[g] # Expected generation
            0. ≤ r_up[g in G,T] ≤ GenData.Rup[g] # Upward reserve capacity
            0. ≤ r_dw[g in G,T] ≤ GenData.Rdw[g] # Downward reserve capacity
            end)

            # 2nd stage
            @variables(mUC, begin
                0. ≤ p[g in G,Ω,T] ≤ GenData.Pmax[g] # Real-time generation
                0. ≤ r_p[g in G,Ω,T] #≤ GenData.Rplus[g] # Real-time up reserve
                0. ≤ r_m[g in G,Ω,T] #≤ GenData.Rminus[g] # Real-time down reserve
                0. ≤ l_s[t in T, Ω] ≤ sum(Pd,dims=1)[t] # load shedding
                0. ≤ w_s[t in T, ω in Ω] ≤ RES[t,ω] # wind spillage
                0. ≤ Θ[Ω] # operation costs

            end)

        ### Constraints
        #### 1st stage
            # Power limits
            @constraints(mUC, begin
                PMax[g in G, t in T], p_exp[g,t] + r_up[g,t]≤ v[g,t]*GenData.Pmax[g]
                PMin[g in G, t in T], p_exp[g,t] - r_dw[g,t]≥ v[g,t]*GenData.Pmin[g]
                RupLim[g in G, ω in Ω, t in T], r_up[g,t] ≤ v[g,t]*GenData.Rup[g]
                RdwLim[g in G, ω in Ω, t in T], r_dw[g,t] ≤ v[g,t]*GenData.Rdw[g]
                end)
            # Costs
            @constraint(mUC, obj1 == sum(GenData.Cg[g]*p_exp[g,t]
                                + c_su[g,t] + c_rc[g,t] for g in G, t in T))

            @constraints(mUC, begin
                CostReserveCap[g in G, t in T],
                        c_rc[g,t] == r_up[g,t]*GenData.Cup[g] + r_dw[g,t]*GenData.Cdw[g]
                #
                CostSU[g in G, t in T],
                        c_su[g,t] ≥ GenData.Cstart[g]*(v[g,t] - v[g,T[t-1]])
                            end)

            # Minimum Online & Offline Times
            @constraints(mUC, begin
                MinOn[g in G, t in T],
                    sum(v[g,n] for n in t:min(t+GenData.UT[g]-1,24)) ≥ GenData.UT[g]*(v[g,t]-v[g,T[t-1]])

                MinOff[g in G, t in T],
                    sum(1-v[g,n] for n in t:min(t+GenData.DT[g]-1,24)) ≥ GenData.DT[g]*(v[g,T[t-1]]-v[g,t])
                end)

            #### 2nd stage
            # Power balance
            @constraint(mUC, PowerBalance_Second[ω in Ω, t in T],
                                sum(Pd,dims=1)[t] - l_s[t,ω]== RES[t,ω] + sum(p[:,ω,t]) - w_s[t,ω])

            # Power limits
            @constraints(mUC, begin
                PMax2nd[g in G, ω in Ω, t in T], p[g,ω,t] ≤ v[g,t]*GenData.Pmax[g]
                PMin2nd[g in G, ω in Ω, t in T], p[g,ω,t] ≥ v[g,t]*GenData.Pmin[g]
                RplusLim[g in G, ω in Ω, t in T], r_p[g,ω,t] ≤ r_up[g,t]
                RminusLim[g in G, ω in Ω, t in T], r_m[g,ω,t] ≤ r_dw[g,t]
                Ramps_2nd[g in G, ω in Ω, t in T], -GenData.Rdw[g] ≤ p[g,ω,t] - p[g,ω,T[t-1]] ≤ GenData.Rup[g]

                end)

            # Generated Power
            @constraint(mUC, PGen_2nd[g in G, ω in Ω, t in T],
                    p[g,ω,t] == p_exp[g,t] + r_p[g,ω,t] - r_m[g,ω,t])

            # Total Cost
            @constraint(mUC, Cost_2nd[ω in Ω],
                    Θ[ω] == sum(C_ls*l_s[t,ω] + C_ws*w_s[t,ω] + sum(GenData.Cplus[g]*r_p[g,ω,t] - GenData.Cminus[g]*r_m[g,ω,t] for g in G) for t in T)
                    )

        ### Objective
        @objective(mUC, Min, obj1 + sum(ρ[ω]*Θ[ω] for ω in Ω))

        optimize!(mUC)
        # println("Gap: $(MOI.get(mUC, MOI.RelativeGap()))")

    V = JuMP.value.(v)
        V= collect(V[i,j] for i in G, j in T)
            Pexp = (JuMP.value.(p_exp))
            Pexp= collect(Pexp[i,j] for i in G, j in T)
            Rup = (JuMP.value.(r_up))
            Rup= collect(Rup[i,j] for i in G, j in T)
            Rdw = (JuMP.value.(r_dw))
            Rdw= collect(Rdw[i,j] for i in G, j in T)
            Obj1 = JuMP.value.(obj1)
            LB = JuMP.objective_value.(mUC)

            # println("l_s:")
            # println(C_ls*sum(JuMP.value.(l_s)))
            # println("w_s")
            # println(C_ws*sum(JuMP.value.(w_s)))
            # println("Reserves:")
            # println(sum(sum(GenData.Cplus[g]*JuMP.value.(r_p[g,ω,t]) - GenData.Cminus[g]*JuMP.value.(r_m[g,ω,t]) for g in G) for t in T, ω in Ω))

    return V, Pexp, Rup, Rdw, Obj1, LB, mUC

    end


export Platte

struct Platte{T<:AbstractFloat, Ti<:Integer}
    xL::T
    yL::T
    nx::Ti
    ny::Ti
    rho::Matrix{T}
    lambda::Matrix{T}
    cp::Matrix{T}
    temp::Matrix{T}
end

function Platte(xL::T, yL::T, nx::Ti, ny::Ti, dens, heat, cap, tv::Union{Number,Function}) where {T,Ti}
    if nx <= 0
        throw(ArgumentError("nx muss > 0 sein, ist aber $nx"))
    end

        function setgrid(f::Number)
            return [f for ix in 1:nx, iy in 1:ny]
        end

        function setgrid(f::Function)
            return [f((ix - 0.5) * xL / nx, (iy - 0.5) * xL / ny) for ix in 1:nx, iy in 1:ny]
        end

    rho = setgrid(dens)
    temp = setgrid(tv)
    lambda= setgrid(heat)
    cp = setgrid(cap)

    return Platte{T, Ti}(xL, yL, nx, ny, rho, lambda, cp, temp)
end

function tempiter(p::Platte, deltat::T) where {T}

    dtemp = [heattransfer(ix, iy, p) for ix in 1:p.nx, iy in 1:p.ny]
    tempn = dtemp * deltat + p.temp

    return Platte(p.xL, p.yL, p.nx, p.ny, p.rho, p.lambda, p.cp, tempn)
end

function heattransfer(ix, iy, p::Platte)

    produkt = [p.lambda[ix, iy] * nabla(p.temp, p)[ix, iy] for ix in 1:p.nx, iy in 1:p.ny]

    tempdiff = nabla(produkt, p)[ix, iy] / (p.cp[ix, iy] * p.rho[ix, iy])

end

function nabla(vm::Matrix{<:Vector}, p::Platte)
    vmx = [vm[ix, iy][1] for ix in 1:p.nx, iy in 1:p.ny]
    vmy = [vm[ix, iy][2] for ix in 1:p.nx, iy in 1:p.ny]
    [nablax(vmx, ix, iy, p) + nablay(vmy, ix, iy, p) for ix in 1:p.nx, iy in 1:p.ny]
end

function nabla(m::Matrix{<:Number}, p::Platte)
    [[nablax(m, ix, iy, p), nablay(m, ix, iy, p)] for ix in 1:p.nx, iy in 1:p.ny]
end

function nablax(m::Matrix, ix, iy, p::Platte)
    dx = p.xL / p.nx
    ix == 1 ? (m[ix + 1, iy] - m[ix, iy]) / dx :
    ix == p.nx ? (m[ix, iy] - m[ix - 1, iy]) / dx :
    (m[ix + 1, iy] - m[ix - 1, iy]) / (2 * dx)
end

function nablay(m::Matrix, ix, iy, p::Platte)
    dy = p.yL / p.ny
    iy == 1 ? (m[ix, iy + 1] - m[ix, iy]) / dy :
    iy == p.ny ? (m[ix, iy] - m[ix, iy - 1]) / dy :
    (m[ix, iy + 1] - m[ix, iy - 1]) / (2 * dy)
end


cd(@__DIR__)
using Pkg
Pkg.activate(@__DIR__); Pkg.instantiate()

include("prinz_neuron_setup.jl")
include("helpers.jl")

system = neuron.sys
@named simulation = ODESystem([D(system.Isyn) ~ 0]; systems = [system])
simplified = structural_simplify(simulation)

sim_exp = ModelingToolkit.build_torn_function(simplified; expression=true)
cleaned = clean_expr!(sim_exp, simplified)

clipboard(format_text(string(prettify(cleaned)))) # can paste readable output below

function prinz_neuron!(du, u, p, t)
    @inbounds let Caᵢ         = u[1],
                  Vₘ          = u[2],
                  NaV₊m       = u[3],
                  NaV₊h       = u[4],
                  CaT₊m       = u[5],
                  CaT₊h       = u[6],
                  CaS₊m       = u[7],
                  CaS₊h       = u[8],
                  KA₊m        = u[9],
                  KA₊h        = u[10],
                  KCa₊m       = u[11],
                  Kdr₊m       = u[12],
                  H₊m         = u[13],

                  cₘ          = p[1],
                  aₘ          = p[2],
                  τCa         = p[3],
                  Ca∞         = p[4],
                  f           = p[5],
                  ENa         = p[6],
                  EK          = p[7],
                  EK          = p[8],
                  EK          = p[9],
                  EH          = p[10],
                  El          = p[11],
                  NaV₊gbar    = p[12],
                  CaT₊gbar    = p[13],
                  CaS₊gbar    = p[14],
                  KA₊gbar     = p[15],
                  KCa₊gbar    = p[16],
                  Kdr₊gbar    = p[17],
                  H₊gbar      = p[18],
                  leak₊g      = p[19]

        let
            du[1] =
                (
                    Ca∞ +
                    -1Caᵢ +
                    -1000 *
                    f *
                    (
                        CaS₊gbar *
                        aₘ *
                        CaS₊h *
                        CaS₊m^3.0 *
                        (
                            Vₘ +
                            12.199306599503455 *
                            log(0.0003333333333333333 * max(Caᵢ, 0.001))
                        ) +
                        CaT₊gbar *
                        aₘ *
                        CaT₊h *
                        CaT₊m^3.0 *
                        (
                            Vₘ +
                            12.199306599503455 *
                            log(0.0003333333333333333 * max(Caᵢ, 0.001))
                        )
                    )
                ) * inv(τCa)
            du[2] =
                (
                    -1 * aₘ * leak₊g * (Vₘ + -1El) +
                    -1 *
                    H₊gbar *
                    aₘ *
                    H₊m *
                    (Vₘ + -1EH) +
                    -1 *
                    KCa₊gbar *
                    aₘ *
                    KCa₊m^4.0 *
                    (Vₘ + -1EK) +
                    -1 *
                    Kdr₊gbar *
                    aₘ *
                    Kdr₊m^4.0 *
                    (Vₘ + -1EK) +
                    -1 *
                    CaS₊gbar *
                    aₘ *
                    CaS₊h *
                    CaS₊m^3.0 *
                    (
                        Vₘ +
                        12.199306599503455 *
                        log(0.0003333333333333333 * max(Caᵢ, 0.001))
                    ) +
                    -1 *
                    CaT₊gbar *
                    aₘ *
                    CaT₊h *
                    CaT₊m^3.0 *
                    (
                        Vₘ +
                        12.199306599503455 *
                        log(0.0003333333333333333 * max(Caᵢ, 0.001))
                    ) +
                    -1 *
                    KA₊gbar *
                    aₘ *
                    KA₊h *
                    KA₊m^3.0 *
                    (Vₘ + -1EK) +
                    -1 *
                    NaV₊gbar *
                    aₘ *
                    NaV₊h *
                    NaV₊m^3.0 *
                    (Vₘ + -1ENa)
                ) *
                inv(aₘ) *
                inv(cₘ)
            du[3] =
                (
                    -1NaV₊m +
                    inv(1.0 + exp(-4.8204158790170135 + -0.1890359168241966Vₘ))
                ) * inv(2.64 + -2.52 * inv(1 + exp(-4.8 + -0.04Vₘ)))
            du[4] =
                0.7462686567164178 *
                (1.0 + exp(-6.29 + -0.1Vₘ)) *
                (
                    -1NaV₊h +
                    inv(1.0 + exp(9.44015444015444 + 0.19305019305019305Vₘ))
                ) *
                inv(1.5 + inv(1.0 + exp(9.694444444444445 + 0.2777777777777778Vₘ)))
            du[5] =
                (
                    -1CaT₊m +
                    inv(1.0 + exp(-3.7638888888888893 + -0.1388888888888889Vₘ))
                ) * inv(
                    43.4 +
                    -42.6 *
                    inv(1.0 + exp(-3.321951219512195 + -0.04878048780487805Vₘ)),
                )
            du[6] =
                (
                    -1CaT₊h +
                    inv(1.0 + exp(5.836363636363637 + 0.18181818181818182Vₘ))
                ) * inv(
                    210.0 +
                    -179.6 *
                    inv(1.0 + exp(-3.2544378698224854 + -0.0591715976331361Vₘ)),
                )
            du[7] =
                (
                    -1CaS₊m +
                    inv(1.0 + exp(-4.074074074074074 + -0.1234567901234568Vₘ))
                ) * inv(
                    2.8 +
                    14.0 * inv(
                        exp(-5.384615384615385 + -0.07692307692307693Vₘ) +
                        exp(2.7 + 0.1Vₘ),
                    ),
                )
            du[8] =
                (
                    -1CaS₊h +
                    inv(1.0 + exp(9.67741935483871 + 0.16129032258064516Vₘ))
                ) * inv(
                    120.0 +
                    300.0 * inv(
                        exp(-4.0625 + -0.0625Vₘ) +
                        exp(6.111111111111111 + 0.1111111111111111Vₘ),
                    ),
                )
            du[9] =
                (
                    -1KA₊m +
                    inv(1.0 + exp(-3.1264367816091956 + -0.1149425287356322Vₘ))
                ) * inv(
                    23.2 +
                    -20.8 *
                    inv(1.0 + exp(-2.164473684210526 + -0.06578947368421052Vₘ)),
                )
            du[10] =
                (
                    -1KA₊h +
                    inv(1.0 + exp(11.612244897959181 + 0.2040816326530612Vₘ))
                ) * inv(
                    77.2 +
                    -58.4 *
                    inv(1.0 + exp(-1.4679245283018867 + -0.03773584905660377Vₘ)),
                )
            du[11] =
                (
                    -1KCa₊m +
                    Caᵢ *
                    inv(3.0 + Caᵢ) *
                    inv(1.0 + exp(-2.246031746031746 + -0.07936507936507936Vₘ))
                ) * inv(
                    180.6 +
                    -150.2 *
                    inv(1.0 + exp(-2.026431718061674 + -0.04405286343612335Vₘ)),
                )
            du[12] =
                (
                    -1Kdr₊m +
                    inv(1.0 + exp(-1.0423728813559323 + -0.0847457627118644Vₘ))
                ) * inv(
                    14.4 +
                    -12.8 *
                    inv(1.0 + exp(-1.4739583333333335 + -0.052083333333333336Vₘ)),
                )
            du[13] =
                1 // 2 *
                (
                    exp(-1.8671328671328669 + 0.06993006993006992Vₘ) +
                    exp(-14.629310344827585 + -0.08620689655172414Vₘ)
                ) *
                (
                    -1H₊m +
                    inv(1.0 + exp(13.636363636363637 + 0.18181818181818182Vₘ))
                )
            nothing
        end
    end
end


module Constants

    export cee,
           emgm,
           pmgm,
           kboltz,
           kboltz_kev2erg,
           gr_to_kev,
           me_kev,
           emerg,
           charg,
           sigtom,
           herg,
           hkev,
           mjy,
           re0,
           gconst,
           sbconst,
           aconst,
           kpc,
           msun

    using GSL
    const GSL_CONST = GSL.Const

    const cee  =           GSL_CONST.CGSM_SPEED_OF_LIGHT
    const emgm=            GSL_CONST.CGSM_MASS_ELECTRON
    const pmgm  =          GSL_CONST.CGSM_MASS_PROTON
    const kboltz    =      GSL_CONST.CGSM_BOLTZMANN
    const kboltz_kev2erg = 1.6022e-9
    const gr_to_kev   =    5.6095883571872e+29
    const me_kev    =      511.
    const emerg   =        (GSL_CONST.CGSM_MASS_ELECTRON*GSL_CONST.CGSM_SPEED_OF_LIGHT^2.0)
    const charg   =        4.8e-10
    const sigtom =         GSL_CONST.CGSM_THOMSON_CROSS_SECTION
    const herg   =         GSL_CONST.CGSM_PLANCKS_CONSTANT_H
    const hkev   =         (GSL_CONST.CGSM_PLANCKS_CONSTANT_H*6.2415e8)
    const mjy   =          1.e-26
    const re0    =         2.81794e-13
    const gconst  =        GSL_CONST.CGSM_GRAVITATIONAL_CONSTANT
    const sbconst  =       GSL_CONST.CGSM_STEFAN_BOLTZMANN_CONSTANT
    const aconst   =       7.56e-15
    const kpc =             1e3*GSL_CONST.CGSM_PARSEC
    const msun =            GSL_CONST.CGSM_SOLAR_MASS

end


using CAGD
using SparseArrays
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

V = [
	0.0 2.0 4.0 1.0 2.0 3.0 1.0 3.0 0.0 4.0
	4.0 4.0 4.0 3.0 3.0 3.0 1.0 1.0 0.0 0.0
];
EV = [
	[ 1, 2], [ 2, 3], [ 1, 9], [ 2, 5],
	[ 3,10], [ 4, 5], [ 5, 6], [ 4, 7],
	[ 6, 8], [ 7, 8], [ 9,10]
];
FV = [
    [1, 2, 2, 3, 4, 5, 5, 6, 7, 8, 9, 10],
    [            4, 5,    6, 7, 8       ],
    [1, 2,    3,                   9, 10]
]
FE = [
	[1, 2, 3, 4,  4, 5, 6, 7,  8, 9, 10, 11],
	[                   6, 7,  8, 9, 10    ],
	[1, 2, 3,        5,                  11]
]

copEV = SparseArrays.sparse(Int8.([
	1  1  0  0  0  0  0  0  0  0
	0  1  1  0  0  0  0  0  0  0
	1  0  0  0  0  0  0  0  1  0
	0  1  0  0  1  0  0  0  0  0
	0  0  1  0  0  0  0  0  0  1
	0  0  0  1  1  0  0  0  0  0
	0  0  0  0  1  1  0  0  0  0
	0  0  0  1  0  0  1  0  0  0
	0  0  0  0  0  1  0  1  0  0
	0  0  0  0  0  0  1  1  0  0
	0  0  0  0  0  0  0  0  1  1
]))

copFE = SparseArrays.sparse(Int8.([
	1  1  1  2  1  1  1  1  1  1  1
	0  0  0  0  0  1  1  1  1  1  0
	1  1  1  0  1  0  0  0  0  0  1
]))

copFV = SparseArrays.sparse(Int8.([
    1  2  1  1  2  1  1  1  1  1
    0  0  0  1  1  1  1  1  0  0
    1  1  1  0  0  0  0  0  1  1
]))

#--- Using CAGD
model = CAGD.Model(V)
CAGD.addModelCells!(model, 1, EV, signed = false)
CAGD.addModelCells!(model, 2, FE, signed = false)
copEV = model.T[1]
copFE = model.T[2]
#--- Otherwise

copFV = map(x -> floor(Int8, x / 2), copFE * copEV)

copFE = map(x -> floor(Int8, x / 2), copFV * copEV')

# EXAMPLE 2
#

V = [
    1.0 3.0  0.0 1.0 3.0 4.0  2.0  2.0  1.0 3.0
    4.0 4.0  3.0 3.0 3.0 3.0  2.0  1.0  0.0 0.0
]

copEV = SparseArrays.sparse(Int8.([
    1 0 1 0 0  0 0 0 0 0
    1 0 0 1 0  0 0 0 0 0
    0 1 0 0 1  0 0 0 0 0
    0 1 0 0 0  1 0 0 0 0
    0 0 1 1 0  0 0 0 0 0
    0 0 0 0 1  1 0 0 0 0
    0 0 0 1 0  0 1 0 0 0
    0 0 0 0 1  0 1 0 0 0
    0 0 0 0 0  0 1 1 0 0
    0 0 0 0 0  0 0 1 1 0
    0 0 0 0 0  0 0 1 0 1
    0 0 0 0 0  0 0 0 1 1
]))

copFE = SparseArrays.sparse(Int8.([
    1 1 0 0 1  0 0 0 0 0  0 0
    0 0 1 1 0  1 0 0 0 0  0 0
    0 0 0 0 0  0 0 0 0 1  1 1
    1 1 1 1 1  1 2 2 2 1  1 1
]))


s = open("input.txt") do file
    readlines(file);
end

function ParseLine(line)
    m = match(r"mem\[(\d+)\] = (\d+)", line);
    if m != nothing
        return map(x->parse(UInt64,x), m.captures);
    end

    m = match(r"mask = ([01X]+)", line);
    if m != nothing
        return (m.captures|>collect);
    end

    return ();
end

function SumMem(program)
    mem = Dict();
    oneMask = 0;
    zeroMask = -1;
    for line in program
        parsed = ParseLine(line);
        if length(parsed) == 0
            println(line);
            continue;
        end

        if length(parsed) == 1
            # new mask
            oneMask = parse(UInt64, String(map(x->ifelse(x=='1', '1', '0'), parsed[1])), base=2);
            zeroMask = parse(UInt64, String(map(x->ifelse(x=='0', '0', '1'), parsed[1])), base=2);
        else
            # put in mem
            val = parsed[2];
            val |= oneMask;
            val &= zeroMask;

            mem[parsed[1]] = val;
        end
    end

    return sum(map(UInt64,values(mem)));
end

println(Int(SumMem(s)))


using Distributions
import Random: rand!, rand
import Distributions: logpdf

struct MALA{T}
    ϵ::T
    MALA(ϵ::T) where T = new{T}(ϵ)
end

function rand!(mala::MALA, θ, ∇ll, ::UpdtIdx) where UpdtIdx
    id = [idx(UpdtIdx())...]
    θ[id] .+= mala.ϵ .* ∇ll .+ sqrt.(2.0.*mala.ϵ).*randn(length(∇ll))
    θ
end

function rand(mala::MALA, θ, ∇ll, ::UpdtIdx) where UpdtIdx
    θc = copy(θ)
    rand!(mala, θc, ∇ll, UpdtIdx())
end

function logpdf(mala::MALA, θ, θᵒ, ∇ll, ::UpdtIdx) where UpdtIdx
    id = [idx(UpdtIdx())...]
    logpdf(MvNormal(θ[id] .+ mala.ϵ.*∇ll, diagm(0=>sqrt.(2.0.*mala.ϵ))), θᵒ[id])
end


using Plots, LaTeXStrings
gr(size=(600,300),colorbar=false,legend=false,grid=false)

using FourierGPE

## initialize
L = (60.0,)
N = (512,)
sim = Sim(L,N)
@unpack_Sim sim

## simulation parameters
μ = 1.0
g = 0.01
γ = 0.0
tf = 20.0pi
Nt = 150
ti = 0.0
t = LinRange(ti,tf,Nt)

## Bogoliubov state
x = X[1]
k = K[1]

## Define Bog state
f(k) = 1 + 4/k^2
u(k) = 0.5*(sqrt(f(k))+1)/f(k)^(1/4)
v(k) = 0.5*(sqrt(f(k))-1)/f(k)^(1/4)
lam = 0.01
bog(x,k) = u(k)*exp(im*k*x) - conj(v(k))*exp(-im*k*x)
ψb(x,k) = sqrt(μ/g)*(complex(1) + lam*bog(x,k))

## choose an allowed k for periodic system
kb = k[20]
ψi = ψb.(x,kb)
ϕi = kspace(ψi,sim)

# Set time evolution and pack
# reltol = 1e-7
# alg = Vern7()

@pack_Sim! sim

## Evolve in k space
@time sol = runsim(sim)
y = abs2.(xspace(sol[end-1],sim)); plot(x,y)

## density movie
anim = @animate for i=1:Nt
    ψ = xspace(sol[i],sim)
    y = g*abs2.(ψ)
    plot(x,y,c=c3,fill=(0,0.5,c3))
    plot!(x,one.(x)*μ,c=:black,alpha=0.3)
    ylims!(0,1.1)
    xlims!(first(x),last(x))
    xlabel!(L"x/\xi"); ylabel!(L"gn(x)/\mu")
end

gif(anim, "./examples/bogoliubov_density.gif", fps = 20)


#Logistic map
T(x) = mod(3.92*x*(1-x), 1)

#Arnold's maps
#epsilon = 0.7
#tau = 1.4
#T(x) = mod(x-(epsilon/(2*pi))*sin(2*pi*x)+tau, 1)

#Tsuda's map (from Chaotic Itinerancy in Random Dynamical System Related to Associative Memory Models)
#C = 0.1
#A = 0.07
#T(x) = mod(x+C+A*sin(4*pi*x), 1)

P = NoiseFourier.Fourier1D.assemble_matrix(T, 128; x_0 = 0, x_1 = 1)
D = NoiseFourier.Fourier1D.noise_matrix(0.005, 128)
M = D*P
v2, vinf = NoiseFourier.Fourier1D.rigorous_norm(M)

print(v2)
print(vinf)

function convert_to_float(M, n)
    Q = zeros(Complex{Float64},n,n)
    ze = 0
    zer = []
    max_nnz_rows = 0
    a = 0
    b = 0
    for i in 1:n
        for j in 1:n
            a = Float64(real((M[i,j])))
            b = Float64(imag((M[i,j])))
            Q[i,j] = a + im*b
            if Q[i,j] == 0
                ze = ze +1
            end
        end
        push!(zer, ze)
        ze = 0
    end
    max_nnz_rows = maximum(zer)
    return Q#, max_nnz_rows
end

function idx_to_freq(j,k)
    # k is the size of the matrix
    if j>=0 && j <=k ÷ 2+1
        return j-1
    end
    if j >k ÷ 2+1
        return j-k-1
    end
end

function measure(x,v)
    t = real(v[1])
    for i in 2:128
        a = 2*real(v[i])
        b = 2*imag(v[i])
        k = idx_to_freq(i,128)
        t += a*cos(2*pi*x*k) - b*sin(2*pi*x*k)
    end
    return t
end

Z = convert_to_float(M, 128)
v = eigvecs(Z)[:,128]
p = 50000
x = range(0,1, length = p)
y = zeros(p)
for i in 1:p
    y[i] = measure(x[i],v)
end
plot(x,y,ylims=(0,5),fmt = :png)#,xticks=25:5:75)


using LinearSpace
using Test

struct MT i::Int end
#scalartype(mt::MT) = Float64

@linearspace MT #Float64

x, y = MT(2), MT(3)

z = x + 3y

@test z.a == [1, 3]
@test z.v == [x, y]

z = 1x
@test z.a == [1]
@test z.v == [x]

@test lctype(MT) == MT_LC

ts = MT[]
as = Float64[]
z = x + 3y
for (a,t) in terms(z)
    @show a, t
    push!(as,a)
    push!(ts,t)
end

@test as == [1.0, 3.0]
@test ts == [x,y]

ts = MT[]
as = Float64[]
z = x
for (a,t) in terms(z)
    @show a, t
    push!(as,a)
    push!(ts,t)
end

@test as == [1.0]
@test ts == [x]


"""
 function for computation of interpolation matrix from FEM mesh to regular grid
    from displacements in tetraedra fem grid to a regular grid mesh

  Inputs:
    - Nodes: matrix with coordinates of nodes of the FEMmesh
    - Conec: vector of vectors. at i-th entry has the vector with the connectivity indexes
    - numComponentes: number of components of the magnitude to be interpolated
    - order: 0 interpolates from elements centroids
            1 interpolates from nodes
"""

function computeFEM2GridInterpMatrix( Nodes, Conec, intGrid, numCompon, order )

  numVoxelsX   =       intGrid.voxelNums[1]
  numVoxelsXY  = prod( intGrid.voxelNums[1:2] )
  numVoxels    = prod( intGrid.voxelNums )

  numNodes  = size( Nodes, 1)
  numElems  = length( Conec )

  @assert typeof(Conec) == Vector{ Vector{Int64} }

  booleanInterFound = zeros( Bool, numVoxels )

  rowIndexes = Vector{Int}()
  colIndexes = Vector{Int}()
  nonZerVals = Vector{Float64}()

  if order == 1
    totalRows = numCompon * numVoxels
    totalCols = numCompon * numNodes
  end

  # loop in elements for see which grid nodes are in an element.
  for i in (1:numElems)

    nodesElem   = Conec[i]
    coordesElem = Nodes[ nodesElem , :]

    # coordenadas de bounding box de elemento
    mins = vec( minimum( coordesElem, dims=1 ) )
    maxs = vec( maximum( coordesElem, dims=1 ) )

    # indexes of grid within bounding box
    indsIni, indsEnd = ranges( mins, maxs, intGrid )

    if length( indsIni ) > 0

      for kg in (indsIni[3]:indsEnd[3])
        for jg in (indsIni[2]:indsEnd[2])
          for ig in (indsIni[1]:indsEnd[1])

            # number of voxel
            ind = (kg-1) * numVoxelsXY + (jg-1)*numVoxelsX + ig ;

            # if an interpolation of the node was not already obtained
            if !(booleanInterFound[ind])

              pointCand = intGrid.startVox + intGrid.voxelWidths .* ([ ig, jg, kg ] .- 1 )

              # verify if it belongs to the current element
              itIsInTetra, shapeFuncs = checkInTetra( coordesElem, pointCand )

              # if it is in the tetrahedron
              if itIsInTetra
                # then we found it
                booleanInterFound[ind] = true

                dofsGridInd = nodes2dofs( ind, numCompon )

                # if order == 0
                #    # assign interpolation of grid node ind to element i
                #    interpMatrix[ind, i] = 1

                if order == 1

                  for indShapeFunc in (1:length(shapeFuncs))

                    dofsNode = nodes2dofs( nodesElem[indShapeFunc], numCompon )

                    append!( rowIndexes, dofsGridInd )
                    append!( colIndexes, dofsNode    )
                    append!( nonZerVals, shapeFuncs[ indShapeFunc ]*ones(numCompon) )

                  end

                end # if order

              end # if it is in tetra
            end # if node already analyzed
          end # loop x index
        end # loop y index
      end # loop z index
    end # if length inds > 0

  end # loop elements

  return sparse( rowIndexes, colIndexes, nonZerVals, totalRows, totalCols )
end

"""
 funcion para calculo rangos de indices xyz de grid que efecitivamente caen dentro del box del elemento
"""

function ranges( mins, maxs, intGrid )

  # if the maximums of the box are higher than all components of the startVoxel
  # then there could be an intersection
  if ( sum( maxs .>= intGrid.startVox ) == 3 ) && ( sum( mins .<= intGrid.endVox ) == 3 )
    # there is intersection
    indsEnd = Int.( min.( ceil.(  ( maxs - intGrid.startVox ) ./ intGrid.voxelWidths .+ 1 ), intGrid.voxelNums ) )
    indsIni = Int.( max.( floor.( ( mins - intGrid.startVox ) ./ intGrid.voxelWidths .+ 1 ), ones(3)           ) )
  else
    indsIni = []
    indsEnd = []
  end

  return indsIni, indsEnd
end

"""
TO DO
"""

function checkInTetra( nodesTetra, pointCand )

  A          = zeros(4,4)
  A[:, 1]   .= 1.0
  A[:, 2:4] .= nodesTetra
  volTot     = det(A) / 6.0

  volsRel = zeros(4)

  boolInTetra = true
  i=0
  while boolInTetra && i<4
    i = i + 1

    B        = copy(A)
    B[i,2:4] = pointCand
    volsRel[i] = ( det( B ) / 6.0 ) / volTot ;

    boolInTetra = boolInTetra && ( volsRel[i] >= -1.0e-10 )
  end

  if boolInTetra
    shapeFuncs = volsRel
  else
    shapeFuncs = []
  end

  return boolInTetra, shapeFuncs
end

"""
TO DO
"""

function vecGridConvert(ind)

  return 1
end

"""

copied from
https://github.com/jorgepz/FEMAssembler.jl/blob/main/src/mesh/dofsComputations.jl

"""
function nodes2dofs( nodes , degreespernode )
    n    = length(nodes);
    dofs = zeros( Int, n*degreespernode ) ;
    for i in (1:n);
        dofs[ (i-1)*degreespernode .+ Vector(1:degreespernode) ] = (degreespernode*(nodes[i]-1)).+Vector(1:degreespernode)  ;
    end
    return dofs;
end