export angular_distance, azimuth, angular_step, surface_distance, forward
export inverse, get_dist, get_azi, get_baz, get_loc

"""
angular_distance, azimuth, angular_step, surface_distance, forward and inverse
functions were written by Andy Nowacki in his (unregistered) Geodesics.jl
package https://github.com/anowacki/Geodesics.jl
"""

"Earth ellipsoid semi-axes in WGS84"
const EARTH_R_MAJOR_WGS84 = 6378137.0000
const EARTH_R_MINOR_WGS84 = 6356752.3142
"Flattening of the Earth in WGS84"
const F_WGS84 = (EARTH_R_MAJOR_WGS84 - EARTH_R_MINOR_WGS84)/EARTH_R_MAJOR_WGS84

"""
    angular_distance(lon0, lat0, lon1, lat1, degrees=true; f=F_WGS84) -> Δ
Return the angular distance between points (`lon0`,`lat0`) and (`lon1`,`lat1`)
on a flattened sphere.  The default flattening is $F_WGS84.  By default,
input and output are in degrees, but specify `degrees` as `false` to use radians.
Use `f=0` to perform computations on a sphere.
"""
function angular_distance(lon0, lat0, lon1, lat1, degrees::Bool=true; f=F_WGS84)
    if degrees
        lon0, lat0, lon1, lat1 = deg2rad(lon0), deg2rad(lat0), deg2rad(lon1), deg2rad(lat1)
    end
    gcarc, az, baz = inverse(lon0, lat0, lon1, lat1, 1.0, f)
    degrees ? rad2deg(gcarc) : gcarc
end

"""
    azimuth(lon0, lat0, lon1, lat1, degrees=true, f=F_WGS84)
Return the azimuth from point (`lon0`,`lat0`) to point (`lon1`,`lat1`) on a
flattened sphere.  The default flattening is $F_WGS84.  By default,
input and output are in degrees, but specify `degrees` as `false` to use radians.
Use `f=0` to perform computations on a sphere.
"""
function azimuth(lon0::T1, lat0::T2, lon1::T3, lat1::T4, degrees::Bool=true; f=F_WGS84) where {T1,T2,T3,T4}
    # Promote to Float64 at least to avoid error when one of the points is on the pole
    T = promote_type(Float64, T1, T2, T3, T4)
    lon0, lat0, lon1, lat1 = T(lon0), T(lat0), T(lon1), T(lat1)
    if degrees
        lon0, lat0, lon1, lat1 = deg2rad(lon0), deg2rad(lat0), deg2rad(lon1), deg2rad(lat1)
    end
    gcarc, az, baz = inverse(lon0, lat0, lon1, lat1, 1.0, f)
    degrees ? rad2deg(az) : az
end

"""
    step(lon, lat, azimuth, distance, degrees=false; f=F_WGS84) -> lon′, lat′, backazimuth
Return the longitude `lon′`, latitude `lat′` and backazimuth `baz` reached by
travelling an angular `distance` along `azimuth` from the starting point at
(`lon`,`lat`).  The default flattening is $F_WGS84.  By default,
input and output are in degrees, but specify `degrees` as `false` to use radians.
Use `f=0` to perform computations on a sphere.
"""
function angular_step(lon, lat, azimuth, distance, degrees::Bool=true; f=F_WGS84)
    if degrees
        lon, lat, azimuth, distance = deg2rad(lon), deg2rad(lat), deg2rad(azimuth), deg2rad(distance)
    end
    lon′, lat′, baz = forward(lon, lat, azimuth, distance, 1.0, f)
    if degrees
        rad2deg(lon′), rad2deg(lat′), rad2deg(baz)
    else
        lon′, lat′, baz
    end
end

"""
    surface_distance(lon0, lat0, lon1, lat1, a, degrees::Bool=false; f=F_WGS84)
Return the physical distance between points (`lon0`,`lat0`) and (`lon1`,`lat1`) on
the flattened sphere with flattening `f` and semimajor radius `a`.  Distance is given
in the same units as `a`.  By default, input angles are in degrees, but specify `degrees`
as `false` to use radians.  Use `f=0` to perform computations on a sphere.
"""
function surface_distance(lon0, lat0, lon1, lat1, a, degrees::Bool=true; f=F_WGS84)
    if degrees
        lon0, lat0, lon1, lat1 = deg2rad(lon0), deg2rad(lat0), deg2rad(lon1), deg2rad(lat1)
    end
    distance, az, baz = inverse(lon0, lat0, lon1, lat1, a, f)
    distance
end

"""
    forward(lon, lat, azimuth, distance, a, f) -> lon′, lat′, backazimuth
Return the longitude `lon′` and latitude `lat′` and `backazimuth` of a projected
point, reached by travelling along an `azimuth` for `distance` from an original
point at (`lon`, `lat`).  Specify the spheroid with the semimajor radius `a` and
flattening `f`.
Coordinates and azimuth are in radians.
Calculations use Vincenty's forward formula [1].
#### References
1. Vincenty, T. (1975). "Direct and Inverse Solutions of Geodesics on the Ellipsoid
   with application of nested equations" (PDF). Survey Review. XXIII (176): 88–93.
   doi:10.1179/sre.1975.23.176.88
"""
function forward(lon, lat, azimuth, distance, a, f)::Tuple{Float64,Float64,Float64}
    abs(lat <= π/2) || throw(ArgumentError("Latitude ($lat) must be in range [-π/2, π/2]"))
    a > 0 || throw(ArgumentError("Semimajor axis ($a) must be positive"))
    abs(f) < 1 || throw(ArgumentError("Magnitude of flattening ($f) must be less than 1"))
    # Calculations are done with Float64s internally as the tolerances are hard-wired
    lambda1, phi1, alpha12, s = Float64(lon), Float64(lat), Float64(azimuth), Float64(distance)
    a, f = Float64(a), Float64(f)
    alpha12 = mod(alpha12, 2π)
    b = a*(1 - f)

    TanU1 = (1 - f)*tan(phi1)
    U1 = atan(TanU1)
    sigma1 = atan( TanU1, cos(alpha12) )
    Sinalpha = cos(U1)*sin(alpha12)
    cosalpha_sq = 1.0 - Sinalpha*Sinalpha

    u2 = cosalpha_sq*(a*a - b*b )/(b*b)
    A = 1.0 + (u2/16384)*(4096 + u2*(-768 + u2*(320 - 175*u2)))
    B = (u2/1024)*(256 + u2*(-128 + u2*(74 - 47*u2)))

    # Starting with the approximation
    sigma = (s/(b*A))

    # Not moving anywhere. We can return the location that was passed in.
    if sigma == 0
        return phi1, lambda1, alpha12
    end

    last_sigma = 2*sigma + 2 # something impossible

    # Iterate the following three equations
    # until there is no significant change in sigma
    # two_sigma_m , delta_sigma
    while abs((last_sigma - sigma)/sigma) > 1.0e-9
        global two_sigma_m = 2*sigma1 + sigma
        delta_sigma = B*sin(sigma)*(cos(two_sigma_m) + (B/4)*(cos(sigma)*(-1 + 2*cos(two_sigma_m)^2 - (B/6)*cos(two_sigma_m)*(-3 + 4*sin(sigma)^2)*(-3 + 4*cos(two_sigma_m)^2 ))))
        last_sigma = sigma
        sigma = (s/(b*A)) + delta_sigma
    end

    phi2 = atan((sin(U1)*cos(sigma) + cos(U1)*sin(sigma)*cos(alpha12)),
        ((1-f)*sqrt(Sinalpha^2 + (sin(U1)*sin(sigma) - cos(U1)*cos(sigma)*cos(alpha12))^2)))

    lambda = atan((sin(sigma)*sin(alpha12)),
        (cos(U1)*cos(sigma) - sin(U1)*sin(sigma)*cos(alpha12)))

    C = (f/16)*cosalpha_sq*(4 + f*(4 - 3*cosalpha_sq))

    omega = lambda - (1-C)*f*Sinalpha*(sigma + C*sin(sigma)*(
        cos(two_sigma_m) + C*cos(sigma)*(-1 + 2*cos(two_sigma_m)^2)))

    lambda2 = lambda1 + omega

    alpha21 = atan(Sinalpha, (-sin(U1)*sin(sigma) + cos(U1)*cos(sigma)*cos(alpha12)))
    alpha21 = mod(alpha21 + π, 2π)

    return lambda2, phi2, alpha21
end

"""
    inverse(lon1, lat1, lon2, lat2, a, f) -> distance, azimuth, backazimuth
Return the `distance`, `azimuth` and `backazimuth` between two points with longitudes
`lon1` and `lon2`, and latitudes `lat1` and `lat2`.  Specify the spheroid with the
semimajor radius `a` and flattening `f`.
Coordinates and angles are in radians, whilst `distance` is in the same units as `a`.
Calculations use Vincenty's inverse formula [1].
#### References
1. Vincenty, T. (1975). "Direct and Inverse Solutions of Geodesics on the Ellipsoid
   with application of nested equations" (PDF). Survey Review. XXIII (176): 88–93.
   doi:10.1179/sre.1975.23.176.88
"""
function inverse(lon1, lat1, lon2, lat2, a, f)::Tuple{Float64,Float64,Float64}
    for lat in (lat1, lat2)
        abs(lat <= π/2) || throw(ArgumentError("Latitude ($lat) must be in range [-π/2, π/2]"))
    end
    a > 0 || throw(ArgumentError("Semimajor axis ($a) must be positive"))
    abs(f) < 1 || throw(ArgumentError("Magnitude of flattening ($f) must be less than 1"))
    lambda1, phi1, lambda2, phi2 = Float64(lon1), Float64(lat1), Float64(lon2), Float64(lat2)
    a, f = Float64(a), Float64(f)
    tol = 1.0e-8
    if (abs(phi2 - phi1) < tol) && (abs(lambda2 - lambda1) < tol)
        return 0.0, 0.0, 0.0
    end

    b = a*(1 - f)

    TanU1 = (1 - f)*tan(phi1)
    TanU2 = (1 - f)*tan(phi2)

    U1 = atan(TanU1)
    U2 = atan(TanU2)

    lambda = lambda2 - lambda1
    last_lambda = -4000000.0 # an impossibe value
    omega = lambda

    # Iterate the following equations until there is no significant change in lambda
    alpha, sigma, Sin_sigma, Cos2sigma_m, Cos_sigma, sqr_sin_sigma =
        -999999., -999999., -999999., -999999., -999999., -999999.
    while ((last_lambda < -3000000.0) || (lambda != 0)) &&
            (abs((last_lambda - lambda)/lambda) > 1.0e-9)
        sqr_sin_sigma = (cos(U2)*sin(lambda))^2 +
                         ((cos(U1)*sin(U2) - sin(U1)*cos(U2)*cos(lambda)))^2
        Sin_sigma = sqrt(sqr_sin_sigma)
        Cos_sigma = sin(U1)*sin(U2) + cos(U1)*cos(U2)*cos(lambda)
        sigma = atan(Sin_sigma, Cos_sigma)

        Sin_alpha = cos(U1)*cos(U2)*sin(lambda)/sin(sigma)

        if Sin_alpha >= 1
            Sin_alpha = 1.0
        elseif Sin_alpha <= -1
            Sin_alpha = -1.0
        end

        alpha = asin(Sin_alpha)
        Cos2sigma_m = cos(sigma) - 2*sin(U1)*sin(U2)/cos(alpha)^2
        C = (f/16)*cos(alpha)^2*(4 + f*(4 - 3*cos(alpha)^2))
        last_lambda = lambda
        lambda = omega + (1 - C)*f*sin(alpha)*(sigma +
            C*sin(sigma)*(Cos2sigma_m + C*cos(sigma)*(-1 + 2*Cos2sigma_m^2)))
    end

    u2 = cos(alpha)^2*(a*a - b*b)/(b*b)
    A = 1 + (u2/16384)*(4096 + u2*(-768 + u2*(320 - 175*u2)))
    B = (u2/1024)*(256 + u2*(-128 + u2*(74 - 47*u2)))
    delta_sigma = B*Sin_sigma*(Cos2sigma_m + (B/4)*(
        Cos_sigma*(-1 + 2*Cos2sigma_m^2) -
        (B/6)*Cos2sigma_m*(-3 + 4*sqr_sin_sigma)*(-3 + 4*Cos2sigma_m^2)))
    s = b*A*(sigma - delta_sigma)

    alpha12 = atan((cos(U2)*sin(lambda)), ( cos(U1)*sin(U2) - sin(U1)*cos(U2)*cos(lambda)))
    alpha21 = atan((cos(U1)*sin(lambda)), (-sin(U1)*cos(U2) + cos(U1)*sin(U2)*cos(lambda)))

    alpha12 = mod(alpha12, 2π)
    alpha21 = mod(alpha21 + π, 2π)

    return s, alpha12, alpha21
end

"""
  get_dist(loc1,loc2)

Get distance bewteen source `loc1` and receiver `loc2` in km.
If either location is empty, return zero.
# Arguments
- `loc1::GeoLoc`: Location of source.
- `loc2::GeoLoc`: Location of receiver.
"""
function get_dist(loc1::GeoLoc, loc2::GeoLoc)
    if isempty(loc1) | isempty(loc2)
        dist =  zero(Float64)
    else
        dist =  surface_distance(loc1.lon,loc1.lat,loc2.lon,loc2.lat,
                               EARTH_R_MAJOR_WGS84/1000)
    end
    return dist
end

"""
  get_dist(FFT1,FFT2)

Get distance bewteen source `FFT1` and receiver `FFT2` in km.
If either location is empty, return zero.
# Arguments
- `FFT1::FFTData`: FFTData with source location.
- `FFT2::FFTData`: FFTData with receiver location.
"""
function get_dist(FFT1::FFTData, FFT2::FFTData)
    return get_dist(FFT1.loc,FFT2.loc)
end

"""
  get_azi(loc1,loc2)

Get azimuth from source `loc1` to receiver `loc2` in degrees.
If either location is empty, return zero.
# Arguments
- `loc1::GeoLoc`: Location of source.
- `loc2::GeoLoc`: Location of receiver.
"""
function get_azi(loc1::GeoLoc, loc2::GeoLoc)
    if isempty(loc1) | isempty(loc2)
        azi =  zero(Float64)
    else
        azi = azimuth(loc1.lon,loc1.lat,loc2.lon,loc2.lat)
    end
    return azi
end

"""
  get_azi(FFT1,FFT2)

Get azimuth from source `FFT1` to receiver `FFT2` in degrees.
If either location is empty, return zero.
# Arguments
- `FFT1::FFTData`: FFTData with source location.
- `FFT2::FFTData`: FFTData with receiver location.
"""
function get_azi(FFT1::FFTData, FFT2::FFTData)
    return get_azi(FFT1.loc,FFT2.loc)
end

"""
  get_baz(loc1,loc2)

Get back azimuth from source `loc1` to receiver `loc2` in degrees.
If either location is empty, return zero.
# Arguments
- `loc1::GeoLoc`: Location of source.
- `loc2::GeoLoc`: Location of receiver.
"""
function get_baz(loc1::GeoLoc, loc2::GeoLoc)
    if isempty(loc1) | isempty(loc2)
        baz =  zero(Float64)
    else
        baz = inverse(deg2rad.((loc1.lon,loc1.lat,loc2.lon,loc2.lat))...,
                          EARTH_R_MAJOR_WGS84, F_WGS84)[3]
        baz = rad2deg(baz)
    end
    return baz
end

"""
  get_baz(FFT1,FFT2)

Get back azimuth from source `FFT1` to receiver `FFT2` in degrees.
If either location is empty, return zero.
# Arguments
- `FFT1::FFTData`: FFTData with source location.
- `FFT2::FFTData`: FFTData with receiver location.
"""
function get_baz(FFT1::FFTData, FFT2::FFTData)
    return get_baz(FFT1.loc,FFT2.loc)
end

"""
  get_loc(loc1,azi, dist)

Get location of receiver given location of source and dist + azi to receiver.
If location is empty, return empty location.
# Arguments
- `loc1::GeoLoc`: Location of source.
- `azi::Float64`: Azimuth from source `loc1` to receiver.
- `dist::Float64`: Distance from source `loc1` to receiver in km.
"""
function get_loc(loc1::GeoLoc, azi::Float64, dist::Float64)
    if isempty(loc1)
        loc2 = GeoLoc()
    else
        lon2,lat2,baz2 = forward(deg2rad.((loc1.lon, loc1.lat,azi))...,
                         dist*1000, EARTH_R_MAJOR_WGS84, F_WGS84)
        lon2, lat2, baz2 = rad2deg.((lon2, lat2, baz2))
        loc2 = GeoLoc("",lat2,lon2, zero(Float64), zero(Float64), zero(Float64),
                      loc1.inc)
    end
    return loc2
end

"""
  get_loc(C)

Get location of receiver and source from CorrData `C`.
If either location is empty, return empty locations.
# Arguments
- `C::CorrData`: CorrData with source location and distance + azimuth to receiver.
"""
function get_loc(C::CorrData)
    if isempty(C.loc)
        loc2 = GeoLoc()
    else
        loc2 = get_loc(C.loc, C.azi, C.dist)
    end

    return C.loc, loc2
end


add_noise(X, d, N, noiseVar = 1e-10) = X .+ (noiseVar .* rand(d, N))

# Return the kNN nearest point.
function knn_nearest_point(X::Array{Float64, 2}, knnTree, point, k)
  indices, distances = nearest(knnTree, point, k+1)
  sorted = sortperm(distances)
  X[:,indices[sorted[end]]]
end

# Return the kNN nearest point.
function knn_nearest_distance(X::Array{Float64, 2}, knnTree, point, k)
  indices, distances = nearest(knnTree, point, k+1)
  sorted = sortperm(distances)
  distances[sorted[end]]
end


"""
    minimum_weight_perfect_matching(g, w::Dict{Edge,Real})
    minimum_weight_perfect_matching(g, w::Dict{Edge,Real}, cutoff)

Given a graph `g` and an edgemap `w` containing weights associated to edges,
returns a matching with the mimimum total weight among the ones containing
exactly `nv(g)/2` edges.

Edges in `g` not present in `w` will not be considered for the matching.

This function relies on the BlossomV.jl package, a julia wrapper
around Kolmogorov's BlossomV algorithm.

Eventually a `cutoff` argument can be given, to the reduce computational time
excluding edges with weights higher than the cutoff.

The returned object is of type `MatchingResult`.

In case of error try to change the optional argument `tmaxscale` (default is `tmaxscale=10`).
"""
function minimum_weight_perfect_matching end

function minimum_weight_perfect_matching(g::Graph, w::Dict{E,U}, cutoff, kws...) where {U<:Real, E<:Edge}
    wnew = Dict{E, U}()
    for (e, c) in w
        if c <= cutoff
            wnew[e] = c
        end
    end
    return minimum_weight_perfect_matching(g, wnew; kws...)
end

function minimum_weight_perfect_matching(g::Graph, w::Dict{E,U}; tmaxscale=10.) where {U<:AbstractFloat, E<:Edge}
    wnew = Dict{E, Int32}()
    cmax = maximum(values(w))
    cmin = minimum(values(w))
    tmax = typemax(Int32)  / tmaxscale # /10 is kinda arbitrary,
                                # hopefully high enough to not occur in overflow problems
    for (e, c) in w
        wnew[e] = round(Int32, (c-cmin) / (cmax-cmin) * tmax)
    end
    match = minimum_weight_perfect_matching(g, wnew)
    weight = zero(U)
    for i=1:nv(g)
        j = match.mate[i]
        if j > i
            weight += w[E(i,j)]
        end
    end
    return MatchingResult(weight, match.mate)
end

function minimum_weight_perfect_matching(g::Graph, w::Dict{E,U}) where {U<:Integer, E<:Edge}
    m = BlossomV.Matching(nv(g))
    for (e, c) in w
        BlossomV.add_edge(m, src(e)-1, dst(e)-1, c)
    end
    BlossomV.solve(m)

    mate = fill(-1, nv(g))
    totweight = zero(U)
    for i=1:nv(g)
        j = BlossomV.get_match(m, i-1) + 1
        mate[i] = j <= 0 ? -1 : j
        if i < j
            totweight += w[Edge(i,j)]
        end
    end
    return MatchingResult(totweight, mate)
end


module ElasticFDA

## Requirements
using Interpolations
using Dierckx
using NLopt
using ProgressMeter
using Distributions
using Printf
using Distributed
using Statistics
using LinearAlgebra
using RecipesBase
using Cxx
using Libdl

import Base. summary

export
    smooth_data!,
    smooth_data,
    f_to_srsf,
    srsf_to_f,
    optimum_reparam,
    warp_q_gamma,
    warp_f_gamma,
    vert_fPCA,
    horiz_fPCA,
    gauss_model,
    group_warping_bayes,
    pair_warping_bayes,
    pair_warping_expomap,
    elastic_distance,
    elastic_regression,
    elastic_logistic,
    elastic_prediction,
    srsf_align,
    align_fPCA,
    trapz,
    gradient,
    rgam,
    curve_karcher_mean,
    curve_srvf_align,
    curve_principal_directions,
    curve_karcher_cov,
    sample_shapes,
    curve_to_q,
    q_to_curve,
    calc_shape_dist,
    curve_pair_align,
    curve_geodesic,
    resamplecurve,
    bs,
    project_curve,
    inverse_exp_coord,
    inverse_exp,
    project_tangent

# load fdasrsf library
if (Sys.isunix())
    path = "../deps/src/fdasrsf/fdasrsf"
else
    path = "..\\deps\\fdasrsf"
end
const libfdasrsf = joinpath(dirname(@__FILE__), path)

# load gropt library
if (Sys.isunix())
    path1 = "../deps/src/gropt/gropt"
else
    path1 = "..\\deps\\gropt"
end
const libgropt = joinpath(dirname(@__FILE__), path1)

# load fdaqmap library
if (Sys.isunix())
    path2 = "../deps/src/fdaqmap/fdaqmap"
else
    path2 = "..\\deps\\fdaqmap"
end
const libfdaqmap = joinpath(dirname(@__FILE__), path2)

# Ensure library is available.
function __init__()
    if (dlopen_e(libfdasrsf) == C_NULL)
        error("libfdasrsf not properly installed. Run build(\"ElasticFDA\")")
    end

   if (dlopen_e(libgropt) == C_NULL)
        error("libgropt not properly installed. Run build(\"ElasticFDA\")")
    end

   if (dlopen_e(libfdaqmap) == C_NULL)
       error("libfdaqmap not properly installed. Run build(\"ElasticFDA\")")
   end

end

### source files
include("types.jl")
include("geometry.jl")
include("summary.jl")
include("plots.jl")
include("curve_funcs.jl")
include("curve_stats.jl")
include("srsf_funcs.jl")
include("misc_funcs.jl")
include("fPCA.jl")
include("bspline.jl")
include("dp_bayes.jl")
include("mcmc_align.jl")
include("gauss_model.jl")
include("group_warping_bayes.jl")
include("pair_warping_bayes.jl")
include("regression.jl")
include("regression_funcs.jl")
include("time_warping.jl")
include("pair_warping_expomap.jl")
include("image_funcs.jl")
include("image_stats.jl")

end # module


mutable struct Bits{N}
  offset::Int
  state::State
end

const Bit = Bits{1}

const states = WeakKeyDict()

function Bits{N}(vs::Integer...) where N
  b = Bits{N}(0, State(N, vs))
  states[b.state] = WeakKeyDict(b=>nothing)
  return b
end

Bits{N}() where N = Bits{N}(0)

bits(b::Bits) = b

nbits(::Bits{N}) where N = N

function Base.show(io::IO, b::Bits)
  println(io, "$(nbits(b)) bits with offset $(b.offset)")
  show(io, b.state)
end

function combine(a::State, b::State)
  a === b && return a
  of = nbits(a)
  a.ψ = a.ψ ⊗ b.ψ
  for buf in keys(states[b])
    buf.state = a
    buf.offset += of
    states[a][buf] = nothing
  end
  return a
end

combine(b::Bits) = b.state
combine(bs::Bits...) = reduce(combine, map(x -> x.state, bs))

bitindices(b::Bits{N}) where N = ntuple(i -> b.offset+i, Val(N))

tcat(x::Tuple) = x
tcat(x::Tuple, y::Tuple, ys::Tuple...) = tcat((x..., y...), ys...)

function apply!(U, bs::Bits...)
  s = combine(bs...)
  is = tcat(bitindices.(bs)...)
  apply!(U, s, is...)
  return
end

apply!(U, bs...) = apply!(U, bits.(bs)...)

measure(b::Bits) = measure!(b.state, bitindices(b)...)

measure(x) = repr(typeof(x),measure(bits(x)))


pprint(it...;sep=" ",ed="\n")=length(it)>0 ? (@inbounds for i in 1:length(it)-1;print(it[i],sep);end;print(last(it),ed)) : print(ed)

function main()
    a,b,c=parse.(Int64, split(readline()))
    count = 0
    while true
        if count>32
            pprint(-1)
            break
        end
        if a%2+b%2+c%2>0
            pprint(count)
            break
        end
        a,b,c=b/2+c/2,a/2+c/2,b/2+c/2
        count+=1
    end
end
main()


using InteractiveDynamics, Observables, Test
using DynamicalSystems

ds = Systems.henonheiles()

potential(x, y) = 0.5(x^2 + y^2) + (x^2*y - (y^3)/3)
energy(x,y,px,py) = 0.5(px^2 + py^2) + potential(x,y)
E = energy(get_state(ds)...)
function complete(y, py, x)
    V = potential(x, y)
    Ky = 0.5*(py^2)
    Ky + V ≥ E && error("Point has more energy!")
    px = sqrt(2(E - V - Ky))
    ic = [x, y, px, py]
    return ic
end

chaotic = get_state(ds)
stable = [0., 0.1, 0.5, 0.]
plane = (1, 0.0)
tf = 1000.0
idxs = SVector{2, Int}(2, 4)
u0 = get_state(ds); D = 4
Ttr = 200.0
rootkw = (xrtol = 1e-6, atol = 1e-6)

integ = integrator(ds, u0)
planecrossing = PlaneCrossing(plane, true)
f = (t) -> planecrossing(integ(t))
i = SVector{2, Int}(idxs)
data = poincaresos(integ, planecrossing, tf, Ttr, i, rootkw)
data0 = copy(data)

# Emulate clicking on the psos
click = Observable([0.0, 0.0])
Data = Observable(data)

on(click) do pos
    x, y = pos; z = plane[2]
    newstate = try
       complete(x, y, z)
    catch err
       @error "Could not get state, got error:" exception=err
       return
    end
    @assert length(newstate) == D

    reinit!(integ, newstate)

    Data[] = poincaresos(integ, planecrossing, tf, Ttr, i, rootkw)
    return data
end

# we click on the psos now
click[] = [0.2, 0.2]

@test Data[] != data0


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 04-2019
# Define operations for 1D spaces
#--------------------------------------------------------------------

using LinearAlgebra
import Base: maximum, minimum, range, size, ndims
export order, identity, incomingboundary, outgoingboundary

ndims(space::S)   where {S<:Cardinal{Tag, N, T}} where {Tag, N, T} = 1
minimum(space::S) where {S<:Cardinal{Tag, N, T}} where {Tag, N, T} = space.min
maximum(space::S) where {S<:Cardinal{Tag, N, T}} where {Tag, N, T} = space.max
order(space::S)   where {S<:Cardinal{Tag, N, T}} where {Tag, N, T} = N-1
size(space::S)    where {S<:Cardinal{Tag, N, T}} where {Tag, N, T} = N
range(space::S)   where {S<:Cardinal{Tag, N, T}} where {Tag, N, T} = (minimum(space), maximum(space))

function Base. identity(space::S)::Operator{S} where {S<:Cardinal{Tag, N, T}} where {Tag, N, T}
    return Operator(space, Diagonal(ones(T, size(space))))
end

function outgoingboundary(space::S)::Operator{S} where {S<:Cardinal{Tag, N, T}} where {Tag, N, T}
    return Operator(space, Diagonal([1, zeros(T, order(space))...]))
end

function incomingboundary(space::S)::Operator{S} where {S<:Cardinal{Tag, N, T}} where {Tag, N, T}
    return Operator(space, Diagonal([zeros(T, order(space))..., 1]))
end

function derivative(space::S)::Operator{S} where {S<:Cardinal{Tag, N, T}} where {Tag, N, T}
    return Operator(space, [derivative(space, i, j) for i in 1:size(space), j in 1:size(space)])
end

function integral(space::S)::Operator{S} where {S<:Cardinal{Tag, N, T}} where {Tag, N, T}
    return Operator(space, Diagonal([integral(space, i) for i in 1:size(space)]))
end

function Field(space::S, umap::Function)::Field{S} where {S<:Cardinal{Tag, N, T}} where {Tag, N, T}
    value = zeros(T, size(space))
    for index in CartesianIndices(value)
        value[index] = umap(collocation(space, index.I[1]))
    end
    return Field(space, value)
end


# SteinProbitRegression
#
# We assume we have the following model:
#
# y_i = \Psi (beta_1 x_1 + ... + beta_p x_p)
#
# where \Psi is assumed to be the noraml CDF (and a flat prior for beta).
# We sample this by introducing the auxiliary variables
# y_i = I{z_i > 0}. Hence we have
#
# beta | Z, y ~ N((X^t X)^-1 (X_t Z), (X^t X)^-1)
# Z | beta, y ~ TN(X beta, I, y)
#
# where TN(mu, sigma, t) is a truncated Gaussian
# with mean mu, variance sigma and truncated to be
# positive if t == 1 and negative if t == 0.

using Distributions: TruncatedNormal, Normal
import Distributions

type SteinProbitRegression <: SteinPosterior
    # the X data
    X::Array{Float64,2}
    # the 0/1 y values
    y::Array{Float64,1}
    # sigma2 is the variance of Z given beta
    sigma2::Float64
end

# Values used in Seth Tribble's thesis:
defaultsigma2 = 1.0

SteinProbitRegression(X, y) =
    SteinProbitRegression(X, y, defaultsigma2)

# Returns lower bound of support
function supportlowerbound(d::SteinProbitRegression, i::Int)
    -Inf
end

# Returns upper bound of support
function supportupperbound(d::SteinProbitRegression, i::Int)
    Inf
end

function rungibbs(d::SteinProbitRegression,
                  n::Int;
                  qmcvariates=[])
    X = d.X
    m, p = size(X)
    betas = Array(Float64, n+1, p)
    Zs = Array(Float64, n+1, m)
    gsn = Normal()
    # work with qmc variates or iid ones
    if size(qmcvariates,1) != 0
        @assert size(qmcvariates) == (n, p+m)
        randvariates = qmcvariates
    else
        randvariates = rand(n, p+m)
    end
    # initialize the samples (z_i = I{y > 0})
    Zs[1,:] = round(Int, d.y .> 0.0)
    betas[1,:] = 0.0
    # do the gibbs!
    for i in 1:n
        # sample the betas
        xtxinv = inv(X' * X)
        xtxinvchol = chol(xtxinv)
        xtz = X' * vec(Zs[i,:])
        beta_mean = vec(xtxinv * xtz)

        noise = zeros(p)
        for j in 1:p
            noise[j] = Distributions.quantile(gsn, randvariates[i,j])
        end
        betas[i+1,:] = beta_mean + vec(xtxinvchol * noise)
        # now sample the z's
        zmean = vec(X * vec(betas[i+1,:]))
        for j in 1:m
            if d.y[j] > 0.0
                g = TruncatedNormal(zmean[j], d.sigma2, 0.0, Inf)
            else
                g = TruncatedNormal(zmean[j], d.sigma2, -Inf, 0.0)
            end
            Zs[i+1,j] = Distributions.quantile(g, randvariates[i,p+j])
        end
    end
    # discard the first point
    betas[2:end,:]
end

# Returns the log prior log pi(beta)
function logprior(d::SteinProbitRegression,
                  beta::Array{Float64, 1})
    0.0
end

# Returns log pi(d.y[idx] | beta)
function loglikelihood(d::SteinProbitRegression,
                       beta::Array{Float64, 1};
                       idx=1:length(d.y))
    y = d.y[idx]
    X = d.X[idx,:]
    ypos = find(y .> 0.0)
    yneg = find(y .<= 0.0)
    u = vec(X * beta)
    gsn = Normal()

    logprobpos = Distributions.logcdf(gsn, u[ypos])
    logprobneg = log(1.0 - Distributions.cdf(gsn, u[yneg]))
    sum(logprobpos) + sum(logprobneg)
end

# Returns the log prior log pi(beta)
function gradlogprior(d::SteinProbitRegression,
                      beta::Array{Float64, 1})
    zeros(size(beta,1))
end

# Returns grad_{beta} log pi(d.y[idx] | beta)
function gradloglikelihood(d::SteinProbitRegression,
                           beta::Array{Float64,1};
                           idx=1:length(d.y))
    y = d.y[idx]
    X = d.X[idx,:]
    ypos = find(y .> 0.0)
    yneg = find(y .<= 0.0)
    u = vec(X * beta)
    gsn = Normal()

    posweights = Distributions.pdf(gsn, u[ypos]) ./ Distributions.cdf(gsn, u[ypos])
    negweights = -Distributions.pdf(gsn, u[yneg]) ./ (1.0 - Distributions.cdf(gsn, u[yneg]))

    xpos = vec(sum(broadcast(*, X[ypos,:], posweights), 1))
    xneg = vec(sum(broadcast(*, X[yneg,:], negweights), 1))

    xneg + xpos
end

function numdimensions(d::SteinProbitRegression)
    size(d.X, 2)
end


#sweep.jl
#Based on https://github.com/joshday/SweepOperator.jl
#Thanks to @joshday and @Hua-Zhou

function nsyrk!(alpha, A, C)
    q = size(C, 1)
    #p = size(A, 2)
    @simd for n ∈ axes(C, 1)
        @simd for m ∈ n:q
            @inbounds @simd for i ∈ axes(A, 2)
                C[n, m] += A[n, i] * A[m, i] * alpha
            end
        end
    end
    C
end
function sweep!(A::AbstractArray{T}, k::Integer, inv::Bool = false; syrkblas::Bool = false) where T
    sweepb!(Vector{T}(undef, size(A, 2)), A, k, inv; syrkblas = syrkblas)
end
function sweepb!(akk::AbstractArray{T, 1}, A::AbstractArray{T, 2}, k::Integer, inv::Bool = false; syrkblas::Bool = false) where T <: Number
    p = checksquare(A)
    #p == length(akk) || throw(DimensionError("incorrect buffer size"))
    @inbounds d = one(T) / A[k, k]
    @simd for j in 1:k
        @inbounds akk[j] = A[j, k]
    end
    @simd for j in (k+1):p
        @inbounds akk[j] = A[k, j]
    end
    #syrk!(uplo, trans, alpha, A, beta, C)
    #Rank-k update of the symmetric matrix C as alpha*A*transpose(A) + beta*C
    #or alpha*transpose(A)*A + beta*C according to trans.
    #Only the uplo triangle of C is used. Returns C.
    if syrkblas
        BLAS.syrk!('U', 'N', -d, akk, one(T), A)
    else
        nsyrk!(-d, akk, A)
    end
    rmul!(akk, d * (-one(T)) ^ inv)
    @simd for i in 1:(k-1)
        @inbounds A[i, k] = akk[i]
    end
    @simd for j in (k+1):p
        @inbounds A[k, j] = akk[j]
    end
    @inbounds A[k, k] = -d
    A
end
function sweep!(A::AbstractArray{T, 2}, ks::AbstractVector{I}, inv::Bool = false; syrkblas::Bool = false, logdet::Bool = false) where {T <: Number, I <: Integer}
    akk = Vector{T}(undef, size(A,2))
    sweepb!(akk, A, ks, inv; syrkblas = syrkblas, logdet = logdet)
end
function sweepb!(akk::AbstractArray{T, 1}, A::AbstractArray{T, 2}, ks::AbstractVector{I}, inv::Bool = false; syrkblas::Bool = false, logdet::Bool = false) where
        {T <: Number, I<:Integer}
    ld = NaN
    noerror = true
    if logdet
        ld = 0
        @inbounds for k in ks
            Akk = A[k,k]
            if Akk > 0
                ld += log(Akk)
            else
                noerror = false
                if Akk < 0
                    ld += log(- Akk * LDCORR)
                else
                    ld += LOGLDCORR
                end
            end
            sweepb!(akk, A, k, inv; syrkblas = syrkblas)
        end
    else
        for k in ks
            sweepb!(akk, A, k, inv; syrkblas = syrkblas)
        end
    end
    A, ld, noerror
end


"""
    cohenES(test::String = "",size::String = "")

Give the conventional effect size (small, medium, large) for the tests available in this package.
Available options for `test` are `p`, `t`, `r`, `anova`, `chisq`, and `f2`.
Available options for `size` are `small`, `medium`, and `large`.
"""
function cohenES(;test::String = "",size::String = "")

    if lower(test) in ("p","t","r","anova","chisq","f2") == false
        error("`test` must be one of `p`, `t`, `r`, `anova`, `chisq`,or `f2`")
    end

    testd = Dict("p"=>1,"t"=>2,"r"=>3,"anova"=>4,"chisq"=>5,"f2"=>6)

    if lower(size) in ("small","medium","large") == false
        error("`size` must be one of `small`, `medium`, or `large`")
    end

    sized = Dict("small" => 1, "medium" => 2, "large" => 3)

    effsize = [[.2 .5 .8],
        [.2 .5 .8],
        [.1 .3 .5],
        [.1 .25 .4],
        [.1 .3 .5],
        [.02 .15 .35]]

    return effsize[testd[test]][sized[size]]
end


# SPDX-License-Identifier: MIT

#======================================================================
Forward recursion
======================================================================#

function αrecursion(π::AbstractVector{SR},
                    Tᵀ::AbstractMatrix{SR},
                    lhs::AbstractMatrix{SR}) where SR <: Semiring
    S, N = length(π), size(lhs, 2)
    α = similar(lhs, SR, S, N)
    buffer = similar(lhs[:, 1])

    @views elmul!(α[:,1], π, lhs[:,1])
    @views for n in 2:N
        matmul!(buffer, Tᵀ, α[:,n-1])
        elmul!(α[:,n], buffer, lhs[:,n])
    end
    α
end

#======================================================================
Backward recursion
======================================================================#

function βrecursion(T::AbstractMatrix{SR},
                    lhs::AbstractMatrix{SR}) where SR <: Semiring
    S, N = size(T, 1), size(lhs, 2)
    β = fill!(similar(lhs, SR, S, N), one(SR))
    buffer = similar(lhs[:, 1], S)

    @views for n in N-1:-1:1
        elmul!(buffer, β[:,n+1], lhs[:,n+1])
        matmul!(β[:,n], T, buffer)
    end
    β
end

# Do the backward recursion and multiply with the α immediately.
# This avoids to allocate a second array for the β messages.
function βrecursion_mulα!(α::AbstractMatrix{SR}, T::AbstractMatrix{SR},
                    lhs::AbstractMatrix{SR}) where SR <: Semiring
    S, N = size(T, 1), size(lhs, 2)
    βₘ = fill!(similar(lhs[:, 1], S), one(SR))
    buffer = similar(lhs[:, 1], S)

    @views for n in N-1:-1:1
        elmul!(buffer, βₘ, lhs[:,n+1])
        matmul!(βₘ, T, buffer)
        elmul!(α[:,n], α[:,n], βₘ)
    end
    α
end

#======================================================================
Specialized algorithms
======================================================================#

function _expand(lhs::AbstractMatrix{T}, seqlength = size(lhs, 2)) where T
    S, N = size(lhs)
    retval = hcat(vcat(lhs, zeros(T, 1, N)), zeros(T, S+1))
    @views fill!(retval[1:end-1,seqlength+1:end], zero(T))
    @views fill!(retval[end,seqlength+1:end], one(T))
    retval
end

function _split(fsm, lhs)
	offset = 0
	retval = []
	for r in fsm.ranges
		nr = (r[1]:r[end]-1) .- offset
		push!(retval, lhs[nr, :])
		offset += 1
	end
	retval
end

function _drop_extradims(fsm, γ)
    [γ[r[1]:r[end]-1, 1:end-1] for r in fsm.ranges]
end

"""
    pdfposteriors(mfsm::MatrixFSM{SF} lhs) where SF <: Semifield
    pdfposteriors(union(mfsm1, mfsm2, ...), batch_lhs)

Calculate the conditional posterior of "assigning" the \$n\$th frame
to the \$i\$th pdf. The output is a tuple `γ, ttl` where `γ` is a matrix
(# pdf x # frames) and `ttl` is the total probability of the sequence.
This function can also be called in "batch-mode" by providing a union
of compiled fsm and a 3D tensor containing the per-pdf, per-frame
and per-batch values.
"""
pdfposteriors

function pdfposteriors(mfsm::MatrixFSM{SR},
                       in_lhs::AbstractMatrix{T}) where {SR<:LogSemifield,T}
    # Convert the likelihood matrix to the mfsm' semiring.
    lhs = copyto!(similar(in_lhs, SR), in_lhs)

    # Add one frame and one dimension for the final state.
    expanded_lhs = _expand(lhs)

    S = size(mfsm.C, 1)
    N = size(expanded_lhs, 2)

    # Get the per-state likelihood matrix.
    state_lhs = matmul!(similar(lhs, S, N), mfsm.C, expanded_lhs)

    α = αrecursion(mfsm.π, mfsm.Tᵀ, state_lhs)
    state_γ = βrecursion_mulα!(α, mfsm.T, state_lhs)

    # Transform the per-state γs to per-likelihoods γs.
    γ = matmul!(expanded_lhs, mfsm.Cᵀ, state_γ) # re-use `expanded_lhs` memory.

    # Re-normalize the γs.
    sums = sum(γ, dims = 1)
    eldiv!(γ, γ, sums)

    # Convert the result to the Real-semiring
    out = copyto!(similar(in_lhs), γ[1:end-1,1:end-1])
    out = map!(exp, out, out)

    out, convert(T, minimum(sums))
end

function pdfposteriors(mfsm::UnionMatrixFSM{SR},
                       in_lhs::AbstractArray{T,3},
                       seqlengths,
                      ) where {SR<:LogSemifield,T}

    # Convert the likelihood tensor to the mfsm's semiring.
    lhs_tensor = copyto!(similar(in_lhs, SR), in_lhs)

    # Reshape the 3D tensor to have a matrix of size BK x N where
    # B is the number of elements in the batch.
    lhs = vcat(
        map(t -> _expand(t...),
            zip(eachslice(lhs_tensor, dims = 3), seqlengths))...
    )

    S = size(mfsm.C, 1)       # number of states
    K = size(in_lhs, 1) + 1   # number of pdfs
    N = size(in_lhs, 2) + 1   # number of frames

    # Get the per-state likelihoods.
    state_lhs = matmul!(similar(lhs, SR, S, N), mfsm.C, lhs)

    α = αrecursion(mfsm.π, mfsm.Tᵀ, state_lhs)
    state_γ = βrecursion_mulα!(α, mfsm.T, state_lhs)

    # Transform the per-state γs to per-likelihoods γs.
    γ = matmul!(lhs, mfsm.Cᵀ, state_γ) # re-use `lhs` memory.
    γ = permutedims(reshape(γ, K, :, N), (1, 3, 2))

    # Re-normalize each element of the batch.
    sums = sum(γ, dims = 1)
    eldiv!(γ, γ, sums)

    # Convert the result to the Real-semiring
    out = copyto!(similar(in_lhs), γ[1:end-1,1:end-1,:])
    out = map!(exp, out, out)

    ttl = dropdims(minimum(sums, dims = (1, 2)), dims = (1, 2))
    out, copyto!(similar(ttl, T), ttl)
end

@deprecate stateposteriors(mfsm, lhs) pdfposteriors(mfsm, lhs)

"""
    maxstateposteriors(mfsm, lhs)

Calculate the posterior of "assigning" the \$n\$th frame
to the \$i\$th state conditioned on all other states maximizing the
likelihood of the sequence. The output is a matrix `μ` (# pdf x # frames).
"""
function maxstateposteriors(mfsm::MatrixFSM{SR},
                            in_lhs::AbstractArray) where SR<:TropicalSemiring
    # Convert the FSM and the likelihood matrix to the Log-semifield.
    lhs = copyto!(similar(in_lhs, SR), in_lhs)

    # Add one frame and one dimension for the final state.
    expanded_lhs = _expand(lhs)

    S = size(mfsm.C, 1)
    N = size(expanded_lhs, 2)

    # Expand the likelihood matrix to get the per-state likelihoods.
    state_lhs = matmul!(similar(lhs, S, N), mfsm.C, expanded_lhs)

    α = αrecursion(mfsm.π, mfsm.Tᵀ, state_lhs)
    β = βrecursion(mfsm.T, state_lhs)
    elmul!(similar(state_lhs), α, β)[1:end-1,1:end-1]
end

function findmaxstate(mfsm::MatrixFSM{SR}, μ::AbstractArray{SR},
                      prevstate::Int) where {SR,T}
    m = typemin(SR)
    s = -1
    for (i, μᵢ) in pairs(μ)
        if μᵢ >= m && mfsm.T[prevstate,i] ≠ zero(SR)
            m = μᵢ
            s = i
        end
    end
    s
end

"""
    bestpath(mfsm, μ)

Return the sequence of state with the highest value from `μ`, where
`μ` is the ouput of [`maxstateposteriors`](@ref).
"""
function bestpath(mfsm::MatrixFSM{SR}, μ::AbstractArray{SR}) where SR
    seq = Vector{Int}(undef, size(μ, 2))
    seq[1] = argmax(μ[:,1])
    for n in 2:size(μ,2)
        seq[n] = findmaxstate(mfsm, μ[:,n], seq[n-1])
    end
    seq
end


using ITensors.NDTensors,
      Test

@testset "DiagTensor basic functionality" begin

  t = tensor(Diag(rand(ComplexF64,100)), (100,100))
  @test conj(data(store(t))) == data(store(conj(t)))
  @test typeof(conj(t)) <: DiagTensor

  d = rand(Float32, 10)
  D = Diag{ComplexF64}(d)
  @test eltype(D) == ComplexF64
  @test Array(dense(D)) == convert.(ComplexF64, d)
  simD = similar(D)
  @test length(simD) == length(D)
  @test eltype(simD) == eltype(D)
  D = Diag(1.0)
  @test eltype(D) == Float64
  @test complex(D) == Diag(one(ComplexF64))
  @test similar(D) == Diag(0.0)

  d = 3
  vr = rand(d)
  D = tensor(Diag(vr), (d,d))
  @test Array(D) == NDTensors.LinearAlgebra.diagm(0=>vr)
  @test matrix(D) == NDTensors.LinearAlgebra.diagm(0=>vr)
  # fails because of missing similar method for NonuniformDiag :(
  #@test permutedims(D, (2, 1)) == tensor(diagITensor(vr, j, i))
  #@test permutedims(tensor(diagITensor(2.0, j, i)), (2, 1)) == tensor(diagITensor(2.0, j, i))
end


@testset "Neighbour joining implementation" begin

    D = DataFrame(CSV.File(joinpath(@__DIR__,"..","examples","caudata_dist.txt")); copycols=false)
    tree = nj(D)
    # read results from ape implementation of nj
    apetree = readTopology(joinpath(@__DIR__, "..", "examples", "caudata_dist_nj.txt"))
    @test hardwiredClusterDistance(tree, apetree, false) == 0
    # also check branch lengths (more or less)
    @test sort!([e.length for e in tree.edge]) ≈ sort!([e.length for e in apetree.edge])

    # example where Ints are converted to Floats, and there's a < 0 edge length
    df = DataFrame(s1=[0, 5,9,9,5], s2=[5,0,10,10,6], s3=[9,10,0,8,4],
                   s4=[9,10,8,0,0], s5=[5,6, 4, 0,0])
    tree = (@test_logs (:info, r"have negative lengths") nj(df))
    @test writeTopology(tree) == "(((s1:2.0,s2:3.0):3.0,s3:4.0):2.0,s4:2.0,s5:-2.0);"
    # example with no names argument, force_nonnegative_edges
    D = [0 5 9 9 6.5; 5 0 10 10 7.5; 9 10 0 8 5.5; 9 10 8 0 1.5; 6.5 7.5 5.5 1.5 0]
    tree = (@test_logs (:info, r"reset to 0") PhyloNetworks.nj!(D, force_nonnegative_edges=true))
    @test writeTopology(tree) == "(((1:2.0,2:3.0):3.0,3:4.0):2.0,4:2.0,5:0.0);"
end


# issue 288
function not_a_reduction!(A, B)
  @turbo for j ∈ 1:size(A.re, 2)
    jre = B.re[j]
    jim = B.im[j]
    for i ∈ 1:size(A.re, 1)
      ire = B.re[i]
      iim = B.im[i]
      cisim = iim * jre - ire * jim
      cisre = ire * jre + iim * jim
      ρre_i = A.re[i,j]
      ρim_i = A.im[i,j]
      re_out = ρre_i * cisre - ρim_i * cisim
      im_out = ρre_i * cisim + ρim_i * cisre
      A.re[i,j] = re_out
      A.im[i,j] = im_out
    end
  end
  return nothing
end
function not_a_reduction_noturbo!(A, B)
  @turbo for j ∈ 1:size(A.re, 2)
    jre = B.re[j]
    jim = B.im[j]
    for i ∈ 1:size(A.re, 1)
      ire = B.re[i]
      iim = B.im[i]
      cisim = iim * jre - ire * jim
      cisre = ire * jre + iim * jim
      ρre_i = A.re[i,j]
      ρim_i = A.im[i,j]
      re_out = ρre_i * cisre - ρim_i * cisim
      im_out = ρre_i * cisim + ρim_i * cisre
      A.re[i,j] = re_out
      A.im[i,j] = im_out
    end
  end
  return nothing
end

@testset "Untangle reductions" begin
  @show @__LINE__
  N = 11
  A1 = (re = rand(N,N), im = rand(N,N))
  A2 = deepcopy(A1)
  B = (re = rand(N), im = rand(N))
  not_a_reduction!(A1, B)
  not_a_reduction_noturbo!(A2, B)
  @test A1.re ≈ A2.re
  @test A1.im ≈ A2.im
end


#import WavesAndEigenvalues.Meshutils: create_rotation_matrix_around_axis, get_rotated_index, get_reflected_index
using WavesAndEigenvalues.Helmholtz
## load mesh from file
case="NTNU_new"
mesh=Mesh("../NTNU/NTNU.msh",scale=1.0)
mesh=octosplit(mesh)
#create unit and full mesh from half-cell
#doms=[("Interior",:full),("Inlet",:full), ("Outlet_high",:full), ("Outlet_low",:full), ("Walls",:full),("Flame",:unit),("Bloch",:unit)]#
doms=[("Interior",:full),("Inlet",:full), ("Outlet_high",:full), ("Outlet_low",:full), ("Flame",:unit),]#

collect_lines!(mesh)
unit_mesh=extend_mesh(mesh,doms,unit=true)
full_mesh=extend_mesh(mesh,doms,unit=false)
D=Dict()
D["mesh"]=zeros(Float64,6044)
vtk_write(case*"_mesh",unit_mesh,D)
## meshing
#naxis=unit_mesh.dos.naxis
#nxbloch=unit_mesh.dos.nxbloch

## Helmholtz solver
#define speed of sound field
function speedofsound(x,y,z)
    if z<0.415
        return 347.0#m/s
    else
        return 850.0#m/s
    end
end

c=generate_field(unit_mesh,speedofsound,0)
C=generate_field(full_mesh,speedofsound,0)
#D=Dict()

#vtk_write("NTNU_c",full_mesh,D)
##describe model
dscrp=Dict()
dscrp["Interior"]=(:interior,())
dscrp["Outlet_high"]=(:admittance, (:Y_in,1E15))
dscrp["Outlet_low"]=(:admittance, (:Y_out,1E15))
##discretize models
#L=discretize(full_mesh, dscrp, C, order=:1)
l=discretize(unit_mesh, dscrp, c, order=:1,b=:b)
## solve model using beyn
#(type Γ by typing \Gamma and Enter)
Γ=[150.0+5.0im, 150.0-5.0im, 1000.0-5.0im, 1000.0+5.0im].*2*pi #corner points for the contour (in this case a rectangle)
#Ω, P = beyn(L,Γ,l=10,N=64, output=true)
l.params[:b]=1
#ω, p = beyn(l,Γ,l=10,N=4*128, output=true)
##
sol, n, flag = householder(l,914*2*pi,maxiter=16, tol=0*1E-10,output = true, n_eig_val=3)

##
# identify surface points
surface_points, tri_mask, tet_mask =get_surface_points(unit_mesh)
# compute normal vectors on surface triangles
normal_vectors=get_normal_vectors(unit_mesh)
## DA approach
blochify_surface_points!(unit_mesh, surface_points, tri_mask, tet_mask)
sens=discrete_adjoint_shape_sensitivity(unit_mesh,dscrp,c,surface_points,tri_mask,tet_mask,l,sol)

## correct sens for bloch formalism
#sens[:,unit_mesh.dos.naxis+1:unit_mesh.dos.naxis+unit_mesh.dos.nxbloch].+=sens[:,end-unit_mesh.dos.nxbloch+1:end]
sens[:,end-unit_mesh.dos.nxbloch+1:end]=sens[:,unit_mesh.dos.naxis+1:unit_mesh.dos.naxis+unit_mesh.dos.nxbloch]

##
# normalize point sensitivity with directed surface triangle area
normed_sens=normalize_sensitivity(surface_points,normal_vectors,tri_mask,sens)
# boundary-mass-based-normalizations
nsens = bound_mass_normalize(surface_points,normal_vectors,tri_mask,unit_mesh,sens)
# compute sensitivity in unit normal direction
normal_sens = normal_sensitivity(normal_vectors, normed_sens)

## save data
DD=Dict()
DD["real normed_sens1"]=real.(normed_sens[1,:])
DD["real normed_sens2"]=real.(normed_sens[2,:])
DD["real normed_sens3"]=real.(normed_sens[3,:])
DD["real normal_sens"]=real.(normal_sens)
DD["imag normed_sens1"]=imag.(normed_sens[1,:])
DD["imag normed_sens2"]=imag.(normed_sens[2,:])
DD["imag normed_sens3"]=imag.(normed_sens[3,:])
DD["imag normal_sens"]=imag.(normal_sens)
vtk_write_tri(case*"_tri",unit_mesh,DD)

##
mode="[$(string(round((sol.params[:ω]/2/pi),digits=2)))]Hz"
data=Dict()
data["real DA x"*mode]=real.(sens[1,:])
data["real DA y"*mode]=real.(sens[2,:])
data["real DA z"*mode]=real.(sens[3,:])
data["imag DA x"*mode]=imag.(sens[1,:])
data["imag DA y"*mode]=imag.(sens[2,:])
data["imag DA z"*mode]=imag.(sens[3,:])
data["real nDA x"*mode]=real.(nsens[1,:])
data["real nDA y"*mode]=real.(nsens[2,:])
data["real nDA z"*mode]=real.(nsens[3,:])
data["imag nDA x"*mode]=imag.(nsens[1,:])
data["imag nDA y"*mode]=imag.(nsens[2,:])
data["imag nDA z"*mode]=imag.(nsens[3,:])
#data["finite difference"*mode]=real.(sens_FD)
#data["central finite difference"]=real.(sens_CFD)
#data["diff DA -- FD"*mode]=abs.((sens.-sens_FD))
vtk_write(case, unit_mesh, data)
data=Dict()
v=bloch_expand(unit_mesh,sol,:b)
data["abs_p"*mode]=abs.(v)./maximum(abs.(v))
data["phase_p"*mode]=angle.(v)
#data["abs_p_adj"*mode]=abs.(sol.v_adj)./maximum(abs.(sol.v_adj))
#data["phase_p_adj"*mode]=angle.(sol.v_adj)
vtk_write(case*"_modes", full_mesh, data)

## FD approach

fd_sens=forward_finite_differences_shape_sensitivity(unit_mesh,dscrp,c,surface_points,tri_mask,tet_mask,l,sol)
fd_sens[:,end-unit_mesh.dos.nxbloch+1:end]=fd_sens[:,unit_mesh.dos.naxis+1:unit_mesh.dos.naxis+unit_mesh.dos.nxbloch]

## write output
mode="[$(string(round((sol.params[:ω]/2/pi),digits=2)))]Hz"
data=Dict()
data["abs_p"*mode]=abs.(sol.v)./maximum(abs.(sol.v))
data["phase_p"*mode]=angle.(sol.v)
data["abs_p_adj"*mode]=abs.(sol.v_adj)./maximum(abs.(sol.v_adj))
data["phase_p_adj"*mode]=angle.(sol.v_adj)
data["real DA x"*mode]=real.(fd_sens[1,:])
data["real DA y"*mode]=real.(fd_sens[2,:])
data["real DA z"*mode]=real.(fd_sens[3,:])
data["imag DA x"*mode]=imag.(fd_sens[1,:])
data["imag DA y"*mode]=imag.(fd_sens[2,:])
data["imag DA z"*mode]=imag.(fd_sens[3,:])
#data["finite difference"*mode]=real.(fd_sens_FD)
#data["central finite difference"]=real.(sens_CFD)
#data["diff DA -- FD"*mode]=abs.((sens.-sens_FD))
vtk_write(case*"_fd", unit_mesh, data)
##
findmax(abs.(sens-fd_sens))


Combinatorics.combinations(a, c, k::Integer) = MultisetCombinations(a, c, k)

struct MultisetCombinations{T,U,V}
    a::T
    c::U
    k::V
end

Base.start(c::MultisetCombinations) = begin
    s = [fill(i, c.c[i]) for i in eachindex(c.a)] |> flatten
    c.k <= length(s) ? s[1:c.k] : [length(c.a) + 1]
end

Base.next(c::MultisetCombinations, s) = begin
    combination = [c.a[si] for si in s]
    c.k > 0 || return combination, [length(c.a) + 1]

    s = copy(s)
    for i in length(s):-1:1
        s[i] += 1
        while s[i] <= length(c.a) && countnz(s .== s[i]) > c.c[s[i]]
            s[i] += 1
        end

        if s[i] <= length(c.a)
            t = [fill(i, c.c[i] - countnz(s .== i)) for i in s[i]:length(c.a)] |> flatten
            if length(t) >= length(s)-i
                s[i+1:end] = t[1:(length(s)-i)]
            else
                s = [length(c.a) + 1]
            end
            break
        end
    end

    combination, s
end

Base.done(c::MultisetCombinations, s) = !isempty(s) && s[1] > length(c.a)

Base.eltype(c::MultisetCombinations) = eltype(typeof(c))
Base.eltype(::Type{MultisetCombinations{T,U,V}}) where {T,U,V} = Array{eltype(T),1}

Base.length(c::MultisetCombinations) = begin
    c.k == 0 && return big(1)

    csum = sum(c.c)
    c.k > csum && return big(0)
    c.k == csum && return big(1)

    # http://www.m-hikari.com/ams/ams-2011/ams-17-20-2011/siljakAMS17-20-2011.pdf
    R, z = PowerSeriesRing(ZZ, c.k + 1, "z")
    gf = prod(sum(z^j for j in 0:m) for m in c.c)
    l = coeff(gf, c.k)

    convert(BigInt, l)
end


abstract type AbstractPenalty{T} end

mutable struct PowerPenalty{T} <: AbstractPenalty{T}
    p::T
end
(P::PowerPenalty)(x) = x^P.p

mutable struct RationalPenalty{T} <: AbstractPenalty{T}
    p::T
end
(R::RationalPenalty)(x) = sinh(R.p*x)/sinh(R.p)

import Base: copy
copy(p::TP) where {TP<:AbstractPenalty} = TP(p.p)


using ModelingToolkit, SparseArrays, Test, GalacticOptim, Optim

@variables x y
@parameters a b
loss = (a - x)^2 + b * (y - x^2)^2
sys1 = OptimizationSystem(loss,[x,y],[a,b],name=:sys1)
sys2 = OptimizationSystem(loss,[x,y],[a,b],name=:sys2)

@variables z
@parameters β
loss2 = sys1.x - sys2.y + z*β
combinedsys = OptimizationSystem(loss2,[z],[β],systems=[sys1,sys2],name=:combinedsys)

equations(combinedsys)
states(combinedsys)
parameters(combinedsys)

calculate_gradient(combinedsys)
calculate_hessian(combinedsys)
generate_function(combinedsys)
generate_gradient(combinedsys)
generate_hessian(combinedsys)
ModelingToolkit.hessian_sparsity(combinedsys)

u0 = [
    sys1.x=>1.0
    sys1.y=>2.0
    sys2.x=>3.0
    sys2.y=>4.0
    z=>5.0
]
p = [
    sys1.a => 6.0
    sys1.b => 7.0
    sys2.a => 8.0
    sys2.b => 9.0
    β => 10.0
]

prob = OptimizationProblem(combinedsys,u0,p,grad=true)
sol = solve(prob,NelderMead())
@test sol.minimum < -1e5

prob2 = remake(prob,u0=sol.minimizer)
sol = solve(prob,BFGS(initial_stepnorm=0.0001),allow_f_increases=true)
@test sol.minimum < -1e8
sol = solve(prob2,BFGS(initial_stepnorm=0.0001),allow_f_increases=true)
@test sol.minimum < -1e9

rosenbrock(x, p) =  (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
x0 = zeros(2)
_p  = [1.0, 100.0]

f = OptimizationFunction(rosenbrock,ModelingToolkit.AutoModelingToolkit(),x0,_p,grad=true,hess=true)
prob = OptimizationProblem(f,x0,_p)
sol = solve(prob,Optim.Newton()


# logarithmic spirals

spiraldata = [
    (10,  0.05),
    (4,   0.10),
    (0.5, 0.17)]

grid = GridRect(O - (200, 0), 175, 50)
@png begin
    background("grey15")
    for aspiral in spiraldata
        @layer begin
            sethue("antiquewhite")
            setcolor(sethue("antiquewhite"))
            Luxor.translate(nextgridpoint(grid))
            setline(.5)
            spiral(first(aspiral), last(aspiral), log=true, period=10π, :stroke)
            fontsize(16)
            label(string(aspiral), :S, offset=100)
        end
    end
end


using DoubleExponentialFormulas: quaddeo
using LinearAlgebra
using Test

# Test functions in
# Ooura, T., Mori, M., 1991.
# The double exponential formula for oscillatory functions over the half infinite interval.
# Journal of Computational and Applied Mathematics 38, 353–360.
# https://doi.org/10.1016/0377-0427(91)90181-I

# 1
let f(x) = exp(-x)*cos(x), expect = 1/2
    I, E = quaddeo(f, 1.0, π/2, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 2
let f(x) = x*sin(x)/(1 + x^2), expect = π/(2*exp(1))
    I, E = quaddeo(f, 1.0, 0.0, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 3
let f(x) = cos(x)/(1 + x^2), expect = π/(2*exp(1))
    I, E = quaddeo(f, 1.0, π/2, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 4
let f(x) = log((x^2 + 4)/(x^2 + 1))*cos(x), expect = (exp(-1) - exp(-2))*π
    I, E = quaddeo(f, 1.0, π/2, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 5 (Dirichlet integral)
let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 6
let f(x) = 1/sqrt(x)*sin(x), expect = sqrt(π/2)
    I, E = quaddeo(f, 1.0, 0.0, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 7
let f(x) = cos(x)/sqrt(x), expect = sqrt(π/2)
    I, E = quaddeo(f, 1.0, π/2, 0.0, Inf)
    @test I isa Float64
    @test_broken I ≈ expect
    @test abs(I - expect) < 1e-6
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# 8
let f(x) = sin(x)*log(x), expect = -Base.MathConstants.eulergamma
    I, E = quaddeo(f, 1.0, 0.0, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Various interval (-∞, ∞)
let f(x) = sin(x)/x, expect = π
    I, E = quaddeo(f, 1.0, 0.0, -Inf, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Various interval (0, ∞)
let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, 0.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Various interval (a, ∞)
let f(x) = sin(x)/x, expect = π/2 - BigFloat("1.851937051982466170361053370157991363345809728981154909804")
    I, E = quaddeo(f, 1.0, 0.0, π, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Various interval (-∞, 0)
let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, -Inf, 0.0)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Various interval (-∞, b)
let f(x) = sin(x)/x, expect = π/2 - BigFloat("1.851937051982466170361053370157991363345809728981154909804")
    I, E = quaddeo(f, 1.0, 0.0, -Inf, -π)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Various interval (a, b)
let f(x) = sin(x)/x, expect = BigFloat("-0.43378547584983772011527320785824193420335637968613863796")
    I, E = quaddeo(f, 1.0, 0.0, π, 2π)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end
let f(x) = sin(x)/x, expect = BigFloat("0.43378547584983772011527320785824193420335637968613863796")
    I, E = quaddeo(f, 1.0, 0.0, -π, -2π)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end
let f(x) = sin(x)/x, expect = BigFloat("3.703874103964932340722106740315982726691619457962309819609")
    I, E = quaddeo(f, 1.0, 0.0, -π, π)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Corner case (a == b)
let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, 0.0, 0.0)
    @test I isa Float64
    @test I ≈ 0
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, 1.0, 1.0)
    @test I isa Float64
    @test I ≈ 0
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Corner case (a > b)
let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, Inf, 0.0)
    @test I isa Float64
    @test I ≈ -π/2
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Multiple integral intervals
let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, 0.0, 1.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

let f(x) = sin(x)/x, expect = π/2
    I, E = quaddeo(f, 1.0, 0.0, 0.0, 1.0, 2.0, Inf)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I))*norm(I))
end

# Vectorized numerical integration
let f(x) = [sin(x)/x, exp(-x)*cos(x)], expect = [π/2, 1/2]
    I, E = quaddeo(f, 1.0, 0.0, 0.0, Inf)
    @test eltype(I) == Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(eltype(I))*norm(I))
end

# Type handling
let f(x) = exp(-x)*cos(x), expect = BigFloat("0.5")
    ω = BigFloat(1.0)
    θ = BigFloat(π)/2
    a = BigFloat(0.0)
    b = BigFloat(Inf)

    for T in [Float32, Float64, BigFloat]
        ω_T, θ_T, a_T, b_T = T.((ω, θ, a, b))
        I, E = quaddeo(f, ω_T, θ_T, a_T, b_T)
        @test I isa T
        @test I ≈ expect
        @test E ≤ sqrt(eps(typeof(I))*norm(I))
    end
end


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Frames (Modia3D/Frames/_module.jl)
#

"""
    const Modia3D.Vector3D = SVector{3,Float64}

Type of a vector in 3D space (e.g. position vector of the origin of a frame)
"""
const Vector3D = SVector{3,Float64}

"""
    const Modia3D.ZeroVector3D = Vector3D(0.0, 0.0, 0.0)

Constant of a Vector3D where all elements are zero
"""
const ZeroVector3D = Vector3D(0.0, 0.0, 0.0)

"""
    vec = Modia3D.axisValue(axis, value)
    vec = Modia3D.axisValue(axis, positive, value)

Return `vec::Vector3D` where all elements are zero with exception of `vec[axis] = value` or
`vec[axis] = positive ? value : -value`.
"""
@inline axisValue(axis::Int, value::Number)::Vector3D = axis==1 ? @SVector([value, 0.0, 0.0]) : (axis==2 ? @SVector([0.0, value, 0.0]) :
                                                       (axis==3 ? @SVector([0.0, 0.0, value]) : error("Bug in Modia3D: axisValue($axis, ...) - argument needs to be 1,2 or 3.")))
@inline axisValue(axis::Int, positive::Bool, value::Number)::Vector3D = positive ? axisValue(axis,value) : axisValue(axis,-value)

"""
    M = Modia3D.skew(e::AbstractVector)

Return the skew symmetric matrix `M::SMatrix{3,3,Float64,9}` of vector `e` (`length(e) = 3`)
"""
skew(e::AbstractVector) = @SMatrix([  0.0   -e[3]    e[2];
                                      e[3]   0.0    -e[1];
                                     -e[2]   e[1]    0.0 ])


#=
include("scripts/example3by3.jl")
=#
addprocs(4)
using WarpFlow
using PyPlot

#---- set the data and the target
dim = 2
tmpx = shuffle([rand(50); randn(80)/10 .+ .8])
tmpy = tmpx + rand(size(tmpx)) .* 1.1
nPhi = length(tmpx)
X = Array{Float64,1}[[(tmpx[i]- minimum(tmpx))/maximum(tmpx), (tmpy[i]-minimum(tmpy))/maximum(tmpy)] for i=1:nPhi]
# target(x) = targetUnif2d(x; targSig = 0.1, limit = 0.4, center = 0.5)
@everywhere target(x) = WarpFlow.targetNormal2d(x; targSig = 0.5, center = [0.5,0.5])

# set an array of lambda, sigmas
matSigmaLambda = [(s,l) for s in (0.2, 0.1), l in (5.0, 0.5)]

#  gradient ascent on kappa and eta_coeff
est_den_container = @parallel (hcat) for (s,l) in matSigmaLambda
	kappa = Array{Float64,1}[X[Int(i)]  for i in 1:round(nPhi/2)]
	y0    = Flow(kappa, WarpFlow.array1(dim, length(kappa)), X, WarpFlow.array2eye(dim, nPhi), dim)
	for counter = 1:300
		z0 = WarpFlow.get_grad(y0, target, l, s)
		y0 = y0 + 0.005 * z0
	end
	#-- now we save the estimated density
	x_grd, y_grd =  WarpFlow.meshgrid(linspace(-0.1, 1.1, 175), linspace(-0.1, 1.1, 175))
	phix_grd_0   = Array{Float64,1}[[x_grd[i]; y_grd[i]] for i=1:length(x_grd)]
	Dphix_grd_0  = Array{Float64,2}[eye(2)               for i=1:length(x_grd)]
	yplt0        = Flow(y0.kappa, y0.eta_coeff, phix_grd_0, Dphix_grd_0, dim)
	dydt(t,y)    = WarpFlow.d_forward_dt(y, s)
 	(t1, yplt1)  = WarpFlow.ode23_abstract_end(dydt, [0,1], yplt0) # Flow y0 forward to time 1
	det_grd = Float64[abs(det(yplt1.Dphix[i])) for i=1:length(x_grd)]
	den,    = target(yplt1.phix)
	det_grd .* den
end

rmprocs(workers())

# writecsv("simulations/example3by3den.csv", est_den_container)
# writecsv("simulations/example3by3X.csv", hcat(X...))

#dim = 2
#matSigmaLambda = [(s,l) for s in (0.4, 0.2, 0.05), l in (10.0, 5.0, 1.0)]
#est_den_container = readcsv("simulations/example3by3den.csv")
#Xcat = readcsv("simulations/example3by3X.csv")
#X =  Array{Float64,1}[Xcat[:,k] for k in 1:size(Xcat,2) ]
x_grd, y_grd =  WarpFlow.meshgrid(linspace(-0.1, 1.1, 175),linspace(-0.1, 1.1, 175))
fig = figure(figsize=(14,14))
for k = 1:4
	subplot(2,2,k)
	title("(sigma, lambda) = $(matSigmaLambda[k])")
	scatter(Float64[point[1] for point in X], Float64[point[2] for point in X], c="b")
	contour(x_grd, y_grd, reshape(est_den_container[:,k], size(x_grd)), 30)
end
# plt[:savefig]("simulations/example3by3.pdf",dpi=180)
# plt[:close](fig)


"""
    IntervalMesh <: AbstractMesh

A 1D mesh on an `IntervalDomain`.

# Constuctors

    IntervalMesh(domain::IntervalDomain, faces::AbstractVector)

Construct a 1D mesh with face locations at `faces`.

    IntervalMesh(domain::IntervalDomain[, stetching=Uniform()]; nelems=)

Constuct a 1D mesh on `domain` with `nelems` elements, using `stretching`. Possible values of `stretching` are:

- [`Uniform()`](@ref)
- [`ExponentialStretching(H)`](@ref)
"""
struct IntervalMesh{I <: IntervalDomain, V <: AbstractVector} <: AbstractMesh1D
    domain::I
    faces::V
end

domain(mesh::IntervalMesh) = mesh.domain
nelements(mesh::IntervalMesh) = length(mesh.faces) - 1
elements(mesh::IntervalMesh) = Base.OneTo(nelements(mesh))

function Base.show(io::IO, mesh::IntervalMesh)
    print(io, nelements(mesh), "-element IntervalMesh of ")
    print(io, mesh.domain)
end

coordinates(mesh::IntervalMesh, elem::Integer, vert::Integer) =
    mesh.faces[elem + vert - 1]

function coordinates(
    mesh::IntervalMesh,
    elem::Integer,
    (ξ,)::Tuple{T},
) where {T}
    ca = mesh.faces[elem]
    cb = mesh.faces[elem + 1]
    Geometry.linear_interpolate((ca, cb), ξ)
end

function containing_element(mesh::IntervalMesh, coord)
    i = min(searchsortedlast(mesh.faces, coord), nelements(mesh))
    lo = Geometry.component(mesh.faces[i], 1)
    hi = Geometry.component(mesh.faces[i + 1], 1)
    val = Geometry.component(coord, 1)
    ξ = ((val - lo) + (val - hi)) / (hi - lo)
    return i, (ξ,)
end
function is_boundary_face(mesh::IntervalMesh, elem::Integer, face)
    !Domains.isperiodic(mesh.domain) &&
        ((elem == 1 && face == 1) || (elem == nelements(mesh) && face == 2))
end

function boundary_face_name(mesh::IntervalMesh, elem::Integer, face)
    if !Domains.isperiodic(mesh.domain)
        if elem == 1 && face == 1
            return mesh.domain.boundary_names[1]
        elseif elem == nelements(mesh) && face == 2
            return mesh.domain.boundary_names[2]
        end
    end
    return nothing
end

abstract type StretchingRule end

"""
    Uniform()

Use uniformly-sized elements.
"""
struct Uniform <: StretchingRule end

function IntervalMesh(
    domain::IntervalDomain{CT},
    ::Uniform = Uniform();
    nelems,
) where {CT <: Geometry.Abstract1DPoint{FT}} where {FT}
    faces = range(domain.coord_min, domain.coord_max; length = nelems + 1)
    IntervalMesh(domain, faces)
end

"""
    ExponentialStretching(H)

Apply exponential stretching to the domain when constructing elements. `H` is
the scale height (a typical atmospheric scale height `H ≈ 7.5`km).

For an interval ``[z_0,z_1]``, this makes the elements uniformally spaced in
``\\zeta``, where
```math
\\zeta = \\frac{1 - e^{-\\eta/h}}{1-e^{-1/h}},
```
where ``\\eta = \\frac{z - z_0}{z_1-z_0}``, and ``h = \\frac{H}{z_1-z_0}`` is
the non-dimensional scale height.
"""
struct ExponentialStretching{FT} <: StretchingRule
    H::FT
end

function IntervalMesh(
    domain::IntervalDomain{CT},
    stretch::ExponentialStretching;
    nelems,
) where {CT <: Geometry.Abstract1DPoint{FT}} where {FT}
    cmin = Geometry.component(domain.coord_min, 1)
    cmax = Geometry.component(domain.coord_max, 1)
    R = cmax - cmin
    h = stretch.H / R
    η(ζ) = -h * log1p((expm1(-1 / h)) * ζ)
    faces =
        [CT(cmin + R * η(ζ)) for ζ in range(FT(0), FT(1); length = nelems + 1)]
    IntervalMesh(domain, faces)
end


@doc "All the abstract domains" ->
all_domains() = Set(Interval,AbstractBool,HyperBox)

@doc "Return a type Union of all the abstract domains which can abstract a type T" ->
abstractdomains(CT::DataType) = Union{filter(AT->isabstract(CT,AT),all_domains())...}

## Random Number Generation
## =========================
@doc "Random number between `a` and `b`" ->
rand_interval(a::Float64, b::Float64) = a + (b - a) * rand()


module TestGradfunc
    using Calculus2
    using Base.Test

    f(x) = sin(x)
    g(x) = @gradfunc(sin(x), x)
    h(x) = @hessfunc(sin(x), x)

    f(1.0)
    g(1.0)
    h(1.0)

    f(z) = 1 / (1 + exp(-z))
    g(z) = @gradfunc(1 / (1 + exp(-z)), z)

    @test isapprox(
        g(0.0),
        grad(f, 0.0)
    )

    @test isapprox(
        g(1.0),
        grad(f, 1.0)
    )

    f(x) = sin(x) + cos(x)^2 + sin(x)^3
    g(x) = @gradfunc(sin(x) + cos(x)^2 + sin(x)^3, x)

    @test isapprox(
        g(0.0),
        grad(f, 0.0)
    )

    @test isapprox(
        g(1.0),
        grad(f, 1.0)
    )

    # NB: This won't work for multivariate functions
    f(x) = sin(x) + cos(x)^2 + sin(x)^3
    grad(f, 1.0, method = :finite)
    grad(f, 1.0, method = :ad)
    gradfunc(f, 1.0, method = :finite)(1.0)
    gradfunc(f, 1.0, method = :ad)(1.0)
    g(x) = @gradfunc(sin(x) + cos(x)^2 + sin(x)^3, x)
    g(1.0)
    hess(f, 1.0, method = :finite)
    hessfunc(f, 1.0, method = :finite)(1.0)
    h(x) = @hessfunc(sin(x) + cos(x)^2 + sin(x)^3, x)
    h(1.0)
end


#Copyright (c) 2015 Rex Computing and Isaac Yonemoto

#see LICENSE.txt

#this work was supported in part by DARPA Contract D15PC00135

#test-warlpiri-division.jl

###############################################################################
##COMPLETE DIVISION TABLE

#column key
#none_ pfew_ pone_ psome ptwo_ pmany pall_ junk_ nnone n_few n_one nsome n_two nmany n_all
wdiv = Utype[
#none_ / line
 junk_ none_ none_ none_ none_ none_ none_ junk_ junk_ none_ none_ none_ none_ none_ none_;
#pfew_ / line
 junk_ pf_pm pfew_ pfew_ pfew_ pfew_ none_ junk_ junk_ nm_nf n_few n_few n_few n_few none_;
#pone_ / line
 junk_ ps_pm pone_ pfew_ pfew_ pfew_ none_ junk_ junk_ nm_ns n_one n_few n_few n_few none_;
#psome / line
 junk_ ps_pm psome pf_ps pfew_ pfew_ none_ junk_ junk_ nm_ns nsome ns_nf n_few n_few none_;
#ptwo_ / line
 junk_ pmany ptwo_ psome pone_ pfew_ none_ junk_ junk_ nmany n_two nsome n_one n_few none_;
#pmany / line
 junk_ pmany pmany ps_pm ps_pm pf_pm none_ junk_ junk_ nmany nmany nm_ns nm_ns nm_nf none_;
#pall_ / line
 junk_ pall_ pall_ pall_ pall_ pall_ junk_ junk_ junk_ n_all n_all n_all n_all n_all junk_;
#junk_ / line
 junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_;
#nnone / line
 junk_ none_ none_ none_ none_ none_ none_ junk_ junk_ none_ none_ none_ none_ none_ none_;
#n_few / line
 junk_ nm_nf n_few n_few n_few n_few none_ junk_ junk_ pf_pm pfew_ pfew_ pfew_ pfew_ none_;
#n_one / line
 junk_ nm_ns n_one n_few n_few n_few none_ junk_ junk_ ps_pm pone_ pfew_ pfew_ pfew_ none_;
#nsome / line
 junk_ nm_ns nsome ns_nf n_few n_few none_ junk_ junk_ ps_pm psome pf_ps pfew_ pfew_ none_;
#n_two / line
 junk_ nmany n_two nsome n_one n_few none_ junk_ junk_ pmany ptwo_ psome pone_ pfew_ none_;
#nmany / line
 junk_ nmany nmany nm_ns nm_ns nm_nf none_ junk_ junk_ pmany pmany ps_pm ps_pm pf_pm none_;
#n_all / line
 junk_ n_all n_all n_all n_all n_all junk_ junk_ junk_ pall_ pall_ pall_ pall_ pall_ junk_
]
testop(/, wdiv)


@testset "Generation tools" begin

    t = (1,2)
    @test ϕ_tuple(t::NTuple) ≈ atan(2, 1)
    t = (1.23,129.23)
    @test ϕ_tuple(t::NTuple) ≈ atan(t[2], t[1])

    @test idx_to_dim([(1, 2) (3, 4); (5, 6) (7, 8)]) == [1 3; 5 7;;; 2 4; 6 8]

    @test idx_to_dim([(1, 2), (3, 4)]) == [1 2; 3 4]

end


using Printf, PyCall, PyPlot
using AbaqusReader, JLD, MAT

using AD4SM

mean(x) = sum(x)/length(x)
norm(x) = sqrt(sum(x.^2))
#
#
# sort_bnd sorts the nodes of the boundary to
# form a path without self-intersections
function sort_bnd(nodes, nid_bnd)
  closest(node0, nodes) = findmin([norm(node0-node) for node in nodes])[2]

  id_srtd   = [nid_bnd[1],]
  lastnode  = nid_bnd[1]
  bnd_nodes = nid_bnd[2:end]

  while !isempty(bnd_nodes)
    idd = closest(nodes[lastnode], nodes[bnd_nodes])
    lastnode = bnd_nodes[idd]
    push!(id_srtd, bnd_nodes[idd])
    deleteat!(bnd_nodes, idd)
  end
  return id_srtd
end

# get_vol(p0) = π/2*sum((p0[2,1:end-1]+p0[2,2:end]) .*
#                      (p0[1,2:end]-p0[1,1:end-1]) .*
#                      (p0[1,2:end]+p0[1,1:end-1]) )
#
# get_vol calculates the volume of the cavity
get_vol(p0) = π/2*sum(
  (p0[2:2:end-2]+p0[4:2:end]) .*
  (p0[3:2:end]-p0[1:2:end-2]) .*
  (p0[3:2:end]+p0[1:2:end-2]) )
#
# functions for plotting
#
function get_I1(elems, u0)
    F   = Elements.getinfo(elems,u0,info=:F)
    [ Materials.getInvariants(transpose(F)*F)[1] for F in F]
end
function get_I3(elems, u0)
    F   = Elements.getinfo(elems,u0,info=:F)
    [ Materials.getInvariants(transpose(F)*F)[3] for F in F]
end
patch = pyimport("matplotlib.patches")
coll  = pyimport("matplotlib.collections")
function plot_model(elems, nodes;
                    u = zeros(length(nodes[1]), length(nodes)),
                    Φ = [],
                    linewidth = 0.25,
                    facecolor = :c,
                    edgecolor = :b,
                    alpha     = 1,
                    cmap      = :hsv,
                    clim      = [],
                    dTol      = 1e-6,
                    cfig      = figure(),
                    ax        = cfig.add_subplot(1,1,1))

  nodes     = [node + u[:,ii] for (ii,node) in enumerate(nodes)]
  patchcoll = coll.PatchCollection([patch.Polygon(nodes[elem.nodes])
                                    for elem ∈ elems], cmap=cmap)
    if !isempty(Φ)
    patchcoll.set_array(Φ)
    cfig.colorbar.(patchcoll, ax=ax)

    if isempty(clim)
      clim = patchcoll.get_clim()
      if abs(clim[2]-clim[1]) < dTol
        clim  = sum(clim)/2*[0.9, 1.1]
      end
    end
    patchcoll.set_clim(clim)
  else
    patchcoll.set_color(facecolor)
  end

  patchcoll.set_edgecolor(edgecolor)
  patchcoll.set_alpha(alpha)
  patchcoll.set_linewidth(linewidth)

  ax.set_aspect("equal")
  ax.add_collection(patchcoll)
  ax.autoscale()
  (cfig, ax, patchcoll)
end
;

sMeshFile = "AxSymDomainj.inp"
mat       = Materials.MooneyRivlin(1e1, 1e1, 1e3)
Δu        = 18
sPosFix   = "NHd"
N         = 200
LF        = range(1/10N, 1, length=N)
dTol      = 1e-6
ballus    = true
;

sFileName = splitext(basename(sMeshFile))[1];
mymodel   = AbaqusReader.abaqus_read_mesh(sMeshFile)
nodes     = [mymodel["nodes"][ii] for ii in 1:mymodel["nodes"].count]
el_nodes  = [item[2]              for item in mymodel["elements"]]

# for node in nodes
#   node[1]-=10
# end

id_btm   = mymodel["node_sets"]["BTM"]
id_top   = mymodel["node_sets"]["TOP"]
id_bnd   = mymodel["node_sets"]["BND"]

elems     = []
for item ∈ el_nodes
  if length(item)==3
    push!(elems, Elements.ASTria(item, nodes[item], mat=mat))
  elseif length(item)==4
    push!(elems, Elements.ASQuad(item, nodes[item], mat=mat))
  end
end

@show (nNodes, nElems) = (length(nodes), length(elems)); flush(stdout)

# sort the internal boundary nodes
#
id_srtd   = sort_bnd(nodes, id_bnd);
push!(id_srtd, id_srtd[1])
pos0      = hcat(nodes[id_srtd]...)
idxu      = LinearIndices((2,nNodes))
@show V0  = get_vol(pos0)

pos_btm   = hcat(nodes[sort_bnd(nodes, id_btm)]...)
pos_top   = hcat(nodes[sort_bnd(nodes, id_top)]...)

plot_model(elems, nodes)
PyPlot.plot(pos0[1,:], pos0[2,:], color=:b)
PyPlot.plot(pos_btm[1,:], pos_btm[2,:], color=:r)
PyPlot.plot(pos_top[1,:], pos_top[2,:], color=:r)
PyPlot.title("Undeformed domain")
;

# displacement boundary conditions
u0     = 1e-5randn(2, nNodes)
idxu   = LinearIndices(u0)
bifree = trues(size(u0))

bifree[:,id_btm] .= false
bifree[2,id_top] .= false

u0[:,id_btm] .= 0
u0[2,id_top] .= 0

;

@time Elements.solvestep!(elems, u0, bifree,
  bprogress=true, becho=false, dTol=dTol)
;

u            = copy(u0)
u[:,id_btm] .= 0
u[2,id_top] .= -Δu

allus_e = Elements.solve(elems, u, ifree=bifree, LF=range(1/10N, 1, length=N),
                     becho=true, bechoi=false, ballus=ballus, bprogress=false,
                     dTolu=1e-6, dTole=1e-2, maxiter=15)
;

cfig = PyPlot.figure()

ax1  = cfig.add_subplot(1,1,1)
plot_model(elems, nodes, cfig=cfig, ax=ax1,
  u=allus_e[end][1], Φ=get_I3(elems, allus_e[end][1]).-1)
;

nid_bnd = length(id_srtd)
nDoFsu  = 2*nNodes
q       = 10
nGropus = ceil(Int, nid_bnd/q)

id_eqs_DoFs = [(ii-1)*q+1:ii*q+1 for ii in 1:nGropus-1]
append!(id_eqs_DoFs, [id_eqs_DoFs[end][end]:nid_bnd])

eqns = vcat(
  [begin
      id_bnd_ii = id_srtd[idx]
      idx_ii = idxu[:, id_bnd_ii][:]
      Elements.ConstEq(x->get_vol(pos0[:,idx][:]+x[1:end-1])-x[end],
                        vcat(idx_ii, nDoFsu+ii), adiff.D2)
      end for (ii,idx) in enumerate(id_eqs_DoFs)] ...,
  Elements.ConstEq(x->sum(x)-V0, nDoFsu+1:nDoFsu+nGropus, adiff.D2))
@show nEqns = length(eqns)
;

# displacement boundary conditions
idxu   = LinearIndices((2,nNodes))
u0     = 1e-4randn(2, nNodes+ceil(Int, nGropus/2))
bifree = trues(size(u0))
λ      = zeros(nEqns)

# make sure that the volume constraint is initially satisfied
for (ii, Eq) in enumerate(eqns[1:end-1])
  u0[nDoFsu+ii] = Eq.func(zeros(size(Eq.iDoFs)))
end

bifree[:,id_btm] .= false
bifree[2,id_top] .= false

u0[:,id_btm]     .= 0
u0[2,id_top]     .= 0
;

@time Elements.solvestep!(elems,
        u0, bifree, λ=λ, bprogress=true, becho=false,
        dTolu=1e-5, dTole=1e-3, eqns=eqns)
;

u           = copy(u0)
u[:,id_btm] .= 0
u[2,id_top] .= -Δu

allus_d = Elements.solve(elems, u, ifree=bifree, λ=λ, eqns=eqns, LF=range(0, 1, length=200),
                     becho=true, bechoi=false, ballus=ballus, bprogress=false,
                     dTolu=1e-5, dTole=5e-2, maxiter=11)
;

cfig = PyPlot.figure()

ax1  = cfig.add_subplot(1,1,1)
plot_model(elems, nodes, cfig=cfig, ax=ax1,
  u=allus_d[end][1], Φ=get_I1(elems, allus_d[end][1]).-3)
;

cfig = PyPlot.figure()

ax1  = cfig.add_subplot(2,1,1)
plot_model(elems, nodes, cfig=cfig, ax=ax1,
  u=allus_d[end][1], Φ=get_I1(elems, allus_d[end][1]).-3)
PyPlot.plot(pos0[1,:]+allus_d[end][1][1,id_srtd],
  pos0[2,:]+allus_d[end][1][2,id_srtd], color=:b)

ax2  = cfig.add_subplot(2,1,2)
plot_model(elems, nodes, cfig=cfig, ax=ax2,
  u=allus_e[end][1], Φ=get_I1(elems, allus_e[end][1]).-3)
PyPlot.plot(pos0[1,:]+allus_e[end][1][1,id_srtd],
  pos0[2,:]+allus_e[end][1][2,id_srtd], color=:b)

ax1.set_xlim([15, 85]); ax1.set_ylim([-26, 10])
ax1.set_title("with internal volume constraint, \$I_1-3\$")
ax1.set_xticks([]); ax1.set_yticks([])
ax2.set_xlim([15, 85]); ax2.set_ylim([-26, 10])
ax2.set_title("without internal volume constraint, \$I_1-3\$")
ax2.set_xticks([]); ax2.set_yticks([])

rf_tot_d = [sum(item[2][2, id_top])  for item in allus_d]
Δu_tot_d = [mean(item[1][2,id_top])  for item in allus_d]

rf_tot_e = [sum(item[2][2, id_top])  for item in allus_e]
Δu_tot_e = [mean(item[1][2,id_top])  for item in allus_e]

PyPlot.figure()
PyPlot.plot(Δu_tot_d/50, rf_tot_d/1e5, color=:r)
PyPlot.plot(Δu_tot_e/50, rf_tot_e/1e5, color=:b)
PyPlot.xlabel("\$\\Delta u_2/L_2\$")
PyPlot.ylabel("\$F_2\\times10^{-5}\$")
PyPlot.title("Total reaction force")
;

JLD.save(sFileName*".jld", Dict(
            "nodes"=>nodes, "elems"=>elems, "el_nodes"=>el_nodes,
            "allus_d"=>allus_d, "allus_e"=>allus_e,
            "LF"=>LF, "mat"=>mat,
            "id_btm"=>id_btm, "id_top"=>id_top, "id_bnd"=>id_bnd,
            "sMeshFile"=>sMeshFile))
@printf("results written to %s.jld\n", sFileName); flush(stdout)
;

MAT.matwrite(sFileName*".mat", Dict(
        "nodes"=>nodes, "elems"=>el_nodes, "nNodes"=>nNodes,
        "u0_d"=>allus_d[end][1], "u0_e"=>allus_e[end][1],
        "I1_d"=>get_I1(elems, allus_d[end][1]),
        "I1_e"=>get_I1(elems, allus_e[end][1]),
        "id_btm"=>id_btm, "id_top"=>id_top, "id_bnd"=>id_bnd,
        "u_tot_d"=>Δu_tot_d, "u_tot_e"=>Δu_tot_e,
        "rf_tot_d"=>rf_tot_d, "rf_tot_e"=>rf_tot_e))
@printf("results written to %s.mat\n", sFileName); flush(stdout)
;


### A Pluto.jl notebook ###
# v0.14.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 3649f170-923a-11eb-321c-cf95849cc044
html"""

Section 2.5

 Random Walks

Lecture Video

"""

# ╔═╡ 97e807b2-9237-11eb-31ef-6fe0d4cc94d3
begin
    import Pkg
    Pkg.activate(mktempdir())

    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="BenchmarkTools", version="0.6"),
    ])

    using Plots, PlutoUI, BenchmarkTools
end

# ╔═╡ 5f0d7a44-91e0-11eb-10ae-d73156f965e6
TableOfContents(aside=true)

# ╔═╡ 9647147a-91ab-11eb-066f-9bc190368fb2
md"""
# Julia concepts

- Benchmarking: BenchmarkTools.jl
- Plotting in a loop
- Generic programming
- Mutable vs immutable structs
- Vectors of vectors
- Aliasing of memory
- `cumsum`
"""

# ╔═╡ ff1aca1e-91e7-11eb-343e-0f89d9570b06
md"""
# Motivation: Dynamics of hard discs
"""

# ╔═╡ 66a2f510-9232-11eb-3be9-131febc0039f
md"""
Brown observed **Brownian motion** in 1827: Large particles like sand or pollen in water move around seemingly at random. Einstein explained this in 1905 as repeated impacts with water molecules.

We can visualise that with a simulation of hard discs bouncing off one another. Even though the dynamics is not random -- each disc follows Newton's laws -- if we just look at a single one of them it *looks* random.
"""

# ╔═╡ bd3170e6-91ae-11eb-06f8-ebb6b2e7869f
md"""
## Visualising random walks
"""

# ╔═╡ a304c842-91df-11eb-3fac-6dd63087f6de
md"""
A **random walk** models random motion in time and space. At each time step an object moves in a random direction.

Let's visualise the result in 2 dimensions.
"""

# ╔═╡ 798507d6-91db-11eb-2e4a-3ba02f12ba65
md"""
N = $(@bind N Slider(1:6, show_value=true, default=1))
"""

# ╔═╡ 3504168a-91de-11eb-181d-1d580d5dc071
md"""
t = $(@bind t Slider(1:10^N, show_value=true, default=1))
"""

# ╔═╡ 4c8d8294-91db-11eb-353d-c3696c615b3d
begin
	plot(traj[1:t], ratio=1, leg=false, alpha=0.5, lw=2)
	scatter!([ traj[1], traj[t] ], c=[:red, :green])

	xlims!(minimum(first.(traj)) - 1, maximum(first.(traj)) + 1)
	ylims!(minimum(last.(traj)) - 1, maximum(last.(traj)) + 1)

end

# ╔═╡ b62c4af8-9232-11eb-2f66-dd27dcb87d20
md"""
We see that the dynamics closely resembles, at least qualitatively, that of the hard disc.
"""

# ╔═╡ 905379ce-91ad-11eb-295d-8354ecf5c5b1
md"""
# Why use random walks?

#### Why should we model with random processes?

- Either we need to get a lot of data to precisely characterise a system, or we assume it's random

- It may well be that *can* have all the details, but by making a simplifying assumption we actually gain more *understanding*.

#### Examples:

- Stock price going up and down

- Pollutants getting dispersed in the air

- Neutral genes moving through a population

"""

# ╔═╡ 5c4f0f26-91ad-11eb-033b-2bd221f0bdba
md"""
# Simple random walk

The simplest version is called **simple random walk**: one object jumps on the integers: at each time step it jumps left or right.

Each such step is a new random variable, taking values $\pm 1$ at each step, each with probability $1/2$ (in the simplest case).
We can think of this as a scaling of a Bernoulli random variable.

In order to simulate this, we need to know how to generate the jumps.
"""

# ╔═╡ da98b676-91e0-11eb-0d97-57b8a8aadf2a
md"""
## Julia: Benchmarking
"""

# ╔═╡ e0b607c0-91e0-11eb-10aa-53ec33570e59
md"""
There are various ways we could generate random values $\pm 1$. Let's use this as an opportunity to learn about **benchmarking** in Julia, i.e. measuring run time to compare performance. In this case we will do "micro-benchmarks", which try to compare tiny pieces of code which we intend to run many millions of times, so that small differences in run time can be significant.

The `BenchmarkTools.jl` package provides relatively easy-to-use tools for measuring this, by running the code many times and calculating statistics. The `@btime` macro is a simple way to estimate actual running time. Each option should be enclosed in a function.
"""

# ╔═╡ fa1635d4-91e3-11eb-31bd-cf61c502ad35
md"""
Here are a few different ways we could generate random steps:
"""

# ╔═╡ f7f9e4c6-91e3-11eb-1a56-8b98f0b09b46
begin
	step1() = rand( (-1, +1) )

	step2() = 2 * (rand() < 0.5) - 1

	step3() = 2 * rand(Bool) - 1

	step4() = sign(randn())
end

# ╔═╡ 5da7b076-91b4-11eb-3eba-b3f5849efabb
with_terminal() do
	@btime step1()
	@btime step2()
	@btime step3()
	@btime step4()
end

# ╔═╡ ea9e77e2-91b1-11eb-185d-cd006db11f60
md"""
## Trajectory of a random walk
"""

# ╔═╡ 12b4d528-9239-11eb-2824-8ddb5e2ba892
md"""
We can now calculate the **trajectory** of a 1D random walk as it takes several steps.
It starts at an initial position, for example, 0, and takes consecutive steps:
"""

# ╔═╡ 2f525796-9239-11eb-1865-9b01eadcf548
function walk1D(N)
	x = 0
	xs = [x]

	for i in 1:N
		x += step1()
		push!(xs, x)
	end

	return xs
end

# ╔═╡ 51abfe6e-9239-11eb-362a-259570250663
begin
	plot()

	for i in 1:10
		plot!(walk1D(100), leg=false, size=(500, 300), lw=2, alpha=0.5)
	end

	plot!()
end

# ╔═╡ b847b5ca-9239-11eb-02fe-db4d9625bc5f
md"""
# Making it more general: Random walks using types
"""

# ╔═╡ c2deb090-9239-11eb-0739-a74379c15ce6
md"""
Now suppose we want to think about more general random walks, for example moving around in 2D. Then we need to *generalise* the above function.

Based on our experience from last time, you should suspect that a good way to do this is with *types*. We will define
"""

# ╔═╡ d420d492-91d9-11eb-056d-33cc8f0aed74
abstract type Walker end

# ╔═╡ ad2d4dd8-91d5-11eb-27af-6f0c6e61a86a
struct Walker1D <: Walker
	pos::Int
end

# ╔═╡ d0f81f28-91d9-11eb-2e79-61461ef5b132
position(w::Walker) = w.pos

# ╔═╡ b8f2c508-91d5-11eb-31b5-61810f171270
step(w::Walker1D) = rand( (-1, +1) )

# ╔═╡ 3c3971e2-91da-11eb-384c-01c627318bdc
update(w::W, step) where {W <: Walker} = W(position(w) + step)

# ╔═╡ cb0ef266-91d5-11eb-314b-0545c0c817d0
function trajectory(w::W, N) where {W}   # W is a type parameter
	ws = [position(w)]

	for i in 1:N
		pos = position(w)
		w = update(w, step(w))

		push!(ws, position(w))
	end

	return ws
end

# ╔═╡ 048fac02-91da-11eb-0d26-4f258b4cd043
trajectory(Walker1D(0), 10)

# ╔═╡ 23b84ce2-91da-11eb-01f8-c308ac4d1c7a
struct Walker2D <: Walker
	x::Int
	y::Int
end

# ╔═╡ 537f952a-91da-11eb-33cf-6be2fd3bc45c
position(w::Walker2D) = (w.x, w.y)

# ╔═╡ 5b972296-91da-11eb-29b1-074f3926181e
step(w::Walker2D) = rand( [ [1, 0], [0, 1], [-1, 0], [0, -1] ] )

# ╔═╡ 3ad5a93c-91db-11eb-3227-c96bf8fd2206
update(w::Walker2D, step::Vector) = Walker2D(w.x + step[1], w.y + step[2])

# ╔═╡ 74182fe0-91da-11eb-219a-01f13b86406d
traj = trajectory(Walker2D(0, 0), 10^N)

# ╔═╡ 57972a32-91e5-11eb-1d62-fbc22c494db9
md"""
## Random walks as sum of random variables
"""

# ╔═╡ 63122dd0-91e5-11eb-34c3-b1b5c87809b8
md"""
We can connect up with the discussion from last lecture about the sum of random variables: The position $S_n$ of a random walk at time $n$ is a random variable that is the sum of $n$ random variables (which are often independent and id

$$S_n = X_1 + \cdots + X_n$$

where $X_i$ is the random variable modelling the $i$th step of the walk.

Often each $X_i$ will be independent random variables of the same type, in which case we talk about an **independent and identically distributed** (IID) collection.

In principle we could use the method from last lecture to model this, but there is a difficulty since we are interested in all intermediate steps: The values $S_{n-1}$ and $S_n$, for example, are *not* independent (even though the $X_i$ are). For example, if $S_{n-1}$ is large, then $S_n$ is also about as large, so they are highly correlated.
"""

# ╔═╡ bd013582-91e7-11eb-2d79-e18e45f6d639
md"""
## Cumulative sum
"""

# ╔═╡ c092ab72-91e7-11eb-2bd1-13c8c8bb30e4
md"""
Suppose we generate steps $X_i$, e.g.
"""

# ╔═╡ cfb0f9ba-91e7-11eb-26b5-5f0f59d03cff
steps = rand( (-1, +1), 10 )

# ╔═╡ d6edb326-91e7-11eb-03b1-d93e3bc83ca6
md"""
The trajectory, or sample path, of the random walk, is the collection of positions over time. These are the **partial sums**

$$\begin{align}
S_1 & = X_1 \\
S_2 &= X_1 + X_2 \\
S_3 &= X_1 + X_2 + X_3
\end{align}$$

etc.
"""

# ╔═╡ 4ef42cec-91e8-11eb-2976-0950ffe5de6c
md"""
We can calculate all of these values using the `cumsum` ("cumulative sum", also called "prefix sum" or "scan") function:
"""

# ╔═╡ ace8658e-91e8-11eb-0b9d-4b759635e417
cumsum(steps)

# ╔═╡ b049ff58-91e8-11eb-203b-4f4b5ee5f01f
md"""
Let's plot this:
"""

# ╔═╡ b6775d3a-91e8-11eb-0187-618cb538d142
plot(cumsum(steps), m=:o, leg=false, size=(500, 300))

# ╔═╡ 8b1441b6-91e4-11eb-16b2-d7eadd3fd69c
md"""
# Trajectories vs evolving probability distributions
"""

# ╔═╡ d1315f94-91e4-11eb-1076-81156e24d2f1
md"""
So far we have looked at single trajectories of random walks. We can think of this as the equivalent of sampling using `rand`. Suppose that we were to (conceptually) run millions or billions of trajectories of our random walk. At each *time* step, we can ask for the probability distribution in *space*, averaged over all those walks.
Note that we now have *two* variables. We could do this by literally running many walks and computing histograms. Instead we can look at how probabilities move around.

Let's call $p_i^t$ the probability of being at site $i$ at time $t$. We can then ask what the probability of being at site $i$ at time $t+1$. It must have been in one of the neighbouring sites at time $t$, so

$$p_i^{t+1} = \textstyle \frac{1}{2} (p_{i-1}^{t} + p_{i+1}^{t}).$$

This is sometimes called the **master equation** (even though that is rather a useless, non-descriptive name). It describes how the probability distribution
"""

# ╔═╡ 97c958c2-91eb-11eb-17cb-410acc7f3e49
md"""
At time $t$ we have a whole probability distribution, which we can think of as a vector $\mathbf{p}^t$.

Let's write a function to **evolve** the vector to the next time step.
"""

# ╔═╡ b7a98dba-91eb-11eb-3c78-074aa835b5fb
function evolve(p)
	p′ = similar(p)   # make a vector of the same length and type
	                  # to store the probability vector at the next time step

	for i in 2:length(p)-1
		p′[i] = 0.5 * (p[i-1] + p[i+1])
	end

	p′[1] = 0
	p′[end] = 0

	return p′

end

# ╔═╡ f7f17c0c-91eb-11eb-3fa5-3bd90bb7044e
md"""
Wait... Do you recognise this?

We have seen this before! This is just a **convolution**, like the blurring kernel from image processing (except that in our particular model `p[i]` itself does not contribute to the value of `new_p[i]`.)
"""

# ╔═╡ f0aed302-91eb-11eb-13fb-d9418ef327a8
md"""
Note that just as we did with images, we have a problem at the edges of the system. We need to specify **boundary conditions**. For now we have put 0s at the boundary. This corresponds to probability *escaping* at the boundary: any probability that arrives at the boundary (i.e. the first or last cell) at a given time step does not stay in the system. We can think of the probability as analogous to a chemical that disappears
"""

# ╔═╡ c1062e00-922b-11eb-1f31-ddbd03f8f986
md"""
We also need to specify an *initial* condition $\mathbf{p}_0$. This tells us where our walker is at time $0$. Suppose that all walkers start at $0$. We will place this in the middle of our vector. Then the probability at that location is 1, while the probability elsewhere is 0:
"""

# ╔═╡ 547188ea-9233-11eb-1a89-5ff9468b31f7
function initial_condition(n)

	p0 = zeros(n)
	p0[n ÷ 2 + 1] = 1

	return p0
end

# ╔═╡ 2920abfe-91ec-11eb-19bc-935fa1ba0a96
md"""
Now let's try and visualise the time evolution.
"""

# ╔═╡ 3fadf88c-9236-11eb-19fa-d191ac5a6191
function time_evolution(p0, N)
	ps = [p0]
	p = p0

	for i in 1:N
		p = evolve(p)
		push!(ps, copy(p))
	end

	return ps
end

# ╔═╡ 58653a70-9236-11eb-3dae-47adc2a77cb4
p0 = initial_condition(101)

# ╔═╡ 5d02f21e-9236-11eb-26ea-6593aa80a2eb
ps = time_evolution(p0, 100)

# ╔═╡ 863caf0a-9236-11eb-1013-ab7d83f3fc0c
ps[2]

# ╔═╡ b803406e-9236-11eb-3aad-056b7f2c9b4b
md"""
t = $(@bind tt Slider(1:length(ps), show_value=true, default=1))
"""

# ╔═╡ cc7aaeea-9236-11eb-3fad-2b5ad3962ec1
plot(ps[tt], ylim=(0, 1), leg=false, size=(500, 300))

# ╔═╡ dabb5766-9236-11eb-3be9-9b33ba5af68a
ps[tt]

# ╔═╡ 6cde6ef4-9236-11eb-219a-4d20adaf9988
M = reduce(hcat, ps)'

# ╔═╡ 7e8c1a2a-9236-11eb-20e9-57f6601f5472
heatmap(M, yflip=true)

# ╔═╡ Cell order:
# ╟─3649f170-923a-11eb-321c-cf95849cc044
# ╠═97e807b2-9237-11eb-31ef-6fe0d4cc94d3
# ╠═5f0d7a44-91e0-11eb-10ae-d73156f965e6
# ╟─9647147a-91ab-11eb-066f-9bc190368fb2
# ╟─ff1aca1e-91e7-11eb-343e-0f89d9570b06
# ╟─66a2f510-9232-11eb-3be9-131febc0039f
# ╟─bd3170e6-91ae-11eb-06f8-ebb6b2e7869f
# ╟─a304c842-91df-11eb-3fac-6dd63087f6de
# ╟─798507d6-91db-11eb-2e4a-3ba02f12ba65
# ╟─3504168a-91de-11eb-181d-1d580d5dc071
# ╟─4c8d8294-91db-11eb-353d-c3696c615b3d
# ╟─b62c4af8-9232-11eb-2f66-dd27dcb87d20
# ╟─905379ce-91ad-11eb-295d-8354ecf5c5b1
# ╟─5c4f0f26-91ad-11eb-033b-2bd221f0bdba
# ╟─da98b676-91e0-11eb-0d97-57b8a8aadf2a
# ╟─e0b607c0-91e0-11eb-10aa-53ec33570e59
# ╟─fa1635d4-91e3-11eb-31bd-cf61c502ad35
# ╠═f7f9e4c6-91e3-11eb-1a56-8b98f0b09b46
# ╠═5da7b076-91b4-11eb-3eba-b3f5849efabb
# ╟─ea9e77e2-91b1-11eb-185d-cd006db11f60
# ╟─12b4d528-9239-11eb-2824-8ddb5e2ba892
# ╠═2f525796-9239-11eb-1865-9b01eadcf548
# ╠═51abfe6e-9239-11eb-362a-259570250663
# ╟─b847b5ca-9239-11eb-02fe-db4d9625bc5f
# ╟─c2deb090-9239-11eb-0739-a74379c15ce6
# ╠═d420d492-91d9-11eb-056d-33cc8f0aed74
# ╠═ad2d4dd8-91d5-11eb-27af-6f0c6e61a86a
# ╠═d0f81f28-91d9-11eb-2e79-61461ef5b132
# ╠═b8f2c508-91d5-11eb-31b5-61810f171270
# ╠═3c3971e2-91da-11eb-384c-01c627318bdc
# ╠═cb0ef266-91d5-11eb-314b-0545c0c817d0
# ╠═048fac02-91da-11eb-0d26-4f258b4cd043
# ╠═23b84ce2-91da-11eb-01f8-c308ac4d1c7a
# ╠═537f952a-91da-11eb-33cf-6be2fd3bc45c
# ╠═5b972296-91da-11eb-29b1-074f3926181e
# ╠═3ad5a93c-91db-11eb-3227-c96bf8fd2206
# ╠═74182fe0-91da-11eb-219a-01f13b86406d
# ╟─57972a32-91e5-11eb-1d62-fbc22c494db9
# ╟─63122dd0-91e5-11eb-34c3-b1b5c87809b8
# ╟─bd013582-91e7-11eb-2d79-e18e45f6d639
# ╟─c092ab72-91e7-11eb-2bd1-13c8c8bb30e4
# ╠═cfb0f9ba-91e7-11eb-26b5-5f0f59d03cff
# ╟─d6edb326-91e7-11eb-03b1-d93e3bc83ca6
# ╟─4ef42cec-91e8-11eb-2976-0950ffe5de6c
# ╠═ace8658e-91e8-11eb-0b9d-4b759635e417
# ╟─b049ff58-91e8-11eb-203b-4f4b5ee5f01f
# ╠═b6775d3a-91e8-11eb-0187-618cb538d142
# ╟─8b1441b6-91e4-11eb-16b2-d7eadd3fd69c
# ╟─d1315f94-91e4-11eb-1076-81156e24d2f1
# ╟─97c958c2-91eb-11eb-17cb-410acc7f3e49
# ╠═b7a98dba-91eb-11eb-3c78-074aa835b5fb
# ╟─f7f17c0c-91eb-11eb-3fa5-3bd90bb7044e
# ╟─f0aed302-91eb-11eb-13fb-d9418ef327a8
# ╟─c1062e00-922b-11eb-1f31-ddbd03f8f986
# ╠═547188ea-9233-11eb-1a89-5ff9468b31f7
# ╟─2920abfe-91ec-11eb-19bc-935fa1ba0a96
# ╠═3fadf88c-9236-11eb-19fa-d191ac5a6191
# ╠═58653a70-9236-11eb-3dae-47adc2a77cb4
# ╠═5d02f21e-9236-11eb-26ea-6593aa80a2eb
# ╠═863caf0a-9236-11eb-1013-ab7d83f3fc0c
# ╟─b803406e-9236-11eb-3aad-056b7f2c9b4b
# ╠═cc7aaeea-9236-11eb-3fad-2b5ad3962ec1
# ╠═dabb5766-9236-11eb-3be9-9b33ba5af68a
# ╠═6cde6ef4-9236-11eb-219a-4d20adaf9988
# ╠═7e8c1a2a-9236-11eb-20e9-57f6601f5472


mutable struct BatchTD <: Optimizer
    α::Float64
    α_mod::Float64
    avg::Bool
    BatchTD(α, α_mod, avg) = new(α, α_mod, avg)
end

function update!(value::VFunction, opt::BatchTD, s_t, s_tp1, r, γ, ρ, terminal)
    α = opt.α*opt.α_mod
    if opt.avg
        α = α/length(ρ)
    end
    δ = r + γ.*(dot.(s_tp1, [value.weights])) - dot.(s_t, [value.weights])
    Δθ = α*sum(ρ.*(δ.*s_t))
    value.weights .+= Δθ
end

mutable struct BatchTDC <: Optimizer
    α::Float64
    β::Float64
    α_mod::Float64
    β_mod::Float64
    avg::Bool
    h::IdDict
    # TDC(α, β) = new(α, β, 1.0, 1.0, true)
    # TDC(α, β, α_mod, β_mod, avg) = new(α, β, α_mod, β_mod, avg)
end

function update!(value::VFunction, opt::BatchTDC, s_t, s_tp1, r, γ, ρ, terminal)
    # α = TD.α*TD.α_mod
    # β = TD.β*TD.β_mod

    h = get!(opt.h, value, zero(value.weight))

    α = opt.α*opt.α_mod
    β = opt.β*opt.β_mod
    if opt.avg
        α = α/length(ρ)
        β = β/length(ρ)
    end

    δ = r + γ.*(dot.(s_tp1, [value.weights])) - dot.(s_t, [value.weights])
    Δθ = α*sum(ρ.*(δ.*s_t - γ.*dot.(s_t, [value.h]).*s_tp1))
    Δh = β*sum(s_t.*(ρ.*δ - dot.(s_t, [value.h])))

    value.weights .+= Δθ
    value.h .+= Δh
end

mutable struct WISBatchTDC <: Optimizer
    α::Float64
    β::Float64
    α_mod::Float64
    β_mod::Float64
    WISTDC(α, β) = new(α, β, 1.0, 1.0)
end

function update!(value::VFunction, opt::WISBatchTDC, s_t, s_tp1, r, γ, ρ, terminal)
    δ = r + γ.*(dot.(s_tp1, [value.weights])) - dot.(s_t, [value.weights])
    Δθ = α*sum(ρ.*(δ.*s_t - γ.*dot.(s_t, [value.h]).*s_tp1))
    Δh = β*sum(s_t.*(ρ.*δ - dot.(s_t, [value.h])))

    value.weights .+= Δθ./sum(ρ)
    value.h .+= Δh./sum(ρ)
end

mutable struct BatchTDC2 <: Optimizer
    α::Float64
    β::Float64
    α_mod::Float64
    β_mod::Float64
    avg::Bool
    TDC2(α, β) = new(α, β, 1.0, 1.0, true)
end

function update!(value::VFunction, opt::BatchTDC2, s_t, s_tp1, r, γ, ρ, terminal)
    # Python code sample = (prev_phi.copy(), phi.copy(), action, reward, state, prev_state, pi[action]/mu[action])
    # td_error = (sample[3] + gamma*(sample[1].dot(weights_tdc_iwer)) - sample[0].dot(weights_tdc_iwer))
    # weights_tdc_iwer = weights_tdc_iwer + alpha*sample[-1]*(td_error*sample[0] - gamma*sample[1]*(sample[0].dot(h_tdc_iwer)))
    # h_tlndc_iwer = h_tdc_iwer + alpha_h*(sample[-1]*td_error - sample[0].dot(h_tdc_iwer))*sample[0]
    α = opt.α*opt.α_mod
    β = opt.β*opt.β_mod
    if opt.avg
        α = α/length(ρ)
        β = β/length(ρ)
    end

    δ = r + γ.*(dot.(s_tp1, [value.weights])) - dot.(s_t, [value.weights])
    Δθ = α*sum(ρ.*(δ.*s_t - γ.*dot.(s_t, [value.h]).*s_tp1))
    Δh = β*sum(s_t.*(ρ.*(δ - dot.(s_t, [value.h]))))

    value.weights .+= Δθ
    value.h .+= Δh
end

mutable struct BatchGTD2 <: Optimizer
    α::Float64
    β::Float64
    α_mod::Float64
    β_mod::Float64
    avg::Bool
    GTD2(α, β) = new(α, β, 1.0, 1.0, true)
    GTD2(α, β, α_mod, β_mod, avg) = new(α, β, α_mod, β_mod, avg)
end

function update!(value::VFunction, opt::BatchGTD2, s_t, s_tp1, r, γ, ρ, terminal)
    # Python code sample = (prev_phi.copy(), phi.copy(), action, reward, state, prev_state, pi[action]/mu[action])
    # td_error = (sample[3] + gamma*(sample[1].dot(weights_gtd2_iwer)) - sample[0].dot(weights_gtd2_iwer))
    # weights_gtd2_iwer = weights_gtd2_iwer + alpha*sample[-1]*(sample[0] - gamma*sample[1])*(sample[0].dot(h_gtd2_iwer))
    # h_gtd2_iwer = h_gtd2_iwer + alpha_h*(sample[-1]*td_error - sample[0].dot(h_gtd2_iwer))*sample[0]
    α = opt.α*opt.α_mod
    β = opt.β*opt.β_mod
    if opt.avg
        α = α/length(ρ)
        β = β/length(ρ)
    end

    δ = r + γ.*(dot.(s_tp1, [value.weights])) - dot.(s_t, [value.weights])
    Δθ = α*sum(ρ.*(s_t - γ.*s_tp1).*(dot.(s_t, [value.h])))
    Δh = β*sum(s_t.*(ρ.*δ - dot.(s_t, [value.h])))

    value.weights .+= Δθ
    value.h .+= Δh
end