module aoc_05
    using Test
    export read_input, part1, part2
    export test_part1, test_part2, test
    export input, demo

    const input_r = r"([FB]{7})([LR]{3})"

    function read_input(filename = "input.txt")
        lines = readlines(filename)
        # for line in eachline(filename)
        # for line in lines
        # end
        [parse_line(line) for line in eachline(filename)]
        # lines
    end

    function parse_line(line)
        m = match( input_r , line )
        ( m.captures[1], m.captures[2] )
    end

    const input = read_input("input.txt")
    const demo = read_input("demo.txt")

    function binspace2int(string::AbstractString, high::Char, shift::Int)
        ba = Vector{Char}(string) .== high
        Int(bitreverse(UInt8(ba.chunks[1])) >> shift)
    end

    function seat2rc(seat::Tuple{AbstractString,AbstractString})
        r = binspace2int(seat[1], 'B', 1)
        c = binspace2int(seat[2], 'R', 5)
        (r,c)
    end

    function rc2seatid(seat::Tuple{Int,Int})
        seat[1]*8 + seat[2]
    end

    seatid(seat) = rc2seatid( seat2rc(seat) )

    function part1(input = input)
        maximum( seatid.( input ) )
    end

    function test_part1()
        @test seatid( demo[1] ) == 357
        @test seatid( demo[2] ) == 567
        @test seatid( demo[3] ) == 119
        @test seatid( demo[4] ) == 820
        @test part1( demo ) == 820
    end

    function part2(input = input)
        seatids = seatid.(input)
        e = extrema(seatids)
        seq = e[1]:e[2]
        (seq)[findfirst(.~in.(seq, (seatid.(input),) ))]
    end

    function test_part2()
        seatids = seatid.(input)
        @test part2()-1 ∈ seatids
        @test part2()+1 ∈ seatids
        @test part2() == 615
    end

    function test()
        test_part1()
        test_part2()
    end
end


# ---
# title: 986. Interval List Intersections
# id: problem986
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Two Pointers
# link: <https://leetcode.com/problems/interval-list-intersections/description/>
# hidden: true
# ---
#
# Given two lists of **closed** intervals, each list of intervals is pairwise
# disjoint and in sorted order.
#
# Return the intersection of these two interval lists.
#
# _(Formally, a closed interval`[a, b]` (with `a <= b`) denotes the set of real
# numbers `x` with `a <= x <= b`.  The intersection of two closed intervals is a
# set of real numbers that is either empty, or can be represented as a closed
# interval.  For example, the intersection of [1, 3] and [2, 4] is [2, 3].)_
#
#
#
# **Example 1:**
#
# **![](https://assets.leetcode.com/uploads/2019/01/30/interval1.png)**
#
#
#
#     Input: A = [[0,2],[5,10],[13,23],[24,25]], B = [[1,5],[8,12],[15,24],[25,26]]
#     Output: [[1,2],[5,5],[8,10],[15,23],[24,24],[25,25]]
#
#
#
#
# **Note:**
#
#   1. `0 <= A.length < 1000`
#   2. `0 <= B.length < 1000`
#   3. `0 <= A[i].start, A[i].end, B[i].start, B[i].end < 10^9`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# Tests for Row type and associated methods

@testset "Row" begin
    @testset "Constructors" begin
        @test_throws DomainError Row(-1, [])

        a = Row(0, [1])
        b = Row(a)
        c = Row(0)
        d = Row{Int}(0, 10)
        @test a == b == c == d

        a = Row(3, [1.0,3.0,3.0,1.0])
        b = Row(a)
        c = Row(3) # different data type
        d = Row{Float64}(3, 10)
        @test a == b == c == d

        a = Row(10)
        b = Row(a)
        b.data[2] = 4
        @test a != b
    end

    @testset "Accessors" begin
        a = Row(5)
        @test rownumber(a) == 5
        @test value(a) == [1,5,10,10,5,1]
    end

    @testset "Array functions" begin
        a = Row(7)
        b = Row{Float64}(7, 8)
        @test sum(a) == 128
        @test sum(x -> x^2, a) == 3432
        @test sum(b) == 128.0
        @test sum(x -> x^2, b) == 3432.0

        @test axes(a) == axes(b) == (0:1:7,)
        @test size(a) == size(b) == (8,)

        @testset "Data access $i" for i ∈ 0:7
            @test a[i] == b[i]
        end

        @test_throws BoundsError a[-1]
        @test_throws BoundsError a[10]

        @test firstindex(a) == firstindex(b) == 0
        @test lastindex(a) == lastindex(b) == 7
    end

    @testset "Checks" begin
        a = Row(0)
        b = Row(4)
        c = Row(4, [7]) # bad data
        d = Row(4, [6.0])

        @test isfirst(a)
        @test !isfirst(b)

        @test PT.isvalid(a)
        @test !PT.isvalid(c)

        @test toarray(b) == toarray(d)
    end

    @testset "Movement" begin
        a = Row(0)
        b = Row(8)
        c = Row(9)
        d = Row(10)

        @test_throws OutOfBoundsError up!(a)
        @test_throws OutOfBoundsError up(a)
        @test_throws OutOfBoundsError prev!(a)
        @test_throws OutOfBoundsError prev(a)

        @test up(c) == prev(c) == b
        @test down(c) == next(c) == d

        e = Row(b)
        f = Row(d)

        down!(b)
        next!(e)
        up!(d)
        prev!(f)
        @test b == c == d == e == f
    end
end


@show 1+1
@show 2+2
@show factorial(3)
fac = factorial
@show fac(3)
@show abs(-6)
@show 3-2(fac(2))-abs(-7)
Lion = 5
Pika = 6
@show Lion + Pika
@show a = 4
@show b = 3
@show a*b == 6
@show a*b == 2*6
nothing


# # Digging deeper

# The core functioning of Recombinase is based on a set of preimplemented analysis
# functions and on the [OnlineStats](https://github.com/joshday/OnlineStats.jl) package
# to compute summary statistics efficiently.

# ## Analysis functions

# Under the hood, the actual computations are implemented in some built-in `Analysis`
# objects. Each `Analysis` takes some vectors (a varying number, one for `density`
# and two for `prediction` for example), an optional `axis` argument to specify
# the `x` axis argument and returns an iterator of `(x, y)` values, representing
# the `y` value corresponding to each point of the `x` axis.

using Recombinase: density, hazard, cumulative, prediction
x = randn(100)
res = density(x)
collect(Iterators.take(res, 5)) # let's see the first few elements

# We can collect the iterator in a table container to see
# that we recover the `x` and `y` axis, ready for plotting

using JuliaDB
s = table(res)

# `prediction` is similar but requrires two arguments:

xs = 10 .* rand(100)
ys = sin.(xs) .+ 0.5 .* rand(100)

res = prediction(xs, ys)
table(res)

# The function `discrete` can turn any analysis into its discrete counter-part
# (analyses are continuous for numerical axes and discrete otherwise by default).
# It also automatically comutes the error of the prediction (s.e.m).

using Recombinase: discrete

x = rand(1:3, 1000)
y = rand(1000) .+ 0.1 .* x

res = discrete(prediction)(x, y)
table(res)

# ## OnlineStats integration

# Summary statistics are computed using the excellent [OnlineStats](https://github.com/joshday/OnlineStats.jl)
# package. First the data is split by the splitting variable (in this case `School`),
# then the analysis is computed on each split chunk on the selected column(s).
# Finally the results from different schools are put together in summary statistics
# (default is `mean` and `s.e.m`):

using Recombinase: datafolder, compute_summary
data = loadtable(joinpath(datafolder, "school.csv"))
compute_summary(density, data, :School; select = :MAch)

# Any summary statistic can be used. If we only want the mean we can use

using OnlineStats
compute_summary(density, data, :School; select = :MAch, stats = Mean())

# We can also pass a `Tuple` of statistics to compute many at the same time:

using OnlineStats
compute_summary(density, data, :School; select = :MAch, stats = Series(Mean(), Variance()))

# One can optionally pass pairs of `OnlineStat` and function to apply the function to the
# `OnlineStat` before plotting (default is just taking the `value`).
# To compute a different error bar (for example just the standard deviation) you can simply do:

stats = (Mean(), Variance() => t -> sqrt(value(t)))
compute_summary(density, data, :School; select = :MAch, stats = stats)


export Descent, Momentum, AdaGrad, Adam

const ε = 1e-8

struct Descent
    eta::Float64
end

Descent() = Descent(0.1)

apply!(opt::Descent, x, Δ) = opt.eta .* Δ

struct Momentum
    eta::Float64
    alpha::Float64
    velocity::Dict
end

Momentum(η = 0.01, α = 0.9) = Momentum(η, α, Dict())

function apply!(opt::Momentum, x, Δ)
    η, α = opt.eta, opt.alpha
    v = get!(() -> zero(x), opt.velocity, x)::typeof(x) #get! () function returns a view.
    v = (α*v - η*Δ)
    Δ = -v
end

struct AdaGrad
    eta::Float64
    h::Dict
end

AdaGrad(eta = 0.01) = AdaGrad(eta, Dict())

function apply!(opt::AdaGrad, x, Δ)
    η = opt.eta
    h = get!(() -> zero(x), opt.h, x)::typeof(x)
    @. h += Δ^2
    @. Δ *= η / (√h + ε)
end

struct Adam
    eta::Float64
    beta::Tuple{Float64, Float64}
    recode::Dict
end

Adam(η = 0.01, β = (0.9, 0.99)) = Adam(η, β, Dict())

function apply!(opt::Adam, x, Δ)
    η, β = opt.eta, opt.beta
    mt, vt , βp = get!(opt.recode, x) do
        (zero(x), zero(x), Float64[β[1], β[2]])
    end
    @. mt = β[1] * mt + (1 - β[1]) * Δ
    @. vt = β[2] * vt + (1 - β[2]) * Δ^2
    @. Δ = η * mt / (√(vt / (1 - βp[2])) + ε)
    βp .= βp .* β
    return Δ
end


"""
    Binion{T <: Real} = Elliptic{T}

A binion is an elliptic Cayley-Dickson construct with real numbers.
"""
const Binion{T <: Real} = Elliptic{T}

function Binion(a::T, b::T) where T <: Real
    Binion{T}(a, b)
end

function Binion(a::T) where T <: Real
    Binion{T}(a, zero(T))
end

function Binion(a::Real, b::Real)
    Binion(promote(a, b)...)
end

"""
    unreal(z::Binion)

The unreal part of a complex number is a 1-vector.
"""
function unreal(z::Binion{T}) where T <: Real
    T[z.r]
end

function show(io::IO, z::Binion)
    print(io, "[1: ")
    print(io, z.l)
    print(io, ", i: ")
    print(io, z.r)
    print(io, "]")
end

function random(::Type{Binion{T}}) where T <: Real
    Binion{T}(rand(T), rand(T))
end

function random(::Type{Binion{BigInt}})
    Binion{BigInt}(randomBigInt(), randomBigInt())
end

function random(::Type{Binion{BigFloat}})
    Binion{BigFloat}(randomBigFloat(), randomBigFloat())
end

function random(::Type{Binion{Rational{BigInt}}})
    Binion{Rational{BigInt}}(randomBigRational(), randomBigRational())
end

function convert(::Type{Binion{T}}, a::Real) where T <: Real
    Binion{T}(a, zero(T))
end

function convert(::Type{Binion{T}}, z::Binion) where T <: Real
    Binion{T}(z.l, z.r)
end

function promote_rule(::Type{Binion{T}}, ::Type{S}) where {T <: Real, S <: Real}
    Binion{promote_type(T, S)}
end

function promote_rule(::Type{Binion{T}}, ::Type{Binion{S}}) where {T <: Real, S <: Real}
    Binion{promote_type(T, S)}
end

function selfstar(z::Binion)
    (z + star(z)) / 2
end

function anti_selfstar(z::Binion)
    (z - star(z)) / 2
end


module TransformerModule

using LinearAlgebra: norm, Transpose, mul!
using FinEtools

"""
    QTEQTransformer

QTEQTransformer of element matrices.

A callable object: computes `Q^T * E * Q`, where `E` the element stiffness
matrix, and `Q` is the transformation matrix.

Buffers the intermediate result. Hence no allocation is incurred.
"""
struct QTEQTransformer
    _buff::FFltMat
    function QTEQTransformer(Q)
        _buff = fill(0.0, size(Q)...);
        return new(_buff)
    end
end

"""
    (o::QTEQTransformer)(E, Q)

Perform the transformation on the matrix `E` with the transformation matrix
`Q`: `Ebar = Q^T * E * Q`.
"""
(o::QTEQTransformer)(E, Q) = begin
    @assert size(o._buff) == size(Q)
    @assert size(E) == size(Q)
    mul!(o._buff, E, Q)
    mul!(E, Transpose(Q), o._buff)
    return E
end

"""
    QEQTTransformer

QEQTTransformer of element matrices.

A callable object: computes `Q * E * Q^T`, where `E` the element stiffness
matrix, and `Q` is the transformation matrix.

Buffers the intermediate result. Hence no allocation is incurred.
"""
struct QEQTTransformer
    _buff::FFltMat
    function QEQTTransformer(Q)
        _buff = fill(0.0, size(Q)...);
        return new(_buff)
    end
end

"""
    (o::QEQTTransformer)(E, Q)

Perform the transformation on the matrix `E` with the transformation matrix
`Q`: `Ebar = Q * E * Q^T`.
"""
(o::QEQTTransformer)(E, Q) = begin
    @assert size(o._buff) == size(Q)
    @assert size(E) == size(Q)
    mul!(o._buff, E, Transpose(Q))
    mul!(E, Q, o._buff)
    return E
end

struct Layup2ElementAngle
    M::FFltMat
    M2D::FFltMat
end

function Layup2ElementAngle()
    Layup2ElementAngle(fill(0.0, 3, 3), fill(0.0, 2, 2))
end

function (o::Layup2ElementAngle)(E_G, lcsmat)
    mul!(o.M, E_G', lcsmat)
    o.M2D[1, 1] = o.M[1, 1]
    o.M2D[1, 2] = o.M[1, 2]
    o.M2D[2, 1] = o.M[2, 1]
    o.M2D[2, 2] = o.M[2, 2]
    o.M2D[:, 1] ./= norm(@view o.M2D[:, 1])
    o.M2D[:, 2] ./= norm(@view o.M2D[:, 2])
    m = (o.M2D[1, 1] + o.M2D[2, 2]) / 2
    n = (o.M2D[1, 2] - o.M2D[2, 1]) / 2
    sn = n >= 0.0 ? +1.0 : -1.0
    n = sn * sqrt(1 - m^2)
    return m, n
end

end # module


using GCoptimization

width = 10
height = 5
pixelNum = width * height
labelNum = 7
result = zeros(pixelNum)

gco = GCoptimizationGridGraph(width, height, labelNum)

function dataFn(pixel::GCOSiteID, label::GCOLabelID)::GCOEnergyTermType
    if pixel < 25
        return label == 0 ? 0 : 10
    else
        return label == 5 ? 0 : 10
    end
end
setDataCost(gco, dataFn)

smoothFn(p1::GCOSiteID, p2::GCOSiteID, l1::GCOLabelID, l2::GCOLabelID)::GCOEnergyTermType = (l1-l2)*(l1-l2) <= 4 ? (l1-l2)*(l1-l2) : 4
setSmoothCost(gco, smoothFn)

println("Before optimization energy is ", compute_energy(gco))
expansion(gco, 2)
println("After optimization energy is ", compute_energy(gco))

result = [whatLabel(gco, i-1) for i = 1:pixelNum]

showall(result)


using Pkg;Pkg.activate(".")
using PairVelocities
using StaticArrays
using BenchmarkTools
using DelimitedFiles

rbins = LinRange(0.001,150.,150)
rbins_c = 0.5*(rbins[2:end] + rbins[1:end-1])

filename = "moments.csv"

positions, velocities, boxsize, redshift = read_abacus()
println(positions[:][1:20])
N = size(positions)[end]
println("N halos")
println(N)
positions = convert(Array{Float64}, positions)
velocities = convert(Array{Float64}, velocities)

moments = PairVelocities.compute_pairwise_velocity_moments(
    positions,
    velocities,
    rbins,
    boxsize,
)
DD = 2 .*moments[1][:]
bin_volume = 4/3 * pi .* diff(rbins.^3)
RR = N * (N-1)/boxsize^3 * bin_volume
println(DD)
println(RR)
xi = DD./RR .- 1.

open(filename; write=true) do f
    write(f, "r_c\txi\tv_r\tsigma_r\tsigma_t\tskewness_r\tskewness_rt\tkurtosis_r\tkurtosis_t\tkurtosis_rt\n")
    writedlm(f,
             zip(rbins_c, xi, moments[2][:],
            moments[3][:], moments[6][:],
            moments[4][:], moments[7][:],
            moments[5][:], moments[8][:], moments[9][:]))
    end

println("Wrote file !")


# ------------------------------------------------------------------------------------------
# # Introduction
#
# **DynamicalSystems.jl** is a Julia software library for the exploration of chaos and
# nonlinear dynamics.
#
# ## Contents of these tutorials
#
# 1. How to create a dynamical system and get a trajectory out of it.
# 3. Orbit diagrams & Poincaré surfaces of section of a dynamical system.
# 4. Lyapunov exponents of a dynamical system.
# 5. Entropies and dimensions of numerical data.
# 6. Delay coordinates embedding & Neighborhoods.
# 7. Quick summary of the rest of the features possible with **DynamicalSystems.jl**:
#     * Lyapunov exponent from numerical data
#     * Broomhead-King coordinates
#     * Finding fixed points of any map of any order
#     * Detecting chaos with GALI
#
# 8. How to predict timeseries (*still in beta!*).
#
# ## About DynamicalSystems.jl
#
# 1. **DynamicalSystems.jl** has been built from the ground up with the principles of
# clarity, intuition and robustness.
#
# 3. Aims to be a "library" in the literal sense: all offered functionality is well
# documented and accompanied with descriptions of the algorithms as well as references to
# original papers. This offers new users plenty of chances to learn new algorithms.
#
# ## Disclaimer
#
# **DynamicalSystems.jl** is in v1.0-alpha! Comments, testing, suggestions etc. are very
# welcome before official release.
#
# Any changes that will happen for v1.0 release will be summarized in the [News
# page](https://juliadynamics.github.io/DynamicalSystems.jl/latest/news/).
#
# ---
#
# # Documentation
#
# https://juliadynamics.github.io/DynamicalSystems.jl/latest/.
#
# In the same documentation page you will find a [list of
# contents](https://juliadynamics.github.io/DynamicalSystems.jl/latest/#contents).
#
# If you have questions, you can open issues in the appropriate [GitHub
# repository](https://github.com/JuliaDynamics); if you think your question is minor, don't
# hesitate to chat with us over [gitter](https://gitter.im/JuliaDynamics/Lobby).
#
# **DynamicalSystems.jl** is a registered package. On Julia v0.6 you only have to type
# `Pkg.add("DynamicalSystems")` to install its features.
#
# ## Documentation Strings
#
# All functions that are exported by **DynamicalSystems.jl** have very detailed
# documentation strings.
# * Simply press `?` followed by the function name to see the documentation strings!
#
# ---
#
# # About this tutorial series
# Tutorials 1-4 are intended for any audience *but* with some prior exposure to math and
# physics, as outlined in the prerequisites. In these tutorials, the methods shown will be
# explained and discussed in detail, so people unfamiliar with nonlinear dynamics and chaos
# can still follow!
#
# Tutorials 5 and onwards will target more a experienced audience. These tutorials will
# provide example applications, but will not explain the methods/algorithms used. This is so
# that we can demo the available features in a relatively short period of time. Therefore,
# don't feel worried if you are not able to fully grasp all concepts at that point!
#
# ## Prerequisites
# This tutorial series assumes some basic knowledge of Julia & physics/mathematics concepts.
#
# * Julia
#     * [Functions](https://docs.julialang.org/en/stable/manual/functions/#man-functions-1).
#     * [Broadcasting](https://docs.julialang.org/en/stable/manual/arrays/#Broadcasting-1)
# ("vectorizing" function application).
#     * [Methods](https://docs.julialang.org/en/stable/manual/methods/#Methods-1) and
# multiple dispatch. This also relies on basic knowledge of the
# [Type](https://docs.julialang.org/en/stable/manual/types/#man-types-1) system.
#     * Basic handling of vectors and matrices.
#
#
# * Mathematics/Physics
#     * State-space & trajectory of a dynamical system
#     * What does it mean for a trajectory to be periodic
#     * Jacobians & linearized dynamics
#
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/introduction-to-dynamicalsystems.jl/0. Introduction, Documentation and Prerequisites.jl


using Flux
using Zygote
using MAT
#using Plots
using Statistics
using Distributions
using LinearAlgebra
using Random
using BSON
using JLD
using PyPlot
using CUDA
#using Base.Iterators: partition

gpu_id = 1
CUDA.device!(gpu_id)

Xtrain = JLD.load("Data/XYtrain_XYvalid_processed_T1w.jld","Xtrain")
Ytrain = JLD.load("Data/XYtrain_XYvalid_processed_T1w.jld","Ytrain")
XValid = JLD.load("Data/XYtrain_XYvalid_processed_T1w.jld","XValid")
YValid = JLD.load("Data/XYtrain_XYvalid_processed_T1w.jld","YValid")

idx_shuffle = shuffle(collect(1:size(Xtrain,4)))
Xtrain = Xtrain[:,:,:,idx_shuffle]
Ytrain = Ytrain[:,:,:,idx_shuffle]

Xtrain = reshape(Xtrain,size(Xtrain,1),size(Xtrain,1),2,size(Xtrain,4))

Xtrain = gpu(Xtrain)
Ytrain = gpu(Ytrain)

batch_size = 10

train_loader = Flux.Data.DataLoader(Xtrain,Ytrain, batchsize=batch_size,shuffle=true)

#XValid = XValid[:,:,1,:]
XValid = reshape(XValid,size(XValid,1),size(XValid,1),2,size(XValid,4))

XValid = gpu(XValid)
YValid = gpu(YValid)

channels_in=2
channels_out=1

# Unet define
function _random_normal(shape...)
  return Float32.(rand(Normal(0.f0,0.02f0),shape...))
end

BatchNormWrap(out_ch) = BatchNorm(out_ch)

ResNetConvBlock(chs, kernel = (3, 3)) =
    Chain(Conv(kernel, chs=>chs,pad = (1, 1);init=_random_normal),
          BatchNorm(chs),
          x->relu.(x),
          Conv(kernel, chs=>chs,pad = (1, 1);init=_random_normal),
          BatchNorm(chs))

r = Chain(
            Conv((4,4),channels_in=>64,pad=(1,1),stride=(2,2);init=_random_normal),
            x->relu.(x),
            SkipConnection(ResNetConvBlock(64),(mx,x)->mx+x),
            SkipConnection(ResNetConvBlock(64),(mx,x)->mx+x),
            SkipConnection(ResNetConvBlock(64),(mx,x)->mx+x),
            SkipConnection(ResNetConvBlock(64),(mx,x)->mx+x),
            ConvTranspose((2,2),64=>channels_out,stride=(2,2);init=_random_normal),
            BatchNorm(channels_out)
        )

r = gpu(r)

#loss(x,y) = mean(abs,(u(x).-y))
loss(x,y) = mean((r(x).-y).^2)
validationloss() = mean(Flux.mse(r(XValid[:,:,:,(idx-1)*100+1:idx*100]),YValid[:,:,:,(idx-1)*100+1:idx*100]) for idx=1:round(Int,size(XValid,4)/100))
#validationloss() = mean(mean(abs,(u(XValid[:,:,:,(idx-1)*100+1:idx*100]).-YValid[:,:,:,(idx-1)*100+1:idx*100])) for idx=1:round(Int,size(XValid,4)/100))

opt = ADAM(0.001)

function my_custom_train!(loss, ps, data, opt)
    # training_loss is declared local so it will be available for logging outside the gradient calculation.
    local training_loss
    ps = Params(ps)
    for d in data
      if d isa AbstractArray{<:Number}
        gs = Flux.gradient(ps) do
          training_loss = loss(d)
          return training_loss
        end
      else
        gs = Flux.gradient(ps) do
          training_loss = loss(d...)
          return training_loss
        end
      end
      # Insert whatever code you want here that needs training_loss, e.g. logging.
      # logging_callback(training_loss)
      # Insert what ever code you want here that needs gradient.
      # E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge.
      #println("training_loss = ",training_loss,"; gradient[1] = ",gs[ps[1]],"; W1 = ",ps.order.data[1][1],"; b1 = ",ps.order.data[2][1])
      println("training_loss = ",training_loss)
      Flux.update!(opt, ps, gs)
      # Here you might like to check validation set accuracy, and break out to do early stopping.
    end
  end

loss_valid = zeros(100)
for epoch in 1:100
  learn_rate = 0.001*exp(log(0.1)/99*(epoch-1))
  opt = ADAM(learn_rate)
  #@time Flux.train!(loss, Flux.params(u), train_loader, opt, cb = cb)
  @time my_custom_train!(loss, Flux.params(r), train_loader, opt)
  #@time Flux.train!(loss, Flux.params(u), train_set, opt)
  #model = cpu(u)
  #BSON.@save string("Results/model_resnet_ptx_processed_7dir_gpu_epoch",string(epoch),".bson") model
  #u = gpu(model)
  loss_valid[epoch] = validationloss()
end

model = cpu(r)
BSON.@save string("Results/model_resnet_ptx_processed_7Vol_gpu_T1w_epoch100.bson") model

#=
loss_valid = zeros(50)
for epoch in 1:50
  learn_rate = 0.001*exp(log(0.1)/99*(epoch-1))
  opt = ADAM(learn_rate)
  @time Flux.train!(loss, Flux.params(r), train_loader, opt)
  #@time my_custom_train!(loss, Flux.params(r), train_loader, opt)
  #@time Flux.train!(loss, Flux.params(u), train_set, opt)
  #model = cpu(u)
  #BSON.@save string("Results/model_resnet_ptx_processed_7dir_gpu_epoch",string(epoch),".bson") model
  #u = gpu(model)
  loss_valid[epoch] = validationloss()
end

model = cpu(r)
BSON.@save string("Results/model_resnet_ptx_processed_7Vol_gpu_T1w_epoch50.bson") model
=#

#test
#=
XValid = JLD.load("XYtrain_XYvalid_processed.jld","XValid")
YValid = JLD.load("XYtrain_XYvalid_processed.jld","YValid")
=#

BSON.@load string("Results/model_resnet_ptx_processed_7Vol_gpu_T1w_epoch50.bson") model
slc_list = collect(1:100)
dir_list = collect(1:7)
XtestOutput4write = zeros(size(XValid,1),size(XValid,2),length(slc_list),length(dir_list))

Xtest = zeros(Float32,size(XValid,1),size(XValid,2),2,1)
Ytest = zeros(Float32,size(XValid,1),size(XValid,2),1,1)
for idx_dir_test in 1:length(dir_list)
  for idx_slc_test in 1:length(slc_list)
    idx_test = idx_slc_test + (idx_dir_test-1)*100
    Xtest[:,:,:,1] = XValid[:,:,:,idx_test]
    #Ytest[:,:,1,1] = YValid[:,:,idx_slc_test,idx_dir_test]
    XtestOutput4write[:,:,idx_slc_test,idx_dir_test] = model(Xtest)
  end
end

file = matopen("Results/XValidOutput_subj5_processed_7Vol_gpu_resnet_T1w_epoch50.mat", "w")
write(file, "XtestOutput4write", XtestOutput4write)
close(file)

# testing on training data
BSON.@load string("model_unet_ptx_processed_7dir_epoch",string(10),".bson") u
XValid = Xtrain[:,:,:,1:700]
XValid = reshape(XValid,size(XValid,1),size(XValid,2),100,7)
XtestOutput4write = zeros(size(XValid))

Xtest = zeros(Float32,size(XValid,1),size(XValid,2),1,1)
Ytest = zeros(Float32,size(XValid,1),size(XValid,2),1,1)
for idx_dir_test in 1:size(XValid,4)
  for idx_slc_test in 1:size(XValid,3)
    Xtest[:,:,1,1] = XValid[:,:,idx_slc_test,idx_dir_test]
    #Ytest[:,:,1,1] = YValid[:,:,idx_slc_test,idx_dir_test]
    XtestOutput4write[:,:,idx_slc_test,idx_dir_test] = u(Xtest)
  end
end

file = matopen("XValidOutput_traingsubj1_processed_7dir_epoch10.mat", "w")
write(file, "XtestOutput4write", XtestOutput4write)
close(file)


include(joinpath(@__DIR__, "pedagogy_graph.jl"))
using Random

dims = [5, 5]
g = LG.grid(dims)

pp = PRS.SinkFreeGraph(g)
rng = Random.MersenneTwister(7)
sink_free_graph_pedagogy(pp, dims, rng)


module kNN
    export knn, kernelregression

    using StatsBase
    using Distances
    using NearestNeighbors
	using SmoothingKernels
    using LinearAlgebra
    using Statistics

    include("bandwidth.jl")
    include("majority_vote.jl")
    include("classifier.jl")
    include("regress.jl")
    include("regression.jl")
    include("anomaly.jl")
end


function grad_y(x_hat_l,grad_x_hat_l,G, dG)
    m = size(G,1)
    return -kron(x_hat_l', I(m))*dG'-G*grad_x_hat_l'
end

function grad_x_hat(x_hat_l, dA,A,grad_x_hat_l, y, grad_y, K,dK)
    n = size(A,1)
    return kron(x_hat_l',I(n))*dA' + A*grad_x_hat_l' + kron(y',I(n))*dK' + K*grad_y'
end

function diff_K(A,P,G,S_l,dA,dG,dP,dS_l)
    m = size(G,1)
    n = size(A,1)
    P_inv = inv(P)
    Kmn = commutation_matrix(m,n)
    return kron(P_inv*G*S_l,I(n))*dA' + kron(P_inv*G,A)*dS_l'+kron(P_inv,A*S_l)*Kmn*dG' - kron(P_inv,A*S_l*G'*P_inv)*dP'
end

function diff_S(A,P,K,S_l,dA,dP,dK,dS_l,dQ)
    m = size(P,1)
    n = size(A,1)
    Knn = commutation_matrix(n,n)
    Kmn = commutation_matrix(m,n)
    b1 = (kron(A*S_l,I(n)) + kron(I(n),S_l*A')*Knn)*dA'
    b2 = (kron(K*P,I(n)) + kron(I(n),K*P)*Kmn)*dK'
    return b1 + kron(A,A)*dS_l' -b2 - kron(K,K)*dP' +dQ'
end

function diff_P(G,S_l,dG,dS_l,dR)
    m = size(G,1)
    n = size(G,2)
    Kmn = commutation_matrix(m,n)
    b1 = (kron(G*S_l,I(m)) + kron(I(m),G*S_l)*Kmn)*dG'
    return b1 + kron(G,G)*dS_l' + dR'
end

function diff_ll(P,y,grad_y,dP)
    P_inv = inv(P)
    return -1/2*det(P)*vec(P_inv)'*dP' - y'*P_inv*grad_y' + 1/2*y'*P_inv*dP'*P_inv*y
end

function diff_S0(A,S0,dA,dQ)
    n = size(A,1)
    Knn = commutation_matrix(n,n)
    b1 = (kron(A*S0,I(n)) + kron(I(n),A*S0)*Knn)*dA'
    ba = I(n^2) - kron(A,A)
    bb = b1 + dQ'
    return inv(ba)*bb
end


"""
    calc_preds(X, Z, B)

Predict values

# Arguments

- X = 2d array of floats consisting of the row covariates, standardized as
  necessary
- Z = 2d array of floats consisting of the column covariates, standardized
  as necessary
- B = 2d array of floats consisting of coefficient estimates

# Value

2d array of floats

"""
function calc_preds(X::AbstractArray{Float64,2},
                    Z::AbstractArray{Float64,2},
                    B::AbstractArray{Float64,2})

    # Predict new values
    X*B*transpose(Z)
end

"""
    calc_preds(preds, X, Z, B)

Predict values in place

# Arguments

- preds = 2d array of floats consisting of the predicted values, to be
  updated in place
- X = 2d array of floats consisting of the row covariates, standardized as
  necessary
- Z = 2d array of floats consisting of the column covariates, standardized
  as necessary
- B = 2d array of floats consisting of coefficient estimates

# Value

None; updates predicted values in place.

"""
function calc_preds!(preds::AbstractArray{Float64,2},
                     X::AbstractArray{Float64,2},
                     Z::AbstractArray{Float64,2},
                     B::AbstractArray{Float64,2})

    # Update predicted values in place
    mul!(preds, X*B, transpose(Z))
end


using FourVectors
using FourVectors.StaticArrays
using Test

@test FourVector(1,2,3; t=4) == FourVector(MVector(1,2,3,4.0))
@test FourVector(1,2,3; t=4) == Particle(p=[1,2,3], E=4.0)
@test Particle(0,0,3; msq=4^2) == Particle(p=[0,0,3], E=5.0)
#
@test mass(Particle(1,2,3; msq=4)) ≈ 2

@test collect(FourVector(1,2,3; t=4)) == [1,2,3,4]

x = Particle(1.1,2.2,3.0; msq=4)
newp = [1.1,1.1,1.1]
x.p .= newp
@test invmasssq(x) == x.E^2 - sum(abs2,newp)

@test x.E == x.T == x[0]== x[4]
@test x.Px == x.X == x.x == x.px
@test x.Py == x.Y == x.y == x.py
@test x.Pz == x.Z == x.z == x.pz

@test length(x) == 4
@test mass(x * 1000) ≈ 1000*mass(x)

@test sphericalangles(Particle(0,4,3; msq=4^2)) == (3/5,π/2)
@test boostfactor(Particle(4,0,0; msq=3^2)) == 5/3

x1,x2 = x[[1,2]]
x[[1,2]] .= x[[2,1]]
@test x[[1,2]] == [x2,x1]


# SteinGMMPosterior
# A Gaussian mixture model posterior distribution p(theta) = pi(theta | y_1,...,y_L) induced by the model
#
#   y_l | theta ~iid 0.5 N(theta1, sigma2y) + 0.5 N(theta1 + theta2, sigma2y)
#   (theta1, theta2) ~ N(0, diag(sigma2theta1, sigma2theta2))
#
# This distribution is analyzed in the paper
#
#   http://www.columbia.edu/~jwp2128/Teaching/E9801/papers/WellingTeh2011.pdf

using Distributions: Beta, Binomial, MvNormal, Normal
import Distributions

# Default hyperparameter value
const defsigma2y = 2.0

type SteinGMMPosterior <: SteinPosterior
    # Vector of observed datapoints
    y::Array{Float64, 1}
    # Model hyperparameters
    sigma2y::Float64
    sigma2theta1::Float64
    sigma2theta2::Float64
    # Stein factors
    c1::Float64
    c2::Float64
    c3::Float64
end

# Constructor with uniform setting of (c1,c2,c3) and
# default settings of sigma2y and sigma2theta
SteinGMMPosterior(y; sigma2y=defsigma2y, sigma2theta1=10.0, sigma2theta2=1.0,
                  c1=1.0, c2=1.0, c3=1.0) =
SteinGMMPosterior(y, sigma2y, sigma2theta1, sigma2theta2, c1, c2, c3)

# Draws L datapoints from the likelihood distribution
#   0.5 N(theta[1], sigma2y) + 0.5 N(theta[1] + theta[2], sigma2y)
function randgmm(theta::Array{Float64, 1}, L::Int64;
                 sigma2y::Float64=defsigma2y)
    bin = Binomial(L)
    k = Distributions.rand(bin)

    phi1 = Normal(theta[1], sqrt(sigma2y))
    phi2 = Normal(theta[1] + theta[2], sqrt(sigma2y))

    vcat(Distributions.rand(phi1, k), Distributions.rand(phi2, L-k))
end

# Returns the log prior log pi(theta)
function logprior(d::SteinGMMPosterior,
                  theta::Array{Float64, 1})
    priornormal = MvNormal(diagm([d.sigma2theta1, d.sigma2theta2]))
    Distributions.logpdf(priornormal, theta)
end

# Returns log pi(d.y[idx] | theta)
function loglikelihood(d::SteinGMMPosterior,
                       theta::Array{Float64, 1};
                       idx=1:length(d.y))
    y = d.y[idx]

    phi1 = Normal(theta[1]           , sqrt(d.sigma2y))
    phi2 = Normal(theta[1] + theta[2], sqrt(d.sigma2y))
    sum(log(0.5 .* (Distributions.pdf(phi1, y) + Distributions.pdf(phi2, y))))
end

# Returns the gradient of the log prior grad log pi(theta)
function gradlogprior(d::SteinGMMPosterior,
                      theta::Array{Float64, 1})
    -(1.0 ./ [d.sigma2theta1, d.sigma2theta2]) .* theta
end

# Returns grad_{theta} log pi(d.y[idx] | theta)
function gradloglikelihood(d::SteinGMMPosterior,
                           theta::Array{Float64, 1};
                           idx=1:length(d.y))
    y = d.y[idx]

    phi1 = Normal(theta[1]           , sqrt(d.sigma2y))
    phi2 = Normal(theta[1] + theta[2], sqrt(d.sigma2y))
    # we have to be careful computing ratios b/c of numerical stability
    maxlogpdf = max(Distributions.logpdf(phi1, y),
                    Distributions.logpdf(phi2, y))
    logratiodenom = log(
        sum(exp([Distributions.logpdf(phi1, y) Distributions.logpdf(phi2, y)]
                .- maxlogpdf), 2)
    )
    ratio1 = exp(Distributions.logpdf(phi1, y) - maxlogpdf - vec(logratiodenom))
    ratio2 = 1 - ratio1
    # we can ignore the 1/2, since we're taking the gradient
    gradtheta1 = (y - theta[1]           ) .* ratio1 +
                 (y - theta[1] - theta[2]) .* ratio2
    gradtheta2 = (y - theta[1] - theta[2]) .* ratio2

    totgradtheta1 = sum(gradtheta1) / d.sigma2y
    totgradtheta2 = sum(gradtheta2) / d.sigma2y

    [totgradtheta1, totgradtheta2]
end

function numdimensions(d::SteinGMMPosterior)
    2
end

function diaghessian(d::SteinGMMPosterior,
                     theta::Array{Float64,1};
                     idx=1:length(d.y))
    y = d.y[idx]

    phi1 = Normal(theta[1], sqrt(d.sigma2y))
    phi2 = Normal(theta[2], sqrt(d.sigma2y))

    dens1 = Distributions.logpdf(phi1, y)
    dens2 = Distributions.logpdf(phi2, y)
    maxlogpdf = max(dens1, dens2)
    logratiodenom = log(sum(exp([dens1 dens2] .- maxlogpdf), 2))

    w1 = exp(dens1 - maxlogpdf - vec(logratiodenom))
    w2 = 1 .- w1

    ptweights = w1 .* w2 ./ (d.sigma2y)^2
    hess11 =  sum(ptweights .* (y .- theta[1]) .^2)
    hess22 =  sum(ptweights .* (y .- theta[2]) .^2)
    hess12 = -sum(ptweights .* (y .- theta[1]) .* (y .- theta[2]))

    # rotate the hessian b/c means are theta1, theta1 + theta2
    [hess11, hess12 + hess22] .- 1 ./ [d.sigma2theta1, d.sigma2theta2]
end


export GSSmoother,smooth
using SparseArrays
using LinearAlgebra
using BenchmarkTools

struct GSSmoother
    invh2::Float64
    bc::NTuple{2,NTuple{2,BoundaryCondition}}
    finest::Bool
    function GSSmoother(invh2,nrows,ncols,bc,finest)
        new(invh2,bc,finest)
    end
end

@inline function sij(denom,i,j,sl,rl,ih2,nx,ny)
    @inbounds denom*(rl[i,j]+ih2*(sl[i,j-1]+sl[i-1,j]+sl[i+1,j]+sl[i,j+1]))
end

using LoopVectorization

@inline function smooth_line(nrm1,j,i1,sl,rl,ih2,denom)
    if nrm1 >= 0
    @fastmath @simd for i=i1:2:nrm1
            sl[i,j]=denom*(rl[i,j]+ih2*(sl[i,j-1]+sl[i-1,j]+sl[i+1,j]+sl[i,j+1]))
        end
    else
        # @show nrm1
        @fastmath @inbounds @simd for i=i1:2:nrm1
            sl[i,j]=denom*(rl[i,j]+ih2*(sl[i,j-1]+sl[i-1,j]+sl[i+1,j]+sl[i,j+1]))
        end
    end
end

@inline function smooth_seq(nrows,ncols,sl,rl,ih2)
    denom=1/(4ih2)
    nrm1=nrows-1

    @inbounds for j=2:2:ncols-1
        smooth_line(nrm1,j,2,sl,rl,ih2,denom)
        smooth_line(nrm1,j+1,3,sl,rl,ih2,denom)
    end

    # @inbounds for j=2:2:ncols-1
    #     smooth_line(nrm1,j,2,sl,rl,ih2,denom)
    #     smooth_line(nrm1,j+1,3,sl,rl,ih2,denom)
    # end
    @inbounds for j=2:2:ncols-1
        smooth_line(nrm1,j,3,sl,rl,ih2,denom)
        smooth_line(nrm1,j+1,2,sl,rl,ih2,denom)
    end

end

@inline function smooth_cart(nrows,ncols,sl,rl,ih2)
    denom=1/(4ih2)
    nrm1=nrows-1
    ncm1=ncols-1
    @avx for j=2:ncm1
         for i=2:nrm1
            sl[i,j]=denom*(rl[i,j]+ih2*(sl[i,j-1]+sl[i-1,j]+sl[i+1,j]+sl[i,j+1]))
        end
    end

    # @inbounds for j=2:2:ncols-1
    #     smooth_line(nrm1,j,2,sl,rl,ih2,denom)
    #     smooth_line(nrm1,j+1,3,sl,rl,ih2,denom)
    # end
    # @inbounds for j=2:2:ncols-1
    #     smooth_line(nrm1,j,3,sl,rl,ih2,denom)
    #     smooth_line(nrm1,j+1,2,sl,rl,ih2,denom)
    # end

end

function smooth(sl,rl,a::GSSmoother)
    ih2=a.invh2
    (nrows,ncols)=size(sl)
    @assert(size(rl)==size(sl))
    denom=1/(4ih2)
    # smooth_cart(nrows,ncols,sl,rl,ih2)
    smooth_seq(nrows,ncols,sl,rl,ih2)
    return
end


using SIAN

println("Setting up the problem")

ode = @ODEmodel(
  x1'(t) = -(a21 + a31 + a01) * x1(t) + a12 * x2(t) + a13 * x3(t) + u(t),
  x2'(t) = a21 * x1(t) - a12 * x2(t),
  x3'(t) = a31 * x1(t) - a13 * x3(t),
  y(t) = x1(t)
)

res = identifiability_ode(ode, get_parameters(ode); p=0.99, p_mod=0, nthrds=1)

println(res)


#####################################################################################
#Forecast using model averaging method with lagged aggregate variable.
#Averging model of size Q
#####################################################################################
function modelavg_oos_forc(J, h, h_lag, Y, Y_lag, X_lag, Q)
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    ŷ = zeros(Float64, size(J - (h_lag + h):length(Y)))
    N = size(X_lag, 2)
    picks = collect(combinations((1:N), Q))
    ŷ_picks = zeros(length(picks))
    for (n, vec_î) in enumerate(J - (h_lag + h):length(Y))
        for (p, pick) in enumerate(picks)
            LHS = Y[1:(vec_î - h)]
            sample_size = length(LHS)
            RHS = hcat([ones(sample_size) Y_lag[1:sample_size]], X_lag[1:sample_size, pick])
            β̂ = OLSestimator(LHS, RHS)
            RHS_crnt = [1.0; Y_lag[vec_î]; X_lag[vec_î, pick]]
            ŷ_picks[p] = dot(β̂, RHS_crnt)
        end
        ŷ[n] = mean(ŷ_picks)
    end
    RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
    return ŷ, RMSE
end


using ReachabilityBenchmarks, MathematicalSystems, Reachability, Plots

include(@relpath "heat_model.jl")
include(@relpath "heat_specifications.jl")

S = heat_model()
X0, options = heat_specification()
options = Dict(options)

# initial value problem
problem = InitialValueProblem(S, X0)

# partition with 1D blocks
partition = [i:i for i in 1:200]

# reachability algorithms
algorithm_reach = BFFPSV18(:δ => 1e-3, :vars => [133], :partition => partition)
algorithm_check = BFFPSV18(:δ => 1e-3, :vars => [133], :partition => partition)

# check property
options[:mode] = "check"
solution = solve(problem, options; op=algorithm_check)
@assert solution.satisfied

# compute flowpipe
options[:mode] = "reach"
solution = solve(problem, options; op=algorithm_reach)

# project flowpipe to time and x₁₃₃
solution.options[:plot_vars] = [0, 133]
solution_proj = project(solution)

# plot projection
plot(solution_proj)
savefig("heat")


# calculation of T-matrix using complex numbers.
# the only drawback of this is that it is not compatible with automatic differentiation

# TODO:
# 1- Use StaticArrays when possible.

export calculate_Tmatrix_for_spheroid

include("utils.jl")
include("geometry.jl")

# import
using VectorSphericalWaves
using StaticArrays
using EllipsisNotation
using LinearAlgebra
using Trapz
using HDF5

"""
    Calculate M_mn_wave as an array
This is the same as VectorSphericalHarmonics.M_mn_wave, but accept kr_array, θ_array, ϕ_array
Parameters
==========
kr_array, θ_array, ϕ_array : arrays of arbitrary shape

return
======
M_mn_wave_array_ : M_mn_wave with shape same as any of kr_array, θ_array, ϕ_array, with an added dimension to represent the three components
"""
function M_mn_wave_array(
        m::Int, n::Int, kr_array::AbstractVecOrMat{<:Complex{<:Real}}, θ_array::AbstractVecOrMat{R},
        ϕ_array::AbstractVecOrMat{R}; kind="regular"
    ) where {R <: Real}
    # Alok way is faster indeed!
    # TODO: @Alok, I think if we use boradcast it would be faster. I think avoiding preallocation makes the code cleaner and faster
    M_mn_wave_array_ = (_ -> zero(SVector{3,Complex})).(kr_array)
    for idx in eachindex(kr_array)
        M_mn_wave_array_[idx] = M_mn_wave(m, n, kr_array[idx], θ_array[idx], ϕ_array[idx], kind=kind)
    end
    return M_mn_wave_array_
end

"""
    Calculate N_mn_wave as an array
This is the same as VectorSphericalHarmonics.N_mn_wave, but accept kr_array, θ_array, ϕ_array
Parameters
==========
kr_array, θ_array, ϕ_array : arrays of arbitrary shape

return
======
N_mn_wave_array_ : N_mn_wave with shape same as any of kr_array, θ_array, ϕ_array, with an added dimension to represent the three components
"""
function N_mn_wave_array(
        m::Int, n::Int, kr_array::AbstractVecOrMat{<:Complex{<:Real}}, θ_array::AbstractVecOrMat{R},
        ϕ_array::AbstractVecOrMat{R}; kind="regular"
    ) where {R <: Real}
    # Alok way
    N_mn_wave_array_ = (_ -> zero(SVector{3,Complex})).(kr_array)
    for idx in eachindex(kr_array)
        N_mn_wave_array_[idx] = N_mn_wave(m, n, kr_array[idx], θ_array[idx], ϕ_array[idx], kind=kind)
    end
    return N_mn_wave_array_
end

#############################################################
# calculate J and Rg J, from equations 5.184 and 5.190
"""
    Calculate the integrand of J (including dS=r²sin(θ)dθdϕ) from equations 5.184 and 5.190.
It calculate all the integrand, and multiply it by r²sin(θ)

Inputs
======
m, n, m_, n_ : int, order and rank of VSWF inside and outside the particle
k1r_array, k2r_array : complex 2D array, product of wavevector and r, evaluated at points on the surface of the particle
θ_array, ϕ_array : complex 2D array, spherical coordinates of points on the particle surface
n̂_array : 3D array, unit vector normal to the surface of the particle
kind : string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
J_superscript : superscript at the top of J, it can be any of [11,12,21,22]

Outputs
=======
J
"""
function J_mn_m_n__integrand(
        m::Int, n::Int, m_::Int, n_::Int,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", J_superscript=11
    ) where {R <: Real, C <: Complex{R}}

    # determining the type of the first the second VSWF
    if J_superscript == 11 # TODO: this if-statement can be done more nicely. We separate J_superscript into two pieces, the number 1 represents M_mn_wave_array, while number 2 represents N_mn_wave_array
        first_function = M_mn_wave_array
        second_function = M_mn_wave_array
    elseif J_superscript == 12
        first_function = M_mn_wave_array
        second_function = N_mn_wave_array
    elseif J_superscript == 21
        first_function = N_mn_wave_array
        second_function = M_mn_wave_array
    elseif J_superscript == 22
        first_function = N_mn_wave_array
        second_function = N_mn_wave_array
    else
        throw(DomainError("J_superscript must be any of [11,12,21,22]"))
    end

    # determining the type of the first and second VSWF
    kind_first_function = "regular"
    if kind == "irregular"
        kind_second_function = "irregular"
    elseif kind == "regular"
        kind_second_function = "regular"
    else
        throw(DomainError("""kind must be any of ["regular", "irregular"]"""))
    end

    # the cross product
    cross_product_MN = cross.(
        first_function(m_, n_, k2r_array, θ_array, ϕ_array, kind=kind_first_function), # I can directly call M,N waves, with dots.
        second_function(-m, n, k1r_array, θ_array, ϕ_array, kind=kind_second_function)
    )

    # dot multiplying the cross product by the unit vector, and multiplying by (-1)^m
    cross_product_MN_dot_n̂ = (-1).^m .* vector_dot_product.(cross_product_MN, n̂_array)

    # multiplying by dS=r²sin(θ)
    J_integrand = surface_integrand(cross_product_MN_dot_n̂, r_array, θ_array)

    return J_integrand
end

"""
    Same as J_mn_m_n__integrand, but uses M_mn_wave_SVector and N_mn_wave_SVector
"""
function J_mn_m_n__integrand_SVector(
    m::Int, n::Int, m_::Int, n_::Int,
    k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
    r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
    kind="regular", J_superscript=11
) where {R <: Real, C <: Complex{R}}

# determining the type of the first the second VSWF
if J_superscript == 11 # TODO: this if-statement can be done more nicely. We separate J_superscript into two pieces, the number 1 represents M_mn_wave_SVector, while number 2 represents N_mn_wave_SVector
    first_function = M_mn_wave_SVector
    second_function = M_mn_wave_SVector
elseif J_superscript == 12
    first_function = M_mn_wave_SVector
    second_function = N_mn_wave_SVector
elseif J_superscript == 21
    first_function = N_mn_wave_SVector
    second_function = M_mn_wave_SVector
elseif J_superscript == 22
    first_function = N_mn_wave_SVector
    second_function = N_mn_wave_SVector
else
    throw(DomainError("J_superscript must be any of [11,12,21,22]"))
end

# determining the type of the first and second VSWF
kind_first_function = "regular"
if kind == "irregular"
    kind_second_function = "irregular"
elseif kind == "regular"
    kind_second_function = "regular"
else
    throw(DomainError("""kind must be any of ["regular", "irregular"]"""))
end

# the cross product
cross_product_MN = cross.(
    first_function.(m_, n_, k2r_array, θ_array, ϕ_array, kind=kind_first_function), # I can directly call M,N waves, with dots.
    second_function.(-m, n, k1r_array, θ_array, ϕ_array, kind=kind_second_function)
)

# dot multiplying the cross product by the unit vector, and multiplying by (-1)^m
cross_product_MN_dot_n̂ = (-1).^m .* vector_dot_product.(cross_product_MN, n̂_array)

# multiplying by dS=r²sin(θ)
J_integrand = surface_integrand(cross_product_MN_dot_n̂, r_array, θ_array)

return J_integrand
end

function J_mn_m_n_(
        m::Int, n::Int, m_::Int, n_::Int,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", J_superscript=11, rotationally_symmetric=false,
    ) where {R <: Real, C <: Complex{R}}
    if rotationally_symmetric
        # make sure that θ_array is 1D
        if length(size(θ_array)) != 1
            throw(DomainError("Since you have indicated << rotationally_symmetric = true >>, θ_array has to be 1D. Now it is $(length(size(θ_array)))D"))
        end
        ϕ_array = convert(typeof(θ_array), zeros(size(θ_array)))
    end

    # getting the integrand
    if rotationally_symmetric && (m != m_)
        # the integral over ϕ is 2π * δ_m_m_, so it is zero if m != m_
        J_integrand_dS = zeros(size(θ_array))

    else
        J_integrand_dS = J_mn_m_n__integrand_SVector(
            m, n,m_,n_,
            k1r_array,k2r_array,
            r_array,θ_array,ϕ_array,n̂_array;
            kind=kind,J_superscript=J_superscript
        )
    end

    # calculate the surface integral
    if rotationally_symmetric
        # integrate over θ only
        J = 2π * trapz((θ_array), J_integrand_dS)
    else
        # integrate over θ and ϕ
        # TODO: replace this integral with surface mesh quadrature, like this one: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5095443/ , https://github.com/jareeger/Smooth_Closed_Surface_Quadrature_RBF-julia
        # assuming that θ_array, ϕ_array were created with meshgrid function
        J = trapz((θ_array[:,1], ϕ_array[1,:]), J_integrand_dS)
end

    return J
end

function Q_mn_m_n_(
        m::Int, n::Int, m_::Int, n_::Int,
        k1::C, k2::C,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", Q_superscript=11, rotationally_symmetric=false,
    ) where {R <: Real, C <: Complex{R}}
    if Q_superscript == 11; J_superscript_1 = 21 ; J_superscript_2 = 12
    elseif Q_superscript == 12; J_superscript_1 = 11 ; J_superscript_2 = 22
    elseif Q_superscript == 21; J_superscript_1 = 22 ; J_superscript_2 = 11
    elseif Q_superscript == 22; J_superscript_1 = 12 ; J_superscript_2 = 21
    end

    Q = (
        -im .* k1 .* k2 .* J_mn_m_n_(m, n, m_, n_, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array; kind=kind, J_superscript=J_superscript_1, rotationally_symmetric=rotationally_symmetric)
        - im .* k1.^2    .* J_mn_m_n_(m, n, m_, n_, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array; kind=kind, J_superscript=J_superscript_2, rotationally_symmetric=rotationally_symmetric)
    )

    return Q
end

function Q_matrix(
        n_max::Int,
        k1::C, k2::C,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false,
        verbose=false,
    ) where {R <: Real, C <: Complex{R}}
    idx_max = get_max_single_index_from_n_max(n_max)
    Q_mn_m_n_11 = zeros(typeof(k1), idx_max, idx_max) # TODO: @Alok, should I replace arrays with SMatrix? I am afraid it may get slower, as I have seen that StaticArray may get slower for arrays larger than 100 elements
    Q_mn_m_n_12 = zeros(typeof(k1), idx_max, idx_max)
    Q_mn_m_n_21 = zeros(typeof(k1), idx_max, idx_max)
    Q_mn_m_n_22 = zeros(typeof(k1), idx_max, idx_max)

    idx = 0;
    for n = 1:n_max
        for m = -n:n
            idx += 1
            idx_ = 0;
            for n_ = 1:n_max
                for m_ = -n_:n_
                    idx_ += 1
                    if verbose; println("n,m,idx = $n,$m,$idx  n_,m_,idx_ = $n_,$m_,$idx_"); end

                    if rotationally_symmetric
                        if m == m_
                            if symmetric_about_plane_perpendicular_z
                                # apply equations 5.208 and 5.209
                                if iseven(n + n_)
                                    Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                                    Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                                else
                                    Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                                    Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                                end
                            else
                                Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                            end
                        end
                    else
                        if symmetric_about_plane_perpendicular_z && (m == m_)
                            # apply equations 5.208 and 5.209
                            if iseven(n + n_)
                                Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                            else
                                Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                            end
                        else
                            Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                            Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                            Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                            Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                        end
            end
        end
    end
        end
    end
    Q = vcat(
        (hcat(Q_mn_m_n_11, Q_mn_m_n_12)),
        (hcat(Q_mn_m_n_21, Q_mn_m_n_22))
    )
    return Q
end

function T_matrix(
        n_max::Int,
        k1::C, k2::C,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false, HDF5_filename=nothing,
        verbose=false, create_new_arrays=false, BigFloat_precision = nothing
    ) where {R <: Real, C <: Complex{R}}

    if BigFloat_precision != nothing
        return setprecision(BigFloat_precision) do
            return T_matrix(
                n_max, big(k1), big(k2), big.(k1r_array), big.(k2r_array), big.(r_array), big.(θ_array), big.(ϕ_array), [big.(n) for n in n̂_array];
                HDF5_filename=HDF5_filename, rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z,
            )
        end
    else
        if create_new_arrays
            RgQ =       Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="regular"  , rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose)
            Q_inv = inv(Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="irregular", rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose))
            T = -1 .* RgQ * Q_inv

        else
            T = (
                -     Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="regular"  , rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose)
                * inv(Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="irregular", rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose))
            )
        end

        if HDF5_filename != nothing
            save_Tmatrix_to_HDF5_file(T, HDF5_filename)
        end

        return T
    end
end

"""
    Calculate T-matrix for a spheroid, with an extra kwarg `BigFloat_precision`

BigFloat_precision : set to `nothing` by default. When BigFloat_precision == nothing, then the function will use input types as they are.
"""
function calculate_Tmatrix_for_spheroid(
        rx::R, rz::R, n_max::Int,
        k1::Complex{R}, k2::Complex{R};
        n_θ_points=10, n_ϕ_points=20, HDF5_filename=nothing,
        rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false,
        BigFloat_precision = nothing
    ) where {R <: Real}

    # create a grid of θ_ϕ
    θ_array, ϕ_array = meshgrid_θ_ϕ(n_θ_points, n_ϕ_points; min_θ=1e-16, min_ϕ=1e-16, rotationally_symmetric=rotationally_symmetric)

    # calculate r and n̂ for the geometry
    r_array, n̂_array = ellipsoid(rx, rz, θ_array);

    # calculate T-matrix
    k1r_array = k1 .* r_array;
    k2r_array = k2 .* r_array;
    T = T_matrix(
        n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;
        HDF5_filename=HDF5_filename, rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z,
        BigFloat_precision = BigFloat_precision
    )
    return T
end

"""
    General and nice wrapper for T-matric
Inputs
surrounding material (e.g., dielectric constant, refractive index, wavevector)
particle material (e.g., dielectric constant, refractive index, wavevector)
angular_resolution
n_max
particle geometry: it can be:
- function ``r(θ)`` (i.e., axi-symmetric particle)
- a function ``r(θ,ϕ)`` (i.e., arbitrary 3D particle)
- a meshfile
- a string defining geometry class (e.g., "cylinder", "spheroid", "ellipsoid", etc.)

# output
# - T_matrix,
"""
function Tmatrix_nice()
    # TODO
end

### precompiling
"""
R =
precompile(T_matrix, (
    Int, Complex{R}, k2::Complex{R},
    k1r_array::AbstractVecOrMat{<:Complex{<:Real}}, k2r_array::AbstractVecOrMat{<:Complex{<:Real}},
    r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
    rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false, HDF5_filename=nothing,
    verbose=false, create_new_arrays=false,)
"""


"""
computeElemsByMEBI
This function computes the vector of the element indexes for a given vector of MEBI parameter
"""
function computeElemsByMEBI( MEBIVals, MEBIVec )

    # get the relevant BCs definitions
    MEBIUnique = sort( unique( MEBIVals ) )     # get and sort them
    MEBIUnique[1]==0 && popfirst!( MEBIUnique ) # and remove zeros

    elements = []
    # loop
    print(MEBIUnique)
    for MEBIval in MEBIUnique
        print("\nMEBI:", MEBIval, "\n" )
        # find the elements with the current MEBIval
        aux = findall( x->x==MEBIval, MEBIVec )

        # if there are any, add them to the list
        length(aux) > 0 && push!( elements, aux )
    end
    return elements
end

function computeNeumDofs( nNodes )

    neumDofs = nodes2dofs( (1:nNodes), 6 )

    neumDofs[2:2:end] .= 0

    return neumDofs

end

"""
nodes2dofs computes the vector of dofs for an input vector of nodes
"""
function nodes2dofs( nodes , degreespernode )
    n    = length(nodes);
    dofs = zeros( Int, n*degreespernode ) ;
    for i in (1:n);
        print("\ni:",i,"\n")
        print( (i-1)*degreespernode .+ Vector(1:degreespernode) )
        print((degreespernode*(nodes[i]-1)).+Vector(1:degreespernode))
        dofs[ (i-1)*degreespernode .+ Vector(1:degreespernode) ] = (degreespernode*(nodes[i]-1)).+Vector(1:degreespernode)  ;
    end
    return dofs;
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.1
#     language: julia
#     name: julia-1.6
# ---

# %%
struct F{p, T<:Integer} <: Integer
    a::T
    F{p, T}(a::Integer) where {p, T<:Integer} = new{p, T}(mod(T(a), p))
end
F{p}(a::Integer) where p = F{p, typeof(a)}(mod(a, p))
F{p, S}(x::F{p, T}) where {S<:Integer, p, T<:Integer} = F{p, S}(x.a)

Base.promote_rule(::Type{F{p, T}}, ::Type{S}) where {p, T<:Integer, S<:Integer} =
    F{p, promote_type(T, S)}

Base.zero(::Type{F{p, T}}) where {p, T<:Integer} = F{p}(mod(zero(T), p))
Base.one(::Type{F{p, T}}) where {p, T<:Integer} = F{p}(mod(one(T), p))
for op in (:-, :+)
    @eval Base.$op(x::F{p}) where p = F{p}(mod($op(x.a), p))
end
for op in (:-, :+, :*)
    @eval Base.$op(x::F{p}, y::F{p}) where p = F{p}(mod($op(x.a, y.a), p))
end
Base.inv(x::F{p}) where p = F{p}(invmod(x.a, p))
Base.:/(x::F{p}, y::F{p}) where p = x * inv(y)
Base.:\(x::F{p}, y::F{p}) where p = inv(x) * y
Base.:(==)(x::F{p}, y::F{p}) where p = x.a == y.a
Base.:<(x::F{p}, y::F{p}) where p = x.a < y.a

Base.show(io::IO, x::F{p}) where p = print(io, "F", p, '(', x.a, ')')

F7 = F{7, BigInt}
x, y = F7(10), F7(-2)
@show(x, y, zero(x), one(x), +x, -x, x + y, x - y, x * y, x / y, x \ y, x^3, x^-5, x == F7(3), x == 3)
println()
A = F7[1 2; 3 4]
@show(A, 4A, A/4)
println()
using LinearAlgebra
L, U = lu(A)
@show(L, U, L * U, det(A), inv(A), inv(A) * A, A * inv(A));

# %%
"""See https://docs.julialang.org/en/v1/manual/interfaces/"""
Base.iterate(Fp::Type{F{p, T}}) where {p, T<:Integer} = (zero(Fp), zero(T))
function Base.iterate(Fp::Type{F{p, T}}, state) where {p, T<:Integer}
    nextstate = state + 1
    nextstate < p ? (Fp(nextstate), nextstate) : nothing
end
Base.IteratorSize(Fp::Type{F{p, T}}) where {p, T<:Integer} = Base.HasLength()
Base.length(Fp::Type{F{p, T}}) where {p, T<:Integer} = p
Base.eltype(Fp::Type{F{p, T}}) where {p, T<:Integer} = Fp

squares(Fp) = Fp[x^2 for x in Fp]
squareroots(k, Fp) = Fp[x for x in Fp if x^2 == k]
@show(collect(F7), squares(F7), squareroots.(0:6, Ref(F7)));

# %%
using GaloisFields, LinearAlgebra
GF7 = @GaloisField 7
B = GF7[1 2; 3 4]
det(B)

# %%
inv(B)

# %%
lu(B)

# %%
using AbstractAlgebra

@show GF7 = GF(7)
@show P, x = PolynomialRing(GF7, "x")
@show C = GF7[1 2; 3 4]
squares(Fp) = [x^2 for x in Fp]
squareroots(k, Fp) = [x for x in Fp if x^2 == k]
println()
@show(det(C), inv(C), lu(C), charpoly(P, C))
println()
@show(collect(GF7), squares(GF7), squareroots.(0:6, Ref(GF7)));

# %%
@which iterate(GF7)

# %%
@code_warntype iterate(F7, big(1))

# %%
[x for x in F7 if x^2 == 4]

# %%
[x for x in F7 if x^2 == 5]

# %%
@which collect(x for x in F7 if x^2 == 5)

# %%
Base.@default_eltype(x for x in F7 if x^2 == 5)

# %%
Base.@default_eltype(x for x in F7 if x^2 == 5)

# %%
(@macroexpand Base.@default_eltype(x for x in F7 if x^2 == 5)) |> Base.remove_linenums!

# %%
iter = (x for x in F7 if x^2 == 5)

# %%
Base._iterator_upper_bound(iter)

# %%
Core.Compiler.return_type(Base._iterator_upper_bound, Tuple{typeof(iter)})

# %%
@which Base.inferencebarrier(nothing)

# %%
iter2 = (x for x in F7 if x^2 == 4)

# %%
Base._iterator_upper_bound(iter2)

# %%
iter3 = (x for x in GF7 if x^2 == 4)
collect(iter3)

# %%
iter4 = (x for x in GF7 if x^2 == 5)
collect(iter4)

# %%
Core.Compiler.return_type(Base._iterator_upper_bound, Tuple{typeof(iter4)})

# %%
Base._iterator_upper_bound(iter4)

# %%
@code_warntype Base._iterator_upper_bound(iter4)

# %%
@code_warntype Base._iterator_upper_bound(iter2)

# %%
@code_warntype iterate(iter4)

# %%
@code_warntype iterate(iter)

# %%
Core._apply_iterate(Base.iterate, Base.iterate, Core.tuple(iter3.iter), ())

# %%
Core._apply_iterate(Base.iterate, Base.iterate, Core.tuple(iter2.iter), ())

# %%
Core.tuple(iter.iter)

# %%
Core.tuple(iter4.iter)

# %%

0010/minimal def of a prime field old.jl


### Instability Plots
# Generates the code for the intability/suboptimality in the paper
# Presentation uses similar plots but with different paramters, see presentation plots.jl

using Distributions, Random, DelimitedFiles
include("../src/JS_SAA_main.jl")

function genSinglePathCurves(subcase; seed=8675409, usePoisson=true)
	K = 10000
	N = 10

	if subcase == :TradeoffBad
		p0 = .3
		s = .5
		path = string("../Results/TradeoffBad", usePoisson, ".csv")
	elseif subcase == :TradeoffGoodP0
		p0 = .75
		s = .5
		path = string("../Results/TradeoffGoodP0", usePoisson, ".csv")
	elseif subcase == :TradeoffGoodS
		p0 = .3
		s = .2
		path = string("../Results/TradeoffGoodS", usePoisson, ".csv")
	end
	Random.seed!(seed)

	#Gen Pks uniformly over interval for now
	ps = rand(K)
	ps = .6 .+ .3 * rand(K)

	f = open(path, "w")

	if usePoisson
		Nhats = rand(Poisson(N), K)
		Nhats[Nhats .== 0] .= 1
	else
		Nhats = fill(N, K)
	end
	mhats = JS.nv_sim_path(ps, Nhats);

	alpha_grid = range(0, stop=2N, length=100)

	#note, this actually computes N * ZLoo
	#loo_perf = map(a->JS.nv_loo(mhats, p0, a, Nhats, s), alpha_grid)
	alphaOR, jstar, outOR = JS.nv_oracle_alpha(mhats, ps, p0, alpha_grid, Nhats, s)
	alphaLOO, jstar, loo_perf = JS.nv_loo_alpha(mhats, p0, alpha_grid, Nhats, s)
	loo_perf /= N

	saa_subopt = map(a-> JS.nv_saa(mhats, p0, a, Nhats, s), alpha_grid) ./ N
	saa_zero = saa_subopt[1]
	saa_subopt = saa_subopt .- saa_zero
	instab = loo_perf - saa_subopt .- saa_zero;

	#Header
	writedlm(f, ["Alpha" "OR" "LOO" "SAA-SubOpt" "Instability"], ',')
	writedlm(f, [vec(alpha_grid) outOR loo_perf saa_subopt instab], ',')
	close(f)
end

genSinglePathCurves(:TradeoffBad)
genSinglePathCurves(:TradeoffGoodP0)
genSinglePathCurves(:TradeoffGoodS)

genSinglePathCurves(:TradeoffBad, usePoisson=false)
genSinglePathCurves(:TradeoffGoodP0, usePoisson=false)
genSinglePathCurves(:TradeoffGoodS, usePoisson=false)


using ClusterManagers
using Distributed

OnCluster = false #set to false to run locally
addWorkers = false #set to false to run serially
println("OnCluster = $(OnCluster)")

# Current number of workers
#--------------------------
currentWorkers = nworkers()
println("Initial number of workers = $(currentWorkers)")

# Increase the number of workers available
#-----------------------------------------
maxNumberWorkers = 1
if addWorkers == true
    if OnCluster == true
        addprocs(SlurmManager(maxNumberWorkers))
    else
        addprocs(maxNumberWorkers)
    end
end

# Sanity checks
#-------------
hosts = []
pids = []
for i in workers()
    host, pid = fetch(@spawnat i (gethostname(), getpid()))
    println("Hello I am worker $(i), my host is $(host)")
    push!(hosts, host)
    push!(pids, pid)
end

currentWorkers = nworkers()
println("Number of workers = $(currentWorkers)")

using Plots
using LaTeXStrings
using ParallelDataTransfer

using MSM
using DataStructures
using OrderedCollections
using Distributions
using Random
using DataStructures
using Statistics
using LinearAlgebra

# Generate simulated data
Random.seed!(1234)  #for replicability reasons
T = 10000         #number of periods
P = 2               #number of dependent variables
beta0 = rand(P)     #choose true coefficients by drawing from a uniform distribution on [0,1]
alpha0 = rand(1)[]  #intercept
theta0 = 0.0        #coefficient to create serial correlation in the error terms
println("True intercept = $(alpha0)")
println("True coefficient beta0 = $(beta0)")
println("Serial correlation coefficient theta0 = $(theta0)")

# Generation of error terms
# row = individual dimension
# column = time dimension
U = zeros(T)
d = Normal()
U[1] = rand(d, 1)[] #first error term
# loop over time periods
for t = 2:T
    U[t] = rand(d, 1)[] + theta0 * U[t-1]
end

# Let's simulate the dependent variables x_t
x = zeros(T, P)

d = Uniform(0, 5)
for p = 1:P
    x[:, p] = rand(d, T)
end

# Let's calculate the resulting y_t
y = zeros(T)

for t = 1:T
    y[t] = alpha0 + x[t, 1] * beta0[1] + x[t, 2] * beta0[2] + U[t]
end

# Visualize data
p1 = scatter(
    x[1:100, 1],
    y[1:100],
    xlabel = "x1",
    ylabel = "y",
    legend = :none,
    smooth = true,
)
p2 = scatter(
    x[1:100, 2],
    y[1:100],
    xlabel = "x2",
    ylabel = "y",
    legend = :none,
    smooth = true,
)
plot(p1, p2)

# Define locally
optionsSMM = MSMOptions(
    maxFuncEvals = 2000,
    globalOptimizer = :dxnes,
    localOptimizer = :NelderMead,
)
myProblem = MSMProblem(options = optionsSMM);

# Priors
dictPriors = OrderedDict{String,Array{Float64,1}}()
dictPriors["alpha"] = [0.5, 0.001, 1.0]
dictPriors["beta1"] = [0.5, 0.001, 1.0]
dictPriors["beta2"] = [0.5, 0.001, 1.0]

# Empirical moments
dictEmpiricalMoments = OrderedDict{String,Array{Float64,1}}()
dictEmpiricalMoments["mean"] = [mean(y)] #informative on the intercept
dictEmpiricalMoments["mean_x1y"] = [mean(x[:, 1] .* y)] #informative on betas
dictEmpiricalMoments["mean_x2y"] = [mean(x[:, 2] .* y)] #informative on betas
dictEmpiricalMoments["mean_x1y^2"] = [mean((x[:, 1] .* y) .^ 2)] #informative on betas
dictEmpiricalMoments["mean_x2y^2"] = [mean((x[:, 2] .* y) .^ 2)] #informative on betas

W = Matrix(1.0 .* I(length(dictEmpiricalMoments)))#initialization
#Special case: diagonal matrix
#(you may choose something else)
for (indexMoment, k) in enumerate(keys(dictEmpiricalMoments))
    W[indexMoment, indexMoment] = 1.0 / (dictEmpiricalMoments[k][1])^2
end

set_priors!(myProblem, dictPriors)
set_empirical_moments!(myProblem, dictEmpiricalMoments)
set_weight_matrix!(myProblem, W)

# x[1] corresponds to the intercept, x[2] corresponds to beta1, x[3] corresponds to beta2
function functionLinearModel(
    x;
    uniform_draws::Array{Float64,1},
    simX::Array{Float64,2},
    nbDraws::Int64 = length(uniform_draws),
    burnInPerc::Int64 = 0,
)
    T = nbDraws
    P = 2       #number of dependent variables

    alpha = x[1]
    beta = x[2:end]
    theta = 0.0     #coefficient to create serial correlation in the error terms

    # Creation of error terms
    # row = individual dimension
    # column = time dimension
    U = zeros(T)
    d = Normal()
    # Inverse cdf (i.e. quantile)
    gaussian_draws = quantile.(d, uniform_draws)
    U[1] = gaussian_draws[1] #first error term

    # loop over time periods
    for t = 2:T
        U[t] = gaussian_draws[t] + theta * U[t-1]
    end

    # Let's calculate the resulting y_t
    y = zeros(T)

    for t = 1:T
        y[t] = alpha + simX[t, 1] * beta[1] + simX[t, 2] * beta[2] + U[t]
    end

    # Get rid of the burn-in phase:
    #------------------------------
    startT = max(1, Int(nbDraws * (burnInPerc / 100)))

    # Moments:
    #---------
    output = OrderedDict{String,Float64}()
    output["mean"] = mean(y[startT:nbDraws])
    output["mean_x1y"] = mean(simX[startT:nbDraws, 1] .* y[startT:nbDraws])
    output["mean_x2y"] = mean(simX[startT:nbDraws, 2] .* y[startT:nbDraws])
    output["mean_x1y^2"] = mean((simX[startT:nbDraws, 1] .* y[startT:nbDraws]) .^ 2)
    output["mean_x2y^2"] = mean((simX[startT:nbDraws, 2] .* y[startT:nbDraws]) .^ 2)
    return output
end

# Let's freeze the randomness during the minimization
d_Uni = Uniform(0, 1)
nbDraws = T #number of draws in the simulated data
uniform_draws = rand(d_Uni, nbDraws)
simX = zeros(length(uniform_draws), 2)
d = Uniform(0, 5)
for p = 1:2
    simX[:, p] = rand(d, length(uniform_draws))
end

# Construct the objective function everywhere
set_simulate_empirical_moments!(
    myProblem,
    x -> functionLinearModel(x, uniform_draws = uniform_draws, simX = simX),
)
construct_objective_function!(myProblem)

println("Global Algorithm")
# Choose a global optimizer that supports parallel evaluations (e.g. xnes or dxnes)
# (see the documentation: https://github.com/robertfeldt/BlackBoxOptim.jl)
msm_optimize!(myProblem, verbose = false)

# Access the results using best_candidate
minimizer = msm_minimizer(myProblem)
minimum_val = msm_minimum(myProblem)

println("Minimum objective function = $(minimum_val)")
println("Estimated value for alpha = $(minimizer[1]). True value for beta1 = $(alpha0[1]) \n")
println("Estimated value for beta1 = $(minimizer[2]). True value for beta1 = $(beta0[1]) \n")
println("Estimated value for beta2 = $(minimizer[3]). True value for beta2 = $(beta0[2]) \n")

println("Multistart Algorithm")
# Start several local optimization algorithms in parallel
# Choose algorithms from the package Optim.jl (https://github.com/JuliaNLSolvers/Optim.jl)
# The "global" mininimum is the minimum of the local minima.
msm_multistart!(myProblem, nums = nworkers(), verbose = false)
minimizer_multistart = msm_multistart_minimizer(myProblem)
minimum_multistart = msm_multistart_minimum(myProblem)

println("Minimum objective function = $(minimum_multistart)")
println("Estimated value for alpha = $(minimizer_multistart[1]). True value for beta1 = $(alpha0[1]) \n")
println("Estimated value for beta1 = $(minimizer_multistart[2]). True value for beta1 = $(beta0[1]) \n")
println("Estimated value for beta2 = $(minimizer_multistart[3]). True value for beta2 = $(beta0[2]) \n")

# Empirical Series
#-----------------
X = zeros(T, length(dictEmpiricalMoments))
X[:, 1] = y
X[:, 2] = (x[:, 1] .* y)
X[:, 3] = (x[:, 2] .* y)
X[:, 4] = (x[:, 1] .* y) .^ 2
X[:, 5] = (x[:, 2] .* y) .^ 2

# "Distance Matrix" (see Duffie and Singleton, 1993)
Sigma0 = cov(X)

# Heatmap to visualize correlation
xs = [string("x", i) for i = 1:length(dictEmpiricalMoments)]
ys = [string("x", i) for i = 1:length(dictEmpiricalMoments)]
z = cor(X)
hh = heatmap(xs, ys, z, aspect_ratio = 1)

set_Sigma0!(myProblem, Sigma0)
# nbDraws = number of draws in the simulated data
# To decrease standard errors, increase nbDraws
calculate_Avar!(myProblem, minimizer_multistart, tau = T / nbDraws)

df = summary_table(myProblem, minimizer_multistart, T, 0.05)
println(df)

# Check the rank condition
# Local identification requires D to be full column rank (in a neighborhood of the solution)
D = calculate_D(myProblem, minimizer_multistart)
println("number of parameters: $(size(D,2))")
println("rank of D is: $(rank(D))")

# Slices
vXGrid, vYGrid = msm_slices(myProblem, minimizer_multistart, nbPoints = 7);
p1 = plot(vXGrid[:, 1],vYGrid[:, 1],title = L"\alpha", label = "",linewidth = 3, xrotation = 45)
plot!(p1, [minimizer_multistart[1]], seriestype = :vline, label = "",linewidth = 1)
p2 = plot(vXGrid[:, 2],vYGrid[:, 2],title = L"\beta_1", label = "",linewidth = 3, xrotation = 45)
plot!(p2, [minimizer_multistart[2]], seriestype = :vline, label = "",linewidth = 1)
p3 = plot(vXGrid[:, 3],vYGrid[:, 3],title = L"\beta_2", label = "",linewidth = 3, xrotation = 45)
plot!(p3, [minimizer_multistart[3]], seriestype = :vline, label = "",linewidth = 1)
plot_combined = plot(p1, p2, p3)
savefig(plot_combined, "slices.png")
display(plot_combined)

# Efficient GMM
#---------------
myProblem.W = inv(Sigma0)
msm_multistart!(myProblem, nums = nworkers(), verbose = false)
minimizer_multistart = msm_multistart_minimizer(myProblem)
minimum_multistart = msm_multistart_minimum(myProblem)

# J-test
J, c = J_test(myProblem, minimizer_multistart, T, nbDraws, 0.05)
println("J-statistic: $(J)")
println("Critical value: $(c)")

# Compare results with GLM
using DataFrames, GLM
data = DataFrame(x1 = x[:, 1], x2 = x[:, 2], y = y[:]);
ols = lm(@formula(y ~ x1 + x2), data)
println(ols)


using IterativeSolvers
using FactCheck

n = 10
m = 6
srand(1234321)

include("common.jl")

##################
# Linear solvers #
##################

facts("Stationary solvers") do
    for T in (Float32, Float64, Complex64, Complex128)
        context("$T") do

        A=convert(Matrix{T}, randn(n,n))
        T<:Complex && (A+=convert(Matrix{T}, im*randn(n,n)))
        A+=A'
        A+=n*maximum(abs(A[:]))*diagm(ones(T,n)) #Diagonally dominant
        b=convert(Vector{T}, randn(n))
        T<:Complex && (b+=convert(Vector{T}, im*randn(n)))
        b=b/norm(b)
        x0=convert(Vector{T}, randn(n))
        T<:Complex && (x0+=convert(Vector{T}, im*randn(n)))
        x = A\b
        for solver in [jacobi, gauss_seidel]
            xi, ci=solver(A, b, maxiter=n^4)
            @fact ci.isconverged --> true
            @fact norm(x-xi) --> less_than(n^3*eps(typeof(real(b[1]))))
        end
        for solver in [jacobi!, gauss_seidel!]
            xi, ci=solver(copy(x0), A, b, maxiter=n^4)
            @fact ci.isconverged --> true
            @fact norm(x-xi) --> less_than(n^3*eps(typeof(real(b[1]))))
        end

        ω = 0.5
        for solver in [sor, ssor]
            xi, ci=solver(A, b, ω, maxiter=n^4)
            @fact ci.isconverged --> true
            @fact norm(x-xi) --> less_than(n^3*eps(typeof(real(b[1]))))
        end
        for solver in [sor!, ssor!]
            xi, ci=solver(copy(x0), A, b, ω, maxiter=n^4)
            @fact ci.isconverged --> true
            @fact norm(x-xi) --> less_than(n^3*eps(typeof(real(b[1]))))
        end

        end
    end
end

#Conjugate gradients
include("cg.jl")

#GMRES
facts("gmres") do

for T in (Float32, Float64, Complex64, Complex128)
    context("Matrix{$T}") do

    A = convert(Matrix{T}, randn(n,n))
    L = convert(Matrix{T}, randn(n,n))
    R = convert(Matrix{T}, randn(n,n))
    b = convert(Vector{T}, randn(n))
    if T <: Complex
        A += im*convert(Matrix{T}, randn(n,n))
        L += im*convert(Matrix{T}, randn(n,n))
        R += im*convert(Matrix{T}, randn(n,n))
        b += im*convert(Vector{T}, randn(n))
    end
    F = lufact(A)
    b = b/norm(b)

    x_gmres, c_gmres = gmres(A, b, L, R)
    @fact c_gmres.isconverged --> true
    @fact norm(A*x_gmres - b) --> less_than(√eps(real(one(T))))

    x_gmres, c_gmres = gmres(A, b, F; maxiter=1, restart=1)
    @fact c_gmres.isconverged --> true
    @fact norm(A*x_gmres - b) --> less_than(√eps(real(one(T))))

    x_gmres, c_gmres = gmres(A, b, 1, F; maxiter=1, restart=1)
    @fact c_gmres.isconverged --> true
    @fact norm(A*x_gmres - b) --> less_than(√eps(real(one(T))))
    end
end

for T in (Float64, Complex128)
    context("SparseMatrixCSC{$T}") do
    A = sprandn(n,n,0.5)+0.001*eye(T,n,n)
    L = sprandn(n,n,0.5)
    R = sprandn(n,n,0.5)
    b = convert(Vector{T}, randn(n))
    if T <: Complex
        A += im*sprandn(n,n,0.5)
        L += im*sprandn(n,n,0.5)
        R += im*sprandn(n,n,0.5)
        b += im*randn(n)
    end
    F = lufact(A)
    b = b / norm(b)

    x_gmres, c_gmres= gmres(A, b, L, R)
    @fact c_gmres.isconverged --> true
    @fact norm(A*x_gmres - b) --> less_than(√eps(real(one(T))))

    x_gmres, c_gmres = gmres(A, b, F; maxiter=1, restart=1)
    @fact c_gmres.isconverged --> true
    @fact norm(A*x_gmres - b) --> less_than(√eps(real(one(T))))

    x_gmres, c_gmres = gmres(A, b, 1, F; maxiter=1, restart=1)
    @fact c_gmres.isconverged --> true
    @fact norm(A*x_gmres - b) --> less_than(√eps(real(one(T))))
    end
end
end

#IDRs
facts("idrs") do

for T in (Float32, Float64, Complex64, Complex128)
    # without smoothing
    context("Matrix{$T}, without residual smoothing") do

    A = convert(Matrix{T}, randn(n,n))
    b = convert(Vector{T}, randn(n))
    if T <: Complex
        A += im*convert(Matrix{T}, randn(n,n))
        b += im*convert(Vector{T}, randn(n))
    end
    b = b/norm(b)

    x_idrs, c_idrs = idrs(A, b)
    @fact c_idrs.isconverged --> true
    @fact norm(A*x_idrs - b) --> less_than(√eps(real(one(T))))
    end

    # with smoothing
    context("Matrix{$T}, with residual smoothing") do

    A = convert(Matrix{T}, randn(n,n))
    b = convert(Vector{T}, randn(n))
    if T <: Complex
        A += im*convert(Matrix{T}, randn(n,n))
        b += im*convert(Vector{T}, randn(n))
    end
    b = b/norm(b)

    x_idrs, c_idrs = idrs(A, b; smoothing=true)
    @fact c_idrs.isconverged --> true
    @fact norm(A*x_idrs - b) --> less_than(√eps(real(one(T))))
    end
end

for T in (Float64, Complex128)
    context("SparseMatrixCSC{$T}") do
    A = sprandn(n,n,0.5)+0.001*eye(T,n,n)
    b = convert(Vector{T}, randn(n))
    if T <: Complex
        A += im*sprandn(n,n,0.5)
        b += im*randn(n)
    end
    b = b / norm(b)

    x_idrs, c_idrs= idrs(A, b)
    @fact c_idrs.isconverged --> true
    @fact norm(A*x_idrs - b) --> less_than(√eps(real(one(T))))
    end
end
end

facts("Chebyshev") do
for T in (Float32, Float64, Complex64, Complex128)
    context("Matrix{$T}") do
    A=convert(Matrix{T}, randn(n,n))
    T<:Complex && (A+=convert(Matrix{T}, im*randn(n,n)))
    A=A+A'
    A=A'*A #Construct SPD matrix
    b=convert(Vector{T}, randn(n))
    T<:Complex && (b+=convert(Vector{T}, im*randn(n)))
    b=b/norm(b)
    tol = 0.1 #For some reason Chebyshev is very slow
    v = eigvals(A)
    mxv = maximum(v)
    mnv = minimum(v)
    x_cheby, c_cheby= chebyshev(A, b, mxv+(mxv-mnv)/100, mnv-(mxv-mnv)/100, tol=tol, maxiter=10^5)
    @fact c_cheby.isconverged --> true
    @fact norm(A*x_cheby-b) --> less_than(tol)
    end
end
end

#######################
# Eigensystem solvers #
#######################

facts("simple eigensolvers") do
for T in (Float32, Float64, Complex64, Complex128)
    context("Matrix{$T}") do
    A=convert(Matrix{T}, randn(n,n))
    T<:Complex && (A+=convert(Matrix{T}, im*randn(n,n)))
    A=A+A' #Symmetric/Hermitian

    tol = (eltype(T) <: Complex ?2:1)*n^2*cond(A)*eps(real(one(T)))
    v = eigvals(A)

    ## Simple methods

    context("Power iteration") do
    eval_big = maximum(v) > abs(minimum(v)) ? maximum(v) : minimum(v)
    eval_pow = eigvals_power(A; tol=sqrt(eps(real(one(T)))), maxiter=2000)[1].val
    @fact norm(eval_big-eval_pow) --> less_than(tol)
    end

    context("Inverse iteration") do
    irnd = ceil(Int, rand()*(n-2))
    eval_rand = v[1+irnd] #Pick random eigenvalue
    # Perturb the eigenvalue by < 1/4 of the distance to the nearest eigenvalue
    eval_diff = min(abs(v[irnd]-eval_rand), abs(v[irnd+2]-eval_rand))
    σ = eval_rand + eval_diff/2*(rand()-.5)
    eval_ii = eigvals_ii(A, σ; tol=sqrt(eps(real(one(T)))), maxiter=2000)[1].val
    @fact norm(eval_rand-eval_ii) --> less_than(tol)
    end

    #context("Rayleigh quotient iteration") do
    #XXX broken?
    #l = eigvals_rqi(A, eigvals_rand, 2000, √eps()).val
    #@fact norm(eigvals_rand-l) --> less_than(tol)
    #end

    end
end
end

#Lanczos methods

type MyOp{T}
    buf::Matrix{T}
end

import Base: *, size, eltype

*(A::MyOp, x::AbstractVector) = A.buf*x
size(A::MyOp, i) = size(A.buf, i)
eltype(A::MyOp) = eltype(A.buf)

facts("eigvals_lanczos") do
for T in (Float32, Float64)
    context("Matrix{$T}") do
    A = convert(Matrix{T}, randn(n,n))
    A = A + A' #Symmetric
    v = eigvals(A)

    eval_lanczos, c_lanczos = eigvals_lanczos(A)
    @fact c_lanczos.isconverged --> true
    @fact norm(v - eval_lanczos) --> less_than(√eps(T))
    end

    context("Op{$T}") do
        A = MyOp(convert(Matrix{T}, randn(5,5)) |> t -> t + t')
        v = eigvals(Symmetric(A.buf))
        eval_lanczos, c_lanczos = eigvals_lanczos(A)
        @fact c_lanczos.isconverged --> true
        @fact norm(v - eval_lanczos) --> less_than(√eps(T))
    end
end
end

#Golub-Kahan-Lanczos singular values computation
include("svdl.jl")

include("lsqr.jl")
include("lsmr.jl")

#Randomized algorithms
include("rlinalg.jl")
include("rsvd.jl")
include("rsvd_fnkz.jl")

#Expensive tests - don't run by default
#include("matrixmarket.jl")
#include("matrixcollection.jl")

exitstatus()


"""

#### References

Smooth function approximation using neural networks.

S Ferrari and R F Stengel.

IEEE Trans Neural Netw, 2005 vol. 16 (1) pp. 24-38.

http://ieeexplore.ieee.org/lpdocs/epic03/wrapper.htm?arnumber=1388456
"""
mutable struct AlgebraicNetwork{TN} <: AbstractSLFN where TN <: Linear
    p::Int  # Number of training points
    q::Int  # Dimensionality of function domain
    s::Int  # number of neurons
    n_train_it::Int  # number of training iterations
    f::Float64
    maxit::Int
    neuron_type::TN
    μx::Vector{Float64}
    σx::Vector{Float64}
    Wt::Matrix{Float64}  # transpose of W matrix
    d::Vector{Float64}
    v::Vector{Float64}

    function AlgebraicNetwork(p::Int, q::Int, s::Int, neuron_type::TN,
                              f::Float64, maxit::Int, μx, σx) where TN
        Wt = Array{Float64}(undef, q, s)
        v = Array{Float64}(undef, s)
        d = Array{Float64}(undef, s)
        new{TN}(p, q, s, 0, f, maxit, neuron_type, μx, σx, Wt, d, v)
    end
end

function AlgebraicNetwork(x::AbstractArray, y::AbstractArray;
                          neuron_type::TN=Linear(Tanh()),
                          s::Int=size(x, 1), f::Float64=0.8,
                          maxit::Int=1000,
                          reg::AbstractLinReg=LSSVD()) where TN <: Linear
    p = size(x, 1)
    q = size(x, 2)

    @assert size(y, 1) == p "x and y must have same number of observations"

    s = min(s, p)
    xn, μx, σx = standardize(x[:, :])
    out = AlgebraicNetwork(p, q, s, neuron_type, f, maxit, μx, σx)
    fit!(out, xn, y, reg)
    out
end

## API methods
isexact(an::AlgebraicNetwork) = an.p == an.s
function fit!(an::AlgebraicNetwork, x::AbstractArray, y::AbstractVector,
              reg::AbstractLinReg)
    i = 0
    while true
        i += 1
        randn!(an.Wt)
        an.Wt ./= an.f
        # fill!(an.Wt, 5.0)
        an.d = -diag(x * an.Wt)
        S = hidden_out(an, x)

        if rank(S) < an.s && i < an.maxit
            an.n_train_it += 1
            continue
        else
            an.v = regress(reg, S, y)
            return an
        end
    end

end

function Base.show(io::IO, an::AlgebraicNetwork{TA}) where TA
    s =
    """
    AlgebraicNetwork with
      - $(TA) Activation function
      - $(an.q) input dimension(s)
      - $(an.s) neuron(s)
      - $(an.p) training point(s)
      - $(an.n_train_it) training iterations
    """
    print(io, s)
end


using LinearAlgebra
using DataStructures

const positions = map(line -> findall("#", line),
    readlines(length(ARGS) > 0 ? ARGS[1] : "example1.txt"))
const asteroids = Array{Vector,1}(undef,0)
for i in 1:length(positions)
    for j in 1:(length(positions[i]))
        push!(asteroids, Vector([positions[i][j][1] - 1, i - 1]))
    end
end

function direction(a, b)
    # Aha! Burned by a rounding error!
    v = [a[2] - b[2], b[1] - a[1]]
    normalized = v / norm(v)
    d = atand(normalized[2], normalized[1])
    # change to north-up azimuth increasing clockwise.
    #d = (450.0 - d) % 360.0
    d = (360 + d) % 360;
    return round(d, digits=3)
end

function directions(asteroids, a)
    return map(b -> direction(a, b), filter(b -> b != a, asteroids))
end

print("Day 10 Part 1: ");
v = findmax(map(a -> length(unique(directions(asteroids, a))), asteroids))
laser = asteroids[v[2]]
println("$(v[1]) are visible from $(laser)")

function print_field(a)
    xmin = findmin([x[1] for x in a])[1]
    xmax = findmax([x[1] for x in a])[1]
    ymin = findmin([y[1] for y in a])[1]
    ymax = findmax([y[1] for y in a])[1]
    for y in ymin:ymax
        for x in xmin:xmax
            print([x,y] in a ? '#' : '.')
        end
        println()
    end
end

#print_field(asteroids)

#for b in filter(x -> x != laser, asteroids)
#    println("angle from $(laser) to $(b) is $(direction(laser, b))")
#end

q = SortedDict()
for b in asteroids
    if b != laser
        d = direction(laser, b)
        if !(d in keys(q))
            q[d] = []
        end
        push!(q[d], b)
    end
end

# Well, that was ugly, but now we have a relation of angle -> [points].
# Next, we need to sort the points by their distance to the laser.
for target in values(q)
    sort!(target, by=(x -> norm(x - laser)))
end

const print_these = [1,2,3,10,20,50,100,199,200,201,299];
const th = Dict(1 => "st", 2 => "nd", 3 => "rd");

# Now spin the laser, remove the closest vertex each time, until we reach i == 200.
i = 0
while length(q) > 0
    for targets in keys(q)
        #println("   $(targets) => $(q[targets]): ");
        target = popfirst!(q[targets])
        if isempty(q[targets])
            delete!(q, targets)
        end
        global i += 1
        if i == 200
            global vaporize = target
        end
        if i in print_these
            println("The $(i)$(get(th, i, "th")) asteroid to be vaporized is at $(target[1]),$(target[2]).")
        end
    end
end

println("Day 10 Part 2: $(vaporize[1] * 100 + vaporize[2])")


function interaction_effect_probabilities(ch, data)
    γpost_x = collect(hcat(data.Xunique...)')
    γpost_mins = [minimum(γpost_x[:, i]) for i in 1:size(γpost_x, 2)]
    γpost_x = γpost_x - ones(size(γpost_x, 1)) * (γpost_mins .- 1)'

    ch_gamma_mat = collect(hcat(ch.gamma...)')
    nvar = size(γpost_x, 2)
    p_effects = zeros(nvar * (nvar - 1) ÷ 2)
    var1 = zeros(Int, nvar * (nvar - 1) ÷ 2)
    var2 = zeros(Int, nvar * (nvar - 1) ÷ 2)
    row = 1
    for i in 1:(nvar - 1), j = (i + 1):nvar
        subset = deepcopy(γpost_x)
        subset[:, i] = 2 .- (subset[:, i] .!= 1)
        subset[:, j] = 2 .- (subset[:, j] .!= 1)
        subset = maximum(subset, dims = 2)[:, 1] .== 1
        if sum(subset) > 0
            p_effects[row] = mean(maximum(ch_gamma_mat[:, subset], dims = 2))
            var1[row] = i
            var2[row] = j
            row += 1
        end
    end
    return DataFrame(
        var1 = "x" .* string.(var1),
        var2 = "x" .* string.(var2),
        prob = p_effects
    )[1:(row - 1), :]
end


using Stheno, Plots
plotly();

############################## Define model and data ################################

# Define additive model.
σ_noise, gpc = 1e-1, GPC();
f₁ = GP(pick_dims(EQ(), 1), gpc);
f₂ = GP(pick_dims(EQ(), 2), gpc);
f₃ = f₁ + f₂;
y = f₃ + GP(Noise(σ_noise^2), gpc);

# Sample from marginal process to generate toy data.
rng, S = MersenneTwister(123456), 25;
X = ColsAreObs(randn(2, 100));
ŷ = rand(rng, y(X));

############################### Set up plotting ##################################

# Define plotting locations.
P, Q, pos, b = 50, 50, 1, 4.0;
Ps, Qs, Xp = linspace(-b, b, P), linspace(-b, b, Q), Matrix{Float64}(2, P * Q);
for p in 1:P, q in 1:Q
    Xp[1, pos] = Ps[p]
    Xp[2, pos] = Qs[q]
    pos += 1
end
Xp = ColsAreObs(Xp);

exact_plot, approx_plot = plot(), plot();

################################ Do exact inference ################################

@show logpdf(y(X), ŷ);
f₁′, f₂′, f₃′ = (f₁, f₂, f₃) | (y(X) ← ŷ);
f₁′Xp, f₂′Xp, f₃′Xp = rand(rng, [f₁′(Xp), f₂′(Xp), f₃′(Xp)], S);
μ₁′, σ₁′ = marginals(f₁′(Xp));
μ₂′, σ₂′ = marginals(f₂′(Xp));
μ₃′, σ₃′ = marginals(f₃′(Xp));

################################ Do approximate inference ###############################

M = 15;
Z₁ = ColsAreObs(hcat(linspace(-b, b, M), zeros(M))');
Z₂ = ColsAreObs(hcat(zeros(M), linspace(-b, b, M))');
u₁, u₂ = f₁(Z₁), f₂(Z₂);
u = BlockGP([u₁, u₂]);
@show elbo(f₃(X), ŷ, u, σ_noise);
μ′u, Σ′u = Stheno.optimal_q(f₃(X), ŷ, u, σ_noise);
conditioner = Stheno.Titsias(u, μ′u, Σ′u, gpc);
f₁′u, f₂′u, f₃′u = (f₁, f₂, f₃) | conditioner;


using GalacticOptim, Optim, Test, ReverseDiff, Tracker

x0 = zeros(2)
rosenbrock(x, p=nothing) =  (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
l1 = rosenbrock(x0)

function g!(G, x)
    G[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
    G[2] = 200.0 * (x[2] - x[1]^2)
end

function h!(H, x)
    H[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
    H[1, 2] = -400.0 * x[1]
    H[2, 1] = -400.0 * x[1]
    H[2, 2] = 200.0
end

G1 = Array{Float64}(undef,2)
G2 = Array{Float64}(undef,2)
H1 = Array{Float64}(undef, 2, 2)
H2 = Array{Float64}(undef, 2, 2)

g!(G1, x0)
h!(H1, x0)

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff())
optprob = GalacticOptim.instantiate_function(optf,x0,GalacticOptim.AutoForwardDiff(),nothing)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, BFGS())
@test 10*sol.minimum < l1

sol = solve(prob, Newton())
@test 10*sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10*sol.minimum < l1

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoZygote())
optprob = GalacticOptim.instantiate_function(optf,x0,GalacticOptim.AutoZygote(),nothing)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, BFGS())
@test 10*sol.minimum < l1

sol = solve(prob, Newton())
@test 10*sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10*sol.minimum < l1

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoReverseDiff())
optprob = GalacticOptim.instantiate_function(optf,x0,GalacticOptim.AutoReverseDiff(),nothing)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2

prob = OptimizationProblem(optprob, x0)
sol = solve(prob, BFGS())
@test 10*sol.minimum < l1

sol = solve(prob, Newton())
@test 10*sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10*sol.minimum < l1

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoTracker())
optprob = GalacticOptim.instantiate_function(optf,x0,GalacticOptim.AutoTracker(),nothing)
optprob.grad(G2, x0)
@test G1 == G2
@test_throws ErrorException optprob.hess(H2, x0)

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, BFGS())
@test 10*sol.minimum < l1

@test_throws ErrorException solve(prob, Newton())

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoFiniteDiff())
optprob = GalacticOptim.instantiate_function(optf,x0,GalacticOptim.AutoFiniteDiff(),nothing)
optprob.grad(G2, x0)
@test G1 ≈ G2 rtol=1e-6
optprob.hess(H2, x0)
@test H1 ≈ H2 rtol=1e-6

prob = OptimizationProblem(optprob, x0)
sol = solve(prob, BFGS())
@test 10*sol.minimum < l1

sol = solve(prob, Newton())
@test 10*sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test sol.minimum < l1 #the loss doesn't go below 5e-1 here