using HypothesisTests, Test

@testset "Clark-West tests" begin

    e1 = [ 9.78432,  12.73500,   8.67224,   2.62740,   5.60947,
           7.57809,   4.53774,   2.51084,   0.49290,   3.48394,
           9.46152,   9.41220,   2.36289,   0.34495,   3.33599,
           5.31357,   4.28219,  -3.74471,  -5.73575,  -3.71781 ]

    e2 =  [ 2.82053,   4.39754,  -1.78647,  -4.30662,   3.69526,
            3.37259,  -1.09002,  -0.50984,  -0.78973,   3.89077,
            7.52849,   2.82373,  -3.95838,  -0.13606,   4.54444,
            4.18216,   1.67993,  -5.38077,  -0.85686,   2.70479 ]

    atol = 1e-4
    cw_test = ClarkWestTest(e1, e2)
    @test pvalue(cw_test) ≈ 0.0002 atol=atol
    @test pvalue(cw_test, tail=:right) ≈ 0.0001 atol=atol

    cw_test = ClarkWestTest(e1, e2, 3)
    @test pvalue(cw_test) ≈ 0.0157 atol=atol
    @test pvalue(cw_test, tail=:right) ≈ 0.0079 atol=atol

    show(IOBuffer(), cw_test)

    @test_throws DimensionMismatch ClarkWestTest(rand(3), rand(4))

end


# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

using MathOptInterface
import GLPK
import Clp
const MOI = MathOptInterface

function example_diet(optimizer, bridge)
    category_data = [
        1800.0 2200.0
          91.0    Inf
           0.0   65.0
           0.0 1779.0
    ]
    cost = [2.49, 2.89, 1.50, 1.89, 2.09, 1.99, 2.49, 0.89, 1.59]
    food_data = [
        410 24 26 730
        420 32 10 1190
        560 20 32 1800
        380  4 19 270
        320 12 10 930
        320 15 12 820
        320 31 12 1230
        100  8 2.5 125
        330  8 10 180
    ]
    bridge_model = if bridge
        MOI.instantiate(optimizer; with_bridge_type=Float64)
    else
        MOI.instantiate(optimizer)
    end
    model = MOI.Utilities.CachingOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
        MOI.Utilities.AUTOMATIC,
    )
    MOI.Utilities.reset_optimizer(model, bridge_model)
    MOI.set(model, MOI.Silent(), true)
    nutrition = MOI.add_variables(model, size(category_data, 1))
    for (i, v) in enumerate(nutrition)
        MOI.add_constraint(model, v, MOI.GreaterThan(category_data[i, 1]))
        MOI.add_constraint(model, v, MOI.LessThan(category_data[i, 2]))
    end
    buy = MOI.add_variables(model, size(food_data, 1))
    MOI.add_constraint.(model, buy, MOI.GreaterThan(0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    f = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(cost, buy), 0.0)
    MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
    for (j, n) in enumerate(nutrition)
        f = MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.(food_data[:, j], buy),
            0.0,
        )
        push!(f.terms, MOI.ScalarAffineTerm(-1.0, n))
        MOI.add_constraint(model, f, MOI.EqualTo(0.0))
    end
    MOI.optimize!(model)
    term_status = MOI.get(model, MOI.TerminationStatus())
    @assert term_status == MOI.OPTIMAL
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.(1.0, [buy[end-1], buy[end]]),
            0.0,
        ),
        MOI.LessThan(6.0),
    )
    MOI.optimize!(model)
    @assert MOI.get(model, MOI.TerminationStatus()) == MOI.INFEASIBLE
    return
end

if length(ARGS) > 0
    bridge = get(ARGS, 2, "") != "--no-bridge"
    println("Running: $(ARGS[1]) $(get(ARGS, 2, ""))")
    if ARGS[1] == "clp"
        @time example_diet(Clp.Optimizer, bridge)
        @time example_diet(Clp.Optimizer, bridge)
    else
        @assert ARGS[1] == "glpk"
        @time example_diet(GLPK.Optimizer, bridge)
        @time example_diet(GLPK.Optimizer, bridge)
    end
    exit(0)
end

# using SnoopCompileCore
# tinf = @snoopi_deep example_diet(Clp.Optimizer, true)
# using SnoopCompile, ProfileView
# fg = flamegraph(tinf)
# ProfileView.view(fg)


import Random

perturbation(step, sig) = -sig + 2*sig*rand()

function initial_point(N, side, tol; seed=123)

	if seed > 0
		Random.seed!(seed)
	else
		Random.seed!()
	end

	p = [ zeros(2) for i in 1:N ]

	l = ceil(Int, sqrt(N))
	step = (side-tol)/l

	if tol > step
		error("tol must be smaller than $step")
	end

	skip = Int[]

	while length(skip) < (l^2-N)

		i = rand(1:l^2)

		if findfirst(isequal(i), skip) == nothing
			push!(skip, i)
		end
	end

	sig = (step-tol)/2

	ip = 0
	igrid = 0
	x = [ tol/2, tol/2 ]

	for i in 1:l
		for j in 1:l

			igrid += 1

			if findfirst(isequal(igrid), skip) != nothing
				continue
			end

			ip += 1
			p[ip] .= x .+ perturbation(step, sig)
			x[2] += step
		end

		x[2] = tol/2
		x[1] += step
	end

	return p
end

initial_point(data :: Data; tol=0.9, seed=123) = initial_point(data.N, data.side, tol, seed=seed)


# ---
# title: Faceting
# cover: assets/faceting.png
# description: Generating a grid of plots
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
set_aog_theme!() #src

# ## Facet grid

df = (x=rand(100), y=rand(100), i=rand(["a", "b", "c"], 100), j=rand(["d", "e", "f"], 100))
plt = data(df) * mapping(:x, :y, col=:i, row=:j)
draw(plt)

# ## Facet wrap

df = (x=rand(100), y=rand(100), l=rand(["a", "b", "c", "d", "e"], 100))
plt = data(df) * mapping(:x, :y, layout=:l)
draw(plt)

# ## Adding traces to only some subplots

df1 = (x=rand(100), y=rand(100), i=rand(["a", "b", "c"], 100), j=rand(["d", "e", "f"], 100))
df2 = (x=[0, 1], y=[0.5, 0.5], i=fill("a", 2), j=fill("e", 2))
layers = data(df1) * visual(Scatter) + data(df2) * visual(Lines)
fg = draw(layers * mapping(:x, :y, col=:i, row=:j))

# save cover image #src
mkpath("assets") #src
save("assets/faceting.png", fg) #src


struct KeypointTensor{N, T, M} <: Block
    sz::NTuple{M, Int}
end

mockblock(block::KeypointTensor{N}) where N = rand(SVector{N, Float32}, block.sz)
function checkblock(block::KeypointTensor{N, T}, data::AbstractArray{T}) where {N, T}
    return length(data) == (prod(block.sz) * N)
end

"""
    KeypointPreprocessing(bounds) <: Encoding

Scale a `Keypoints` block falling in a rectangle of `bounds` so
that they lie between -1 and 1.
"""
struct KeypointPreprocessing{N, T<:Number} <: Encoding
    bounds::NTuple{N, Int}
end
KeypointPreprocessing(bounds::NTuple{N, Int}) where N = KeypointPreprocessing{N, Float32}(bounds)

function encode(enc::KeypointPreprocessing{N, T}, context, block::Keypoints{N}, data) where {N, T}
    ks = map(k -> (SVector{N, T}(k) .* (convert(T, 2) ./ enc.bounds)) .- one(T), data)
    return reinterpret(T, ks)
end

function decode(enc::KeypointPreprocessing{N, T}, context, block::KeypointTensor{N}, data) where {N, T}
    ks = reshape(reinterpret(SVector{N, T}, data), block.sz)
    return map(k -> ((k) .+ one(T)) ./ (convert(T, 2) ./ SVector{N, T}(enc.bounds)), ks)
end

encodedblock(::KeypointPreprocessing{N, T}, block::Keypoints{N, M}) where {N, T, M} = KeypointTensor{N, T, M}(block.sz)
decodedblock(::KeypointPreprocessing{N}, block::KeypointTensor{N}) where N = Keypoints{N}(block.sz)


"""
    vehicle_miles_data()

Returns a 15 x 12 matrix with the vehicle miles traveled in the US
(in millions), per month, for the years 2000, ..., 2014 (discussed on page 252).
"""
function vehicle_miles_data()
# Vehicle Miles Traveled (VMT) (Millions)
# US Dept of Transportation, Bureau of Transportation Statistics
# www.transtats.bts.gov

# 15 x 12 matrix: monthly value for 15 years 2000-2014
vmt = [
203442 199261 232490 227698 242501 242963 245140 247832 227899 236491 222819 218390;  # 2000
209685 200876 232587 232513 245357 243498 250363 253274 226312 241050 230511 229584;  # 2001
215215 208237 236070 237226 251746 247868 256392 258666 233625 245556 230648 234260;  # 2002
218534 203677 236679 239415 253244 252145 262105 260687 237451 254048 233698 238538;  # 2003
222450 213709 251403 250968 257235 257383 265969 262836 243515 254496 239796 245029;  # 2004
224072 219970 253182 250860 262678 263816 267025 265323 242240 251419 243056 245787;  # 2005
233302 220730 256645 250665 263393 263805 263442 265229 245624 257961 245367 248208;  # 2006
233799 219221 259740 252734 267646 265475 267179 271401 246050 261505 245928 240444;  # 2007
233469 221728 252773 252699 261890 256152 262152 261228 238701 256402 237009 242326;  # 2008
224840 218031 247433 251481 258793 258487 265026 260838 242034 252683 237342 239774;  # 2009
220177 210968 251858 254014 257401 260159 265861 264358 244712 256867 239656 240932;  # 2010
222724 213547 250410 249309 254145 258025 260317 260623 241764 252058 238278 244615;  # 2011
226834 218714 253785 249567 261355 260534 260880 264983 239001 254170 240734 238876;  # 2012
228607 216306 250496 252116 263923 260023 264570 268609 242582 259281 240146 241365;  # 2013
226444 215166 252089 257947 268075 264868 272335 271018 249125 267185 242816 253618 ]   # 2014

return vmt
end


# MobilityTheories.jl

"""
    polaronmobility(Trange,
                    ε_Inf, ε_S, freq, effectivemass;
                    verbose::Bool=false)

    Solves the Feynman polaron problem variationally with finite temperature
    Osaka energies.  From the resulting v, and w parameters, calculates polaron
    structure (wave function size, etc.).  Uses FHIP, Kadanoff (Boltzmann
    relaxation time) and Hellwarth direct contour integration to predict
    a temperature-dependent mobility for the material system.
    Input is a temperature range (e.g. 10:50:1000),
    reduced dielectric constants (e.g. 5, 20),
    characteristic dielectric phonon frequency (e.g. 2.25E12) - units Hertz
    bare-band effective-mass (e.g. 012) - units electron mass.

    Returns a structure of type Polaron, containing arrays of useful
    information.  Also prints a lot of information to the standard out - which
    may be more useful if you're just inquiring as to a particular data point,
    rather than plotting a temperature-dependent parameter.

    As an example, to calculate the electron polaron in MAPI at 300 K:
    polaronmobility(300, 4.5, 24.1, 2.25E12, 0.12)
"""
function polaronmobility(Trange, ε_Inf, ε_S, freq, effectivemass; verbose::Bool=false)
    println("\n\nPolaron mobility for system ε_Inf=$ε_Inf, ε_S=$ε_S, freq=$freq,
                 effectivemass=$effectivemass; with Trange $Trange ...")

    # Internally we we use 'mb' for the 'band mass' in SI units, of the effecitve-mass of the electron
    mb=effectivemass*MassElectron
    ω = (2*pi)*freq # angular-frequency, of the phonon mode

    α=frohlichalpha(ε_Inf, ε_S,  freq,    effectivemass)
    #α=2.395939683378253 # Hard coded; from MAPI params, 4.5, 24.1, 2.25THz, 0.12me

    v=7.1 # starting guess for v,w variational parameters
    w=6.5

    @printf("Polaron mobility input parameters: ε_Inf=%f ε_S=%f freq=%g α=%f \n",ε_Inf, ε_S, freq, α )
    @printf("Derived params in SI: ω =%g mb=%g \n",ω ,mb)

    # Empty struct for storing data
    # A slightly better way of doing this ф_ф ...
    p=Polaron()
    # populate data structure with (athermal) parameters supplied...
    append!(p.α,α) # appending so as not to mess with type immutability
    append!(p.mb,mb)
    append!(p.ω,ω)

    # We define βred as the subsuming the energy of the phonon; i.e. kbT c.f. ħω
    for T in Trange
        β=1/(kB*T)
        βred=ħ*ω*β
        append!(p.βred,βred)
        @printf("T: %f β: %.3g βred: %.3g ħω  = %.3g meV\t",T,β,βred, 1E3* ħ*ω  / q)

        if T==0
            v,w=feynmanvw(α, v=v,w=w)
        else
            v,w=feynmanvw(α,βred, v=v,w=w)
        end

        @printf("\n Polaron Parameters:  v= %.4f  w= %.4f",v,w)

        # From 1962 Feynman, definition of v and w in terms of the coupled Mass and spring-constant
        # See Page 1007, just after equation (18)
        # Units of M appear to be 'electron masses'
        # Unsure of units for k, spring coupling constant
        k=(v^2-w^2)
        M=(v^2-w^2)/w^2
        append!(p.k,k)
        append!(p.M,M)

        @printf("  ||   M=%f  k=%f\t",M,k)
        @printf("\n Bare-band effective mass: %f Polaron effective mass: %f Polaron mass enhancement: %f%%",effectivemass,effectivemass*(1+M),M*100)

        @printf("\n Polaron frequency (SI) v= %.2g Hz  w= %.2g Hz",
                v*ω /(2*pi),w*ω /(2*pi))

		# (46) in Feynman1955
		meSmallAlpha(α )=α /6 + 0.025*α ^2
		# (47) In Feynman1955
		meLargeAlpha(α )=16*α ^4 / (81*π ^4)
		#meLargeAlpha(α )=202*(α /10)^4
        if (verbose) # asymptotic solutions - not that interesting when you have the actual ones!
		    @printf("\n Feynman1955(46,47): meSmallAlpha(α)= %.3f meLargeAlpha(α)= %.3f",
                    meSmallAlpha(α),meLargeAlpha(α))
            @printf("\n Feynman1962: Approximate ~ Large alpha limit, v/w = %.2f  =~approx~= alpha^2 = %.2f ",
                    v/w,α ^2)
        end

        # POLARON SIZE
        @printf("\n Polaron size (rf), following Schultz1959. (s.d. of Gaussian polaron ψ )")
        # Schultz1959 - rather nicely he actually specifies everything down into units!
        # just before (2.4) in Schultz1959
        mu=((v^2-w^2)/v^2)
        # (2.4)
        rf=sqrt(3/(2*mu*v))
        # (2.4) SI scaling inferred from units in (2.5a) and Table II
        rfsi=rf*sqrt(2*me*ω )
        @printf("\n\t Schultz1959(2.4): rf= %g (int units) = %g m [SI]",rf,rfsi )
        append!(p.rfsi,rfsi)

        if (verbose)
        scale=sqrt(2*mb*ω) # Note we're using mb;
        #band effective-mass in SI units (i.e. meff*melectron)

        rfa=(3/(0.44*α ))^0.5 # As given in Schultz1959(2.5a), but that 0.44 is actually 4/9
        @printf("\n\t Schultz1959(2.5a) with 0.44: Feynman α→0 expansion: rfa= %g (int units) = %g m [SI]",rfa,scale*rfa )
        rfa=(3/((4/9)*α ))^0.5 # Rederived from Feynman1955, 8-8-2017; Yellow 2017.B Notebook pp.33-34
        @printf("\n\t Schultz1959(2.5a) with 4/9 re-derivation: Feynman α→0 expansion: rfa= %g (int units) = %g m [SI]",rfa,scale*rfa )
        append!(p.rfsmallalpha,scale*rfa)

        rfb=3*(pi/2)^0.5 * α
        @printf("\n\t Schultz1959(2.5b): Feynman α→∞ expansion: rf= %g (int units) = %g m [SI]",rfb,scale*rfb )

        # Schultz1959 - Between (5.7) and (5.8) - resonance of Feynman SHM system
        phononfreq=sqrt(k/M)
        @printf("\n\t Schultz1959(5.7-5.8): fixed-e: phononfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            phononfreq,phononfreq*ω /(2*pi), phononfreq*hbar*ω *1000/q)

        phononfreq=sqrt(k/mu) # reduced mass
        @printf("\n\t Schultz1959(5.7-5.8): reducd mass: phononfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            phononfreq,phononfreq*ω /(2*pi), phononfreq*hbar*ω *1000/q)

        @printf("\n\t Schultz1959: electronfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            sqrt(k/1),sqrt(k/1)*ω /(2*pi), sqrt(k/1)*hbar*ω *1000/q)
        @printf("\n\t Schultz1959: combinedfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            sqrt(k/(1+M)),sqrt(k/(1+M))*ω /(2*pi), sqrt(k/(1+M))*hbar*ω *1000/q)

        # Devreese1972: 10.1103/PhysRevB.5.2367
        # p.2371, RHS.
        @printf("\n Devreese1972: (Large Alpha) Franck-Condon frequency = %.2f", 4*α ^2/(9*pi))
        end

        # F(v,w,β,α)=-(A(v,w,β)+B(v,w,β,α)+C(v,w,β)) #(62a) - Hellwarth 1999
        if (T>0)
            @printf("\n Polaron Free Energy: A= %f B= %f C= %f F= %f",A(v,w,βred),B(v,w,βred,α),C(v,w,βred),F(v,w,βred,α))
            @printf("\t = %f meV",1000.0 * F(v,w,βred,α) * ħ*ω  / q)
            append!(p.A,A(v,w,βred))
            append!(p.B,B(v,w,βred,α))
            append!(p.C,C(v,w,βred))
            append!(p.F,F(v,w,βred,α))
        else # Athermal case; Enthalpy
            @printf("\n Polaron Enthalpy: F= %f = %f meV \n",F(v,w,α), 1_000*F(v,w,α) * ħ*ω/q)

            return # return early, as if T=0, all mobility theories = infinite / fall over
        end

        # FHIP
        #    - low-T mobility, final result of Feynman1962
        # [1.60] in Devreese2016 page 36; 6th Edition of Frohlich polaron notes (ArXiv)
        # I believe here β is in SI (expanded) units
        @printf("\n Polaron Mobility theories:")
        μ=(w/v)^3 * (3*q)/(4*mb*ħ*ω^2*α*β) * exp(ħ*ω*β) * exp((v^2-w^2)/(w^2*v))
        @printf("\n\tμ(FHIP)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)
        append!(p.T,T)
        append!(p.FHIPμ,μ*100^2)

        # Kadanoff
        #     - low-T mobility, constructed around Boltzmann equation.
        #     - Adds factor of 3/(2*beta) c.f. FHIP, correcting phonon emission behaviour
        # [1.61] in Devreese2016 - Kadanoff's Boltzmann eqn derived mob
        μ=(w/v)^3 * (q)/(2*mb*ω*α) * exp(ħ*ω*β) * exp((v^2-w^2)/(w^2*v))
        @printf("\n\tμ(Kadanoff,via Devreese2016)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)

        append!(p.Kμ,μ*100^2)

        ######
        # OK, now deep-diving into Kadanoff1963 itself to extract Boltzmann equation components
        # Particularly right-hand-side of page 1367
        #
        # From 1963 Kadanoff, (9), define eqm. number of phonons (just from T and phonon omega)
        Nbar=(exp(βred) - 1)^-1
        @printf("\n\t\tEqm. Phonon. pop. Nbar: %f ",Nbar)
        if (verbose)
            @printf("\n\texp(Bred): %f exp(-Bred): %f exp(Bred)-1: %f",exp(βred),exp(-βred),exp(βred)-1)
        end
        Nbar=exp(-βred)
        #Note - this is only way to get Kadanoff1963 to be self-consistent with
        #FHIP, and later statements (Devreese) of the Kadanoff mobility.
        #It suggests that Kadanoff used the wrong identy for Nbar in 23(b) for
        #the Gamma0 function, and should have used a version with the -1 to
        #account for Bose / phonon statistics

        #myv=sqrt(k*(1+M)/M) # cross-check maths between different papers
        #@printf("\nv: %f myv: %f\n",v,myv)

        # Between 23 and 24 in Kadanoff 1963, for small momenta skip intergration --> Gamma0
        Gamma0=2*α * Nbar * (M+1)^(1/2) * exp(-M/v)
        Gamma0*=ω  #* (ω *hbar) # Kadanoff 1963 uses hbar=omega=mb=1 units
            # Factor of omega to get it as a rate relative to phonon frequency
            # Factor of omega*hbar to get it as a rate per energy window
        μ=q/( mb*(M+1) * Gamma0 ) #(25) Kadanoff 1963, with SI effective mass

        if (verbose) # these are cross-checks
            @printf("\n\tμ(Kadanoff1963 [Eqn. 25]) = %f m^2/Vs \t = %.2f cm^2/Vs",μ,μ*100^2)
            @printf("\n\t\t Eqm. Phonon. pop. Nbar: %f ",Nbar)
        end

        @printf("\n\t\tGamma0 = %g rad/s = %g /s ",
                Gamma0, Gamma0/(2*pi))
        @printf(" \n\t\tTau=1/Gamma0 = %g s = %f ps",
            2*pi/Gamma0, 2*pi*1E12/Gamma0)
        Eloss=hbar*ω * Gamma0/(2*pi) # Simply Energy * Rate
        @printf("\n\t\tEnergy Loss = %g J/s = %g meV/ps",Eloss,Eloss * 1E3 / (q*1E12) )
        append!(p.Tau, 2*pi*1E12/Gamma0) # Boosted into ps ?

        # Hellwarth1999 - directly do contour integration in Feynman1962, for
        # finite temperature DC mobility
        # Hellwarth1999 Eqn (2) and (1) - These are going back to the general
        # (pre low-T limit) formulas in Feynman1962.  to evaluate these, you
        # need to do the explicit contour integration to get the polaron
        # self-energy
        R=(v^2-w^2)/(w^2*v) # inline, page 300 just after Eqn (2)

        #b=R*βred/sinh(b*βred*v/2) # This self-references b! What on Earth?
        # OK! I now understand that there is a typo in Hellwarth1999 and
        # Biaggio1997. They've introduced a spurious b on the R.H.S. compared to
        # the original, Feynman1962:
        b=R*βred/sinh(βred*v/2) # Feynman1962 version; page 1010, Eqn (47b)

        a=sqrt( (βred/2)^2 + R*βred*coth(βred*v/2))
        k(u,a,b,v) = (u^2+a^2-b*cos(v*u))^(-3/2)*cos(u) # integrand in (2)
        K=quadgk(u->k(u,a,b,v),0,Inf)[1] # numerical quadrature integration of (2)

        #Right-hand-side of Eqn 1 in Hellwarth 1999 // Eqn (4) in Baggio1997
        RHS= α/(3*sqrt(π)) * βred^(5/2) / sinh(βred/2) * (v^3/w^3) * K
        μ=RHS^-1 * (q)/(ω*mb)
        @printf("\n\tμ(Hellwarth1999)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)
        append!(p.Hμ,μ*100^2)

        if (verbose)
        # Hellwarth1999/Biaggio1997, b=0 version... 'Setting b=0 makes less than 0.1% error'
        # So let's test this
        R=(v^2-w^2)/(w^2*v) # inline, page 300 just after Eqn (2)
        b=0
        a=sqrt( (βred/2)^2 + R*βred*coth(βred*v/2))
        #k(u,a,b,v) = (u^2+a^2-b*cos(v*u))^(-3/2)*cos(u) # integrand in (2)
        K=quadgk(u->k(u,a,b,v),0,Inf)[1] # numerical quadrature integration of (2)

        #Right-hand-side of Eqn 1 in Hellwarth 1999 // Eqn (4) in Baggio1997
        RHS= α/(3*sqrt(π)) * βred^(5/2) / sinh(βred/2) * (v^3/w^3) * K
        μ=RHS^-1 * (q)/(ω*mb)
        @printf("\n\tμ(Hellwarth1999,b=0)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)
        @printf("\n\tError due to b=0; %f",(100^2*μ-p.Hμ[length(p.Hμ)])/(100^2*μ))
        #append!(Hμs,μ*100^2)
        end
        @printf("\n") # blank line at end of spiel.

        # Recycle previous variation results (v,w) as next guess
        initial=[v,w] # Caution! Might cause weird sticking in local minima

        append!(p.v,v)
        append!(p.w,w)

    end

    return(p)
end

"""
    savepolaron(fileprefix, p::Polaron)

Saves data from polaron 'p' into file "fileprefix".
This is a simple space-delimited text file, with each entry a separate temperature, for plotting with Gnuplot or similar.

Structure of file is written to the header:
# Ts, βreds, Kμs, Hμs, FHIPμs, vs, ws, ks, Ms, As, Bs, Cs, Fs, Taus, rfsis
# 1    2     3    4     5      6   7   8   9  10  11  12  13    14     15
"""
function savepolaron(fileprefix, p::Polaron)
    println("Saving data to $fileprefix.dat ...")
    f=open("$fileprefix.dat","w")

    @printf(f,"# %s\n",fileprefix) # put name / material at header
    @printf(f,"# Params in SI: ω =%g mb=%g \n",p.ω[1] ,p.mb[1])
    @printf(f,"# Alpha parameter: α = %f  \n",p.α[1] )

    @printf(f,"# Ts, βreds, Kμs, Hμs, FHIPμs, vs, ws, ks, Ms, As, Bs, Cs, Fs, Taus, rfsis\n")
    @printf(f,"#  1    2     3    4     5      6   7   8   9  10  11  12  13    14     15\n") # columns for GNUPLOT etc.

    for i in 1:length(p.T)
        @printf(f,"%d %03f %g %g %g %g %g %g %g %g %g %g %g %g %g \n",
        p.T[i], p.βred[i], p.Kμ[i], p.Hμ[i], p.FHIPμ[i],
        p.v[i], p.w[i],
        p.k[i], p.M[i], p.A[i], p.B[i], p.C[i], p.F[i],
        p.Tau[i], p.rfsi[i])
    end
    close(f)
end


using Plots,ApproxFun

#The following samples eigenvalues of an n = 2 Gaussian Unitary Ensemble:

ff=(x,y)->(x-y)^2*exp(-x^2/2.-y^2/2)
f=Fun(ff,[-4.,4.],[-4.,4.])
r=ApproxFun.sample(f,5000)

#We can compare the histogram to the 1-point correlation
plot(sum(f,1)/sum(f))
histogram!(vec(r);normed=true)

#We can compare the histograms of x with the GUE:
using RandomMatrices

histogram(vcat([sqrt(2)eigvalrand(GaussianHermite(2),2) for k=1:10000]...);normed=true)
plot!(sum(f,1)/sum(f))


# Pkg.add.(["TensorFlow", "Distributions", "ProgressMeter", "MLLabelUtils", "MLDataUtils"])
# # Pkg.add("MLDatasets")  # not registered yet
# Pkg.clone("https://github.com/JuliaML/MLDatasets.jl.git")
# Pkg.checkout("MLDataUtils", "dev")
# using Conda
# # Conda.add("tensorflow")
# # using PyCall
# # @pyimport Conda
# # @show Conda.PYTHONDIR

using Base.Test
using Distributions
using MLLabelUtils
using MLDataUtils
using MLDatasets
using ProgressMeter
using TensorFlow

#Training Hyper Parameter
const learning_rate = 0.001
const training_iters = 2 #Just two, becuase I don't have anything to stop overfitting and I don't got all day
const batch_size = 256
const display_step = 100 #How often to display the

# Network Parameters
const n_input = 28 # MNIST data input (img shape: 28*28)
const n_steps = 28 # timesteps
const n_hidden = 128 # hidden layer num of features
const n_classes = 10; # MNIST total classes (0-9 digits)

const traindata_raw, trainlabels_raw = MNIST.traindata();
@show size(traindata_raw)
@show size(trainlabels_raw)

imshow(x) = join(mapslices(join, (x->x ? 'X': ' ').(x'.>0), 2), "\n") |> print
@show trainlabels_raw[8]
imshow(traindata_raw[:,:,8])

"""Makes 1 hot, row encoded labels."""
encode(labels_raw) = convertlabel(LabelEnc.OneOfK, labels_raw, LabelEnc.NativeLabels(collect(0:9)),  LearnBase.ObsDim.First())

"""Prepares the data by encoding the labels and batching"""
prepared_batchs(data_raw, labels_raw) = eachbatch((data_raw, encode(labels_raw)), #Will zip these
                               batch_size,
                               (MLDataUtils.ObsDim.Last(), MLDataUtils.ObsDim.First())) #Slicing dimentions for each

@testset "data prep" begin

    @test encode([4,1,2,3,0]) == [0 0 0 0 1 0 0 0 0 0
                                  0 1 0 0 0 0 0 0 0 0
                                  0 0 1 0 0 0 0 0 0 0
                                  0 0 0 1 0 0 0 0 0 0
                                  1 0 0 0 0 0 0 0 0 0]

    data_b1, labels_b1 = first(prepared_batchs(traindata_raw, trainlabels_raw))
    @test size(data_b1) == (n_steps, n_input, batch_size)
    @test size(labels_b1) == (batch_size, n_classes)
end;

sess = Session(Graph())
X = placeholder(Float32, shape=[n_steps, n_input, batch_size])
Y_obs = placeholder(Float32, shape=[batch_size, n_classes])

variable_scope("model", initializer=Normal(0, 0.5)) do
    global W = get_variable("weights", [n_hidden, n_classes], Float32)
    global B = get_variable("bias", [n_classes], Float32)
end;

# Prepare data shape to match `rnn` function requirements
# Current data input shape: (n_steps, n_input, batch_size) from the way we declared X (and the way the data actually comes)
# Required shape: 'n_steps' tensors list of shape (batch_size, n_input)

x = transpose(X, Int32.([1, 3, 2].-1)) # Permuting batch_size and n_steps. (the -1 is to use 0 based indexing)
x = reshape(x, [n_steps*batch_size, n_input]) # Reshaping to (n_steps*batch_size, n_input)
x = split(1, n_steps, x) # Split to get a list of 'n_steps' tensors of shape (batch_size, n_input)
@show get_shape.(x);

Hs, states = nn.rnn(nn.rnn_cell.LSTMCell(n_hidden), x; dtype=Float32);
Y_pred = nn.softmax(Hs[end]*W + B)

@show get_shape(Y_obs)
@show get_shape(Y_pred);

cost = reduce_mean(-reduce_sum(Y_obs.*log(Y_pred), reduction_indices=[1])) #cross entropy
@show get_shape(Y_obs.*log(Y_pred))
@show get_shape(cost) #Should be [] as it should be a scalar

optimizer = train.minimize(train.AdamOptimizer(learning_rate), cost)

correct_prediction = indmax(Y_obs, 2) .== indmax(Y_pred, 2)
@show get_shape(correct_prediction)
accuracy = reduce_mean(cast(correct_prediction, Float32));

run(sess, initialize_all_variables())

kk=0
for jj in 1:training_iters
    for (xs, ys) in prepared_batchs(traindata_raw, trainlabels_raw)
        run(sess, optimizer,  Dict(X=>xs, Y_obs=>ys))
        kk+=1
        if kk % display_step == 1
            train_accuracy, train_cost = run(sess, [accuracy, cost], Dict(X=>xs, Y_obs=>ys))
            info("step $(kk*batch_size), loss = $(train_cost),  accuracy $(train_accuracy)")
        end
    end
end


export HMat;

@enum ADMISSTYPES ADMISS_STRONG ADMISS_STANDARD ADMISS_MILD ADMISS_WEAK;
@enum BLOCKTYPES LOWRANK DENSE HMAT;

type HMat{T<:Number}
    # variables
    height::    Int
    width::     Int
    level::     Int
    trg::       Array{Int,1}
    src::       Array{Int,1}
    blockType:: BLOCKTYPES
    childHMat:: Array{HMat{T},2}
    UMat::      Array{T,2}
    VMat::      Array{T,2}
    DMat::      Array{T,2}
    # global settings
    EPS::       Float64
    MAXRANK::   Int
    MINN::      Int

    HMat() = new()
end

function admiss(x,y,type_admiss)
    if ((type_admiss == ADMISS_STRONG) || (type_admiss == ADMISS_STANDARD))
        if maximum(x-y) > 1
            return true;
        else
            return false;
        end
    elseif type_admiss == ADMISS_MILD
        if sum(abs(x-y)) > 1
            return true;
        else
            return false;
        end
    elseif type_admiss == ADMISS_WEAK
        if x != y
            return true;
        else
            return false;
        end
    end
end

function Z2Cmapper(n,minn,tensor)
    if maximum(n) <= minn
        return reshape(tensor,prod(n));
    end
    nl = ifloor(n/2);
    nr = n-nl;
    nv = zeros(n);
    nrange = Array(Any,2,1);
    offset = 0;
    order = vec(zeros(Int64,prod(n),1));
    for i = 0:2^2-1
        for d = 0:2-1
            nv[d+1] = nl[d+1]*(1-mod(i>>d,2))+nr[d+1]*mod(i>>d,2);
            nrange[d+1] = nl[d+1]*mod(i>>d,2)+(1:ifloor(nv[d+1]));
        end
        order[offset+(1:prod(nv))] = Z2Cmapper(nv,minn,tensor[nrange[1],nrange[2]]);
        offset = offset + prod(nv);
    end
    return order;
end

function svdtrunc(A,eps,mR)
    (U, S, V) = svd(full(A),thin=true);
    if minimum(size(A))>1 && S[1] > 1e-15
        idx = find(find(S/S[1].>eps).<=mR);
        return U[:,idx], S[idx], V[:,idx];
    else
        return U, S, V;
    end
end

function HMatd2h{T}(D::AbstractMatrix{T}, nTrg, nSrc, type_admiss, idxTrg, idxSrc, level, EPS, MaxRank, minn)
    node = HMat{T}();
    node.height = prod(nTrg);
    node.width = prod(nSrc);
    node.trg = idxTrg;
    node.src = idxSrc;
    node.level = level;

    node.EPS = EPS;
    node.MAXRANK = MaxRank;
    node.MINN = minn;

    if admiss(idxTrg,idxSrc,type_admiss)
        (Utmp, Stmp, Vtmp) = svdtrunc(D,EPS,MaxRank);
        node.blockType = LOWRANK;
        if length(Stmp) > 0
            node.UMat = Utmp.*sqrt(Stmp)';
            node.VMat = Vtmp.*sqrt(Stmp)';
        else
            node.UMat = Utmp;
            node.VMat = Vtmp;
        end
    elseif maximum(nTrg) <= minn || maximum(nSrc) <= minn
        node.blockType = DENSE;
        node.DMat = full(D);
    else
        node.blockType = HMAT;
        node.childHMat = Array(HMat{T},4,4);
        trg = 2*idxTrg;
        src = 2*idxSrc;
        toffset = 0;
        for tx = 0:1, ty = 0:1
            trg = 2*idxTrg + [tx,ty];
            (txlen,tylen) = ifloor(nTrg/2).*(1-[tx,ty]) + iceil(nTrg/2).*[tx,ty];
            tRange = toffset + (1:txlen*tylen);
            toffset = toffset + txlen*tylen;

            soffset = 0;
            for sx = 0:1, sy = 0:1
                src = 2*idxSrc + [sx,sy];
                (sxlen,sylen) = ifloor(nSrc/2).*(1-[sx,sy]) + iceil(nSrc/2).*[sx,sy];
                sRange = soffset + (1:sxlen*sylen);
                soffset = soffset + sxlen*sylen;
                node.childHMat[tx*2+ty+1,sx*2+sy+1] = HMatd2h(D[tRange,sRange],[txlen,tylen],[sxlen,sylen],type_admiss,trg,src,level+1,EPS,MaxRank,minn);
            end
        end
    end
    return node;
end

include("../src/hcopy.jl");
include("../src/hadjoint.jl");
include("../src/hempty.jl");
include("../src/hidentity.jl");
include("../src/hmatvec.jl");
include("../src/hcompress.jl");
include("../src/huncompress.jl");
include("../src/hscale.jl");
include("../src/hadd.jl");
include("../src/hnorm.jl");
include("../src/hmul.jl");
include("../src/hinv.jl");
include("../src/hNewton.jl");


module Corners

export Corner, Wedge

using ..Elements
using ..Motions

"""
    Corner{N} <: Elements.Element

An immutable structure representing a corner (`N=1`)/wedge (`N=2`)

## Fields

- `σ::Real`: signed intensity
- `ν::Real`: interior angle factor (angle is νπ)
- `n₀::ComplexF64`: vector of bisector of interior angle
- `L::Real` : characteristic length of corner

## Constructors

- `Corner(σ,ν,θ)` : set up a corner with a bisector at angle `θ` and length `L=1`
- `Wedge(σ,ν,θ)`  : set up a wedge with a bisector at angle `θ` and length `L=1`
"""
struct Corner{N} <: Element
    σ::Real
    ν::Real
    n₀::ComplexF64
    L::Real
end

Corner(σ,ν,n₀,L) = Corner{1}(σ,ν,n₀,L)
Wedge(σ,ν,n₀,L) =  Corner{2}(σ,ν,n₀,L)

Corner(σ::Real,ν::Real,θ̄::Real) = Corner(σ,ν,exp(im*θ̄),1.0)
Wedge(σ::Real,ν::Real,θ̄::Real) = Wedge(σ,ν,exp(im*θ̄),1.0)

Elements.kind(::Corner) = Singleton
Elements.kind(::Type{Corner}) = Singleton

Base.angle(c::Corner) = angle(c.n₀)

function Elements.complexpotential(z::ComplexF64, c::Corner{N}) where {N}
    fact = exp(im*N*π/2)
    return fact*c.ν*c.σ*c.L^(1-1/c.ν).*(z*c.n₀').^(1/c.ν)
end

Elements.streamfunction(z::ComplexF64, c::Corner) = imag(complexpotential(z,c))

function Motions.induce_velocity(z::ComplexF64, c::Corner{N}, t) where {N}
    fact = exp(im*N*π/2)
    return conj(fact*c.σ*(z*c.n₀'/c.L).^(1/c.ν-1)*c.n₀')
end

end


###########Zero Rate Curve Types
#Abstract
abstract type AbstractZeroRateCurve end

#Scalar
struct ZeroRate{T2 <: Number} <: AbstractZeroRateCurve
    r::T2
    function ZeroRate(r::T2) where {T2 <: Number}
        return new{T2}(r)
    end
end

#Curve
struct ZeroRateCurve{num1 <: Number, num2 <: Number} <: AbstractZeroRateCurve
    r::Curve{num1, num2}
    function ZeroRateCurve(r_::Array{num1}, T::num2) where {num1 <: Number, num2 <: Number}
        new{num2, num1}(Curve(r_, T))
    end

    function ZeroRateCurve(r_::Curve{num1, num2}) where {num1 <: Number, num2 <: Number}
        new{num2, num1}(r_)
    end
end

#Common Constructors
function ZeroRate(r_::Array{num1}, T::num2) where {num1 <: Number, num2 <: Number}
    return ZeroRateCurve(r_, T)
end
function ImpliedZeroRate(r_::Array{num1}, T::num2) where {num1 <: Number, num2 <: Number}
    return ZeroRateCurve(ImpliedCurve(r_, T))
end
function ImpliedZeroRate(r_::Array{num1}, T::Array{num2}) where {num1 <: Number, num2 <: Number}
    return ZeroRateCurve(ImpliedCurve(r_, T))
end

export ZeroRateCurve;
export ZeroRate;


#=

# 115: 1D heterogeneous catalysis
 ([source code](SOURCE_URL))

Let $\Omega=(0,1)$, $\Gamma_1=\{0\}$, $\Gamma_2=\{1\}$
Regard a system of three species: $A,B,C$ and let
$u_A=[A]$, $u_B=[B]$ and $u_C=[C]$ be their corresponding concentrations.

Species $A$ and $B$ exist in the interior of the domain, species $C$
lives a the boundary $\Gamma_1$.  We assume a heterogeneous reaction scheme
where $A$ reacts to $C$ and $C$ reacts to $B$:

```math
\begin{aligned}
      A &\leftrightarrow C\\
      C &\leftrightarrow B
\end{aligned}
```
with reaction constants $k_{AC}^\pm$ and k_{BC}^\pm$.

In $\Omega$, both $A$ and $B$ are transported through diffusion:

```math
\begin{aligned}
\partial_t u_B - \nabla\cdot D_A \nabla u_A & = f_A\\
\partial_t u_B - \nabla\cdot D_B \nabla u_B & = 0\\
\end{aligned}
```
Here, $f(x)$ is a source term creating $A$.
On $\Gamma_2$, we set boundary conditions
```math
\begin{aligned}
D_A \nabla u_A & = 0\\
u_B&=0
\end{aligned}
```
describing no normal flux for $A$ and zero concentration of $B$.
On $\Gamma_1$, we use the mass action law to describe the boundary reaction and
the evolution of the boundary concentration $C$. We assume that there is a limited
amount of surface sites $S$ for species C, so in fact A has to react with a free
surface site in order to become $C$ which reflected by the factor $1-u_C$. The same
is true for $B$.
```math
\begin{aligned}
R_{AC}(u_A, u_C)&=k_{AC}^+ u_A(1-u_C) - k_{AC}^-u_C\\
R_{BC}(u_C, u_B)&=k_{BC}^+ u_B(1-u_C) - k_{BC}^-u_C\\
- D_A \nabla u_A  + S R_{AC}(u_A, u_C)& =0 \\
- D_B \nabla u_B  + S R_{BC}(u_B, u_C)& =0 \\
\partial_t C  - R_{AC}(u_A, u_C) - R_{BC}(u_B, u_C) &=0
\end{aligned}
```

=#

module Example115_HeterogeneousCatalysis1D
using Printf
using VoronoiFVM
using ExtendableGrids
using GridVisualize
using LinearAlgebra

function main(;n=10,Plotter=nothing,verbose=false,tend=1, unknown_storage=:sparse)

    h=1.0/convert(Float64,n)
    X=collect(0.0:h:1.0)
    N=length(X)

    grid=VoronoiFVM.Grid(X)
    ## By default, \Gamma_1 at X[1] and \Gamma_2 is at X[end]

    ## Species numbers
    iA=1
    iB=2
    iC=3

    ## Diffusion flux for species A and B
    D_A=1.0
    D_B=1.0e-2
    function flux!(f,u0,edge)
        u=unknowns(edge,u0)
        f[iA]=D_A*(u[iA,1]-u[iA,2])
        f[iB]=D_B*(u[iB,1]-u[iB,2])
    end

    ## Storage term of species A and B
    function storage!(f,u,node)
        f[iA]=u[iA]
        f[iB]=u[iB]
    end

    ## Source term for species a around 0.5
    function source!(f,node)
        x1=node[1]-0.5
        f[iA]=exp(-100*x1^2)
    end

    ## Reaction constants (p = + , m = -)
    ## Choosen to prefer path A-> C -> B
    ## More over, A reacts faster than to C than C to B
    ## leading to "catalyst poisoning", i.e. C taking up most of the
    ## available catalyst sites
    kp_AC=100.0
    km_AC=1.0

    kp_BC=0.1
    km_BC=1.0

    S=0.01

    R_AC(u_A, u_C)=kp_AC*u_A*(1-u_C) - km_AC*u_C
    R_BC(u_B, u_C)=kp_BC*u_B*(1-u_C) - km_BC*u_C

    function breaction!(f,u,node)
        if  node.region==1
            f[iA]=S*R_AC(u[iA], u[iC])
            f[iB]=S*R_BC(u[iB], u[iC])
            f[iC]=-R_BC(u[iB], u[iC])-R_AC(u[iA], u[iC])
        end
    end

    ## This is for the term \partial_t u_C at the boundary
    function bstorage!(f,u,node)
        if  node.region==1
            f[iC]=u[iC]
        end
    end

    physics=VoronoiFVM.Physics(
        num_species=3,
        breaction=breaction!,
        bstorage=bstorage!,
        flux=flux!,
        storage=storage!,
        source=source!
    )

    sys=VoronoiFVM.System(grid,physics,unknown_storage=unknown_storage)

    ## Enable species in bulk resp
    enable_species!(sys,iA,[1])
    enable_species!(sys,iB,[1])

    ## Enable surface species
    enable_boundary_species!(sys,iC,[1])

    ## Set Dirichlet bc for species B on \Gamma_2
    boundary_dirichlet!(sys,iB,2,0.0)

    ## Initial values
    inival=unknowns(sys)
    inival.=0.0
    U=unknowns(sys)

    tstep=0.01
    time=0.0

    ## Data to store surface concentration vs time

    p=GridVisualizer(Plotter=Plotter,layout=(3,1))

    control=fixed_timesteps!(VoronoiFVM.NewtonControl(),tstep)
    tsol=solve(inival,sys,[0,tend],control=control)

    p=GridVisualizer(Plotter=Plotter,layout=(3,1),fast=true)
    for it=1:length(tsol)
        time=tsol.t[it]
        scalarplot!(p[1,1],grid,tsol[iA,:,it],clear=true,title=@sprintf("[A]: (%.3f,%.3f)",extrema(tsol[iA,:,it])...))
        scalarplot!(p[2,1],grid,tsol[iB,:,it],clear=true,title=@sprintf("[B]: (%.3f,%.3f)",extrema(tsol[iB,:,it])...))
        scalarplot!(p[3,1],tsol.t[1:it],tsol[iC,1,1:it],title=@sprintf("[C]"),clear=true,show=true)
    end

    return tsol[iC,1,end]
end

function test()
    testval=0.87544440641274
    isapprox(main(unknown_storage=:sparse),testval,rtol=1.0e-12)&&
    isapprox(main(unknown_storage=:dense),testval,rtol=1.0e-12)
end
end


function scalar_field()
    # isoparametric approximation
    mesh = generate_grid(QuadraticQuadrilateral, (20, 20))
    f(x) = x[1]^2
    nodal_vals = [f(p.x) for p in mesh.nodes]

    ip_f = Lagrange{2,RefCube,2}() # function interpolation
    ip_g = Lagrange{2,RefCube,2}() # geometry interpolation

    # compute values in quadrature points
    qr = QuadratureRule{2, RefCube}(3) # exactly approximate quadratic field
    cv = CellScalarValues(qr, ip_f, ip_g)
    qp_vals = [Vector{Float64}(undef, getnquadpoints(cv)) for _ in 1:getncells(mesh)]
    for cellid in eachindex(mesh.cells)
        xe = getcoordinates(mesh, cellid)
        reinit!(cv, xe)
        for qp in 1:getnquadpoints(cv)
            qp_vals[cellid][qp] = f(spatial_coordinate(cv, qp, xe))
        end
    end

    # do a L2Projection for getting values in dofs
    projector = L2Projector(ip_f, mesh)
    projector_vals = project(projector, qp_vals, qr; project_to_nodes=false)

    # points where we want to retrieve field values
    points = [Vec((x, 0.52)) for x in range(0.0; stop=1.0, length=100)]

    # set up PointEvalHandler and retrieve values
    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, projector, projector_vals)
    @test f.(points) ≈ vals

    # alternatively retrieve vals from nodal values TODO: make this work?
    # vals = get_point_values(ph, nodal_vals)
    # @test f.(points) ≈ vals
end

function vector_field()
    ## vector field
    # isoparametric approximation
    mesh = generate_grid(QuadraticQuadrilateral, (20, 20))
    f(x) = Vec((x[1]^2, x[1]))
    nodal_vals = [f(p.x) for p in mesh.nodes]

    ip_f = Lagrange{2,RefCube,2}() # function interpolation
    ip_g = Lagrange{2,RefCube,2}() # geometry interpolation

    # compute values in quadrature points
    qr = QuadratureRule{2, RefCube}(3) # exactly approximate quadratic field
    cv = CellScalarValues(qr, ip_f, ip_g)
    qp_vals = [Vector{Vec{2,Float64}}(undef, getnquadpoints(cv)) for i=1:getncells(mesh)]
    for cellid in eachindex(mesh.cells)
        xe = getcoordinates(mesh, cellid)
        reinit!(cv, xe)
        for qp in 1:getnquadpoints(cv)
            qp_vals[cellid][qp] = f(spatial_coordinate(cv, qp, xe))
        end
    end

    # do a L2Projection for getting values in dofs
    projector = L2Projector(ip_f, mesh)
    projector_vals = project(projector, qp_vals, qr; project_to_nodes=false)
    # TODO: project_to_nodes should probably return dof values and not Vecs for vector fields
    # projector_vals = convert(Vector{Float64}, reinterpret(Float64, projector_vals))

    # points where we want to retrieve field values
    points = [Vec((x, 0.52)) for x in range(0.0; stop=1.0, length=100)]

    # set up PointEvalHandler and retrieve values
    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, projector, projector_vals)
    @test f.(points) ≈ vals

    # alternatively retrieve vals from nodal values# TODO
    # vals = get_point_values(ph, nodal_vals)
    # @test f.(points) ≈ vals
end

function superparametric()
    # superparametric approximation
    mesh = generate_grid(Quadrilateral, (20, 20))
    f(x) = x*x[1]
    ip_f = Lagrange{2,RefCube,2}() # function interpolation
    ip_g = Lagrange{2,RefCube,1}() # geometry interpolation

    # compute values in quadrature points
    qr = QuadratureRule{2, RefCube}(3) # exactly approximate quadratic field
    cv = CellScalarValues(qr, ip_f, ip_g)
    qp_vals = [Vector{Vec{2,Float64}}(undef, getnquadpoints(cv)) for i=1:getncells(mesh)]
    for cellid in eachindex(mesh.cells)
        xe = getcoordinates(mesh, cellid)
        reinit!(cv, xe)
        for qp in 1:getnquadpoints(cv)
            qp_vals[cellid][qp] = f(spatial_coordinate(cv, qp, xe))
        end
    end

    # do a L2Projection for getting values in dofs
    projector = L2Projector(ip_f, mesh)
    projector_vals = project(projector, qp_vals, qr; project_to_nodes=false)

    # points where we want to retrieve field values
    points = [Vec((x, 0.52)) for x in range(0.0; stop=1.0, length=100)]

    # set up PointEvalHandler and retrieve values
    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, projector, projector_vals)

    # can recover a quadratic field by a quadratic approximation
    @test f.(points) ≈ vals
end

function dofhandler()
    mesh = generate_grid(Quadrilateral, (2,2))
    dof_vals = [1., 2., 5., 4., 3., 6., 8., 7., 9.]
    points = [node.x for node in mesh.nodes] # same as nodes

    dh = DofHandler(mesh)
    push!(dh, :s, 1) # a scalar field
    close!(dh)

    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, dh, dof_vals, :s)
    @test vals ≈ 1.0:9.0

    # TODO
    # vals = get_point_values(ph, collect(1.0:9.0))
    # @test vals ≈ 1.0:9.0
end

function mixed_grid()
    ## Mixed grid where not all cells have the same fields

    # 5_______6
    # |\      |
    # |   \   |
    # 3______\4
    # |       |
    # |       |
    # 1_______2

    nodes = [Node((0.0, 0.0)),
            Node((1.0, 0.0)),
            Node((0.0, 1.0)),
            Node((1.0, 1.0)),
            Node((0.0, 2.0)),
            Node((1.0, 2.0))]

    cells = Ferrite.AbstractCell[Quadrilateral((1,2,4,3)),
            Triangle((3,4,6)),
            Triangle((3,6,5))]

    mesh = Grid(cells, nodes)
    addcellset!(mesh, "quads", Set{Int}((1,)))
    addcellset!(mesh, "tris", Set{Int}((2, 3)))

    ip_quad = Lagrange{2,RefCube,1}()
    ip_tri = Lagrange{2,RefTetrahedron,1}()

    f(x) = x[1]

    # compute values in quadrature points for quad
    qr = QuadratureRule{2, RefCube}(2)
    cv = CellScalarValues(qr, ip_quad)
    qp_vals_quads = [Vector{Float64}(undef, getnquadpoints(cv)) for cell in getcellset(mesh, "quads")]
    for (local_cellid, global_cellid) in enumerate(getcellset(mesh, "quads"))
        xe = getcoordinates(mesh, global_cellid)
        reinit!(cv, xe)
        for qp in 1:getnquadpoints(cv)
            qp_vals_quads[local_cellid][qp] = f(spatial_coordinate(cv, qp, xe))
        end
    end

    # construct projector
    projector = L2Projector(ip_quad, mesh; set=getcellset(mesh, "quads"))

    points = [Vec((x, 2x)) for x in range(0.0; stop=1.0, length=10)]

    # first alternative: L2Projection to dofs
    projector_values = project(projector, qp_vals_quads, qr; project_to_nodes = false)
    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, projector, projector_values)
    @test vals[1:5] ≈ f.(points[1:5])
    @test all(isnan, vals[6:end])
    # TODO
    # # second alternative: L2Projection to nodes
    # nodal_vals = project(projector, qp_vals_quads, qr; project_to_nodes = true)
    # vals = get_point_values(ph, nodal_vals)
    # @test vals[1:5] ≈ f.(points[1:5])
    # @test all(isnan.(vals[6:end]))

    # second alternative: assume a vector field :v
    dh = MixedDofHandler(mesh)
    field = Field(:v, ip_quad, 2)
    fh_quad = FieldHandler([field], getcellset(mesh, "quads"))
    push!(dh, fh_quad)
    field = Field(:v, ip_tri, 2)
    fh_tri = FieldHandler([field], getcellset(mesh, "tris"))
    push!(dh, fh_tri)
    close!(dh)

    dof_vals = [1., 1., 2., 2., 4., 4., 3., 3., 6., 6., 5., 5.]
    points = [node.x for node in mesh.nodes]
    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, dh, dof_vals, :v)
    @test vals == [Vec((i, i)) for i=1.0:6.0]
end

function oneD()
    # isoparametric approximation
    mesh = generate_grid(Line, (2,))
    f(x) = x[1]
    nodal_vals = [f(p.x) for p in mesh.nodes]

    ip_f = Lagrange{1,RefCube,1}() # function interpolation
    ip_g = Lagrange{1,RefCube,1}() # geometry interpolation

    # compute values in quadrature points
    qr = QuadratureRule{1, RefCube}(2)
    cv = CellScalarValues(qr, ip_f, ip_g)
    qp_vals = [Vector{Float64}(undef, getnquadpoints(cv)) for i=1:getncells(mesh)]
    for cellid in eachindex(mesh.cells)
        xe = getcoordinates(mesh, cellid)
        reinit!(cv, xe)
        for qp in 1:getnquadpoints(cv)
            qp_vals[cellid][qp] = f(spatial_coordinate(cv, qp, xe))
        end
    end

    # do a L2Projection for getting values in dofs
    projector = L2Projector(ip_f, mesh)
    projector_values = project(projector, qp_vals, qr; project_to_nodes=false)

    # points where we want to retrieve field values
    points = [Vec((x,)) for x in range(-1.0; stop=1.0, length=5)]

    # set up PointEvalHandler and retrieve values
    ph = PointEvalHandler(mesh, points)
    vals = get_point_values(ph, projector, projector_values)
    @test f.(points) ≈ vals

    # alternatively retrieve vals from nodal values
    # TODO
    # vals = get_point_values(ph, nodal_vals)
    # @test f.(points) ≈ vals
end

@testset "PointEvalHandler" begin
    scalar_field()
    vector_field()
    dofhandler()
    superparametric()
    mixed_grid()
    oneD()
end


module signalobject

include("utils/utils.jl")
include("config/ecgconfig.jl")

import Base

import .ecgconfig as cfg
using .utils

export SignalObject
mutable struct SignalObject{T<:Real}
    signals::Vector{T}
    signal_mean_absolute::Union{Vector{T}, UndefInitializer}
    signal_second_derivative::Union{Vector{T}, UndefInitializer}
    signal_enhance::Union{Vector{T}, UndefInitializer}
    SignalObject(signals::Vector{T}) where T = new{T}(signals)
end

function signal_mean_absolute(signals::Vector{<:Real}, windows_msec::Real = 20)

    windows = msec_to_sample(windows_msec, cfg.FS)
    return  utils.ma(abs.(signals), windows)
end

function signal_second_derivative(signals::Vector{T}) where T
    second_dif_signals = zeros(T, length(signals))
    second_dif_signals[2:end-1] = diff(diff(signals))
    return second_dif_signals
end

function Base.getproperty(sig::SignalObject, s::Symbol)
    if s == :signal_mean_absolute
        if isdefined(sig, s)
            return getfield(sig, s)
        else
            setfield!(sig, s, signal_mean_absolute(sig.signals))
            return getfield(sig, s)
        end
    elseif s == :signal_enhance
        if isdefined(sig, s)
            return getfield(sig, s)
        else
            return sig.signals
        end
    elseif s == :signal_second_derivative
        if isdefined(sig, s)
            return getfield(sig, s)
        else
            setfield!(sig, s, signal_second_derivative(sig.signals))
            return getfield(sig, s)
        end
    else
        return getfield(sig, s)
    end
end

Base.length(sig::SignalObject) = length(sig.signals)

end


"""
    PoissonLikelihood()

Poisson likelihood with rate as exponential of samples from GP `f`. This is to be used if
we assume that the uncertainity associated with the data follows a Poisson distribution.
"""
struct PoissonLikelihood end

(l::PoissonLikelihood)(f::Real) = Poisson(exp(f))

(l::PoissonLikelihood)(fs::AbstractVector{<:Real}) = Product(Poisson.(exp.(fs)))


struct HH3DPlaneWave{T,P} <: Functional
    direction::P
    wavenumber::T
    amplitude::T
end

function HH3DPlaneWave(direction, wavenumber; amplitude=1)
    w, a = promote(wavenumber, amplitude)
    HH3DPlaneWave(direction, w, a)
end

function (f::HH3DPlaneWave)(r)
    d = f.direction
    k = f.wavenumber
    a = f.amplitude
    a * exp(-im*k*dot(d,r))
end

struct NormalDerivative{F} <: Functional
    field::F
end

const ∂n = Val{:normalderivative}
(::Type{Val{:normalderivative}})(f) = NormalDerivative(f)

function (f::NormalDerivative{T})(manipoint) where T<:HH3DPlaneWave
    d = f.field.direction
    k = f.field.wavenumber
    a = f.field.amplitude
    n = normal(manipoint)
    r = cartesian(manipoint)
    -im*k*a * dot(d,n) * exp(-im*k*dot(d,r))
end

integrand(::NormalDerivative, test_vals, field_vals) = dot(test_vals[1], field_vals)


# functions related to binomial distribution

import .RFunctions:
    binompdf,
    binomlogpdf,
    binomcdf,
    binomccdf,
    binomlogcdf,
    binomlogccdf,
    binominvcdf,
    binominvccdf,
    binominvlogcdf,
    binominvlogccdf


export SGD, solve,getSGDsolver

"""
Stochastic Gradient Descent
"""
mutable struct SGD{T}
    maxEpochs::Int
    miniBatch::Int
    out::Bool
    learningRate::T
    momentum::T
    nesterov::Bool
	ADAM::Bool
end

function getSGDsolver(TYPE::Type ;maxEpochs=10,miniBatch=16,out=true,learningRate=0.1,momentum=0.9,nesterov=false,ADAM=false)
	if ADAM && nesterov
		warn("sgd(): ADAM and nestrov together - choosing ADAM");
		nesterov  = false;
	end
	return SGD{TYPE}(maxEpochs,miniBatch,out,convert(TYPE,learningRate),convert(TYPE,momentum), nesterov, ADAM)
end

Base.display(this::SGD)=println("SGD(maxEpochs=$(this.maxEpochs),miniBatch=$(this.miniBatch),learningRate=$(this.learningRate),momentum=$(this.momentum),nesterov=$(this.nesterov),ADAM=$(this.ADAM))")

function solve(this::SGD{T},objFun::dnnObjFctn,xc::Array{T},Y::Array{T},C::Array{T},Yv::Array{T},Cv::Array{T}) where {T<:Number}

    # evaluate training and validation
    epoch = 1;
    xOld = copy(xc);
    dJ = zeros(T,size(xc));
    mJ = zeros(T,size(xc));
    vJ = zeros(T,size(xc));
    if this.ADAM
        mJ = zeros(T,size(xc));
        vJ = zeros(T,size(xc));
    end
    beta2 = convert(T,0.999);
    beta1 = this.momentum;

    lr    = this.learningRate

    if this.out; display(this); end

    # Declare tmp - We know nothing about its shape or datatypes
    tmp = Array{Any}(0,0)

    while epoch <= this.maxEpochs
        nex = size(Y,2)
        ids = randperm(nex)

        for k=1:ceil(Int64,nex/this.miniBatch)
            idk = ids[(k-1)*this.miniBatch+1: min(k*this.miniBatch,nex)]
            if this.nesterov && !this.ADAM
                Jk,dummy,dJk,tmp = evalObjFctn(objFun,xc-this.momentum*dJ,Y[:,idk],C[:,idk],tmp);
            else
                Jk,dummy,dJk,tmp = evalObjFctn(objFun,xc,Y[:,idk],C[:,idk],tmp);
            end

            if this.ADAM
               mJ = beta1*mJ + (one(T)-beta1)*(dJk)
               vJ = beta2*vJ + (one(T)-beta2)*(dJk.^2)
#              dJ = lr*((mJ./(1-beta1^(epoch)))./sqrt.((vJ./(1-beta2^(epoch)))+1e-8))
            else
               dJ = lr*dJk + this.momentum*dJ
            end
            xc = xc - dJ
        end
        # we sample 2^12 images from the training set for displaying the objective.
        idt     = ids[1:min(nex,2^12)]
        Jtrain,ptrain   = getMisfit(objFun,xc,Y[:,idt],C[:,idt],tmp,false);
        Jval,pVal = getMisfit(objFun,xc,Yv,Cv,tmp,false);

        if this.out;
            @printf "%d\t%1.2e\t%1.2f\t%1.2e\t%1.2e\t%1.2f\n" epoch Jtrain 100*(1-ptrain[3]/ptrain[2]) norm(xOld-xc) Jval 100*(1-pVal[3]/pVal[2])
        end

        xOld       = copy(xc);
        epoch = epoch + 1;
    end
    return xc
end


using MAT
using Dates
using GLPK
using JuMP
using Statistics
using CartesianGeneticProgramming
using Cambrian
import Cambrian.mutate
using Plots
using Plots.PlotMeasures
using StatsPlots
using DataFrames
using GLM
include("utils.jl")

function get_wave_t_in_range(param, min, max)
    indices_more = findall(>=(min), param)
    indices_less = findall(<(max), param)

    indices_to_keep = intersect(indices_more, indices_less)
    param = param[indices_to_keep]

    param, indices_to_keep
end

function fit_poly_inf_norm(x,y,n)
	m = Model(GLPK.Optimizer)

    @variable(m, alpha[1:n+1])
	@variable(m, t >= 0.)
	@constraint(m, [i=1:length(y)], sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i] <= t);
	@constraint(m, [i=1:length(y)], -t <= sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i]);
	@objective(m, Min, t)

	optimize!(m) # added new 2020-04-08

    return JuMP.value.(alpha)
end

f2(x, args) = (args[2] * x) .+ args[1]

const MATLAB_EPOCH = Dates.DateTime(-1,12,31)

date2num(d::Dates.DateTime) = Dates.value(d-MATLAB_EPOCH)/(1000*60*60*24)
num2date(n::Number) =  MATLAB_EPOCH + Dates.Millisecond(round(Int64, n*1000*60*60*24))

function plot_bathy(xs, profile; show_points=false, x_last_positive=NaN, x_first_negative=NaN, value_last_positive=NaN, value_first_negative=NaN, x_at_zero=NaN, title=NaN)
    p1 = plot(xs, profile, label="profile_t0")
    if show_points
        scatter!([x_last_positive, x_first_negative], [value_last_positive, value_first_negative], label="A, B")
        scatter!([x_at_zero], [0], label="Shoreline")
    end
    display(plot(p1, title=title))
end

function find_x_at_zero(profile, xs, i)
    t0_index_last_positive = findlast(profile.>0)
    if t0_index_last_positive == nothing t0_index_last_positive = 1 end

    t0_index_first_negative = findfirst(profile.<0)
    if t0_index_first_negative == nothing t0_index_first_negative = 1 end

    t0_x_last_positive = xs[t0_index_last_positive]
    t0_x_first_negative = xs[t0_index_first_negative]
    t0_value_last_positive = profile[t0_index_last_positive]
    t0_value_first_negative = profile[t0_index_first_negative]

    t0_coeffs = fit_poly_inf_norm([t0_value_last_positive, t0_value_first_negative], [t0_x_last_positive, t0_x_first_negative], 1)
    t0_x_at_zero = f2([0], t0_coeffs)[1]
    t0_x_at_zero
end

dir_path = get_dataset_path(;dataset_name="grand_popo")
kalman_bathy_vars = matread(string(dir_path, "Kalman_bathy.mat"))
bathy_xb = kalman_bathy_vars["xb"][1, :]# 1×213 Matrix{Float64}
# bathy_tn = kalman_bathy_vars["tn"]      # 1×1214 Matrix{Float64}
bathy_hbK = kalman_bathy_vars["hbK"]    # 213×1214 Matrix{Float64}

min_target_matdate = minimum(filter(!isnan, kalman_bathy_vars["tn"]))
max_target_matdate = maximum(filter(!isnan, kalman_bathy_vars["tn"]))

wave_vars = matread(string(dir_path, "Wave_GPP_ERA.mat"))
wave_t = wave_vars["t"]
wave_Tm = wave_vars["Tm"]
wave_hs = wave_vars["hs"]
wave_dir = wave_vars["dir"]

wave_t, indices = get_wave_t_in_range(wave_t, min_target_matdate, max_target_matdate)
wave_Tm = wave_Tm[indices]
wave_hs = wave_hs[indices]
wave_dir = wave_dir[indices]

bathy_dates = map(num2date, collect(min_target_matdate:max_target_matdate))  # collect() instead of using bathy_tn because bathy_tn includes NaN's, even though bathy_hbK has no NaN
wave_dates = map(num2date, wave_t)

days = []
n_offshore = 0
n_landward = 0
n_none = 0

mean_tp = []
min_tp = []
max_tp = []
std_tp = []

mean_hs = []
min_hs = []
max_hs = []
std_hs = []

mean_dir = []
min_dir = []
max_dir = []
std_dir = []

changes = []
for i in 2:length(bathy_dates)
    _d0 = bathy_dates[i-1]
    _d1 = bathy_dates[i]

    _t0 = date2num(_d0)
    _t1 = date2num(_d1)

    _bathy_t0 = bathy_hbK[:, i-1]
    _bathy_t1 = bathy_hbK[:, i]

    _wave_t = deepcopy(wave_t)
    _wave_Tm = deepcopy(wave_Tm)
    _wave_dir = deepcopy(wave_dir)
    _wave_hs = deepcopy(wave_hs)

    _wave_t, _indices = get_wave_t_in_range(_wave_t, _t0, _t1)
    _wave_Tm = _wave_Tm[_indices]
    _wave_hs = _wave_hs[_indices]
    _wave_dir = _wave_dir[_indices]

    t0_x_at_zero = find_x_at_zero(_bathy_t0, bathy_xb, i)
    t1_x_at_zero = find_x_at_zero(_bathy_t1, bathy_xb, i)
    change = (t1_x_at_zero - t0_x_at_zero)# / 100

    # if change > 1 || change < -1
    #     println(i, "    ", change)
    #     exit()
    # end

    if change > 0
        class = (1, 0, 0)  # (offshore, landward, none)
    elseif change < 0
        class = (0, 1, 0)
    else
        class = (0, 0, 1)
    end

    # _wave_Tm = _wave_Tm ./ 100
    # _wave_hs = _wave_hs ./ 10
    # _wave_dir = _wave_dir ./ 100

    push!(mean_tp, mean(_wave_Tm))
    push!(min_tp, minimum(_wave_Tm))
    push!(max_tp, maximum(_wave_Tm))
    push!(std_tp, std(_wave_Tm))

    push!(mean_hs, mean(_wave_hs))
    push!(min_hs, minimum(_wave_hs))
    push!(max_hs, maximum(_wave_hs))
    push!(std_hs, std(_wave_hs))

    push!(mean_dir, mean(_wave_dir))
    push!(min_dir, minimum(_wave_dir))
    push!(max_dir, maximum(_wave_dir))
    push!(std_dir, std(_wave_dir))

    push!(changes, change)
end

train_df = DataFrame(
    Tp_mean=convert(Vector{Float64}, mean_tp[1:800]),
    Tp_min=convert(Vector{Float64}, min_tp[1:800]),
    Tp_max=convert(Vector{Float64}, max_tp[1:800]),
    Tp_std=convert(Vector{Float64}, std_tp[1:800]),

    Hs_mean=convert(Vector{Float64}, mean_hs[1:800]),
    Hs_min=convert(Vector{Float64}, min_hs[1:800]),
    Hs_max=convert(Vector{Float64}, max_hs[1:800]),
    Hs_std=convert(Vector{Float64}, std_hs[1:800]),

    Dir_mean=convert(Vector{Float64}, mean_dir[1:800]),
    Dir_min=convert(Vector{Float64}, min_dir[1:800]),
    Dir_max=convert(Vector{Float64}, max_dir[1:800]),
    Dir_std=convert(Vector{Float64}, std_dir[1:800]),

    Change=convert(Vector{Float64}, changes[1:800])
)

test_df = DataFrame(
    Tp_mean=convert(Vector{Float64}, mean_tp[800:end]),
    Tp_min=convert(Vector{Float64}, min_tp[800:end]),
    Tp_max=convert(Vector{Float64}, max_tp[800:end]),
    Tp_std=convert(Vector{Float64}, std_tp[800:end]),

    Hs_mean=convert(Vector{Float64}, mean_hs[800:end]),
    Hs_min=convert(Vector{Float64}, min_hs[800:end]),
    Hs_max=convert(Vector{Float64}, max_hs[800:end]),
    Hs_std=convert(Vector{Float64}, std_hs[800:end]),

    Dir_mean=convert(Vector{Float64}, mean_dir[800:end]),
    Dir_min=convert(Vector{Float64}, min_dir[800:end]),
    Dir_max=convert(Vector{Float64}, max_dir[800:end]),
    Dir_std=convert(Vector{Float64}, std_dir[800:end]),

    Change=convert(Vector{Float64}, changes[800:end])
)

println(describe(train_df))

model = lm(@formula(Change ~ Tp_mean + Tp_min + Tp_max + Tp_std + Hs_mean + Hs_min + Hs_max + Hs_std + Dir_mean + Dir_min + Dir_max + Dir_std), train_df)
println(model)
println("R2 score: ", r2(model))

ypredicted_test = predict(model, test_df)
ypredicted_train = predict(model, train_df)

rmse(target, preds) = mean((target .- preds).^2)

println("Train RMSE: ", rmse(changes[1:800], ypredicted_train))
println("Test RMSE: ", rmse(changes[800:end], ypredicted_test))

corr_train = round(cor(changes[1:800], ypredicted_train), digits=2)
corr_test = round(cor(changes[800:end], ypredicted_test), digits=2)
println("Train corr: ", corr_train)
println("Test corr: ", corr_test)

change_xs = collect(1:length(changes))
train_xs = collect(1:800)
test_xs = collect(800:length(changes))

p1 = plot(change_xs, changes, label="Raw data")
plot!(train_xs, ypredicted_train, label="Train set preds")
plot!(test_xs, ypredicted_test, label="Test set preds")
# display(plot(p1))
# readline()

plot!(size=(1920, 1080))
savefig("/Users/mn/JuliaProjects/cgp-morphcast/output/grand_popo/linear_regression.pdf")


# sum_n
@tensordiff (Z[i,j] = Espresso.sum_1(X[:,j])) (dZ[i,j]/dX[m,n] = 1.0 * (j == n))
@tensordiff (Z[i,j] = Espresso.sum_2(X[i,:])) (dZ[i,j]/dX[m,n] = 1.0 * (i == m))

# sum_n + sum
@tensordiff (Z = Espresso.sum_1(X[:,j])) (dZ[i,j]/dX[m,n] = 1.0)
@tensordiff (Z = Espresso.sum_2(X[i,:])) (dZ[i,j]/dX[m,n] = 1.0)

# matrix-by-matrix product
@tensordiff (Z[i,j] = X[i,k] * Y[k,j]) (dZ[i,j]/dX[m,n] = Y[n,j] * (i == m))
@tensordiff (Z[i,j] = X[i,k] * Y[k,j]) (dZ[i,j]/dY[m,n] = X[i,m] * (n == j))

# matrix-by-vector product
@tensordiff (Z[i] = X[i,k] * Y[k]) (dZ[i]/dX[m,n] = Y[n] * (i == m))
@tensordiff (Z[i] = X[i,k] * Y[k]) (dZ[i]/dY[m] = X[i,m])

# inner product of 2 vectors
@tensordiff (Z = X[i] * Y[i]) (dZ/dX[i] = Y[i])
@tensordiff (Z = X[i] * Y[i]) (dZ/dY[i] = X[i])

# outer product of 2 vectors
@tensordiff (Z[i,j] = X[i] * Y[j]) (dZ[i,j]/dX[m] = Y[j] * (i == m))
@tensordiff (Z[i,j] = X[i] * Y[j]) (dZ[i,j]/dY[m] = X[i] * (j == m))

# inner product of 2 matrices
@tensordiff (Z = X[i,j] * Y[i,j]) (dZ/dX[i,j] = Y[i,j])
@tensordiff (Z = X[i,j] * Y[i,j]) (dZ/dY[i,j] = X[i,j])

# index permutation (broken)
# @tensordiff (Z[i,j] = X[j,i]) (dZ[i,j]/dX[m,n] = 1 * (i == n) * (j == m))

# length
@tensordiff (Z = length(X[:])) (dZ/dX[i] = 0.0)

# convolution
@tensordiff (Z[i,j] = X[i+m-1, j+n-1] * W[m,n]) (dZ[i,j] / dW[m,n] = X[i+m-1, j+n-1])
@tensordiff (Z[i,j] = X[i+m-1, j+n-1] * W[m,n]) (dZ[i,j] / dX[p,q] = W[p-i+1, q-j+1])

@tensordiff (Z = X[i+m-1, j+n-1] * W[m,n]) (dZ / dW[m,n] = X[i+m-1, j+n-1])
@tensordiff (Z = X[i+m-1, j+n-1] * W[m,n]) (dZ / dX[p,q] = W[p-i+1, q-j+1])


#################### Gelman, Rubin, and Brooks Diagnostics ####################

function gelmandiag(chn::Chains;
        alpha::Real=0.05,
        mpsrf::Bool=false,
        transform::Bool=false,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
        showall=false,
        sorted=true)
    c = showall ?
        sorted ? sort(chn) : chn :
        Chains(chn, _clean_sections(chn, sections); sorted=sorted)
    @assert !any(ismissing.(c.value)) "Gelman diagnostics doesn't support missing values"

    n, p, m = size(c.value)
    m >= 2 ||
    throw(ArgumentError("less than 2 chains supplied to gelman diagnostic"))

    psi = transform ? link(c) : c.value.data

    S2 = mapslices(cov, psi, dims = [1, 2])
    W = dropdims(mapslices(mean, S2, dims = [3]), dims = 3)

    psibar = reshape(mapslices(mean, convert(Array, psi), dims = [1]), p, m)'
    B = n * cov(psibar)

    w = diag(W)
    b = diag(B)
    s2 = reshape(mapslices(diag, S2, dims = [1, 2]), p, m)'
    psibar2 = vec(mapslices(mean, psibar, dims = [1]))

    var_w = vec(mapslices(var, s2, dims = [1])) / m
    var_b = (2.0 / (m - 1)) * b.^2
    var_wb = (n / m) * (diag(cov(s2, psibar.^2))
            - 2.0 * psibar2 .* diag(cov(s2, psibar)))

    V = ((n - 1) / n) * w + ((m + 1) / (m * n)) * b
    var_V = ((n - 1)^2 * var_w + ((m + 1) / m)^2 * var_b +
        (2.0 * (n - 1) * (m + 1) / m) * var_wb) / n^2
    df = 2.0 * V.^2 ./ var_V
    B_df = m - 1
    W_df = 2.0 * w.^2 ./ var_w

    psrf = Array{Float64}(undef, p, 2)
    R_fixed = (n - 1) / n
    R_random_scale = (m + 1) / (m * n)
    q = 1.0 - alpha / 2.0
    for i in 1:p
        correction = (df[i] + 3.0) / (df[i] + 1.0)
        R_random = R_random_scale * b[i] / w[i]
        psrf[i, 1] = sqrt(correction * (R_fixed + R_random))
        if !isnan(R_random)
            R_random *= quantile(FDist(B_df, W_df[i]), q)
        end
        psrf[i, 2] = sqrt(correction * (R_fixed + R_random))
    end
    psrf_labels = Symbol.([:parameters, :PSRF, Symbol(string(100 * q) * "%")])
    psrf_names = Symbol.(names(c))

    if mpsrf
        x = isposdef(W) ?
        R_fixed + R_random_scale * eigmax(inv(cholfact(W)) * B) : NaN
        psrf = vcat(psrf, [x NaN])
        psrf_names = [psrf_names; "Multivariate"]
    end

    columns = [psrf_names, psrf[:,1], psrf[:,2]]
    name = "Gelman, Rubin, and Brooks Diagnostic"
    nt = NamedTuple{tuple(psrf_labels...)}(tuple(columns...))
    return ChainDataFrame(name, nt)
end


using Pseudospectra, Test

@testset "Sparse-ARPACK" begin
    # 400x400 real sparse
    A = Pseudospectra.convdiff_fd(10)
    opts = Dict{Symbol,Any}(:ARPACK_plots => false,:npts=>40)
    opts[:arpack_opts] = ArpackOptions{eltype(A)}(which=:SR,nev=20,ncv=60)

    ps_data = new_matrix(A,opts)
    driver!(ps_data,opts,gs)
    @test iscomputed(ps_data)
end


function forward_stagewise_regression(X, y; epsilon = 0.01, numIterations = 100)
  n, p = size(X)
  ny = length(y)
  @assert n == ny
  forward_stagewise_regression_inner(X, y, n, p, epsilon, numIterations)
end

function forward_stagewise_regression_inner(X, y, n, p, epsilon, numIterations)
  betas = zeros(p, numIterations)
  for k in 2:numIterations
    maxabsinnerproduct = 0.0
    bestinnerproduct = 0.0
    residuals = y .- X * betas[:, k-1]
    bestcolumn = rand(1:p)
    for j in 1:p
      innerproduct = sum( X[:,j] .* residuals )
      if abs(innerproduct) > maxabsinnerproduct
        bestinnerproduct = innerproduct
        maxabsinnerproduct = abs(innerproduct)
        bestcolumn = j
      end
    end
    betas[:, k] = betas[:, k-1]
    betas[bestcolumn, k] += epsilon * sign(bestinnerproduct)
  end
  betas
end

N = 20
P = 5
X = randn(N, P)
y = 1.0 * X[:,1] + 2.0 * X[:,3]
betas = forward_stagewise_regression(X, y; epsilon = 0.10)


__precompile__()

module Quaternions

  import Base: +, -, *, /, ^, ==
  import Base: abs, abs2, angle, conj, cos, exp, inv, isreal, isfinite, isinf, iszero, isnan, log, real, sin, sqrt
  import Base: convert, promote_rule, float
  import Base: rand, randn
  import LinearAlgebra: lyap, norm, normalize, sylvester
  using LinearAlgebra: cross, dot
  using Random

  Base.@irrational INV_SQRT_EIGHT 0.3535533905932737622004 sqrt(big(0.125))

  include("Quaternion.jl")
  include("Octonion.jl")
  include("DualQuaternion.jl")

  export Quaternion,
    QuaternionF16,
    QuaternionF32,
    QuaternionF64,
    Octonion,
    OctonionF16,
    OctonionF32,
    OctonionF64,
    DualQuaternion,
    DualQuaternionF16,
    DualQuaternionF32,
    DualQuaternionF64
  export quat
  export octo
  export imag_part
  export dualquat
  export angleaxis
  export angle
  export axis
  export linpol
  export normalize
  export normalizea
  export dconj
  export quatrand
  export nquatrand
  export octorand
  export dualquatrand
  export ndualquatrand
  export qrotation
  export rotationmatrix
  export slerp
end


# This module is a collection of functions used to scan different parameters in CarrierCapture.
export getHarmonicQ_m, getMorseQ_m, fitHarmonicParams, fitMorseParams, getHarmonicCapture, getMorseCapture

"""
    getHarmonicQ_m(ħω,E0)

This function gets the value of ΔQ corresponding to an activationless
Marcus regime (Q_m) for two harmonic potential energy surfaces

# Parameters
`ħω` : energy of vibration of the initial state, in eV
`E0` : vertical shift between PES, in eV

# Returns
`Q_m` : The value of ΔQ corresponding to an activationless Marcus regime, in amu^0.5 Å^-1

"""
function getHarmonicQ_m(ħω,E0)
    amu = 931.4940954E6   # eV / c^2
    ħc = 0.19732697E-6    # eV m

    a = amu/2 * (ħω/ħc/1E10)^2
    return (E0/a)^0.5
end

"""
    getMorseQ_m(a, b, E0)

This function gets the value of Q corresponding to an activationless
Marcus regime (Q_m) for two Morse potentials

# Parameters
`a` : first parameter of the Morse potential of the initial state
`b` : second parameter of the Morse potential of the final state
`ΔE` : vertical shift between PES, in eV

# Returns
`Q_m` : The value of ΔQ corresponding to an activationless Marcus regime, in amu^0.5 Å^-1

"""
function getMorseQ_m(a, b, E0)
    Q_m = (1/b)*log(1-(E0/a)^0.5)
    return Q_m
end

"""
    function fitHarmonicParams(ħω_i, ħω_f, ΔQ, ΔE)
This function assigns the parameters to the harmonic potential energy surfaces and
solves the potentials using a 1D S.E. solver

# Parameters
`ħω_i` : energy of vibration of the initial state
`ħω_f` : energy of vibration of the final state

`ΔQ` : horizontal shift between the PES', in amu^0.5 Å^-1
`ΔE` : vertical shift between the PES', in eV

# Returns
`poti`, `potf` : potential of initial and final states

"""
function fitHarmonicParams(ħω_i, ħω_f, ΔQ, ΔE)
        # set nev
        nev_excited = 180
        nf_min = (1/(ħω_f))*((nev_excited+0.5)ħω_i+2ΔE)-0.5
        nev_relaxed = floor(nf_min+1)

        Q = range(-20-ΔQ, stop=20+ΔQ, length=5000) # ensure high enough to get good harmonic fit

        # potential 1
        potf = Potential(); potf.name = "Relaxed state"
        potf.Q0 = ΔQ; potf.E0 = 0
        potf.nev = nev_relaxed # must overlap with initial excited state
        potf.func = x -> harmonic(x, ħω_f; E₀ = potf.E0, Q₀ = potf.Q0)
        potf.params["ħω"]= ħω_f
        potf.Q = Q
        potf.E = potf.func.(Q)
        solve_pot!(potf)

        # potential 2
        poti = Potential(); poti.name = "Excited state"
        poti.Q0 = 0; poti.E0 = ΔE;
        poti.nev = nev_excited # must be converged for partition function
        poti.func = x -> harmonic(x, ħω_i; E₀ = poti.E0, Q₀ = poti.Q0)
        poti.Q = Q
        poti.E = poti.func.(Q)
        poti.params["ħω"]= ħω_i
        solve_pot!(poti)
        println("parameters fit to potentials")
        return poti, potf
end

"""
    fitMorseParams(a_i, a_f, b_i, b_f, Q0, E0)

This function fits the parameters to the Morse potential energy surfaces and
solves the potentials using a 1D S.E. solver

# Parameters
`a_i` : first parameter of the Morse potential of the initial state
`a_f` : first parameter of the Morse potential of the final state
`b_i` : second parameter of the Morse potential of the initial state
`b_f` : second parameter of the Morse potential of the final state
`Q0` : horizontal shift between the PES', in amu^0.5 Å^-1
`E0` : vertical shift between the PES', in eV

# Returns
`poti`, `potf` : potential of initial and final states

"""
function fitMorseParams(a_i, a_f, b_i, b_f, Q0, E0)
    # set number of eigenvalues (nev) to solve for in potential
    nev_excited = 180 # must be converged for partition function
    nev_relaxed = 800 # ensure max(ε_f) - max(ε_i) > E0

    Q = range(-20-Q0, stop=20+Q0, length=5000) # ensure high enough to get good harmonic fit

    # potential 1
    poti = Potential(); poti.name = "Excited state"
    poti.Q0 = 0; poti.E0 = E0;
    poti.nev = nev_excited
    poti.func = x -> morse(x, [a_i, b_i]; E₀ = poti.E0, Q₀ = poti.Q0)
    poti.Q = Q
    poti.E = poti.func.(Q)
    poti.params["a"]= a_i
    poti.params["b"]= b_i
    solve_pot!(poti)

    # potential 2
    potf = Potential(); potf.name = "Relaxed state"
    potf.Q0 = Q0; potf.E0 = 0
    potf.nev = nev_relaxed # must overlap with initial excited state
    potf.func = x -> morse(x, [a_f, b_f]; E₀ = potf.E0, Q₀ = potf.Q0)
    potf.params["a"]= a_f
    potf.params["b"]= b_f
    potf.Q = Q
    potf.E = potf.func.(Q)
    solve_pot!(potf)

    println("parameters fit to potentials")
    return poti, potf
end

"""
    getHarmonicCapture(poti, potf)

Get the capture coefficient in a charge trapping process from an initial to
final state described by harmonic potential energy surfaces.

# Parameters
`poti` : initial potential
`potf` : final potential

# Returns
`capt_coeff` : capture coefficient in units of cm^3 /s
`barrier_height` : classical barrier to a capture process in eV
"""
function getHarmonicCapture(poti, potf)
    Volume = 1e-21 # cm^3
    σ = 0.8*potf.params["ħω"]
    g = 1 # degeneracy
    temperature = [300]

    ħω= potf.params["ħω"]
    E0 = poti.E0
    println("E0: $E0")
    Q_m = getHarmonicQ_m(ħω, E0)
    Q0 = potf.Q0
    W = 0.068/(Q0-Q_m) # e-ph coupling

    cc = conf_coord(poti, potf) # i, f
    println("conf coordinate fitted")
    cc.W = W
    cc.g = g
    calc_overlap!(cc; σ = σ) # assigns overlap matrix and δ matrix to cc
    println("overlap calculated")
    calc_capt_coeff!(cc, Volume, temperature)
    println("capture coefficient calculated")

    cross = try
                find_crossing(potf,poti)[2]
            catch
                50.0 # if it can't converge, set barrier very high
            end

    if cross==50.0
        barrier_height = 50.0
    else
        barrier_height = cross-poti.E0
    end
    capt_coeff = cc.capt_coeff[1]

    return capt_coeff, barrier_height
end

"""
    getMorseCapture(poti, potf)

Get the capture coefficient in a charge trapping process from an initial to
final state described by harmonic potential energy surfaces.

# Parameters
`poti` : initial potential
`potf` : final potential

# Returns
`capt_coeff` : capture coefficient in units of cm^3 /s
`barrier_height` : classical barrier to a capture process in eV
"""
function getMorseCapture(poti, potf)
    Volume = 1e-21 # cm^3
    σ = 0.01
    g = 1 # degeneracy
    temperature = [300]

    Q0 = potf.Q0
    E0 = poti.E0
    a = potf.params["a"]
    b = potf.params["b"]
    Q_m = getMorseQ_m(a, b, E0)

    W = 0.068/(Q0-Q_m) # e-ph coupling

    cc = conf_coord(poti, potf) # i, f
    println("conf coordinate fitted")
    cc.W = W
    cc.g = g
    calc_overlap!(cc; σ = σ) # assigns overlap matrix and δ matrix to cc
    println("overlap calculated")
    calc_capt_coeff!(cc, Volume, temperature)
    println("capture coefficient calculated")

    cross = try
                find_crossing(potf,poti)[2]
            catch
                50.0 # if it can't converge, set barrier very high
            end

    if cross==50.0
        barrier_height = 50.0
    else
        barrier_height = cross-poti.E0
    end
    capt_coeff = cc.capt_coeff[1]

    return capt_coeff, barrier_height
end


include("SLPDE.jl")
include("etdrk4.jl")
include("evaluate_lambda.jl")

srand(0)

println("Testing Local Allen--Cahn")

n = 25
U0 = zeros(4n, 8n-1);
U0[1:n, 1:2n-1] = sphrandn(Float64, n, n)/n

SLPDE = SemiLinearPDE(LaplaceBeltrami(0.1), NonlinearOperator(u->u-u^3), U0)

IMAX = 6

U = Vector{Matrix{Float64}}(IMAX)
for i = 1:IMAX
    U[i] = ETDRK4(SLPDE, 1.0, 10*2^(i-1))
    println("Done i = ", i)
end

err = zeros(IMAX-1)
for i = 1:IMAX-1
    err[i] = vecnorm(U[i] - U[end], Inf)
    println("i = ",i,", and the ∞-norm absolute error: ", err[i])
end

for i = 1:IMAX-2
    println(err[i]/err[i+1])
end

println("Testing Nonlocal Allen--Cahn")

SLPDE = SemiLinearPDE(NonlocalLaplaceBeltrami(0.1, -0.5, 0.2), NonlinearOperator(u->u-u^3), U0)

U = Vector{Matrix{Float64}}(IMAX)
for i = 1:IMAX
    U[i] = ETDRK4(SLPDE, 1.0, 10*2^(i-1))
    println("Done i = ", i)
end

err = zeros(IMAX-1)
for i = 1:IMAX-1
    err[i] = vecnorm(U[i] - U[end], Inf)
    println("i = ",i,", and the ∞-norm absolute error: ", err[i])
end

for i = 1:IMAX-2
    println(err[i]/err[i+1])
end

println("Testing Local Gray--Scott")

n = 25
U0 = zeros(4n, 8n-1), zeros(4n, 8n-1);
U0[1][1:n, 1:2n-1] = sphrandn(Float64, n, n)/n
U0[2][1:n, 1:2n-1] = sphrandn(Float64, n, n)/n

SLPDE = SemiLinearPDE((LaplaceBeltrami(sqrt(0.00002)), LaplaceBeltrami(sqrt(0.00001))), (NonlinearOperator((u,v)->0.04*(1-u)-u*v^2), NonlinearOperator((u,v)->-0.1*v+u*v^2)), U0)

U = Vector{Tuple{Matrix{Float64},Matrix{Float64}}}(IMAX)
for i = 1:IMAX
    U[i] = ETDRK4(SLPDE, 1.0, 10*2^(i-1))
    println("Done i = ", i)
end

err = zeros(IMAX-1)
for i = 1:IMAX-1
    err[i] = vecnorm(U[i][1] - U[end][1], Inf) + vecnorm(U[i][2] - U[end][2], Inf)
    println("i = ",i,", and the ∞-norm absolute error: ", err[i])
end

for i = 1:IMAX-2
    println(err[i]/err[i+1])
end

println("Testing Nonlocal Gray--Scott")

SLPDE = SemiLinearPDE((NonlocalLaplaceBeltrami(sqrt(0.00002), -0.5, 0.2), NonlocalLaplaceBeltrami(sqrt(0.00001), -0.5, 0.2)), (NonlinearOperator((u,v)->0.04*(1-u)-u*v^2), NonlinearOperator((u,v)->-0.1*v+u*v^2)), U0)

U = Vector{Tuple{Matrix{Float64},Matrix{Float64}}}(IMAX)
for i = 1:IMAX
    U[i] = ETDRK4(SLPDE, 1.0, 10*2^(i-1))
    println("Done i = ", i)
end

err = zeros(IMAX-1)
for i = 1:IMAX-1
    err[i] = vecnorm(U[i][1] - U[end][1], Inf) + vecnorm(U[i][2] - U[end][2], Inf)
    println("i = ",i,", and the ∞-norm absolute error: ", err[i])
end

for i = 1:IMAX-2
    println(err[i]/err[i+1])
end


using jInv.Mesh;
using Helmholtz
using Helmholtz.PointSourceADR
using Multigrid

plotting = true;

if plotting
	using PyPlot;
	close("all");
end

include("getModels.jl");

NeumannAtFirstDim = false;
Sommerfeld = true;
gamma_val = 0.01;

model = "const";
domain = [0.0,1.0,0.0,1.0];
n_cells = [512,512];
src = div.(n_cells,2).+1;
# src[2] = 1;

if model=="const"
	NeumannAtFirstDim = false;
	Sommerfeld = true;
	f = 40.0;
	src = div.(n_cells,2).+1;
	m_coarse = getConstModel(n_cells.+1);
	pad_cells = [20;20];
else
	error("unknown model");
end
Minv = getRegularMesh(domain,n_cells);

w = 2*pi*f

if plotting
	figure()
	imshow(m_coarse'); colorbar();
	title(string("The ",model," model"))
end

n_tup = tuple((n_cells.+1)...);
q_coarse = zeros(ComplexF64,n_tup)
src2 = div.(n_cells,2).+1;
q_coarse[src2[1],src2[2]] = 1/(Minv.h[2]^2);

# src2 = div(3*n_cells,4)+1;
# q_coarse[src2[1],src2[2]] = 1/(Minv.h[2]^2);
# src = src2;

println("omega*h:");
println(w*Minv.h*sqrt(maximum(m_coarse)));

maxOmega = getMaximalFrequency(m_coarse,Minv);
ABLamp = maxOmega;

m = m_coarse;

gamma = (w*gamma_val)*ones(Float32,size(m)) + getABL(Minv.n.+1,NeumannAtFirstDim,pad_cells.+1,ABLamp);

if plotting
	figure();
	imshow(reshape(gamma,n_tup));
end

# getHelmholtzADR(central::Bool, Mesh::RegularMesh, mNodal::Array{Float64},omega::Float64, gamma::Array,T,G,LT,
									# NeumannAtFirstDim::Bool,src::Array{Int64},Sommerfeld::Bool,single::Bool = true)

X1,X2 = ndgrid(collect(0:Minv.n[1])*Minv.h[1],collect(0:Minv.n[2])*Minv.h[2]);

T  = X1;
Tx = ones(size(X1));
Ty = zeros(size(X1));
# G = [Tx,Ty];
G =  Array{Array{Float64}}(undef, Minv.dim);
G[1] = Tx;
G[2] = Ty;
LT = zeros(size(X1));
(ADR_long,ADR_short,T) = getHelmholtzADR(true,Minv, m, w, gamma,T,G,LT, NeumannAtFirstDim,[3,3],Sommerfeld);

# ADR_short = 0.5*ADR_long + 0.5*ADR_short;

if plotting
	figure();
	imshow(real(exp.(1im*w*reshape(T,n_tup)))');
end

H, = GetHelmholtzOperator(Minv, m, w, gamma_val*ones(Float32,size(m)),NeumannAtFirstDim,pad_cells.+1,ABLamp,Sommerfeld);
S =  GetHelmholtzShiftOP(m, w,0.2);

q = q_coarse;
xh = (H\q[:])
q = q[:].*exp.(1im*w*T[:]);
xs = ((ADR_short )\q[:]).*exp.(-1im*w*T[:]);
us = (H+S)\q[:];

# if plotting
# figure();
# subplot(1,4,1);
# imshow(real(I2'));
# subplot(1,4,2);
# imshow(reshape(T,n_tup)');
# subplot(1,4,3);
# imshow(reshape(real(alu),n_tup)');
# subplot(1,4,4);
# imshow(reshape(real(exp.(-1im*w*T)),n_tup)');

if plotting
	figure();
	subplot(1,3,1);
	imshow(reshape(real(xh),n_tup)',cmap = "jet"); title("Full Helmholtz solution"); colorbar();
	subplot(1,3,2);
	imshow(reshape(real(-us),n_tup)',cmap = "jet"); title("Shifted Helmholtz solution"); colorbar();
	subplot(1,3,3);
	imshow(reshape(real(xs),n_tup)',cmap = "jet"); title("First order upwind ADR solution");colorbar();
end
a = 1;

levels      = 6;
numCores 	= 2;
maxIter     = 10;
relativeTol = 1e-10;
relaxType   = "Jac-GMRES";
relaxParam  = 0.7;
relaxPre 	= 3;
relaxPost   = 3;
cycleType   ='W';
coarseSolveType = "NoMUMPS";

MG = getMGparam(ComplexF64,Int64,levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0);
nrhs = 1;
MGsetup(sparse(ADR_short'),Minv,MG,nrhs,true);

println("****************************** Stand-alone GMG RAP for ADR 1st: ******************************")
a = zeros(ComplexF64, size(q[:]));
solveMG(MG,q,a,true);

println("****************************** GMRES GMG RAP for ADR 1st: ******************************")
F = getAfun(sparse(ADR_short'),zeros(eltype(q),size(q)),1)
MG.maxOuterIter = 2;
solveGMRES_MG(F,MG,q,a,true,10,true);

z = 1;

# ADRparam = getADRparam(Minv,gamma,m,w,T,G,LT,src,NeumannAtFirstDim,true,M);
	# getOperator = (mesh,ADRparam)->((ADR2,ADR1,~) = getHelmholtzADR(true,ADRparam.Minv, ADRparam.m, ADRparam.w, ADRparam.gamma,
									# ADRparam.T,ADRparam.G,ADRparam.LT,ADRparam.NeumannAtFirstDim,ADRparam.src,ADRparam.Sommerfeld);
									# H = ADR2 + GetHelmholtzShiftOP(ADRparam.m, ADRparam.w,shift);
									# # M = exp(-1im*w*T);
									# # M = ADRparam.scaling;
									# # rescaleMatrix(H,M);
									# return H);

	# MG = getMGparam(levels,numCores,maxIter,relativeTol/relres,relaxType,relaxParam,
				# relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0,"FullWeighting");
	# MGsetup(getMultilevelOperatorConstructor(ADRparam,getOperator,restrictADRParam),ADRparam.Minv,MG,eltype(q),1,true);


module LegendrePolynomials

using OffsetArrays

export Pl
export collectPl
export collectPl!
export dnPl
export collectdnPl
export collectdnPl!

function checkdomain(x)
	abs(x) > 1 && throw(DomainError(x,
		"Legendre Polynomials are defined for arguments lying in -1 ⩽ x ⩽ 1"))
end

assertnonnegative(l) = (l >= 0 || throw(ArgumentError("l must be >= 0, received " * string(l))))

function checklength(arr, minlength)
	length(arr) >= minlength || throw(ArgumentError(
		"array is not large enough to store all values, require a minimum length of " * string(minlength)))
end

@inline function Pl_recursion(::Type{T}, ℓ, Plm1, Plm2, x) where {T}
	# relation is valid from ℓ = 1
	Pl = ((2ℓ-1) * x * Plm1 - (ℓ-1) * Plm2)/ℓ
	convert(T, Pl)
end

@inline function dPl_recursion(::Type{T}, ℓ, n, P_n_lm1, P_nm1_lm1, P_n_lm2, x) where {T}
	P_n_l = ((2ℓ - 1) * (x * P_n_lm1 + n * P_nm1_lm1) - (ℓ - 1) * P_n_lm2)/ℓ
	convert(T, P_n_l)
end

# special cases
# case 1 : l == n, in which case there are no (l-1,n) and (l-2,n) terms
@inline function dPl_recursion(::Type{T}, ℓ, n, ::Nothing, P_nm1_lm1, ::Nothing, x) where {T}
	P_n_l = (2ℓ-1) * P_nm1_lm1
	convert(T, P_n_l)
end
# case 1 : l == n + 1, in which case there's no (l-2,n) term
@inline function dPl_recursion(::Type{T}, ℓ, n, P_n_lm1, P_nm1_lm1, ::Nothing, x) where {T}
	P_n_l = ((2ℓ - 1) * (x * P_n_lm1 + n * P_nm1_lm1))/ℓ
	convert(T, P_n_l)
end

polytype(x) = typeof(float(x*x))

"""
	LegendrePolynomialIterator(x, [lmax::Integer])

Return an iterator that generates the values of the Legendre polynomials ``P_\\ell(x)`` for the given `x`.
If `lmax` is specified then only the values of ``P_\\ell(x)`` from `0` to `lmax` are returned.

# Examples
```jldoctest
julia> import LegendrePolynomials: LegendrePolynomialIterator

julia> iter = LegendrePolynomialIterator(0.5, 4);

julia> collect(iter)
5-element OffsetArray(::Array{Float64,1}, 0:4) with eltype Float64 with indices 0:4:
  1.0
  0.5
 -0.125
 -0.4375
 -0.2890625

julia> iter = LegendrePolynomialIterator(0.5);

julia> collect(Iterators.take(iter, 5)) # evaluete 5 elements (l = 0:4)
5-element Array{Float64,1}:
  1.0
  0.5
 -0.125
 -0.4375
 -0.2890625

julia> collect(Iterators.take(Iterators.drop(iter, 100), 5)) # evaluate Pl for l = 100:104
5-element Array{Float64,1}:
 -0.0605180259618612
  0.02196749072249231
  0.08178451892628381
  0.05963329258495025
 -0.021651535258316177
```
"""
struct LegendrePolynomialIterator{T, L <: Union{Integer, Nothing}, V}
	x :: V
	lmax :: L
	function LegendrePolynomialIterator{T,L,V}(x::V, lmax::L) where {T, L <: Union{Integer, Nothing}, V}
		checkdomain(x)
		new{T,L,V}(x, lmax)
	end
end
LegendrePolynomialIterator(x) = LegendrePolynomialIterator{polytype(x), Nothing, typeof(x)}(x, nothing)
function LegendrePolynomialIterator(x, lmax)
	assertnonnegative(lmax)
	LegendrePolynomialIterator{polytype(x), typeof(lmax), typeof(x)}(x, lmax)
end

Base.eltype(::Type{<:LegendrePolynomialIterator{T}}) where {T} = T
Base.IteratorSize(::Type{<:LegendrePolynomialIterator{<:Any,Nothing}}) = Base.IsInfinite()

function Base.iterate(::LegendrePolynomialIterator{T}) where {T}
	return one(T), (1, (zero(T), one(T)))
end
function Base.iterate(iter::LegendrePolynomialIterator{T,Nothing}, state) where {T}
	l, (Plm2, Plm1) = state
	x = iter.x
	Pl = Pl_recursion(T, l, Plm1, Plm2, x)
	return Pl, (l+1, (Plm1, Pl))
end
function Base.iterate(iter::LegendrePolynomialIterator{T}, state) where {T}
	l, (Plm2, Plm1) = state
	l > iter.lmax && return nothing
	x = iter.x
	Pl = Pl_recursion(T, l, Plm1, Plm2, x)
	return Pl, (l+1, (Plm1, Pl))
end

Base.length(iter::LegendrePolynomialIterator{<:Any,<:Integer}) = iter.lmax + 1
Base.size(iter::LegendrePolynomialIterator{<:Any,<:Integer}) = (length(iter),)
Base.axes(iter::LegendrePolynomialIterator{<:Any,<:Integer}) = (0:iter.lmax, )
Base.keys(iter::LegendrePolynomialIterator{<:Any,<:Integer}) = 0:iter.lmax

function Base.collect(iter::LegendrePolynomialIterator{T,<:Integer}) where {T}
	arr = zeros(T, 0:iter.lmax)
	collectPl!(arr, iter.x; lmax = iter.lmax)
end

Base.copy(iter::LegendrePolynomialIterator) = typeof(iter)(iter.x, iter.lmax)

"""
	Pl(x, l::Integer)

Compute the Legendre Polynomial ``P_\\ell(x)`` for the argument `x` and the degree `l`

# Examples
```jldoctest
julia> Pl(1, 2)
1.0

julia> Pl(0.5, 20)
-0.04835838106737356
```
"""
function Pl(x, l::Integer)
	iter = LegendrePolynomialIterator(x)
	d = Iterators.drop(iter, l) # 0 to l-1
	first(d)
end

function doublefactorial(T, n)
	p = convert(T, one(n))
	for i in n:-2:1
		p *= convert(T, i)
	end
	convert(T, p)
end

function _checkvalues(x, l, n)
	assertnonnegative(l)
	checkdomain(x)
	n >= 0 || throw(ArgumentError("order of derivative n must be >= 0"))
end

Base.@propagate_inbounds function _unsafednPl!(cache, x, l, n)
	# unsafe, assumes 1-based indexing
	checklength(cache, l - n + 1)
	if n == l # may short-circuit this
		cache[1] = doublefactorial(eltype(cache), 2l-1)
	else
		collectPl!(cache, x, lmax = l - n)

		for ni in 1:n
			# We denote the terms as P_ni_li

			# li == ni
			P_nim1_nim1 = cache[1]
			P_ni_ni = dPl_recursion(eltype(cache), ni, ni, nothing, P_nim1_nim1, nothing, x)
			cache[1] = P_ni_ni

			# li == ni + 1
			P_nim1_ni = cache[2]
			P_ni_nip1 = dPl_recursion(eltype(cache), ni + 1, ni, P_ni_ni, P_nim1_ni, nothing, x)
			cache[2] = P_ni_nip1

			for li in ni+2:min(l, l - n + ni)
				P_ni_lim2 = cache[li - ni - 1]
				P_ni_lim1 = cache[li - ni]
				P_nim1_lim1 = cache[li - ni + 1]
				P_ni_li = dPl_recursion(eltype(cache), li, ni, P_ni_lim1, P_nim1_lim1, P_ni_lim2, x)
				cache[li - ni + 1] = P_ni_li
			end
		end
	end
	nothing
end

"""
	dnPl(x, l::Integer, n::Integer, [cache::AbstractVector])

Compute the ``n``-th derivative ``d^n P_\\ell(x)/dx^n`` of the Legendre polynomial ``P_\\ell(x)``.
Optionally a pre-allocated vector `cache` may be provided, which must have a minimum length of `l - n + 1`
and may be overwritten during the computation.

The order of the derivative `n` must be non-negative. For `n == 0` this function just returns
Legendre polynomials.

# Examples

```jldoctest
julia> dnPl(0.5, 3, 2) # second derivative of P3(x) at x = 0.5
7.5

julia> dnPl(0.5, 4, 0) == Pl(0.5, 4) # zero-th order derivative == Pl(x)
true
```
"""
Base.@propagate_inbounds function dnPl(x, l::Integer, n::Integer,
	A = begin
		_checkvalues(x, l, n)
		# do not allocate A if the value is trivially zero
		if l < n
			return zero(polytype(x))
		end
		zeros(polytype(x), l - n + 1)
	end
	)

	_checkvalues(x, l, n)
	# check if the value is trivially zero in case A is provided in the function call
	if l < n
		return zero(eltype(A))
	end

	cache = OffsetArrays.no_offset_view(A)
	# function barrier, as no_offset_view may be type-unstable
	_unsafednPl!(cache, x, l, n)

	return cache[l - n + 1]
end

"""
	collectPl!(v::AbstractVector, x; [lmax::Integer = length(v) - 1])

Compute the Legendre Polynomials ``P_\\ell(x)`` for the argument `x` and all degrees `l` in `0:lmax`,
and store them in `v`.

At output `v[firstindex(v) + l] == Pl(x,l)`.

# Examples
```jldoctest
julia> v = zeros(4);

julia> collectPl!(v, 0.5)
4-element Array{Float64,1}:
  1.0
  0.5
 -0.125
 -0.4375

julia> v = zeros(0:4);

julia> collectPl!(v, 0.5, lmax = 3) # only l from 0 to 3 are populated
5-element OffsetArray(::Array{Float64,1}, 0:4) with eltype Float64 with indices 0:4:
  1.0
  0.5
 -0.125
 -0.4375
  0.0
```
"""
function collectPl!(v::AbstractVector, x; lmax::Integer  = length(v) - 1)
	checklength(v, lmax + 1)

	iter = LegendrePolynomialIterator(x, lmax)

	for (l, Pl) in pairs(iter)
		v[l + firstindex(v)] = Pl
	end

	v
end

"""
	collectPl(x; lmax::Integer)

Compute the Legendre Polynomial ``P_\\ell(x)`` for the argument `x` and all degrees `l` in `0:lmax`.
Return `P` with indices `0:lmax`, where `P[l] == Pl(x,l)`

# Examples
```jldoctest
julia> collectPl(0.5, lmax = 4)
5-element OffsetArray(::Array{Float64,1}, 0:4) with eltype Float64 with indices 0:4:
  1.0
  0.5
 -0.125
 -0.4375
 -0.2890625
```
"""
collectPl(x; lmax::Integer) = collect(LegendrePolynomialIterator(x, lmax))

"""
	collectdnPl(x; n::Integer, lmax::Integer)

Compute the ``n``-th derivative of a Legendre Polynomial ``P_\\ell(x)`` for the argument `x` and all degrees `l = 0:lmax`.

The order of the derivative `n` must be greater than or equal to zero.

Returns `v` with indices `0:lmax`, where `v[l] == dnPl(x, l, n)`.

# Examples

```jldoctest
julia> collectdnPl(0.5, lmax = 3, n = 2)
4-element OffsetArray(::Array{Float64,1}, 0:3) with eltype Float64 with indices 0:3:
 0.0
 0.0
 3.0
 7.5
```
"""
function collectdnPl(x; lmax::Integer, n::Integer)
	assertnonnegative(lmax)
	n >= 0 || throw(ArgumentError("order of derivative n must be >= 0"))
	v = zeros(polytype(x), 0:lmax)
	if lmax >= n
		collectdnPl!(parent(v), x; lmax = lmax, n = n)
	end
	v
end

"""
	collectdnPl!(v::AbstractVector, x; lmax::Integer, n::Integer)

Compute the ``n``-th derivative of a Legendre Polynomial ``P_\\ell(x)`` for the argument `x` and all degrees `l = 0:lmax`,
and store the result in `v`.

The order of the derivative `n` must be greater than or equal to zero.

At output, `v[l + firstindex(v)] == dnPl(x, l, n)` for `l = 0:lmax`.

# Examples

```jldoctest
julia> v = zeros(4);

julia> collectdnPl!(v, 0.5, lmax = 3, n = 2)
4-element Array{Float64,1}:
 0.0
 0.0
 3.0
 7.5
```
"""
function collectdnPl!(v, x; lmax::Integer, n::Integer)
	assertnonnegative(lmax)
	n >= 0 || throw(ArgumentError("order of derivative n must be >= 0"))
	checklength(v, lmax + 1)

	# trivially zero for l < n
	fill!((@view v[(0:n-1) .+ firstindex(v)]), zero(eltype(v)))
	# populate the other elements
	@inbounds dnPl(x, lmax, n, @view v[(n:lmax) .+ firstindex(v)])

	v
end

end


# ---
# title: 757. Set Intersection Size At Least Two
# id: problem757
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Greedy
# link: <https://leetcode.com/problems/set-intersection-size-at-least-two/description/>
# hidden: true
# ---
#
# An integer interval `[a, b]` (for integers `a < b`) is a set of all
# consecutive integers from `a` to `b`, including `a` and `b`.
#
# Find the minimum size of a set S such that for every integer interval A in
# `intervals`, the intersection of S with A has a size of at least two.
#
#
#
# **Example 1:**
#
#
#
#     Input: intervals = [[1,3],[1,4],[2,5],[3,5]]
#     Output: 3
#     Explanation: Consider the set S = {2, 3, 4}.  For each interval, there are at least 2 elements from S in the interval.
#     Also, there isn't a smaller size set that fulfills the above condition.
#     Thus, we output the size of this set, which is 3.
#
#
# **Example 2:**
#
#
#
#     Input: intervals = [[1,2],[2,3],[2,4],[4,5]]
#     Output: 5
#     Explanation: An example of a minimum sized set is {1, 2, 3, 4, 5}.
#
#
#
#
# **Constraints:**
#
#   * `1 <= intervals.length <= 3000`
#   * `intervals[i].length == 2`
#   * `0 <= ai < bi <= 108`
#
#
## @lc code=start
using LeetCode

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