# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Desription
# ==============================================================================
#
#   Tests related to conversion from Euler angles to direction cosine matrices.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File: ./src/conversions/angle_to_dcm.jl
# =======================================

# Functions: angle_to_dcm
# -----------------------

@testset "Angle about an axis => DCM (Float64)" begin
    # Sample an angle to use in all rotations.
    ang = _rand_ang()

    # Rotations
    # =========

    # X Axis
    # ------

    # Create a DCM that rotates about X axis.
    dcm = angle_to_dcm(ang, :X)
    @test eltype(dcm) === Float64

    # Create a vector that does not have X component.
    v = SVector(0, randn(), randn())

    # Rotate the reference using the DCM.
    v_r = dcm * v
    @test eltype(v_r) === Float64

    # Get the sine of the angle between the representations.
    sin_ang = ((v_r × v) / norm(v)^2)[1]

    # Test the angle between the two representations.
    @test abs(sin(ang) - sin_ang) ≈ 0 atol = 1e-12

    # Y Axis
    # ------

    # Create a DCM that rotates about Y axis.
    dcm = angle_to_dcm(ang, :Y)
    @test eltype(dcm) === Float64

    # Create a vector that does not have Y component.
    v = SVector(randn(), 0, randn())

    # Rotate the reference using the DCM.
    v_r = dcm * v
    @test eltype(v_r) === Float64

    # Get the sine of the angle between the representations.
    sin_ang = ((v_r × v) / norm(v)^2)[2]

    # Test the angle between the two representations.
    @test abs(sin(ang) - sin_ang) ≈ 0 atol = 1e-12

    # Z Axis
    # ------

    # Create a DCM that rotates about Z axis.
    dcm = angle_to_dcm(ang, :Z)
    @test eltype(dcm) === Float64

    # Create a vector that does not have Z component.
    v = SVector(randn(), randn(), 0)

    # Rotate the reference using the DCM.
    v_r = dcm * v
    @test eltype(v_r) === Float64

    # Get the sine of the angle between the representations.
    sin_ang = ((v_r × v) / norm(v)^2)[3]

    # Test the angle between the two representations.
    @test abs(sin(ang) - sin_ang) ≈ 0 atol = 1e-12
end

@testset "Angle about an axis => DCM (Float32)" begin
    # Sample an angle to use in all rotations.
    ang = _rand_ang(Float32)

    # Rotations
    # =========

    # X Axis
    # ------

    # Create a DCM that rotates about X axis.
    dcm = angle_to_dcm(ang, :X)
    @test eltype(dcm) === Float32

    # Create a vector that does not have X component.
    v = SVector{3, Float32}(0, randn(), randn())

    # Rotate the reference using the DCM.
    v_r = dcm * v
    @test eltype(v_r) === Float32

    # Get the sine of the angle between the representations.
    sin_ang = ((v_r × v) / norm(v)^2)[1]

    # Test the angle between the two representations.
    @test abs(sin(ang) - sin_ang) ≈ 0 atol = 1e-6

    # Y Axis
    # ------

    # Create a DCM that rotates about Y axis.
    dcm = angle_to_dcm(ang, :Y)
    @test eltype(dcm) === Float32

    # Create a vector that does not have Y component.
    v = SVector{3, Float32}(randn(), 0, randn())

    # Rotate the reference using the DCM.
    v_r = dcm * v
    @test eltype(v_r) === Float32

    # Get the sine of the angle between the representations.
    sin_ang = ((v_r × v) / norm(v)^2)[2]

    # Test the angle between the two representations.
    @test abs(sin(ang) - sin_ang) ≈ 0 atol = 1e-6

    # Z Axis
    # ------

    # Create a DCM that rotates about Z axis.
    dcm = angle_to_dcm(ang, :Z)
    @test eltype(dcm) === Float32

    # Create a vector that does not have Z component.
    v = SVector{3, Float32}(randn(), randn(), 0)

    # Rotate the reference using the DCM.
    v_r = dcm * v
    @test eltype(v_r) === Float32

    # Get the sine of the angle between the representations.
    sin_ang = ((v_r × v) / norm(v)^2)[3]

    # Test the angle between the two representations.
    @test abs(sin(ang) - sin_ang) ≈ 0 atol = 1e-6
end

@testset "Angle about an axis => DCM (Errors)" begin
    @test_throws ArgumentError angle_to_dcm(0, :A)
end

@testset "Euler angles => DCM (Float64)" begin
    T = Float64

    # Two rotations
    # ==========================================================================

    for rot_seq in valid_rot_seqs_2angles
        # Sample Euler angles.
        θ₁ = _rand_ang(T)
        θ₂ = _rand_ang(T)

        # Convert to DCM.
        D = angle_to_dcm(θ₁, θ₂, rot_seq)
        @test eltype(D) === T

        # Create the DCM by composing the rotations.
        rot_seq_str = string(rot_seq)
        D₁ = angle_to_dcm(θ₁, Symbol(rot_seq_str[1]))
        D₂ = angle_to_dcm(θ₂, Symbol(rot_seq_str[2]))
        De = D₂ * D₁

        # Compare.
        @test D ≈ De
    end

    # Three rotations
    # ==========================================================================

    for rot_seq in valid_rot_seqs
        # Sample Euler angles.
        ea = EulerAngles(_rand_ang(T), _rand_ang(T), _rand_ang(T), rot_seq)

        # Convert to DCM.
        D = angle_to_dcm(ea)
        @test eltype(D) === T

        # Create the DCM by composing the rotations.
        rot_seq_str = string(rot_seq)
        D₁ = angle_to_dcm(ea.a1, Symbol(rot_seq_str[1]))
        D₂ = angle_to_dcm(ea.a2, Symbol(rot_seq_str[2]))
        D₃ = angle_to_dcm(ea.a3, Symbol(rot_seq_str[3]))
        De = D₃ * D₂ * D₁

        # Compare.
        @test D ≈ De
    end
end

@testset "Euler angles => DCM (Float32)" begin
    T = Float32

    # Two rotations
    # ==========================================================================

    for rot_seq in valid_rot_seqs_2angles
        # Sample Euler angles.
        θ₁ = _rand_ang(T)
        θ₂ = _rand_ang(T)

        # Convert to DCM.
        D = angle_to_dcm(θ₁, θ₂, rot_seq)
        @test eltype(D) === T

        # Create the DCM by composing the rotations.
        rot_seq_str = string(rot_seq)
        D₁ = angle_to_dcm(θ₁, Symbol(rot_seq_str[1]))
        D₂ = angle_to_dcm(θ₂, Symbol(rot_seq_str[2]))
        De = D₂ * D₁

        # Compare.
        @test D ≈ De
    end

    # Three rotations
    # ==========================================================================

    for rot_seq in valid_rot_seqs
        # Sample Euler angles.
        ea = EulerAngles(_rand_ang(T), _rand_ang(T), _rand_ang(T), rot_seq)

        # Convert to DCM.
        D = angle_to_dcm(ea)
        @test eltype(D) === T

        # Create the DCM by composing the rotations.
        rot_seq_str = string(rot_seq)
        D₁ = angle_to_dcm(ea.a1, Symbol(rot_seq_str[1]))
        D₂ = angle_to_dcm(ea.a2, Symbol(rot_seq_str[2]))
        D₃ = angle_to_dcm(ea.a3, Symbol(rot_seq_str[3]))
        De = D₃ * D₂ * D₁

        # Compare.
        @test D ≈ De
    end
end

@testset "Euler angles => DCM (Promotion)" begin
    # Check if promotion is working as intended.
    dcm = angle_to_dcm(Int64(1), 0.0f0, Float64(0))
    @test eltype(dcm) === Float64

    dcm = angle_to_dcm(Int64(1), 0.0f0, 0.0f0)
    @test eltype(dcm) === Float32

    dcm = angle_to_dcm(Int64(1), Int32(0), Float16(0))
    @test eltype(dcm) === Float16
end

@testset "Euler angles => DCM (Errors)" begin
    @test_throws ArgumentError angle_to_dcm(1, :V)
    @test_throws ArgumentError angle_to_dcm(1, 2, :ZZ)
    @test_throws ArgumentError angle_to_dcm(1, 2, 3, :ZZY)
end

# Functions: smallangle_to_dcm
# ----------------------------

@testset "Small Euler angles => DCM (Float64)" begin
    T = Float64

    # Sample three small angles.
    θx = _rand_ang(T) * T(0.001)
    θy = _rand_ang(T) * T(0.001)
    θz = _rand_ang(T) * T(0.001)

    # Create the rotation matrix without normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = false)
    @test eltype(D) === T

    # Expected result.
    De = DCM(
          1, +θz, -θy,
        -θz,   1, +θx,
        +θy, -θx,   1
    )'

    @test De ≈ D

    # Create the rotation matrix with normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = true)
    @test eltype(D) === T

    @test orthonormalize(De) ≈ D
end

@testset "Small Euler angles => DCM (Float32)" begin
    T = Float32

    # Sample three small angles.
    θx = _rand_ang(T) * T(0.001)
    θy = _rand_ang(T) * T(0.001)
    θz = _rand_ang(T) * T(0.001)

    # Create the rotation matrix without normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = false)
    @test eltype(D) === T

    # Expected result.
    De = DCM(
          1, +θz, -θy,
        -θz,   1, +θx,
        +θy, -θx,   1
    )'

    @test De ≈ D

    # Create the rotation matrix with normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = true)
    @test eltype(D) === T

    @test orthonormalize(De) ≈ D
end

@testset "Small Euler angles => DCM (Promotion)" begin
    # Check if promotion is working as intended.
    dcm = smallangle_to_dcm(Int64(1), 0.0f0, Float64(0))
    @test eltype(dcm) === Float64

    dcm = smallangle_to_dcm(Int64(1), 0.0f0, 0.0f0)
    @test eltype(dcm) === Float32

    dcm = smallangle_to_dcm(Int64(1), Int32(0), Float16(0))
    @test eltype(dcm) === Float16
end


import Base: rand,
             ∈,
             isempty

export Universe

"""
    Universe{N<:Real} <: AbstractPolyhedron{N}

Type that represents the universal set, i.e., the set of all elements.
"""
struct Universe{N<:Real} <: AbstractPolyhedron{N}
    dim::Int
end

isoperationtype(::Type{<:Universe}) = false

# default constructor of type Float64
Universe(dim::Int) = Universe{Float64}(dim)

# --- AbstractPolyhedron interface functions ---

"""
    constraints_list(U::Universe{N}) where {N<:Real}

Return the list of constraints defining a universe.

### Input

- `U` -- universe

### Output

The empty list of constraints, as the universe is unconstrained.
"""
function constraints_list(U::Universe{N}) where {N<:Real}
    return LinearConstraint{N, Vector{N}}[]
end

"""
    constrained_dimensions(U::Universe)::Vector{Int}

Return the indices in which a universe is constrained.

### Input

- `U` -- universe

### Output

The empty vector, as the universe is unconstrained in every dimension.
"""
function constrained_dimensions(U::Universe)::Vector{Int}
    return Int[]
end

# --- LazySet interface functions ---

"""
    dim(U::Universe)

Return the dimension of a universe.

### Input

- `U` -- universe

### Output

The dimension of a universe.
"""
function dim(U::Universe)::Int
    return U.dim
end

"""
    ρ(d::AbstractVector{N}, U::Universe{N}) where {N<:Real}

Return the support function of a universe.

### Input

- `d` -- direction
- `U` -- universe

### Output

The support function in the given direction.

### Algorithm

If the direction is all zero, the result is zero.
Otherwise, the result is `Inf`.
"""
function ρ(d::AbstractVector{N}, U::Universe{N}) where {N<:Real}
    return iszero(d) ? zero(N) : N(Inf)
end

"""
    σ(d::AbstractVector{N}, U::Universe{N}) where {N<:Real}

Return the support vector of a universe.

### Input

- `d` -- direction
- `U` -- universe

### Output

A vector with infinity values, except in dimensions where the direction is zero.
"""
function σ(d::AbstractVector{N}, U::Universe{N}) where {N<:Real}
    return [v == zero(N) ? v : v > zero(N) ? N(Inf) : N(-Inf) for v in d]
end

"""
    ∈(x::AbstractVector{N}, U::Universe{N})::Bool where {N<:Real}

Check whether a given point is contained in a universe.

### Input

- `x` -- point/vector
- `U` -- universe

### Output

The output is always `true`.

### Examples

```jldoctest
julia> [1.0, 0.0] ∈ Universe(2)
true
```
"""
function ∈(x::AbstractVector{N}, U::Universe{N})::Bool where {N<:Real}
    @assert length(x) == dim(U)
    return true
end

"""
    an_element(U::Universe{N}) where {N<:Real}

Return some element of a universe.

### Input

- `U` -- universe

### Output

The origin.
"""
function an_element(U::Universe{N}) where {N<:Real}
    return zeros(N, dim(U))
end

"""
    rand(::Type{Universe}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing
        )::Universe{N}

Create a universe (note that there is nothing to randomize).

### Input

- `Universe` -- type for dispatch
- `N`        -- (optional, default: `Float64`) numeric type
- `dim`      -- (optional, default: 2) dimension
- `rng`      -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`     -- (optional, default: `nothing`) seed for reseeding

### Output

The (only) universe of the given numeric type and dimension.
"""
function rand(::Type{Universe};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing
             )::Universe{N}
    rng = reseed(rng, seed)
    return Universe{N}(dim)
end

"""
    isempty(U::Universe)::Bool

Return if a universe is empty or not.

### Input

- `U` -- universe

### Output

`false`.
"""
function isempty(U::Universe)::Bool
    return false
end

"""
    isbounded(U::Universe)::Bool

Determine whether a universe is bounded.

### Input

- `U` -- universe

### Output

`false` as the universe is unbounded.
"""
function isbounded(U::Universe)::Bool
    return false
end

"""
    isuniversal(U::Universe{N}, [witness]::Bool=false
               )::Union{Bool, Tuple{Bool, Vector{N}}} where {N<:Real}

Check whether a universe is universal.

### Input

- `U`       -- universe
- `witness` -- (optional, default: `false`) compute a witness if activated

### Output

* If `witness` option is deactivated: `true`
* If `witness` option is activated: `(true, [])`
"""
function isuniversal(U::Universe{N}, witness::Bool=false
                    )::Union{Bool, Tuple{Bool, Vector{N}}} where {N<:Real}
    return witness ? (true, N[]) : true
end

"""
    norm(U::Universe, [p]::Real=Inf)

Return the norm of a universe.
It is the norm of the enclosing ball (of the given ``p``-norm) of minimal volume
that is centered in the origin.

### Input

- `U` -- universe
- `p` -- (optional, default: `Inf`) norm

### Output

An error.
"""
function norm(U::Universe, p::Real=Inf)
    error("a universe does not have a norm")
end

"""
    radius(U::Universe, [p]::Real=Inf)

Return the radius of a universe.
It is the radius of the enclosing ball (of the given ``p``-norm) of minimal
volume with the same center.

### Input

- `U` -- universe
- `p` -- (optional, default: `Inf`) norm

### Output

An error.
"""
function radius(U::Universe, p::Real=Inf)
    error("a universe does not have a radius")
end

"""
    diameter(U::Universe, [p]::Real=Inf)

Return the diameter of a universe.
It is the maximum distance between any two elements of the set, or,
equivalently, the diameter of the enclosing ball (of the given ``p``-norm) of
minimal volume with the same center.

### Input

- `U` -- universe
- `p` -- (optional, default: `Inf`) norm

### Output

An error.
"""
function diameter(U::Universe, p::Real=Inf)
    error("a universe does not have a diameter")
end

"""
    translate(U::Universe{N}, v::AbstractVector{N}) where {N<:Real}

Translate (i.e., shift) a universe by a given vector.

### Input

- `U` -- universe
- `v` -- translation vector

### Output

The universe.
"""
function translate(U::Universe{N}, v::AbstractVector{N}) where {N<:Real}
    @assert length(v) == dim(U) "cannot translate a $(dim(U))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    return U
end


function distances(x,y)
    xij = @. x - x'
    yij = @. y - y'
    return xij,yij
end

function periodic_distances!(x,y,Lx,Ly)
    make_periodic!(x,Lx)
    make_periodic!(y,Ly)
end

function make_periodic!(x,L)
    @. x[x>L/2]  -= L
    @. x[x<-L/2] += L
end

"""
    rij = sparse_distances(x,y)

Sparse form of unique lower triangular distance matrix for vectors `x`, `y` specifying points on the plane.
"""
function sparse_distances(x,y)
    xij = @. x - x'
    yij = @. y - y'
    rij = hypot.(xij,yij) |> tril |> sparse
end

"""
    tree = spanning_tree(x,y)

Minimal spanning tree for points specified by vectors `x`, `y`.
"""
function spanning_tree(x,y)
    rij = sparse_distances(x,y)
    # graph representation
    g = SimpleWeightedGraph(length(rij))
    for edge in eachindex(rij)
        n1,n2 = edge[1],edge[2]
        add_edge!(g,n1,n2,rij[n1,n2])
    end
    # get the spanning tree
    tree = prim_mst(g)
end

"""
    vortex_marker(c)

Define vortex marker using sign.
"""
function vortex_marker(c)
    if c==1.0
        return (4,0.6,:green, stroke(1.3, 1.0, :white))
    elseif c==-1.0
        return (4,0.6,:blue, stroke(1.3, 1.0, :white))
    else
    end
end

"""
    plot_vortices!(p,vdata)

Plot vortices defined in `vdata` on figure `p`.
"""
function plot_vortices!(p,vdata)
    for j in 1:size(vdata)[1]
        vx,vy = vdata[j,1],vdata[j,2]
        scatter!(p1,[vx],[vy],marker=vortex_marker(vdata[j,3]),label=false)
    end
    return p1
end

"""
    plot_cluster_vortices!(p,cluster::Cluster)

Plot cluster in figure `p` using `vortex_marker`.
"""
function plot_cluster_vortices!(p,cluster::Cluster)
    vdata = vortex_array(cluster.vortices)
    for j in 1:size(vdata)[1]
        vx,vy = vdata[j,1],vdata[j,2]
        scatter!(p,[vx],[vy],marker=vortex_marker(vdata[j,3]),label=false)
    end
    return p
end

"""
    plot_cluster_tree!(p,cluster::Cluster)

Plot minimal spanning tree on figure `p` for `cluster`.
"""
function plot_cluster_tree!(p,cluster::Cluster)
    vortices = cluster.vortices
    tree = cluster.tree
    va = vortex_array(vortices)
    xi = @view va[:,1]
    yi = @view va[:,2]
    ci = @view va[:,3]

    for i in 1:length(tree)
        edge = tree[i]
        src,dst = edge.src,edge.dst
        x1,y1 = xi[src],yi[src]
        x2,y2 = xi[dst],yi[dst]
        csign = vortices[1].qv
        col = csign == 1.0 ? :green : :blue
        plot!(p,[x1,x2],[y1,y2],label=false,c=col,alpha=0.4)
    end
    return p
end

"""
    plot_custer!(p,c)

Plot the vortices and minimal spanning tree.
"""
function plot_cluster!(p,c)
    plot_cluster_vortices!(p,c)
    plot_cluster_tree!(p,c)
    return p
end


import Base: ==

struct LagrangeBasis{NF,T} <: AbstractBasis{1,NF}
    funcs::SP.PolynomialSystem{NF,1}
    points::Vector{T}
    function LagrangeBasis(funcs::Vector{DP.Polynomial{C,R}},
        points::Vector{T}) where {C,R,T}

        NF = length(funcs)
        order = NF - 1
        npoints = length(points)

        @assert NF > 0 "Vector of functions must be non-empty"
        @assert NF == npoints "Number of functions must be equal to number of interpolation points"
        @assert isapprox(sum(funcs),one(R)) "Polynomial basis must sum to unity"

        var = funcs[1].x.vars[1]
        for f in funcs
            @assert length(f.x.vars) == 1 "Require univariate polynomials"
            @assert var in f.x.vars "Require same variable in all polynomials"
            ord = maximum(maximum.(f.x.Z))
            @assert ord == order "Require all polynomials of same order"
        end
        polysystem = SP.PolynomialSystem(funcs)
        new{NF,T}(polysystem,points)
    end
end

function type_of_interpolation_points(B::LagrangeBasis{NF,T}) where {NF,T}
    return T
end

function Base.show(io::IO,basis::LagrangeBasis{NF,T}) where {NF,T}
    p = order(basis)
    print(io, "1-D LagrangeBasis\n\tOrder: $p")
end

function order(basis::LagrangeBasis{NF,T}) where {NF,T}
    return NF-1
end

function LagrangeBasis(order::Z;start = -1.0, stop = 1.0) where {Z<:Integer}
    @assert start < stop
    @assert order >= 0

    DP.@polyvar x
    NF = order+1
    points = NF == 1 ? [0.5*(start+stop)] : Array(range(start,stop=stop,length=NF))
    funcs = lagrange_polynomials(x,points)
    return LagrangeBasis(funcs,points)
end

function (B::LagrangeBasis)(x)
    return SP.evaluate(B.funcs,[x])
end

function derivative(B::LagrangeBasis{NF,T},x) where {NF,T}
    vals = SP.jacobian(B.funcs,[x])
    return vals
end

function gradient(B::LagrangeBasis{NF,T},x) where {NF,T}
    return derivative(B,x)
end

function hessian(B::LagrangeBasis{NF,T},x) where {NF,T}
    h = zeros(NF,1,1)
    SP.hessian!(h,B.funcs,[x])
    return h[:,1,1]
end


module VisualTests

using Gadfly, Interpolations
const nx = 10
xg = 1:nx
xf = -1.2nx:.1:3.3nx

f1(x) = sin(2pi/nx * (x-1))

y1 = map(f1,xg)

p = plot()

if true

for (IT, EB) in (
        (Constant{OnCell}, ExtrapConstant),
        (Constant{OnGrid}, ExtrapConstant),
        (Linear{OnCell}, ExtrapNaN),
        (Linear{OnGrid}, ExtrapNaN),
        (Quadratic{Flat,OnCell}, ExtrapConstant),
        (Quadratic{Flat,OnGrid}, ExtrapConstant),
        (Quadratic{Line,OnCell}, ExtrapLinear),
        (Quadratic{Line,OnGrid}, ExtrapLinear),
        (Quadratic{Line,OnGrid}, ExtrapReflect),
        (Quadratic{Flat,OnCell}, ExtrapReflect),
        (Quadratic{Free,OnCell}, ExtrapReflect),
        (Quadratic{Free,OnGrid}, ExtrapReflect),
        (Quadratic{Periodic,OnGrid}, ExtrapPeriodic),
        (Quadratic{Periodic,OnCell}, ExtrapPeriodic),
    )
    itp = Interpolation(y1, IT(), EB())

    stuff = Any[]

    push!(stuff, layer(x=xg,y=y1,Geom.point,Theme(default_color=color("green"))))
    push!(stuff, layer(x=xf,y=[itp[x] for x in xf],Geom.path,Theme(default_point_size=2px)))
    title = "$(IT.name.name){$(join(map(t -> t.name.name, IT.parameters), ','))}, $(EB.name.name)"
    push!(stuff, Guide.title(title))
    display(plot(stuff...))
end

nx = 10
xg = 1:nx
ny = 6
yg = 1:ny

f2(x,y) = sin(2pi/nx * (x-1)) * cos(2pi/ny * (y-1))
zg = Float64[f2(x,y) for x in xg, y in yg]

xf = -1:.1:nx+1
yf = -1:.1:ny+1

itp2 = Interpolation(zg, Quadratic(Flat(),OnCell()),ExtrapLinear())

display(plot(
    layer(x=xf,y=yf,z=[itp2[x,y] for x in xf, y in yf], Geom.contour),
    Guide.title("Quadratic{Flat,Oncell}, ExtrapLinear")
))

end

end


# ## Objectives
#
# Here we will show you how to perform an aerodynamic analysis of an airfoil.
# > **Recipe**
# > 1. Compute the coordinates of a NACA 4-digit series airfoil.
# > 2. Plot its camber, thickness, upper and lower surface representations.
# > 3. Perform an aerodynamic analysis at a given angle of attack.
# > 4. Plot its aerodynamic properties.
# For this, we will need to import some packages which will be convenient for plotting.
using AeroMDAO      # Main package
using Plots         # Plotting library
gr(dpi = 300)       # Plotting backend
using LaTeXStrings  # For LaTeX printing in plots

# ## Your First Airfoil
#
# You can define a NACA 4-digit series airfoil using the following function.
airfoil = naca4((2,4,1,2), 60)

# You can access the $x$- and $y$-coordinates as the following fields.
airfoil.x, airfoil.y

# You can obtain the coordinates as an array by calling the following function.
coordinates(airfoil)

# ### Geometric Representations
# You can convert these coordinates into the camber-thickness representation.
xcamthick = camber_thickness(airfoil, 60)

# You can split the coordinates into their upper and lower surfaces.
upper, lower = split_surface(airfoil);

x_upper, y_upper = @views upper[:,1], upper[:,2]
x_lower, y_lower = @views lower[:,1], lower[:,2];

# ### Plotting
# Let's plot everything!

## Plot object
af_plot = plot(aspect_ratio = 1, xlabel=L"(x/c)", ylabel = L"y")

## Upper surface
plot!(upper[:,1], upper[:,2], label = "$(airfoil.name) Upper",
      ls = :solid, lw = 2, c = :cornflowerblue)

## Lower surface
plot!(lower[:,1], lower[:,2], label = "$(airfoil.name) Lower",
      ls = :solid, lw = 2, c = :orange)

## Camber
plot!(xcamthick[:,1], xcamthick[:,2], label = "$(airfoil.name) Camber",
      ls = :dash, lw = 2, c = :burlywood3)

## Thickness
plot!(xcamthick[:,1], xcamthick[:,3], label = "$(airfoil.name) Thickness",
      ls = :dash, lw = 2, c = :grey)

# ## Your First Doublet-Source Analysis
#
# Now we have an airfoil, and we would like to analyze its aerodynamic characteristics. The potential flow panel method for inviscid analyses of airfoils, which you may have studied in your course in aerodynamics, provides decent estimations of the lift generated by the airfoil.

# Our analysis also requires boundary conditions, which is the freestream flow defined by a magnitude ``V_\infty`` and angle of attack ``\alpha``. We provide these to the analysis by defining variables and feeding them to a `Uniform2D` type, corresponding to uniform flow in 2 dimensions.
V       = 1.0
alpha   = 4.0 # degrees
uniform = Uniform2D(V, alpha)

# You can analyze this airfoil for the given freestream flow as shown.
system  = solve_case(
                     airfoil, uniform;
                     num_panels = 80
                    );

# This will run the analysis and return a system which can be used to obtain the aerodynamic quantities of interest and post-processing. The panels used for the analysis can be accessed as follows.
panels = system.surface_panels

# You can compute the lift coefficient from the system.
cl = lift_coefficient(system)

# You can also compute the sectional lift, moment and pressure coefficients.
cls, cms, cps = surface_coefficients(system)

# Note the difference between the lift coefficient computed and the sum of the sectional lift coefficients; this is due to numerical errors in the solution procedure and modeling.
cl, sum(cls)

#md # !!! info
#md #     Support for drag prediction with boundary layer calculations will be added soon. For now, try out the amazing [Webfoil](http://webfoil.engin.umich.edu/) developed by the MDOLab at University of Michigan -- Ann Arbor!

# ### Visualization
#
# Let's see what the pressure and lift distribution curves look like over the airfoil. AeroMDAO provides more helper functions for post-processing data. For example, you can make your fancy plots by segregating the values depending on the locations of the panels by defining the following function.
cp_upper, cp_lower = get_surface_values(panels, cps)

# Now let's plot the results!

## Pressure coefficients
plot(yflip = true, xlabel = L"(x/c)", ylabel = L"C_p", lw = 2)
plot!(cp_upper, label = L"$C_p$ Upper",
      ls = :dash, lw = 2, c = :cornflowerblue)
plot!(cp_lower, label = L"$C_p$ Lower",
      ls = :dash, lw = 2, c = :orange)
plot!(x_upper, -y_upper, label = "$(airfoil.name) Upper",
      ls = :solid, lw = 2, c = :cornflowerblue)
plot!(x_lower, -y_lower, label = "$(airfoil.name) Lower",
      ls = :solid, lw = 2, c = :orange)

#
## Lift coefficients
cl_upper, cl_lower = get_surface_values(panels, cls)

cl_plot = plot(xlabel = L"(x/c)", ylabel = L"C_l", lw = 2)
plot!(cl_upper, label = L"$C_l$ Upper",
      ls = :dash, lw = 2, c = :cornflowerblue)
plot!(cl_lower, label = L"$C_l$ Lower",
      ls = :dash, lw = 2, c = :orange)
plot!(x_lower, y_lower, label = "$(airfoil.name) Lower",
      ls = :solid, lw = 2, c = :orange)
plot!(x_upper, y_upper, label = "$(airfoil.name) Upper",
      ls = :solid, lw = 2, c = :cornflowerblue)

#
## Moment coefficients
cm_upper, cm_lower = get_surface_values(panels, cms)

cm_plot = plot(xlabel = L"(x/c)", ylabel = L"C_m")
plot!(cm_upper, label = L"$C_m$ Upper",
      ls = :dash, lw = 2, c = :cornflowerblue)
plot!(cm_lower, label = L"$C_m$ Lower",
      ls = :dash, lw = 2, c = :orange)
plot!(x_upper, y_upper, label = L"$C_m$ Upper",
      ls = :solid, lw = 2, c = :cornflowerblue)
plot!(x_lower, y_lower, label = L"$C_m$ Lower",
      ls = :solid, lw = 2, c = :orange)

# Great! We've created our first airfoil and run an aerodynamic analysis in 2 dimensions. Now we can move on to analyzing an aircraft configuration in the [Aircraft Aerodynamic Analysis](tutorials-aircraft.md) tutorial.


using HartreeFock: gaussian_product_center, norm2, boys, Vec3
using Test

@test gaussian_product_center(1.0, Vec3([0.0, 0.0, 0.0]), 1.5, Vec3([1.0, 1.0, 1.0])) == [0.6, 0.6, 0.6]
@test norm2(Vec3([1.0, 2.0, 3.0])) == 14.0
@test boys(0, 1.0) ≈ 0.7468241328


"""
    Plate(normal::AbstractVector{T}, width::T, [, origin::AbstractVector{T}=zeros()])

A plate defined by all the points ``\\mathbf x`` that satify ``|(\\mathbf x - \\mathbf o) \\cdot \\mathbf n| < w /2`` where ``\\mathbf n`` is the unit normal, ``\\mathbf o`` is the origin, and ``w`` is the width.
"""
struct Plate{T,Dim} <: Shape{T,Dim}
    normal::SVector{Dim,T} # a unit vector which is orthogonal to the plate
    width::T # the width
    origin::SVector{Dim,T}
end

# Constructors
Plate(normal::AbstractVector{T}, width::T, origin::AbstractVector{T} = zeros(T,length(normal))) where {T} = Plate{T,length(normal)}(normal ./ norm(normal), width, origin)

Plate(normal::NTuple{Dim,T}, width::T, origin::AbstractVector{T} = zeros(T,Dim)) where {T,Dim} = Plate{T,Dim}(normal ./ norm(normal), width, origin)

name(shape::Plate) = "Plate"

volume(shape::Plate) = Inf
outer_radius(hs::Plate{T}) where T = T(Inf)

# import Base.issubset
# function issubset(inner_rect::Plate{T}, outer_rect::Plate{T}) where T
#     all(topright(inner_rect) .<= topright(outer_rect)) &&
#     all(bottomleft(inner_rect) .>= bottomleft(outer_rect))
# end

import Base.in
function in(x::AbstractVector{T}, p::Plate)::Bool where T
    abs(dot(x - p.origin, p.normal)) < p.width / T(2)
end

import Base.(==)
function ==(h1::Plate{T}, h2::Plate{T}) where T
    h1.origin == h2.origin &&
    h1.width == h2.width &&
    h1.normal == h2.normal
end

import Base.isequal
function isequal(h1::Plate{T}, h2::Plate{T}) where T
    isequal(h1.origin, h2.origin) &&
    isequal(h1.width, h2.width) &&
    isequal(h1.normal, h2.normal)
end

function iscongruent(h1::Plate{T}, h2::Plate{T}) where T
    (h1.normal  == h2.normal) && (h1.width  == h2.width)
end

function congruent(h::Plate{T}, x) where T
    Plate(h.normal, h.width, x)
end

# function boundary_functions(h::Plate{T,2}, scale::T = 10.0) where T
#     surface_vec = [-h.normal[2], h.normal[1]]
#
#     function x(t)
#         check_boundary_coord_range(t)
#         h.origin[1] + surface_vec[1] * width/2 + t * scale * surface_vec[1]
#     end
#
#     function y(t)
#         check_boundary_coord_range(t)
#         h.origin[2] + t * scale * surface_vec[2]
#     end
#
#     return x, y
# end


#= functions for the power series exterior map from unit disk =#

struct PowerSeries{N,T}
  ccoeff :: Vector{T}
  dcoeff :: Vector{T}
end

function PowerSeries(ccoeff::Vector{T},dcoeff::Vector{T}) where T
  ncoeff = length(ccoeff)-2
  PowerSeries{ncoeff,T}(ccoeff,dcoeff)
end

function Base.show(io::IO,ps::PowerSeries{N,T}) where {N,T}
  println(io, "multipole coefficients:")
  println(io, "  c₁ = $(ps.ccoeff[1]), ")
  println(io, "  c₀ = $(ps.ccoeff[2]), ")
  print(io,   "  c₋ᵢ = ")
  for i = 1:N
    print(io,"$(ps.ccoeff[2+i]), ")
  end
  println(io, "i = 1:$(N)")
end

function (ps::PowerSeries)(ζ::T) where T<:Number
  ζⁿ = ζ
  z = zero(ζ)
  for c in ps.ccoeff
    z += c*ζⁿ
    ζⁿ /= ζ
  end
  return z
end

(ps::PowerSeries)(ζ::Vector{T}) where T<:Number = ps.(ζ)

struct PowerSeriesDerivatives
  ps :: PowerSeries
end

function (dps::PowerSeriesDerivatives)(ζ::T) where T<:Number
  C = dps.ps.ccoeff
  dz = C[1]
  ζⁿ = 1/ζ^2
  ddz = ComplexF64(0)
  for n in 1:length(C)-2
    dz -= n*C[n+2]*ζⁿ
    ζⁿ /= ζ
    ddz += n*(n+1)*C[n+2]*ζⁿ
  end
  return dz, ddz
end

function (dps::PowerSeriesDerivatives)(ζs::Vector{T}) where T<:Number
  dz = zero(ζs)
  ddz = zero(ζs)
  for (i,ζ) in enumerate(ζs)
    dz[i], ddz[i] = dps(ζ)
  end
  return dz, ddz
end

function evalinv_exterior(z::Vector{ComplexF64},ps::PowerSeries,
                                dps::PowerSeriesDerivatives)
#=
Evaluates the inverse of the exterior power series mapping, using a combination
of integration and Newton iteration.
=#

   ζ = zeros(ComplexF64,size(z))
   lenz = length(z)
   ζ0 = []
   maxiter = 10
   tol = 1.0e-8

   # Find z values close to vertices and set ζ to the corresponding
   # prevertices
   done = zeros(Bool,size(z))

   # Now, for remaining z values, first try to integrate
   #  dζ/dt = (z - z(ζ₀))/z'(ζ) from t = 0 to t = 1,
   # with the initial condition ζ(0) = ζ₀.
   if isempty(ζ0)
     # choose a point on the unit circle
     ζ0 = exp.(im*zeros(lenz))
     ζ0[isapprox.(angle.(z),π;atol=eps())] .= exp(im*π)
     dz0,ddz0 = dps(ζ0)
     # check for starting points on edges of the body, and rotate them
     # a bit if so
     onedge = isapprox.(abs.(dps(ζ0)[1]),0.0;atol=eps())
     ζ0[onedge] .*= exp(im*π/20)
     z0 = ps(ζ0)
   else
     z0 = ps(ζ0)
     if length(ζ0)==1 && lenz > 1
       ζ0 = repeat(transpose(ζ0), lenz)
       z0 = repeat(transpose(z0), lenz)
     end
     z0 = z0[.~done]
     ζ0 = ζ0[.~done]
   end
   odetol = max(tol,1e-3)
   scale = z[.~done] - z0

   ζ0 = [real(ζ0);imag(ζ0)]

   f(ζ,p,t) = invfunc(ζ,scale,dps)
   tspan = (0.0,1.0)
   prob = ODEProblem(f,ζ0,tspan)
   sol = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
   lenu = length(ζ0)
   ζ[.~done] = sol.u[end][1:lenz]+im*sol.u[end][lenz+1:lenu];
   out = abs.(ζ) .> 1
   ζ[out] = sign.(ζ[out])

   # Now use Newton iterations to improve the solution
   ζn = ζ
   k = 0
   while ~all(done) && k < maxiter
     F = z[.~done] - ps(ζn[.~done])
     M = length(F)
     dF, ddz = dps(ζn[.~done])
     ζn[.~done] = ζn[.~done] + F./dF

     done[.~done] = abs.(F).< tol
     k += 1
   end
   F = z[.~done] - ps(ζn[.~done])
   if any(abs.(F).> tol)
     error("Check solution")
   end
   ζ = ζn

end

function invfunc(u,scale,dps::PowerSeriesDerivatives)
    lenu = length(u)
    lenzp = Int(lenu/2)
    ζ = u[1:lenzp]+im*u[lenzp+1:lenu]

    dz, ddz = dps(ζ)
    f = scale./dz
    zdot = [real(f);imag(f)]
end


using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(file::String)::Int
    chars = Dict(Char(Int('A') + i) => i + 1 for i ∈ 0:25)
    return sum(i * chars[c] for (i, name) ∈ enumerate(sort(split(replace(file, "\"" => ""), ","))) for c ∈ name)
end

file = read("problems/0022/p022_names.txt", String)

compute(file)

@benchmark compute(file)


"""
Module ParallelRandomFields

Enables to sample spatial realisations of 2D and 3D Gaussian random fields with given power spectrum using following covariance functions:
- anisotropic exponential
- isotropic Gaussian

ParallelRandomFields can be deployed on both (multi-) GPUs and CPUs.

# General overview and examples
https://github.com/luraess/ParallelRandomFields.jl

Functions are defined in the submodules.

# Submodules
- [`ParallelRandomFields.grf2D_Threads`](@ref)
- [`ParallelRandomFields.grf3D_Threads`](@ref)
- [`ParallelRandomFields.grf2D_CUDA`](@ref)
- [`ParallelRandomFields.grf3D_CUDA`](@ref)

To see a description of a function type `?<functionname>`.
"""
module ParallelRandomFields

using ParallelStencil

# include submodules and reset ParallelStencil prior
include("grf2D_Threads.jl")

ParallelStencil.@reset_parallel_stencil()
include("grf3D_Threads.jl")

ParallelStencil.@reset_parallel_stencil()
include("grf2D_CUDA.jl")

ParallelStencil.@reset_parallel_stencil()
include("grf3D_CUDA.jl")

ParallelStencil.@reset_parallel_stencil()

end # Module ParallelRandomFields


using LinearAlgebra
using CompScienceMeshes
using SauterSchwab3D
using StaticArrays

const pI   = point(1,0,0)
const pII  = point(0,1,0)
const pIII = point(0,0,0)
const pIV  = point(0,0,1)
const qIV  = point(0,0,-1)

const P = simplex(pI,pII,pIV,pIII)
const Q = simplex(pI,pIII,pII)

const sing = SauterSchwab3D.singularity_detection(P,Q)

Accuracy2 = 18
cf_ref = CommonFace5D(sing,SauterSchwab3D._legendre(Accuracy2,0.0,1.0))

function integrand(x,y)
      return ((x-pI)'*(y-pIV))*exp(-im*1*norm(x-y))/(4pi*norm(x-y))
end

function INTEGRAND(u,v)
   j1 = volume(P) * factorial(dimension(P))
   j2 = volume(Q) * factorial(dimension(Q))
   x = barytocart(P,u)
   y = barytocart(Q,v)
   output = integrand(x,y)*j1*j2

   return output
end

print("Ref: ")
ref = sauterschwab_parameterized(INTEGRAND, cf_ref)
println(ref)
println()

res_tp =[]
res_sp =[]
res_gm =[]
n1 = []
n2 = []
n3 = []

for i in 2:1:15
   Accuracy = i
   cf = CommonFace5D(sing,SauterSchwab3D._legendre(Accuracy,0.0,1.0))

   int_tp = sauterschwab_parameterized(INTEGRAND, cf)

   num_pts = 9*length(cf.qps)^5
   push!(n1,num_pts)

   push!(res_tp,int_tp)
end

for i in 2:1:7
   Accuracy = i
   cf_s = CommonFace5D_S(sing,(SauterSchwab3D._legendre(Accuracy,0.0,1.0),
                          SauterSchwab3D._shunnham2D(Accuracy),
                          SauterSchwab3D._shunnham3D(Accuracy)))

   int_sp = sauterschwab_parameterized(INTEGRAND, cf_s)

   num_pts = length(cf_s.qps[1])*length(cf_s.qps[2])^2+8*length(cf_s.qps[2])*length(cf_s.qps[3])
   push!(n2,num_pts)
   push!(res_sp,int_sp)
end

for i in 2:1:12
   Accuracy = i
   cf_gm = CommonFace5D_S(sing,(SauterSchwab3D._legendre(Accuracy,0.0,1.0),
                          SauterSchwab3D._grundmannMoeller2D(2*Accuracy-1),
                          SauterSchwab3D._grundmannMoeller3D(2*Accuracy-1)))

   int_gm = sauterschwab_parameterized(INTEGRAND, cf_gm)

   num_pts = length(cf_gm.qps[1])*length(cf_gm.qps[2])^2+8*length(cf_gm.qps[2])*length(cf_gm.qps[3])
   push!(n3,num_pts)
   push!(res_gm,int_gm)
end

err_tp = norm.(res_tp.-ref)/norm(ref)
err_sp = norm.(res_sp.-ref)/norm(ref)
err_gm = norm.(res_gm.-ref)/norm(ref)

println(err_tp)
println(err_sp)
println(err_gm)

using Plots
plot( yaxis=:log, xaxis=:log, fontfamily="Times")
plot!(n1,err_tp, label="Gauss Tensor-Product",markershape=:circle)
plot!(n2,err_sp, label="Simplex Tensor-Product",markershape=:rect)
#plot!(n3,err_gm, label="Simplex-Product GM",markershape=:x)
plot!(xlims=(1e2,1e5),ylims=(1e-7,1))
plot!(xlabel="#Quad. pts/Func. evals", ylabel="Rel. Error.", title="Common Face 5D",legend=:bottomleft)

using BenchmarkTools

ref = -0.00463274359881397 + 0.002581200521272345im

cf = CommonFace5D(sing,SauterSchwab3D._legendre(4,0.0,1.0))

int_tp = sauterschwab_parameterized(INTEGRAND, cf)

num_pts = 9*length(cf.qps)^5
println("#Pts: ",num_pts)
err_tp = norm.(int_tp.-ref)/norm(ref)
println("Rel. Err: ",err_tp)
@time for i in 1:100 sauterschwab_parameterized(INTEGRAND, cf) end

cf_s = CommonFace5D_S(sing,(SauterSchwab3D._legendre(4,0.0,1.0),
                       SauterSchwab3D._shunnham2D(4),
                       SauterSchwab3D._shunnham3D(4)))

int_sp = sauterschwab_parameterized(INTEGRAND, cf_s)

num_pts = length(cf_s.qps[1])*length(cf_s.qps[2])^2+8*length(cf_s.qps[2])*length(cf_s.qps[3])

println("#Pts: ",num_pts)
err_sp = norm.(int_sp.-ref)/norm(ref)
println("Rel Err: ",err_sp)
@time  for i in 1:100  sauterschwab_parameterized(INTEGRAND, cf_s) end

#=
cf_gm = CommonFace5D_S((SauterSchwab3D._legendre(4,0.0,1.0),
                        SauterSchwab3D._grundmannMoeller2D(2*4-1),
                        SauterSchwab3D._grundmannMoeller3D(2*4-1)))

int_gm = sauterschwab_parameterized(INTEGRAND, cf_gm)

num_pts = length(cf_gm.qps[1])*length(cf_gm.qps[2])^2+8*length(cf_gm.qps[2])*length(cf_gm.qps[3])

println("#Pts: ",num_pts)
err_gm = norm.(int_gm.-ref)/norm(ref)
println("Rel Err: ",err_gm)
=#


using Random
using Plots
using ImageView
using TestImages
using ImageCore
using ProgressBars

function flatten(X)
    return reshape(X, (:, last(size(X))))
end

function euclidean_distance(A, B)
    return sum((A-B).^2)^0.5
end

function find_neighbours(p, neighbourhood, ϵ, distanceFunc)
    points = flatten(neighbourhood)
    neighbours::Vector{Int} = []
    for i in range(1, size(points)[1])
        q = points[i, :]
        distance = distanceFunc(p, q)
        if (0 < distance) & (distance <= ϵ)
            push!(neighbours, i)
        end
    end
    return neighbours
end

function dbscan(X, ϵ, minNeighbourhoodSize, distanceFunc)
    points = flatten(X)
    n_points = length(points[:,1])
    cluster::Int = 0
    neighbourhoods = Dict()
    labels = Dict()
    for i in tqdm(range(1, n_points))
        if ~haskey(labels, i)
            get!(neighbourhoods, i) do
                find_neighbours(points[i,:], points, ϵ, distanceFunc)
            end
            if length(neighbours_i) < minNeighbourhoodSize
                labels[i] = -1 # Noise
            else
                cluster+=1
                labels[i] = cluster

                local_neighbourhood = copy(neighbours_i)
                while ~isempty(local_neighbourhood)
                    j = pop!(local_neighbourhood)
                    if ~haskey(labels, j) || (labels[j] == -1)
                        labels[j] = cluster
                        neighbours_j = find_neighbours(points[j,:], points, ϵ, distanceFunc) # TODO store as a dict
                        if length(neighbours_j) >= minNeighbourhoodSize
                            append!(local_neighbourhood, neighbours_j)
                        end

                    end
                end
            end
        end
    end
    label_map = Array{Int}(undef, n_points, 1)
    for (k, v) in labels
        label_map[k] = v
    end
    label_map = reshape(label_map, size(X)[1:end-1])
    return label_map
end

img = testimage("lighthouse")

data = permutedims(channelview(img), [2,3,1]);
data = convert(Array{Float64, size(data)[end]}, data)

mask = dbscan(data, 0.1, 1000, euclidean_distance)

imshow(mask)


# !INPUT/OUTPUT PARAMETERS:
#   sol  : speed of light in atomic units (in,real)
#   kpa  : quantum number kappa (in,integer)
#   e    : energy (in,real)
#   nr   : number of radial mesh points (in,integer)
#   r    : radial mesh (in,real(nr))
#   vr   : potential on radial mesh (in,real(nr))
#   nn   : number of nodes (out,integer)
#   g0   : m th energy derivative of the major component multiplied by r
#          (out,real(nr))
#   g1   : radial derivative of g0 (out,real(nr))
#   f0   : m th energy derivative of the minor component multiplied by r
#          (out,real(nr))
#   f1   : radial derivative of f0 (out,real(nr))
function rdiracint!(sol, kpa, e, nr, r, vr, g0, g1, f0, f1)

    # rescaling limit
    rsc = 1.0e100

    @assert nr >= 4

    # inverse speed of light
    ci = 1.0/sol

    # electron rest energy
    e0 = sol^2
    t1 = 2.0*e0 + e

    # determine the r -> 0 boundary values of F and G
    t2 = kpa/r[1]
    t3 = ci*(t1 - vr[1])
    t4 = ci*(vr[1] - e)
    f0[1] = 1.0
    f1[1] = 0.0
    g0[1] = (f1[1] - t2*f0[1])/t4
    g1[1] = t3*f0[1] - t2*g0[1]

    # extrapolate to the first four points
    g1[2:4] .= g1[1]
    f1[2:4] .= f1[1]

    nn = 0 # number of nodes
    for ir in 2:nr
        t2 = kpa/r[ir]
        t3 = ci*(t1 - vr[ir])
        t4 = ci*(vr[ir] - e)
        ir0 = ir - 3
        if ir0 < 1
            ir0 = 1
        end
        @views g1[ir] = poly3(r[ir0:ir0+2], g1[ir0:ir0+2], r[ir])
        @views f1[ir] = poly3(r[ir0:ir0+2], f1[ir0:ir0+2], r[ir])
        # integrate to find wavefunction
        @views g0[ir] = poly4i( r[ir0:ir0+3], g1[ir0:ir0+3], r[ir]) + g0[ir0]
        @views f0[ir] = poly4i( r[ir0:ir0+3], f1[ir0:ir0+3], r[ir]) + f0[ir0]
        # compute the derivatives
        g1[ir] = t3*f0[ir] - t2*g0[ir]
        f1[ir] = t4*g0[ir] + t2*f0[ir]
        # integrate for correction
        @views g0[ir] = poly4i(r[ir0:ir0+3], g1[ir0:ir0+3], r[ir]) + g0[ir0]
        @views f0[ir] = poly4i(r[ir0:ir0+3], f1[ir0:ir0+3], r[ir]) + f0[ir0]
        # compute the derivatives again
        g1[ir] = t3*f0[ir] - t2*g0[ir]
        f1[ir] = t4*g0[ir] + t2*f0[ir]
        # check for overflow
        if ( (abs(g0[ir]) > rsc) || (abs(g1[ir]) > rsc) ||
             (abs(f0[ir]) > rsc) || (abs(f1[ir]) > rsc) )
            # set the remaining points and return
            g0[ir:nr] .= g0[ir]
            g1[ir:nr] .= g1[ir]
            f0[ir:nr] .= f0[ir]
            f1[ir:nr] .= f1[ir]
            return nn, e
        end
        # check for node
        if g0[ir-1]*g0[ir] < 0.0
            nn = nn + 1
        end
    end
    return nn, e
end

function poly3(xa, ya, x)
    # arguments
    # real(8) xa(3),ya(3),x
    # local variables

    # evaluate the polynomial coefficients
    x0 = xa[1]
    x1 = xa[2] - x0
    x2 = xa[3] - x0
    y0 = ya[1]
    y1 = ya[2] - y0
    y2 = ya[3] - y0
    t0 = 1.0/(x1*x2*(x2-x1))
    t1 = x1*y2
    t2 = x2*y1
    c1 = x2*t2 - x1*t1
    c2 = t1 - t2
    t1 = x - x0

    # evaluate the polynomial
    return y0 + t0*t1*(c1 + c2*t1)
end

function poly4i(xa, ya, x)
    # arguments
    #real(8), intent(in) :: xa(4),ya(4),x

    # evaluate the polynomial coefficients
    x0 = xa[1]
    x1 = xa[2] - x0
    x2 = xa[3] - x0
    x3 = xa[4] - x0
    y0 = ya[1]
    y1 = ya[2] - y0
    y2 = ya[3] - y0
    y3 = ya[4] - y0
    t4 = x1 - x2
    t5 = x1 - x3
    t6 = x2 - x3
    t1 = x1*x2*y3
    t2 = x2*x3*y1
    t3 = x1*x3
    t0 = 1.0/(x2*t3*t4*t5*t6)
    t3 = t3*y2
    c3 = t1*t4 + t2*t6 - t3*t5
    t4 = x1^2
    t5 = x2^2
    t6 = x3^2
    c2 = t1*(t5 - t4) + t2*(t6 - t5) + t3*(t4 - t6)
    c1 = t1*(x2*t4 - x1*t5) + t2*(x3*t5 - x2*t6) + t3*(x1*t6 - x3*t4)
    t1 = x - x0
    # integrate the polynomial
    return t1*(y0 + t0*t1*(0.5*c1 + t1*(c2/3 + 0.25*c3*t1) ) )
end


struct ModifiedEquinoctialElements <: AbstractElements
    p::Float64 # Parameter
    f::Float64 # These
    g::Float64 # Don't
    h::Float64 # Have
    k::Float64 # Names
    l::Float64 # True Longitude
end

# Faster than defaults
parameter(m::ModifiedEquinoctialElements) = m.p
true_longitude(m::ModifiedEquinoctialElements) = m.l

# Conversion from Keplerian Elements
function ModifiedEquinoctialElements(ke::KeplerianElements)
    ecc = eccentricity(ke)
    inc = inclination(ke)
    aop = argument_of_periapsis(ke)
    raan = right_ascension(ke)
    sr, cr = sincos(raan)
    sar, car = sincos(aop + raan)
    ti = tan(inc / 2)
    ModifiedEquinoctialElements(
        parameter(ke),
        ecc * car,
        ecc * sar,
        ti * cr,
        ti * sr,
        true_longitude(ke)
    )
end

# Interface
semi_major_axis(m::ModifiedEquinoctialElements) = parameter(m) / (1 - m.f^2 - m.g^2)
eccentricity(m::ModifiedEquinoctialElements) = sqrt(m.f^2  + m.g^2)
inclination(m::ModifiedEquinoctialElements) = atan(2*sqrt(m.h^2 + m.k^2), 1-m.h^2-m.k^2)
right_ascension(m::ModifiedEquinoctialElements) = atan(m.k, m.h)
true_anomaly(m::ModifiedEquinoctialElements) = m.l - atan(m.g / m.f)

function argument_of_periapsis(m::ModifiedEquinoctialElements)
    num = m.g * m.h - m.f * m.k
    den = m.f * m.h + m.g * m.k
    return atan(num, den)
end


@testset "FPGroups" begin
    A = Alphabet([:a, :A, :b, :B, :c, :C], [2,1,4,3,6,5])

    @test FreeGroup(A) isa FreeGroup
    @test sprint(show, FreeGroup(A)) == "free group on 3 generators"

    F = FreeGroup([:a, :b, :c], A)
    @test sprint(show, F) == "free group on 3 generators"

    a,b,c = gens(F)
    @test c*b*a isa FPGroupElement

    # quotient of F:
    G = FPGroup(F, [a*b=>b*a, a*c=>c*a, b*c=>c*b])

    @test G isa FPGroup
    @test sprint(show, G) == "⟨ a  b  c | \n\t  a*b => b*a  a*c => c*a  b*c => c*b ⟩"
    @test rand(G) isa FPGroupElement

    f = a*c*b
    @test word(f) isa Word{UInt8}

    aG,bG,cG = gens(G)

    @test aG isa FPGroupElement
    @test_throws AssertionError aG == a
    @test word(aG) == word(a)

    g = aG*cG*bG

    @test_throws AssertionError f == g
    @test word(f) == word(g)
    @test word(g) == [1, 5, 3]
    Groups.normalform!(g)
    @test word(g) == [1, 3, 5]

    let g = aG*cG*bG
        # test that we normalize g before printing
        @test sprint(show, g) == "a*b*c"
    end

    # quotient of G
    H = FPGroup(G, [aG^2=>cG, bG*cG=>aG], maxrules=200)

    h = H(word(g))

    @test h isa FPGroupElement
    @test_throws AssertionError h == g
    @test_throws MethodError h*g

    H′ = FPGroup(G, [aG^2=>cG, bG*cG=>aG], maxrules=200)
    @test_throws AssertionError one(H) == one(H′)

    Groups.normalform!(h)
    @test h == H([5])

    @testset "GroupsCore conformance: H" begin
        test_Group_interface(H)
        test_GroupElement_interface(rand(H, 2)...)
    end

end


using Plots,FFTW
using JLD2

for (ind, arg) in enumerate(ARGS)
    t = typeof(arg)
    msg = "$(ind) -> $arg :$t"
    println(msg)
end

function sqaure_windowing_signal(signal::Array{Float64})
    windowed_signal = signal
    window_correction_factor = 1
    return (windowed_signal,window_correction_factor)
end

function hanning_windowing_signal(signal::Array{Float64})
    windowed_signal::Array{Float64} = []
    window::Array{Float64} = []
    total_window_power = 0
    j = 1
    for w = signal
        window_val = 0.5*(1 - cos(2*pi*j/length(signal)))
        push!(window,window_val)
        windowed = window_val*w
        push!(windowed_signal, windowed)
        total_window_power += window_val
        j += 1
    end
    window_correction_factor = length(signal)/total_window_power
    return (windowed_signal,window_correction_factor)
end

function main()
    sampling_rate = 0.0
    window_size = 64*1024*1024;
    w_m_vals::Array{Float64} = []

    jldopen(ARGS[1], "r") do f
        sampling_rate = f["samplingrate"]
        w_m_vals = f["w_m"]
    end

    println(sampling_rate)
    println(length(w_m_vals))

    #(windowed_signal, window_correction_factor)  = sqaure_windowing_signal(w_m_vals[1:window_size])
    (windowed_signal, window_correction_factor) = hanning_windowing_signal(w_m_vals[1:window_size])
    println(window_correction_factor)
    powers = map(x->20*log10(abs(x)*window_correction_factor),fft(windowed_signal)|>fftshift)
    freqs = fftfreq( window_size, sampling_rate )|>fftshift

    plot_graph=plot(freqs, powers, xlims=(0,2000*1000), fmt =:png)
    savefig(plot_graph, "fft1.png")
    plot_graph=plot(freqs, powers, xlims=(944*1000,964*1000), fmt =:png)
    savefig(plot_graph, "fft2.png")
end

main()


"""
Function object to map points on the cube surface to the sphere of radius r.
We need sides of length 2 centred at the origin using the approach
in http://mathproofs.blogspot.com/2005/07/mapping-cube-to-sphere.html.

Construct arguments:
  r  : sphere radius
Function object arguments:
  xyz: 3D cartesian coordinates of a point on the cube surface
"""
struct MapCubeToSphere{T} <: Function
  radius::T
end

function (map::MapCubeToSphere{T})(xyz) where T
  x,y,z = xyz
  xₛ = x*sqrt(1.0-y^2/2-z^2/2+y^2*z^2/(3.0))
  yₛ = y*sqrt(1.0-z^2/2-x^2/2+x^2*z^2/(3.0))
  zₛ = z*sqrt(1.0-x^2/2-y^2/2+x^2*y^2/(3.0))
  map.radius*Point(xₛ,yₛ,zₛ)
end

function generate_Γface_to_bgface(model)
  panel_to_entity=[23,26,24,25,21,22]
  entity_to_panel=Dict(23=>1,26=>2,24=>3,25=>4,21=>5,22=>6)
  ptr_ncells_panel=zeros(Int64,length(panel_to_entity)+1)
  labels = get_face_labeling(model)
  bgface_to_mask = Gridap.Geometry.get_face_mask(labels,"boundary",2)
  Γface_to_bgface = findall(bgface_to_mask)
  face_to_entity  = labels.d_to_dface_to_entity[3][Γface_to_bgface]
  for entity in face_to_entity
    panel=entity_to_panel[entity]
    ptr_ncells_panel[panel+1]=ptr_ncells_panel[panel+1]+1
  end
  ptr_ncells_panel[1]=1
  for i=1:length(ptr_ncells_panel)-1
    ptr_ncells_panel[i+1]=ptr_ncells_panel[i+1]+ptr_ncells_panel[i]
  end
  Γface_to_bgface_panelwise=similar(Γface_to_bgface)
  for (i,entity) in enumerate(face_to_entity)
    panel=entity_to_panel[entity]
    Γface_to_bgface_panelwise[ptr_ncells_panel[panel]]=Γface_to_bgface[i]
    ptr_ncells_panel[panel]=ptr_ncells_panel[panel]+1
  end
  Γface_to_bgface_panelwise
end

function CubedSphereDiscreteModel(n,order; radius=1)
  function _cell_vector_to_dof_vector!(dof_vector,cell_node_ids, cell_vector)
    cache_cell_node_ids = array_cache(cell_node_ids)
    cache_cell_vector   = array_cache(cell_vector)
    for k=1:length(cell_node_ids)
       current_node_ids = getindex!(cache_cell_node_ids,cell_node_ids,k)
       current_values   = getindex!(cache_cell_vector,cell_vector,k)
       for (i,id) in enumerate(current_node_ids)
        dof_vector[current_node_ids[i]]=current_values[i]
       end
    end
  end

  domain = (-1,1,-1,1,-1,1)
  cells  = (n,n,n)
  model  = CartesianDiscreteModel(domain,cells)

  # Restrict model to cube surface
  Γface_to_bgface=generate_Γface_to_bgface(model)
  cube_surface_model = Gridap.Geometry.BoundaryDiscreteModel(Polytope{2},model,Γface_to_bgface)

  # Generate high-order FE map and ordering
  vector_reffe=ReferenceFE(lagrangian,VectorValue{3,Float64},order)
  V = FESpace(cube_surface_model,vector_reffe; conformity=:H1)
  vh = interpolate(MapCubeToSphere(radius),V)
  scalar_reffe=ReferenceFE(QUAD,lagrangian,Float64,order)
  xref=Gridap.ReferenceFEs.get_node_coordinates(scalar_reffe)
  xrefₖ=Fill(xref,num_cells(cube_surface_model))
  vhx=lazy_map(evaluate,Gridap.CellData.get_data(vh),xrefₖ)
  V = FESpace(cube_surface_model,scalar_reffe; conformity=:H1)
  node_coordinates = Vector{Point{3,Float64}}(undef,num_free_dofs(V))
  cell_node_ids    = get_cell_dof_ids(V)
  _cell_vector_to_dof_vector!(node_coordinates,cell_node_ids,vhx)
  cell_types  = collect(Fill(1,num_cells(cube_surface_model)))
  cell_reffes = [scalar_reffe]

  cube_surface_grid = Gridap.Geometry.UnstructuredGrid(node_coordinates,
                                                       Gridap.Arrays.Table(cell_node_ids),
                                                       cell_reffes,
                                                       cell_types,
                                                       Gridap.Geometry.Oriented())
  topology=Gridap.Geometry.get_grid_topology(cube_surface_model)
  face_labeling=Gridap.Geometry.get_face_labeling(cube_surface_model)
  Gridap.Geometry.UnstructuredDiscreteModel(cube_surface_grid,topology,face_labeling)
end

struct AnalyticalMapCubedSphereDiscreteModel{T,B,C} <: Gridap.Geometry.DiscreteModel{2,3}
  cell_map::T
  cubed_sphere_model::B
  trian::C
  function AnalyticalMapCubedSphereDiscreteModel(n;radius=1)
    domain = (-1,1,-1,1,-1,1)
    cells  = (n,n,n)
    model  = CartesianDiscreteModel(domain,cells)

    # Restrict model to cube surface
    Γface_to_bgface=generate_Γface_to_bgface(model)
    cube_surface_model = Gridap.Geometry.BoundaryDiscreteModel(Polytope{2},model,Γface_to_bgface)

    m1=Fill(Gridap.Fields.GenericField(MapCubeToSphere(radius)),num_cells(cube_surface_model))
    m2=get_cell_map(cube_surface_model)
    m=lazy_map(∘,m1,m2)

    cubed_mesh_model=CubedSphereDiscreteModel(n,1)

    # Wrap up BoundaryTriangulation
    btrian=Triangulation(cubed_mesh_model)
    trian=AnalyticalMapCubedSphereTriangulation(m,btrian)

    # Build output object
    T=typeof(m)
    B=typeof(cubed_mesh_model)
    C=typeof(trian)
    GC.gc()
    new{T,B,C}(m,cubed_mesh_model,trian)
  end
end

Gridap.Geometry.get_cell_map(model::AnalyticalMapCubedSphereDiscreteModel) = model.cell_map
Gridap.Geometry.get_grid(model::AnalyticalMapCubedSphereDiscreteModel) = Gridap.Geometry.get_grid(model.cubed_sphere_model)
Gridap.Geometry.get_grid_topology(model::AnalyticalMapCubedSphereDiscreteModel) = Gridap.Geometry.get_grid_topology(model.cubed_sphere_model)
Gridap.Geometry.get_face_labeling(model::AnalyticalMapCubedSphereDiscreteModel) = Gridap.Geometry.get_face_labeling(model.cubed_sphere_model)
Gridap.Geometry.get_triangulation(a::AnalyticalMapCubedSphereDiscreteModel) = a.trian
Gridap.Geometry.Triangulation(a::AnalyticalMapCubedSphereDiscreteModel) = a.trian

function CubedSphereDiscreteModel(n;radius=1)
  AnalyticalMapCubedSphereDiscreteModel(n;radius)
end

const CSDMT = Union{AnalyticalMapCubedSphereDiscreteModel,
                 <:Gridap.Geometry.UnstructuredDiscreteModel{2,3}}

function Gridap.CellData.get_normal_vector(model::CSDMT)
    cell_normal = Gridap.Geometry.get_facet_normal(model)
    Gridap.CellData.GenericCellField(cell_normal,Triangulation(model),ReferenceDomain())
end

function Gridap.Geometry.get_facet_normal(model::CSDMT)
  function _unit_outward_normal(v::Gridap.Fields.MultiValue{Tuple{2,3}},sign_flip::Bool)
    n1 = v[1,2]*v[2,3] - v[1,3]*v[2,2]
    n2 = v[1,3]*v[2,1] - v[1,1]*v[2,3]
    n3 = v[1,1]*v[2,2] - v[1,2]*v[2,1]
    n = VectorValue(n1,n2,n3)
    (-1)^sign_flip*n/norm(n)
  end

  # Get the Jacobian of the cubed sphere mesh
  trian = Triangulation(model)
  map   = get_cell_map(trian)
  Jt    = lazy_map(∇,map)

  # Get the index of the panel for each element
  fl = get_face_labeling(model)
  panel_id  = fl.d_to_dface_to_entity[3]
  sign_flip = [panel_id[i] == 25 || panel_id[i] == 21 || panel_id[i] == 24 for i=1:length(panel_id)]
  fsign_flip = lazy_map(Gridap.Fields.ConstantField,sign_flip)

  lazy_map(Operation(_unit_outward_normal),Jt,fsign_flip)
end


# Accelerated row matrix product

C = zeros(m,p)

for j=1:m
    C[j,:] = BLAS.gemv(’T’,B,A[j,:])
end

Code/Numerical Analysis & Machine Learning/Numerical Linear Algebra/RowMatrixProduct.jl


@testset "Primitives" begin
  @testset "Lines" begin
    l = Line(P2(0,0), P2(1,1))
    @test paramdim(l) == 1
    @test isconvex(l)

    l = Line(P2(0,0), P2(1,1))
    @test (l(0), l(1)) == (P2(0,0), P2(1,1))
  end

  @testset "Rays" begin
    r = Ray(P2(0,0), V2(1,1))
    @test paramdim(r) == 1
    @test isconvex(r)

    r = Ray(P2(0,0), V2(1,1))
    @test r(T(0.)) == P2(0,0)
    @test r(T(1.)) == P2(1,1)
    @test r(T(Inf)) == P2(Inf,Inf)
    @test_throws DomainError(T(-1), "r(t) is not defined for t < 0.") r(T(-1))
  end

  @testset "Planes" begin
    p = Plane(P3(0, 0, 0), V3(1, 0, 0), V3(0, 1, 0))
    @test p(T(1.0), T(0.0)) == P3(1, 0, 0)
    @test paramdim(p) == 2
    @test embeddim(p) == 3
    @test isconvex(p)
    @test origin(p) == P3(0, 0, 0)
    @test normal(p) == Vec(0, 0, 1)
  end

  @testset "Bezier curves" begin
    # fix import conflict with Plots
    BezierCurve = Meshes.BezierCurve

    b = BezierCurve(P2(0,0),P2(0.5,1),P2(1,0))
    for method in [DeCasteljau(), Horner()]
      @test b(T(0), method) == P2(0,0)
      @test b(T(1), method) == P2(1,0)
      @test b(T(0.5), method) == P2(0.5,0.5)
      @test b(T(0.5), method) == P2(0.5,0.5)
      @test_throws DomainError(T(-0.1), "b(t) is not defined for t outside [0, 1].") b(T(-0.1), method)
      @test_throws DomainError(T(1.2), "b(t) is not defined for t outside [0, 1].") b(T(1.2), method)
    end

    b = BezierCurve(P2.(randn(100), randn(100)))
    t1 = @timed b(T(0.2))
    t2 = @timed b(T(0.2), Horner())
    @test t1.time > t2.time
    @test t2.bytes < 100
  end

  @testset "Boxes" begin
    b = Box(P2(0,0), P2(1,1))
    @test embeddim(b) == 2
    @test paramdim(b) == 2
    @test coordtype(b) == T
    @test minimum(b) == P2(0,0)
    @test maximum(b) == P2(1,1)
    @test extrema(b) == (P2(0,0), P2(1,1))
    @test isconvex(b)

    b = Box(P2(0,0), P2(1,1))
    @test measure(b) == T(1)
    @test P2(1,1) ∈ b

    b = Box(P2(1,1), P2(2,2))
    @test sides(b) == T[1,1]
    @test Meshes.center(b) == P2(1.5,1.5)
    @test diagonal(b) == √T(2)

    b = Box(P2(1,2), P2(3,4))
    @test vertices(b) == P2[(1,2),(3,2),(3,4),(1,4)]

    b = Box(P3(1,2,3), P3(4,5,6))
    @test vertices(b) == P3[(1,2,3),(4,2,3),(4,5,3),(1,5,3),(1,2,6),(4,2,6),(4,5,6),(1,5,6)]

    b = Box(P2(0,0), P2(1,1))
    @test boundary(b) == Chain(P2[(0,0),(1,0),(1,1),(0,1),(0,0)])

    b = Box(P3(0,0,0), P3(1,1,1))
    m = boundary(b)
    @test m isa Mesh
    @test nvertices(m) == 8
    @test nelements(m) == 6

    # subsetting with boxes
    b1 = Box(P2(0,0), P2(0.5,0.5))
    b2 = Box(P2(0.1,0.1), P2(0.5,0.5))
    b3 = Box(P2(0,0), P2(1,1))
    @test b1 ⊆ b3
    @test b2 ⊆ b3
    @test !(b1 ⊆ b2)
    @test !(b3 ⊆ b1)
    @test !(b3 ⊆ b1)
  end

  @testset "Ball" begin
    b = Ball(P3(1,2,3), T(5))
    @test embeddim(b) == 3
    @test paramdim(b) == 3
    @test coordtype(b) == T
    @test Meshes.center(b) == P3(1,2,3)
    @test radius(b) == T(5)
    @test isconvex(b)

    b = Ball(P2(0,0), T(2))
    @test measure(b) ≈ π*(2^2)
    b = Ball(P3(0,0,0), T(2))
    @test measure(b) ≈ (4/3)*π*(2^3)

    b = Ball(P2(0,0), T(2))
    @test P2(1,0) ∈ b
    @test P2(0,1) ∈ b
    @test P2(2,0) ∈ b
    @test P2(0,2) ∈ b
    @test P2(3,5) ∉ b

    b = Ball(P3(0,0,0), T(2))
    @test P3(1,0,0) ∈ b
    @test P3(0,0,1) ∈ b
    @test P3(2,0,0) ∈ b
    @test P3(0,0,2) ∈ b
    @test P3(3,5,2) ∉ b
  end

  @testset "Sphere" begin
    s = Sphere(P3(0,0,0), T(1))
    @test embeddim(s) == 3
    @test paramdim(s) == 2
    @test coordtype(s) == T
    @test Meshes.center(s) == P3(0, 0, 0)
    @test radius(s) == T(1)
    @test !isconvex(s)

    s = Sphere(P2(0,0), T(2))
    @test measure(s) ≈ 2π*2
    s = Sphere(P3(0,0,0), T(2))
    @test measure(s) ≈ 4π*(2^2)

    s = Sphere(P2(0,0), T(2))
    @test P2(1,0) ∉ s
    @test P2(0,1) ∉ s
    @test P2(2,0) ∈ s
    @test P2(0,2) ∈ s
    @test P2(3,5) ∉ s

    s = Sphere(P3(0,0,0), T(2))
    @test P3(1,0,0) ∉ s
    @test P3(0,0,1) ∉ s
    @test P3(2,0,0) ∈ s
    @test P3(0,0,2) ∈ s
    @test P3(3,5,2) ∉ s

    # 2D sphere passing through 3 points
    s = Sphere(P2(0,0), P2(0.5,0), P2(1,1))
    @test Meshes.center(s) == P2(0.25, 0.75)
    @test radius(s) == T(0.7905694150420949)
    s = Sphere(P2(0,0), P2(1,0), P2(0,1))
    @test Meshes.center(s) == P2(0.5, 0.5)
    @test radius(s) == T(0.7071067811865476)
    s = Sphere(P2(0,0), P2(1,0), P2(1,1))
    @test Meshes.center(s) == P2(0.5, 0.5)
    @test radius(s) == T(0.7071067811865476)
  end

  @testset "Cylinder" begin
    c = Cylinder(P3(1,2,3), P3(4,5,6), T(5))
    @test embeddim(c) == 3
    @test paramdim(c) == 3
    @test coordtype(c) == T
    @test radius(c) == T(5)
    @test height(c) ≈ √27
    @test isconvex(c)

    @test measure(c) ≈ π*5.0^2*√27
  end
end


"""
    module Interpolation

Module implementing various interpolation schemes commonly used in finite
element and boundary element methods.
"""
module Interpolation

using StaticArrays
using LinearAlgebra

using WavePropBase
using WavePropBase.Utils
using WavePropBase.Geometry

WavePropBase.@import_interface

export
    # abstract types
    AbstractElement,
    AbstractHyperRectangle,
    # structs
    ParametricElement,
    LagrangeElement,
    LagrangePoint,
    LagrangeLine,
    LagrangeTriangle,
    LagrangeSquare,
    LagrangeTetrahedron,
    Pk,
    TensorLagInterp,
    ChebInterp,
    MonomialBasis,
    PolynomialBasis,
    LagrangeBasis,
    HyperRectangle,
    HyperCube,
    # methods
    reference_nodes,
    monomial_basis,
    lagrange_basis,
    barycentric_lagrange_matrix,
    barycentric_lagrange_weights,
    vandermond,
    lagrange_basis,
    gradLagrangeBasis,
    interpolation_nodes,
    trace,
    cheb1nodes,
    cheb1nodes_iter,
    cheb1weights,
    cheb2nodes,
    cheb2nodes_iter,
    cheb2weights,
    grad,
    vals,
    low_corner,
    high_corner,
    precompute_nodes_and_weights!

include("polynomials.jl")
include("element.jl")
include("hyperrectangle.jl")
include("tensorlaginterp.jl")

end # module Interpolation


using SurrogateModelOptim
using PlotlyJS
using Statistics
using Distributed

# Discontinuous Forrester function
Forrester(x) = (6*x[1] - 2)^2 * sin(12*x[1] - 4)
Forrester_disc(x) = x[1]>0.15 ? Forrester(x[1])+15 : Forrester(x[1])
search_range=[(0.0,1.0)]

# Optimize the test function
result = smoptimize(Forrester_disc, search_range;
                    options=SurrogateModelOptim.Options(
                    iterations=25,
                    num_interpolants=20, #Preferably even number of added processes
                    num_start_samples=5,
                        ));

function plot_fun_1D(fun_original,fun_estimate,search_range)
    N = 101
    x = range(search_range[1][1], length=N,stop=search_range[1][2])
    y_original = fun_original.(x)
    y_estimate = [median(fun_estimate([x])) for x in x]

    trace1 = scatter(;x=x, y=y_original, mode="lines",name="Original function")
    trace2 = scatter(;x=x, y=y_estimate, mode="lines",name="Estimated function")
    p = plot([trace1,trace2])
end

# Plot the results
display(plot_fun_1D(Forrester_disc,result.sm_interpolant,search_range))


using CUDA
using CUDA.CUSPARSE
using ExaPF
using KernelAbstractions
using LinearAlgebra
using Random
using SparseArrays
using Test
using TimerOutputs

import ExaPF: Parse, PowerSystem

Random.seed!(2713)

case = "case14.raw"
# case = "ACTIVSg70K.raw"

@testset "Powerflow residuals and Jacobian" begin
    # read data
    to = TimerOutputs.TimerOutput()
    datafile = joinpath(dirname(@__FILE__), case)
    data = Parse.parse_raw(datafile)
    BUS_B, BUS_AREA, BUS_VM, BUS_VA, BUS_NVHI, BUS_NVLO, BUS_EVHI,
    BUS_EVLO, BUS_TYPE = Parse.idx_bus()
    bus = data["BUS"]
    nbus = size(bus, 1)

    # obtain V0 from raw data
    V = Array{Complex{Float64}}(undef, nbus)
    T = Vector
    for i in 1:nbus
        V[i] = bus[i, BUS_VM]*exp(1im * pi/180 * bus[i, BUS_VA])
    end

    # form Y matrix
    Ybus, Yf_br, Yt_br, Yf_tr, Yt_tr = PowerSystem.makeYbus(data);

    Vm = abs.(V)
    Va = angle.(V)
    bus = data["BUS"]
    gen = data["GENERATOR"]
    load = data["LOAD"]
    nbus = size(bus, 1)
    ngen = size(gen, 1)
    nload = size(load, 1)

    ybus_re, ybus_im = ExaPF.Spmat{T}(Ybus)
    SBASE = data["CASE IDENTIFICATION"][1]
    Sbus = PowerSystem.assembleSbus(gen, load, SBASE, nbus)
    pbus = real(Sbus)
    qbus = imag(Sbus)

    ref, pv, pq = PowerSystem.bustypeindex(bus, gen)
    npv = size(pv, 1);
    npq = size(pq, 1);

    @testset "Computing residuals" begin
        F = zeros(Float64, npv + 2*npq)
        # First compute a reference value for resisual computed at V
        F♯ = ExaPF.residualFunction(V, Ybus, Sbus, pv, pq)
        # residual_polar! uses only binary types as this function is meant
        # to be deported on the GPU
        ExaPF.residualFunction_polar!(F, Vm, Va,
            ybus_re.nzval, ybus_re.colptr, ybus_re.rowval,
            ybus_im.nzval, ybus_im.colptr, ybus_im.rowval,
            pbus, qbus, pv, pq, nbus)
        @test F ≈ F♯
    end
    @testset "Computing Jacobian of residuals" begin
        F = zeros(Float64, npv + 2*npq)
        # Compute Jacobian at point V manually and use it as reference
        J = ExaPF.residualJacobian(V, Ybus, pv, pq)
        J♯ = copy(J)

        # Then, create a JacobianAD object
        coloring = ExaPF.matrix_colors(J)
        jacobianAD = ExaPF.AD.JacobianAD(J, coloring, F, Vm, Va, pv, pq)
        # and compute Jacobian with ForwardDiff
        ExaPF.AD.residualJacobianAD!(
            jacobianAD, ExaPF.residualFunction_polar!, Vm, Va,
            ybus_re, ybus_im, pbus, qbus, pv, pq, nbus, to)
        @test jacobianAD.J ≈ J♯
    end
end

@testset "Wrapping of iterative solvers" begin
    nblocks = 2
    n, m = 32, 32
    to = TimerOutputs.TimerOutput()

    # Add a diagonal term for conditionning
    A = randn(n, m) + 15I
    x♯ = randn(m)
    b = A * x♯

    # Be careful: all algorithms work with sparse matrix
    A = sparse(A)
    # Init iterators for KernelAbstractions
    if has_cuda_gpu()
        iterators = zip([CPU(), CUDADevice()],
                        [Array{Float64}, CuVector{Float64}],
                        [SparseMatrixCSC, CuSparseMatrixCSR])
    else
        iterators = zip([CPU()], [Array{Float64}], [SparseMatrixCSC])
    end
    @testset "Run iterative solvers on device $(device)" for (device, V, SM) in iterators
        As = SM(A)
        bs = convert(V, b)
        x0 = convert(V, zeros(m))
        xs♯ = convert(V, x♯)

        # TODO: currently Preconditioner takes as input only sparse matrix
        # defined on the main memory.
        precond = ExaPF.Precondition.Preconditioner(A, nblocks, device)

        # First test the custom implementation of BICGSTAB
        @testset "BICGSTAB" begin
            # Need to update preconditioner before resolution
            ExaPF.Precondition.update(As, precond, to)
            P = precond.P
            fill!(x0, 0.0)
            x_sol, n_iters = ExaPF.bicgstab(As, bs, P, x0, to)
            @test n_iters <= m
            @test x_sol ≈ xs♯ atol=1e-6
        end
        @testset "Interface for iterative algorithm ($algo)" for algo in ExaPF.list_solvers(device)
            fill!(x0, 0.0)
            n_iters = ExaPF.Iterative.ldiv!(x0, As, bs, algo, precond, to)
            @test n_iters <= m
            @test x0 ≈ xs♯ atol=1e-6
        end
    end
end

@testset "Powerflow solver" begin
    datafile = joinpath(dirname(@__FILE__), case)
    nblocks = 8
    # Create a network object:
    pf = ExaPF.PowerSystem.PowerNetwork(datafile)
    x = ExaPF.PowerSystem.get_x(pf)
    u = ExaPF.PowerSystem.get_u(pf)
    p = ExaPF.PowerSystem.get_p(pf)
    target = CPU()
    @testset "[CPU] Powerflow solver $precond" for precond in ExaPF.list_solvers(target)
        sol, has_conv, res = solve(pf, x, u, p, nblocks, precond, device=target)
        @test has_conv
        @test res < 1e-6
    end

    if has_cuda_gpu()
        target = CUDADevice()
        @testset "[GPU] Powerflow solver $precond" for precond in ExaPF.list_solvers(target)
            sol, conv, res = solve(pf, x, u, p, nblocks, precond, device=target)
            @test conv
            @test res < 1e-6
        end
    end
end

# # Not working yet. Will check whether Ipopt and reduced method match in objective
# @testset "rgm_3bus" begin
#    include("../scripts/rgm_3bus.jl")
#    @show red_cost = cfun(xk, uk, p)
#    include("../scripts/ipopt.jl")
#    @show ipopt_cost = cfun(xk, uk, p)
#    gap = abs(red_cost - ipopt_cost)
#    println("gap = abs(red_cost - ipopt_cost): $gap = abs($red_cost - $ipopt_cost)")
#    @test gap ≈ 0.0
# end

# @testset "rgm_3bus_ref" begin
#    include("../scripts/rgm_3bus_ref.jl")
#    @show red_cost = cfun(xk, uk, p)
#    include("../scripts/ipopt_ref.jl")
#    @show ipopt_cost = cfun(xk, uk, p)
#    gap = abs(red_cost - ipopt_cost)
#    println("gap = abs(red_cost - ipopt_cost): $gap = abs($red_cost - $ipopt_cost)")
#    @test gap ≈ 0.0
# end


"""
    Sedenion{T <: Real} = Elliptic{Octonion{T}}

A sedenion is an elliptic Cayley-Dickson construct with Octonion octonions.
"""
const Sedenion{T <: Real} = Elliptic{Octonion{T}}

function Sedenion(z::Octonion{T}) where T <: Real
    Sedenion{T}(z, zero(Octonion{T}))
end

function Sedenion(z::Quaternion{T}) where T <: Real
    Sedenion{T}(
        Octonion{T}(z, zero(Quaternion{T})),
        zero(Octonion{T})
    )
end

function Sedenion(z::Binion{T}) where T <: Real
    Sedenion{T}(
        Octonion{T}(Quaternion{T}(z, zero(Binion{T})),
        zero(Quaternion{T})),
        zero(Octonion{T})
    )
end

function Sedenion(v::T...) where T <: Real
    n = length(v)
    if n <= 8
        return Sedenion{T}(Octonion(v...), zero(Octonion{T}))
    end

    if n > 16
        return Sedenion{T}(Octonion(v[1:8]...), Octonion(v[9:16]...))
    end

    Sedenion{T}(Octonion(v[1:8]...), Octonion(v[9:end]...))
end

"""
    unreal(z::Sedenion)

The unreal part of a sedenion is a 15-vector.
"""
function unreal(z::Sedenion)
    vcat(unreal(z.l), array(z.r))
end

function show(io::IO, z::Sedenion)
    print(io, "[1: ")
    print(io, z.l.l.l.l)
    print(io, ", i: ")
    print(io, z.l.l.l.r)
    print(io, ", j: ")
    print(io, z.l.l.r.l)
    print(io, ", ij: ")
    print(io, z.l.l.r.r)
    print(io, ", k: ")
    print(io, z.l.r.l.l)
    print(io, ", ik: ")
    print(io, z.l.r.l.r)
    print(io, ", jk: ")
    print(io, z.l.r.r.l)
    print(io, ", (ij)k: ")
    print(io, z.l.r.r.r)
    print(io, ", m: ")
    print(io, z.r.l.l.l)
    print(io, ", im: ")
    print(io, z.r.l.l.r)
    print(io, ", jm: ")
    print(io, z.r.l.r.l)
    print(io, ", (ij)m: ")
    print(io, z.r.l.r.r)
    print(io, ", km: ")
    print(io, z.r.r.l.l)
    print(io, ", (ik)m: ")
    print(io, z.r.r.l.r)
    print(io, ", (jk)m: ")
    print(io, z.r.r.r.l)
    print(io, ", ((ij)k)m: ")
    print(io, z.r.r.r.r)
    print(io, "]")
end

function random(::Type{Sedenion{T}}) where T <: Real
    Sedenion{T}(random(Octonion{T}), random(Octonion{T}))
end


using NURBS.Bases

bs = BSplineBasis(0, 1, 1, 4)
pts = linspace(0, 1, 100000)

bs(pts)
@time bs(pts)

@profile bs(pts)
Profile.print()

Profile.clear_malloc_data()
bs(pts)


using StatsBase
using LinearAlgebra
using PyCall

import Distributions.dim

####################
# Sensor definitions
####################

export ProbabilisticAgentSpecification, ProbabilisticSensor

struct ProbabilisticAgentSpecification
  max_success_probability::Float64
  sensor_radius::Float64
  station_radius::Float64
  num_sensors::Float64
end

# detection probability is of the form:
# max_success_probability * exp(-||x-center||^4 / sensor_radius^4)
struct ProbabilisticSensor
  center::Array{Float64, 1}
  sensor_radius::Float64
  max_success_probability::Float64
end

function make_agent(agent_specification::ProbabilisticAgentSpecification)
  agent_center = rand(2)

  make_sensor() = ProbabilisticSensor(agent_center
                                    + agent_specification.station_radius *
                                    rand_in_circle(),
                                    agent_specification.sensor_radius,
                                    agent_specification.max_success_probability)

  sensors = [make_sensor() for agent in 1:agent_specification.num_sensors]

  Agent(agent_center, agent_specification.station_radius, sensors)
end

##############
# Sensor model
##############

export detection_probability, mean_detection_probability

function detection_probability(sensor::ProbabilisticSensor, x)
  square_dist = sum((x - sensor.center).^2)

  sensor.max_success_probability * exp(-square_dist^2 / sensor.sensor_radius^4)
end

function detection_probability(sensors::Array{ProbabilisticSensor,1}, x)
  p_fail = reduce(sensors; init=1.0) do p, sensor
    p * (1.0 - detection_probability(sensor, x))
  end

  1.0 - p_fail
end

function mean_detection_probability(sensors, points)
  sum(map(1:size(points, 2)) do ii
        detection_probability(sensors, points[:,ii])
      end) / length(points)
end

###################
# Environment model
###################

export Gaussian, GaussianMixture, sample_pdf, evaluate_pdf, dim,
in_limits, sample_reject, standard_mixture, generate_events

pdf_coefficient(covariance) = sqrt((2*pi)^size(covariance,1) * det(covariance))

struct Gaussian
  mean::Array{Float64,1}
  covariance::Array{Float64,2}
  chol_covariance # result of cholesky decomposition
  inv_covariance # result of inverse
  det_covariance::Float64 # result of determinant
  pdf_coefficient::Float64 # result of pdf_coefficient

  function Gaussian(mean, covariance)
    new(mean, covariance, cholesky(covariance), inv(covariance), det(covariance),
        pdf_coefficient(covariance))
  end
end

struct GaussianMixture
  weights::Array{Float64,1}
  components::Array{Gaussian,1}
end

dim(g::Gaussian) = length(g.mean)

dim(gm::GaussianMixture) = dim(gm.components[1])

function sample_pdf(g::Gaussian)
  g.mean + (randn(1,size(g.covariance,1)) * g.chol_covariance.U)'
end

function sample_pdf(gm::GaussianMixture)
  normalized_weights = gm.weights / sum(gm.weights)

  sample_pdf(sample(gm.components, Weights(normalized_weights)))
end

in_limits(x, limits) = all(map(1:length(x)) do ii
                             limits[ii,1] <= x[ii] && x[ii] <= limits[ii,2]
                           end)

function sample_reject(p, limits)
  ret = zeros(dim(p), dim(p))
  found_valid = false

  while ~found_valid
    ret = sample_pdf(p)
    found_valid = in_limits(ret, limits)
  end

  ret
end

function evaluate_pdf(g::Gaussian, x::Array{T,1}) where T <: Number
  d = x - g.mean

  exp(- d' * g.inv_covariance * d / 2)[1] * g.pdf_coefficient
end

function evaluate_pdf(gm::GaussianMixture, x::Array{T,1}) where T <: Number
  reduce(zip(gm.weights, gm.components); init=0.0) do value, w_c
    value + w_c[1] * evaluate_pdf(w_c[2], x)
  end
end

function standard_mixture()
  w = [0.30, 0.6, 0.1]
  g1 = Gaussian([0.2, 0.8], Diagonal([0.004, 0.1]))
  g2 = Gaussian([0.8, 0.2], Diagonal([0.1, 0.01]))
  g3 = Gaussian([0.7, 0.7], Diagonal([0.03, 0.03]))

  GaussianMixture(w, [g1, g2, g3])
end

function generate_events(n = 500; limits = [0 1; 0 1])
  dist = standard_mixture()

  hcat(map(x->sample_reject(dist, limits), 1:n)...)
end

###############
# visualization
###############
export visualize_events

function plot_element(circle::ProbabilisticSensor; color = "k")
  center = circle.center
  scatter([center[1]], [center[2]], color = color, s = 4*agent_scale,
          marker = sensor_center)
end

function visualize_solution(sensors::Array{ProbabilisticSensor}, colors;
                            limits = [0 1; 0 1], n = 1000, cmap = "viridis")
  xlim = limits[1,:]
  ylim = limits[2,:]

  # Highlight the sensors
  map(sensors, colors) do sensor, color
    center = sensor.center
    scatter([center[1]], [center[2]], color = rgb_tuple(color), s = 4*agent_scale,
            marker = selected_sensor, edgecolors="k")
  end

  # Plot the actual distribution
  probabilities = [detection_probability(sensors, [x,y])
                           for x in range(xlim[1], stop=xlim[2], length=n),
                           y in range(ylim[1], stop=ylim[2], length=n)]

  # detection probability is a direct probability of detection at a given
  # locationand so varies from zero to one
  image = imshow(probabilities', cmap=cmap, vmin=0.0, vmax=1.0,
                 extent=[xlim[1], xlim[2], ylim[1], ylim[2]],
                 interpolation="nearest", origin="lower")
  make_tight_colorbar(image)

  nothing
end

function visualize_events(events)
  scatter(events[1,:]', events[2,:]', color = "k", marker = event_marker,
          s = 2.0 * agent_scale, linewidth = 1.0)
end


using JLD2, Attitude, MATLAB, SatelliteDynamics, Dierckx
# include ksfunctions.jl
include(joinpath(dirname(dirname(@__FILE__)),"ks_functions.jl"))
@load "rate_control_transfers/60_day_v2_transfer.jld2" X U t_hist
dscale = 1e7 # m
tscale = 20000 # s
uscale = 10000.0
dt = 3e-2

r_eci_hist = dscale*mat_from_vec([x_from_u(X[i][1:4]) for i = 1:length(X)])
v_eci_hist = (dscale/tscale)*mat_from_vec([xdot_from_u(X[i][1:4],X[i][5:8]) for i = 1:length(X)])

# TODO: Spline this in a data efficient way
mat"
figure
hold on
plot3($r_eci_hist(1,:),$r_eci_hist(2,:),$r_eci_hist(3,:))
axis equal
hold off
"

Xm = mat_from_vec(X)
U_real = Xm[12:14,:]

mat"
figure
hold on
title('Thrust Vector')
plot($t_hist,$U_real')
ylabel('Thrust Acceleration (0.0001m/s^2)')
xlabel('Time (days)')
hold off
xlim([0,$t_hist(end)])
legend('u_x','u_y','u_z')
%saveas(gcf,'long_traj.png')
"

U_real_vec = [X[i][12:14] for i = 1:length(X)]
Unorm = [norm(U_real_vec[i]) for i = 1:length(X)]

mat"
figure
hold on
plot($t_hist,$Unorm/$uscale)
xlim([$t_hist(4),$t_hist(end)])
xlabel('Time (days)')
ylabel('Thrust Magnitude (m/s^2)')
hold off
"

println("Time in days:  $t_days")

# now let's get orbital elements
const SD = SatelliteDynamics
oe_hist = [SD.sCARTtoOSC([r_eci_hist[:,i];v_eci_hist[:,i]]) for i = 1:size(r_eci_hist,2)]

a_hist = [oe_hist[i][1] for i = 1:length(oe_hist)]
e_hist = [oe_hist[i][2] for i = 1:length(oe_hist)]
i_hist = [rad2deg(oe_hist[i][3]) for i = 1:length(oe_hist)]

mat"
figure
hold on
subplot(3,1,1)
hold on
plot($t_hist,$a_hist/1000)
h = plot($t_hist, 42e3*ones(length($t_hist),1),'r');
legend([h(1)],'GEO SMA','Location','SouthEast')
ylabel('Semi-Major Axis (km)')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
ylim([20e3,46e3])
subplot(3,1,2)
hold on
plot($t_hist,$e_hist)
ylabel('Eccentricity')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
subplot(3,1,3)
hold on
plot($t_hist,$i_hist)
ylabel('Inclination (deg)')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
hold off
"

angle_hist = get_rtn_in_out(U_real,r_eci_hist,v_eci_hist)
angles = mat_from_vec(angle_hist)

sixty = (t_hist = t_hist, a_hist = a_hist, e_hist = e_hist, i_hist = i_hist, Unorm = Unorm,angles = angles)

using JLD2
@save "plotting_data/plotting_data_60.jld2" sixty


# SimpleMPSJulia
# Density matrix renormalization group (DMRG) in matrix product state (MPS)

# This file contains the definition of matrix product state
# For theoretical backgrounds, see the [reference]:
# Ulrich Schollwöck, The density-matrix renormalization group in the age of matrix product states,
# Annals of Physics, 326 (2011), 96-192

module SimpleMPS

export MatrixProductOperator, MatrixProductState, search_groundstate, expectation

using TensorOperations
using LinearAlgebra
using LinearMaps
using IterativeSolvers
using Statistics

"""
Base class for MPS and MPO
"""
abstract type MatrixProduct end

Base.length(mp::MatrixProduct) = length(mp.mp)
Base.iterate(mp::MatrixProduct, state=1) = iterate(mp.mp, state)
Base.show(io::IO, mp::MatrixProduct) = print(io, [size(m) for m in mp.mp])

struct MatrixProductOperator <: MatrixProduct
    mp::Array{Array}
    MatrixProductOperator(mp) = new(mp)
    function MatrixProductOperator(single_mo::Array, site_num::Int)
        # the first MPO site, only contains the last row
        mpo_1 = single_mo[end:end, :, :, :]
        # the last MPO site, only contains the first column
        mpo_n = single_mo[:, :, :, 1:1]
        new([[mpo_1]; fill(single_mo, site_num - 2); [mpo_n]])
    end
end

struct MatrixProductState <: MatrixProduct
    mp::Array{Array}
    MatrixProductState(mp) = new(mp)
    """
    Construct MPS with random matrices.
    The physical dimension is obtained from `mpo` and bond dimension is set to `m`
    """
    function MatrixProductState(mpo::MatrixProductOperator, m::Int)
        bond_dim = [[1]; fill(m, length(mpo) - 1); [1]]
        pdim = [size(m)[2] for m in mpo]
        # Firstly randomly fill in the matrices
        mp = [rand(bond_dim[i], pdim[i], bond_dim[i+1]) for i in 1:length(mpo)]
        mps = new(mp)
        # Then perform canonicalisation (and normalization)
        for i in 1:length(mpo)-1
            push_center!(mps, i)
        end
        return mps
    end
end

"""
Flip MPO and MPS. After the operation the left side of MPO (MPS) becomes the
right side and vice versa.
Note that <mps|mpo|mps> conserves after the operation.
"""
Base.reverse(mpo::MatrixProductOperator) = MatrixProductOperator(reverse([permutedims(m, (4, 2, 3, 1)) for m in mpo]))
Base.reverse(mps::MatrixProductState) = MatrixProductState(reverse([permutedims(m, (3, 2, 1)) for m in mps]))

"""
Perform left-canonicalisation on the `idx`th site of `mps` so that the canonical center is moved to `idx+1`.
The matrices are normalized by the way.
"""
function push_center!(mps::MatrixProductState, idx::Int)
    orig_shape = size(mps.mp[idx])
    ms = reshape(mps.mp[idx], (prod(orig_shape[1:2]), :))
    f = qr(ms)
    mps.mp[idx] = reshape(Array(f.Q), orig_shape[1], orig_shape[2], :)
    @tensor mps.mp[idx+1][a, c, d] := f.R[a, b] * mps.mp[idx+1][b, c, d]
    mps.mp[idx+1] /= norm(mps.mp[idx+1])
end

"""
Calculate environment on the next site based on previous environment `res`
`ms` and `mo` are local sites of the MPS and MPO.
"""
function update_env(res, ms, mo)
    """
    Graphical notation: (* for MPS and # for MPO)
    *--a--*--e
    |     |
    |     d
    |     |
    #--b--#--g
    |     |
    |     f
    |     |
    *--c--*--h
    """
    @tensor res[e, g, h] := res[a, b, c] * conj(ms)[a, d, e] * mo[b, d, f, g] * ms[c, f, h]
end

"""
Calculate an array of left-side enrivonment tensor based on `mps` and `mpo`.
The index `i` of the array represents the left environment of the `i`th site of the MPS.
"""
function get_env(mps::MatrixProductState, mpo::MatrixProductOperator)::Array{Array}
    sentinel = ones(1, 1, 1)
    env = [sentinel]
    for (ms, mo) in zip(mps, mpo)
        res = update_env(env[length(env)], ms, mo)
        push!(env, res)
    end
    push!(env, sentinel)
    return env
end

"""
Calculate the expectation value <mps|mpo|mps> assuming `mps` is normalized
"""
function expectation(mps::MatrixProductState, mpo::MatrixProductOperator)::AbstractFloat
    return get_env(mps, mpo)[end-1][1][1][1]
end

"""
Sweep `mps` from left to right to find the ground state of the Hamiltonian `mpo`.
The `mps` should be right-canonical
The energies calculatd during the sweep are appended to `energies`
and `env` is passed in as the environment from the right side
"""
function optimize_oneround!(mps::MatrixProductState, mpo::MatrixProductOperator, energies::Array, env::Array)
    sentinel = ones(1, 1, 1)
    # left-side environment and right-side environment. Corresponds to $L$ and $R$ in the reference literature.
    l_env = [sentinel]
    r_env = env

    for cur_idx in 1:length(mps)
        l, mo, r = l_env[cur_idx], mpo.mp[cur_idx], r_env[end-cur_idx-1]

        orig_shape = size(mps.mp[cur_idx])
        # don't optimize on the first site because it has probably been optmized during the previous sweep
        if cur_idx != 1
            # find the lowest eigenvalue of the local Hamitonian operator.
            # The computation can be accelerated by using a "sparse" representation of the operator
            # because some eigensolvers (such as `lobpcg` in this case) only need to know what would happen
            # if the operator(matrix) is applied on a certain state(vector).
            # See [power method](https://en.wikipedia.org/wiki/Power_iteration) as an example
            function h_map_f(ms)
                ms = reshape(ms, orig_shape)
                """
                Graphical notation: (* for MPS and # for MPO)
                *--a-- --h--*
                |     |     |
                |     f     |
                |     |     |
                #--b--#--g--#
                |     |     |
                |     d     |
                |     |     |
                *--c--*--e--*
                """
                @tensor res[a, f, h] := l[a, b, c] * ms[c, d, e] * mo[b, f, d, g] * r[h, g, e]
                return res
            end
            # some boilerplate to call the eigensolver
            new_dim = prod(orig_shape)
            h_map = LinearMap(h_map_f, new_dim, new_dim, issymmetric=true)
            # find the eigenvalue and eigenstate
            lobpcg_res = lobpcg(h_map, false, reshape(mps.mp[cur_idx], new_dim), 1)
            # update variables based on the result
            new_energy = lobpcg_res.λ[1]
            push!(energies, new_energy)
            mps.mp[cur_idx] = reshape(lobpcg_res.X, orig_shape)
        end

        # move the canonical center to the right side
        if cur_idx != length(mps)
            push_center!(mps, cur_idx)
        end
        # update the left-side environment
        push!(l_env, update_env(l, mps.mp[cur_idx], mpo.mp[cur_idx]))
    end

    push!(l_env, sentinel)

    return l_env
end

"""
Find the ground state of the Hamiltonian `mpo`
"""
function search_groundstate(mpo::MatrixProductOperator, bond_dimension::Int=20)
    # random initial guess
    mps = MatrixProductState(mpo, bond_dimension)
    l_env = get_env(mps, mpo)
    energies = [expectation(mps, mpo)]
    while true
        # do a right to left and then left to right sweep
        for i in 1:2
            # the mps is left canonical, so flip over to make it right canonical
            # so that we can achieve sweeping from right to left by sweeping from left to right
            mps, mpo = reverse.([mps, mpo])
            l_env = optimize_oneround!(mps, mpo, energies, l_env)
        end
        # consider it as converged if the lastest energies do not vary much
        if 10 < length(energies) && isapprox(std(energies[end-10:end]), 0, atol=1e-8)
            break
        end
    end
    return mps, energies
end

end


int_rules_1_1_1_1 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 1.1.1.1*(a+b*x)^m =#
    @apply_utils Antiderivative(1 / ~x, ~x) => log(~x)
    @apply_utils Antiderivative((~x) ^ ~(m'), ~x) => (~x) ^ (~m + 1) / (~m + 1) <-- FreeQ(~m, ~x) && NeQ(~m, -1)
    @apply_utils Antiderivative(1 / (~a + ~(b') * ~x), ~x) => log(RemoveContent(~a + ~b * ~x, ~x)) / ~b <-- FreeQ([~a, ~b], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * ~x) ^ ~m, ~x) => (~a + ~b * ~x) ^ (~m + 1) / (~b * (~m + 1)) <-- FreeQ([~a, ~b, ~m], ~x) && NeQ(~m, -1)
    @apply_utils Antiderivative((~(a') + ~(b') * ~u) ^ ~m, ~x) => (1 / Coefficient(~u, ~x, 1)) * Subst(Antiderivative((~a + ~b * ~x) ^ ~m, ~x), ~x, ~u) <-- FreeQ([~a, ~b, ~m], ~x) && (LinearQ(~u, ~x) && NeQ(~u, ~x))
end

src/rules/1 Algebraic functions/1.1 Binomial products/1.1.1 Linear/1.1.1.1 (a+b x)^m.jl


tol = 2e-2
tolP = 1e-3
maxnit = 2000

@testset "String with ODE" begin

preconI = SaddleSearch.localPrecon(precon = [I], precon_prep! = (P, x) -> P)

heading1("TEST: String type methods with ode12r solver")

heading2("Muller potential")
V = MullerPotential()
x = ic_path(V, :near, 15)
E, dE = objective(V)
precon = x-> hessianprecond(V, x)

path = StaticString(0.0009, tol, maxnit, preconI, serial(), false,
                        verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODEString(reltol=0.1, tol = tol, maxnit = maxnit,
                        precon_scheme = preconI,
                        path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = SaddleSearch.run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = StaticNEB(0.0007, 0.0002, 3, tol, maxnit, preconI, serial(),
                        false, verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODENEB(reltol=1e-2, k=0.0002, interp=3, tol = tol,
                        maxnit = maxnit,
                        precon_scheme = preconI, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

heading2("Double well potential")
c = 16.0
V = DoubleWell(diagm([.5, c*c*(0.02^(c-1)), c]))
x = ic_path(V, :near, 11)
E, dE = objective(V)

P = copy(V.A); P[1] = 1.0
precon = x->[P]

path = StaticString(1./c, tol, maxnit, preconI, serial(), false,
                        verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODEString(reltol=1e-3, tol = tol, maxnit = maxnit,
                        precon_scheme = preconI, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = SaddleSearch.run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = StaticNEB(0.07, 0.1, 3, tol, maxnit, preconI, serial(),
                        false, verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODENEB(reltol=1e-1, k=0.1, interp=3, tol = tol,
                        maxnit = maxnit,
                        precon_scheme = preconI, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

preconP = SaddleSearch.localPrecon(precon = [P], precon_prep! = (P, x) -> precon(x))

path = StaticString(0.25, tolP, maxnit, preconP, serial(), false,
                        verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODEString(reltol=0.1, tol = tolP, maxnit = maxnit,
                        precon_scheme = preconP, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = SaddleSearch.run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = StaticNEB(0.6, 0.00001, 1, tolP, maxnit, preconP, serial(),
                        false, verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODENEB(reltol=0.1, k=0.00001, interp=1, tol = tolP,
                        maxnit = maxnit,
                        precon_scheme = preconP, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

heading2("Vacancy migration potential")
V = LJVacancy2D(R = 3.1)
x = ic_path(V, :min, 9)
E, dE = objective(V)

precon = x->[copy(precond(V, xn)) for xn in x]

path = StaticString(0.001, tol, maxnit, preconI, serial(), false,
                        verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODEString(reltol=1e-2, threshold=1e-5, tol = tol,
                        maxnit = maxnit,
                        precon_scheme = preconI,
                        path_traverse = serial(),
                        verbose = verbose)   # allow failure on 600 iterations
PATHx, PATHlog, _ = SaddleSearch.run!(path, E, dE, Path(x))
@show PATHlog[:maxres][end]
println("[allowed test failure: target is $(path.tol)]")
# @test PATHlog[:maxres][end] <= path.tol

path = StaticNEB(0.001, 0.01, 3, tol, maxnit, preconI, serial(),
                        false, verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODENEB(reltol=1e-2, k=0.00001, interp=3, tol = tol,
                        maxnit = maxnit,
                        precon_scheme = preconI, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

preconP = SaddleSearch.localPrecon(precon = precon(x),
            precon_prep! = (P, x) -> precon(x))

path = StaticString(1.55, tolP, maxnit, preconP, serial(), false,
                        verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODEString(reltol=1e-1, tol = tolP, maxnit = maxnit,
                        precon_scheme = preconP,
                        path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = SaddleSearch.run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = StaticNEB(1.0, 0.001, 3, tolP, maxnit, preconP, serial(),
                        false, verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

path = ODENEB(reltol=1e-2, k=0.001, interp=3, tol = tolP,
                        maxnit = maxnit,
                        precon_scheme = preconP, path_traverse = serial(),
                        verbose = verbose)
PATHx, PATHlog, _ = run!(path, E, dE, Path(x))
@test PATHlog[:maxres][end] <= path.tol

end


function pooled_pava_isotonic_regression!(y::Vector{Float64}, weights::Vector{Float64})

    n = length(y)
    if n <= 1
        return y
    end
    if n != length(weights)
        throw(DimensionMismatch("Lengths of values and weights mismatch"))
    end

    j = 1
    S = Dict(0 => 0, 1 => 1)
    ydash = Dict(1 => y[1])
    wdash = Dict(1 => weights[1])
    @inbounds begin
        for i in 2 : n
            j = j + 1
            ydash[j] = y[i]
            wdash[j] = weights[i]
            while j > 1 && ydash[j] < ydash[j-1]
                ydash[j-1] = (wdash[j] * ydash[j] + wdash[j-1] * ydash[j-1]) /
                   (wdash[j] + wdash[j-1])
                wdash[j-1] = wdash[j] + wdash[j-1]
                j = j-1
            end
            S[j] = i
        end
        for k in 1 : j
            for l in S[k-1] + 1 : S[k]
                y[l] = ydash[k]
            end
        end
    end
    return y
end

pooled_pava_isotonic_regression!(y::Vector{Float64}) = pooled_pava_isotonic_regression!(y, ones(size(y, 1)))

# non-mutating versions
pooled_pava_isotonic_regression(y::Vector{Float64}, weights::Vector{Float64}) = pooled_pava_isotonic_regression!(copy(y), weights)
pooled_pava_isotonic_regression(y::Vector{Float64}) = pooled_pava_isotonic_regression(y, ones(size(y,1)))