using Test, Random, Zygote, Distributions, SpecialFunctions, FiniteDifferences
import TruncatedNormal as TruncNorm

for a = -10:10
    d = truncated(Normal(), a, Inf)
    @test TruncNorm.tnmean(a) ≈ mean(d)
    @test TruncNorm.tnstd(a)  ≈ std(d)
    @test TruncNorm.tnvar(a)  ≈ var(d)
    @test a ≤ TruncNorm.tnmean(a) < Inf
    @test 0 ≤ TruncNorm.tnvar(a) ≤ 1
end
@test 1e20 ≤ TruncNorm.tnmean(1e20) < Inf
@test_broken 0 ≤ TruncNorm.tnvar(1e80) ≤ 1

for a = -2:2
    @test mean(TruncNorm.randnt'(a) for _ = 1:10^6) ≈ TruncNorm.tnmean'(a) rtol=0.1
    z, da = TruncNorm.∇randnt(Random.GLOBAL_RNG, a)
    @test da ≈ erfc(z/√2) / erfc(a/√2) * exp((z^2 - a^2)/2)
end

@testset "randnt_half gradient" begin
    for μ = -1:1, σ = 1:2
        dμ_ = mean(gradient(TruncNorm.randnt_half, μ, σ)[1] for _ = 1:10^6)
        dσ_ = mean(gradient(TruncNorm.randnt_half, μ, σ)[2] for _ = 1:10^6)
        dμ, dσ = gradient(μ, σ) do μ, σ
            μ + σ * TruncNorm.tnmean(-μ/σ)
        end
        @test dμ ≈ dμ_ rtol=0.1
        @test dσ ≈ dσ_ rtol=0.1
    end
end

@testset "sqrt1half" begin
    @test TruncNorm.sqrt1half(5) ≈ 5.1925824035672520156
    @test TruncNorm.sqrt1half(0) == 1
    @test TruncNorm.sqrt1half(-1) == TruncNorm.sqrt1half(1) ≈ 1.6180339887498948482
    @test isnan(TruncNorm.sqrt1half(NaN))
    @test TruncNorm.sqrt1half(Inf) == TruncNorm.sqrt1half(-Inf) == Inf
    @test TruncNorm.sqrt1half(1e300) ≈ 1e300
end

@testset "randnt" begin
    @test TruncNorm.randnt(0) > 0
    @test TruncNorm.randnt(1e300) == 1e300
    @test TruncNorm.randnt(Inf) == Inf
    @test isnan(TruncNorm.randnt(NaN))
    @test TruncNorm.randnt(floatmax(Float64)) == floatmax(Float64)
end


"""
Based on
https://hyperopt.github.io/hyperopt/
"""
module TestQuantUniform

using Test
using TreeParzen

# Simple quantuniform
# Minimize the objective over the space
space = Dict(:x => HP.QuantUniform(:x, 0.0, 1.0, 2.0))
best = fmin(params -> params[:x]^2, space, 100)
@test 0 <= best[:x] <= 1

end # module TestQuantUniform
true


using Test
@testset "Tests for triangulation" begin
# import Triangle
@testset "Test function is_left_of" begin
    points = [
         0  0
         1  0
         1  1
         1 -1
        -1  1
        -1 -1
    ]
    @test is_left_of(points, 1, 2, 3) == true
    @test is_left_of(points, 1, 2, 4) == false
    @test is_left_of(points, 1, 2, 5) == true
    @test is_left_of(points, 1, 2, 6) == false
end
@testset "Test function is_right_of" begin
    points = [
         0  0
         1  0
         1  1
         1 -1
        -1  1
        -1 -1
    ]
    @test is_right_of(points, 1, 2, 3) == false
    @test is_right_of(points, 1, 2, 4) == true
    @test is_right_of(points, 1, 2, 5) == false
    @test is_right_of(points, 1, 2, 6) == true
end

@testset "Test of Triangle constructors" begin
    points = [
         0  0
         1  0
         1  1
    ]
    @testset "For two vertices" begin
        t = Triangle(1, 2)
        @test t.v3 == 0
        @test t.t1 == t.t2 == t.t3
        next = t.t3
        @test next.v3 == 0
        @test next.v1 == t.v2
        @test next.v2 == t.v1
        @test next.t1 == next.t2 == next.t3 == t
    end
    @testset "For three vertices" begin
        t = Triangle(points, 1, 2, 3)
        # v1 -> v2 -> v3 should be clockwise
        @test t.v1 == 1
        @test t.v2 == 3
        @test t.v3 == 2
        @test is_left_of(points, t.v1, t.v3, t.v2)
        @test t.v1 == 1

        @test t.t1.t3 == t.t2.t3 == t.t3.t3 == t
        @test t.t1.t1 == t.t3.t2 == t.t2
        @test t.t1.t2 == t.t2.t1 == t.t3
        @test t.t2.t2 == t.t3.t1 == t.t1

        @test t.t2.v2 == t.t3.v1 == t.v1
        @test t.t1.v1 == t.t3.v2 == t.v2
        @test t.t1.v2 == t.t2.v1 == t.v3
    end
end
@testset "Test is_hull" begin
    points = [
         0  0
         1  0
         1  1
    ]
    t = Triangle(points, 1, 2, 3)
    @test is_hull(t.t1)
    @test !is_hull(t)
end
@testset "Test hull navigation function" begin
    points = [
         0  0
         1  0
         1  1
    ]
    t = Triangle(points, 1, 2, 3)
    @test next(t.t1) == t.t2
    @test prev(t.t1) == t.t3
end
@testset "Test find_base_tangente" begin
    points = [
         0  0
         1  0
         1  1
         2  0
         3  0
         3  1
    ]
    t1 = Triangle(points, 1, 2, 3)
    t2 = Triangle(points, 4, 5, 6)
    t3 = Triangle(2, 3)
    t4 = Triangle(4, 5)

    t_left, t_right = find_base_tangente(points, t1.t1, t2.t2)
    @test (t_left.v1, t_right.v2) == (1, 4)
end
@testset "Test _is_above_base" begin
    @testset "Test _is_above_base_left" begin
        @testset "For two edges triangles" begin
            points = [
                0 0
                0 1
                1 0
            ]
            t = Triangle(1, 2)
            @test _is_above_base_left(t, points, 1, 3) == true
            points = [
                0 -1
                0 0
                1 0
            ]
            t = Triangle(2, 1)
            @test _is_above_base_left(t, points, 2, 3) == false
        end
        @testset "For classic triangles" begin
            points = [
                0 0
                1 2
                2 1
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 1, 4) == true
            points = [
                0 0
                1 -1
                1 -2
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 1, 4) == false
            points = [
                0 1
                1 0
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == true
            points = [
                0 -1
                1 0
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == false
            points = [
                0 1
                1 1
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == true
            points = [
                0 -1
                1 1
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == false
        end
    end
    @testset "Test _is_above_base_right" begin
        @testset "For two edges triangles" begin
            points = [
                0 0
                1 0
                1 1
            ]
            t = Triangle(2, 3)
            @test _is_above_base_right(t, points, 1, 2) == true
            points = [
                0 0
                1 0
                1 -1
            ]
            t = Triangle(3, 2)
            @test _is_above_base_right(t, points, 1, 2) == false
        end
        @testset "For classic triangles" begin
            points = [
                0 0
                1 3
                2 1
                3 0
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 4) == true
            points = [
                0 0
                1 -3
                2 -2
                3 0
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 4) == false
            points = [
                0 0
                1 1
                2 0
                3 1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 3) == true
            points = [
                0 0
                1 -1
                2 0
                3 -1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 3) == false
            points = [
                0 0
                1 0
                2 1
                3 1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 2) == true
            points = [
                0 0
                1 0
                2 -1
                3 -1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 2) == false
        end
    end
end
@testset "Test _next" begin
    @testset "For two edges triangles" begin
        t = Triangle(1, 2)
        @test _next_left(t, 1) == t.t3
        @test _next_left(_next_left(t, 1), 1) == t
        @test _next_right(t, 1) == t.t3
        @test _next_right(_next_right(t, 1), 1) == t
    end
    @testset "For classic triangles" begin
        points = [
            0 0
            1 0
            1 1
        ]
        t = Triangle(points, 1, 2, 3)
        @test _next_left(t, 1) == t.t3
        @test _next_left(t, 3) == t.t1
        @test _next_left(t, 2) == t.t2
        @test _next_right(t, 1) == t.t2
        @test _next_right(t, 3) == t.t3
        @test _next_right(t, 2) == t.t1
    end
end
end


using SymbolicRegression, Test
include("test_params.jl")

_inv(x::Float32)::Float32 = 1.0f0 / x
X = rand(Float32, 5, 100) .+ 1
y = 1.2f0 .+ 2 ./ X[3, :]

options = SymbolicRegression.Options(;
    default_params..., binary_operators=(+, *), unary_operators=(_inv,), npopulations=8
)
hallOfFame = EquationSearch(X, y; niterations=8, options=options, numprocs=4)

dominating = calculate_pareto_frontier(X, y, hallOfFame, options)

best = dominating[end]
# Test the score
@test best.loss < maximum_residual / 10


# Method when using Boykov-Kolmogorov as a subroutine
# Kishimoto algorithm

@traitfn function kishimoto(
    flow_graph::::IsDirected,                       # the input graph
    source::Integer,                           # the source vertex
    target::Integer,                           # the target vertex
    capacity_matrix::AbstractMatrix,      # edge flow capacities
    flow_algorithm::BoykovKolmogorovAlgorithm, # keyword argument for algorithm
    routes::Int                                # keyword argument for routes
    )
    # Initialisation
    flow, F, labels = maximum_flow(flow_graph, source, target,
    capacity_matrix, algorithm = flow_algorithm)
    restriction = flow / routes
    flow, F, labels = maximum_flow(flow_graph, source, target, capacity_matrix,
    algorithm = flow_algorithm, restriction = restriction)

    # Loop condition : approximatively not equal is enforced by floating precision
    i = 1
    while flow < routes * restriction && flow ≉ routes * restriction
        restriction = (flow - i * restriction) / (routes - i)
        i += 1
        flow, F, labels = maximum_flow(flow_graph, source, target, capacity_matrix,
        algorithm = flow_algorithm, restriction = restriction)
    end

    # End
    return flow, F, labels
end

"""
    kishimoto(flow_graph, source, target, capacity_matrix, flow_algorithm, routes)

Compute the maximum multiroute flow (for an integer number of `route`s)
between `source` and `target` in `flow_graph` with capacities in `capacity_matrix`
using the [Kishimoto algorithm](http://dx.doi.org/10.1109/ICCS.1992.255031).
Return the value of the multiroute flow as well as the final flow matrix,
along with a multiroute cut if Boykov-Kolmogorov is used as a subroutine.
"""
function kishimoto end
@traitfn function kishimoto(
    flow_graph::::IsDirected,                   # the input graph
    source::Integer,                       # the source vertex
    target::Integer,                       # the target vertex
    capacity_matrix::AbstractMatrix,  # edge flow capacities
    flow_algorithm::AbstractFlowAlgorithm, # keyword argument for algorithm
    routes::Int                            # keyword argument for routes
    )
    # Initialisation
    flow, F = maximum_flow(flow_graph, source, target,
    capacity_matrix, algorithm = flow_algorithm)
    restriction = flow / routes

    flow, F = maximum_flow(flow_graph, source, target, capacity_matrix,
    algorithm = flow_algorithm, restriction = restriction)

    # Loop condition : approximatively not equal is enforced by floating precision
    i = 1
    while flow < routes * restriction && flow ≉ routes * restriction
        restriction = (flow - i * restriction) / (routes - i)
        i += 1
        flow, F = maximum_flow(flow_graph, source, target, capacity_matrix,
        algorithm = flow_algorithm, restriction = restriction)
    end

    # End
    return flow, F
end


export impedanceM1, impedanceM2, impedanceM3, impedanceM4,
        impedanceM5, impedanceM6, impedanceM7, impedanceM8,
        impedanceM9, impedanceM9b

include("constant.jl")
#using .constant

"""

        impedanceM1(f,σ)

Delany and Bazley model

Compute characteristic ground impedance with 1 input parameter

# Arguments
- `σ`: Pa.s.m-2, flow resitivity (realistic values [5x10^3 3x10^7])

# Example:
```julia-repl
julia> impedanceM1(100, 200_000)
```
---
[1]. Denaly & Bazley, Acoustical properties of fibrous absorbent materials,
Appl. Acoustisc, 1970

[2]. Attenborough, Keith, Imran Bashir, and Shahram Taherzadeh. "Outdoor ground impedance models."
The Journal of the Acoustical Society of America 129.5 (2011): 2806-2819.

"""
function impedanceM1(f,σ)
    X = f*ϱ0/σ # in Ref. [2], this equation is typo error
    Zc = 1 + 0.0571*X^(-0.754) + 1im*0.087*X^(-0.732) # Eq. (1b) Ref. [2]
    return Zc
end

"""

        impedanceM2(f,σ)

        impedanceM2(f,σ,α)

Variable prositty models

Compute characteristic ground impedance with 2 input parameters

# Arguments

- `σ`: Pa.s.m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `α`: /m, rate of porosity variation (realistic values [1 10]), defaul α = 8 m-1

# Example:
```julia-repl
julia> impedanceM2(100,200_000)
julia> impedanceM2(100, 200_000,5)
```

---
[1]. Attenborough et al., Outdoor ground impedance models EuroNoise, 2015

"""
function impedanceM2(f,σ,α::Any=8)
    # default value for α = 8m-1 as used for institutional grass (Donato, 1977, JASA)
    Zc =    (1+1im)/sqrt(π*γa*ϱ0)*sqrt(σ/f) +
            1im*c0*α/(8*π*f)
    return Zc
end

"""

        impedanceM3(f,σ,Ω)

        impedanceM3(f,σ,Ω,svor,sent)

Wilson relaxation model

Compute characteristic ground impedance with 2 input parameters

# Arguments
- `σ`: Pa.s.m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω`: prosity (realistic values [0.1 0.8]) (Table 10.1, ref. [1])

- `svor`, `sent`: shape factor, default values = 1

# Example:
```julia-repl
julia> impedanceM3(100,200_000,0.5)
```

---
[1]. Ostashev, Vladimir E., and D. Keith Wilson.
Acoustics in moving inhomogeneous media. CRC Press, 2015.

"""
function impedanceM3(f,σ,Ω, svor::Any=1,sent::Any=1)
    ω = 2π*f
    q = Ω^-0.25 # page 371
    τ0 = (2*ϱ0*q^2)/(σ*Ω) # Eq 10.53

    term1 = 2im/(ω*τ0) + 1 + 1/sqrt(1-1im*svor^2*ω*τ0)
    term2 = 1 + (γa-1)/sqrt(1 - 1im*sent^2*Npr*ω*τ0)
    Zc = (q/Ω)*(term1/term2)^0.5 # Eq 10.54

    return Zc
end

"""

        impedanceM4(f,σ,Ω)

        impedanceM4(f,σ,Ω,q)

Zwikker and Kosten model

Compute characteristic ground impedance with 3 input parameters

# Arguments

- `σ`: Pa.s.m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω`: prosity (realistic values [0.1 0.8])

- `q`: tortuosity (realistic values [1.8 1.0]), default `q` = 1.0

# Example:
```julia-repl
julia> impedanceM4(100,200_000,0.5)
julia> impedanceM4(100,200_000,0.5,1.0)
```

---
[1]. Zwikker, Cornelis, and Cornelis Willem Kosten. Sound absorbing materials. Elsevier publishing company, 1949.

[2]. Attenborough, Keith, Imran Bashir, and Shahram Taherzadeh. "Outdoor ground impedance models."
The Journal of the Acoustical Society of America 129.5 (2011): 2806-2819.

"""
function impedanceM4(f,σ,Ω,q::Any=1.0)
    # the default value for q is for soft soil
    T = q*q
    ω = 2π*f

    Zc = (1/Ω) * ( T + 1im*Ω*σ/(ω*ϱ0) )^0.5 # Eq. 2a (ref. [2])
    return Zc
end

"""

        impedanceM5(f,σ,Ω)

        impedanceM5(f,σ,Ω,q)

Taraldsen model

Compute characteristic ground impedance with 3 input parameters

# Arguments

- `σ`: Pa.s.m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω``: prosity (realistic values [0.1 0.8])

- `q`: tortuosity (realistic values [1.8 1.0]), default `q` = 1.0

# Example:
```julia-repl
julia> impedanceM5(100,200_000,0.5)
julia> impedanceM5(100,200_000,0.5,1.0)
```

---
[1]. Taraldsen, G. "The Delany-Bazley impedance model and Darcy's law."
Acta Acustica united with Acustica 91.1 (2005): 41-50.

"""
function impedanceM5(f,σ,Ω,q::Any=1.0)
    β✶ = sqrt(γa)*Ω/q
    f✶ = 1/(2π*ϱ0)*(Ω/q)*(σ/q) # Eq. 14 note √T = q
    β = β✶*(1+1im*f✶/f)^-0.5
    Zc = 1/β

    return Zc
end

"""

        impedanceM6(f,σ,Ω)

        impedanceM6(f,σ,Ω,q)

Hamet phenomenological models

Compute characteristic ground impedance with 3 input parameters

# Arguments

- `σ`: Pa.s.m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω`: prosity (realistic values [0.1 0.8])

- `q`: tortuosity (realistic values [1.8 1.0]), default `q` = 1.0

# Example:
```julia-repl
julia> impedanceM6(100,200_000,0.5)
julia> impedanceM6(100,200_000,0.5,1.0)
```

---
[1]. Bérengier, M. C., et al. "Porous road pavements: Acoustical characterization and propagation effects."
The Journal of the Acoustical Society of America 101.1 (1997): 155-162.
"""
function impedanceM6(f,σ,Ω,q::Any=1.0)

    q² = q*q
    fμ = Ω*σ/(2π*ϱ0*q²) # Eq. 19
    fθ = σ/(2π*ϱ0*Npr)  # Eq. 20
    Fμ = 1+1im*fμ/f     # Eq. 21
    Fθ = 1 + 1im*fθ/f   # Eq. 22
    Zc = (q/Ω)*Fμ^0.5*( γa - (γa-1)/Fθ )^-0.5   # Eq. 23, without ρ0c term

    return Zc
end

"""

        impedanceM7(f,σ,Ω)

        impedanceM7(f,σ,Ω,q)

Identical tortuous pores model

Compute characteristic ground impedance with 3 input parameters

# Arguments

- `σ`: Pa s m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω`: prosity (realistic values [0.1 0.8])

- `q`: tortuosity (realistic values [1.8 1.0]), default `q` = 1.0

# Example:
```julia-repl
julia> impedanceM7(100,200_000,0.5)
julia> impedanceM7(100,200_000,0.5,1.0)
```

---
[1]. Attenborough, Keith, Imran Bashir, and Shahram Taherzadeh. "Outdoor ground impedance models."
The Journal of the Acoustical Society of America 129.5 (2011): 2806-2819.

"""
function impedanceM7(f,σ,Ω,q::Any=1.0)
    ω = 2π*f
    q² = q*q # or = T
    GS(x) = 1-tanh(x*sqrt(-1im)) / (x*sqrt(-1im)) # Eq. 6e

    λ = sqrt((3*ϱ0*ω*q²)/(Ω*σ)) # Eq. 6f
    ρ_λ = ϱ0/GS(λ) # Eq. 6c
    C_λ = 1/(γa*P0) * ( γa - (γa-1)*GS(λ*sqrt(Npr)) ) # Eq. 6d

    Zc = (ϱ0*c0)^-1 * ( (q²/Ω^2)*ρ_λ/C_λ )^0.5 # Eq. 6b

    return Zc
end

"""

        impedanceM8(f,σ,Ω)

        impedanceM8(f,σ,Ω,sf)

        impedanceM8(f,σ,Ω,sf,q)

Attenborough four parameter model

Compute characteristic ground impedance with 4 input parameters

# Arguments

- `σ`: Pa s m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω`: prosity (realistic values [0.1 0.8])

- `q`: tortuosity (realistic values [1.8 1.0]), default `q` = 1.0

- `sf`: pore shape factor (realistic values [0.6 1]), default `sf` = 0.75

# Example:
```julia-repl
julia> impedanceM8(100,200_000,0.5)
julia> impedanceM8(100,200_000,0.5,0.75)
julia> impedanceM8(100,200_000,0.5,0.75,1.0)
```

---
[1]. Attenborough, Keith. "Acoustical impedance models for outdoor ground surfaces."
Journal of Sound and Vibration 99.4 (1985): 521-544.

"""
function impedanceM8(f,σ,Ω,sf::Any=0.75,q::Any=1.0)
    ω = 2π*f
    q² = q*q

    kb =    (γa*Ω)^0.5 *
            ( (4/3 - (γa-1)/γa*Npr)*q²/Ω + 1im*sf^2*σ/(ω*ϱ0) )^0.5 # Eq. 12

    Zc = ( 4*q²/(3*Ω) + 1im*sf^2*σ/(ω*ϱ0) ) / kb # Eq. 11

    return Zc
end

"""

        impedanceM9(f,s̄,Ω)

        impedanceM9(f,s̄,Ω,σs)

Horoshenkov porous media model

Compute characteristic ground impedance with 3 input parameters

# Arguments

- `s̄`: median pore size (realistic values [2 2000 μm]), this parameter linked
to flow resistivity

- `Ω`: prosity (equivalent to Ω of other models) realistics values [0.1 0.8]

- `σs`: standard deviation of pore size distribution (realistic values [0 0.5]),
default `σs` = 0.3

# Example:
```julia-repl
julia> impedanceM9(100,200*10^-6,0.5)
julia> impedanceM9(100,200*10^-6,0.5,0.3)
```

---
[1]. Horoshenkov, Kirill V., Alistair Hurrell, and Jean-Philippe Groby. "A three-parameter analytical model for the acoustical properties of porous media."
The Journal of the Acoustical Society of America 145.4 (2019): 2512-2517.
"""
function impedanceM9(f,s̄,Ω,σs::Any=0.3)
    𝜙 = Ω # for consistent with original notations (ref. [1])

    ω = 2π*f
    α∞ = exp( 4*(σs*log(2))^2 ) # tourtuosity See Eq. 12

    θρ1 = 1/3
    θρ2 = exp( -1/2*(σs*log(2))^2 ) / sqrt(2)
    θρ3 = θρ1/θρ2

    σ = 8*η0*α∞ / (s̄^2*𝜙) * exp( 6*(σs*log(2))^2 )    # Eq. 15
    σ′ = 8*η0*α∞ / (s̄^2*𝜙) * exp( -6*(σs*log(2))^2 )  # Eq. 18

    ϵρ = sqrt( -1im*ω*ϱ0*α∞/(𝜙*σ) )
    F̃ρ = (1 + θρ3*ϵρ + θρ1*ϵρ^2) / (1 + θρ3*ϵρ)    # Eq. 14

    θc1 = 1/3
    θc2 = exp( 3/2*(σs*log(2))^2 ) / sqrt(2)
    θc3 = θc1/θc2

    ϵc = sqrt( -1im*ω*ϱ0*Npr*α∞ / (𝜙*σ′) )
    F̃c = (1 + θc3*ϵc + θc1*ϵc^2) / (1+ θc3*ϵc)     # Eq. 17

    C̃ω = 𝜙/(γa*P0) * ( γa - (γa-1)/(1+ϵc^-2*F̃c) )    # Eq 16
    ρ̃ω = ϱ0*α∞/𝜙 * (1 + ϵρ^-2*F̃ρ)    # Eq. 13

    Zc = 1/(ϱ0*c0) * sqrt( ρ̃ω / C̃ω)   # Eq. 19 characteristic impedance

    return Zc
end

"""

        impedanceM9(f,σ,Ω)

        impedanceM9(f,σ,Ω,σs)

Horoshenkov porous media model (σ input)

Compute characteristic ground impedance with 3 input parameters

# Arguments

- `σ`: Pa s m-2, flow resitivity (realistic values [5x10^3 3x10^7])

- `Ω`: prosity (equivalent to Ω of other models) realistics values [0.1 0.8]

- `σs`: standard deviation of pore size distribution (realistic values [0 0.5]),
default `σs` = 0.3

# Example:
```julia-repl
julia> impedanceM9(100,200_000,0.5)
julia> impedanceM9(100,200_000,0.5,0.3)
```

---
[1]. Horoshenkov, Kirill V., Alistair Hurrell, and Jean-Philippe Groby. "A three-parameter analytical model for the acoustical properties of porous media."
The Journal of the Acoustical Society of America 145.4 (2019): 2512-2517.
"""
function impedanceM9b(f,σ,Ω,σs::Any=0.3)
    𝜙 = Ω # for consistent with original notations (ref. [1])

    ω = 2π*f
    α∞ = exp( 4*(σs*log(2))^2 ) # tourtuosity See Eq. 12

    θρ1 = 1/3
    θρ2 = exp( -1/2*(σs*log(2))^2 ) / sqrt(2)
    θρ3 = θρ1/θρ2

    # convert σ to s̄ using equation 15
    s̄ = ( 8*η0*α∞ / (σ*𝜙) * exp( 6*(σs*log(2))^2 ) )^0.5

    #σ = 8*η0*α∞ / (s̄^2*𝜙) * exp( 6*(σs*log(2))^2 )    # Eq. 15
    σ′ = 8*η0*α∞ / (s̄^2*𝜙) * exp( -6*(σs*log(2))^2 )  # Eq. 18

    ϵρ = sqrt( -1im*ω*ϱ0*α∞/(𝜙*σ) )
    F̃ρ = (1 + θρ3*ϵρ + θρ1*ϵρ^2) / (1 + θρ3*ϵρ)    # Eq. 14

    θc1 = 1/3
    θc2 = exp( 3/2*(σs*log(2))^2 ) / sqrt(2)
    θc3 = θc1/θc2

    ϵc = sqrt( -1im*ω*ϱ0*Npr*α∞ / (𝜙*σ′) )
    F̃c = (1 + θc3*ϵc + θc1*ϵc^2) / (1+ θc3*ϵc)     # Eq. 17

    C̃ω = 𝜙/(γa*P0) * ( γa - (γa-1)/(1+ϵc^-2*F̃c) )    # Eq 16
    ρ̃ω = ϱ0*α∞/𝜙 * (1 + ϵρ^-2*F̃ρ)    # Eq. 13

    Zc = 1/(ϱ0*c0) * sqrt( ρ̃ω / C̃ω)   # Eq. 19 characteristic impedance

    return Zc
end


"""
   magsq!(magusq::Nodes{Primal/Dual},u::Edges{Primal/Dual})

Calculate the in-placed squared magnitude of vector grid data `u`, placing the result on
the cell centers in `magusq`.
"""
function magsq!(umagsq::Nodes{C},u::Edges{C}) where {C <: CellType}

  usq = u∘u
  usq_nodes = Nodes(C,u)

  grid_interpolate!(usq_nodes,usq.u)
  umagsq .= usq_nodes

  grid_interpolate!(usq_nodes,usq.v)
  umagsq .+= usq_nodes

  return umagsq
end

"""
   magsq(u::Edges{Primal/Dual}) -> Nodes{Primal/Dual}

Calculate the squared magnitude of vector grid data `u`, placing the result on
the cell centers.
"""
magsq(u::Edges{C}) where {C<:CellType} = magsq!(Nodes(C,u),u)

"""
    mag!(magu::Nodes{Primal/Dual},u::Edges{Primal/Dual})

Calculate the in-placed magnitude of vector grid data `u`, placing the result on
the cell centers in `magu`.
"""
function mag!(magu::Nodes{C},u::Edges{C}) where {C <: CellType}
    magsq!(magu,u)
    @. magu = sqrt(magu)
end

"""
   mag(u::Edges{Primal/Dual}) -> Nodes{Primal/Dual}

Calculate the magnitude of vector grid data `u`, placing the result on
the cell centers.
"""
mag(u::Edges{C}) where {C<:CellType} = mag!(Nodes(C,u),u)

"""
    directional_derivative!(out,f,q)

Compute the directional derivative of `f` in the direction of `q`, ``q\\cdot\\nabla f``,
and put the result into `out`. Note that the result is not scaled by any grid spacing.
"""
function directional_derivative!(out::Edges{C},f::Edges{C},q::Edges{C}) where {C <: CellType}

    qdf = EdgeGradient(C,f)
    df = EdgeGradient(C,f)
    qt = EdgeGradient(C,f)

    out .= 0.0
    product!(qdf,transpose(grid_interpolate!(qt,q)),grad!(df,f))
    grid_interpolate!(out,qdf)

    #=
    # some temp storage
    q_dual = Nodes(othertype(C),f)
    q_primal = Nodes(C,f)
    outx = XEdges(C,f)
    outy = YEdges(C,f)

    # compute the gradient of f
    df = EdgeGradient(C,f)
    grad!(df,f)

    # compute qy*(dfx/dy) after first transforming q.v to dual nodes and then transforming product to primal x-edges
    grid_interpolate!(outx,grid_interpolate!(q_dual,q.v) ∘ df.dudy)
    out.u .= outx

    # compute qx*(dfx/dx) after first transforming q.u to primal nodes
    # and then interpolating product to primal x-edges
    grid_interpolate!(outx,grid_interpolate!(q_primal,q.u) ∘ df.dudx)
    out.u .+= outx

    # compute qx*(dfy/dx) after first transforming qx to dual nodes and then transforming product to primal y-edges
    grid_interpolate!(outy,grid_interpolate!(q_dual,q.u) ∘ df.dvdx)
    out.v .= outy

    # compute qy*(dfy/dy) after first transforming qy to primal nodes
    # and then interpolating product to primal y-edges
    grid_interpolate!(outy,grid_interpolate!(q_primal,q.v) ∘ df.dvdy)
    out.v .+= outy
    =#

    return out

end

"""
    directional_derivative_conserve!(out,f,q)

Compute the conservative form of the directional derivative of `f` in the direction
of `q`, ``\\nabla\\cdot(qf)``, and put the result
into `out`. Note that the result is not scaled by any grid spacing. This form is only
appropriate if `q` is divergence free.
"""
function directional_derivative_conserve!(out::Edges{C},f::Edges{C},q::Edges{C}) where {C <: CellType}

    # some temp storage
    q_dual = Nodes(othertype(C),f)
    f_dual = Nodes(othertype(C),f)

    qf = q ∘ f
    qfnode = NodePair(C,f)

    # interpolate both fx and qx to primal nodes, multiply element by element, assign to NodePair.u
    grid_interpolate!(qfnode.u,qf.u)

    # interpolate both f.u and q.v to dual nodes, multiply element by element, assign to NodePair.v
    qfnode.v .= grid_interpolate!(q_dual,q.v) .* grid_interpolate!(f_dual,f.u)

    # take divergence and assign to out.u
    divergence!(out.u,qfnode)

    qfnode = NodePair(othertype(C),f)

    # interpolate both fy and qx to dual nodes, multiply element by element, assign to NodePair.u
    qfnode.u .= grid_interpolate!(q_dual,q.u) .* grid_interpolate!(f_dual,f.v)

    # interpolate both fy and qy to primal nodes, multiply element by element, assign to NodePair.v
    grid_interpolate!(qfnode.v,qf.v)

    # take divergence and assign to out.v
    divergence!(out.v,qfnode)

    return out

end

"""
    curl_cross!(out,a,b)

Compute the curl of the cross product of `a` and `b`,
``\\nabla\\times(a\\times b)``,
and put the result into `out`. Note that the result is not scaled by any grid spacing.
"""
function curl_cross!(out::Edges{C},a::Edges{C},b::Edges{C}) where {C <: CellType}

    # some temp storage
    bx_dual = Nodes(othertype(C),a)
    by_dual = Nodes(othertype(C),a)

    ax_dual = Nodes(othertype(C),a)
    ay_dual = Nodes(othertype(C),a)

    acrossb = Nodes(othertype(C),a)

    grid_interpolate!((ax_dual,ay_dual),a)
    grid_interpolate!((bx_dual,by_dual),b)

    acrossb .= ax_dual.*by_dual - ay_dual.*bx_dual
    curl!(out,acrossb)

    return out

end

"""
    convective_derivative!(out,q)

Compute the convective derivative of `q` in the form
``u\\cdot\\nabla u``
and put the result into `out`. Note that the result is not scaled by any grid spacing.
"""
convective_derivative!(out::Edges{C},q::Edges{C}) where {C<:CellType} = directional_derivative!(out,q,q)

"""
    convective_derivative_rot!(out,q)

Compute the rotational form of the convective derivative of `q` in the form
``\\frac{1}{2}\\nabla|u|^2-u\\times(\\nabla\\times u)``
and put the result into `out`. Note that the result is not scaled by any grid spacing.
"""
function convective_derivative_rot!(out::Edges{C},q::Edges{C}) where {C<:CellType}

    # some temp storage
    qx_primal = Nodes(C,q)
    qy_primal = Nodes(C,q)
    ω = curl(q)

    grid_interpolate!((qx_primal,qy_primal),q)
    grad!(out,qx_primal ∘ qx_primal + qy_primal ∘ qy_primal)
    out .*= 0.5

    q_dual = Nodes(othertype(C),q)
    qcrossω = Edges(C,q)
    grid_interpolate!(qcrossω.u,grid_interpolate!(q_dual, q.v) ∘ ω)
    grid_interpolate!(qcrossω.v,grid_interpolate!(q_dual,-q.u) ∘ ω)

    out .-= qcrossω

    return out

end


module CartPendulum

global const g = 9.82

using DifferentialEquations
using ControlSystems
mutable struct Pendulum{T<:Real}
    M::T
    m::T
    l::T
    b::T
    ulims::Tuple{T,T}
    nx::T
    nθ::T
    sgn::T
end

function Pendulum()
    Pendulum{Float64}(1., 0.2, 0.6, 0.1, (-10., 10.), 0.1, 0.01,1)
end

function I(pen::Pendulum{T}) where T<:Real
    return pen.m*pen.l^2/3
end

function linearize(pen::Pendulum{<:Real})
    (M, m, l, b) = (pen.M, pen.m, pen.l, pen.b)
    denom = I(pen) * (M + m) + M * m * l^2
    a22 = -(I(pen) + m * l^2) * b / denom
    a23 = m^2 * g * l^2 / denom;
    a42 = -m * l * b / denom;
    a43 = m * g * l * (M + m) / denom;
    b2 = (I(pen) + m * l^2) / denom;
    b4 = m * l / denom;

    A = [0 1 0 0
    0 a22 a23 0
    0 0 0 1
    0 a42 a43 0]

    B = [0;b2;0;b4];
    # For now, assume full state feedback
    # C = diagm(ones(4));
    C = [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
    return ss(A, B, C, 0)
end

# Deterministic part
function f!(dx::AbstractVector{T},
    x::AbstractVector{T},
    pen::Pendulum{T}, t::T) where T<:Real
    (M, m, l, b, ulims) = (pen.M, pen.m, pen.l, pen.b, pen.ulims)

    lump1 = m^2 * l^2 / (I(pen) + m * l^2)
    c = cos(x[3])
    s = sin(x[3])

    dx[1] = x[2];
    dx[2] = (-b * x[2] + g * lump1 * c * s + m * l * s * x[4]^2 + pen.sgn*sat(x[5], ulims)) / (M + m - lump1 * c^2)
    dx[3] = x[4];
    dx[4] = (-m * l * c * dx[2] - m * g * l * s) / (I(pen) + m * l^2);
    dx[5:end] .= 0; # Controller state is set by callback and held constant otherwise
end

# Stochastic noise
function g!(dx::AbstractVector{T},
    x::AbstractVector{T},
    pen::Pendulum{T}, t::T) where T<:Real

    (M, m, l, b) = (pen.M, pen.m, pen.l, pen.b)
    (nx, nθ) = (pen.nx, pen.nθ)
    lump1 = m^2 * l^2 / (I(pen) + m * l^2)
    c = cos(x[3])
    s = sin(x[3])

    dx[1] = 0;
    dx[2] = pen.nx/ (M + m - lump1 * c^2)
    dx[3] = 0;
    dx[4] = nθ/(I(pen) + m * l^2)
    dx[5:end] .= 0
end

function sat(val::T, lims::Tuple{T,T}) where T<:Real
    return min(max(val, lims[1]), lims[2])
end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    RegularSampling(n1, n2, ..., np)

Generate samples regularly using `n1` points along the first
parametric dimension, `n2` points along the second parametric
dimension, ..., `np` poitns along the last parametric dimension.

## Example

Sample sphere regularly with 360 longitudes and 180 latitudes:

```julia
sample(Sphere((0,0,0), 1), RegularSampling(360, 180))
```
"""
struct RegularSampling{N} <: ContinuousSamplingMethod
  sizes::NTuple{N,Int}
end

RegularSampling(sizes::Vararg{Int,N}) where {N} = RegularSampling(sizes)

function sample(segment::Segment, method::RegularSampling)
  n = first(method.sizes)
  a, b = extrema(segment)

  δ = (b - a) / (n - 1)

  (a + (i-1)*δ for i in 1:n)
end

function sample(box::Box, method::RegularSampling)
  sz = _adjust_sizes(method.sizes, paramdim(box))
  l, u = extrema(box)

  # origin and spacing
  or, sp = l, (u - l) ./ (sz .- 1)

  ivec(or + (ind.I .- 1) .* sp for ind in CartesianIndices(sz))
end

function sample(sphere::Sphere{2,T}, method::RegularSampling) where {T}
  sz = _adjust_sizes(method.sizes, paramdim(sphere))
  c, r = center(sphere), radius(sphere)

  V = T <: AbstractFloat ? T : Float64
  θmin, θmax = V(0), V(2π)
  δθ = (θmax - θmin) / sz[1]
  θrange = range(θmin, stop=θmax-δθ, length=sz[1])

  r⃗(θ) = Vec{2,V}(r*cos(θ), r*sin(θ))

  ivec(c + r⃗(θ) for θ in θrange)
end

# spherical coordinates in ISO 80000-2:2019 convention
function sample(sphere::Sphere{3,T}, method::RegularSampling) where {T}
  sz = _adjust_sizes(method.sizes, paramdim(sphere))
  c, r = center(sphere), radius(sphere)

  V = T <: AbstractFloat ? T : Float64
  θmin, θmax = V(0), V(π)
  φmin, φmax = V(0), V(2π)
  δθ = (θmax - θmin) / (sz[1] + 1)
  δφ = (φmax - φmin) / (sz[2]    )
  θrange = range(θmin+δθ, stop=θmax-δθ, length=sz[1])
  φrange = range(φmin, stop=φmax-δφ, length=sz[2])

  r⃗(θ, φ) = Vec{3,V}(r*sin(θ)*cos(φ), r*sin(θ)*sin(φ), r*cos(θ))

  ivec(c + r⃗(θ, φ) for θ in θrange, φ in φrange)
end

function sample(ball::Ball{Dim,T}, method::RegularSampling) where {Dim,T}
  sz = _adjust_sizes(method.sizes, paramdim(ball))
  c, r = center(ball), radius(ball)

  V = T <: AbstractFloat ? T : Float64
  smin, smax = V(0), V(1)
  δs = (smax - smin) / (last(sz) - 1)
  srange = range(smin+δs, stop=smax, length=last(sz))

  # reuse samples on the boundary
  points = sample(Sphere(c, r), RegularSampling(sz[1:Dim-1]))

  scale(p, s) = Point(s * coordinates(p))

  ivec(scale(p, s) for p in points, s in srange)
end

# helper function to adjust sizes to a given length
function _adjust_sizes(sizes, len)
  N, S = len, length(sizes)
  S != N ? ntuple(i -> i ≤ S ? sizes[i] : last(sizes), N) : sizes
end


struct SymEGRQSCholesky{T,UT<:AbstractMatrix,WT<:AbstractMatrix,dT<:AbstractArray} <: AbstractMatrix{T}
	Ut::UT
    Wt::WT
	d::dT
    n::Int
    p::Int
    function SymEGRQSCholesky{T,UT,WT,dT}(Ut,Wt,d,n,p) where
		{T,UT<:AbstractMatrix,WT<:AbstractMatrix,dT<:AbstractArray}
		Up, Un = size(Ut)
		Wp, Wn = size(Wt)
		(Up == Wp && Un == Wn && Un == length(d)) || throw(DimensionMismatch())
        new(Ut,Wt,d,n,p)
    end
end

#### Creating W and dbar s.t. L = tril(UW',-1) + diag(dbar) ####
function dss_create_wdbar(Ut::AbstractArray, Vt::AbstractArray, d::AbstractArray)
    p, n = size(Ut);
    P  = zeros(p, p);
    Wt  = similar(Vt);
    dbar = zeros(n);
    @inbounds for i = 1:n
        tmpU  = Ut[:,i]
        tmpW  = Vt[:,i] - P*tmpU;
		tmpds = sqrt(abs(tmpU'*tmpW + d[i]));
        #tmpds = sqrt(tmpU'*tmpW + d[i]);
        tmpW  = tmpW/tmpds
        Wt[:,i] = tmpW;
        dbar[i] = tmpds;
        P += tmpW*tmpW';
    end
    return Wt, dbar
end

function cholesky(K::SymEGRQSMatrix{T,UT,VT,dT}) where
	{T,UT<:AbstractMatrix,VT<:AbstractMatrix,dT<:AbstractArray}
	Wt, dbar = dss_create_wdbar(K.Ut,K.Vt,K.d)
	SymEGRQSCholesky{T,typeof(K.Ut),typeof(Wt),typeof(dbar)}(K.Ut,Wt,dbar,K.n,K.p)
end

function cholesky(K::SymEGRSSMatrix{T,UT,VT}, σ::Number) where
	{T,UT<:AbstractMatrix,VT<:AbstractMatrix}
	Wt, dbar = dss_create_wdbar(K.Ut,K.Vt,ones(K.n)*σ)
	SymEGRQSCholesky{T,typeof(K.Ut),typeof(Wt),typeof(dbar)}(K.Ut,Wt,dbar,K.n,K.p)
end

function SymEGRQSCholesky(K::SymEGRSSMatrix{T,UT,VT}, σ::Number) where
	{T,UT<:AbstractMatrix,VT<:AbstractMatrix}
	Wt, dbar = dss_create_wdbar(K.V,K.V,ones(K.n)*σ)
	SymEGRQSCholesky{T,typeof(K.Ut),typeof(Wt),typeof(dbar)}(K.Ut,Wt,dbar,K.n,K.p)
end

function SymEGRQSCholesky(Ut::UT, Wt::WT, d::AbstractArray) where
	{T,UT<:AbstractMatrix,WT<:AbstractMatrix}
	p, n = size(Ut);
	SymEGRQSCholesky{eltype(Ut),typeof(Ut),typeof(Wt),typeof(d)}(Ut,Wt,d,n,p)
end

########################################################################
#### Helpful properties. Not nessecarily computionally efficient    ####
########################################################################
Matrix(K::SymEGRQSCholesky) = getproperty(K,:L)
size(K::SymEGRQSCholesky) = (K.n, K.n)
size(K::SymEGRQSCholesky,d::Int) = (1 <= d && d <=2) ? size(K)[d] : throw(ArgumentError("Invalid dimension $d"))

function getindex(K::SymEGRQSCholesky, i::Int, j::Int)
	i > j && return dot(K.Ut[:,i], K.Wt[:,j])
	i == j && return K.d[i]
	return 0
end

function getproperty(K::SymEGRQSCholesky, d::Symbol)
    if d === :U
        return UpperTriangular(triu(K.Wt'*K.Ut,1) + Diagonal(K.d))
    elseif d === :L
        return LowerTriangular(tril(K.Ut'*K.Wt,-1) + Diagonal(K.d))
    else
        return getfield(K, d)
    end
end

Base.propertynames(F::SymEGRQSCholesky, private::Bool=false) =
    (:U, :L, (private ? fieldnames(typeof(F)) : ())...)

function Base.show(io::IO, mime::MIME{Symbol("text/plain")},
		 K::SymEGRQSCholesky{<:Any,<:AbstractArray,<:AbstractArray,<:AbstractArray})
    summary(io, K); println(io)
    show(io, mime, K.L)
end

########################################################################
#### Cholesky factorization of Higher-order quasiseparable matrices ####
########################################################################

#### Forward substitution ####
function dss_forward!(X::AbstractArray,Ut::AbstractArray,Wt::AbstractArray,
                     ds::AbstractArray,B::AbstractArray)
    p, n = size(Ut)
    mx = size(B,2)
    Wbar = zeros(p,mx);
    @inbounds for i = 1:n
        tmpU = Ut[:,i];
        tmpW = Wt[:,i];
        X[i:i,:] = (B[i:i,:] - tmpU'*Wbar)/ds[i];
        Wbar += tmpW .* X[i:i,:];
    end
end
#### Backward substitution ####
function dssa_backward!(X::AbstractArray,Ut::AbstractArray,Wt::AbstractArray,
                       ds::AbstractArray,B::AbstractArray)
    p, n = size(Ut)
    mx = size(B,2)
    Ubar = zeros(p,mx);
    @inbounds for i = n:-1:1
        tmpU = Ut[:,i];
        tmpW = Wt[:,i];
        X[i:i,:] = (B[i:i,:] - tmpW'*Ubar)/ds[i];
        Ubar += tmpU .* X[i:i,:];
    end
end

#### Multiplying with L ####
function dss_tri_mul!(Y::AbstractArray,K::SymEGRQSCholesky,X::AbstractArray)
     p, n = size(K.Ut)
     mx = size(X, 2)
     Wbar = zeros(p, mx)
     @inbounds for i = 1:n
         tmpW = K.Wt[:,i]
         tmpU = K.Ut[:,i]
         tmpX = X[i:i,:];
         Y[i,:] = tmpU'*Wbar + K.d[i]*tmpX;
         Wbar  +=  tmpW*tmpX;
     end
	 return Y
end

#### Multiplying with L' ####
function dssa_tri_mul!(Y::AbstractArray,K::SymEGRQSCholesky,X::AbstractArray)
     p, n = size(K.Ut)
     mx = size(X, 2)
     Ubar = zeros(p,mx)
     @inbounds for i = n:-1:1
         tmpW = K.Wt[:,i]
         tmpU = K.Ut[:,i]
         tmpX = X[i:i,:]
         Y[i,:] = tmpW'*Ubar + K.d[i]*tmpX;
         Ubar = Ubar + tmpU*tmpX;
     end
	 return Y
end

#### Squared norm of columns of L = tril(UW',-1) + diag(dbar) ####
function squared_norm_cols(Ut::AbstractArray,Wt::AbstractArray,
                         dbar::AbstractArray)
    p, n = size(Ut)
    P = zeros(p, p)
    c = zeros(n)
    @inbounds for i = n:-1:1
        tmpW = Wt[:,i]
        tmpU = Ut[:,i]
        c[i] = dbar[i]^2 + tmpW'*P*tmpW
        P += tmpU*tmpU'
    end
    return c
end
#### Implicit inverse of  L = tril(UW',-1) + diag(dbar) ####
function dss_create_yz(Ut::AbstractArray, Wt::AbstractArray,
                     dbar::AbstractArray)
    p, n = size(Ut)
    Y = zeros(n,p)
    Z = zeros(n,p)
    dss_forward!(Y, Ut, Wt, dbar, Ut')
    dssa_backward!(Z, Ut, Wt, dbar, Wt')
    # Probably best not to use inv
    return copy(Y'), copy((Z*inv(Ut*Z - Diagonal(ones(p))))')
	#return Y, Z*((U'*Z - Diagonal(ones(m)))\Diagonal(ones(m)))
end

function fro_norm_L(L::SymEGRQSCholesky)
	return sum(squared_norm_cols(L.Ut, L.Wt, L.d))
end

function trinv(L::SymEGRQSCholesky)
	Yt, Zt = dss_create_yz(L.Ut, L.Wt, L.d)
	return sum(squared_norm_cols(Yt, Zt, L.d.^(-1)))
end

function tr(Ky::SymEGRQSCholesky, K::SymEGRSSMatrix)
	n = Ky.n;
	p = Ky.p;
	c = Ky.d;
	Ut = K.Ut;
	Vt = K.Vt;
	Yt, Zt = dss_create_yz(Ky.Ut, Ky.Wt, Ky.d);
	b = 0;
	P = zeros(p,p);
	R = zeros(p,p);
	@inbounds for k = 1:Ky.n
		yk = Yt[:,k];
		zk = Zt[:,k];
		uk = Ut[:,k];
		vk = Vt[:,k];
		cki = c[k]^(-1);
		b += yk'*P*yk + 2*yk'*R*uk*cki + uk'*vk*(cki^2);
		P += ((uk'*vk)*zk)*zk' + zk*(R*uk)' + (R*uk)*zk';
		R += zk*vk';
	end
	return b
end

#### Log-determinant ####
function logdet(K::SymEGRQSCholesky)
	dd = sum(log.(K.d))
    return dd
end

#### Determinant ####
function det(K::SymEGRQSCholesky)
    return exp(logdet(K))
end

mul!(y::AbstractVecOrMat, K::SymEGRQSCholesky, x::AbstractVecOrMat) =
	dss_tri_mul!(y,K,x)
mul!(y::AbstractVecOrMat, K::Adjoint{<:Any,<:SymEGRQSCholesky}, x::AbstractVecOrMat) =
	dssa_tri_mul!(y,K.parent,x)
function (\)(F::SymEGRQSCholesky, B::AbstractVecOrMat)
	X = similar(B)
	dss_forward!(X,F.Ut,F.Wt,F.d,B)
	return X
end

function (\)(F::Adjoint{<:Any,<:SymEGRQSCholesky}, B::AbstractVecOrMat)
	Y = similar(B)
	dssa_backward!(Y,F.parent.Ut,F.parent.Wt,F.parent.d,B)
	return Y
end


using LinearAlgebra
using Printf
using Random

"""
    ABCDs, umax, S = scaleABCD(ABCD, nlev=2, f=0, xlim=1, ymax=nlev+5, umax, N=100000, N0=10)

Scale the loop filter of a general delta-sigma modulator for dynamic range.

Inputs:
- `ABCD`:   The state-space description of the loop filter.
- `nlev`:   The number of levels in the quantizer.
- `xlim`:   A vector or scalar specifying the limit for each state variable.
- `ymax`:   The stability threshold. Inputs that yield quantizer inputs above `ymax`
            are considered to be beyond the stable range of the modulator.
- `umax`:   The maximum allowable input amplitude. `umax` is calculated if it
            is not supplied.
Outputs:
- `ABCDs`:  The state-space description of the scaled loop filter.
- `S`:      The diagonal scaling matrix, `ABCDs` = [`S`A`S`⁻¹ `S`B; C`S`⁻¹ D]; xs = `S`x;
"""
function scaleABCD(ABCD, nlev=2, f=0, xlim=1, ymax=NaN, umax=NaN, N_sim=100000, N0=10)
    if isnan(ymax)
        ymax = nlev + 5
    end
    order = size(ABCD, 1) - 1
    if size(xlim) == ()
        xlim = xlim * ones(order)   # Convert scalar xlim to a vector
    end

    # Make this function repeatable
    rng = MersenneTwister(0)

    # Envelope for smooth start-up
    raised_cosine = 0.5 * (1 .- cos.(pi/N0 * (0:N0-1)))
    if isnan(umax)
        # Simulate the modulator with DC or sine wave inputs to detect its
        # stable input range.
        # First get a rough estimate of umax.
        ulist = (0.1:0.1:1.0) * (nlev-1)
        umax = nlev - 1
        N = 1000
        u0 = [exp.(2im*pi*f*(-N0:-1)) .* raised_cosine; exp.(2im*pi*f*(0:N-1))]
        u0 = u0 .+ 0.01 * randn(rng, N+N0, 2) * [1, 1im]
        u0 = real(u0)

        for u = ulist
            _, _, _, y = simulateDSM(u*u0, ABCD, nlev)
            if maximum(abs.(y)) > ymax
                umax = u    # umax is the smallest input found which causes 'instability'
                break
            end
        end

        if umax == ulist[1]
            error(@sprintf("Modulator is unstable even with an input amplitude of %.1f.", umax))
        end
    end

    # More detailed simulation
    N = N_sim
    u0 = [exp.(2im*pi*f*(-N0:-1)) .* raised_cosine; exp.(2im*pi*f*(0:N-1))]
    u0 = u0 .+ 0.01 * randn(rng, N+N0, 2) * [1, 1im]
    u0 = real(u0)
    maxima = zeros(order) .- 1
    ulist = range(0.7*umax, umax, length=10)
    for u = ulist
        _, _, xmax, y = simulateDSM(u*u0, ABCD, nlev)
        if maximum(abs.(y)) > ymax
            break
        end
        # We need to do this at least once.
        umax = u    # umax is the largest input which keeps |y| < ymax
        maxima, = findmax([maxima xmax], dims=2)
    end

    # Scale the modulator so that all states are at most xlim.
    scale = maxima[:] ./ xlim
    S = diagm(1 ./ scale)
    Sinv = diagm(scale)
    A, B, C, D = partitionABCD(ABCD)
    ABCDs = [S*A*Sinv S*B; C*Sinv D]

    return ABCDs, umax, S
end


export msec_to_sample
export sec_to_sample
export sample_to_sec
export sample_to_msec
export sample_to_hr

msec_to_sample(msec:: Real, fs::Real, to_int::Bool = true) =
    to_int ? convert(Integer, msec*fs/1000) : msec*fs/1000
sec_to_sample(sec:: Real, fs::Real, to_int::Bool = true) =
    to_int ? convert(Integer, sec*fs) : sec*fs
sample_to_sec(sample:: Real, fs::Real) = sample/fs
sample_to_msec(sample:: Real, fs::Real) = sample/fs * 1000
hr_to_sample(hr::Real, fs::Real, to_int::Bool = true) = sec_to_sample(60/hr, fs, to_int)
sample_to_hr(sample::Real, fs::Real) = 60/sample_to_sec(sample, fs)

msec_to_sample(500.0, 500)


using Flux, MNIST
using Flux: onehotbatch, argmax, mse, throttle
using Base.Iterators: repeated

x, y = traindata()
y = onehotbatch(y, 0:9)

m = Chain(
    Dense(28^2, 32, σ),
    Dense(32, 10),
    softmax
)

using CLArrays
using GPUArrays: JLArray
using CLArrays: GlobalArray, GlobalPointer, PreDeviceArray
using Flux: OneHotMatrix

cl(x) = x
cl(x::JLArray) = x
cl(xs::AbstractArray) = isbits(xs) ? xs : JLArray(xs)
cl(xs::Flux.TrackedArray) = Flux.TrackedArray(xs.f, cl(xs.data), Base.RefValue(cl(Flux.Tracker.grad(xs))))
cl(xs::Flux.OneHotMatrix) = Flux.OneHotMatrix(cl(xs.data))

CLArrays.kernel_convert(x::OneHotMatrix{T}) where T <: CLArray = OneHotMatrix(CLArrays.kernel_convert(x.data))
CLArrays.predevice_type(::Type{OneHotMatrix{T}}) where T <: GlobalArray = OneHotMatrix{CLArrays.predevice_type(T)}
CLArrays.device_type(x::OneHotMatrix{T}) where T <: CLArray = OneHotMatrix{CLArrays.device_type(x.data)}
CLArrays.reconstruct(x::OneHotMatrix{T}, ptr::GlobalPointer) where T <: PreDeviceArray = OneHotMatrix(CLArrays.reconstruct(x.data, ptr))
GPUArrays.arg_length(x::OneHotMatrix{T}) where T <: GPUArrays.GPUArray = UInt32.(size(x))
GPUArrays.to_device(state, x::OneHotMatrix{<: JLArray}) = OneHotMatrix(GPUArrays.to_device(state, x.data))
x, y = cl(x), cl(y)
m = mapparams(cl, m)
loss(x, y) = mse(m(x), y)

dataset = repeated((x, y), 500)
evalcb = () -> @show(loss(x, y))
opt = SGD(params(m), 1)

Flux.train!(loss, dataset, opt, cb = throttle(evalcb, 10))

# Check the prediction for the first digit
argmax(m(x[:,1]), 0:9) == argmax(y[:,1], 0:9)


export RegGrid, gridtopology

"Specification for M by N grid with node degree D"
struct RegGrid{M,N,D}
    width::Float64
end

"Generate nodes of a regular grid"
function gridnodes(grid::RegGrid{M,N,D}) where {M,N,D}
    w = grid.width
    [Node(w*(j - 1), w*(M - i)) for j=1:N for i=1:M]
end

"2d array filled with indices to corresponding 1d array"
nodeidx(m,n) = reshape(1:m*n, m,n)

"generate coordinates for grid neighbors"
function candneighbors(i::Int, j::Int, grid::RegGrid{M,N,4}) where {M,N}
    [          (i-1, j),
     (i, j-1),           (i, j+1),
               (i+1, j)          ]
end
function candneighbors(i::Int, j::Int, grid::RegGrid{M,N,8}) where {M,N}
    [(i-1, j-1), (i-1, j), (i-1, j+1),
     (i, j-1),             (i, j+1),
     (i+1, j-1), (i+1, j), (i+1, j+1)]
end

"valid grid neighbors"
function gridneighbors(i::Int, j::Int, grid::RegGrid{M,N,D}) where {M,N,D}
    filter(t -> t[1] in 1:M && t[2] in 1:N,
           candneighbors(i, j, grid))
end

"Generate (anti-parallel) arcs of a regular grid"
function gridarcs(grid::RegGrid{M,N,D}) where {M,N,D}
    arcs = Arc[]
    idx = nodeidx(M,N)
    for j=1:N
        for i=1:M
            for n in gridneighbors(i, j, grid)
                push!(arcs, Arc(idx[i,j], idx[n...]))
            end
        end
    end
    arcs
end

function gridtopology(grid::RegGrid{M,N,D}) where {M,N,D}
    Topology(gridnodes(grid), gridarcs(grid))
end


"""
    conv_encoder(xsize, zsize, kernelsizes, channels, scalings[; activation, densedims])

Constructs a convolutional encoder.

# Arguments
- `xsize`: size of input - (h,w,c)
- `zsize`: latent space dimension
- `kernelsizes`: kernelsizes for conv. layers (only odd numbers are allowed)
- `channels`: channel numbers
- `scalings`: scalings vector
- `activation`: default relu
- `densedims`: if set, more than one dense layers are used

# Example
```julia-repl
julia> encoder = conv_encoder((64, 48, 1), 2, (3, 5, 5), (2, 4, 8), (2, 2, 2), densedims = (256))
Chain(Conv((3, 3), 1=>2, relu), MaxPool((2, 2), pad = (0, 0, 0, 0), stride = (2, 2)), Conv((5, 5), 2=>4, relu), MaxPool((2, 2), pad = (0, 0, 0, 0), stride = (2, 2)), Conv((5, 5), 4=>8, relu), MaxPool((2, 2), pad = (0, 0, 0, 0), stride = (2, 2)), #44, Dense(384, 2))

julia> encoder(randn(Float32, 64, 48, 1, 2))
2×2 Array{Float32,2}:
  0.247844   0.133781
 -0.605763  -0.494911
```
"""
function conv_encoder(xsize::Union{Tuple, Vector}, zsize::Int, kernelsizes::Union{Tuple, Vector},
    channels::Union{Tuple, Vector}, scalings::Union{Tuple, Vector};
    activation = relu, densedims::Union{Tuple, Vector} = [])
    nconv = length(kernelsizes)
    (nconv == length(channels) == length(scalings)) ? nothing : error("incompatible input dimensions")
    (length(xsize) == 3) ? nothing : error("xsize must be (h, w, c)")
    # also check that kernelsizes are all odd numbers
    (all(kernelsizes .% 2 .== 1)) ? nothing : error("not implemented for even kernelsizes")

    # initialize some stuff
    cins = vcat(xsize[3], channels[1:end-1]...) # channels in
    couts = channels # channels out
    ho = xsize[1]/(reduce(*, scalings)) # height before reshaping
    wo = xsize[2]/(reduce(*, scalings)) # width before reshaping
    (ho == floor(Int, ho)) ? ho = floor(Int, ho) : error("your input size and scaling is not compatible")
    (wo == floor(Int, wo)) ? wo = floor(Int, wo) : error("your input size and scaling is not compatible")
    din = ho*wo*channels[end]

    # now build a vector of layers to be used later
    layers = Array{Any,1}()
    # first add the convolutional and maxpooling layers
    for (k, ci, co, s) in zip(kernelsizes, cins, couts, scalings)
        pad = Int((k-1)/2)
        # paddding so that input and output size are same
        push!(layers, Conv((k,k), ci=>co, activation; pad = (pad, pad)))
        push!(layers, MaxPool((s,s)))
    end

    # reshape
    push!(layers, x -> reshape(x, din, :))

    # and dense layers
    ndense = length(densedims)
    dins = vcat(din, densedims...)
    douts = vcat(densedims..., zsize)
    dacts = vcat([activation for _ in 1:ndense]..., identity)
    for (_di, _do, _da) in zip(dins, douts, dacts)
        push!(layers, Dense(_di, _do, _da))
    end

    Flux.Chain(layers...)
end

"""
    conv_decoder(xsize, zsize, kernelsizes, channels, scalings[; activation, densedims, vec_output,
        vec_output_dim])

Constructs a convolutional encoder.

# Arguments
- `xsize`: size of input - (h,w,c)
- `zsize`: latent space dimension
- `kernelsizes`: kernelsizes for conv. layers (only odd numbers are allowed)
- `channels`: channel numbers
- `scalings`: scalings vector
- `activation`: default relu
- `densedims`: if set, more than one dense layers are used
- `vec_output`: output is vectorized (default false)
- `vec_output_dim`: determine what the final size of the vectorized input should be (e.g. add extra dimensions for variance estimation)

# Example
```julia-repl
julia> decoder = conv_decoder((64, 48, 1), 2, (5, 5, 3), (8, 4, 2), (2, 2, 2); densedims = (256))
Chain(Dense(2, 256, relu), Dense(256, 384, relu), #19, ConvTranspose((2, 2), 8=>8, relu), Conv((5, 5), 8=>4, relu), ConvTranspose((2, 2), 4=>4, relu), Conv((5, 5), 4=>2, relu), ConvTranspose((2, 2), 2=>2, relu), Conv((3, 3), 2=>1))

julia> y = decoder(randn(Float32, 2, 2));

julia> size(y)
(64, 48, 1, 2)
```
"""
function conv_decoder(xsize::Union{Tuple, Vector}, zsize::Int, kernelsizes::Union{Tuple, Vector},
    channels::Union{Tuple, Vector}, scalings::Union{Tuple, Vector};
    activation = relu, densedims::Union{Tuple, Vector} = [],
    vec_output = false, vec_output_dim = nothing)
    nconv = length(kernelsizes)
    (nconv == length(channels) == length(scalings)) ? nothing : error("incompatible input dimensions")
    (length(xsize) == 3) ? nothing : error("xsize must be (h, w, c)")
    # also check that kernelsizes are all odd numbers
    (all(kernelsizes .% 2 .== 1)) ? nothing : error("not implemented for even kernelsizes")

    # initialize some stuff
    cins = channels # channels in
    couts = vcat(channels[2:end]..., xsize[3]) # channels out
    ho = xsize[1]/(reduce(*, scalings)) # height after reshaping
    wo = xsize[2]/(reduce(*, scalings)) # width after reshaping
    (ho == floor(Int, ho)) ? ho = floor(Int, ho) : error("your input size and scaling is not compatible")
    (wo == floor(Int, wo)) ? wo = floor(Int, wo) : error("your input size and scaling is not compatible")
    dout = ho*wo*channels[1]

    # now build a vector of layers to be used later
    layers = Array{Any,1}()

    # start with dense layers
    ndense = length(densedims)
    dins = vcat(zsize, densedims...)
    douts = vcat(densedims..., dout)
    for (_di, _do) in zip(dins, douts)
        push!(layers, Dense(_di, _do, activation))
    end

    # reshape
    push!(layers, x -> reshape(x, ho, wo, channels[1], :))

    # add the transpose and convolutional layers
    acts = vcat([activation for _ in 1:nconv-1]..., identity)
    for (k, ci, co, s, act) in zip(kernelsizes, cins, couts, scalings, acts)
        pad = Int((k-1)/2)
        # use convtranspose for upscaling - there are other posibilities, however this seems to be ok
        push!(layers, ConvTranspose((s,s), ci=>ci, activation, stride = (s,s)))
        push!(layers, Conv((k,k), ci=>co, act; pad = (pad, pad)))
    end

    # if you want vectorized output
    if vec_output
        push!(layers, x -> reshape(x, reduce(*, xsize), :))
        # if you want the final vector to be a different size than the output
        # (e.g. you want extra dimensions for variance estimation)
        (vec_output_dim == nothing) ? nothing : push!(layers, Dense(reduce(*, xsize), vec_output_dim))
    end

    Flux.Chain(layers...)
end


using Optim
using PyPlot
using StatsBase
using ForwardDiff
using celerite
include("regress.jl")

function test_co2_v03()

function log_like_derivative_wrapper(p0,p,x,t,y)
  n = length(t)
  nex = (4(p-p0)+2p0+1)*(n-1)+1
  m1 = 2(p-p0)+p0+2
  if p0 == p
    m1 = p0+ 1
  end
  width = 2m1+1
  indx::Vector{Int64} = collect(1:nex)
#  aex = zeros(Number,width,nex)
#  al_small= zeros(Number,m1,nex)
#  bex = zeros(Number,nex)

  function log_like(x::Vector)
    w0 = exp(x[1])
    alpha = exp(x[2:p+1])
    beta_real = exp(x[p+2:2p+1])
    beta_imag = zeros(eltype(x),p)
    if p > p0
      beta_imag[p0+1:p] = exp(x[2p+2:3p-p0+1])
    end
    aex = zeros(eltype(x),width,nex)
    al_small= zeros(eltype(x),m1,nex)
    bex = zeros(eltype(x),nex)
    logdeta = compile_matrix_symm(alpha,beta_real,beta_imag,w0,t,nex,aex,al_small,indx)
    log_like = compute_likelihood(p,p0,y,aex,al_small,indx,logdeta,bex)
#    println("alpha: ",alpha," beta_real: ",beta_real," white: ",w0," log_like: ",log_like)
    return -log_like
  end

## Optimize the likelihood:
  xdiff = copy(x)
## Now perturb x:
#  dlogdx = zeros(eltype(x),length(x))
#  log_like0=log_like(x)
#  println(log_like0,x)
#  log_like0=log_like(x)
#  x0 = copy(x)
#  for i=1:length(x)
#    x    = copy(x0)
#    if (x[i] == 0.) then
#      h = 1e-6
#    else
#      h = 1e-6*x[i]
#    end
#    x[i] = x[i]+h
#    println(i," ",x,size(x))
#    log_like_h = log_like(x)
#    dlogdx[i]=(log_like_h-log_like0)/h
#  end
#  result=ForwardDiff.gradient(log_like,xdiff)
#  result=ForwardDiff.gradient(log_like,xdiff)
#  println("Gradient AutoDiff:    ",result)
#  println("Gradient Finite diff: ",dlogdx)

  lower = log([1e-15, 1e-8,1e-8,1e-8,1e-8,1e-8, 1e-8,1e-8,1e-8,1e-8,1e-8, 1e-3,1e-3,1e-3,1e-3,1e-3])
  upper = log([1e4, 1e8,1e8,1e8,1e8,1e8, 1e4,1e4,1e4,1e4,1e4, 1e3,1e3,1e3,1e3,1e3])
#  result = optimize(log_like, xdiff, BFGS(), OptimizationOptions(autodiff = true))
#  result = optimize(DifferentiableFunction(log_like), xdiff, lower, upper, Fminbox(), optimizer=GradientDescent)
#  result = optimize(DifferentiableFunction(log_like), xdiff, lower, upper, Fminbox(), optimizer=NelderMead)
  result = optimize(DifferentiableFunction(log_like), xdiff, lower, upper, Fminbox(), optimizer=BFGS, optimizer_o= OptimizationOptions(autodiff = true))
#  result = optimize(DifferentiableFunction(log_like), xdiff, lower, upper, Fminbox(), optimizer=ConjugateGradient)
#  result = optimize(log_like, xdiff, BFGS())
#  result = optimize(log_like, x0, BFGS())
#  result = optimize(log_like, xdiff)
  println(result)
  return exp(Optim.minimizer(result)),Optim.minimum(result)
#  return result
end

# Read in CO2 data:

data = readdlm("data/CO2_data.csv.txt",',')
co2 = vec(data[:,2])
println("Length of data: ",length(co2))
time = vec(data[:,1] - data[1,1])
# Carry out quadratic regression:
nt = length(time)
ord = 2
fn = zeros(ord+1,nt)
tnorm = (time-time[1])/(time[nt]-time[1])
for i=1:ord+1
  fn[i,:]=tnorm.^(i-1)
end
sig = ones(nt)
coeff,cov = regress(fn,co2,sig)
println("Coeff: ",coeff)
co2_trend = zeros(nt)
co2_sub = copy(co2)
for i=1:nt
  for j=1:ord+1
    co2_trend[i] += coeff[j]*fn[j,i]
  end
  co2_sub[i] -= co2_trend[i]
end

plot(time,co2)
plot(time,co2_trend)
plot(time,co2_sub)
#read(STDIN,Char)

# Compute ACF:
nt = length(time)
dt = median(time[2:nt]-time[1:nt-1])
lags = collect(0:1:200)
println(typeof(lags))
acf = autocov(co2_sub,lags)
lags *=dt
clf()
# Now, model covariance:
acf_model = 3.5.*exp(-lags./100.).*cos(2.*pi.*lags)+0.5*cos(2.*pi*lags/0.5)
clf()
plot(lags,acf)
plot(lags,acf_model)
plot(lags,acf-acf_model)
#read(STDIN,Char)

# Initialize model with guessed parameters:

nperiod = 1
#period = logspace(log10(5.0),log10(5.0),nperiod)
#period = logspace(-1.0,2.0,nperiod)
period = logspace(0.0,0.0,nperiod)
#period = logspace(log10(1.1),log10(1.1),nperiod)
# Best fit from initial run with 3 kernels:
#[7.50113753448555e-14,0.052634997384683246,3.814476305265215,0.2990524739656785,0.43730650302365087,19.982885984516415,0.00012381766899862475,0.0001249074549875943,0.2638033900492329,6.285942400966896,12.5689911272956,0.2470381285466161]

nkernel=5
xopt = Float64[]
log_like_opt = Inf

# Best fit with 4 kernels:
# Results of Optimization Algorithm
#  * Algorithm: Fminbox with BFGS
#  * Starting Point: [-30.221136621610558,-2.944374030966639, ...]
#  * Minimizer: [-3.300945139395281,-0.8237995678563096, ...]
#  * Minimum: -4.816914e+02
#  * Iterations: 4
#  * Convergence: true
#    * |x - x'| < 1.0e-32: false
#    * |f(x) - f(x')| / |f(x)| < 1.0e-32: true
#    * |g(x)| < 1.0e-08: false
#    * Reached Maximum Number of Iterations: false
#  * Objective Function Calls: 636
#  * Gradient Calls: 636
# New best fit- initial period 3.3729337950720426 best params: [0.03684832413517181,0.4387613797192722,3.809611698603267,0.2993294592999687,0.0033183947487366045,0.36055491419151137,0.00013119968869186474,0.00015241021114397978,0.00010000000928728985,6.285932185183963,12.569069334961931,18.8434942480683] log like: -481.69138538434845

#xsave = log([7.50113753448555e-14,0.052634997384683246,3.814476305265215,0.2990524739656785,0.43730650302365087,19.982885984516415,0.00012381766899862475,0.0001249074549875943,0.2638033900492329,6.285942400966896,12.5689911272956,0.2470381285466161])
#xsave = log([7.50113753448555e-14,0.052634997384683246,3.814476305265215,0.2990524739656785,0.1,19.982885984516415,0.00012381766899862475,0.0001249074549875943,0.2638033900492329,6.285942400966896,12.5689911272956,0.2470381285466161])
#xsave = log([0.03684832413517181,0.4387613797192722,3.809611698603267,0.2993294592999687,0.0033183947487366045,0.5,0.36055491419151137,0.00013119968869186474,0.00015241021114397978,0.00010000000928728985,1e-3,6.285932185183963,12.569069334961931,18.8434942480683,2pi])
#xsave = log([0.03684832413517181, 0.4387613797192722,3.809611698603267, 0.36055491419151137,0.00013119968869186474, 6.285932185183963])
#New best fit- initial params: [0.24525060661563458,-0.6706280702467335,0.6135393449157531,0.5158110704105079,-0.19071204291502974,-0.8683608405275405,-0.4454421228919916,0.19497110540072837,-0.10251527601637767,-0.9098131451360563,-0.011106363961333798,-0.9510402952632626] best params: [0.04029466597429114,0.2526358743611703,0.2986196855835003,3.8223408614702126,0.19617575450045882,0.6614637313826826,0.0001429291061867777,0.00012651638000879924,1.2683737146520694e-6,12.569019942371668,6.285907554790972,0.21506018780384212] log like: 160.69228348034142
#New best fit- initial params: [-0.7235566627564012,0.6194689646767091,0.029787409423436273,0.06264056996187195,0.7721099787691372,-0.031346813435856546,0.9944816013289715,-0.9923234906574838,0.6680014200732902,0.05233606985509143,0.9914124992763558,0.8947403433841048,0.7972837525544665,-0.10271485834584615] best params: [2.6711800255620015e-6,0.43519427651205367,0.050770915734427255,0.29775630072792336,3.8047644172968242,0.2930869808770129,20.768233217196343,0.0001266517006615938,0.00012184198184990059,0.010015841478464087,12.568981432423138,6.2859299962988935,0.9999971834821783,1.0000011746976747] log like: 160.73525576118857
# New best fit- initial params: [0.061136562954420715,-0.05913570088861064,-0.6077859327932353,0.8006118980060903,0.05891929293106468,0.34715208605827863,0.17937202290847631,0.09639794885504882,-0.34528679613275415,0.5769921986616793,0.337034498127855,-0.5627620564531655,0.016106597035654868,-0.7390684897388655] best params: [3.163633940791601e-6,0.44345590215659636,3.8160116462380707,0.29897108013090246,0.04226615839014087,0.32474949169199874,0.0001351318599722061,0.00012927971290721362,10.842253788156388,6.28592852289604,12.56897388114301,17.48063508199799,0.9999990035091292,0.999998751695153] log like: 158.09364527030715  CO2_spectrum_v04.png
#New best fit- initial params: [0.19456431327663282,0.17777008451335163,0.8403990148156821,0.6839614249344197,-0.07929762858610623,0.7793808090360561,0.7073480572067927,0.14315807344378895,0.33877630747091825,0.4072966264675575,-0.6676937819100126,-0.21481808817726478] best params: [2.7957389459161393e-6,0.5766912572663334,3.8223871455915246,0.04749499292515523,0.2993418814033055,0.20558714343221735,0.00016680586922843604,22.476650723282553,0.0001476106178401725,6.2869384732932,0.010009637397587567,87.96822286229002] log like: 104.73308577578865  (Fewer data points to match Foreman-Mackey's model CO2_spectrum_v05.png)
#New best fit- initial params: [-0.8685768119401973,-0.2132623980155186,-0.07868415414696539,0.7476202152461235,-0.342931132412406,0.6347252166656303,0.8488001935291196,-0.2975105273303984,-0.6143468418867672,0.09299778736926534,-0.06292736152275547,-0.5110956767023276] best params: [3.654198299839067e-7,0.5814441799891235,0.30029024967528783,3.830771738177122,0.04055130349945212,0.22506337451478275,0.00015594402269633877,0.00017916252815285494,11.521068977115174,12.570037522503661,6.286957191746393,17.736327369565533] log like: 103.22857603764857 (Ditto, v06.png)

xsave = -1.0+2.0*rand(15)
nx = length(xsave)
for iperiod=1:nperiod
#  alpha   = [0.0526,3.814,0.3,0.5]
#  beta_real = [20.,1e-4,1e-4,0.26]
#  alpha   = [0.1,0.1,0.1,0.1]
#  beta_real = [1e-2,1e-2,1e-2,1e-2]
#  beta_imag = 2pi./[1.0,0.5,period[iperiod]]
  x0 = copy(xsave)
#  x0[nx]=log(2pi/period[iperiod])
  p = nkernel
  p0 = 1
  xbest,log_like_best = log_like_derivative_wrapper(p0,p,x0,time,co2_sub)
  if log_like_best < log_like_opt
    println("New best fit- initial params: ",x0," best params: ",xbest," log like: ",log_like_best)
    log_like_opt = log_like_best
    xopt = xbest
# Plot the optimized ACF:
    acf_model = xbest[2].*exp(-lags.*xbest[nkernel+2])
    for j=2:nkernel
      acf_model += xbest[j+1]*exp(-lags.*xbest[nkernel+j+1]).*cos(lags.*xbest[2*nkernel+j])
    end
    acf_model[1] += xbest[1]
    clf()
    fig, axes = subplots(2,1)
    ax = axes[1,1]
    ax[:plot](lags,acf_model,"b-")
    ax[:plot](lags,acf,"r-")
    ax[:plot](lags,acf-acf_model,"g-")
    ax[:set_xlabel]("Time [yr]")
    ax[:set_ylabel]("Covariance [(ppm CO2)^2]")
#    read(STDIN,Char)
# Plot the power spectrum:
    nfreq = 1000
    omega = collect(logspace(log10(0.001),log10(1000.),nfreq))
    psd = zeros(nfreq)
    for i=1:nfreq
      psd[i] = 2.0*xbest[2]*xbest[nkernel+2]/(omega[i]^2+xbest[nkernel+2]^2)
      for j=2:nkernel
        psd[i] += 2.0*xbest[j+1]*xbest[nkernel+j+1]/((omega[i]-xbest[2*nkernel+j])^2+xbest[nkernel+j+1]^2)
      end
    end
#    clf()
    ax = axes[2,1]
    ax[:loglog](omega,psd,"k-")
    ax[:set_xlabel]("Frequency/omega [rad/yr]")
    ax[:set_ylabel]("Power")
    read(STDIN,Char)
  end
end
return
end


using GNSSTools
using FFTW
using AllanDeviations
using Statistics
# using LsqFit
# using NNLS
using PyPlot
pygui(true)

# function psd_at_f(h_parms, f, v_0)
#     h₋₂, h₋₁, h₀, h₁, h₂ = h_parms
#     return v_0^2 * (h₋₂*f^(-4) + h₋₁*f^(-3) + h₀*f^(-2) + h₁*f^(-1) + h₂*f^0)
# end

# function fit_h_parms_to_signal_PSD(psd_dBc, f_c=10e6,
#                                    freqs=[1,10,100,1e3,1e4,1e5])
#     N = length(freqs)
#     @assert N == length(psd_dBc)
#     b = zeros(N-1)
#     A = zeros(N-1)
#     for i in 1:N-1
#         j = N - i
#         A[i] = (psd_dBc[j] + psd_dBc[j+1])/2 + 10*log10(freqs[j+1] - freqs[j])
#         if i == 1
#             b[i] = 2*10^(A[i]/10) / (freqs[j+1] - freqs[j])  # b₀
#         elseif i == 2
#             b[i] = 2*10^(A[i]/10) / log(freqs[j+1] - freqs[j])  # b₋₁
#         else
#             b[i] = 2*10^(A[i]/10) / ((2-i)*(freqs[j+1]^(2-i) - freqs[j]^(2-i)))
#         end
#     end
#     # order is h2, h1, h0, h-1, h-2
#     h = b ./ f_c.^2
#     # reverse order to h-2, h-1, h₀, h₁, h₂
#     return reverse(h)
# end

function calc_Sy(x, f_s)
    N = length(x)
    X = fft(x)[1:floor(Int, N/2)+1]
    X = (1/(f_s*N)) .* abs2.(X)
    X[2:end-1] .= 2 .* X[2:end-1]
    X_log10 = 10 .* log10.(X)
    frequencies = Array(range(0, f_s/2, length=length(X)))
    return (X, X_log10, frequencies)
end

# f_s = 36e6
f_carrier1 = 16.3676e6
f_carrier2 = 10e6
f_s = 1*f_carrier1
t_length = 1
N = floor(Int, t_length*f_s)
N_over_2 = floor(Int, N/2)
t = Array(range(0, t_length, length=N))
frequencies = Array(range(0, f_s/2, length=N_over_2+1))
# carrier1 = cis.(2π.*f_carrier1.*t)
# carrier2 = cis.(2π.*f_carrier2.*t)
phase_noise = zeros(Complex{Float64}, N)
# signal = zeros(Complex{Float64}, N)
dBc_freq = [1., 1e1, 1e2, 1e3, 1e4]
X = zeros(Complex{Float64}, N_over_2+1)
x = zeros(Complex{Float64}, N)
rakon_it5300B = Dict("h_parms"=>[1.9e-6, 2e-6, 9.7e-8, 9.1e-11, 9.9e-15],
                     "dBc"=>[-57, -88, -112, -130, -140],
                     "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
                     "name"=>"TCXO, Rakon IT5300B",
                     "f_0"=>f_carrier1)
agilent_e4424b = Dict("h_parms"=>[2.8e-9, 1.8e-11, 1.2e-10, 1e-11, 1.8e-15],
                      "dBc"=>[-88, -118, -128, -135, -147],
                      "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
                      "name"=>"OCXO, Agilent E4424B",
                      "f_0"=>10e6)
isotemp_91_1 = Dict("h_parms"=>[2.2e-9, 1.6e-10, 1.6e-10, 7.1e-13, 7.2e-16],
                    "dBc"=>[-89, -120, -140, -151, -154],
                    "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
                    "name"=>"OCXO, ISOTEMP 91-1",
                    "f_0"=>f_carrier2)
oscilloquartz_8607 = Dict("h_parms"=>[3.9e-14, 8.7e-13, 2.2e-14, 3.2e-13, 6.4e-15],
                          "dBc"=>[-122, -137, -143, -145, -145],
                          "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
                          "name"=>"OCXO, Oscilloquartz 8607",
                          "f_0"=>f_carrier2)
srs_prs10 = Dict("h_parms"=>[5.3e-11, 4.7e-11, 1.6e-15, 9.6e-14, 9.9e-16],
                 "dBc"=>[-103, -135, -150, -152, -153],
                 "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
                 "name"=>"Rubidum, SRS PRS10",
                 "f_0"=>f_carrier2)
cesium = Dict("h_parms"=>[1.5e-12, 6.2e-11, 7.6e-23, 1.4e-14, 6.2e-19],
              "dBc"=>[-105, -135, -160, -170, -190],
              "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
              "name"=>"Cesium, Typical Values",
              "f_0"=>f_carrier2)
mti_250l = Dict("h_parms"=>[1.7e-9, 1e-10, 1.6e-10, 3.5e-13, 1.8e-16],
              "dBc"=>[-90, -120, -140, -155, -160],
              "dBc_freq"=>[1., 1e1, 1e2, 1e3, 1e4],
              "name"=>"OCXO, MTI 250L",#, Low g-sensitivity",
              "f_0"=>f_carrier2)

oscillators = [rakon_it5300B, agilent_e4424b, isotemp_91_1, oscilloquartz_8607,
               srs_prs10, cesium, mti_250l]
# figsize = (6.5, 4)
colorz = ["k", "b", "g", "r", "y", "m", "c", "gray"]
fig = figure(figsize=(6.5,4))
ax1 = fig.add_subplot(1,1,1)
legend_markers = []
marker_names = []
for i in 1:length(oscillators)
# for i in 1:1
    color = colorz[i]
    oscillator = oscillators[i]
    if oscillator["name"] == "TCXO, Rakon IT5300B"
        # carrier = carrier1
        f_carrier = f_carrier1
    else
        # carrier = carrier2
        f_carrier = f_carrier2
    end
    h_parms = oscillator["h_parms"] ./ oscillator["f_0"]^2
    generate_phase_noise!(phase_noise, t_length, f_carrier, h_parms)
    # ϕ = real.(phase_noise)
    # signal = carrier .* cis.(ϕ)
    # signal .= carrier .* cis.(real.(phase_noise))
    # x = signal
    x .= real.(phase_noise) .+ 0im
    # X, X_log, freqs = calc_Sy(x, f_s)
    fft!(x);
    X .= abs2.(x[1:N_over_2+1]) / (f_s*N);
    X[2:end-1] .= 2 .* X[2:end-1];
    X .= 10 .* log10.(X);
    # X_max = maximum(real.(X));
    # X = X .- X_max;
    # ax1.plot(frequencies[2:end] .- f_carrier, X[2:end], color,
    #          label=string(oscillator["name"], " from h Parms"))
    ax1.plot(frequencies[2:end], real.(X[2:end]), color)#),
            #  label=oscillator["name"])
    # ax1.plot(frequencies[2:end], real.(X[2:end]), color,
            #  label=string(oscillator["name"], " from h Parms"))
    ax1.scatter(dBc_freq, oscillator["dBc"].+(10*log10(2)),
                c=color)#, label=string(oscillator["name"], " Truth"))
    legend_marker = plt.Line2D((0,1),(0,0), color=color, marker="o", linestyle="-")
    push!(legend_markers, legend_marker)
    push!(marker_names, oscillator["name"])
end
xscale("log")
xlim([1, 1e6])
xlabel("Frequency (Hz)")
ylabel("dBc/Hz")
subplots_adjust(top=0.96, left=0.12, right=0.88, bottom=0.13)
ax1.legend(legend_markers, marker_names)

# savefig("figures/ch2_oscillator_psds.pdf", dpi=300)


using Images, FileIO, Printf

absdiff(a::RGB{T}, b::RGB{T}) where T = sum(abs(col(a) - col(b)) for col in (red, green, blue))
function pctdiff(A::Matrix{Color{T}}, B::Matrix{Color{T}}) where T
    size(A) != size(B) && throw(ArgumentError("images must be same-size"))
    s = zero(T)
    for (a, b) in zip(A, B)
        s += absdiff(a, b)
    end
    return 100s / 3prod(size(A))
end

img50  = load("data/lenna50.jpg") |> Matrix{RGB{Float64}};
img100 = load("data/lenna100.jpg") |> Matrix{RGB{Float64}};

d = pctdiff(img50, img100)
@printf("Percentage difference: %.4f%%\n", d)


function sum_diff(x)
n = length(x); d = 1/(n-1)
s = zero(eltype(x))
s = s + (x[2] - x[1]) / d
for i = 2:length(x)-1
s = s + (x[i+1] - x[i+1]) / (2*d)
end
s = s + (x[n] - x[n-1])/d
end
function sum_diff_fast(x)
n=length(x); d = 1/(n-1)
s = zero(eltype(x))
@fastmath s = s + (x[2] - x[1]) / d
@fastmath for i = 2:n-1
s = s + (x[i+1] - x[i+1]) / (2*d)
end
@fastmath s = s + (x[n] - x[n-1])/d
end


# # Orthocube Examples (Bravais Lattice)

# ## Preamble
using TightBindingLattice
using Plots
mkpath("example_orthocube")

function draw_orthocube(orthocube::OrthoCube, coordinates::AbstractMatrix{<:Integer})
    xlim = (minimum(coordinates[1,:]) - 3.5, maximum(coordinates[1,:]) + 3.5)
    ylim = (minimum(coordinates[2,:]) - 3.5, maximum(coordinates[2,:]) + 3.5)
    boundary = let
        r0 = [0,0]
        r1 = orthocube.shape_matrix[:,1]
        r2 = orthocube.shape_matrix[:,2]
        hcat(r0, r1, r1 .+ r2, r2, r0)
    end
    fig = plot(boundary[1,:], boundary[2,:], label="", size=(400, 400))
    for Ri in Iterators.product(-1:1, -1:1)
        alpha = (Ri == (0,0)) ? 1.0 : 0.1
        Rr = orthocube.shape_matrix * [Ri...]
        scatter!(coordinates[1,:] .+ Rr[1], coordinates[2,:] .+ Rr[2], aspect_ratio=1, markersize=12, xlim=xlim, ylim=ylim, markeralpha=alpha, label="")
        annotation = []
        for i in 1:size(coordinates, 2)
            x = coordinates[1,i] + Rr[1]
            y = coordinates[2,i] + Rr[2]
            if xlim[1] < x < xlim[2] && ylim[1] < y < ylim[2]
                push!(annotation, (x, y, text("$i", 8, :black, :center)))
            end
        end
        if !isempty(annotation)
            annotation = [annotation...]
            annotate!(annotation)
        end
    end
    fig
end

# ## (4,-4) x (4,4)
size_matrix = [ 4 4; -4 4]
orthocube = OrthoCube(size_matrix)
generator_translations = find_generators(orthocube)
coordinates = generate_coordinates(orthocube, generator_translations)
coordmat = hcat(coordinates...)
println("All elements")
for (i, c) in enumerate(coordinates)
    println("$i : $c")
end
println("Generator translations")
for (it, t) in enumerate(eachcol(generator_translations))
    println("t($it) = $t")
end
draw_orthocube(orthocube, coordmat)
savefig("example_orthocube/bravais_(4,-4)x(4,4).svg")

# ![](example_orthocube//bravais_(4,-4)x(4,4).svg)

# ## (2,-2) x (2,4)
size_matrix = [ 2 2; -2 4]
orthocube = OrthoCube(size_matrix)
generator_translations = find_generators(orthocube)
coordinates = generate_coordinates(orthocube, generator_translations)
coordmat = hcat(coordinates...)
draw_orthocube(orthocube, coordmat)
println("All elements")
for (i, c) in enumerate(coordinates)
    println("$i : $c")
end
println("Generator translations")
for (it, t) in enumerate(eachcol(generator_translations))
    println("t($it) = $t")
end
draw_orthocube(orthocube, coordmat)
savefig("example_orthocube/bravais_(2,-2)x(2,4).svg")

# ![](example_orthocube//bravais_(2,-2)x(2,4).svg)


using Benchmark
################################################################################
# Benchmark to see whether loop is faster than matrix operation
#
# | Row | Function     | Average   | Relative | Replications |
# |-----|--------------|-----------|----------|--------------|
# | 1   | "relu_loop"  | 0.0283602 | 1.0      | 30           |
# | 2   | "relu_whole" | 0.0611887 | 2.15755  | 30           |
#
# | Row | Function         | Average   | Relative | Replications |
# |-----|------------------|-----------|----------|--------------|
# | 1   | "relu_bwd_loop"  | 0.0314793 | 1.0      | 30           |
# | 2   | "relu_bwd_whole" | 0.0696191 | 2.21158  | 30           |
################################################################################

function relu_loop(input)
  input = copy(input)
  for i = 1:length(input)
    input[i] = max(input[i], 0)
  end
  return input
end

function relu_whole(input)
  input = copy(input)
  input = max(input, 0)
  return input
end

function relu_bwd_loop(input, grad)
  grad = copy(grad)
  for i = 1:length(input)
    if input[i] <= 0
      grad[i] = 0
    end
  end
  return grad
end
function relu_bwd_whole(input, grad)
  grad = copy(grad)
  grad .*= (input .> 0)
  return grad
end

input = rand(28, 28, 50, 128)
i1 = relu_loop(input)
i2 = relu_whole(input)
@assert all(abs(i1-i2) .< 1e-10)

grad = rand(size(input))

relu_loop() = relu_loop(input)
relu_whole() = relu_whole(input)
relu_bwd_loop() = relu_bwd_loop(i1, grad)
relu_bwd_whole() = relu_bwd_whole(i1, grad)

df = compare([relu_loop, relu_whole], 30)
println("$df")

df = compare([relu_bwd_loop, relu_bwd_whole], 30)
println("$df")


myParse(x) = parse.(Int, x)
fishes = myParse.(split.(readline("input.txt"), ','))

function simulateFishes!(fishVec, nDays)
    bucket = [0, 0, 0, 0, 0, 0, 0, 0, 0]

    for fish in fishVec
        bucket[fish+1] += 1
    end

    for _ in 1:nDays
        bucket = circshift(bucket, -1)
        bucket[6+1] += bucket[8+1] # For each fish that generates a new fish (bucket tracking 8), add count to the 6 counter as the fish resets
    end
    return sum(bucket)
end

#! Part 1

simulateFishes!(fishes, 80)

#! Part 2

simulateFishes!(fishes, 256)


#
# Factors Reduce
#

"""
    reducedim(op, ϕ::Factor, dims, [v0])

Reduces a dimension in `ϕ` using function `op`.
See `Base.reducedim` for more details
"""
Base.reducedim(op, ϕ::Factor, dim::Symbol, v0=nothing) = reducedim(op, ϕ, [dim], v0)

# a more efficient version than `reducedim!(deepcopy(ϕ))` ?
function Base.reducedim(op, ϕ::Factor, dims::Vector{Symbol}, v0=nothing)
    _check_dims_valid(dims, ϕ)
    # needs to be a tuple for squeeze
    inds = (sort!(indexin(dims, ϕ))...)

    isempty(inds) && return deepcopy(ϕ)

    dims_new = deepcopy(ϕ.dimensions)
    deleteat!(dims_new, inds)
    pot_new = _red_dim(op, ϕ, inds, v0)

    return Factor(dims_new, pot_new)
end

# reduce the dimension and then squeeze them out
_red_dim(op, ϕ::Factor, inds::Tuple, ::Void) = squeeze(reducedim(op, ϕ.potential, inds), inds)
_red_dim(op, ϕ::Factor, inds::Tuple, v0) = squeeze(reducedim(op, ϕ.potential, inds, v0), inds)

"""
    Z(ϕ)

Return the partition function (sum of all elements).
"""
Z(ϕ) = sum(ϕ)

Base.sum(ϕ::Factor, dims=names(ϕ)) = reducedim(+, ϕ, dims)
Base.prod(ϕ::Factor, dims=names(ϕ)) = reducedim(*, ϕ, dims)
Base.maximum(ϕ::Factor, dims=names(ϕ)) = reducedim(max, ϕ, dims)
Base.minimum(ϕ::Factor, dims=names(ϕ)) = reducedim(min, ϕ, dims)

"""
    marginalize(ϕ, dims)

Sum out all but the specified dimensions
"""
marginalize(ϕ::Factor, dim::Symbol) = marginalize(ϕ, [dim])
marginalize(ϕ::Factor, dims::AbstractVector{Symbol}) = sum(ϕ, setdiff(names(ϕ), dims))


"""
Make a grid with the given size, for which the given domain is the covering domain.
"""
function discretize_togrid end

discretize_togrid(d::AbstractInterval, n::Int) = EquispacedGrid(n, d)
discretize_togrid(d::Interval{:closed,:open}, n::Int) = PeriodicEquispacedGrid(n, d)
discretize_togrid(d::Interval{:open,:open}, n::Int) = MidpointEquispacedGrid(n, d)
discretize_togrid(d::UnitInterval{T}, n::Int) where {T} = UnitEquispacedGrid{T}(n)
discretize_togrid(d::ChebyshevInterval{T}, n::Int) where {T} = ChebyshevNodes{T}(n)

discretize_togrid(d::ProductDomain, dims) = productgrid(map(discretize_togrid, components(d), dims)...)

discretize_togrid(d::MappedDomain, dims) =
    MappedGrid(forward_map(d), discretize_togrid(superdomain(d), dims))

discretize_togrid(d::UnitCircle{T}, n::Int) where {T} =
    MappedGrid(mapfrom_parameterdomain(d), UnitPeriodicEquispacedGrid{T}(n))