##################
## LLA <-> ECEF ##
##################

struct LLAfromECEF <: Transformation
    el::Ellipsoid

    a::Float64      # major axis
    f::Float64      # flattening
    e2::Float64     # Eccentricity squared
    e2m::Float64    # 1 - e2
    e2a::Float64    # |e2|
    e4a::Float64    # e2^2

    function LLAfromECEF(el,a,f,e2,e2m,e2a,e4a)
        if !(isfinite(a) && a > 0)
            error("Major radius is not positive")
        end
        if !(isfinite(f) && f < 1)
            error("Minor radius is not positive")
        end
        return new(el,a,f,e2,e2m,e2a,e4a)
    end
end
Base.show(io::IO, trans::LLAfromECEF) = print(io, "LLAfromECEF($(trans.el))")

"""
    LLAfromECEF(ellipsoid)

Construct a `Transformation` object to convert from ECEF coordinates
to LLA coordinates. Pre-caches ellipsoidal parameters for efficiency.
"""
function LLAfromECEF(el::Ellipsoid)
    a = el.a
    b = el.b
    f = el.f
    e2 = f*(2-f)
    e2m = (1-f)*(1-f)  #1 - e2
    e2a = abs(e2)
    e4a = e2*e2

    return LLAfromECEF(el, a, f, e2, e2m, e2a, e4a)
end

LLAfromECEF(datum::Datum) = LLAfromECEF(ellipsoid(datum))

function (trans::LLAfromECEF)(ecef::ECEF)
    # Ported to Julia by Andy Ferris, 2016 and re-released under MIT license.
    #/**
    # * \file Geocentric.cpp
    # * \brief Implementation for GeographicLib::Geocentric class
    # *
    # * Copyright (c) Charles Karney (2008-2015) <charles@karney.com> and licensed
    # * under the MIT/X11 License.  For more information, see
    # * http://geographiclib.sourceforge.net/
    # **********************************************************************/
    R = hypot(ecef.x, ecef.y)
    if R == 0
        slam = 0.0
        clam = 1.0
    else
        slam = ecef.y / R
        clam = ecef.x / R
    end
    h = hypot(R, ecef.z)    # Distance to center of earth

    if (trans.e4a == 0)
        # Treat the spherical case.  Dealing with underflow in the general case
        # with _e2 = 0 is difficult.  Origin maps to north pole, same as the
        # ellipsoidal case below
        if h == 0
            sphi = 1.0
            cphi = 0.0
        else
            sphi = ecef.z / h
            cphi = R / h
        end
        h -= trans.a
    else # Ellipsoidal
        # Treat prolate spheroids by swapping r and z here and by switching
        # the arguments to phi = atan(...) at the end.
        p = (R / trans.a) * (R / trans.a)
        q = trans.e2m * (ecef.z / trans.a) * (ecef.z / trans.a)
        r = (p + q - trans.e4a) / 6
        if (trans.f < 0)
            tmp = p
            p = q
            q = tmp
        end

        if ( !(trans.e4a * q == 0 && r <= 0) )

            # Avoid possible division by zero when r = 0 by multiplying
            # equations for s and t by r^3 and r, resp.
            S = trans.e4a * p * q / 4 # S = r^3 * s
            r2 = r * r
            r3 = r * r2
            disc = S * (2 * r3 + S)
            u = r
            if (disc >= 0)
                T3 = S + r3
                # Pick the sign on the sqrt to maximize abs(T3).  This minimizes
                # loss of precision due to cancellation.  The result is unchanged
                # because of the way the T is used in definition of u.
                T3 += (T3 < 0 ? -sqrt(disc) : sqrt(disc)) # T3 = (r * t)^3
                # N.B. cbrt always returns the real root.  cbrt(-8) = -2.
                T = cbrt(T3) # T = r * t
                # T can be zero; but then r2 / T -> 0.
                u += T + (T != 0 ? r2 / T : 0.0)
            else
                # T is complex, but the way u is defined the result is real.
                ang = atan(sqrt(-disc), -(S + r3))
                # There are three possible cube roots.  We choose the root which
                # avoids cancellation.  Note that disc < 0 implies that r < 0.
                u += 2 * r * cos(ang / 3)
            end

            v = sqrt(u*u + trans.e4a * q) # guaranteed positive
            # Avoid loss of accuracy when u < 0.  Underflow doesn't occur in
            # e4 * q / (v - u) because u ~ e^4 when q is small and u < 0.
            uv = (u < 0 ? trans.e4a * q / (v - u) : u + v) # u+v, guaranteed positive
            # Need to guard against w going negative due to roundoff in uv - q.
            w = max(0.0, trans.e2a * (uv - q) / (2 * v))
            # Rearrange expression for k to avoid loss of accuracy due to
            # subtraction.  Division by 0 not possible because uv > 0, w >= 0.
            k = uv / (sqrt(uv + w*w) + w)
            k1 = (trans.f >= 0 ? k : k - trans.e2)
            k2 = (trans.f >= 0 ? k + trans.e2 : k)
            d = k1 * R / k2
            H = hypot(ecef.z/k1, R/k2)
            sphi = (ecef.z/k1) / H
            cphi = (R/k2) / H
            h = (1 - trans.e2m/k1) * hypot(d, ecef.z)
        else  # e4 * q == 0 && r <= 0
            # This leads to k = 0 (oblate, equatorial plane) and k + e^2 = 0
            # (prolate, rotation axis) and the generation of 0/0 in the general
            # formulas for phi and h.  using the general formula and division by 0
            # in formula for h.  So handle this case by taking the limits:
            # f > 0: z -> 0, k      ->   e2 * sqrt(q)/sqrt(e4 - p)
            # f < 0: R -> 0, k + e2 -> - e2 * sqrt(q)/sqrt(e4 - p)
            zz = sqrt((trans.f >= 0 ? trans.e4a - p : p) / trans.e2m)
            xx = sqrt( trans.f <  0 ? trans.e4a - p : p)
            H = hypot(zz, xx)
            sphi = zz / H
            cphi = xx / H
            if (ecef.z < 0)
                sphi = -sphi # for tiny negative Z (not for prolate)
            end
            h = - trans.a * (trans.f >= 0 ? trans.e2m : 1.0) * H / trans.e2a
        end
    end
    lat = 180 * atan(sphi, cphi) / pi
    lon = 180 * atan(slam, clam) / pi

    return LLA(lat, lon, h)
end

"""
    ECEFfromLLA(ellipsoid_or_datum)

Construct a `Transformation` object to convert from LLA coordinates
to ECEF coordinates.
"""
struct ECEFfromLLA <: Transformation
    el::Ellipsoid
end
Base.show(io::IO, trans::ECEFfromLLA) = print(io, "ECEFfromLLA($(trans.el))")

ECEFfromLLA(datum::Datum) = ECEFfromLLA(ellipsoid(datum))

function (trans::ECEFfromLLA)(lla::LLA)
    ϕdeg, λdeg, h = lla.lat, lla.lon, lla.alt

    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)
    sinλ, cosλ = sind(λdeg), cosd(λdeg)

    N = trans.el.a / sqrt(1 - trans.el.e2 * sinϕ^2)  # Radius of curvature (meters)

    x = (N + h) * cosϕ * cosλ
    y = (N + h) * cosϕ * sinλ
    z = (N * (1 - trans.el.e2) + h) * sinϕ

    return ECEF(x, y, z)
end

Base.inv(trans::LLAfromECEF) = ECEFfromLLA(trans.el)
Base.inv(trans::ECEFfromLLA) = LLAfromECEF(trans.el)

##################
## ENU <-> NED ##
##################

"""
    ENUfromNED(ned::NED)

Construct a `Transformation` object to convert from local `NED` coordinates
to local `ENU` coordinates centered at the same origin. This is a simple
permutation of coordinates and sign change for the altitude.
"""
struct ENUfromNED <: Transformation end     # singleton type

Base.show(io::IO, ::ENUfromNED) = print(io, "ENUfromNED()")

function (::ENUfromNED)(ned::NED)
    ENU(ned.e, ned.n, -ned.d)
end

"""
    NEDfromENU(enu::ENU)

Construct a `Transformation` object to convert from local `ENU` coordinates
to local `NED` coordinates centered at the same origin. This is a simple
permutation of coordinates and sign change for the altitude.
"""
struct NEDfromENU <: Transformation end     # singleton type

Base.show(io::IO, ::NEDfromENU) = print(io, "NEDfromENU()")

function (::NEDfromENU)(enu::ENU)
    NED(enu.n, enu.e, -enu.u)
end

Base.inv(::ENUfromNED) = NEDfromENU()
Base.inv(::NEDfromENU) = ENUfromNED()

##################
## ECEF <-> ENU ##
##################

"""
    ENUfromECEF(origin, datum)
    ENUfromECEF(origin::UTM, zone, isnorth, datum)
    ENUfromECEF(origin::ECEF, lat, lon)

Construct a `Transformation` object to convert from global `ECEF` coordinates
to local `ENU` coordinates centered at the `origin`. This object pre-caches both the
ECEF coordinates and latitude and longitude of the origin for maximal efficiency.
"""
struct ENUfromECEF{T} <: Transformation
    origin::ECEF{T}
    lat::T
    lon::T
end
ENUfromECEF(origin::LLA, datum) = ENUfromECEF(ECEFfromLLA(datum)(origin), origin.lat, origin.lon)
function ENUfromECEF(origin::ECEF, datum)
    origin_lla = LLAfromECEF(datum)(origin)
    ENUfromECEF(origin, origin_lla.lat, origin_lla.lon)
end
Base.show(io::IO, trans::ENUfromECEF) = print(io, "ENUfromECEF($(trans.origin), lat=$(trans.lat)°, lon=$(trans.lon)°)")
Base.isapprox(t1::ENUfromECEF, t2::ENUfromECEF; kwargs...) = isapprox(t1.origin, t2.origin; kwargs...) && isapprox(t1.lat, t2.lat; kwargs...) && isapprox(t1.lon, t2.lon; kwargs...)

function (trans::ENUfromECEF)(ecef::ECEF)
    ϕdeg, λdeg = trans.lat, trans.lon

    ∂x = ecef.x - trans.origin.x
    ∂y = ecef.y - trans.origin.y
    ∂z = ecef.z - trans.origin.z

    # Compute rotation matrix
    sinλ, cosλ = sind(λdeg), cosd(λdeg)
    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)

    # R = [     -sinλ       cosλ  0.0
    #      -cosλ*sinϕ -sinλ*sinϕ cosϕ
    #       cosλ*cosϕ  sinλ*cosϕ sinϕ]
    #
    # east, north, up = R * [∂x, ∂y, ∂z]
    east  = ∂x * -sinλ      + ∂y * cosλ       + ∂z * 0.0
    north = ∂x * -cosλ*sinϕ + ∂y * -sinλ*sinϕ + ∂z * cosϕ
    up    = ∂x * cosλ*cosϕ  + ∂y * sinλ*cosϕ  + ∂z * sinϕ

    return ENU(east, north, up)
end

"""
    ECEFfromENU(origin, datum)
    ECEFfromENU(origin::UTM, zone, isnorth, datum)
    ECEFfromENU(origin::ECEF, lat, lon)

Construct a `Transformation` object to convert from local `ENU` coordinates
centred at `origin` to global `ECEF` coodinates. This object pre-caches both the
ECEF coordinates and latitude and longitude of the origin for maximal efficiency.
"""
struct ECEFfromENU{T} <: Transformation
    origin::ECEF{T}
    lat::T
    lon::T
end
ECEFfromENU(origin::LLA, datum) = ECEFfromENU(ECEFfromLLA(datum)(origin), origin.lat, origin.lon)
function ECEFfromENU(origin::ECEF, datum)
    origin_lla = LLAfromECEF(datum)(origin)
    ECEFfromENU(origin, origin_lla.lat, origin_lla.lon)
end
Base.show(io::IO, trans::ECEFfromENU) = print(io, "ECEFfromENU($(trans.origin), lat=$(trans.lat)°, lon=$(trans.lon)°)")
Base.isapprox(t1::ECEFfromENU, t2::ECEFfromENU; kwargs...) = isapprox(t1.origin, t2.origin; kwargs...) && isapprox(t1.lat, t2.lat; kwargs...) && isapprox(t1.lon, t2.lon; kwargs...)

function (trans::ECEFfromENU)(enu::ENU)
    ϕdeg, λdeg = trans.lat, trans.lon

    # Compute rotation matrix
    sinλ, cosλ = sind(λdeg), cosd(λdeg)
    sinϕ, cosϕ = sind(ϕdeg), cosd(ϕdeg)

    # Rᵀ = [-sinλ -cosλ*sinϕ cosλ*cosϕ
    #        cosλ -sinλ*sinϕ sinλ*cosϕ
    #         0.0       cosϕ      sinϕ]
    # Δx, Δy, Δz = Rᵀ * [east, north, up]
    Δx = -sinλ * enu.e + -cosλ*sinϕ * enu.n + cosλ*cosϕ * enu.u
    Δy =  cosλ * enu.e + -sinλ*sinϕ * enu.n + sinλ*cosϕ * enu.u
    Δz =   0.0 * enu.e +       cosϕ * enu.n +      sinϕ * enu.u

    X = trans.origin.x + Δx
    Y = trans.origin.y + Δy
    Z = trans.origin.z + Δz

    return ECEF(X,Y,Z)
end

Base.inv(trans::ECEFfromENU) = ENUfromECEF(trans.origin, trans.lat, trans.lon)
Base.inv(trans::ENUfromECEF) = ECEFfromENU(trans.origin, trans.lat, trans.lon)

##################
## ECEF <-> NED ##
##################

"""
    NEDfromECEF(origin, datum)
    NEDfromECEF(origin::UTM, zone, isnorth, datum)
    NEDfromECEF(origin::ECEF, lat, lon)

Construct a composite transformation NEDfromENU() ∘ ENUfromECEF(origin, datum)
to convert from global `ECEF` coordinates to local `NED` coordinates centered at the `origin`.
This object pre-caches both the ECEF coordinates and latitude and longitude of the origin for maximal efficiency.
"""
NEDfromECEF(origin, datum) = NEDfromENU() ∘ ENUfromECEF(origin, datum)

"""
    ECEFfromNED(origin, datum)
    ECEFfromNED(origin::UTM, zone, isnorth, datum)
    ECEFfromNED(origin::ECEF, lat, lon)

Construct a composite transformation ECEFfromENU(origin,datum) ∘ ENUfromNED()
to convert from local `NED` coordinates centred at `origin` to global `ECEF` coodinates.
This object pre-caches both the ECEF coordinates and latitude and longitude of the origin for maximal efficiency.
"""
ECEFfromNED(origin, datum) = ECEFfromENU(origin,datum) ∘ ENUfromNED()

#################
## LLA <-> ENU ##
#################

"""
    ENUfromLLA(origin, datum)

Creates composite transformation `ENUfromECEF(origin, datum) ∘ ECEFfromLLA(datum)`.
"""
ENUfromLLA(origin, datum) = ENUfromECEF(origin, datum) ∘ ECEFfromLLA(datum)

"""
    LLAfromENU(origin, datum)

Creates composite transformation `LLAfromECEF(datum) ∘ ECEFfromENU(origin, datum)`.
"""
LLAfromENU(origin, datum) = LLAfromECEF(datum) ∘ ECEFfromENU(origin, datum)

#################
## LLA <-> NED ##
#################

"""
    NEDfromLLA(origin, datum)

Creates composite transformation `NEDfromECEF(origin, datum) ∘ ECEFfromLLA(datum)`.
"""
NEDfromLLA(origin, datum) = NEDfromECEF(origin, datum) ∘ ECEFfromLLA(datum)

"""
    LLAfromNED(origin, datum)

Creates composite transformation `LLAfromECEF(datum) ∘ ECEFfromNED(origin, datum)`.
"""
LLAfromNED(origin, datum) = LLAfromECEF(datum) ∘ ECEFfromNED(origin, datum)

#################
## LLA <-> UTM ##
#################

"""
    LLAfromUTM(zone, isnorth::Bool, datum)

Construct a `Transformation` object to convert from `UTM` coordinates in
the specified zone and hemisphere (`isnorth = true` or `false`) to global `LLA`
coordinates. Pre-caches ellipsoidal parameters for efficiency and performs
Charles Karney's accurate 6th-order series expansion algorithm.

(See also `LLAfromUTMZ`)
"""
struct LLAfromUTM{Datum,Order} <: Transformation
    zone::UInt8
    isnorth::Bool # hemisphere
    tm::TransverseMercator{Order}
    datum::Datum
end
LLAfromUTM(zone::UInt8, h, d) = LLAfromUTM(UInt8(zone), h, TransverseMercator(d), d)
LLAfromUTM(zone::Integer, h, d) = LLAfromUTM(UInt8(zone), h, d)
Base.show(io::IO, trans::LLAfromUTM) = print(io, "LLAfromUTM(zone=$(trans.zone == 0 ? "polar" : trans.zone) ($(trans.isnorth ? "north" : "south")), $(trans.datum))")

function (trans::LLAfromUTM)(utm::UTM)
    if trans.zone == 0
        # Do inverse steriographic projection
        k0 = 0.994
        x = (utm.x - 2e6)
        y = (utm.y - 2e6)
        (lat,lon,γ,k) = polarst_inv(trans.isnorth, k0, trans.tm, x, y)

        return LLA(lat, lon, utm.z)
    else
        lon_ref = Float64(utm_meridian(trans.zone))
        k0 = 0.9996 # Horizontal scaling factor
        x = (utm.x - 5e5) # Convention has 500km offset for easting
        y = (utm.y - (trans.isnorth ? 0.0 : 1e7)) # Northing offset for southern hemisphere
        #(lat,lon,k,γ) = transverse_mercator_inv(lat_ref, lon_ref, x, y, ellipsoid(trans.datum))
        (lat,lon,γ,k) = transverse_mercator_reverse(lon_ref, x, y, k0, trans.tm)
        return LLA(lat, lon, utm.z) # Note: scaling not applied to vertical dimension
    end
end

"""
    UTMfromLLA(zone, isnorth::Bool, datum)

Construct a `Transformation` object to convert from global `LLA` coordinates
to `UTM` coordinates in the specified zone and hemisphere (`isnorth = true` or `false`). Pre-caches
ellipsoidal parameters for efficiency and performs Charles Karney's accurate
6th-order series expansion algorithm.

(See also `UTMZfromLLA`)
"""
struct UTMfromLLA{Datum,Order} <: Transformation
    zone::UInt8
    isnorth::Bool # true = north, false = south
    tm::TransverseMercator{Order}
    datum::Datum
end
UTMfromLLA(zone::UInt8, h, d) = UTMfromLLA(zone, h, TransverseMercator(d), d)
UTMfromLLA(zone::Integer, h, d) = UTMfromLLA(UInt8(zone), h, d)
Base.show(io::IO, trans::UTMfromLLA) = print(io, "UTMfromLLA(zone=$(trans.zone == 0 ? "polar" : trans.zone) ($(trans.isnorth ? "north" : "south")), $(trans.datum))")

function (trans::UTMfromLLA)(lla::LLA)
    if trans.zone == 0
        # Do polar steriographic projection
        k0 = 0.994
        (x,y,γ,k) = polarst_fwd(trans.isnorth, k0, trans.tm, lla.lat, lla.lon)
        x = x + 2e6
        y = y + 2e6

        return UTM(x, y, lla.alt)
    else
        lon_ref = Float64(utm_meridian(trans.zone))
        k0 = 0.9996 # Horizontal scaling factor
        #(x,y,k,γ) = transverse_mercator(lat_ref, lon_ref, lla.lat, lla.lon, ellipsoid(trans.datum))
        (x,y,γ,k) = transverse_mercator_forward(lon_ref, lla.lat, lla.lon, k0, trans.tm)
        x = 5e5 + x # Convention has 500km offset for easting
        y = (trans.isnorth ? 0.0 : 1e7) + y # Northing offset for southern hemisphere
        # also, k = k * k0
        return UTM(x, y, lla.alt) # Note: scaling not applied to vertical dimension
    end
end

Base.inv(trans::LLAfromUTM) = UTMfromLLA(trans.zone, trans.isnorth, trans.tm, trans.datum)
Base.inv(trans::UTMfromLLA) = LLAfromUTM(trans.zone, trans.isnorth, trans.tm, trans.datum)

##################
## ECEF <-> UTM ##
##################

"""
    UTMfromECEF(zone, isnorth, datum)

Creates composite transformation `UTMfromLLA(zone, isnorth, datum) ∘ LLAfromECEF(datum)`.
"""
UTMfromECEF(zone, isnorth, datum) = UTMfromLLA(zone, isnorth, datum) ∘ LLAfromECEF(datum)

"""
    ECEFfromUTM(zone, isnorth, datum)

Creates composite transformation `ECEFfromLLA(datum) ∘ LLAfromUTM(zone, isnorth, datum)`.
"""
ECEFfromUTM(zone, isnorth, datum) = ECEFfromLLA(datum) ∘ LLAfromUTM(zone, isnorth, datum)

##################
## LLA <-> UTMZ ##
##################

"""
    LLAfromUTMZ(datum)

Construct a `Transformation` object to convert from global `UTMZ`
coordinates to global `LLA` coordinates. Pre-caches
ellipsoidal parameters for efficiency and performs Charles Karney's accurate
6th-order series expansion algorithm.

(See also `LLAfromUTM`)
"""
struct LLAfromUTMZ{Datum,Order} <: Transformation
    tm::TransverseMercator{Order}
    datum::Datum
end
LLAfromUTMZ(datum) = LLAfromUTMZ(TransverseMercator(datum), datum)
Base.show(io::IO, trans::LLAfromUTMZ) = print(io, "LLAfromUTMZ($(trans.datum))")

function (trans::LLAfromUTMZ)(utm::UTMZ)
    if utm.zone == 0
        # Do inverse steriographic projection
        k0 = 0.994
        x = (utm.x - 2e6)
        y = (utm.y - 2e6)
        (lat,lon,γ,k) = polarst_inv(utm.isnorth, k0, trans.tm, x, y)

        return LLA(lat, lon, utm.z)
    else
        lon_ref = Float64(utm_meridian(utm.zone))
        k0 = 0.9996 # Horizontal scaling factor
        x = (utm.x - 5e5) # Convention has 500km offset for easting
        y = (utm.y - (utm.isnorth ? 0.0 : 1e7)) # Northing offset for southern hemisphere
        #(lat,lon,k,γ) = transverse_mercator_inv(lat_ref, lon_ref, x, y, ellipsoid(trans.datum))
        (lat,lon,γ,k) = transverse_mercator_reverse(lon_ref, x, y, k0, trans.tm)
        return LLA(lat, lon, utm.z) # Note: scaling not applied to vertical dimension
    end
end

"""
    UTMZfromLLA(datum)

Construct a `Transformation` object to convert from global `LLA` coordinates
to global `UTMZ` coordinates. The zone and hemisphere is automatically calculated
following the standard definitions (including exceptions in Norway). Pre-caches
ellipsoidal parameters for efficiency and performs Charles Karney's accurate
6th-order series expansion algorithm.

(See also `UTMfromLLA`)
"""
struct UTMZfromLLA{Datum,Order} <: Transformation
    tm::TransverseMercator{Order}
    datum::Datum
end
UTMZfromLLA(datum) = UTMZfromLLA(TransverseMercator(datum), datum)
Base.show(io::IO, trans::UTMZfromLLA) = print(io, "UTMZfromLLA($(trans.datum))")

function (trans::UTMZfromLLA)(lla::LLA)
    (zone, isnorth) = utm_zone(lla)
    if zone == 0
        # Do polar steriographic projection
        k0 = 0.994
        (x,y,γ,k) = polarst_fwd(isnorth, k0, trans.tm, lla.lat, lla.lon)
        x = x + 2e6
        y = y + 2e6

        return UTMZ(x, y, lla.alt, 0, isnorth)
    else
        lon_ref = Float64(utm_meridian(zone))
        k0 = 0.9996 # Horizontal scaling factor
        #(x,y,k,γ) = transverse_mercator(lat_ref, lon_ref, lla.lat, lla.lon, ellipsoid(trans.datum))
        (x,y,γ,k) = transverse_mercator_forward(lon_ref, lla.lat, lla.lon, k0, trans.tm)
        x = 5e5 + x # Convention has 500km offset for easting
        y = (isnorth ? 0.0 : 1e7) + y # Northing offset for southern hemisphere
        # also, k = k * k0
        return UTMZ(x, y, lla.alt, zone, isnorth) # Note: scaling not applied to vertical dimension
    end
end

Base.inv(trans::LLAfromUTMZ) = UTMZfromLLA(trans.tm, trans.datum)
Base.inv(trans::UTMZfromLLA) = LLAfromUTMZ(trans.tm, trans.datum)

###################
## UTM <-> UTMZ ##
###################

"""
    UTMZfromUTM(zone, isnorth, datum)

Transformation to append the UTM/UPS zone and hemisphere to a raw `UTM` data point.
"""
struct UTMZfromUTM{Datum} <: Transformation
    zone::Int
    isnorth::Bool
    datum::Datum
end
UTMZfromUTM(zone::Integer, h, d::D) where {D} = UTMZfromUTM{D}(UInt8(zone), h, d)
Base.show(io::IO, trans::UTMZfromUTM) = print(io, "UTMZfromUTM(zone=$(trans.zone == 0 ? "polar" : trans.zone) ($(trans.isnorth ? "north" : "south")), $(trans.datum))")
#Base.show(io::IO, trans::UTMZfromUTM) = print(io, "UTMZfromUTM(zone=$(trans.zone == 0 ? "polar" : trans.zone) ($(trans.isnorth ? "north" : "south")))")

"""
    UTMfromUTMZ(zone, isnorth, datum)

Transformation to remove the zone and hemisphere from `UTMZ` data point, and
automatically convert the data to the specified zone if necessary.
"""
struct UTMfromUTMZ{Datum} <: Transformation
    zone::Int
    isnorth::Bool
    datum::Datum
end
UTMfromUTMZ(zone::Integer, h, d::D) where {D} = UTMfromUTMZ{D}(UInt8(zone), h, d)
Base.show(io::IO, trans::UTMfromUTMZ) = print(io, "UTMfromUTMZ(zone=$(trans.zone == 0 ? "polar" : trans.zone) ($(trans.isnorth ? "north" : "south")), $(trans.datum))")
#Base.show(io::IO, trans::UTMfromUTMZ) = print(io, "UTMfromUTMZ(zone=$(trans.zone == 0 ? "polar" : trans.zone) ($(trans.isnorth ? "north" : "south")))")

(trans::UTMZfromUTM)(utm::UTM) = UTMZ(utm.x, utm.y, utm.z, trans.zone, trans.isnorth)
function (trans::UTMfromUTMZ)(utm::UTMZ)
    if trans.zone == utm.zone && trans.isnorth == utm.isnorth
        UTM(utm.x, utm.y, utm.z)
    else
        # Should this be an error or an automatic transformation to the correct zone?
        #error("Incorrect UTM zone")
        UTMfromLLA(trans.zone, trans.isnorth, trans.datum)(LLAfromUTMZ(trans.datum)(utm))
    end
end

Base.inv(trans::UTMfromUTMZ) = UTMZfromUTM(trans.zone, trans.isnorth, trans.datum)
Base.inv(trans::UTMZfromUTM) = UTMfromUTMZ(trans.zone, trans.isnorth, trans.datum)

###################
## ECEF <-> UTMZ ##
###################

"""
    UTMZfromECEF(datum)

Creates composite transformation `UTMZfromLLA(datum) ∘ LLAfromECEF(datum)`.
"""
UTMZfromECEF(datum) = UTMZfromLLA(datum) ∘ LLAfromECEF(datum)

"""
    ECEFfromUTMZ(datum)

Creates composite transformation `ECEFfromLLA(datum) ∘ LLAfromUTMZ(datum)`.
"""
ECEFfromUTMZ(datum) = ECEFfromLLA(datum) ∘ LLAfromUTMZ(datum)

##################
## ENU <-> UTMZ ##
##################

ENUfromECEF(origin::UTMZ, datum) = ENUfromECEF(LLAfromUTMZ(datum)(origin), datum)
ECEFfromENU(origin::UTMZ, datum) = ECEFfromENU(LLAfromUTMZ(datum)(origin), datum)

"""
    ENUfromUTMZ(origin, datum)

Creates composite transformation `ENUfromLLA(origin, datum) ∘ LLAfromUTMZ(datum)`.
"""
ENUfromUTMZ(origin, datum) = ENUfromLLA(origin, datum) ∘ LLAfromUTMZ(datum)

"""
    UTMZfromENU(origin, datum)

Creates composite transformation `UTMZfromLLA(datum) ∘ LLAfromENU(origin, datum)`.
"""
UTMZfromENU(origin, datum) = UTMZfromLLA(datum) ∘ LLAfromENU(origin, datum)

#################
## ENU <-> UTM ##
#################
ENUfromECEF(origin::UTM, zone::Integer, isnorth::Bool, datum) = ENUfromECEF(LLAfromUTM(zone, isnorth, datum)(origin), datum)
ECEFfromENU(origin::UTM, zone::Integer, isnorth::Bool, datum) = ECEFfromENU(LLAfromUTM(zone, isnorth, datum)(origin), datum)

# Assume origin and utm point share the same zone and hemisphere
UTMfromENU(origin::UTM, zone::Integer, isnorth::Bool, datum) = UTMfromLLA(zone, isnorth, datum) ∘ LLAfromENU(UTMZ(origin, zone, isnorth), datum)
ENUfromUTM(origin::UTM, zone::Integer, isnorth::Bool, datum) = ENUfromLLA(UTMZ(origin, zone, isnorth), datum) ∘ LLAfromUTM(zone, isnorth, datum)

"""
    UTMfromENU(origin, zone, isnorth, datum)

Creates composite transformation `UTMfromLLA(zone, isnorth, datum) ∘ LLAfromENU(origin, datum)`.
If `origin` is a `UTM` point, then it is assumed it is in the given specified zone and hemisphere.
"""
UTMfromENU(origin, zone::Integer, isnorth::Bool, datum) = UTMfromLLA(zone, isnorth, datum) ∘ LLAfromENU(origin, datum)

"""
    ENUfromUTM(origin, zone, isnorth, datum)

Creates composite transformation `UTMfromLLA(zone, isnorth, datum) ∘ LLAfromENU(origin, datum)`.
If `origin` is a `UTM` point, then it is assumed it is in the given specified zone and hemisphere.
"""
ENUfromUTM(origin, zone::Integer, isnorth::Bool, datum) = ENUfromLLA(origin, datum) ∘ LLAfromUTM(zone, isnorth, datum)

##################
## NED <-> UTMZ ##
##################

NEDfromECEF(origin::UTMZ, datum) = NEDfromECEF(LLAfromUTMZ(datum)(origin), datum)
ECEFfromNED(origin::UTMZ, datum) = ECEFfromNED(LLAfromUTMZ(datum)(origin), datum)

"""
    NEDfromUTMZ(origin, datum)

Creates composite transformation `ENUfromLLA(origin, datum) ∘ LLAfromUTMZ(datum)`.
"""
NEDfromUTMZ(origin, datum) = NEDfromLLA(origin, datum) ∘ LLAfromUTMZ(datum)

"""
    UTMZfromNED(origin, datum)

Creates composite transformation `UTMZfromLLA(datum) ∘ LLAfromNED(origin, datum)`.
"""
UTMZfromNED(origin, datum) = UTMZfromLLA(datum) ∘ LLAfromNED(origin, datum)

#################
## NED <-> UTM ##
#################
NEDfromECEF(origin::UTM, zone::Integer, isnorth::Bool, datum) = NEDfromECEF(LLAfromUTM(zone, isnorth, datum)(origin), datum)
ECEFfromNED(origin::UTM, zone::Integer, isnorth::Bool, datum) = ECEFfromNED(LLAfromUTM(zone, isnorth, datum)(origin), datum)

# Assume origin and utm point share the same zone and hemisphere
UTMfromNED(origin::UTM, zone::Integer, isnorth::Bool, datum) = UTMfromLLA(zone, isnorth, datum) ∘ LLAfromNED(UTMZ(origin, zone, isnorth), datum)
NEDfromUTM(origin::UTM, zone::Integer, isnorth::Bool, datum) = NEDfromLLA(UTMZ(origin, zone, isnorth), datum) ∘ LLAfromUTM(zone, isnorth, datum)

"""
    UTMfromNED(origin, zone, isnorth, datum)

Creates composite transformation `UTMfromLLA(zone, isnorth, datum) ∘ LLAfromNED(origin, datum)`.
If `origin` is a `UTM` point, then it is assumed it is in the given specified zone and hemisphere.
"""
UTMfromNED(origin, zone::Integer, isnorth::Bool, datum) = UTMfromLLA(zone, isnorth, datum) ∘ LLAfromNED(origin, datum)

"""
    NEDfromUTM(origin, zone, isnorth, datum)

Creates composite transformation `UTMfromLLA(zone, isnorth, datum) ∘ LLAfromNED(origin, datum)`.
If `origin` is a `UTM` point, then it is assumed it is in the given specified zone and hemisphere.
"""
NEDfromUTM(origin, zone::Integer, isnorth::Bool, datum) = NEDfromLLA(origin, datum) ∘ LLAfromUTM(zone, isnorth, datum)


using Formatting
numbers = map(x -> parse(Int64,x),split(first(readlines(open("input"))), ","))
fuel_cost(pivot, numbers, distance) = sum(map(x -> distance(pivot, x), numbers))
minimum_fuel_cost(numbers, metric) = format(minimum(map(x -> fuel_cost(x, numbers, metric), 0:maximum(numbers))))
minimum_fuel_cost(numbers, (x, y) -> abs(x-y))
minimum_fuel_cost(numbers, (x, y) -> (abs(x-y)*(abs(x-y)+1)) / 2) # small gauss


include("../diagonal_sbp.jl")
include("../deriv_ops_new.jl")
include("../split_matrix_free_dev.jl")
include("../split_matrix_free.jl")

if length(ARGS) != 0
    level = parse(Int,ARGS[1])
    Iterations = parse(Int,ARGS[2])
else
    level = 10
    Iterations = 10000
end

using LinearAlgebra
using SparseArrays
using Plots
using CUDA
using IterativeSolvers
using BenchmarkTools
using MAT

function e(i,n)
    # A = Matrix{Float64}(I,n,n)
    # return A[:,i]
    out = spzeros(n)
    out[i] = 1.0
    return out
end

function eyes(n)
    # return Matrix{Float64}(I,n,n)
    out = spzeros(n,n)
    for i in 1:n
        out[i,i] = 1.0
    end
    return out
end

function u(x,y)
    return sin.(π*x .+ π*y)
end

function Diag(A)
    # Self defined function that is similar to Matlab Diag
    return Diagonal(A[:])
end

function Operators_2d(i, j, p=2, h_list_x = ([1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]),
			 h_list_y = ([1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13])
			 )
    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    (D1x, HIx, H1x, r1x) = diagonal_sbp_D1(p,N_x,xc=(0,1));
    (D2x, S0x, SNx, HI2x, H2x, r2x) = diagonal_sbp_D2(p,N_x,xc=(0,1));

    (D1y, HIy, H1y, r1y) = diagonal_sbp_D1(p,N_y,xc=(0,1));
    (D2y, S0y, SNy, HI2y, H2y, r2y) = diagonal_sbp_D2(p,N_y,xc=(0,1));

    # BSx = sparse(SNx - S0x);
    # BSy = sparse(SNy - S0y);
    BSx = SNx - S0x
    BSy = SNy - S0y

    # Forming 2d Operators
    # e_1x = sparse(e(1,N_x+1));
    # e_Nx = sparse(e(N_x+1,N_x+1));
    # e_1y = sparse(e(1,N_y+1));
    # e_Ny = sparse(e(N_y+1,N_y+1));
    e_1x = e(1,N_x+1);
    e_Nx = e(N_x+1,N_x+1);
    e_1y = e(1,N_x+1);
    e_Ny = e(N_y+1,N_y+1);

    # I_Nx = sparse(eyes(N_x+1));
    # I_Ny = sparse(eyes(N_y+1));
    I_Nx = eyes(N_x+1);
    I_Ny = eyes(N_y+1);

    e_E = kron(e_Nx,I_Ny);
    e_W = kron(e_1x,I_Ny);
    e_S = kron(I_Nx,e_1y);
    e_N = kron(I_Nx,e_Ny);

    E_E = kron(sparse(Diag(e_Nx)),I_Ny);   # E_E = e_E * e_E'
    E_W = kron(sparse(Diag(e_1x)),I_Ny);
    E_S = kron(I_Nx,sparse(Diag(e_1y)));
    E_N = sparse(kron(I_Nx,sparse(Diag(e_Ny))));

    D1_x = kron(D1x,I_Ny);
    D1_y = kron(I_Nx,D1y);

    D2_x = kron(D2x,I_Ny);
    D2_y = kron(I_Nx,D2y);
    D2 = D2_x + D2_y

    HI_x = kron(HIx,I_Ny);
    HI_y = kron(I_Nx,HIy);

    H_x = kron(H1x,I_Ny);
    H_y = kron(I_Nx,H1y);

    BS_x = kron(BSx,I_Ny);
    BS_y = kron(I_Nx,BSy);

    HI_tilde = kron(HIx,HIx);
    H_tilde = kron(H1x,H1y);

    return (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N)
end

h_list_x = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]
h_list_y = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]

let
    k = level
    i = j = k

    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    Nx = N_x + 1;
    Ny = N_y + 1;

    println("Nx = $Nx, Ny=$Ny")
    (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N) = Operators_2d(i,j);
    analy_sol = u(x,y');

    tau_E = -13/hx;
    tau_W = -13/hx;
    tau_N = -1;
    tau_S = -1;

    beta = 1;

     # Forming SAT terms

    ## Formulation 1
    SAT_W = tau_W*HI_x*E_W + beta*HI_x*BS_x'*E_W;
    SAT_E = tau_E*HI_x*E_E + beta*HI_x*BS_x'*E_E;

    # SAT_S = tau_S*HI_y*E_S*D1_y
    # SAT_N = tau_N*HI_y*E_N*D1_y

    SAT_S = tau_S*HI_y*E_S*BS_y;
    SAT_N = tau_N*HI_y*E_N*BS_y;

    SAT_W_r = tau_W*HI_x*E_W*e_W + beta*HI_x*BS_x'*E_W*e_W;
    SAT_E_r = tau_E*HI_x*E_E*e_E + beta*HI_x*BS_x'*E_E*e_E;
    SAT_S_r = tau_S*HI_y*E_S*e_S;
    SAT_N_r = tau_N*HI_y*E_N*e_N;

    (alpha1,alpha2,alpha3,alpha4,beta) = (tau_N,tau_S,tau_W,tau_E,beta);

    g_W = sin.(π*y);
    g_E = -sin.(π*y);
    g_S = -π*cos.(π*x);
    # g_N = -π*cos.(π*x)
    g_N = π*cos.(π*x .+ π);

    # Solving with CPU
    A = D2 + SAT_W + SAT_E + SAT_S + SAT_N;

    b = -2π^2*u(x,y')[:] + SAT_W_r*g_W + SAT_E_r*g_E + SAT_S_r*g_S + SAT_N_r*g_N;

    A = - H_tilde*A;
    b = - H_tilde*b;

    CUDA.reclaim()
    A_d = CUDA.CUSPARSE.CuSparseMatrixCSC(A);
    b_d = CuArray(b);

    for _ = 1:Iterations
        A_d * b_d
    end
    CUDA.memory_status()
    CUDA.unsafe_free!(A_d)
    CUDA.unsafe_free!(b_d)
    CUDA.reclaim()
    return 0
end


using FinancialMonteCarlo, Test;
@show "Test Parameters"
S0 = 100.0;
K = 100.0;
r = 0.02;
Tneg = -1.0;
d = 0.01;
T = -Tneg;
Nsim = 10000;
Nstep = 30;
mc1 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.AntitheticMC());
rfCurve2 = ZeroRate([0.00, 0.02], T);
sigma = 0.2;

McConfig = MonteCarloConfiguration(Nsim, Nstep);
rfCurve = ZeroRate(r);

@show "Test Brownian Motion Parameters"
drift = 0.0
@test_throws(ErrorException, BrownianMotion(-sigma, drift))
@test_throws(ErrorException, BrownianMotion(-sigma, rfCurve2.r))


using Gnuplot
function ngonShape(h, k, r, n)
    x = [h .+ r*sin.(m*2π/n) for m in 1:n]
    y = [k .+ r*cos.(m*2π/n) for m in 1:n]
    x, y
end
function strpolygon(x, y)
    pt0 = "$(x[1]),$(y[1]) to"
    for p in 2:length(x)
       pt0=pt0*" $(x[p]),$(y[p]) to"
    end
    pt0[1:end-3]
end
#rectangle with filledcu
x1 = [-0.5,0.5,0.5,-0.5]
y1 = [0.55,0.55,3.45,3.45]
@gp x1 y1 "w filledcu notit lc 'grey'" yrange=(0.5,3.5) xrange=(-0.7,6.5)
@gp :- "set style fill transparent solid 0.5" "set size ratio -1"
#polygon from gnuplot
c = 1
for i in 1:6, j in 1:3
    global c
    x, y = ngonShape(i, j, 0.45, 2 + c)
    @gp :- "set obj $(c) polygon from $(strpolygon(x, y))" :-
    @gp :- "set obj $(c) fc lt $(c) fill transparent solid 0.3" :-
    @gp :- "set label font ',12' '$(2+c)' at $(i-0.05), $(j)" :-
    c += 1
end
@gp
save(term ="pngcairo size 600,400", output ="polygons.png")


module PointsReconstruction

using LinearAlgebra: norm, diagind, I, dot
using Distances: Euclidean, pairwise, PeriodicEuclidean, colwise
# using Plots: plot, scatter, plot!, scatter!
using Statistics: mean, var
using StaticArrays: SVector
using StatsBase: sample, sqL2dist
using FFTW: ifft, ifftshift, fftshift, fftfreq, fft
# using Interpolations: interpolate, Gridded, Linear, extrapolate
using Distributions: IsoNormal, MvNormal, pdf
# using ThreadTools: tmap1

include("generate_processes.jl")
include("summary_characteristics.jl")
# include("plot_functions.jl")
# include("analyze.jl")
include("wavelets.jl")
include("wavelet_sc.jl")

export uniform, true_uniform, net
export circle_cluster, normal_cluster, cross_cluster
export clusters, voronoi

export distance, distances, K, D, T, H

# export plotfdistr, plot_point_process

# export analyze, lossK, lossD, lossH, default_loss_params

export WaveletParams
export lossW, lossW_params

end


module GfxBase

export Vec4, Vec3, Vec2, Ray, distance

using Images
using StaticArrays
import LinearAlgebra.cross
import LinearAlgebra.normalize

export RGB

# some type aliases:
const Vec4 = SVector{4, Float32} # 4-vector of floats
const Vec3 = SVector{3, Float32} # 3-vector of floats
const Vec2 = SVector{2, Float32} # 2-vector of floats

struct Ray
    origin::Vec3
    direction::Vec3
end

function distance(a::Vec3, b::Vec3)
    sqrt((b[1] - a[1])^2 + (b[2] - a[2])^2 + (b[3] - a[3])^2)
end

end # module GfxBase


abstract type AbstractModel end

struct PoissonLapseUniform <: AbstractModel
    params::Vector{Float64}
    max::Int
end
function init(::Type{<:PoissonLapseUniform}, cm::Dict)
    m = maximum(key.Omissions_plus_one for key in keys(cm))
    return PoissonLapseUniform([10, 2, 0.1], m)
end

Distributions.params(p::PoissonLapseUniform) = p.params
Base.maximum(m::PoissonLapseUniform) = m.max

function Distributions.cdf(p::PoissonLapseUniform, x::Int)
    T, r, ϵ = params(p)
    d = Gamma(T, 1/r)
    # ϵ = 0
    # 1-ϵ e faccio Gamma, ϵ e faccio uniform, sommare e' uguale a "oppure" in probabilita'
    return (1-ϵ)*cdf(d, x) + ϵ*x/p.max
end
##
struct PoissonLapseExponential <: AbstractModel
    params::Vector{Float64}
    max::Int
end
function init(::Type{<:PoissonLapseExponential}, cm::Dict)
    m = maximum(key.Omissions_plus_one for key in keys(cm))
    return PoissonLapseExponential([10, 2, 3], m)
end

Distributions.params(p::PoissonLapseExponential) = p.params
Base.maximum(m::PoissonLapseExponential) = m.max

function Distributions.cdf(p::PoissonLapseExponential, x::Int)
    T, r, ϵ = params(p)
    g = Gamma(T, 1/r)
    l = Exponential(ϵ)
    #ϵ = 0
    # 1-ϵ e faccio Gamma, ϵ e faccio uniform, sommare e' uguale a "oppure" in probabilita'
    #return (1-cdf(l,x))*cdf(d, x) + cdf(l,x)
    return (1-0.01)*cdf(g,x) + 0.01*cdf(l,x)
end


path = @get_scratch!("planarpush")

nq = planarpush.nq
nu = planarpush.nu
nc = planarpush.nc
nz = num_var(planarpush)
nθ = 2 * planarpush.nq + planarpush.nu + planarpush.nw + 1

# Declare variables
@variables z[1:nz]
@variables θ[1:nθ]
@variables κ[1:1]

# Residual
r_pp = residual(planarpush, z, θ, κ)
rz_pp = Symbolics.jacobian(r_pp, z)
rθ_pp = Symbolics.jacobian(r_pp, θ)

# Build function
r_pp_func = build_function(r_pp, z, θ, κ)[2]
rz_pp_func = build_function(rz_pp, z, θ)[2]
rθ_pp_func = build_function(rθ_pp, z, θ)[2]
rz_pp_array = similar(rz_pp, Float64)
rθ_pp_array = similar(rθ_pp, Float64)

@save joinpath(path, "residual.jld2") r_pp_func rz_pp_func rθ_pp_func rz_pp_array rθ_pp_array
@load joinpath(path, "residual.jld2") r_pp_func rz_pp_func rθ_pp_func rz_pp_array rθ_pp_array


import Base: +, -, *, /
import LinearAlgebra: ×

"""
    struct Vec{N, T<:Real} <: AbstractVector{T}
    const Vec2{T} = Vec{2, T}
    const Vec3{T} = Vec{3, T}
    const Vec4{T} = Vec{4, T}
笛卡尔坐标系向量。
"""
struct Vec{N,T<:Real} <: AbstractVector{T}
    data::NTuple{N,T}
    Vec{N,T}(data::NTuple{N,T}) where {N,T<:Real} = new(data)
end

Vec(v::Vec) = v
Vec(data::NTuple{N,T}) where {N,T<:Real} = Vec{N,T}(data)
Vec(data...) = Vec(promote(data...))

+(va::Vec{N,Ta}, vb::Vec{N,Tb}) where {N,Ta,Tb} = Vec(ntuple(i -> va.data[i] + vb.data[i], N))
-(va::Vec{N,Ta}, vb::Vec{N,Tb}) where {N,Ta,Tb} = Vec(ntuple(i -> va.data[i] - vb.data[i], N))
*(k::Tk, va::Vec{N,Ta}) where {N,Ta,Tk<:Real} = Vec(ntuple(i -> k * va.data[i], N))
/(va::Vec{N,Ta}, k::Tk) where {N,Ta,Tk<:Real} = inv(k) * va
*(va::Vec{N,Ta}, k::Tk) where {N,Ta,Tk<:Real} = k * va

Base.iterate(v::Vec{N,T}, state = 1) where {N,T} = state > N ? nothing : (v.data[state], state + 1)
Base.size(v::Vec{N,T}) where {N,T} = (N,)
Base.size(v::Vec{N,T}, dim) where {N,T} = dim == 1 ? N : dim > 1 ? 1 : error("Vec size: dimension out of range")
Base.length(v::Vec{N,T}) where {N,T} = N
Base.eltype(::Type{Vec{N,T}}) where {N,T} = T
Base.getindex(v::Vec{N,T}, i::Integer) where {N,T} = v.data[i]
Base.firstindex(v::Vec{N,T}) where {N,T} = 1
Base.lastindex(v::Vec{N,T}) where {N,T} = N

const Vec2{T} = Vec{2,T}
const Vec3{T} = Vec{3,T}
const Vec4{T} = Vec{4,T}
Vec2(data::NTuple{2,T}) where {T<:Real} = Vec{2,T}(data)
Vec3(data::NTuple{3,T}) where {T<:Real} = Vec{3,T}(data)
Vec4(data::NTuple{4,T}) where {T<:Real} = Vec{4,T}(data)
Vec2(x, y) = Vec2(promote(x, y))
Vec3(x, y, z) = Vec3(promote(x, y, z))
Vec4(x, y, z, w) = Vec4(promote(x, y, z, w))

Vec3(v::Vec2, z = 0) = Vec3(promote(v.data..., z))
Vec4(v::Vec2, z = 0, w = 0) = Vec4(promote(v.data..., z, w))
Vec4(v::Vec3, w = 0) = Vec3(promote(v.data..., w))

×(va::Vec3{Ta}, vb::Vec3{Tb}) where {Ta,Tb} = begin
    Vec3(
        va.data[2] * vb.data[3] - va.data[3] * vb.data[2],
        va.data[3] * vb.data[1] - va.data[1] * vb.data[3],
        va.data[1] * vb.data[2] - va.data[2] * vb.data[1]
    )
end

Base.getproperty(v::Vec{N,T}, sym::Symbol) where {N,T<:Real} = begin
    if N <= 0 || N > 4
        return getfield(v, sym)
    end
    if N >= 1 && sym == :x
        return v.data[1]
    elseif N >= 2 && sym == :y
        return v.data[2]
    elseif N >= 3 && sym == :z
        return v.data[3]
    elseif N >= 4 && sym == :w
        return v.data[4]
    end
    getfield(v, sym)
end


"""
    fr = metagauss(f)

Generate a metagaussian field multiplying the spectralfield `f` with random phases and performing an inverse FFT transform to real space.
"""
function metagauss(f)

    (ns, ns, nt) = size(f)
    #	phases as fft of a gaussian noise random field
    #   ph=zeros(ns,ns,nt,'single')

    ph = zeros(ns, ns, nt)
    for i = 1:nt
        ph[:,:,i] = randn(ns,ns)
    end
    ph = fft(ph)
    ph = ph./abs.(ph); ph[1,1,1] = 0
    ph = f.*ph
    ph = real(ifft(ph))
end


module Scratchy

module Utils

export clip, decision, accuracy

function clip(x::Real, low=1e-15, hi=1 - 1e-15)
    if x < low
        return low
    elseif x > hi
        return hi
    else
        return x
    end
end

function decision(logits::Array)
	# 0.5 decision threshold
	return Int.(logits .>= 0.5)
end

function accuracy(y::Array, y_hat::Array)
	# Compute accuracy
    return sum(y .== y_hat) / size(y, 1)
end

end  # Utils

module LayerOps

export Operation

struct Operation
    name::String       # The name of this layer operation
    forward::Function  # The function itself
    backward::Function # The derivative of this function
end

end  # LayerOps

module Activations

using ..LayerOps

export Sigmoid, ReLU, Tanh

# Sigmoid (Logistic)
function sigmoid(h::Array{Float64})
    return @. 1.0 / (1.0 + exp(-h))
end

function d_sigmoid(h::Array{Float64})
    forward_vals = sigmoid(h)
    return @. forward_vals * (1.0 - forward_vals)
end

Sigmoid = Operation("Sigmoid", sigmoid, d_sigmoid)

# ReLU
function relu(x)
    return replace(a -> a<0 ? 0 : a, x)
end

function d_relu(x)
    return replace(a -> a<0 ? 0 : 1, x)
end

ReLU = Operation("ReLU", relu, d_relu)

# Tanh
function tanh(x)
    ex = exp.(x)
    enx = exp.(-x)
    return @. (ex - enx) / (ex + enx)
end

function d_tanh(x)
    return @. 1 - (tanh(x)^2)
end

Tanh = Operation("Tanh", tanh, d_tanh)

end  # Activations

module Losses

using ..LayerOps
import ..Utils

export BinaryCrossEntropy

# Binary cross entropy
function binary_crossentropy(y::Array, o::Array{Float64})
    # y: the gold-standard labels 0 or 1
    # o: the real-valued output [0,1]
    n = size(o, 1)
    y_hat = Utils.clip.(o)
    ces = @. y * log(o) + (1 - y) * log(1 - o)
    return -(1 / n) * sum(ces)
end

function d_binary_crossentropy(y::Array, o::Array{Float64})
    # y: the gold-standard labels 0 or 1
    # o: the real-valued output [0,1]
    y_hat = Utils.clip.(o)
    return @. (o - y) / (o * (1 - o))
end

BinaryCrossEntropy = Operation("BinaryCrossEntropy", binary_crossentropy,
                               d_binary_crossentropy)

end  # Losses

module Layers

using ..LayerOps
import Random

export Dense, forward!

# Dense feed-forward layer
mutable struct Dense
    name::String            # What we'll call this layer.
    W::Array{Float64}       # The weight matrix
    b::Array{Float64}       # The bias term
    σ::Operation            # The activation function
    δ::Array{Float64}       # Used in backprop
    input::Array{Float64}   # The inputs from the previous layer
    z::Array{Float64}       # Weighted inputs XW + b
    output::Array{Float64}  # The output of this layer, i.e. the actvations

    # Constructor v1: When we have already initialized W and b
    function Dense(name::String, W::Array{Float64},
				   b::Array{Float64}, σ::Operation)
        if !isa(σ, Operation)
            error("Activation function '$σ' in layer '$name' is not an Operation")
        end
        if size(W, 2) != size(b, 2)
            error("Incompatible W and b shapes in layer '$name'")
        else
            # Since δ, input, and output are computed later,
			# we init empty arrays.
            new(name, W, b, σ, [], [], [], [])
        end
    end

    # Constructor v2: When we want to randomly initialize W and
	# b given the shape of W.
    function Dense(name::String, size::Tuple{Int,Int}, σ::Operation)
        W = Random.rand(Float64, size)
        b = Random.rand(Float64, (1, size[2]))
        # Since δ, input, and output are computed later, we init as empty arrays.
        new(name, W, b, σ, [], [], [], [])
    end
end

# Dense layer forward op
function forward!(l::Dense, X::Array)
    # It's always a good idea to double check that the shapes are compatible.
    if size(X, 2) != size(l.W, 1)
        error("Input shape $(size(X)) incompatible with layer '$(l.name)' shape $(size(l.W))")
    end
    l.input = X
    l.z = (X * l.W) .+ l.b  # Weighted inputs
    l.output = l.σ.forward(l.z)
    return l.output
end

end  # Layers

module Models

using ..LayerOps
import ..Layers

export Network, forward!, backward!, fit!, display

struct Network
    name::String
    layers::Array
    loss::Operation
    η::Real  # The learning rate

    # Constructor: Check that all layer shapes are compatible
    function Network(name::String, layers::Array, loss::Operation, η::Real)
        l1 = layers[1]
        for l2 in layers[2:end]
            l1_outdim = size(l1.W, 2)
            l2_indim = size(l2.W, 1)
            if l1_outdim != l2_indim
                error("""Output dimension $(size(l1.W, 2)) of layer
                      '$(l1.name)' incompatible with input dimension
                      $(size(l2.W, 1)) of layer '$(l2.name)'.""")
            end
        end
        new(name, layers, loss, η)
    end
end

# Network forward
function forward!(network::Network, X::Array)
    layer_input = X
    for l in network.layers
        layer_input = Layers.forward!(l, layer_input)
    end
    return layer_input
end

# Helper function for backward!
function compute_δ(network::Network, y::Array, i::Int)
    l = network.layers[i]
    dσ_dz = l.σ.backward(l.z)
    # The base case of the iteration: the last layer in the network.
    if i == length(network.layers)
        o = network.layers[i].output
        δ = network.loss.backward(y, o) .* dσ_dz
        return δ
    end
    l_next = network.layers[i+1]
    # l_next.δ should have already been computed because we're iterating
    # backwards through the network's layers in backward!() below.
    δ = (l_next.δ * transpose(l_next.W)) .* dσ_dz
    return δ
end

# Network backprop
function backward!(net::Network, y::Array)
    num_layers = length(net.layers)
    for i in num_layers:-1:1  # Iterate backwards through the network
        l = net.layers[i]
        l.δ = compute_δ(net, y, i)
        ∂C_∂W = transpose(l.δ) * l.input
        l.W = l.W .- (net.η * transpose(∂C_∂W))
    end
end

# Model estimation
function fit!(net::Network, X::Array, y::Array, epochs::Int)
    for i=1:epochs
        o = forward!(net, X)
        backward!(net, y)
        if i == 1 || i % 100 == 0
            loss = net.loss.forward(y, o)
            println("Loss ($i): $(loss)")
        end
    end
end

# Summarize a network's structure.
import Base.display
function display(nn::Network)
    println("Network: $(nn.name)")
    for layer in nn.layers
        println("$(layer.name) : $(size(layer.W)) : $(layer.σ.name)")
    end
    println("Loss: $(nn.loss.name)")
end

end  # Models

end  # Scratchy


using ImplicitGlobalGrid

@views d_xa(A) = A[2:end  , :     , :     ] .- A[1:end-1, :     , :     ];
@views d_xi(A) = A[2:end  ,2:end-1,2:end-1] .- A[1:end-1,2:end-1,2:end-1];
@views d_ya(A) = A[ :     ,2:end  , :     ] .- A[ :     ,1:end-1, :     ];
@views d_yi(A) = A[2:end-1,2:end  ,2:end-1] .- A[2:end-1,1:end-1,2:end-1];
@views d_za(A) = A[ :     , :     ,2:end  ] .- A[ :     , :     ,1:end-1];
@views d_zi(A) = A[2:end-1,2:end-1,2:end  ] .- A[2:end-1,2:end-1,1:end-1];
@views  inn(A) = A[2:end-1,2:end-1,2:end-1]

@views function diffusion3D()
    # Physics
    lam        = 1.0;                                       # Thermal conductivity
    cp_min     = 1.0;                                       # Minimal heat capacity
    lx, ly, lz = 10.0, 10.0, 10.0;                          # Length of computational domain in dimension x, y and z

    # Numerics
    nx, ny, nz = 128, 128, 128;                             # Number of gridpoints in dimensions x, y and z
    nt         = 10000;                                     # Number of time steps
    me, dims   = init_global_grid(nx, ny, nz);              # Initialize the implicit global grid
    dx         = lx/(nx_g()-1);                             # Space step in dimension x
    dy         = ly/(ny_g()-1);                             # ...        in dimension y
    dz         = lz/(nz_g()-1);                             # ...        in dimension z

    # Array initializations
    T     = zeros(nx,   ny,   nz  );
    Cp    = zeros(nx,   ny,   nz  );
    dTedt = zeros(nx-2, ny-2, nz-2);
    qx    = zeros(nx-1, ny-2, nz-2);
    qy    = zeros(nx-2, ny-1, nz-2);
    qz    = zeros(nx-2, ny-2, nz-1);

    # Initial conditions (heat capacity and temperature with two Gaussian anomalies each)
    Cp .= cp_min .+ [5*exp(-((x_g(ix,dx,Cp)-lx/1.5))^2-((y_g(iy,dy,Cp)-ly/2))^2-((z_g(iz,dz,Cp)-lz/1.5))^2) +
                     5*exp(-((x_g(ix,dx,Cp)-lx/3.0))^2-((y_g(iy,dy,Cp)-ly/2))^2-((z_g(iz,dz,Cp)-lz/1.5))^2) for ix=1:size(T,1), iy=1:size(T,2), iz=1:size(T,3)]
    T  .= [100*exp(-((x_g(ix,dx,T)-lx/2)/2)^2-((y_g(iy,dy,T)-ly/2)/2)^2-((z_g(iz,dz,T)-lz/3.0)/2)^2) +
            50*exp(-((x_g(ix,dx,T)-lx/2)/2)^2-((y_g(iy,dy,T)-ly/2)/2)^2-((z_g(iz,dz,T)-lz/1.5)/2)^2) for ix=1:size(T,1), iy=1:size(T,2), iz=1:size(T,3)]

    # Time loop
    dt = min(dx*dx,dy*dy,dz*dz)*cp_min/lam/8.1;                                               # Time step for the 3D Heat diffusion
    for it = 1:nt
        qx    .= -lam.*d_xi(T)./dx;                                                           # Fourier's law of heat conduction: q_x   = -λ δT/δx
        qy    .= -lam.*d_yi(T)./dy;                                                           # ...                               q_y   = -λ δT/δy
        qz    .= -lam.*d_zi(T)./dz;                                                           # ...                               q_z   = -λ δT/δz
        dTedt .= 1.0./inn(Cp).*(-d_xa(qx)./dx .- d_ya(qy)./dy .- d_za(qz)./dz);               # Conservation of energy:           δT/δt = 1/cₚ (-δq_x/δx - δq_y/dy - δq_z/dz)
        T[2:end-1,2:end-1,2:end-1] .= inn(T) .+ dt.*dTedt;                                    # Update of temperature             T_new = T_old + δT/δt
        update_halo!(T);                                                                      # Update the halo of T
    end

    finalize_global_grid();                                                                   # Finalize the implicit global grid
end

diffusion3D()


using Test
using LinearAlgebra
using CartesianMesh
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

function displacement(alpha, x)
    u1 = alpha * x[2] * sin(pi * x[1])
    u2 = alpha * (x[1]^3 + cos(pi * x[2]))
    return [u1, u2]
end

function stress_field(lambda, mu, alpha, x)
    s11 =
        (lambda + 2mu) * alpha * pi * x[2] * cos(pi * x[1]) -
        lambda * alpha * pi * sin(pi * x[2])
    s22 =
        -(lambda + 2mu) * alpha * pi * sin(pi * x[2]) +
        lambda * alpha * pi * x[2] * cos(pi * x[1])
    s12 = alpha * mu * (3x[1]^2 + sin(pi * x[1]))
    return [s11, s22, s12]
end

function body_force(lambda, mu, alpha, x)
    b1 = alpha * (lambda + 2mu) * pi^2 * x[2] * sin(pi * x[1])
    b2 =
        -alpha * (6mu * x[1] + (lambda + mu) * pi * cos(pi * x[1])) +
        alpha * (lambda + 2mu) * pi^2 * cos(pi * x[2])
    return [b1, b2]
end

function onleftboundary(x, L, W)
    return x[1] ≈ 0.0
end

function onbottomboundary(x, L, W)
    return x[2] ≈ 0.0
end

function onrightboundary(x, L, W)
    return x[1] ≈ L
end

function ontopboundary(x, L, W)
    return x[2] ≈ W
end

function onboundary(x, L, W)
    return onbottomboundary(x, L, W) ||
           onrightboundary(x, L, W) ||
           ontopboundary(x, L, W) ||
           onleftboundary(x, L, W)
end

function solve_plane_interface(x0, normal, nelmts, polyorder, numqp, penaltyfactor)
    L = 1.0
    W = 1.0
    lambda, mu = 1.0, 2.0

    dx = 1.0 / nelmts
    penalty = penaltyfactor / dx * (lambda + mu)
    alpha = 0.1
    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [L, W], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> plane_distance_function(x, normal, x0), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.coherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)

    bottomdisplacementbc = CutCell.DisplacementComponentCondition(
        x -> displacement(alpha, x)[2],
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onbottomboundary(x, L, W),
        [0.0, 1.0],
        penalty,
    )
    rightdisplacementbc = CutCell.DisplacementComponentCondition(
        x -> displacement(alpha, x)[1],
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onrightboundary(x, L, W),
        [1.0, 0.0],
        penalty,
    )
    topdisplacementbc = CutCell.DisplacementComponentCondition(
        x -> displacement(alpha, x)[1],
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> ontopboundary(x, L, W),
        [1.0, 0.0],
        penalty,
    )
    leftdisplacementbc = CutCell.DisplacementComponentCondition(
        x -> displacement(alpha, x)[1],
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onleftboundary(x, L, W),
        [1.0, 0.0],
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)
    CutCell.assemble_body_force_linear_form!(
        sysrhs,
        x -> body_force(lambda, mu, alpha, x),
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        bottomdisplacementbc,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        rightdisplacementbc,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(sysmatrix, sysrhs, topdisplacementbc, cutmesh)
    CutCell.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        leftdisplacementbc,
        cutmesh,
    )

    CutCell.assemble_traction_force_component_linear_form!(
        sysrhs,
        x -> stress_field(lambda, mu, alpha, x)[3],
        basis,
        facequads,
        cutmesh,
        x -> onrightboundary(x, L, W),
        [0.0, 1.0],
    )
    CutCell.assemble_traction_force_component_linear_form!(
        sysrhs,
        x -> -stress_field(lambda, mu, alpha, x)[3],
        basis,
        facequads,
        cutmesh,
        x -> onbottomboundary(x, L, W),
        [1.0, 0.0],
    )
    CutCell.assemble_traction_force_component_linear_form!(
        sysrhs,
        x -> stress_field(lambda, mu, alpha, x)[2],
        basis,
        facequads,
        cutmesh,
        x -> ontopboundary(x, L, W),
        [0.0, 1.0],
    )
    CutCell.assemble_traction_force_component_linear_form!(
        sysrhs,
        x -> -stress_field(lambda, mu, alpha, x)[3],
        basis,
        facequads,
        cutmesh,
        x -> onleftboundary(x, L, W),
        [0.0, 1.0],
    )

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    err = mesh_L2_error(disp, x -> displacement(alpha, x), basis, cellquads, cutmesh)

    return err
end

x0 = [0.5, 0.0]
normal = [1.0, 0.0]
penaltyfactor = 1e2
polyorder = 1
numqp = required_quadrature_order(polyorder) + 2
powers = [3,4,5]
nelmts = [2^p + 1 for p in powers]

err = [
    solve_plane_interface(x0, normal, ne, polyorder, numqp, penaltyfactor) for ne in nelmts
]

dx = 1.0 ./ nelmts
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]

u1rate = convergence_rate(dx, u1err)
u2rate = convergence_rate(dx, u2err)

@test allapprox(u1rate,repeat([2.0],length(u1rate)),0.1)
@test allapprox(u2rate,repeat([2.0],length(u2rate)),0.1)

x0 = [0.5, 0.0]
normal = [1.0, 0.0]
penaltyfactor = 1e2
polyorder = 2
numqp = required_quadrature_order(polyorder) + 2
powers = [3,4,5]
nelmts = [2^p + 1 for p in powers]

err = [
    solve_plane_interface(x0, normal, ne, polyorder, numqp, penaltyfactor) for ne in nelmts
]

dx = 1.0 ./ nelmts
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]

u1rate = convergence_rate(dx, u1err)
u2rate = convergence_rate(dx, u2err)

@test allapprox(u1rate,repeat([3.0],length(u1rate)),1e-1)
@test allapprox(u2rate,repeat([3.0],length(u1rate)),1e-1)


using ForwardDiff, OrdinaryDiffEq
using LinearAlgebra: I

include("utils.jl")

struct TestEuclidean{N} <: Manifold end
struct TestEuclideanMetric <: Metric end

@testset "scaled Euclidean metric" begin
    Manifolds.manifold_dimension(::TestEuclidean{N}) where {N} = N
    function Manifolds.local_metric(M::MetricManifold{<:TestEuclidean,<:TestEuclideanMetric}, x)
        return Diagonal(1.0:manifold_dimension(M))
    end

    n = 3
    E = TestEuclidean{n}()
    g = TestEuclideanMetric()
    M = MetricManifold(E, g)
    G = Diagonal(1.0:n)
    invG = inv(G)
    @test manifold_dimension(M) == n
    @test base_manifold(M) === E
    @test metric(M) === g

    @test_throws ErrorException local_metric_jacobian(E, zeros(3))
    @test_throws ErrorException christoffel_symbols_second_jacobian(E, zeros(3))

    for vtype in (Vector, SVector{n}, MVector{n})
        x, v, w = vtype(randn(n)), vtype(randn(n)), vtype(randn(n))

        @test local_metric(M, x) ≈ G
        @test inverse_local_metric(M, x) ≈ invG
        @test det_local_metric(M, x) ≈ *(1.0:n...)
        @test log_local_metric_density(M, x) ≈ sum(log.(1.0:n)) / 2
        @test inner(M, x, v, w) ≈ dot(v, G * w)
        @test norm(M, x, v) ≈ sqrt(dot(v, G * v))

        T = 0:.5:10
        @test exp(M, x, v, T) ≈ [x + t * v for t in T]
        @test geodesic(M, x, v, T) ≈ [x + t * v for t in T]

        @test christoffel_symbols_first(M, x) ≈ zeros(n, n, n)
        @test christoffel_symbols_second(M, x) ≈ zeros(n, n, n)
        @test riemann_tensor(M, x) ≈ zeros(n, n, n, n)
        @test ricci_tensor(M, x) ≈ zeros(n, n)
        @test ricci_curvature(M, x) ≈ 0
        @test gaussian_curvature(M, x) ≈ 0
        @test einstein_tensor(M, x) ≈ zeros(n, n)
    end
end

struct TestSphere{N,T} <: Manifold
    r::T
end

struct TestSphericalMetric <: Metric end

@testset "scaled Sphere metric" begin
    Manifolds.manifold_dimension(::TestSphere{N}) where {N} = N
    function Manifolds.local_metric(M::MetricManifold{<:TestSphere,<:TestSphericalMetric}, x)
        r = base_manifold(M).r
        d = similar(x)
        d[1] = r^2
        d[2] = d[1] * sin(x[1])^2
        return Diagonal(d)
    end
    sph_to_cart(θ, ϕ) = [cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)]

    n = 2
    r = 10 * rand()
    θ, ϕ = π * rand(), 2π * rand()
    Sr = TestSphere{n,Float64}(r)
    S = Manifolds.Sphere(n)
    g = TestSphericalMetric()
    M = MetricManifold(Sr, g)

    @test manifold_dimension(M) == n
    @test base_manifold(M) === Sr
    @test metric(M) === g

    for vtype in (Vector, SVector{n}, MVector{n})
        x = vtype([θ, ϕ])
        G = Diagonal(vtype([1, sin(θ)^2])) .* r^2
        invG = Diagonal(vtype([1, 1 / sin(θ)^2])) ./ r^2
        v, w = normalize(randn(n)), normalize(randn(n))

        @test local_metric(M, x) ≈ G
        @test inverse_local_metric(M, x) ≈ invG
        @test det_local_metric(M, x) ≈ r^4 * sin(θ)^2
        @test log_local_metric_density(M, x) ≈ 2log(r) + log(sin(θ))
        @test inner(M, x, v, w) ≈ dot(v, G * w)
        @test norm(M, x, v) ≈ sqrt(dot(v, G * v))

        xcart = sph_to_cart(θ, ϕ)
        vcart = [cos(ϕ)*cos(θ) -sin(ϕ)*sin(θ);
                 sin(ϕ)*cos(θ)  cos(ϕ)*sin(θ);
                 -sin(θ) 0] * v

        if !Sys.iswindows() || Sys.ARCH == :x86_64
            @testset "numerically integrated geodesics for $vtype" begin
                T = 0:.1:1
                @test isapprox([sph_to_cart(yi...) for yi in exp(M, x, v, T)],
                               exp(S, xcart, vcart, T); atol=1e-3, rtol=1e-3)
                @test isapprox([sph_to_cart(yi...) for yi in geodesic(M, x, v, T)],
                               geodesic(S, xcart, vcart, T); atol=1e-3, rtol=1e-3)
            end
        end

        Γ₁ = christoffel_symbols_first(M, x)
        for i=1:n, j=1:n, k=1:n
            if (i,j,k) == (1,2,2) || (i,j,k) == (2,1,2)
                @test Γ₁[i,j,k] ≈ r^2*cos(θ)*sin(θ)
            elseif (i,j,k) == (2,2,1)
                @test Γ₁[i,j,k] ≈ -r^2*cos(θ)*sin(θ)
            else
                @test Γ₁[i,j,k] ≈ 0
            end
        end

        Γ₂ = christoffel_symbols_second(M, x)
        for l=1:n, i=1:n, j=1:n
            if (l,i,j) == (1,2,2)
                @test Γ₂[l,i,j] ≈ -cos(θ)*sin(θ)
            elseif (l,i,j) == (2,1,2) || (l,i,j) == (2,2,1)
                @test Γ₂[l,i,j] ≈ cot(θ)
            else
                @test Γ₂[l,i,j] ≈ 0
            end
        end

        R = riemann_tensor(M, x)
        for l=1:n, i=1:n, j=1:n, k=1:n
            if (l,i,j,k) == (2,1,1,2)
                @test R[l,i,j,k] ≈ -1
            elseif (l,i,j,k) == (2,1,2,1)
                @test R[l,i,j,k] ≈ 1
            elseif (l,i,j,k) == (1,2,1,2)
                @test R[l,i,j,k] ≈ sin(θ)^2
            elseif (l,i,j,k) == (1,2,2,1)
                @test R[l,i,j,k] ≈ -sin(θ)^2
            else
                @test R[l,i,j,k] ≈ 0
            end
        end

        @test ricci_tensor(M, x) ≈ G ./ r^2
        @test ricci_curvature(M, x) ≈ 2 / r^2
        @test gaussian_curvature(M, x) ≈ 1 / r^2
        @test einstein_tensor(M, x) ≈ ricci_tensor(M, x) - gaussian_curvature(M, x)  .* G
    end
end

struct BaseManifold{N} <: Manifold end
struct BaseManifoldMetric{M} <: Metric end
struct DefaultBaseManifoldMetric <: Metric end

@testset "Metric decorator" begin
    Manifolds.manifold_dimension(::BaseManifold{N}) where {N} = N
    Manifolds.inner(::BaseManifold, x, v, w) = 2 * dot(v,w)
    Manifolds.exp!(::BaseManifold, y, x, v) = y .= x + 2 * v
    Manifolds.log!(::BaseManifold, v, x, y) = v .= (y - x) / 2
    Manifolds.project_tangent!(::BaseManifold, w, x, v) = w .= 2 .* v
    Manifolds.project_point!(::BaseManifold, y, x) = (y .= x)
    Manifolds.local_metric(::MetricManifold{BaseManifold{N},BaseManifoldMetric{N}},x) where N = 2*one(x*x')
    Manifolds.exp!(::MetricManifold{BaseManifold{N},BaseManifoldMetric{N}}, y, x, v) where N = exp!(base_manifold(M), y, x, v)
    Manifolds.vector_transport_to!(::BaseManifold, vto, x, v, y, ::ParallelTransport) = (vto .= v)
    Manifolds.is_default_metric(M::BaseManifold,G::DefaultBaseManifoldMetric) = Val(true)
    Manifolds.tangent_orthonormal_basis(M::BaseManifold{N},x,v) where {N} = ( [(Matrix(I, N, N)[:,i]) for i in 1:N], zeros(N))
    Manifolds.projected_distribution(M::BaseManifold, d) = ProjectedPointDistribution(M, d, project_point!, rand(d))
    Manifolds.projected_distribution(M::BaseManifold, d, x) = ProjectedPointDistribution(M, d, project_point!, x)

    function Manifolds.flat!(::BaseManifold, v::FVector{Manifolds.CotangentSpaceType}, x, w::FVector{Manifolds.TangentSpaceType})
        v.data .= 2 .* w.data
        return v
    end
    function Manifolds.sharp!(::BaseManifold, v::FVector{Manifolds.TangentSpaceType}, x, w::FVector{Manifolds.CotangentSpaceType})
        v.data .= w.data ./ 2
        return v
    end

    M = BaseManifold{3}()
    g = BaseManifoldMetric{3}()
    MM = MetricManifold(M, g)
    g2 = DefaultBaseManifoldMetric()
    MM2 = MetricManifold(M,g2)

    @test is_decorator_manifold(MM) == Val{true}()
    @test is_decorator_manifold(MM2) == Val{true}()
    @test is_decorator_manifold(M) == Val{false}()
    @test is_default_metric(MM) == is_default_metric(base_manifold(MM),metric(MM))
    @test is_default_metric(MM2) == is_default_metric(base_manifold(MM2),metric(MM2))
    @test is_default_metric(MM2) == Val{true}()

    @test convert(typeof(MM2),M) == MM2
    @test_throws ErrorException convert(typeof(MM),M)
    x = [0.1 0.2 0.4]
    v = [0.5 0.7 0.11]
    w = [0.13 0.17 0.19]
    y = similar(x)

    # Test fallbacks
    @test_throws ErrorException vee!(M,w,x,v)
    @test_throws ErrorException hat!(M,w,x,v)

    @test inner(M, x, v, w) == 2 * dot(v,w)
    @test inner(MM, x, v, w) === inner(M, x, v, w)
    @test norm(MM, x, v) === norm(M, x, v)
    @test exp(M, x, v) == x + 2 * v
    @test exp(MM2, x, v) == exp(M, x, v)
    @test exp!(MM, y, x, v) === exp!(M, y, x, v)
    @test retract!(MM, y, x, v) === retract!(M, y, x, v)
    @test retract!(MM, y, x, v, 1) === retract!(M, y, x, v, 1)
    # without a definition for the metric from the embedding, no projection possible
    @test_throws ErrorException log!(MM, w, x, y) === project_tangent!(M, w, x, y)
    @test_throws ErrorException project_tangent!(MM, w, x, v) === project_tangent!(M, w, x, v)
    @test_throws ErrorException project_point!(MM, y, x) === project_point!(M, y, x)
    @test_throws ErrorException vector_transport_to!(MM, w, x, v, y) === vector_transport_to!(M, w, x, v, y)
    # without DiffEq, these error
    # @test_throws ErrorException exp(MM,x, v, 1:3)
    # @test_throws ErrorException exp!(MM, y, x, v)
    # these always fall back anyways.
    @test zero_tangent_vector!(MM, v, x) === zero_tangent_vector!(M, v, x)
    @test injectivity_radius(MM, x) === injectivity_radius(M, x)
    @test injectivity_radius(MM) === injectivity_radius(M)
    @test is_manifold_point(MM, x) === is_manifold_point(M, x)
    @test is_tangent_vector(MM, x, v) === is_tangent_vector(M, x, v)

    @test_throws ErrorException local_metric(MM2,x)
    @test_throws ErrorException local_metric_jacobian(MM2,x)
    @test_throws ErrorException christoffel_symbols_second_jacobian(MM2,x)
    # MM falls back to nondefault error
    @test_throws ErrorException projected_distribution(MM,1,x)
    @test_throws ErrorException projected_distribution(MM,1)
    @test_throws ErrorException normal_tvector_distribution(MM,x,0.2)

    @test inner(MM2, x, v, w) === inner(M, x, v, w)
    @test norm(MM2, x, v) === norm(M, x, v)
    @test distance(MM2, x, y) === distance(M, x, y)
    @test exp!(MM2, y, x, v) === exp!(M, y, x, v)
    @test log!(MM2, v, x, y) === log!(M, v, x, y)
    @test retract!(MM2, y, x, v) === retract!(M, y, x, v)
    @test retract!(MM2, y, x, v, 1) === retract!(M, y, x, v, 1)

    @test project_point!(MM2, y, x) === project_point!(M, y, x)
    @test project_tangent!(MM2, w, x, v) === project_tangent!(M, w, x, v)
    @test vector_transport_to!(MM2, w, x, v, y) == vector_transport_to!(M, w, x, v, y)
    @test zero_tangent_vector!(MM2, v, x) === zero_tangent_vector!(M, v, x)
    @test injectivity_radius(MM2, x) === injectivity_radius(M, x)
    @test injectivity_radius(MM2) === injectivity_radius(M)
    @test is_manifold_point(MM2, x) === is_manifold_point(M, x)
    @test is_tangent_vector(MM2, x, v) === is_tangent_vector(M, x, v)

    a = Manifolds.projected_distribution(M, Distributions.MvNormal(zero(zeros(3)), 1.0))
    b = Manifolds.projected_distribution(MM2, Distributions.MvNormal(zero(zeros(3)), 1.0))
    @test isapprox(Matrix(a.d.Σ), Matrix(b.d.Σ))
    @test isapprox(a.d.μ, b.d.μ)
    @test tangent_orthonormal_basis(M,x,v) == tangent_orthonormal_basis(MM2,x,v)
    @test_throws ErrorException tangent_orthonormal_basis(MM,x,v)
    cov = flat(M, x, FVector(TangentSpace, v))
    cow = flat(M, x, FVector(TangentSpace, w))
    @test cov.data ≈ flat(MM, x, FVector(TangentSpace, v)).data
    cotspace = CotangentBundleFibers(M)
    cotspace2 = CotangentBundleFibers(MM)
    @test cov.data ≈ 2 * v
    @test inner(M, x, v, w) ≈ inner(cotspace, x, cov.data, cow.data)
    @test inner(MM, x, v, w) ≈ inner(cotspace, x, cov.data, cow.data)
    @test inner(MM, x, v, w) ≈ inner(cotspace2, x, cov.data, cow.data)
    @test sharp(M, x, cov).data ≈ v
end


@testset "Right Division" begin
    mat = diagm([1,2,3])
    v = [2,3,4]
    tv = v.'

    @test (tv/mat)::TransposedVector ≈ [2/1  3/2  4/3]

    @test (v.'/mat)::TransposedVector ≈ [2/1  3/2  4/3]
    @test (v.'/mat.')::TransposedVector ≈ [2/1  3/2  4/3]
    @test (tv/mat.')::TransposedVector ≈ [2/1  3/2  4/3]

    @test (v'/mat)::TransposedVector ≈ [2/1  3/2  4/3]
    @test (v'/mat')::TransposedVector ≈ [2/1  3/2  4/3]
    @test (tv/mat')::TransposedVector ≈ [2/1  3/2  4/3]
end


using Primes
using OffsetArrays

struct primePower
    prime::Int
    power::Int
end

function memoizeany(func)
    thedict = Dict()
    return (a)-> memoizeanyaux(a, func, thedict)
end

function memoizeanyaux(a, func, thedict)
    if haskey(thedict, a)
        return thedict[a]
    end
    res = func(a)
    thedict[a] =res
    return res
end

function maxpow(n::Int, p::Int)
    pow = 1
    n /= p
    while true
        rem = n%p
        if rem != 0
            return pow
        end
        pow+=1
        n /=p
    end
end

function factorall(n)
    dictlist = [Dict{Int, Int}() for i in 1:n]
    theprimes = primes(Int(ceil(sqrt(n))))
    for i in theprimes
        for j in i:i:n
            @inbounds dictlist[j][i]=maxpow(j, i)
        end
    end
    for i in 2:n
        dl = dictlist[i]
        if isempty(dl)
            dl[i] = 1
        end
        pfacs = prod(i^dl[i] for i in keys(dl))
        therat = i/pfacs
        if therat > 1
            dl[therat]=1
        end
    end
    return dictlist
end

function factorall(m,n)
    dictlist = [Dict{Int, Int}() for i in m:n]
    dictlist = OffsetArray(dictlist, m:n)
    theprimes = primes(Int(ceil(sqrt(n))))
    for i in theprimes
        begprime = Int(ceil(m/i))*i
        for j in begprime:i:n
            @inbounds dictlist[j][i]=maxpow(j, i)
        end
    end
    thetwo = max(2, m)
    for i in thetwo:n
        dl = dictlist[i]
        if isempty(dl)
            dl[i] = 1
        end
        pfacs = prod(i^dl[i] for i in keys(dl))
        therat = i/pfacs
        if therat > 1
            dl[therat]=1
        end
    end
    return dictlist
end

function domult(ndict, func)
    if isempty(ndict)
        return 1
    end
    return prod(func(i, ndict[i]) for i in keys(ndict))
end

newfact = memoizeany(factorall)

function doallmult(n, func)
    return [domult(i, func) for i in newfact(n)]
end

function doallmult(m, n, func)
    tmpar = collect(factorall(m, n)) #compiler bug
    res = [domult(i, func) for i in tmpar]
    return OffsetArray(res, m:n)
end

function makesigmab(k)
    thedict = Dict((0, 0)=>BigInt(0))
    return (p, n) -> sigkb(p, n, k, thedict)
end

function makesigma(k)
    thedict = Dict((0, 0)=>0)
    if k==0
        return (p, n)-> n+1
    end
    return (p, n) -> sigk(p, n, k, thedict)
end

function sigkb(p, n, k,thedict)
    if haskey(thedict, (p, n))
        return thedict[(p, n)]
    end
    bp = BigInt(p)
    tmp = sum(bp^(i*k) for i in 0:n)
    thedict[(p, n)] = tmp
    return tmp
end

function sigk(p, n, k,thedict)
    if haskey(thedict, (p, n))
        return thedict[(p, n)]
    end
    @inbounds tmp = sum(p^(i*k) for i in 0:n)
    thedict[(p, n)] = tmp
    return tmp
end

function memoize( func)
    thedict = Dict((0, 0)=>0)
    return (p, n)-> memoizeaux(p, n, func, thedict)
end

function memoizeaux(p, n, func, thedict)
    if haskey(thedict, (p, n))
        return thedict[(p, n)]
    end
    res = func(p, n)
    thedict[(p, n)] = res
    return res
end

mobius(p, n) = n==1?-1:0

phifunc(p, n) = p^(n-1)*(p-1)

function tally(xs)
  freqs = Dict{eltype(xs),Int}()
  for x in xs
    freqs[x] = get(freqs, x, 0) + 1
  end
  return freqs
end


# patches are characterized the by following four parameters
type Patch
    β::Float64 # MC's probability of predation
    λ::Float64 # MC's probability of finding food
    α::Int64 # MC's cost to visit
    Y::Int64 # MC's value of food
end

# an Ecosystem is essentially an array of patches
typealias Ecosystem Array{Patch,1};

# the patches from page 54
myplace = Ecosystem([
    Patch(0., 0., 1, 0),
    Patch(0.004, 0.4, 1, 3),
    Patch(0.02, 0.6, 1, 5)
    ]);

x_c = 3::Int64; # critical energy threshold
C = 10::Int64; # max capacity
T = 20::Int64; # time horizon

function xprime(x::Int64, α::Int64, y::Int64, x_critical::Int64, C::Int64)
    return clamp(x - α + y, x_critical, C)
end;

function xprimeprime(x::Int64, α::Int64, x_critical::Int64, C::Int64)
    return clamp(x - α, x_critical, C)
end;

function findOptimalSelections(
    aplace::Ecosystem,
    C::Int64,
    T::Int64,
    x_c::Int64)
    Fitness = zeros(Float64, C+1, T)
    Selection = zeros(Int64, C+1, T)
    for t = T:-1:1
        for x = 0:C
            if t == T
                Fitness[x+1, t] = x > x_c ? 1.0 : 0.0
            elseif x <= x_c
                Fitness[x+1, t] = 0.0
            else
                patchoptions = [(1-P.β)*
                    (P.λ*Fitness[xprime(x, P.α, P.Y, x_c, C)+1,t+1] +
                    (1.0-P.λ)*Fitness[xprimeprime(x, P.α, x_c, C)+1,t+1])
                    for P in aplace]
                (Fitness[x+1, t], Selection[x+1, t]) = findmax(patchoptions)
            end
        end
    end
    return Selection, Fitness
end;

(Selection, Fitness) = findOptimalSelections(myplace, C, T, x_c);

## compare to Mangel and Clark page 55
Fitness[5:11,19:-1:1]


################################################################################
### UDT Decomposition (no pivoting)
################################################################################

# fully in-place
# Probably don't use this as it has a larger error than the pivoted version

# Taken from julia Base
@inline function reflector!(x::AbstractVector)
    n = length(x)
    @inbounds begin
        ξ1 = x[1]
        normu = abs2(ξ1)
        @turbo for i = 2:n
            normu += abs2(x[i])
        end
        if iszero(normu)
            return zero(ξ1/normu) #zero(ξ1) ?
        end
        normu = sqrt(normu)
        ν = LinearAlgebra.copysign(normu, real(ξ1))
        ξ1 += ν
        x[1] = -ν
        @turbo for i = 2:n
            x[i] /= ξ1
        end
    end
    ξ1/ν
end

# Needed for pivoted as well
@inline function reflectorApply!(x::AbstractVector{<: Real}, τ::Real, A::StridedMatrix{<: Real})
    m, n = size(A)
    @inbounds for j = 1:n
        # dot
        vAj = A[1, j]
        @turbo for i = 2:m
            vAj += conj(x[i]) * A[i, j]
        end

        vAj = conj(τ)*vAj

        # ger
        A[1, j] -= vAj
        @turbo for i = 2:m
            A[i, j] -= x[i]*vAj
        end
    end
    return A
end

"""
    udt_AVX!(U::Matrix, D::Vector, T::Matrix)

In-place calculation of a UDT (unitary - diagonal - (upper-)triangular)
decomposition. The matrix `T` is simultaniously the input matrix that is
decomposed and the triangular output matrix.

This assumes correctly sized square matrices as inputs.
"""
function udt_AVX!(U::AbstractMatrix{C}, D::AbstractVector{C}, input::AbstractMatrix{C}) where {C<:Number}
    # Assumptions:
    # - all matrices same size
    # - input can be changed (input becomes T)

    # @bm "Compute QR decomposition" begin
        n, _ = size(input)
        @inbounds D[n] = zero(C)
        @inbounds for k = 1:(n - 1 + !(C<:Real))
            x = LinearAlgebra.view(input, k:n, k)
            τk = reflector!(x)
            D[k] = τk
            reflectorApply!(x, τk, view(input, k:n, k + 1:n))
        end
    # end

    # @bm "Calculate Q" begin
        copyto!(U, I)
        @inbounds begin
            U[n, n] -= D[n]
            for k = n-1:-1:1
                for j = k:n
                    vBj = U[k,j]
                    @turbo for i = k+1:n
                        vBj += conj(input[i,k]) * U[i,j]
                    end
                    vBj = D[k]*vBj
                    U[k,j] -= vBj
                    @turbo for i = k+1:n
                        U[i,j] -= input[i,k]*vBj
                    end
                end
            end
        end
        # U done
    # end

    # @bm "Calculate R" begin
        @inbounds for j in 1:n-1
            @turbo for i in max(1, j + 1):n
                input[i,j] = zero(input[i,j])
            end
        end
    # end

    # @bm "Calculate D" begin
        @inbounds for i in 1:n
            D[i] = abs(real(input[i, i]))
        end
    # end

    # @bm "Calculate T" begin
        @turbo for i in 1:n
            d = 1.0 / D[i]
            for j in 1:n
                input[i, j] = d * input[i, j]
            end
        end
    # end

    nothing
end

################################################################################
### UDT Decomposition (with pivoting)
################################################################################

# this method skips using a view(x, j:n, j) and requires passing of
# normu = dot(view(x, j:n, j), view(x, j:n, j))
@inline function reflector!(x::Matrix{C}, normu, j=1, n=size(x, 1)) where {C <: Real}
    @inbounds begin
        ξ1 = x[j, j]
        if iszero(normu)
            return zero(ξ1) #zero(ξ1/normu)
        end
        normu = sqrt(normu)
        ν = LinearAlgebra.copysign(normu, real(ξ1))
        ξ1 += ν
        x[j, j] = -ν
        @turbo for i = j+1:n
            x[i, j] /= ξ1
        end
    end
    ξ1/ν
end

function indmaxcolumn(A::Matrix{C}, j=1, n=size(A, 1)) where {C <: Real}
    squared_norm = 0.0
    @turbo for k in j:n
        squared_norm += abs2(A[k, j])
    end
    ii = j
    @inbounds for i in j+1:n
        mi = 0.0
        @turbo for k in j:n
            mi += abs2(A[k, i])
        end
        if abs(mi) > squared_norm
            squared_norm = mi
            ii = i
        end
    end
    return ii, squared_norm
end

# Much higher accuracy, but a bit slower
"""
    udt_AVX_pivot!(
        U::Matrix, D::Vector, T::Matrix[,
        pivot::Vector, temp::Vector, apply_pivoting = Val(true)
    ])

In-place calculation of a UDT (unitary - diagonal - (upper-)triangular)
decomposition. The matrix `T` is simultaniously the input matrix that is
decomposed and the triangular output matrix.

If `apply_pivoting = Val(true)` the `T` matrix will be pivoted such that
`U * Diagonal(D) * T` matches the input.
If `apply_pivoting = Val(false)` `T` will be a "dirty" upper triangular matrix
(i.e. with random values elsewhere) which still requires pivoting.
`rdivp!(A, T, temp, pivot)` is built explicitly for this case - it applies the
pivoting while calculating `A T^-1`. Warning: BlockDiagonal matrices use
per-block pivoting.

This assumes correctly sized square matrices as inputs.
"""
function udt_AVX_pivot!(
        U::AbstractArray{C, 2},
        D::AbstractArray{C, 1},
        input::AbstractArray{C, 2},
        pivot::AbstractArray{Int64, 1} = Vector(UnitRange(1:size(input, 1))),
        temp::AbstractArray{C, 1} = Vector{C}(undef, length(D)),
        apply_pivot::Val = Val(true)
    ) where {C<:Real}
    # Assumptions:
    # - all matrices same size
    # - input can be mutated (input becomes T)

    # @bm "reset pivot" begin
        n = size(input, 1)
        @inbounds for i in 1:n
            pivot[i] = i
        end
    # end

    # @bm "QR decomposition" begin
        @inbounds for j = 1:n
            # Find column with maximum norm in trailing submatrix
            # @bm "get jm" begin
                jm, maxval = indmaxcolumn(input, j, n)
            # end

            # @bm "pivot" begin
                if jm != j
                    # Flip elements in pivoting vector
                    tmpp = pivot[jm]
                    pivot[jm] = pivot[j]
                    pivot[j] = tmpp

                    # Update matrix with
                    @turbo for i = 1:n
                        tmp = input[i,jm]
                        input[i,jm] = input[i,j]
                        input[i,j] = tmp
                    end
                end
            # end

            # Compute reflector of columns j
            # @bm "Reflector" begin
                τj = reflector!(input, maxval, j, n)
                temp[j] = τj
            # end

            # Update trailing submatrix with reflector
            # @bm "apply" begin
                # TODO optimize?
                x = LinearAlgebra.view(input, j:n, j)
                MonteCarlo.reflectorApply!(x, τj, LinearAlgebra.view(input, j:n, j+1:n))
            # end
        end
    # end

    # @bm "Calculate Q" begin
        copyto!(U, I)
        @inbounds begin
            U[n, n] -= temp[n]
            for k = n-1:-1:1
                for j = k:n
                    vBj = U[k,j]
                    @turbo for i = k+1:n
                        vBj += conj(input[i,k]) * U[i,j]
                    end
                    vBj = temp[k]*vBj
                    U[k,j] -= vBj
                    @turbo for i = k+1:n
                        U[i,j] -= input[i,k]*vBj
                    end
                end
            end
        end
        # U done
    # end

    # @bm "Calculate D" begin
        @inbounds for i in 1:n
            D[i] = abs(real(input[i, i]))
        end
    # end

    # @bm "pivoted zeroed T w/ inv(D)" begin
        _apply_pivot!(input, D, temp, pivot, apply_pivot)
    # end

    nothing
end

function _apply_pivot!(input::Matrix{C}, D, temp, pivot, ::Val{true}) where {C <: Real}
    n = size(input, 1)
    @inbounds for i in 1:n
        d = 1.0 / D[i]
        @inbounds for j in 1:i-1
            temp[pivot[j]] = zero(C)
        end
        @turbo for j in i:n
            temp[pivot[j]] = d * input[i, j]
        end
        @turbo for j in 1:n
            input[i, j] = temp[j]
        end
    end
end
function _apply_pivot!(input::Matrix{C}, D, temp, pivot, ::Val{false}) where {C <: Real}
    n = size(input, 1)
    @inbounds for i in 1:n
        d = 1.0 / D[i]
        @turbo for j in i:n
            input[i, j] = d * input[i, j]
        end
    end
end

################################################################################
### Complex pivoted UDT
################################################################################

@inline function reflector!(x::Matrix{C}, normu, j=1, n=size(x, 1)) where {C <: ComplexF64}
    @inbounds begin
        ξ1 = x[j, j]
        if iszero(normu)
            return zero(ξ1) #zero(ξ1/normu)
        end
        normu = sqrt(normu)
        ν = LinearAlgebra.copysign(normu, real(ξ1))
        ξ1 += ν
        x[j, j] = -ν
        for i = j+1:n
            x[i, j] /= ξ1
        end
    end
    ξ1/ν
end

@inline function reflectorApply!(x::AbstractVector, τ::Number, A::StridedMatrix)
    m, n = size(A)
    @inbounds for j = 1:n
        # dot
        vAj = A[1, j]
        @inbounds for i = 2:m
            vAj += conj(x[i]) * A[i, j]
        end

        vAj = conj(τ)*vAj

        # ger
        A[1, j] -= vAj
        @inbounds for i = 2:m
            A[i, j] -= x[i]*vAj
        end
    end
    return A
end

function indmaxcolumn(A::AbstractMatrix{C}, j=1, n=size(A, 1)) where {C <: ComplexF64}
    max = 0.0
    for k in j:n
        max += abs2(A[k, j])
    end
    ii = j
    @inbounds for i in j+1:n
        mi = 0.0
        for k in j:n
            mi += abs2(A[k, i])
        end
        if abs(mi) > max
            max = mi
            ii = i
        end
    end
    return ii, max
end

# avx-less method for compatability with ComplexF64
function udt_AVX_pivot!(
        U::AbstractArray{C, 2},
        D::AbstractArray{Float64, 1},
        input::AbstractArray{C, 2},
        pivot::AbstractArray{Int64, 1} = Vector(UnitRange(1:size(input, 1))),
        temp::AbstractArray{C, 1} = Vector{C}(undef, length(D)),
        apply_pivot::Val = Val(true)
    ) where {C <: Complex}
    # Assumptions:
    # - all matrices same size
    # - input can be changed (input becomes T)

    # @bm "reset pivot" begin
        n = size(input, 1)
        @inbounds for i in 1:n
            pivot[i] = i
        end
    # end

    # @bm "QR decomposition" begin
        @inbounds for j = 1:n
            # Find column with maximum norm in trailing submatrix
            # @bm "get jm" begin
                jm, maxval = indmaxcolumn(input, j, n)
            # end

            # @bm "pivot" begin
                if jm != j
                    # Flip elements in pivoting vector
                    tmpp = pivot[jm]
                    pivot[jm] = pivot[j]
                    pivot[j] = tmpp

                    # Update matrix with
                    for i = 1:n
                        tmp = input[i,jm]
                        input[i,jm] = input[i,j]
                        input[i,j] = tmp
                    end
                end
            # end

            # Compute reflector of columns j
            # @bm "Reflector" begin
                τj = reflector!(input, maxval, j, n)
                temp[j] = τj
            # end

            # Update trailing submatrix with reflector
            # @bm "apply" begin
                # TODO optimize?
                x = LinearAlgebra.view(input, j:n, j)
                MonteCarlo.reflectorApply!(x, τj, LinearAlgebra.view(input, j:n, j+1:n))
            # end
        end
    # end

    # @bm "Calculate Q" begin
        copyto!(U, I)
        @inbounds begin
            U[n, n] -= temp[n]
            for k = n-1:-1:1
                for j = k:n
                    vBj = U[k,j]
                    for i = k+1:n
                        vBj += conj(input[i,k]) * U[i,j]
                    end
                    vBj = temp[k]*vBj
                    U[k,j] -= vBj
                    for i = k+1:n
                        U[i,j] -= input[i,k]*vBj
                    end
                end
            end
        end
        # U done
    # end

    # @bm "Calculate D" begin
        @inbounds for i in 1:n
            D[i] = abs(real(input[i, i]))
        end
    # end

    # @bm "pivoted zeroed T w/ inv(D)" begin
        _apply_pivot!(input, D, temp, pivot, apply_pivot)
    # end

    nothing
end

function _apply_pivot!(input::Matrix{C}, D, temp, pivot, ::Val{true}) where {C <: Complex}
    n = size(input, 1)
    @inbounds for i in 1:n
        d = 1.0 / D[i]
        @inbounds for j in 1:i-1
            temp[pivot[j]] = zero(C)
        end
        for j in i:n
            temp[pivot[j]] = d * input[i, j]
        end
        for j in 1:n
            input[i, j] = temp[j]
        end
    end
end
function _apply_pivot!(input::Matrix{C}, D, temp, pivot, ::Val{false}) where {C <: Complex}
    n = size(input, 1)
    @inbounds for i in 1:n
        d = 1.0 / D[i]
        for j in i:n
            input[i, j] = d * input[i, j]
        end
    end
end


"""
Regression tests

Tests compare results with results retrieved from an ealier implementation of the integration method
(data and previous results from https://doi.org/10.1039/C9CP00435A).
"""

function read_test_data(fn)
    return open(fn, "r") do io
        lines = readlines(io)
        N = length(lines)
        X = Array{Float64}(undef, N, 2)
        for i ∈ 1:N
            l = lines[i]
            x = map(x -> parse(Float64, x), split(l))
            X[i, 1] = x[1]
            X[i, 2] = x[2]
        end
        I = sortperm(X[:, 1])
        X[I, :]
    end
end

@testset "regression test with legacy data (1)" begin
    # test dataset: methanol-3-methylphenylacetylene
    seed!(1)
    dat = joinpath(@__DIR__, "t1.dat") |> read_test_data
    x, y = dat[:, 1], dat[:, 2]
    # remove numeric noise from x and make grid uniform
    x = x[1]:(diff(x) |> mean):x[end] |> collect
    # integration workflow
    spec = crop(Curve(x, y), 3600.0, 3680.0)
    noise = NoiseSample(crop(Curve(x, y), 3750.0, 3850.0))
    nm = fit_noise(noise)
    uspec = add_noise(spec, nm)
    bnds = UncertainBound([3620.0, 3639.0], scale_shift_beta(2, 2, 4.75, 5.25), uspec)
    areas = mc_integrate(uspec, bnds, subtract_baseline=true)
    result = areas[2]/areas[1]
    # compare percentiles to previous results
    prev_result = [3.79e-01, 5.92e-01, 8.60e-01] # 2.5, 50 and 97.5 percentiles from https://doi.org/10.1039/C9CP00435A
    for (p, pr) in zip([2.5, 50, 97.5], prev_result)
        r = percentile(result, p)
        @test abs(r - pr)/r < 0.025 # here, the new implementation reproduces the previous results within 2.5%
    end
end

@testset "regression test with legacy data (2)" begin
    # test dataset: methanol-d-3-methylphenylacetylene
    seed!(1)
    dat = joinpath(@__DIR__, "t2.dat") |> read_test_data
    x, y = dat[:, 1], dat[:, 2]
    # remove numeric noise from x and make grid uniform
    x = x[1]:(diff(x) |> mean):x[end] |> collect
    # integration workflow
    spec = crop(Curve(x, y), 2650.0, 2710.0)
    noise = NoiseSample(crop(Curve(x, y), 2500.0, 2600.0), 1)
    nm = fit_noise(noise; α_guess=2.5e-6, λ_guess=1.0)
    uspec = add_noise(spec, nm)
    bnds = UncertainBound([2671.0, 2685.0], scale_shift_beta(2, 2, 4.75, 5.25), uspec)
    areas = mc_integrate(uspec, bnds, subtract_baseline=true)
    result = areas[2]/areas[1]
    # compare percentiles to previous results
    p25_pr = 2.24e-01 # previous 2.5 percentile ...
    p50_pr = 5.27e-01 # previous 50 percentile ...
    p975_pr = 9.18e-01 # previous 97.5 percentiles from https://doi.org/10.1039/C9CP00435A
    # in case of methanol-d-3-methylphenylacetylene, the percentiles are less well reproduced
    # but the agreement is still better than 7%
    r = percentile(result, 2.5)
    @test abs(r - p25_pr)/r ≈ 0.0739333525 rtol = 1e-7 # worst disagreement (but for smallest number, 0.24 vs. 0.22 is not that bad ...)
    r = percentile(result, 50)
    @test abs(r - p50_pr)/r ≈ 0.000140994154 rtol = 1e-6 # the median is very well reproduced
    r = percentile(result, 97.5)
    @test abs(r - p975_pr)/r ≈ 0.04008646654 rtol = 1e-7
end


using SpecialFunctions: besseli, logfactorial, loggamma
import Distributions
using Distributions: ContinuousMultivariateDistribution, AbstractRNG

log_besseli_term(order, m, x) = -logfactorial(m) - loggamma(m + order + 1) + (2*m + order)*log(x/2)

isunitvec(v) = isapprox(sum(x^2 for x in v), 1.)

function logsumexp(arr)
    maxarr = maximum(arr)
    return maxarr + log(sum(exp.(arr .- maxarr)))
end

function _log_besseli(order, x; n_terms=100)
    log_terms = Vector{Float64}(undef, n_terms)
    for m in 0:n_terms-1
        log_terms[m+1] = log_besseli_term(order, m, x)
    end
    return logsumexp(log_terms)
end

const max_lookup_input = 50000
const log_besseli_0_lookup = Vector{Float64}(undef, max_lookup_input) # order 0 (used by plane rotation)
const log_besseli_1_lookup = Vector{Float64}(undef, max_lookup_input) # order 1 (used by 3D rotation)

function log_besseli(order, x; n_terms=Int(floor(x/2)) + 20)
    if x > 500.0 && (order == 0 || order == 1)
        if x > max_lookup_input
            error("log besseli not available for x > $max_loopup_input, got: x = $x")
        end
        floorx = Int(floor(x))
        if order == 0
            if !isassigned(log_besseli_0_lookup, floorx)
                log_besseli_0_lookup[floorx] = _log_besseli(0, floorx; n_terms=n_terms)
            end
            return log_besseli_0_lookup[floorx]
        elseif order == 1
            if !isassigned(log_besseli_1_lookup, floorx)
                log_besseli_1_lookup[floorx] = _log_besseli(1, floorx; n_terms=n_terms)
            end
            return log_besseli_1_lookup[floorx]
        else
            @assert false
        end
    else
        return log(besseli(order, x))
    end
end

function Distributions._vmflck(p, κ)
    T = typeof(κ)
    hp = T(p/2)
    q = hp - 1
    return q * log(κ) - hp * log(2π) - log_besseli(q, κ)
end

import Distributions: VonMisesFisherSampler, _rand!
import LinearAlgebra: qr!, mul!

struct VonMisesFisherSamplerAlternate <: Distributions.Sampleable{Distributions.Multivariate,Distributions.Continuous}
    p::Int          # the dimension
    κ::Float64
    b::Float64
    x0::Float64
    c::Float64
    Q::Matrix{Float64}
end

function Distributions.VonMisesFisherSampler(μ::Vector{Float64}, κ::Float64)
    p = length(μ)
    b = _vmf_bval(p, κ)
    x0 = (1.0 - b) / (1.0 + b)
    c = κ * x0 + (p - 1) * log1p(-abs2(x0))
    Q = _vmf_rotmat(μ)
    VonMisesFisherSamplerAlternate(p, κ, b, x0, c, Q)
end

function Distributions._rand!(rng::AbstractRNG, spl::VonMisesFisherSamplerAlternate,
                x::AbstractVector, t::AbstractVector)
    w = _vmf_genw(rng, spl)
    p = spl.p
    t[1] = w
    s = 0.0
    for i = 2:p
        t[i] = ti = randn(rng)
        s += abs2(ti)
    end

    # normalize t[2:p]
    r = sqrt((1.0 - abs2(w)) / s)
    for i = 2:p
        t[i] *= r
    end

    # rotate
    mul!(x, spl.Q, t)
    return x
end

Distributions._rand!(rng::AbstractRNG, spl::VonMisesFisherSamplerAlternate, x::AbstractVector) =
    _rand!(rng, spl, x, Vector{Float64}(undef, length(x)))

function Distributions._rand!(rng::AbstractRNG, spl::VonMisesFisherSamplerAlternate, x::AbstractMatrix)
    t = Vector{Float64}(undef, size(x, 1))
    for j = 1:size(x, 2)
        _rand!(rng, spl, view(x,:,j), t)
    end
    return x
end

### Core computation

_vmf_bval(p::Int, κ::Real) = (p - 1) / (2.0κ + sqrt(4 * abs2(κ) + abs2(p - 1)))

function _vmf_genw(rng::AbstractRNG, p, b, x0, c, κ)
    # generate the W value -- the key step in simulating vMF
    #
    #   following movMF's document
    #

    r = (p - 1) / 2.0
    betad = Distributions.Beta(r, r)
    z = rand(rng, betad)
    w = (1.0 - (1.0 + b) * z) / (1.0 - (1.0 - b) * z)
    while κ * w + (p - 1) * log(1 - x0 * w) - c < log(rand(rng))
        z = rand(rng, betad)
        w = (1.0 - (1.0 + b) * z) / (1.0 - (1.0 - b) * z)
    end
    return w::Float64
end

_vmf_genw(rng::AbstractRNG, s::VonMisesFisherSamplerAlternate) =
    _vmf_genw(rng, s.p, s.b, s.x0, s.c, s.κ)

function _vmf_rotmat(u::Vector{Float64})
    # construct a rotation matrix Q
    # s.t. Q * [1,0,...,0]^T --> u
    #
    # Strategy: construct a full-rank matrix
    # with first column being u, and then
    # perform QR factorization
    #

    p = length(u)
    A = zeros(p, p)
    copyto!(view(A,:,1), u)

    # let k the be index of entry with max abs
    k = 1
    a = abs(u[1])
    for i = 2:p
        @inbounds ai = abs(u[i])
        if ai > a
            k = i
            a = ai
        end
    end

    # other columns of A will be filled with
    # indicator vectors, except the one
    # that activates the k-th entry
    i = 1
    for j = 2:p
        if i == k
            i += 1
        end
        A[i, j] = 1.0
    end

    # perform QR factorization
    Q = Matrix(qr!(A).Q)
    if dot(view(Q,:,1), u) < 0.0  # the first column was negated
        for i = 1:p
            @inbounds Q[i,1] = -Q[i,1]
        end
    end
    return Q
end


#!/usr/bin/env julia

# this test isn't run by Pkg.test

# it can be run to test Juno/Jupyter integration...

using Luxor

using Test

using Random
Random.seed!(42)

function draw_hypotrochoids()
    background("grey10")
    sethue("antiquewhite")
    setopacity(0.8)
    tiles = Tiler(1200, 1200, 6, 6)
    for (pos, n) in tiles
        gsave()
        setline(.5)
        translate(pos)

        R = 45; r = n; d = 20
        sethue("antiquewhite")
        poly(hypotrochoid(R, n, d, :stroke, stepby=pi/300, vertices=true), :stroke, close=true)
        text(string("R = $(R), r = $(r), d = $(d)"), halign=:center, O.x, O.y+tiles.tileheight/2)

        # test perimeterlength
        ptslist = hypotrochoid(R, n, d, :stroke, stepby=pi/300, vertices=true)
        polyperimeter(ptslist)

        R = 2n; r = 12; d = 5
        sethue("orange")
        (R == r) && (r += 1)
        p = hypotrochoid(R, r, d, :none, stepby=pi/305, vertices=true)
        polydistances(p) # just testing
        poly(p, :stroke, close=true)
        text(string("R = $(R), r = $(r), d = $(d)"), halign=:center, O.x, O.y+tiles.tileheight/2 + 10)

        # testing portions of polys
        portion = rand()
        setline(.25)
        sethue("magenta")
        section = polyportion(p, portion)
        poly(section, :stroke, close=false)
        sethue("green")
        circle(section[end], 2, :stroke)
        poly(polyremainder(p, portion), :stroke, close=false)

        text(string("portion ", round(portion, digits=2)), halign=:center, O.x, O.y+tiles.tileheight/2 + 30)

        sethue("cyan")
        setline(0.3)
        poly(offsetpoly(p, 15), :stroke, close=true)
        text("offset 15", halign=:center, O.x, O.y+tiles.tileheight/2 + 20)

        grestore()
    end
end

function test_hypotrochoid_at_png()
    @png begin
        draw_hypotrochoids()
    end 1200 1200
end
function test_hypotrochoid_at_svg()
    @svg begin
        draw_hypotrochoids()
    end 1200 1200
end
function test_hypotrochoid_at_pdf()
    @pdf begin
        draw_hypotrochoids()
    end 1200 1200
end

Random.seed!(42)

test_hypotrochoid_at_png() ; println("output PNG")
test_hypotrochoid_at_svg() ; println("output SVG")
test_hypotrochoid_at_pdf() ; println("output PDF")

#testing @manipulate

using Interact

function makecircle(r)
    d = Drawing(300, 300, :svg)
    sethue("black")
    setline(0.1)
    origin()
    hypotrochoid(150, r, 35, :stroke)
    finish()
    return d
end

@manipulate for r in 5:150
    makecircle(r)
end