using Statistics

@testset "Velocity timeseries" begin

    # test object: onbeat notes with velocity ascending from 0 to 159
    notes = Notes{MoreVelNote}(Vector{MoreVelNote}(), 960)
    for i = 0:159
        push!(notes.notes, MoreVelNote(1, i, 960*i, 2))
    end

    #simple time series
    tvec, tseries = timeseries(notes, :velocity, maximum, 0:1)
    @test tvec == collect(0:960:159*960)
    @test tseries == collect(0:159)

    # time series for eights with zeros
    grid = 0:1//2:1
    tvec, tseries = timeseries(notes, :velocity, maximum, grid)

    @test tvec == collect(0:480:159*960+480)
    @test [tseries[i] for i in 1:2:159*2+1] == collect(Float64, 0:159)
    @test [tseries[i] for i in 2:2:159*2+2] == zeros(160)

    # add notes which have been quantized to same bin
    for i = 0:159
        push!(notes.notes, MoreVelNote(1, i+1, 960*i, 2))
    end
    sort!(notes.notes, lt=((x, y)->x.position<y.position))

    # simple time series choosing maximum in bin
    tvec, tseries = timeseries(notes, :velocity, maximum, 0:1)

    @test tvec == collect(Float64, 0:960:159*960)
    @test tseries == collect(Float64, 1:160)

    # simple time series choosing mean value in bin
    tvec, tseries = timeseries(notes, :velocity, mean, 0:1)

    @test tvec == collect(Float64, 0:960:159*960)
    @test tseries == collect(Float64, 1:160) .- 0.5

    # time series for eights with zeros and average for bins
    grid = 0:1//2:1
    tvec, tseries = timeseries(notes, :velocity, mean, grid)

    @test tvec == collect(Float64, 0:480:159*960+480)
    @test [tseries[i] for i in 1:2:159*2+1] == collect(Float64, 0:159) .+ 0.5
    @test [tseries[i] for i in 2:2:159*2+2] == zeros(Int,160)

end

@testset "Pitch timeseries" begin
    # test object: onbeat (quarter) notes with pitch ascending
    notes = Notes{Note}(Vector{Note}(), 960)
    vals = 10:20
    L = length(vals)
    for (i, val) in enumerate(vals)
        push!(notes.notes, Note(val, 5, 960*(i-1), 2))
    end

    #simple time series
    tvec, tseries = timeseries(notes, :pitch, maximum, 0:1)
    @test tvec == collect(0:960:(L-1)*960)
    @test tseries == collect(vals)

    # add notes which have been quantized to same bin
    for  (i, val) in enumerate(vals)
        push!(notes.notes, Note(5, 5, 960*(i-1), 2))
    end
    sort!(notes.notes, lt=((x, y)->x.position<y.position))

    # time series for eights with zeros and average for bins
    grid = 0:1//2:1
    tvec, tseries = timeseries(notes, :pitch, mean, grid)

    @test tvec == collect(0:480:(L)*960-480)
    @test [tseries[i] for i in 1:2:(L-1)*2+1] == (collect(Float64, vals) .+ 5)./2
    @test [tseries[i] for i in 2:2:(L-1)*2+2] == zeros(Int,L)
end


using ExploreFunction, Test

dataset = [ 0 0;
            0 1;
            0 3]

true_distance = [   0 1 3;
                    1 0 2;
                    3 2 0]

@testset "euclideandistance" begin
    @test euclideandistance(dataset) == true_distance
end

true_prune = [  0 3;
                3 0]
true_remaining = [1 3]

@testset "prunedataset" begin
    prune_dist, remaining = prunedataset(true_distance,1.5)
    @test prune_dist == true_prune
    @test true_remaining == true_remaining
end

dataset_ext = vcat(dataset, [0 6])
true_distance_ext = [   0 1 3 6;
                        1 0 2 5;
                        3 2 0 3;
                        6 5 3 0]
@testset "H0Barcode" begin
    barcode = H0Barcode(true_distance_ext)
    @test barcode[2,:] == [1, 2, 3]
end

@testset "lowerneighbors" begin
    adjacency_test = zeros(4,4)
    adjacency_test[1,2] = adjacency_test[2,1] = 1
    adjacency_test[2,3] = adjacency_test[3,2] = 1
    adjacency_test[3,4] = adjacency_test[4,3] = 1
    adjacency_test[2,4] = adjacency_test[4,2] = 1

    n1 = ExploreFunction.lowerneighbors(1,adjacency_test)
    @test n1 == []

    n2 = ExploreFunction.lowerneighbors(2,adjacency_test)
    @test n2 == [1]

    n3 = ExploreFunction.lowerneighbors(3,adjacency_test)
    @test n3 == [2]

    n4 = ExploreFunction.lowerneighbors(4,adjacency_test)
    @test n4 == [2,3]
end

@testset "edgelist" begin
    adjacency_test = zeros(4,4)
    adjacency_test[1,2] = adjacency_test[2,1] = 1
    adjacency_test[2,3] = adjacency_test[3,2] = 1
    adjacency_test[3,4] = adjacency_test[4,3] = 1
    adjacency_test[2,4] = adjacency_test[4,2] = 2

    edges, weights = ExploreFunction.edgelist(adjacency_test)
    true_edges = [  1 2 2 3;
                    2 3 4 4]
    @test true_edges == edges
    @test weights == [1, 1, 2, 1]
end

@testset "neighborintersect" begin
    adjacency_test = zeros(4,4)
    adjacency_test[1,2] = adjacency_test[2,1] = 1
    adjacency_test[2,3] = adjacency_test[3,2] = 1
    adjacency_test[3,4] = adjacency_test[4,3] = 1
    adjacency_test[2,4] = adjacency_test[4,2] = 1

    s1 = [3, 4]
    i1 = ExploreFunction.neighborintersect(s1,adjacency_test)
    @test [2] == i1

    s2 = [2, 3]
    i2 = ExploreFunction.neighborintersect(s2,adjacency_test)
    @test [] == i2
end

@testset "inductiveVR" begin
    adjacency_test = zeros(4,4)
    adjacency_test[1,2] = adjacency_test[2,1] = 1
    adjacency_test[2,3] = adjacency_test[3,2] = 1
    adjacency_test[3,4] = adjacency_test[4,3] = 2
    adjacency_test[2,4] = adjacency_test[4,2] = 1

    VR, weights = inductiveVR(adjacency_test,2)
    # Test the 1-simplices
    true_s1 = [ 1 2 2 3;
                2 3 4 4]
    @test VR[1] == true_s1

    # Test the 2-simplices
    true_s2 = [2, 3, 4]
    @test VR[2][:] == true_s2

    # Test the 3-simplices
    @test VR[3][:] == []

    # Test Weights
    @test weights[1] == [1, 1, 1, 2]
    @test weights[2] == [2]
end


function LJForceMagnitude(r::T,d::T) where T<:Real
    if r > d || r < 0.8*d
        return zero(T)
    else
        return 2*d/r^2*(1 - d/r)
    end
end

function ForceCalculation(d::T,Δx::T,Δy::T) where T<:Real
    r = sqrt(Δx^2 + Δy^2)
    tx = Δx/r
    ty = Δy/r

    LJmag = LJForceMagnitude(r,d)

    Fx = tx*LJmag
    Fy = ty*LJmag

    return Fx,Fy
end

function compute_body_repulsion(ind::Int,olist)
    FX = 0.0
    FY = 0.0

    for ti=1:(ind-1)
        Δx = olist[ind].center.x - olist[ti].center.x
        Δy = olist[ind].center.y - olist[ti].center.y
        d = olist[ti].radius + olist[ind].radius +
            (1 + olist[ti].buffer_percent + olist[ind].buffer_percent)/100

        fx,fy = ForceCalculation(d,Δx,Δy)
        FX += fx; FY += fy
    end

    return FX, FY
end

function compute_wall_repulsion(ind,olist,wlist)
    FX = 0.0
    FY = 0.0

    for tw=1:length(wlist)
        pointOnWall = get_nearest_point(olist[ind].center,wlist[tw])
        if is_point_in_wall(pointOnWall,wlist[tw])
            Δx = pointOnWall.x - olist[ind].center.x
            Δy = pointOnWall.y - olist[ind].center.y

            d = olist[ind].radius + wlist[tw].thickness/2

            fx,fy = ForceCalculation(d,Δx,Δy)
            FX += fx
            FY += fy
        end
    end

    return FX, FY
end

function shuffle_object!(ind::Int,olist,wlist)
    fXr,fYr = compute_body_repulsion(ind,olist)
    fXw,fYw = compute_wall_repulsion(ind,olist,wlist)
    olist[ind].center.x += (fXr + fXw)*1e-6
    olist[ind].center.y += (fYr + fYw)*1e-6
    nothing
end


@testset "Fitting data" begin
	x = collect(1.:1:10*365)./365;
	y = 1869.9 .- 782.0 .* exp.(-0.085 .* x);
	par,er = fitexp(x,y);
	@test round(sum(par*10)) ≈ round(18699-7820-0.85)
	@test sum(er) < 1e-2

	@test evalexp([3.],[10.,0.5,0.05]) ≈ [10. + 0.5*exp(0.05*3.)]

	x = collect(1.:1:365);
	y = 10.0 .+ 0.1 .* x + rand(length(x))./20;
	par,er = fitpoly(x,y,deg=1);
	@test par[1] ≈ 0.1 atol=0.001
	@test par[2] ≈ 10. atol=0.1
	@test sum(er) < 1e-2

	@test evalpoly([10.],[0.01, 0.1, 1.0]) ≈ [1+0.1*10+0.01*10*10]

	@test evalpoly([DateTime(1,1,1,0,1)],[0.01, 0.1, 1.0]) ≈ [1+0.1*86460000+0.01*86460000.0^2]

	## Test model creation
	m = ResampleAndFit.getpolymodel(0)
	@test m([10.,20.],2.0) == [2.0, 2.0]
	m = ResampleAndFit.getpolymodel(3)
	@test m(1,[4.,3.,2.,1]) == sum([4.,3.,2.,1])
	m = ResampleAndFit.getpolymodel(4)
	@test m(1,[5.,4.,3.,2.,1]) == sum([5.,4.,3.,2.,1])
	@test isnan(ResampleAndFit.getpolymodel(10))
end


using PyPlot, transD_GP, Random, LinearAlgebra, Statistics

## model
Random.seed!(3)
nlayers, dz = 50, 5.
ρbg        = 1000
zfixed     = [-1e5,   (0:dz:(nlayers-1)*dz )...]
rho        = [1e12,   10 .^(log10(ρbg) .+ randn(nlayers))...]
ianom      = 5
rho[ianom] = 10.
δ          = 1e-4  # fd step size in log10

##  geometry and modeling parameters
rRx = 13.
zRx = -42.0
zTx = -40.0
freqlow = 1e-3
include("electronics_halt.jl")
calcjacobian = true

Fhm = transD_GP.AEM_VMD_HMD.HFieldDHT(
                      lowpassfcs = lowpassfcs,
                      nmax   = length(zfixed),
                      times  = HM_times,
                      ramp   = HM_ramp,
                      zTx    = zTx,
                      rRx    = rRx,
                      rTx    = rTx,
                      zRx    = zRx,
                      freqlow = freqlow,
                      calcjacobian = calcjacobian,
                      )
transD_GP.AEM_VMD_HMD.getfieldTD!(Fhm, zfixed, rho)

## plot HM
f, ax = plt.subplots(1,3, sharey=true, figsize=(8,6))
delz, delt = diff(zfixed[2:end]), diff(log10.(Fhm.times))
tlabel = [(log10.(Fhm.times) .- delt[1]/2)..., log10.(Fhm.times)[end]+delt[end]/2]
zlab = [zfixed[2:end]..., zfixed[end]+delz[end]]
ax[1].step(log10.(rho[2:end]), zfixed[2:end])
im2 = ax[2].pcolormesh(tlabel, zlab, Fhm.dBzdt_J[2:end,:], shading="auto")
cb2 = colorbar(im2, ax=ax[2])
im3 = ax[3].pcolormesh(tlabel, zlab, log10.(abs.(Fhm.dBzdt_J[2:end,:])))
cb2 = colorbar(im3, ax=ax[3])
ax[1].set_xlabel("log10 Ω-m")
ax[1].invert_xaxis()
ax[2].invert_yaxis()
ax[2].set_xlabel("log10 s")
ax[3].set_xlabel("log10 s")
ax[1].set_ylabel("Depth m")
ax[2].set_title("∂f/∂(log10σ)")
ax[3].set_title("∂f/∂(log10σ)")

## approximate partials with central differences and compare
J = copy(Fhm.dBzdt_J)

function getJ_TD!(Fin, zin, ρ, Δ)
    # sensitivities computed in log10σ
    # same as in AEM_VMD_HMD code
    Jout = similar(Fin.dBzdt_J)
    modelfd = copy(ρ)
    jflag = Fin.calcjacobian
    Fin.calcjacobian = false
    for i = 2:length(zin)
        modelfd[:] = ρ
        modelfd[i] = 10 .^-(-log10(ρ[i]) .+ Δ/2)
        transD_GP.AEM_VMD_HMD.getfieldTD!(Fin, zin, modelfd)
        HTD_z1 = copy(Fin.dBzdt)
        modelfd[i] = 10 .^-(-log10(ρ[i]) .- Δ/2)
        transD_GP.AEM_VMD_HMD.getfieldTD!(Fin, zin, modelfd)
        Jout[i,:] = (HTD_z1 - Fin.dBzdt) / Δ
    end
    Fin.calcjacobian = jflag
    Jout
end
J_fd = getJ_TD!(Fhm, zfixed, rho, δ)

## comparison plots
f, ax = plt.subplots(1,3, sharey=true, figsize=(10,6))
ax[1].step(log10.(rho[2:end]), zfixed[2:end])
im2 = ax[2].pcolormesh(tlabel, zlab, J[2:end,:])
cb2 = colorbar(im2, ax=ax[2])
im3 = ax[3].pcolormesh(tlabel, zlab, J_fd[2:end,:])
cb2 = colorbar(im3, ax=ax[3])
ax[1].set_xlabel("log10 Ω-m")
ax[1].invert_xaxis()
ax[2].invert_yaxis()
ax[2].set_xlabel("log10 s")
ax[3].set_xlabel("log10 s")
ax[1].set_ylabel("Depth m")
ax[2].set_title("∂f/∂(log10σ) analytic")
ax[3].set_title("∂f/∂(log10σ) FD")

## more detailed
figure()
nlayer = 40
semilogx(HM_times, J_fd[nlayer,:], "--k")
semilogx(HM_times, J[nlayer,:], "-")

## Benchmark timings
using BenchmarkTools
@info "analytic timings"
@btime transD_GP.AEM_VMD_HMD.getfieldTD!($Fhm, $zfixed, $rho)
@info "Cent Diff timing"
@btime J_fd = getJ_TD!($Fhm, $zfixed, $rho, $δ);


## ExoplanetsSysSim/src/transit_detection_model.jl
## (c) 2015 Eric B. Ford

# Several functions below based on https://github.com/christopherburke/KeplerPORTs/blob/master/KeplerPORTs_utils.py
# That follows the procedure outlined in Burke et al.(2015).
# However we don't currently interpolate the CDDP or mesthreshold to the relevant duration

function real_log_choose(m::Real, n::Real)
  lgamma(m+1)-lgamma(n+1)-lgamma(m-n+1.0)
end

function real_binom(k::Real, BigM::Real, f::Real)
    F1 = real_log_choose(BigM,k)
    F2 = k*log(f)
    F3 = (BigM-k)*log(1.0-f)
    x = exp(F1+F2+F3)
    return x
end

function kepler_window_function(num_transits_no_gaps::Real, duty_cycle::Real; min_transits::Real = 3.0)
  if num_transits_no_gaps < min_transits
     return 0.0
  else
     return max(1.0 - real_binom(min_transits,num_transits_no_gaps,duty_cycle), 0.0)
  end
end

# See Christiansen et al. (2015)  This assumes a linear limbdarkening coefficient of 0.6
function frac_depth_to_tps_depth(frac_depth::Real)
    const alp = 1.0874
    const bet = 1.0187
    const REALDEPTH2TPSSQUARE = 1.0  # WARNING: Waiting for this to be confirmed
    k = sqrt(frac_depth)
    tps_depth = min( (alp-bet*k) * frac_depth* REALDEPTH2TPSSQUARE, 1.0)   # NOTE: I added the max based on common sense
    return tps_depth
end

function detection_efficiency_theory(mes::Real; min_pdet_nonzero::Float64 = 0.0)
   const muoffset =  0.0
   const sig =  1.0
   const mesthresh = 7.1
   if mes > (9.0 - mesthresh - muoffset)
      return 1.0
   else
      pdet = 0.5 + 0.5*erf((mes - mesthresh - muoffset) / sqrt(2.0*sig*sig))
      pdet = pdet >= min_pdet_nonzero ? pdet : 0.0
      return pdet
   end
end

function detection_efficiency_fressin2013(mes::Real)
   const mesmin =  6.0
   const mesmax =  16.0
   if mes <= mesmin
      return 0.0
   elseif mes >= mesmax
      return 1.0
   else
     return (mes - mesmin) / (mesmax - mesmin)
  end
end

function detection_efficiency_christiansen2015(mes::Real; mes_threshold::Real = 7.1, min_pdet_nonzero::Float64 = 0.0)
   const a =  4.65  # from code for detection_efficiency(...) at https://github.com/christopherburke/KeplerPORTs/blob/master/KeplerPORTs_utils.py
   const b =  0.98
   #const a =  4.35 # from arxiv abstract. TODO DETAIL: Figure out which is better.  Informal testing showed it didn't matter
   #const b =  1.05
   usemes = max(0.0,mes - 7.1 - (mes_threshold - 7.1))
   pdet = cdf(Gamma(a,b), usemes)
   pdet = pdet >= min_pdet_nonzero ? pdet : 0.0
   return pdet
end

function detection_efficiency_dr25_simple(mes::Real)
   const a =  30.87  # from pg 16 of https://exoplanetarchive.ipac.caltech.edu/docs/KSCI-19110-001.pdf
   const b =  0.271
   const c = 0.940
   pdet = c*cdf(Gamma(a,b), mes)
   return pdet
end

# Resume code original to SysSim

function calc_snr_if_transit(t::KeplerTarget, depth::Real, duration::Real, sim_param::SimParam; num_transit::Real = 1)
  cdpp = t.cdpp[1,1]                                                            # TODO: Make CDPP lookup based on season/quarter/month and timescale.  Is CDPP timescale IMPORTANT or DETAIL?
  depth_tps = frac_depth_to_tps_depth(depth)                                        # WARNING: Hardcoded this conversion
  snr = depth_tps*sqrt(num_transit*duration*LC_rate)/cdpp              # WARNING: Assumes measurement uncertainties are uncorrelated & CDPP based on LC
end

function calc_snr_if_transit(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  depth = calc_transit_depth(t,s,p)
  duration = calc_transit_duration(t,s,p)
  num_transit = calc_expected_num_transits(t,s,p,sim_param)
  calc_snr_if_transit(t,depth,duration,sim_param,num_transit=num_transit)
end

function calc_prob_detect_if_transit(t::KeplerTarget, snr::Real, sim_param::SimParam; num_transit::Real = 1)
  const min_transits = 3.0                                                    # WARNING: Hard coded 3 transit minimum
  const mes_threshold = 7.1                                                   # WARNING: Assuming 7.1 for all stars, durations
  const min_pdet_nonzero = 0.0                                                # TODO OPT: Figure out how to prevent a plethora of planets that are very unlikely to be detected due to using 0.0
  wf = kepler_window_function(num_transit, t.duty_cycle, min_transits=min_transits)   # TODO SCI DETAIL: Replace statistical model with checking actual transit times for long period planets
  return wf*detection_efficiency_christiansen2015(snr, mes_threshold=mes_threshold, min_pdet_nonzero=min_pdet_nonzero)        # WARNING: I've hardcoded this choice, since this is called so frequently
  #return wf*detection_efficiency_theory(snr, min_pdet_nonzero=min_pdet_nonzero)	     # WARNING: Hardcoded version of ideal detection efficiency
end

function calc_prob_detect_if_transit(t::KeplerTarget, depth::Real, duration::Real, sim_param::SimParam; num_transit::Real = 1)
  snr = calc_snr_if_transit(t,depth,duration,sim_param, num_transit=num_transit)
  return calc_prob_detect_if_transit(t, snr, sim_param, num_transit=num_transit)
end

function calc_prob_detect_if_transit(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  depth = calc_transit_depth(t,s,p)
  duration = calc_transit_duration(t,s,p)
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_prob_detect_if_transit(t,depth,duration, sim_param, num_transit=ntr)
end

# Compute probability of detection if we average over impact parameters b~U[0,1)
function calc_ave_prob_detect_if_transit(t::KeplerTarget, snr_central::Real, sim_param::SimParam; num_transit::Real = 1)
  const min_transits = 3.0                                                    # WARNING: Hard coded 3 transit minimum
  const mes_threshold = 7.1                                                   # WARNING: Assuming 7.1 for all stars, durations
  const min_pdet_nonzero = 0.0                                                # TODO OPT: Figure out how to prevent a plethora of planets that are very unlikely to be detected due to using 0.0
  wf = kepler_window_function(num_transit, t.duty_cycle, min_transits=min_transits)
  num_impact_param = 5
  ave_detection_efficiency = 0.5*detection_efficiency_christiansen2015(snr_central, mes_threshold=mes_threshold, min_pdet_nonzero=min_pdet_nonzero)    # WARNING: Hardcoded Gamma CDF detection efficiency
  #ave_detection_efficiency = 0.5*detection_efficiency_theory(snr_central, min_pdet_nonzero=min_pdet_nonzero)	  # WARNING: Hardcoded ideal detection efficiency
  for i in 1:(num_impact_param-1)
    b = i/num_impact_param
    snr = snr_central*sqrt(sqrt((1-b)*(1+b)))
    ave_detection_efficiency += detection_efficiency_christiansen2015(snr, mes_threshold=mes_threshold, min_pdet_nonzero=min_pdet_nonzero)   # WARNING: Hardcoded Gamma CDF detection efficiency
    #ave_detection_efficiency += detection_efficiency_theory(snr, min_pdet_nonzero=min_pdet_nonzero)   # WARNING: Hardcoded ideal detection efficiency
  end
  ave_detection_efficiency /= num_impact_param
  return wf*ave_detection_efficiency
end

function calc_ave_prob_detect_if_transit(t::KeplerTarget, depth::Real, duration_central::Real, sim_param::SimParam; num_transit::Real = 1)
  snr_central = calc_snr_if_transit(t,depth,duration_central,sim_param, num_transit=num_transit)
  return calc_ave_prob_detect_if_transit(t, snr_central, sim_param, num_transit=num_transit)
end

function calc_ave_prob_detect_if_transit(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  depth = calc_transit_depth(t,s,p)
  duration_central = calc_transit_duration_central(t,s,p)
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_ave_prob_detect_if_transit(t,depth,duration_central, sim_param, num_transit=ntr)
end


#
# Implementation of LDpred Gibbs sampler
#

#
# true effects μ distributed as
#
#        N(0, Nh^2/Mp)   with probability p
# μ_i ~
#        0 with probability 1-p
#
# define σ2 := Nh^2/M
# Nh²/M can be obtained from LD-score regression
#
# σ2_μ = σ2/p

#
# observed (marginal) z-scores distributed as
#
# z_i|μ_i = N(μ_i, 1)
#
# σ2_z = 1 + σ2_μ
# σ_z = sqrt(1 + σ2_μ)

#
# residualized marginal effect
# i.e. correct SNP for effect of all other SNPs in LD
#
# z̃ = ẑ - (Dμ - μ)
#

#
# probability that variant j is causal
# is probability that z̃_i ~ N(0, 1 + Nh^2/Mp) vs z̃_i ~  N(0, 1)
#
#
# p_j = 1 / (1 + (1-p)/p * sqrt(1+σ2_μ) * exp(-0.5 * z̃^2 * σ2_μ / (1 + σ2_μ)))
#

#
# Posterior for μ_j is then
#
#             N(σ2_μ/(1 + σ2_μ) z̃_j, σ2_μ / (1 + σ2_μ))  with probability p_j
# μ_j | z̃_j ~
#             0                                      with probability 1 - p_j
#

#
# posterior mean of μ_j is then
#
# ω_j = p_j * z̃_j * σ^2 / (1 + σ^2)
#

#
# Gibbs sampler
#
# for iter = 1:(n_burnin + n_iter)
#   for j = 1:n_variants
#      compute z̃_j
#      compute p_j
#      sample μ_j
#      compute ω_j
#      if iter > n_burnin
#        Ω_j += ω_j
#      end
#   end
# end
# μ = Ω / n_iter
#

using StatsBase

export ldpred_gibbs,
    ldpred_gibbs_EB,
    ldpred_gibbs_acc

"""
    ldpred_gibbs(z, D, p, σ2, μ0; n_burnin = 100, n_iter = 500, verbose = false)

Run LDpred Gibbs sampler.

# Arguments

- `z::Vector`: z-scores
- `D::LDMatrix`: LD matrix
- `p::Real` : proportion of variants deemed to be causal
- `σ2::Real`: ``Nh²/M`` as estimated from LDscore regression (`estimate_h2`, or `mean(z^2)/mean(lds)`)
              Note: prior variance of non-null component of ``μ``, ``σ2_μ = σ2/p``
- `μ0::Vector`: starting estimate (e.g. from infinite model)
"""
function ldpred_gibbs(z, D, p, σ2, μ0; n_burnin = 100, n_iter = 500, verbose = false)

    n_variants = length(z)
    μ = μ0[:]
    Ω = zeros(Float64, n_variants)

    σ2_μ = σ2 / p
    σ_z = sqrt(1.0 + σ2_μ)
    η = σ2_μ / (1.0 + σ2_μ)  # shrinkage factor

    for iter = 1:(n_burnin + n_iter)
        verbose && println("Iteration $iter/$(n_burnin + n_iter)")
        for j = 1:n_variants
            z̃_j = z[j] - μ'*D[:,j] + μ[j]
            p_j = 1.0 / (1.0 + (1.0-p)/p * σ_z * exp(-0.5 * z̃_j * z̃_j * η))
            μ[j] = (rand() <= p_j) ?  rand(Normal(η * z̃_j, sqrt(η))) : 0.0
            if iter > n_burnin
                Ω[j] += p_j * z̃_j * η
            end
        end
    end

    res = Ω ./ n_iter
    res[(!isfinite).(res)] .= 0.0 # DON'T DO THIS!!! THIS HIDES ERRORS
    res
end

# version that stores the simulated μ for each iteration
function ldpred_gibbs_acc(z, D, p, σ2, μ0; n_burnin = 100, n_iter = 500, verbose = false)

    n_variants = length(z)
    μ = μ0[:]
    Ω = zeros(Float64, n_variants)

    σ2_μ = σ2 / p
    σ_z = sqrt(1.0 + σ2_μ)
    η = σ2_μ / (1.0 + σ2_μ)  # shrinkage factor

    res = []

    for iter = 1:(n_burnin + n_iter)
        verbose && println("Iteration $iter/$(n_burnin + n_iter)")
        for j = 1:n_variants
            z̃_j = z[j] - μ'*D[:,j] + μ[j]
            p_j = 1.0 / (1.0 + (1.0-p)/p * σ_z * exp(-0.5 * z̃_j * z̃_j * η))
            μ[j] = (rand() <= p_j) ?  rand(Normal(η * z̃_j, sqrt(η))) : 0.0
        end

        if iter > n_burnin
            push!(res, copy(μ))
        end
    end

    res
end

#
# Empirical Bayes extension
#

# instead of using same prior for all μ_i, use prior as estimated from EB
#

function ldpred_gibbs_EB(z, D, π, σ_z, μ0, scale = 1.0; n_burnin = 100, n_iter = 500, verbose = false)

    K = length(π)
    n_variants = length(z)
    μ = μ0[:]
    Ω = zeros(Float64, n_variants)

    σ2_μ = max.(σ_z .^ 2 .- 1, 0.0)/scale
    σ_μ = sqrt.(σ2_μ)
    σ_z_scaled = sqrt.(1.0 .+ σ2_μ)
    η = σ2_μ ./ (1.0 .+ σ2_μ) # η[1,k] : shrinkage factor for component k

    for iter = 1:(n_burnin + n_iter)
        verbose && println("Iteration $iter/$(n_burnin + n_iter)")
        for j = 1:n_variants
            z̃_j = z[j] - μ'*D[:,j] + μ[j]
            p_j = estimate_component_prob([z̃_j], π, σ_z_scaled)

            k = sample(1:K, Weights(p_j[:]))
            μ[j] = rand(Normal(η[k] * z̃_j, sqrt(η[k])))
            #println("$j\t$(z[j])\t$(z̃_j)\t$(p_j)\t$(μ[j])")
            #if !isfinite(p_j[1])
            #    break
            #end
            if iter > n_burnin
                Ω[j] += sum(p_j .* η) * z̃_j
            end
        end
    end

    res = Ω ./ n_iter
    res[(!isfinite).(res)] .= 0.0 # DON'T DO THIS!!! THIS HIDES ERRORS
    res
end


A = sprand(100, 100, 0.2)
Am = Matrix(A)
l = QuadLoss()
rx = QuadReg()
ry = QuadReg()
obs = [(f[1], f[2]) for f in zip(findnz(A)[1], findnz(A)[2])]

gg = GGLRM(A, l, rx, ry, 10, obs=obs)
ggs = GGLRM(A, l, rx, ry, 10, obs=obs)

fit!(gg); fit_sparse!(ggs);
ggobj = whole_objective(gg, gg.X'gg.Y)
ggsobj = whole_objective(ggs, ggs.X'ggs.Y)
@test_approx_eq_eps(ggobj, ggsobj, 0.1*length(A))


include(joinpath(@__DIR__, "../../src/FLS.jl"))

"""
Permutations of a set, i.e. invertible endo-functions.
"""
permschema = @acset LabeledGraph begin
    V = 1
    E = 2
    vlabel = [:X]
    elabel = [:f, :f⁻¹]
    src = [1,1]
    tgt = [1,1]
end

perm = FLS(:perm, permschema, [], [(:inv, [:f, :f⁻¹], [])])


using RecyclePermanent
using MonteCarloMarkovKernels: simulateChainProgress, makeAMKernel, estimateBM,
  kde
using StaticArrays
using StatsBase
using LinearAlgebra
using JLD2
using Plots
Plots.gr()

include("../models/gaussianModel.jl")

@load "data/gaussianData.jld" ys θ
ys = ys[1:10]

d = 2

@inline function toGaussianθ(v::SVector{d, Float64})
  return Gaussianθ(v[1], v[2], θ.C, θ.D)
end

θ0 = SVector{d, Float64}(θ.A, θ.B)

const sigmaProp = SMatrix{d, d, Float64}(Matrix{Float64}(I, d, d))

@inline function lglogprior(theta::Gaussianθ)
  if theta.A < 0 || theta.A > 5 return -Inf end
  if theta.B < 0 || theta.B > 5 return -Inf end
  return 0.0
end

function makeSimpleLTD(ys::Vector{Float64}, N::Int64)
  rpio::RPIO{Float64Particle} = RPIO{Float64Particle}(N, length(ys), 1)
  function ltd(in::SVector{d, Float64})
    theta::Gaussianθ = toGaussianθ(in)
    lp::Float64 = lglogprior(theta)
    if lp == -Inf return -Inf end
    model::Model = makeGaussianModel(theta, ys)
    return simpleEstimate(model, rpio)
  end
end

function makeRecycleLTD(ys::Vector{Float64}, N::Int64)
  rpio::RPIO{Float64Particle} = RPIO{Float64Particle}(N, length(ys), 1)
  function ltd(in::SVector{d, Float64})
    theta::Gaussianθ = toGaussianθ(in)
    lp::Float64 = lglogprior(theta)
    if lp == -Inf return -Inf end
    model::Model = makeGaussianModel(theta, ys)
    return rpEstimate(model, rpio)
  end
end

function makeExactLTD(ys::Vector{Float64})
  function ltd(in::SVector{d, Float64})
    theta::Gaussianθ = toGaussianθ(in)
    lp::Float64 = lglogprior(theta)
    if lp == -Inf return -Inf end
    return gaussianExactLL(theta, ys)
  end
end

N = 40 * length(ys)
chainLength = 1000000

logtargetSimple = makeSimpleLTD(ys, N)
logtargetRecycle = makeRecycleLTD(ys, N)
logtargetExact = makeExactLTD(ys)

PSimple = makeAMKernel(logtargetSimple, sigmaProp)
PRecycle = makeAMKernel(logtargetRecycle, sigmaProp)
PExact = makeAMKernel(logtargetExact, sigmaProp)

Random.seed!(12345)

@time chainSimple = simulateChainProgress(PSimple, θ0, chainLength)
sar = PSimple(:acceptanceRate)

@time chainRecycle = simulateChainProgress(PRecycle, θ0, chainLength)
rar = PRecycle(:acceptanceRate)

@time chainExact = simulateChainProgress(PExact, θ0, chainLength)
ear = PExact(:acceptanceRate)

vsSimple = (i->(x->x[i]).(chainSimple)).(1:d)
vsRecycle = (i->(x->x[i]).(chainRecycle)).(1:d)
vsExact = (i->(x->x[i]).(chainExact)).(1:d)

plot(kde(vsSimple[1], sar))
plot!(kde(vsExact[1], ear))
plot!(kde(vsRecycle[1], rar))

plot(kde(vsSimple[2], sar))
plot!(kde(vsExact[2], ear))
plot!(kde(vsRecycle[2], rar))

contour(kde(vsSimple[1], vsSimple[2], sar))
contour!(kde(vsExact[1], vsExact[2], ear))
contour!(kde(vsRecycle[1], vsRecycle[2], rar))

plot(autocor(vsSimple[1]))
plot!(autocor(vsExact[1]))
plot!(autocor(vsRecycle[1]))

plot(autocor(vsSimple[2]))
plot!(autocor(vsExact[2]))
plot!(autocor(vsRecycle[2]))

avarsSimple = estimateBM.(vsSimple)
avarsRecycle = estimateBM.(vsRecycle)
avarsExact = estimateBM.(vsExact)


using VectorSphericalWaves

"""
## Failure, doesn't support ^
using MultiFloat
MultiFloats.use_bigfloat_transcendentals()
"""
n_max, kr_r, kr_i, θ, ϕ = 2, 1e6, 1e2, 0.2, 0.3;
kr = complex(kr_r, kr_i);
M_N_wave_all_m_n(n_max, kr, θ, ϕ; kind="regular");

# BigFloat
n_max, kr_r, kr_i, θ, ϕ = 2, BigFloat(1e6), BigFloat(1e2), BigFloat(0.2), BigFloat(0.3);
kr = complex(kr_r, kr_i);
M_N_wave_all_m_n(n_max, kr, θ, ϕ; kind="regular");


using Base.Test, GaussianProcesses
using GaussianProcesses: distance, KernelData, StationaryARDData, SEArd

srand(1)

d, n = 5, 10
logℓ, logσ = rand(d), rand()
X = rand(d,n)
kern = SEArd(logℓ, logσ)
data = KernelData(kern, X)

@test_approx_eq distance(kern, X, data) distance(kern, X)


function HMMforward(v,phghm,ph1,pvgh;UseLogArray=true)
    #function [logalpha,loglik]=HMMforward(v,phghm,ph1,pvgh)
    #%HMMFORWARD HMM Forward Pass
    #% [alpha,loglik]=HMMforward(v,phghm,ph1,pvgh)
    #%
    #% Inputs:
    #% v : visible (observation) sequence being a vector v=[2 1 3 3 1 ...]
    #% phghm : homogeneous transition distribution phghm(i,j)=p(h(t)=i|h(t-1)=j)
    #% ph1 : initial distribution
    #% pvgh : homogeneous emission distribution pvgh(i,j)=p(v(t)=i|h(t)=j)
    #%
    #% Outputs:
    #% logalpha : log alpha messages: log p(h(t),v(1:t))
    #% loglik : sequence log likelihood log p(v(1:T))
    #% See also HMMbackward.m, HMMviterbi.m,  HMMsmooth.m, demoHMMinference.m

    T=length(v); H=length(ph1);

    if !UseLogArray # alpha recursion (not recommended due to numerical underflow)
        z=zeros(1,T)
        alpha=zeros(H,T)
        alpha[:,1] = pvgh[v[1],:]'.*ph1
        z[1]=sum(alpha[:,1])
        alpha[:,1]=condp(alpha[:,1])
        for t=2:T
            alpha[:,t]=pvgh[v[t],:]'.*(phghm*alpha[:,t-1])
            z[t]=sum(alpha[:,t])
            alpha[:,t]=condp(alpha[:,t])
        end
        loglik = sum(log(z)) # log likelihood
    end

    if UseLogArray
        logpvgh=LogArray(log(pvgh))
        logphghm=LogArray(log(phghm))
        logph1=LogArray(log(ph1))
        logalpha=Dict{Integer,LogArray}()
        logalphaoutput=zeros(H,T)
        logalpha[1] = logpvgh[v[1],:].'.*logph1
        logalphaoutput[:,1]=(logalpha[1]).content
        for t=2:T
            logalpha[t]=logpvgh[v[t],:].'.*(logphghm*logalpha[t-1])
            #if any(isnan(logalpha[t].content))
            #    return t,v,logpvgh,logphghm,logalpha
            #end

            logalphaoutput[:,t]=(logalpha[t]).content
        end
        loglik = (sum(logalpha[T])).content # log likelihood
    end

    return logalphaoutput, loglik

end


# using Base.Test
# using MIToS.MSA
# using MIToS.Information

let s = res"ARNDCQEGHILKMFPSTWYV-",
  r = reverse(s),
  g = res"GGGGGGGGGGGGGGGGGGGG",
  Pg = probabilities(BigFloat, g),
  Pgg = probabilities(BigFloat, g,g),
  Pggg = probabilities(BigFloat, g,g,g),
  Ps = probabilities(BigFloat, s),
  Pss = probabilities(BigFloat, s,s),
  Psr = probabilities(BigFloat, s,r),
  Psss = probabilities(BigFloat, s,s,s),
  prob_random = probabilities(BigFloat, AdditiveSmoothing(1.0), Residue[], Residue[]),
  count_random = count(AdditiveSmoothing(1.0), Residue[], Residue[])

  print("""

Tests for Information Measures
==============================
""")

  print("""

Entropy
-------
""")

  print("""
Entropy: H(X)
""")

  @test estimate(Entropy{BigFloat}(), Pg) == zero(BigFloat)
  @test_approx_eq estimate(Entropy{BigFloat}(), Ps) log(big"20")
  @test_approx_eq estimate(Entropy{BigFloat}(2), Ps) log(2, big"20")

  print("""
Joint Entropy: H(X,Y)
""")

  @test estimate(Entropy{BigFloat}(), Pgg) == zero(BigFloat)
  @test_approx_eq estimate(Entropy{BigFloat}(), Pss) log(big"20")
  @test_approx_eq estimate(Entropy{BigFloat}(), Psr) log(big"19")
  @test_approx_eq estimate(Entropy{BigFloat}(2), Pss) log(2, big"20")
  @test_approx_eq estimate(Entropy{BigFloat}(), count_random) log(big"400")
  @test_approx_eq estimate(Entropy{BigFloat}(),  prob_random) log(big"400")

  print("""
Joint Entropy: H(X,Y,Z)
""")

  @test estimate(Entropy{BigFloat}(), Pggg) == zero(BigFloat)
  @test_approx_eq estimate(Entropy{BigFloat}(), Psss) log(big"20")

  print("""
Entropy using ResidueCount
""")
  @test estimate(Entropy{BigFloat}(), count(g,g)) == estimate(Entropy{BigFloat}(), Pgg)
  @test estimate(Entropy{BigFloat}(), count(s,s)) == estimate(Entropy{BigFloat}(), Pss)
  @test_approx_eq estimate(Entropy{BigFloat}(), count(s,r)) estimate(Entropy{BigFloat}(), Psr)

  print("""

Kullback-Leibler
----------------
""")

  @test estimate(Information.KullbackLeibler{Float64}(Float64[1/20 for i in 1:20]), Ps) == 0.0
  @test estimate(Information.KullbackLeibler{Float64}(BLOSUM62_Pi), Ps) != 0.0
  @test_approx_eq estimate(Information.KullbackLeibler{Float64}(BLOSUM62_Pi), Ps) sum( map(i -> 0.05 * log(0.05 / BLOSUM62_Pi[i]), 1:20) )

  @test estimate(Information.KullbackLeibler{BigFloat}(Ps), Ps) == zero(BigFloat)
  @test estimate(Information.KullbackLeibler{BigFloat}(Psr), Psr) == zero(BigFloat)

  print("""

Mutual Information
------------------
""")

  @test estimate(MutualInformation{BigFloat}(), Pgg) == zero(BigFloat)
  @test estimate(MutualInformation{BigFloat}(), probabilities(BigFloat, g,s)) == zero(BigFloat)
  @test_approx_eq estimate(MutualInformation{BigFloat}(),  prob_random) zero(BigFloat) # diff for Float64 is 1e-15

  print("""
MI(X,Y) = H(X) + H(Y) - H(X,Y)
""")
  @test estimate(MutualInformation{BigFloat}(), Psr) == estimate_on_marginal(Entropy{BigFloat}(), Psr, 1) + estimate_on_marginal(Entropy{BigFloat}(), Psr, 2) - estimate(Entropy{BigFloat}(), Psr)
  @test estimate(MutualInformation{BigFloat}(), Pss) == estimate_on_marginal(Entropy{BigFloat}(), Pss, 1) + estimate_on_marginal(Entropy{BigFloat}(), Pss, 2) - estimate(Entropy{BigFloat}(), Pss)
  @test estimate(MutualInformation{BigFloat}(2), Psr) == estimate_on_marginal(Entropy{BigFloat}(2), Psr, 1) + estimate_on_marginal(Entropy{BigFloat}(2), Psr, 2) - estimate(Entropy{BigFloat}(2), Psr)
  @test estimate(MutualInformation{BigFloat}(20), Pss) == estimate_on_marginal(Entropy{BigFloat}(20), Pss, 1) + estimate_on_marginal(Entropy{BigFloat}(20), Pss, 2) - estimate(Entropy{BigFloat}(20), Pss)
  @test_approx_eq_eps estimate(MutualInformation{BigFloat}(), prob_random) ( estimate_on_marginal(Entropy{BigFloat}(), prob_random, 1) +
                                                                              estimate_on_marginal(Entropy{BigFloat}(),  prob_random, 2) -
                                                                              estimate(Entropy{BigFloat}(), prob_random) ) 1e-74 # diff for Float64 is 1e-15

  print("""
MI using ResidueCount
""")
  @test_approx_eq estimate(MutualInformation{BigFloat}(), count_random) zero(BigFloat)
  @test estimate(MutualInformation{BigFloat}(), count(g,g)) == estimate(MutualInformation{BigFloat}(), Pgg)
  @test estimate(MutualInformation{BigFloat}(), count(s,s)) == estimate(MutualInformation{BigFloat}(), Pss)
  @test_approx_eq estimate(MutualInformation{BigFloat}(), count(s,r)) estimate(MutualInformation{BigFloat}(), Psr)

  print("""
MI(X,Y,Z)
""")
  @test estimate(MutualInformation{BigFloat}(), Pggg) == zero(BigFloat)
  @test estimate(MutualInformation{BigFloat}(), probabilities(BigFloat, g,s,r)) == zero(BigFloat)
  @test_approx_eq estimate(MutualInformation{BigFloat}(), probabilities(BigFloat, s,s,r)) estimate(MutualInformation{BigFloat}(), count(s,s,r))

  print("""
MI(X,Y,Z) = H(X) + H(Y) + H(Z) - H(X,Y) - H(X,Z) - H(Y,Z) + H(X,Y,Z)
""")
  @test estimate(MutualInformation{BigFloat}(), Psss) == ( estimate_on_marginal(Entropy{BigFloat}(), Psss, 1) + estimate_on_marginal(Entropy{BigFloat}(), Psss, 2) + estimate_on_marginal(Entropy{BigFloat}(), Psss, 3) -
                                                            estimate(Entropy{BigFloat}(), Pss) - estimate(Entropy{BigFloat}(), Pss) - estimate(Entropy{BigFloat}(), Pss) + estimate(Entropy{BigFloat}(), Psss) )

  print("""
MI(X,Y,Z) <= min{ H(X,Y), H(X,Z), H(Y,Z) }
""")
  @test_approx_eq estimate(MutualInformation{BigFloat}(), Psss) estimate(MutualInformation{BigFloat}(),Pss)

  print("""

Pairwise Gap Percentage
-----------------------
""")

  @test estimate(GapUnionPercentage{Float64}(),        count(res"AA--", res"--AA", usegap=true)) == 100.0
  @test estimate(GapIntersectionPercentage{Float64}(), count(res"AA--", res"--AA", usegap=true)) == 0.0

  @test estimate(GapUnionPercentage{Float64}(),        count(res"AA--", res"--AA", usegap=true, weight=Float64[.25, .25, .25, .25])) == 100.0
  @test estimate(GapIntersectionPercentage{Float64}(), count(res"AA--", res"--AA", usegap=true, weight=Float64[.25, .25, .25, .25])) == 0.0

  @test estimate(GapUnionPercentage{Float64}(),        count(res"AAA-", res"AA--", usegap=true)) == 50.0
  @test estimate(GapIntersectionPercentage{Float64}(), count(res"AAA-", res"AA--", usegap=true)) == 25.0

  @test estimate(GapUnionPercentage{Float64}(),        count(res"AAA-", res"AA--", usegap=true, weight=Float64[.2, .2, .2, .4])) == 60.0
  @test estimate(GapIntersectionPercentage{Float64}(), count(res"AAA-", res"AA--", usegap=true, weight=Float64[.2, .2, .2, .4])) == 40.0
end


module PhononAction

using ..Models: HolsteinModel, SSHModel
using ..Utilities: get_index

export calc_Sb, calc_dSbdxdx!

"""
Calculates the pure phonon action Sb such that exp{-Sb}.
"""
function calc_Sb(holstein::HolsteinModel{T1,T2}, shifted::Bool=false) where {T1,T2}

    x   = holstein.x::Vector{T1}
    N   = holstein.Nsites::Int
    Lτ  = holstein.Lτ::Int
    Δτ  = holstein.Δτ::T1
    ω   = holstein.ω::Vector{T1}
    ω₄  = holstein.ω₄::Vector{T1}
    λ   = holstein.λ::Vector{T1}
    Sb = 0.0::T1

    # iterate over sites in lattice
    for i in 1:N
        # iterate over time slice in lattice
        for τ in 1:Lτ
            # get τ-1 accounting for periodic boundary conditions
            τm1 = mod1(τ-1,Lτ)
            # xᵢ(τ)
            xᵢτ = x[get_index(τ,i,Lτ)]
            # xᵢ(τ-1)
            xᵢτm1 = x[get_index(τm1,i,Lτ)]
            # calculate potential energy
            Sb += ω[i]^2*xᵢτ^2/2 + ω₄[i]*xᵢτ^4 - λ[i]*xᵢτ*shifted
            # calculate kintetic energy
            Sb += (xᵢτ-xᵢτm1)^2/Δτ^2/2
        end
    end

    # calculate phonon potential energy associated with phonon dispersion
    if length(holstein.ωᵢⱼ)>0
        ωᵢⱼ                = holstein.ωᵢⱼ::Vector{T2}
        sign_ωᵢⱼ           = holstein.sign_ωᵢⱼ::Vector{Int}
        neighbor_table_ωᵢⱼ = holstein.neighbor_table_ωᵢⱼ::Matrix{Int}
        # iterate over dispersive phonon modes
        for m in 1:length(ωᵢⱼ)
            ωᵢⱼ² = ωᵢⱼ[m] * ωᵢⱼ[m]
            sgn = sign_ωᵢⱼ[m]
            # getting pair of neighboring sites
            i = neighbor_table_ωᵢⱼ[1,m]
            j = neighbor_table_ωᵢⱼ[2,m]
            # iterating over time slices
            for τ in 1:L
                # indexing offset into vectors associated with τ time slice
                indx_i = get_index(τ,i,Lτ)
                indx_j = get_index(τ,j,Lτ)
                # updating partial derivative
                Sb += ωᵢⱼ² * ( x[indx_i] + sgn*x[indx_j] )^2/2
            end
        end
    end

    # necessary scaling factor as defintiion of Sb includes a Δτ out front
    Sb = Δτ * Sb

    return Sb
end

function calc_Sb(ssh::SSHModel{T1,T2,T3}, shifted::Bool=false) where {T1,T2,T3}

    x  = ssh.x::Vector{T1}
    N  = ssh.Nph::Int
    Lτ = ssh.Lτ::Int
    Δτ = ssh.Δτ::T1
    ω  = ssh.ω::Vector{T1}
    ω₄ = ssh.ω₄::Vector{T1}
    Sb = 0.0

    # iterate over phonons in lattice
    for i in 1:N
        # τ=0 field associated with phonon
        field_τ0 = get_index(1,i,Lτ)
        # if primary phonon
        if ssh.primary_field[field_τ0] == field_τ0
            # action associated with current phonon
            sb = 0.0
            # iterate over time slice in lattice
            for τ in 1:Lτ
                # get τ-1 accounting for periodic boundary conditions
                τm1   = mod1(τ-1,Lτ)
                # get field
                field = get_index(τ,i,Lτ)
                # xᵢ(τ)
                xᵢτ   = x[field]
                # xᵢ(τ-1)
                xᵢτm1 = x[get_index(τm1,i,Lτ)]
                # calculate potential energy
                sb   += Δτ*ω[i]^2*xᵢτ^2/2 + Δτ*ω₄[i]*xᵢτ^4
                # calculate kintetic energy
                sb   += (xᵢτ-xᵢτm1)^2/Δτ/2
            end
            # add to Sb total, normalizing by the number of equivalent fields there are
            Sb += sb
        end
    end

    return Sb
end

"""
Calculates the dervative phonon action with respect to each phonon field and adds that value in place
to the vector dSbdx.
"""
function calc_dSbdx!(dSbdx::Vector{T2}, holstein::HolsteinModel{T1,T2}, shifted::Bool=false)  where {T1<:AbstractFloat,T2<:Number}

    @assert length(dSbdx)==holstein.Ndof

    x          = holstein.x::Vector{T1}
    nsites     = holstein.Nsites::Int
    Lτ         = holstein.Lτ::Int
    Δτ         = holstein.Δτ::T1
    ω          = holstein.ω::Vector{T1}
    ω₄         = holstein.ω₄::Vector{T1}
    λ          = holstein.λ::Vector{T1}

    #####################################################
    ## Calculating Derivative Phonon Action Associated ##
    ## With Local Phonon Frequency And Phonon Momentum ##
    #####################################################]

    # iterating over site
    @fastmath @inbounds for site in 1:nsites
        Δτω²  = Δτ * ω[site] * ω[site]
        Δτ4ω₄ = Δτ * 4 * ω₄[site]
        Δτλ   = Δτ * λ[site]
        # iterating over time slices
        for τ in 1:Lτ
            # get τ+1 accounting for periodic boundary conditions
            τp1 = τ%Lτ+1
            # get τ-1 accounting for periodic boundary conditions
            τm1 = (τ-2+Lτ)%Lτ+1
            # indexing offset into vectors associated with τ time slice
            indx_τ = get_index(τ,site,Lτ)
            # indexing offset into vectors associated with τ+1 time slice
            indx_τp1 = get_index(τp1,site,Lτ)
            # indexing offset into vectors associated with τ-1 time slice
            indx_τm1 = get_index(τm1,site,Lτ)
            # phonon field at current time slice
            xτ = x[indx_τ]
            # updating partial derivative
            dSbdx[indx_τ] += Δτω² * xτ - Δτλ * shifted # derivative of Δτ⋅ω²/2⋅x² term
            dSbdx[indx_τ] += Δτ4ω₄ * xτ * xτ * xτ      # derivative of Δτ⋅ω₄⋅x⁴ term
            dSbdx[indx_τ] -= ( x[indx_τp1] + x[indx_τm1] - 2.0*xτ )/Δτ
        end
    end

    #############################################################
    ## Calculating Derivative Of Phonon Action Associated With ##
    ##               Dispersive Phonon Modes                   ##
    #############################################################

    if length(holstein.ωᵢⱼ)>0
        ωᵢⱼ                = holstein.ωᵢⱼ::Vector{T2}
        sign_ωᵢⱼ           = holstein.sign_ωᵢⱼ::Vector{Int}
        neighbor_table_ωᵢⱼ = holstein.neighbor_table_ωᵢⱼ::Matrix{Int}
        # iterate over dispersive phonon modes
        for m in 1:length(ωᵢⱼ)
            Δτωᵢⱼ² = Δτ * ωᵢⱼ[m] * ωᵢⱼ[m]
            sgn = sign_ωᵢⱼ[m]
            # getting pair of neighboring sites
            i = neighbor_table_ωᵢⱼ[1,m]
            j = neighbor_table_ωᵢⱼ[2,m]
            # iterating over time slices
            for τ in 1:Lτ
                # indexing offset into vectors associated with τ time slice
                indx_i = get_index(τ,i,Lτ)
                indx_j = get_index(τ,j,Lτ)
                # updating partial derivative
                Δ = Δτωᵢⱼ² * ( x[indx_i] + sgn*x[indx_j] )
                dSbdx[indx_i] += Δ
                dSbdx[indx_j] += sgn*Δ
            end
        end
    end

    return nothing
end

function calc_dSbdx!(dSbdx::Vector{T2}, ssh::SSHModel{T1,T2,T3}, shifted::Bool=false)  where {T1,T2,T3}

    @assert length(dSbdx)==ssh.Ndof

    x   = ssh.x::Vector{T1}
    Nph = ssh.Nph::Int
    Lτ  = ssh.Lτ::Int
    Δτ  = ssh.Δτ::T1
    ω   = ssh.ω::Vector{T1}
    ω₄  = ssh.ω₄::Vector{T1}

    #####################################################
    ## Calculating Derivative Phonon Action Associated ##
    ## With Local Phonon Frequency And Phonon Momentum ##
    #####################################################

    # iterating over phonons in lattice
    @fastmath @inbounds for i in 1:Nph
        Δτω²  = Δτ * ω[i] * ω[i]
        Δτ4ω₄ = Δτ * 4 * ω₄[i]
        # iterating over time slices
        for τ in 1:Lτ
            # get τ+1 accounting for periodic boundary conditions
            τp1       = mod1(τ+1,Lτ)
            # get τ-1 accounting for periodic boundary conditions
            τm1       = mod1(τ-1,Lτ)
            # indexing offset into vectors associated with τ time slice
            field_τ   = get_index(τ,i,Lτ)
            # indexing offset into vectors associated with τ+1 time slice
            field_τp1 = get_index(τp1,i,Lτ)
            # indexing offset into vectors associated with τ-1 time slice
            field_τm1 = get_index(τm1,i,Lτ)
            # phonon field at current time slice
            xτ        = x[field_τ]
            # updating partial derivative
            dsbdx     = Δτω² * xτ # derivative of Δτ⋅ω²/2⋅x² term
            dsbdx    += Δτ4ω₄ * xτ * xτ * xτ # derivative of Δτ⋅ω₄⋅x⁴ term.
            dsbdx    -= ( x[field_τp1] + x[field_τm1] - 2.0*xτ )/Δτ # kinetic energy term
            # increment total derivative value
            dSbdx[field_τ] += dsbdx
        end
    end

    return nothing
end

end


###########################################################
# Waveguide eigenvalue problem (WEP)
# Sum of products of matrices and functions (SPMF)
"""
Waveguide eigenvalue problem (WEP)
Sum of products of matrices and functions (SPMF)
"""
function assemble_waveguide_spmf_fd(nx::Integer, nz::Integer, hx, Dxx::SparseMatrixCSC, Dzz::SparseMatrixCSC, Dz::SparseMatrixCSC, C1::SparseMatrixCSC, C2T::SparseMatrixCSC, K::Union{Matrix{ComplexF64},Matrix{Float64}}, Km, Kp, pre_Schur_fact::Bool)
    Ix = sparse(ComplexF64(1)I, nx, nx)
    Iz = sparse(ComplexF64(1)I, nz, nz)
    Q0 = kron(Ix, Dzz) + kron(Dxx, Iz) + sparse(Diagonal(vec(K)))
    Q1 = kron(Ix, 2*Dz)
    Q2 = kron(Ix, Iz)

    A = Vector{SparseMatrixCSC}(undef, 3+2*nz)
    A[1] = hvcat((2,2), Q0, C1, C2T, spzeros(ComplexF64,2*nz, 2*nz) )
    A[2] = hvcat((2,2), Q1, spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), spzeros(ComplexF64,2*nz, 2*nz) )
    A[3] = hvcat((2,2), Q2, spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), spzeros(ComplexF64,2*nz, 2*nz) )

    f = Vector{Function}(undef, 3+2*nz)
    f[1] = λ -> one(λ)
    f[2] = λ -> λ
    f[3] = λ -> λ^2

    R, Rinv = generate_R_matvecs(nz)
    S = generate_S_function(nz, hx, Km, Kp)

    for j = 1:nz
        f[j+3] = λ -> S(λ, j)
        e_j = zeros(nz)
        e_j[j] = 1
        E_j = [R(e_j); spzeros(ComplexF64,nz)]
        E_j = E_j * (E_j/nz)'
        A[j+3] =  hvcat((2,2), spzeros(ComplexF64,nx*nz,nx*nz), spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), E_j)
    end
    for j = 1:nz
        f[j+nz+3] = λ -> S(λ, nz+j)
        e_j = zeros(nz)
        e_j[j] = 1
        E_j = [spzeros(ComplexF64,nz); R(e_j)]
        E_j = E_j * (E_j/nz)'
        A[j+nz+3] =  hvcat((2,2), spzeros(ComplexF64,nx*nz,nx*nz), spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), E_j)
    end
    return SPMF_NEP(A,f;Schur_fact=pre_Schur_fact)
end

###########################################################
# Generate R-matrix
# Part of defining the P-matrix, see above, Jarlebring-(1.6) and Ringh-(2.8) and Remark 1
# OBS: R' = nz * Rinv, as noted in Ringh between (2.8) and Remark 1
function generate_R_matvecs(nz)
    # The scaled FFT-matrix R
    p = (nz-1)/2
    bb = exp.(-2im*pi*((1:nz).-1)*(-p)/nz)  # scaling to do after FFT
    function R(X) # Note! Only works for vectors or one-dim matrices
        return reverse(bb .* fft(vec(X)), dims = 1)
    end
    bbinv = 1 ./ bb # scaling to do before inverse FFT
    function Rinv(X)
        return ifft(bbinv .* reverse(vec(X), dims = 1))
    end
    return R, Rinv
end

###########################################################
# Generate S-function for matrix argument
# Part of defining the P-matrix, see above, Jarlebring-(2.4) and Ringh-(2.8)(2.3)
function generate_S_function(nz, hx, Km, Kp)
    # Constants from the problem
    p = (nz-1)/2
    d0 = -3/(2*hx)
    b = 4*pi*1im * (-p:p)
    cM = Km^2 .- 4*pi^2 * ((-p:p).^2)
    cP = Kp^2 .- 4*pi^2 * ((-p:p).^2)

    # Note: λ should be scalar or matrix (not vector)
    betaM = function(λ, j)
        return λ^2 + b[j]*λ + cM[j]*one(λ)
    end
    betaP = function(λ, j)
        return λ^2 + b[j]*λ + cP[j]*one(λ)
    end

    sM = function(λ, j)
        return  1im*sqrt_schur_pos_imag(betaM(λ, j)) + d0*one(λ)
    end
    sP = function(λ, j)
        return  1im*sqrt_schur_pos_imag(betaP(λ, j)) + d0*one(λ)
    end

    S = function(λ, j)
        if j <= nz
            return sM(λ,j)
        elseif j <= 2*nz
            return sP(λ,(j-nz))
        else
            error("The chosen j = ", j, "but the setup nz = ", nz, ". Hence j>2nz which is illegal.")
        end
    end

    return S
end

###########################################################
"""
    sqrt_schur_pos_imag(A::AbstractMatrix)
 Computes the matrix square root on the 'correct branch',
 that is, with positivt imaginary part. Similar to Schur method
 in Algorithm 6.3 in Higham matrix functions.
"""
sqrt_schur_pos_imag(A::Number) = sqrt_pos_imag(A)
function sqrt_schur_pos_imag(A::AbstractMatrix)
    n = size(A,1);
    AA = Matrix{ComplexF64}(A);
    (T, Q, ) = schur(AA)
    U = zeros(ComplexF64,n,n);
    for i = 1:n
        U[i,i] = sqrt_pos_imag(T[i,i])
    end
    private_inner_loops_sqrt!(n, U, T)
    return Q*U*Q'
end
#Helper function executing the inner loop (more Juliaesque)
function private_inner_loops_sqrt!(n, U, T)
    temp = zero(ComplexF64);
    for j = 2:n
        for i = (j-1):-1:1
            temp *= zero(ComplexF64);
            for k = (i+1):(j-1)
                temp += U[i,k]*U[k,j]
            end
            U[i,j] = (T[i,j] - temp)/(U[i,i]+U[j,j])
        end
    end
end

"""
    sqrt_pos_imag(a::ComplexF64) and sqrt_pos_imag(a::Float64)
 Helper function: Computes the scalar square root on the 'correct branch',
 that is, with positivt imaginary part.
"""
function sqrt_pos_imag(a::ComplexF64)
    imag_sign = sign(imag(a))
    if imag_sign == 0 #Real value in complex type
        sqrt(a)
    else
        sign(imag(a))*sqrt(a)
    end
end
function sqrt_pos_imag(a::Float64)
    return sqrt(a)
end

    function Pinv(nep,λ,x)
        return vec(  [R(nep,Rinv(nep,x[1:Int(end/2)]) ./ sM(nep,λ));
                      R(nep,Rinv(nep,x[Int(end/2)+1:end]) ./ sP(nep,λ))  ]  )
    end

    function R(nep,X) # Note! Only works for vectors or one-dim matrices
        return reverse(nep.bb .* fft(vec(X)), dims = 1)
    end

    function Rinv(nep,X)
        return ifft(nep.bbinv .* reverse(vec(X), dims = 1))
    end

    function betaM(nep,λ)
        return λ^2 .+ nep.b*λ + nep.cM
    end

    function betaP(nep,λ)
        return λ^2 .+ nep.b*λ + nep.cP
    end

    function sM(nep,λ)
        bbeta = betaM(nep,λ)
        return 1im*sign.(imag(bbeta)).*sqrt.(bbeta) .+ nep.d0
    end

    function sP(nep,λ)
        bbeta = betaP(nep,λ)
        return 1im*sign.(imag(bbeta)).*sqrt.(bbeta) .+ nep.d0
    end

###########################################################
# Waveguide eigenvalue problem - WEP
# A more optimized (native) implementation of the WEP with FD discretization
"""
    WEP_FD

Waveguide eigenvalue problem

A more optimized implementation of the WEP for FD-discretization.\\
Closer to what is done in the article:
''E. Ringh, and G. Mele, and J. Karlsson, and E. Jarlebring,
Sylvester-based preconditioning for the waveguide eigenvalue problem,
Linear Algebra and its Applications''
"""
    struct WEP_FD <: WEP
        nx::Int
        nz::Int
        hx::Float64
        hz::Float64
        Dxx::SparseMatrixCSC{Float64,Int}
        Dzz::SparseMatrixCSC{Float64,Int}
        Dz::SparseMatrixCSC{Float64,Int}
        Iz::SparseMatrixCSC{Float64,Int}
        C1::SparseMatrixCSC{Float64,Int}
        C2T::SparseMatrixCSC{Float64,Int}
        k_bar::ComplexF64
        K::Matrix{ComplexF64}
        p::Integer
        d0::Float64
        d1::Float64
        d2::Float64
        b::Vector{ComplexF64}
        cM::Vector{ComplexF64}
        cP::Vector{ComplexF64}
        bb::Vector{ComplexF64}
        bbinv::Vector{ComplexF64}

        function WEP_FD(nx, nz, hx, hz, Dxx, Dzz, Dz, C1, C2T, K, Km, Kp)
            n = nx*nz + 2*nz
            k_bar = mean(K)
            K_scaled = K-k_bar*ones(ComplexF64,nz,nx)

            p = (nz-1)/2

            d0 = -3/(2*hx)
            d1 = 2/hx
            d2 = -1/(2*hx)

            b = 4*pi*1im * (-p:p)
            cM = Km^2 .- 4*pi^2 * ((-p:p).^2)
            cP = Kp^2 .- 4*pi^2 * ((-p:p).^2)

            bb = exp.(-2im*pi*((1:nz).-1)*(-p)/nz)  # scaling to do after FFT
            bbinv = 1 ./ bb # scaling to do before inverse FFT

            Iz = sparse(ComplexF64(1)I, nz, nz)

            return new(nx, nz, hx, hz, Dxx, Dzz, Dz, Iz, C1, C2T, k_bar, K_scaled, p, d0, d1, d2, b, cM, cP, bb, bbinv)
        end
    end

    function A(nep::WEP_FD, λ, d=0)
        nz = nep.nz
        if(d == 0)
            return nep.Dzz + 2*λ*nep.Dz + (λ^2 + nep.k_bar)*nep.Iz
        elseif(d == 1)
            return 2*nep.Dz + 2*λ*nep.Iz
        elseif(d == 2)
            return 2*nep.Iz
        else
            return spzeros(ComplexF64, nz, nz)
        end
    end

    function B(nep::WEP_FD, λ, d=0)
        if(d == 0)
            return nep.Dxx
        else
            return spzeros(ComplexF64, nx, nx)
        end
    end

    #Inverses of the boundary operators, Ringh - (2.8)
    #To be used in the Schur-complement- and SMW-context.
    function P_inv_m(nep::WEP_FD, λ, v)
        nz = nep.nz
        coeffs = zeros(ComplexF64, nz)
            aa = 1.0
        for j = 1:nz
            bb = nep.b[j]
            cc = nep.cM[j]
            coeffs[j] = 1im*sqrt_derivative(aa, bb, cc, 0, λ) + nep.d0
        end
        return R(nep,Rinv(nep,v) ./ coeffs)
    end

    function P_inv_p(nep::WEP_FD, λ, v)
        nz = nep.nz
        coeffs = zeros(ComplexF64, nz)
            aa = 1.0
        for j = 1:nz
            bb = nep.b[j]
            cc = nep.cP[j]
            coeffs[j] = 1im*sqrt_derivative(aa, bb, cc, 0, λ) + nep.d0
        end
        return R(nep,Rinv(nep,v) ./ coeffs)
    end

    function size(nep::WEP_FD)
        n = nep.nx*nep.nz + 2*nep.nz
        return (n,n)
    end
    function size(nep::WEP_FD, dim)
        n = nep.nx*nep.nz + 2*nep.nz
        return n
    end

    function issparse(nep::WEP_FD)
        return false
    end

"""
    compute_Mlincomb(nep::WEP_FD, λ::Number, V; a=ones(ComplexF64,size(V,2)))

Specialized for Waveguide Eigenvalue Problem discretized with Finite Difference\\\\
 Computes the linear combination of derivatives\\
 ``Σ_i a_i M^{(i)}(λ) v_i``
"""
    function compute_Mlincomb(nep::WEP_FD, λ::Number, V::AbstractVecOrMat,
                              a::Vector=ones(ComplexF64,size(V,2)))
        na = size(a,1)
        nv = size(V,1)
        mv = size(V,2)
        n_nep = size(nep,1)
        if(na != mv)
            error("Incompatible sizes: Number of coefficients = ", na, ", number of vectors = ", mv, ".")
        end
        if(nv != n_nep)
            error("Incompatible sizes: Length of vectors = ", nv, ", size of NEP = ", n_nep, ".")
        end
        nx = nep.nx
        nz = nep.nz
        max_d = na - 1 #Start on 0:th derivative

        V1 = view(V, 1:nx*nz, :)
        V1_mat = reshape(V1, nz, nx, na)
        V2 = view(V, nx*nz+1:n_nep, :)

        # Compute the top part (nx*nz)
        y1_mat::Matrix{ComplexF64} = (A(nep,λ) * V1_mat[:,:,1] + V1_mat[:,:,1] * B(nep,λ)  +  nep.K .* V1_mat[:,:,1])*a[1]
        for d = 1:min(max_d,3)
            y1_mat[:,:] += A(nep,λ,d) * V1_mat[:,:,d+1] * a[d+1];
        end
        y1::Vector{ComplexF64} = y1_mat[:]
        y1[:] += nep.C1 * V2[:,1] * a[1]

        # Compute the bottom part (2*nz)
        D::Matrix{ComplexF64} = zeros(ComplexF64, 2*nz, na)
        cMP = vcat(nep.cM, nep.cP)
        for j = 1:2*nz
            aa = 1
            bb = nep.b[rem(j-1,nz)+1]
            cc = cMP[j]
            der_coeff = 1im*sqrt_derivative(aa, bb, cc, max_d, λ)
            for jj = 1:na
                D[j, jj] = der_coeff[jj]
            end
        end

        #Multiplication with diagonal matrix optimized by working "elementwise" Jarlebring-(4.6)
        y2_temp::Vector{ComplexF64} =
            (D[:,1] .+ nep.d0) .* [Rinv(nep,V2[1:nz,1]);
                                   Rinv(nep,V2[nz+1:2*nz,1])]*a[1]
        for jj = 2:na
            #Multiplication with diagonal matrix optimized by working "elementwise" Jarlebring-(4.6)
            y2_temp[:] += D[:,jj] .* [Rinv(nep,V2[1:nz,jj]);
                                      Rinv(nep,V2[nz+1:2*nz,jj])] *a[jj]
        end
        y2::Vector{ComplexF64} = [R(nep,y2_temp[1:nz,1]);
                                  R(nep,y2_temp[nz+1:2*nz,1])]
        y2[:] += nep.C2T * V1[:,1]*a[1] #Action of C2T. OBS: Add last because of implcit storage in R*D_i*R^{-1}*v_i

        return vcat(y1, y2)
    end

"""
    compute_Mder(nep::WEP,λ::Number,i::Integer=0)
"""
    function compute_Mder(nep::WEP,λ::Number,i::Integer=0)
        error("The WEP does not implement this function. If this was called in a situation where you want to solve linear systems please look at `WEPLinSolverCreator`")
    end

###########################################################
# Linear Solvers for WEP

    # Matrix vector operations for the Schur complement (to be used in GMRES call)
    # Matrix-vector product according to Ringh (2.13) and (3.3)
    struct SchurMatVec
        nep::WEP_FD
        λ::ComplexF64
        SchurMatVec(nep::WEP_FD, λ::Union{ComplexF64,Float64}) = new(nep, λ)
    end

    function *(M::SchurMatVec,v::AbstractVector)
        λ = M.λ
        nep = M.nep

        X = reshape(v, nep.nz, nep.nx)
        return vec(  vec( A(nep,λ)*X + X*B(nep,λ) + nep.K.*X ) - nep.C1 * Pinv(nep, λ, nep.C2T*v)  )
    end

    function (M::SchurMatVec)(v::AbstractVector) #Overload the ()-function so that a SchurMatVec struct can act and behave like a function
        return M*v
    end

    function size(M::SchurMatVec, dim=-1)
        n = M.nep.nx*M.nep.nz
        if (dim==-1)
            return (n,n)
        else
            return n
        end
    end

    function eltype(M::SchurMatVec)
        return ComplexF64
    end

    # GMRES Solver
    struct WEPGMRESLinSolver<:LinSolver
        schur_comp::LinearMap{ComplexF64}
        kwargs
        gmres_log::Bool
        nep::WEP_FD
        λ::ComplexF64

        function WEPGMRESLinSolver(nep::WEP_FD,λ::Union{ComplexF64,Float64},kwargs)
            f = SchurMatVec(nep, λ)
            schur_comp = LinearMap{ComplexF64}(f, nep.nx*nep.nz, ismutating=false, issymmetric=false, ishermitian=false);
            gmres_log = false
            for elem in kwargs
                gmres_log |= ((elem[1] == :log) && elem[2])
            end
            return new(schur_comp, kwargs, gmres_log,nep,λ)
        end
    end

    function WEP_inner_lin_solve(solver::WEPGMRESLinSolver, rhs::Vector, tol)
        if( solver.gmres_log )
            q, convhist = gmres(solver.schur_comp, rhs; tol=tol, solver.kwargs...)
        else
            q = gmres(solver.schur_comp, rhs; tol=tol, solver.kwargs...)
        end
        return q
    end

    # Direct Backslash solver
    struct WEPBackslashLinSolver<:LinSolver
        schur_comp::SparseMatrixCSC{ComplexF64,Int}
        nep::WEP_FD
        λ::ComplexF64

        function WEPBackslashLinSolver(nep::WEP_FD, λ::Union{ComplexF64,Float64}, kwargs=())
            schur_comp = construct_WEP_schur_complement(nep, λ)
            return new(schur_comp, nep, λ)
        end
    end

    function WEP_inner_lin_solve(solver::WEPBackslashLinSolver, rhs::Vector, tol)
        return solver.schur_comp \ rhs
    end

    # Direct pre-factorized solver
    struct WEPFactorizedLinSolver<:LinSolver
        schur_comp_fact
        nep::WEP_FD
        λ::ComplexF64

        function WEPFactorizedLinSolver(nep::WEP_FD, λ::Union{ComplexF64,Float64}, kwargs=())
            schur_comp_fact = factorize(construct_WEP_schur_complement(nep, λ))
            return new(schur_comp_fact, nep, λ)
        end
    end

    function WEP_inner_lin_solve(solver::WEPFactorizedLinSolver, rhs::Vector, tol)
        return solver.schur_comp_fact \ rhs
    end

"""
    struct WEPLinSolverCreator <: LinSolverCreator
    function WEPLinSolverCreator(;solver_type=:backslash,kwargs=())

The a linsolver creator for the waveguide eigenvalue problem. The
kwarg `solver_type` is one of `:backslash`, `:factorized`, `:gmres`.
The `kwargs` keyword argument is passed to the solver.
"""
    struct WEPLinSolverCreator <: LinSolverCreator
        solver_type::Symbol
        kwargs;
        function WEPLinSolverCreator(;solver_type=:factorized,kwargs=())
            return new(solver_type,kwargs);
        end
    end

    function create_linsolver(creator::WEPLinSolverCreator,nep,λ)
        if (!(nep isa WEP_FD))
            t=typeof(nep)
            error("WEPLinSolver can only be used in combination with WEPs: type(nep)=$t");
        end
        if (creator.solver_type==:backslash)
            return WEPBackslashLinSolver(nep, λ, creator.kwargs)
        elseif (creator.solver_type==:gmres)
            return WEPGMRESLinSolver(nep, λ, creator.kwargs)
        elseif (creator.solver_type==:factorized)
            return WEPFactorizedLinSolver(nep, λ, creator.kwargs)
        else
            s=creator.solver_type
            error("Unknown type of solver_type in linsolvercreator:$s");
        end
    end

    # Helper functions for WEP LinSolvers. To avoid code repetition.
    # Assembls the full Schur-complement, used in both Backslash and LU solvers
    function construct_WEP_schur_complement(nep::WEP_FD, λ::Union{ComplexF64,Float64})
        nz = nep.nz
        nx = nep.nx
        Inz = sparse(ComplexF64(1)I, nz, nz)
        Inx = sparse(ComplexF64(1)I, nx, nx)

        Pinv_minus = Matrix{ComplexF64}(undef, nz, nz)
        Pinv_plus = Matrix{ComplexF64}(undef, nz, nz)
        e = zeros(ComplexF64,nz)
        for i = 1:nz
            e[:] .= 0
            e[i] = 1
            Pinv_minus[:,i] = P_inv_m(nep, λ, e)
            Pinv_plus[:,i] = P_inv_p(nep, λ, e)
        end

        E = spzeros(nx,nx)
        E[1,1] = nep.d1/(nep.hx^2)
        E[1,2] = nep.d2/(nep.hx^2)
        EE = spzeros(nx,nx)
        EE[nx,nx] = nep.d1/(nep.hx^2)
        EE[nx,nx-1] = nep.d2/(nep.hx^2)

        # Kronecker product form of Ringh - Proposition 3.1
        return kron(copy(B(nep,λ)'), Inz) + kron(Inx, A(nep,λ)) + sparse(Diagonal(nep.K[:])) - kron(E, Pinv_minus) - kron(EE, Pinv_plus)
    end

    # lin_solve function to wrapp all the WEP linear solvers.
    # Since Schur-complement transformations are the same.
    # Does transforming between that and the full system.
    # Ringh - Proposition 2.1, see also Algorithm 2, step 10-11.
    function lin_solve(solver::Union{WEPBackslashLinSolver,WEPGMRESLinSolver,WEPFactorizedLinSolver}, x::Vector; tol=eps(Float64))
    # Ringh - Proposition 2.1
        λ = solver.λ
        nep = solver.nep

        x_int = x[1:(nep.nx*nep.nz)]
        x_ext = x[((nep.nx*nep.nz)+1):((nep.nx*nep.nz) + 2*nep.nz)]
        rhs =  vec(  x_int - nep.C1*Pinv(nep, λ, x_ext))

        q = WEP_inner_lin_solve(solver, rhs, tol)

        return [q; vec(Pinv(nep, λ, -nep.C2T * q + x_ext))]

    end

###########################################################
#Square root of second degree polynomial (Gegenbauer polynomials)
#Jarlebring - Appendix C
"""
    sqrt_derivative(a,b,c, d=0, x=0)

Computes all d derivatives of sqrt(a*z^2 + b*z + c)
in the point z = x.
Returns a d+1 vector with all numerical values
"""
function sqrt_derivative(a,b,c, d=0, x=0)
    if(d<0)
        error("Cannot take negative derivative. d = ", d)
    end

    aa = a
    bb = b + 2*a*x
    cc = c + a*x^2 + b*x

    derivatives = zeros(ComplexF64,d+1)

    yi = sqrt_pos_imag(cc)
    derivatives[1] = yi
    if( d==0 )
        return derivatives[1] #OBS: If only function is sought, return it as an integer and not array
    end

    yip1 = bb/(2*sqrt_pos_imag(cc))
    fact = Float64(1)
    derivatives[2] = yip1 * fact
    if( d==1 )
        return derivatives
    end

    yip2 = zero(ComplexF64)
    for i = 2:d
        m = i - 2
        yip2 = - (2*aa*(m-1)*yi  +  bb*(1+2*m)*yip1) / (2*cc*(2+m))
        fact *= i

        yi = yip1
        yip1 = yip2

        derivatives[i+1] = yip2 * fact
    end
    return derivatives
end


# Find a way to run these example in isolation
using Femshop

init_femshop("poisson1d")

useLog("poisson1dlog")

domain(1)

functionSpace(order=3)
mesh(LINEMESH, elsperdim=20)

u = variable("u")
testSymbol("v")
boundary(u, 1, DIRICHLET, 0)

coefficient("f", "-100*pi*pi*sin(10*pi*x)*sin(pi*x) -
             pi*pi*sin(10*pi*x)*sin(pi*x) + 20*pi*pi*cos(10*pi*x)*cos(pi*x)")
weakForm(u, "-grad(u)*grad(v) - f*v")

import Femshop: grid_data, mesh_data, refel
dofs_per_node = 1

N1 = size(grid_data.allnodes,2);
Nn = dofs_per_node * N1;
Np = refel.Np;
nel = mesh_data.nel;

rhsvec = zeros(Nn);
lhsmatI = zeros(Int, nel*dofs_per_node*Np*dofs_per_node*Np);
lhsmatJ = zeros(Int, nel*dofs_per_node*Np*dofs_per_node*Np);
lhsmatV = zeros(nel*dofs_per_node*Np*dofs_per_node*Np);
allocated_vecs = [rhsvec, lhsmatI, lhsmatJ, lhsmatV];

Femshop.CGSolver.assemble(u, Femshop.bilinears[u.index], Femshop.linears[u.index], allocated_vecs)

print("END")

#SUITE["poisson1d-assemble"] = @benchmarkable Femshop.CGSolver.assemble($u,
                         #$Femshop.bilinears[u.index], $Femshop.linears[u.index])


#= Finite field aritmhetic implemented in unsigned integers =#
#= It is generally faster than AbstractAlgebra.GF interface =#

# returns x - y*a modulo ch
@inline function umultsubmod(
        x::T, y::T, a::T, ch::T) where {T<:Unsigned}
    ya = y*a % ch
    (x + ch - ya) % ch
end

# returns x + y*a modulo ch
@inline function umultsummod(
        x::T, y::T, a::T, ch::T) where {T<:Unsigned}
    ya = y*a % ch
    (x + ya) % ch
end

# returns x⁻¹ modulo ch
@inline function uinvmod(x::T, ch::T) where {T<:Unsigned}
    # seems like we dont need % ch
    invmod(x, ch) % ch
end

# returns -x modulo ch
@inline function ucompmod(x::T, ch::T) where {T<:Unsigned}
    (ch - x) % ch
end

# returns x*y modulo ch
@inline function umultmod(x::T, y::T, ch::T) where {T<:Unsigned}
    (x*y) % ch
end

# returns x + y modulo ch
@inline function usummod(x::T, y::T, ch::T) where {T<:Unsigned}
    (x + y) % ch
end


function levMar(z::Array{T,1}, alpha::Array{T,1}, c)
    n :: UInt64 = length(z);
    maxiter :: UInt64 = 1000000;
    mu :: Float64 = 200;
    v :: UInt64 = 2;

    H = Array{T}(8*n,8*n);

    resid_old = residual(z, alpha, c);
    H,d =  hessianApprox(z, alpha, c );

    for j = 1:maxiter

        alpha_old = alpha
        z_old = z

        if isinf( mu ) || isnan( mu )
            @printf "mu too large or too small: %e\n" mu
            break;
        end

        s = (im * imag(H) - diagm(diag(H))*mu)\d;

        z_cand = z + s[1:length(z)];
        alpha_cand = alpha + s[length(z)+1:2*length(z)];
        resid_cand = residual(z_cand, alpha_cand, c);

        # this should be positive
        #LL =  real(dot(s , ( mu * s + d )));

        gain = (resid_old - resid_cand)

        if gain[1] > 0
            v = 2;

            mu = mu * 1/2;
            z = z_cand;
            alpha = alpha_cand;
            resid_old = residual(z, alpha, c);

            H,d =  hessianApprox( z, alpha, c );

        else
            mu = mu * v; v = 2*v;
        end

        if mod(j,10000) == 0
            @printf "iter: %d\tresid: %1.15e\tmu: %e\tv: %e\n" j resid_old[1] mu v
        end
    end

    @printf "resid: %e\tmu: %e\tv: %e\n" resid_old[1] mu v
    return z, alpha;
end


struct PosteriorPsi
   beta::Distributions.Beta
   emperical::Vector{Float64}
   psi::Float64
end

PosteriorPsi() = PosteriorPsi( Beta(1,1), Float64[], 0.0 )

# create one PosteriorPsi object given psi, N, and priors α and β
# Posterior                      Likelihood         Prior
# Beta( Ψ*N + α, (1-Ψ)*N + β ) ∝ Binomial( Ψ, N ) * Beta( α, β )
function PosteriorPsi( Ψ::Float64, N::Float64, α=1.0, β=1.0; size=1000 )
   # Posterior Beta is k + α, n-k + β
   beta = Beta( Ψ*N + α, (1-Ψ)*N + β )
   emperical = rand(beta, size)
   psi = mean(beta)
   PosteriorPsi( beta, emperical, Ψ )
end

# combine multiple beta posteriors for replicates into one
# if unpaired, we don't need to maintain sample1_rep1 to sample2_rep1
# so we can re-sample from the new joint beta
function PosteriorPsi( set::Vector{PosteriorPsi}; paired=false, size=1000 )
   emperical = vcat( map( x->x.emperical, set )... )
   beta = fit(Beta, emperical)
   psi  = mean(beta)
   if !paired
      emperical = rand(beta, size)
   end
   PosteriorPsi( beta, emperical, psi )
end

# Calculates the probability, P( a-b > amt)
function probability( a::PosteriorPsi, b::PosteriorPsi; amt=0.0 )
   success = 0
   total = min(length(a.emperical),length(b.emperical))
   for i in 1:total
      success += a.emperical[i] - b.emperical[i] > amt ? 1 : 0
   end
   @fastmath success / total
end

function open_stream( filename )
   fopen = open( filename, "r" )
   if isgzipped( filename )
      stream = ZlibInflateInputStream( fopen, reset_on_end=true )
   else
      stream = BufferedStreams.BufferedInputStream( fopen )
   end
   stream
end

function open_streams( files::Vector{String} )
   buf = Vector{BufferedStreams.BufferedInputStream}(length(files))
   for i in 1:length(files)
      buf[i] = open_stream( files[i] )
   end
   buf
end

parse_float_omit_text( str::S, header::String ) where {S <: AbstractString} = str != "NA" && str != header ? parse(Float64, str) : 0.0

function parse_psi_line( line::String; min_num=5, size=1000 )
   res  = split( line, '\t' )
   psi  = parse_float_omit_text( res[6], "Psi" )
   num  = parse_float_omit_text( res[9], "Total_Reads" )
   if psi < 0 || num < min_num
      post = PosteriorPsi()
      bool = false
   else
      post = PosteriorPsi( psi, num, size=size )
      bool = true
   end
   res,post,bool
end

mutable struct PosteriorEvent
   event::Vector{SubString{String}}
   complexity::String
   a::PosteriorPsi
   b::PosteriorPsi
end

parse_complexity( c::S ) where {S <: AbstractString} = split( c, COMPLEX_CHAR, keep=false )[1] |> x->parse(Int,x)

function process_psi_line( streams::Vector{BufferedStreams.BufferedInputStream}; min_reads=5, size=1000 )
   postvec = Vector{PosteriorPsi}()
   event   = split( "", "" )
   complex = 0
   entropy = 0.0
   i = 1
   while i <= length(streams)
      line = readline( streams[i] )
      i += 1
      if line != ""
         par,post,isok = parse_psi_line( line, min_num=min_reads, size=size )
         event = par[1:5]
         event[5] == "BS" && (i -= 1; continue)
         !isok && continue
         push!( postvec, post )
         parcomplex = parse_complexity( par[10] )
         parentropy = parse_float_omit_text( par[11], "Entropy" )
         complex = parcomplex > complex ? parcomplex : complex
         entropy = parentropy > entropy ? parentropy : entropy
      end
   end
   event,complex,entropy,postvec
end

function process_psi_files( outfile, a::Vector{BufferedStreams.BufferedInputStream},
                                     b::Vector{BufferedStreams.BufferedInputStream};
                                     min_samp=1, min_reads=5, amt=0.0, size=1000 )
   io = open( outfile, "w" )
   stream = ZlibDeflateOutputStream( io )
   output_diff_header( stream )
   while true # go through all lines until we hit eof
      a_event,a_complex,a_entropy,a_post = process_psi_line( a, min_reads=min_reads, size=size )
      b_event,b_complex,b_entropy,b_post = process_psi_line( b, min_reads=min_reads, size=size )
      if a_event == split( "", "" ) || b_event == split( "", "" )
         break # eof
      end
      complex = a_complex > b_complex ? a_complex : b_complex
      entropy = a_entropy > b_entropy ? a_entropy : b_entropy
      @assert( a_event == b_event, "Incorrect events matched!!" )
      if length(a_post) >= min_samp && length(b_post) >= min_samp
         a_post   = PosteriorPsi( a_post ) # fit new posterior
         b_post   = PosteriorPsi( b_post )
         fwdprob  = probability( a_post, b_post, amt=amt )
         revprob  = probability( b_post, a_post, amt=amt )
         prob     = fwdprob > revprob ? fwdprob : revprob
         psi_a    = mean(a_post.beta)
         psi_b    = mean(b_post.beta)
         deltapsi = psi_a - psi_b
         output_diff( stream, a_event, complex, entropy, psi_a, psi_b, deltapsi, prob )
      end
   end
   close(stream)
   close(io)
end


function areaBC(r)
    return pi.*r.^2
end


##
using IFGF
using LinearAlgebra
using StaticArrays
using Random

include(joinpath(IFGF.PROJECT_ROOT,"test","testutils.jl"))

# parameters
p        = (3,5,5)
ppw      = 10 # point per wavelength
T        = ComplexF64
r        = 1 # sphere radius
area     = 4π*r^2
nmax     = 100
Δs       = (0.5,π/2,π/2)

# loop over wavenumbers
kvec = [2π,2π,4π,8π]
# kvec = [2π,2π,4π,8π,16π,32π,64π]

for k in kvec
    n        = ceil(Int,ppw^2*k^2*area/(4π^2))
    K        = HelmholtzKernel(k)
    Xpts = Ypts = sphere_uniform(n,r)
    I   = collect(1:1000)
    B   = ones(T,n)
    C   = zeros(T,n)
    exa = [sum(K( Xpts[i],Ypts[j])*B[j] for j in 1:n) for i in I]
    ta  = @elapsed A = assemble_ifgf(K,Xpts,Ypts;nmax,order=p.-1,Δs,threads=true)
    tp = @elapsed mul!(C,A,B,1,0;threads=true)
    er = norm(C[I]-exa,2) / norm(exa,2)
    println("="^80)
    println(A)
    println("lambda=$(2*pi/k),N=$n,checksum=$(exa[1]),tassemble = $ta, tprod=$tp,
    error=$er")
end


## Farfield dynamics
#==========================================================================================#

"""
    farfield_velocity(r_i, r_j, Γ_j)

Compute the induced velocity at point ``r_i`` of the wake in the Trefftz plane due to the strength ``Γ_j`` at ``r_j``.
"""
farfield_velocity(r_i, r_j, Γ_j) = let r = r_i - r_j; Γ_j / 2π * SVector(1., 0., 0.) × r / norm(r)^2 end

"""
    farfield_influence_matrix(centers, normals, points)

Compute the aerodynamic influence coefficient matrix of the wake in the Trefftz plane given the center points, normal vectors, and the points of the wake panels.
"""
farfield_influence_matrix(centers, normals, points) = [ dot(farfield_velocity(r_i, r_j2, 1.) - farfield_velocity(r_i, r_j1, 1.), n_i) for (r_i, n_i) in zip(centers, normals), (r_j2, r_j1) in zip(points[2:end], points) ]

"""
    doublet_normal_derivatives(wake_points, Δφs, normals)

Compute the normal derivative strengths of the doublets given the wake points, net doublet strengths ``Δφ``s, and associated normal vectors.
"""
function doublet_normal_derivatives(wake_points, Δφs, normals)
    centers = @. (wake_points[1:end-1] + wake_points[2:end]) / 2
    AIC     = farfield_influence_matrix(centers, normals, wake_points)
    ∂φ_∂n   = AIC * Δφs
end

project_vector(vector, U) = vector - dot(U, vector) * U
normal(wake_proj_vector, U) = normalize(U × wake_proj_vector)
dihedral(wake_proj_vec)  = atan(wake_proj_vec[3], wake_proj_vec[2])

"""
    trefftz_plane_quantities(horseshoes :: AbstractArray{<: Horseshoe}, α, β)

Project `Horseshoe`s into the Trefftz plane aligned with the wind axes angles ``α,~β``, and compute normal vectors, projection angles and lengths.
"""
function trefftz_plane_quantities(horseshoes :: AbstractArray{<: Horseshoe}, α, β)
    # Reference velocity for broadcasting
    U_ref = (Ref ∘ SVector)(1, 0, 0)

    # Transform wake to wind axes
    wake_horsies = @views vec(horseshoes[end,:])                  # Trailing edge
    wake_points  = [ r1.(wake_horsies); [r2(wake_horsies[end])] ] # Trailing edge points
    wake_wind    = geometry_to_wind_axes.(wake_points, α, β)      # Wind axis transformation

    # Project trailing edge horseshoes' bound legs into Trefftz plane along wind axes
    wake_vectors   = @. wake_wind[2:end] - wake_wind[1:end-1] # Vectors
    wake_proj_vecs = @. project_vector(wake_vectors, U_ref)   # Project to wind axes

    # Normals, dihedral angles, and projection lengths
    wake_normals   = @. normal(wake_proj_vecs, U_ref)
    wake_dihedrals = @. dihedral(wake_proj_vecs)
    wake_lengths   = @. norm(wake_proj_vecs)

    wake_points, wake_normals, wake_dihedrals, wake_lengths
end

"""
    compute_farfield_forces(Δφs, Δs, ∂φ_∂n, θs, V, ρ)

Compute the aerodynamic forces in the Trefftz plane given cumulative doublet strengths ``Δφ``s, Trefftz panel lengths ``Δs``, doublet-normal directional derivatives ``∂φ/∂n``, Trefftz panel angles ``θ``s, the freestream speed ``V`` and density ``ρ``.
"""
function compute_farfield_forces(Δφs, Δs, ∂φ_∂n, θs, V, ρ)
    D_i = - 1/2 * ρ * sum(@. Δφs * Δs * ∂φ_∂n)
    Y   = - ρ * V * sum(@. Δφs * Δs * sin(θs))
    L   = ρ * V * sum(@. Δφs * Δs * cos(θs))

    SVector(D_i, Y, L)
end

"""
    farfield_forces(Γs, horseshoes, freestream, ρ)

Compute the aerodynamic forces in the Trefftz plane normal to the freestream given horseshoes, their associated strengths ``Γ``s, and a density ``ρ``.
"""
function farfield_forces(Γs, horseshoes, speed, α, β, ρ)
    # Get projections of horseshoes into Trefftz plane with the associated normals, dihedral angles and lengths
    wake_points, normals, dihedrals, Δs = trefftz_plane_quantities(horseshoes, α, β)

    # Compute directional derivatives of doublets in the normal direction
    Δφs   = vec(sum(Γs, dims = 1))
    ∂φ_∂n = doublet_normal_derivatives(wake_points, Δφs, normals)

    # Compute forces
    compute_farfield_forces(Δφs, Δs, ∂φ_∂n, dihedrals, speed, ρ)
end


function set_volume(nodes,cellxmax,cellymax)
    volume = zeros(cellxmax,cellymax)
    for i in 1:cellxmax
        for j in 1:cellymax
            vec_r1x = nodes[i+1,j+1,1] - nodes[i,j,1]
            vec_r1y = nodes[i+1,j+1,2] - nodes[i,j,2]
            vec_r2x = nodes[i,j+1,1] - nodes[i+1,j,1]
            vec_r2y = nodes[i,j+1,2] - nodes[i+1,j,2]

            volume[i,j] = abs(vec_r1x*vec_r2y - vec_r1y*vec_r2x) /2
        end
    end
    return volume
end

function setup_cell_flux_hat(Qcon,cellxmax,cellymax,cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus)
    cell_E_hat_plas  = zeros(cellxmax,cellymax,4)
    cell_E_hat_minus = zeros(cellxmax,cellymax,4)
    cell_F_hat_plas  = zeros(cellxmax,cellymax,4)
    cell_F_hat_minus = zeros(cellxmax,cellymax,4)

    for i in 1:cellxmax
        for j in 1:cellymax
            for l in 1:4
                for m in 1:4
                    # Ehat = Ahat * Qhat = Ahat * Q/J = 1/J*Ahat * Q
                    # 1/J*Ahatの計算をvecAを使ったため，ここではQconを掛けてる
                    cell_E_hat_plas[i,j,l]  += cell_Ahat_plas[i,j,l,m]*Qcon[i,j,m]
                    cell_E_hat_minus[i,j,l] += cell_Ahat_minus[i,j,l,m]*Qcon[i,j,m]
                    cell_F_hat_plas[i,j,l]  += cell_Bhat_plas[i,j,l,m]*Qcon[i,j,m]
                    cell_F_hat_minus[i,j,l] += cell_Bhat_minus[i,j,l,m]*Qcon[i,j,m]
                end
            end
        end
    end

    return cell_E_hat_plas,cell_E_hat_minus,cell_F_hat_plas,cell_F_hat_minus
end

function cal_jacobi(Qbase,Qcon,cellxmax,cellymax,specific_heat_ratio,vecAx,vecAy)
    """
    Qbase=[rho,u,v,p]
    Qcon=[rho,rhou,rhov,e]

    q=u^2+v^2
    c=(rp/rho)^0.5
    theta=Z=U= A^xi_x * u   +   A^xi_y * v
    H=(e+p)/rho
    b1=q^2/2 * (r-1)/(c^2)
    b2=(r-1)/(c^2)

    xi_x_bar= xi_x / (xi_x^2+xi_y^2)^0.5
    xi_y_bar= xi_y / (xi_x^2+xi_y^2)^0.5
    U_bar=Z_bar=theta_bar=

    A = Rleft * Lambda * Rright
    """

    cell_Ahat_plas  = zeros(cellxmax,cellymax,4,4)
    cell_Ahat_minus = zeros(cellxmax,cellymax,4,4)
    cell_Bhat_plas  = zeros(cellxmax,cellymax,4,4)
    cell_Bhat_minus = zeros(cellxmax,cellymax,4,4)
    r=specific_heat_ratio

    for i in 1:cellxmax
        for j in 1:cellymax
            q=(Qbase[i,j,2]^2+Qbase[i,j,3]^2)^0.5
            c=0
            try c=(r*Qbase[i,j,4]/Qbase[i,j,1])^0.5
            catch
                println("\n"*string(i)*","*string(j)*" Qbase error")
                println(Qbase[i,j,:])
                println("\n")
                throw(UndefVarError(:x))
            end

            H  = (Qcon[i,j,4]+Qbase[i,j,4])/Qbase[i,j,1]
            b1 = q^2/2 * (r-1)/(c^2)
            b2 = (r-1)/(c^2)
            u  = Qbase[i,j,2]
            v  = Qbase[i,j,3]

            k_x = 0.5*(vecAx[i,j,1]+vecAx[i+1,j,1])
            k_y = 0.5*(vecAx[i,j,2]+vecAx[i+1,j,2])

            Rleft_E,Lambda_E,Rright_E = Eigenvalue_vector(q,c,H,b1,b2,u,v,k_x,k_y)
            A_plas, A_minus = RLmbdaR(Rleft_E,Lambda_E,Rright_E)

            for l in 1:4
                for m in 1:4
                    cell_Ahat_plas[i,j,l,m]  = A_plas[l,m]
                    cell_Ahat_minus[i,j,l,m] = A_minus[l,m]
                end
            end

            k_x = 0.5*(vecAy[i,j,1]+vecAy[i,j+1,1])
            k_y = 0.5*(vecAy[i,j,2]+vecAy[i,j+1,2])

            Rleft_E,Lambda_E,Rright_E = Eigenvalue_vector(q,c,H,b1,b2,u,v,k_x,k_y)
            A_plas, A_minus = RLmbdaR(Rleft_E,Lambda_E,Rright_E)

            for l in 1:4
                for m in 1:4
                    cell_Bhat_plas[i,j,l,m]  = A_plas[l,m]
                    cell_Bhat_minus[i,j,l,m] = A_minus[l,m]
                end
            end
        end
    end

    return cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus
end

function Eigenvalue_vector(q,c,H,b1,b2,u,v,k_x,k_y)
    Z= k_x*u + k_y*v
    k_x_bar=k_x / (k_x^2+k_y^2)^0.5
    k_y_bar=k_y / (k_x^2+k_y^2)^0.5
    Z_bar=Z / (k_x^2+k_y^2)^0.5

    Lambda = [Z-c*(k_x^2+k_y^2)^0.5
              Z
              Z+c*(k_x^2+k_y^2)^0.5
              Z]

    Rleft = [[1 1 1 0]
             [u-k_x_bar*c u u+k_x_bar*c -k_y_bar]
             [v-k_y_bar*c v v+k_y_bar*c k_x_bar]
             [H-c*Z_bar q^2/2 H+c*Z_bar -(k_y_bar*u-k_x_bar*v)]]

    Rright = [[(b1+Z_bar/c)/2 -(k_x_bar/c+b2*u)/2 -(k_y_bar/c+b2*v)/2 b2/2]
              [1-b1 b2*u b2*v -b2]
              [(b1-Z_bar/c)/2 (k_x_bar/c-b2*u)/2 (k_y_bar/c-b2*v)/2 b2/2]
              [k_y_bar*u-k_x_bar*v -k_y_bar k_x_bar 0]]

    return Rleft,Lambda,Rright
end

function RLmbdaR(R,Lam,Rm)
    Lamp=zeros(4,4)
    Lamm=zeros(4,4)
    for i in 1:length(Lam)
        Lamp[i,i]=(Lam[i]+abs(Lam[i]))/2
        Lamm[i,i]=(Lam[i]-abs(Lam[i]))/2
    end
    Ap=nn_inner_product(nn_inner_product(R,Lamp),Rm)
    Am=nn_inner_product(nn_inner_product(R,Lamm),Rm)
    return Ap,Am
end

function nn_inner_product(a,b)
    temp=zeros(size(a)[1],size(a)[1])
    for i in 1:size(a)[1]
        for j in 1:size(a)[1]
            for k in 1:size(a)[1]
                temp[i,j] += a[i,k]*b[k,j]
            end
        end
    end
    return temp
end

function FVS(cell_E_hat_plas,cell_E_hat_minus,cell_F_hat_plas,cell_F_hat_minus,cellxmax,cellymax)
    E_hat = zeros(cellxmax+1,cellymax,4)
    F_hat = zeros(cellxmax,cellymax+1,4)

    Threads.@threads for i in 2:cellxmax+1-1
        for j in 2:cellymax-1
            for k in 1:4
                E_hat[i,j,k] = cell_E_hat_plas[i-1,j,k] + cell_E_hat_minus[i,j,k]
            end
        end
    end
    Threads.@threads for i in 2:cellxmax-1
        for j in 2:cellymax+1-1
            for k in 1:4
                F_hat[i,j,k] = cell_F_hat_plas[i,j-1,k] + cell_F_hat_minus[i,j,k]
            end
        end
    end

    return E_hat,F_hat
end

function setup_RHS(cellxmax,cellymax,E_hat,F_hat,Qcon_hat)
    RHS = zeros(cellxmax,cellymax,4)

    Threads.@threads for i in 2:cellxmax-1
        for j in 2:cellymax-1
            for k in 1:4
                RHS[i,j,k] = -(E_hat[i+1,j,k]-E_hat[i,j,k]+F_hat[i,j+1,k]-F_hat[i,j,k])
            end
        end
    end
    return RHS
end

function set_RHS_inner(cellxmax,cellymax,dt,Qcon_hat_n,Qcon_hat_m1,E_hat,F_hat)
    RHS = zeros(cellxmax,cellymax,4)

    Threads.@threads for i in 2:cellxmax-1
        for j in 2:cellymax-1
            for k in 1:4
                RHS[i,j,k] = -(Qcon_hat_m1[i,j,k]-Qcon_hat_n[i,j,k]) -dt*(E_hat[i+1,j,k]-E_hat[i,j,k]+F_hat[i,j+1,k]-F_hat[i,j,k])
            end
        end
    end
    return RHS
end

function setup_Qcon_hat(Qcon,cellxmax,cellymax,volume)
    Qcon_hat=ones(cellxmax,cellymax,4)
    Threads.@threads for i in 1:cellxmax
        for j in 1:cellymax
            for k in 1:4
                Qcon_hat[i,j,k] = Qcon[i,j,k]*volume[i,j]
            end
        end
    end
    return Qcon_hat
end

function Qhat_to_Q(Qcon_hat,cellxmax,cellymax,volume)
    Qcon = ones(cellxmax,cellymax,4)
    Threads.@threads for i in 1:cellxmax
        for j in 1:cellymax
            for k in 1:4
                Qcon[i,j,k] = Qcon_hat[i,j,k]/volume[i,j]
            end
        end
    end
    return Qcon
end


using CairoMakie
x(t; a = 5, b = 7, c = 1.5) = (a - b)*cos(t) + c*cos(t*(a/b - 1))
y(t; a = 5, b = 7, c = 1.5) = (a - b)*sin(t) - c*sin(t*(a/b - 1))
t = 0:0.001:20π
with_theme(theme_dark()) do
    fig, = lines(x.(t), y.(t); color = ("#08F7FE",0.8), linewidth = 0.8,
        axis = (; aspect = DataAspect(), xlabel = "x", ylabel = "y"),
        figure = (; resolution = (600,600)))
    [lines!(x.(t), y.(t); color = ("#08F7FE", 0.08), linewidth = 1 + 0.8*i)
        for i in 1:13]
    save("lines_hypotrochoid.png", fig, px_per_unit = 2)
    display(fig)
end


# thsi script calculates the barycentric coordinates for 2D and 3D
# SBP operator nodes, and prints the case statement needed by
# apf::FieldShape

using SummationByParts
include("nodecalc3.jl")

function nodecalc(sbp::TriSBP, isDG::Bool)
  vtx = [0.0 0; 1 0; 0 1]

  if isDG
    coords = SummationByParts.SymCubatures.calcnodes(sbp.cub, vtx)
  else
    coords = calcnodes(sbp, vtx)
  end

  xi = convertToBary(coords)

  return xi, coords
end

function convertToBary(coords_xy::Array)
# convert x-y coordinates to barycentric coordinates
  dim = size(coords_xy, 1)

  if dim == 2
    vtx = [0.0 0; 1 0; 0 1]
    r1 = vtx[1, :]
    r2 = vtx[2, :]
    r3 = vtx[3, :]

    T = zeros(2,2)
    T[:, 1] = r2 - r1
    T[:, 2] = r3 - r1
  elseif dim == 3

    vtx = [-1.0 -1.0 -1.0
          1.0 -1.0 -1.0
         -1.0 1.0 -1.0
         -1.0 -1.0 1.0]
    #=
    vtx = [0 0  0;
           1. 0 0;
           0 1 0;
           0 0 1]
    =#
    r1 = vtx[1, :]
    r2 = vtx[2, :]
    r3 = vtx[3, :]
    r4 = vtx[4, :]
    T = zeros(3,3)
    T[:, 1] = r2 - r1
    T[:, 2] = r3 - r1
    T[:, 3] = r4 - r1
    println("T = \n", T)
    println("r1 = \n", r1)
    println("size(T) = ", size(T))
    println("size(r1 = ", size(r1))
  end

  xi = zeros(coords_xy)
  for i=1:size(coords_xy,2)
    xi[:, i] = T\(coords_xy[:, i] - r1)
  end

  return xi
end

function nodecalc(sbp::TetSBP, isDG::Bool)

  vtx = [-1.0 -1.0 -1.0
        1.0 -1.0 -1.0
       -1.0 1.0 -1.0
       -1.0 -1.0 1.0]

  if isDG
    coords = SummationByParts.SymCubatures.calcnodes(sbp.cub, vtx)
  else
    coords = calcnodes(sbp, vtx)
  end

  xi = convertToBary(coords)

  return xi, coords
end

function minNodeDist(sbp, isDG::Bool)
# get the minimum distance between nodes on a reference element of degree p

    xi, coords = nodecalc(sbp, isDG)
    min_dist = typemax(Float64)
#    println("coords = ", coords)
    for i=1:size(coords, 2)
      for j=(i+1):size(coords, 2)
	# calculate distance between node i and node j
	dist_j = norm(coords[:, i] - coords[:, j])

	if dist_j < min_dist
	  min_dist = dist_j
	end

      end  # end loop j
    end  # end loop i

#    println("for p=$p elements, min node distance = ", min_dist)

  return min_dist
end

#sbp = getTriSBPWithDiagE(degree=4, vertices=false)
#sbp = getTetSBPWithDiagE(degree=2)a
#sbp = getTriSBPOmega2(degree=4)
#sbp = TriSBP{Float64}(degree=1, reorder=false, internal=false)
#sbp = TetSBP{Float64}(degree=3, reorder=false, internal=true)
#sbp = getTetSBPDiagE(degree=4)
sbp = getTriSBPDiagE(degree=4, mincond=false)
#sbp = getTriSBPOmega(degree=3)
xi, coords = nodecalc(sbp, true)
printCaseStatement(xi)
printBaryCoords(xi, coords)
#writedlm("coordsout.dat", coords, ' ')
#minNodeDist(2)

#=
coords = [0.5 0 0;
          0 1 0;
          0 0 1]
xi = convertToBary(coords)
println("xi = \n", xi)
=#


using Printf
using FiniteDiff
using ForwardDiff
using KernelAbstractions
using LinearAlgebra
using Random
using SparseArrays
using Test
using TimerOutputs
using ExaPF
import ExaPF: PowerSystem, AutoDiff

const PS = PowerSystem
const KA = KernelAbstractions

@testset "Compute reduced gradient on CPU" begin
    @testset "Case $case" for case in ["case9.m", "case30.m"]
        datafile = joinpath(dirname(@__FILE__), "..", "..", "data", case)
        tolerance = 1e-8
        pf = PS.PowerNetwork(datafile)

        polar = PolarForm(pf, CPU())
        cache = ExaPF.get(polar, ExaPF.PhysicalState())
        ExaPF.init_buffer!(polar, cache)

        xk = ExaPF.initial(polar, State())
        u = ExaPF.initial(polar, Control())

        jx = AutoDiff.Jacobian(polar, ExaPF.power_balance, State())
        ju = AutoDiff.Jacobian(polar, ExaPF.power_balance, Control())
        ∂obj = ExaPF.AdjointStackObjective(polar)
        pbm = AutoDiff.TapeMemory(ExaPF.active_power_generation, ∂obj, nothing)

        # solve power flow
        conv = powerflow(polar, jx, cache, NewtonRaphson(tol=1e-12))
        # No need to recompute ∇gₓ
        ∇gₓ = jx.J
        ∇gᵤ = AutoDiff.jacobian!(polar, ju, cache)
        # test jacobian wrt x
        ∇gᵥ = AutoDiff.jacobian!(polar, jx, cache)
        @test isequal(∇gₓ, ∇gᵥ)

        # Test with Matpower's Jacobian
        V = cache.vmag .* exp.(im * cache.vang)
        J = ExaPF.matpower_jacobian(polar, State(), ExaPF.power_balance, V)
        @test isapprox(∇gₓ, J)

        # Test gradients
        @testset "Reduced gradient" begin
            # Refresh cache with new values of vmag and vang
            ExaPF.update!(polar, PS.Generators(), PS.ActivePower(), cache)
            # We need uk here for the closure
            uk = copy(u)
            ExaPF.gradient_objective!(polar, pbm, cache)
            ∇fₓ = ∂obj.∇fₓ
            ∇fᵤ = ∂obj.∇fᵤ

            ## ADJOINT
            # lamba calculation
            λk  = -(∇gₓ') \ ∇fₓ
            grad_adjoint = ∇fᵤ + ∇gᵤ' * λk
            # ## DIRECT
            S = - inv(Array(∇gₓ)) * ∇gᵤ
            grad_direct = ∇fᵤ + S' * ∇fₓ
            @test isapprox(grad_adjoint, grad_direct)

            # Compare with finite difference
            function reduced_cost(u_)
                # Ensure we remain in the manifold
                ExaPF.transfer!(polar, cache, u_)
                convergence = powerflow(polar, jx, cache, NewtonRaphson(tol=1e-14))
                ExaPF.update!(polar, PS.Generators(), PS.ActivePower(), cache)
                return ExaPF.cost_production(polar, cache)
            end

            grad_fd = FiniteDiff.finite_difference_gradient(reduced_cost, uk)
            @test isapprox(grad_fd, grad_adjoint, rtol=1e-4)
        end
        @testset "Gradient of line-flow constraints" begin
            # Adjoint of flow_constraints()
            nbus = length(cache.vmag)
            M = typeof(u)
            m = ExaPF.size_constraint(polar, ExaPF.flow_constraints)
            x = ExaPF.xzeros(M, 2 * nbus)
            x[1:nbus] .= cache.vmag
            x[1+nbus:2*nbus] .= cache.vang
            bus_gen = polar.indexing.index_generators
            VI = typeof(bus_gen)

            ## Example with using sum as a sort of lumping of all constraints
            function sum_constraints(x)
                VT = typeof(x)
                # Needed for ForwardDiff to have a cache with the right active type VT
                adcache = ExaPF.PolarNetworkState{VI, VT}(
                    cache.vmag, cache.vang, cache.pinj, cache.qinj,
                    cache.pg, cache.qg, cache.balance, cache.dx, bus_gen,
                )
                adcache.vmag .= x[1:nbus]
                adcache.vang .= x[1+nbus:2*nbus]
                g = VT(undef, m) ; fill!(g, 0)
                ExaPF.flow_constraints(polar, g, adcache)
                return sum(g)
            end
            adgradg = ForwardDiff.gradient(sum_constraints,x)
            fdgradg = FiniteDiff.finite_difference_gradient(sum_constraints,x)
            ## We pick sum() as the reduction function. This could be a mask function for active set or some log(x) for lumping.
            m_flows = ExaPF.size_constraint(polar, ExaPF.flow_constraints)
            weights = ones(m_flows)
            gradg = ExaPF.xzeros(M, 2 * nbus)
            ExaPF.flow_constraints_grad!(polar, gradg, cache, weights)
            @test isapprox(adgradg, fdgradg)
            # TODO: The gradient is slightly off with the handwritten adjoint
            @test_broken isapprox(gradg, fdgradg)
        end
    end
end


using Distances
using UMAP
using NearestNeighbors
using Statistics

import MultivariateStats
import MultivariateStats.transform
import MultivariateStats.fit

struct UmapFit{TT <: NNTree}
    embedding::Matrix{<:Real};
    nn_tree::TT;
    pca_transform::MultivariateStats.PCA;
    nn_interpolate::Int;
end

function fit(::Type{UmapFit}, x::Array{Float64, 2}, pca::MultivariateStats.PCA; n_components::Int=2, nn_interpolate::Int=5,
        spread=2.0, min_dist=0.1, metric::MT where MT <: NearestNeighbors.MinkowskiMetric = Euclidean(), kwargs...)::UmapFit
    x_pc = transform(pca, x)
    nn_tree = KDTree(x_pc, metric)
    embedding = umap(x_pc, n_components; spread=spread, min_dist=min_dist, metric=metric, kwargs...)

    return UmapFit(embedding, nn_tree, pca, nn_interpolate)
end

fit(::Type{UmapFit}, x::Array{Float64, 2}; n_pcs::Int=3, kwargs...)::UmapFit =
    fit(UmapFit, x, fit(MultivariateStats.PCA, x, maxoutdim=n_pcs); kwargs...)

function transform(transformation::UmapFit, x::Array{Float64, 2}; dist_offset::Float64=1e-10)::Array{Float64, 2}
    x_pc = transform(transformation.pca_transform, x)

    # This transformation is much faster than the implementation from UMAP.jl, even for input dimensionality = 50.
    # Probably, because of the slow NearestNeighborDescent package.
    indices, distances = knn(transformation.nn_tree, x_pc, transformation.nn_interpolate)
    res = [mapslices(x -> wmean(x, 1 ./ (dists .+ dist_offset)), transformation.embedding[:, ids], dims=2) for (ids, dists) in zip(indices, distances)]
    return hcat(res...)
end


module NuclearAttraction

#Julia external modules
using GSL

using TypesBasis,TypesParticles
using UtilityFunctions

export buildnuclrattract

@doc raw"""
function buildnuclrattract(system::)

This function applies the Coulomb operator between nuclei and electronic volume, i.e., the volume
integral of the electron density.

```
$V = int d\tau \frac{-Z_{i}}{r_{i,n}\phi^2_{n}(r_{j})$
```

This function will only support atom centered basis functions and thus the potential is:

```
V_{i,j} = \sum_N - Z_n \int\frac{\phi_i \phi_j}{r_{i,j}} d\tau

```

The overlap $phi_i \phi_j$ can be written as $\phi_p = N_p e^{-pr_{p}^2} and the integral becomes
$\int\frac{\phi_p}{r_{i,j}}$

The approach to solving this for spatial integral of gaussian functions is to use Boys functions:

```
#\int \frac{\phi_p}{r_{i,j}} = \frac{2\pi C_p}{p} F_o(pR^2_{NP})
```

Input:
Results:

module requirements:

"""
function buildnuclrattract(system::SystemOfAtoms,basisfunc::Array{GaussOrbitals},basis::Basis)
    natoms = system.natoms;
    numbasisfunc = basis.nbasisfunc; #number of basis functions
    basissize = natoms*numbasisfunc;
    nuclrattract = zeros(Float64,basissize,basissize);

    flatbasisfunc = flattenbasisfunc(natoms,numbasisfunc,basisfunc);

    for i=1:natoms
        atom = system.atoms[i];
        for n=1:basissize
            for m=1:basissize

                #Basis functions with primitives
                basisn = flatbasisfunc[n];
                basism = flatbasisfunc[m];

                #Contraction coeffs.
                cn = basisn.coefs;
                cm = basism.coefs;

                #Number of primitive Gaussians
                primsn = length(cn);
                primsm = length(cm);

                for np=1:primsn
                    for mp=1:primsm
                        cnm = cn[np]*cm[mp];
                        nuclrattract[n,m] += cnm * coulomb(np,mp,basisn,basism,atom);
                    end
                end #primsn

            end
        end #basissize
    end #system.natoms
    return nuclrattract

end #buildnuclrattract

@doc raw"""
function coulomb(np::Int,mp::Int,basisn::Gaussian
coulomb potential

"""
function coulomb(np::Int,mp::Int,basisn::GaussOrbitals,
                 basism::GaussOrbitals,atom::Atom)

    xn,yn,zn = basisn.ax,basisn.ay,basisn.az;
    xm,ym,zm = basism.ax,basism.ay,basism.az;

    alphanp = basisn.alphas[np];
    alphamp = basism.alphas[mp];

    normcont = basisn.norms[np] * basism.norms[mp];

    p = alphanp + alphamp;
    px = (alphanp*xn + alphamp*xm)/p;
    py = (alphanp*yn + alphamp*ym)/p;
    pz = (alphanp*zn + alphamp*zm)/p;
    orbatomdist = (atom.x-px)^2 + (atom.y-py)^2 + (atom.z-pz)^2;

    spatialx = gaussprod1D(xn,alphanp,xm,alphamp);
    spatialy = gaussprod1D(yn,alphanp,ym,alphamp);
    spatialz = gaussprod1D(zn,alphanp,zm,alphamp);

    spatial = spatialx*spatialy*spatialz;
    prefact = (-1*atom.Z*2*pi) / p;

    ##WARNING! call is for spherical Gaussians only, x=0.
    potential = prefact * normcont * spatial * boysfunction(p*orbatomdist,n=0);

    return potential

end #coulomb

end #NuclearAtraction


estimates = collect(KalmanFilter(lssm, lobs), measures)
mu = [e[1] for e in estimates]
emp_var = empirical_var(mu-states,n,k)
@test emp_var < n*2.