"""
    This provides units and creates a non-dimensionalization object
"""
module Units
using Unitful
import Unitful: superscript
using Parameters

import Base.show
using GeoParams: AbstractMaterialParam, AbstractMaterialParamsStruct, AbstractPhaseDiagramsStruct, PerpleX_LaMEM_Diagram

# Define additional units that are useful in geodynamics
@unit    Myrs  "Myrs"   MillionYears    1000000u"yr"    false

function __init__()
    Unitful.register(Units)
end

Unitful.register(Units)

# Define a number of useful units
const km    = u"km"
const m     = u"m"
const cm    = u"cm"
const mm    = u"mm"
const Myrs  = u"Myrs"
const yr    = u"yr"
const s     = u"s"
const kg    = u"kg"
const g     = u"g"
const Pa    = u"Pa"
const MPa   = u"MPa"
const kbar  = u"kbar"
const Pas   = u"Pa*s"
const K     = u"K"
const C     = u"°C"
const mol   = u"mol"
const kJ    = u"kJ"
const J     = u"J"
const Watt  = u"W"
const μW    = u"μW"

export
    km, m, cm, mm, Myrs, yr, s, MPa, Pa, kbar, Pas, K, C, g, kg, mol, J, kJ, Watt, μW,
    GeoUnit, GeoUnits, GEO_units, SI_units, NO_units, AbstractGeoUnits,
    Nondimensionalize, Nondimensionalize!, Dimensionalize, Dimensionalize!,
    superscript, upreferred, GEO, SI, NONE, isDimensional, Value, NumValue, Unit

"""
AbstractGeoUnits

Abstract supertype for geo units.
"""
abstract type AbstractGeoUnits{TYPE} end

abstract type AbstractUnitType end

struct GEO <: AbstractUnitType end
struct SI  <: AbstractUnitType end
struct NONE<: AbstractUnitType end

"""
    Structure that holds a GeoUnit parameter and their dimensions

    Having that is useful, as non-dimensionalization removes the units from a number
    and we thus no longer know how to transfer it back to the correct units.

"""
mutable struct GeoUnit
    val                 # the actual value ()
    unit :: Unitful.FreeUnits
end

GeoUnit(v::Unitful.Quantity)            =   GeoUnit(v, unit(v))     # store the units
GeoUnit(v::Number)                      =   GeoUnit(v, NoUnits)     # in case we just have a number with no units
GeoUnit(v::Array)                       =   GeoUnit(v, NoUnits)     # array, no units
GeoUnit(v::Array{Unitful.Quantity})     =   GeoUnit(v, unit.(v))    # with units
GeoUnit(v::StepRange)                   =   GeoUnit(v, unit.(v))    # with units
Value(v::GeoUnit)                       =   v.val                   # get value of GeoUnit
NumValue(v::GeoUnit)                    =   ustrip(v.val)           # numeric value, with no units
Unit(v::GeoUnit)                        =   v.unit                  # extract unit

Base.convert(::Type{Float64}, v::GeoUnit)       =   v.val
Base.convert(::Type{GeoUnit}, v::Quantity)      =   GeoUnit(v)
Base.convert(::Type{GeoUnit}, v::Number)        =   GeoUnit(v, NoUnits)
Base.convert(::Type{GeoUnit}, v::Vector)        =   GeoUnit(v, unit(v[1]))
Base.convert(::Type{GeoUnit}, v::Array)         =   GeoUnit(v, unit(v[1]))
Base.convert(::Type{GeoUnit}, v::StepRangeLen)  =   GeoUnit(v, unit(v[1]))
Base.convert(::Type{GeoUnit}, v::StepRange)     =   GeoUnit(v, unit(v[1]))

# define a few basic routines so we can easily operate with GeoUnits
Base.show(io::IO, x::GeoUnit)   =   println(x.val)
Base.length(v::GeoUnit)         =   length(v.val)
Base.size(v::GeoUnit)           =   size(v.val)

Base.:*(x::GeoUnit, y::Number)  = x.val*y
Base.:+(x::GeoUnit, y::Number)  = x.val+y
Base.:/(x::GeoUnit, y::Number)  = x.val/y
Base.:-(x::GeoUnit, y::Number)  = x.val-y

Base.:*(x::GeoUnit, y::GeoUnit) = x.val*y.val
Base.:+(x::GeoUnit, y::GeoUnit) = x.val+y.val
Base.:/(x::GeoUnit, y::GeoUnit) = x.val/y.val
Base.:-(x::GeoUnit, y::GeoUnit) = x.val-y.val

Base.:*(x::Number, y::GeoUnit)  = y.val*x
Base.:+(x::Number, y::GeoUnit)  = y.val+x
Base.:/(x::Number, y::GeoUnit)  = x/y.val
Base.:-(x::Number, y::GeoUnit)  = x-y.val

Base.:*(x::GeoUnit, y::Quantity)  = GeoUnit(x.val*y, x.unit)
Base.:+(x::GeoUnit, y::Quantity)  = GeoUnit(x.val+y, x.unit)
Base.:/(x::GeoUnit, y::Quantity)  = GeoUnit(x.val/y, x.unit)
Base.:-(x::GeoUnit, y::Quantity)  = GeoUnit(x.val-y, x.unit)

Base.:*(x::GeoUnit, y::Array)   = GeoUnit(x.val*y, x.unit)
Base.:/(x::GeoUnit, y::Array)   = GeoUnit(x.val/y, x.unit)
Base.:+(x::GeoUnit, y::Array)   = GeoUnit(x.val+y, x.unit)
Base.:-(x::GeoUnit, y::Array)   = GeoUnit(x.val-y, x.unit)

Base.:*(x::Array, y::GeoUnit)   = GeoUnit(x*y.val, y.unit)
Base.:/(x::Array, y::GeoUnit)   = GeoUnit(x/y.val, y.unit)
Base.:+(x::Array, y::GeoUnit)   = GeoUnit(x +y.val, y.unit)
Base.:-(x::Array, y::GeoUnit)   = GeoUnit(x -y.val, y.unit)

Base.:*(x::GeoUnit, y::Vector)   = GeoUnit(x.val*y, x.unit)
Base.:/(x::GeoUnit, y::Vector)   = GeoUnit(x.val/y, x.unit)
Base.:+(x::GeoUnit, y::Vector)   = GeoUnit(x.val+y, x.unit)
Base.:-(x::GeoUnit, y::Vector)   = GeoUnit(x.val-y, x.unit)

Base.:*(x::GeoUnit, y::StepRange)   = GeoUnit(x.val*y, x.unit)
Base.:/(x::GeoUnit, y::StepRange)   = GeoUnit(x.val/y, x.unit)
Base.:+(x::GeoUnit, y::StepRange)   = GeoUnit(x.val+y, x.unit)
Base.:-(x::GeoUnit, y::StepRange)   = GeoUnit(x.val-y, x.unit)

Base.getindex(x::GeoUnit, i::Int64, j::Int64, k::Int64) = x.val[i,j,k]
Base.getindex(x::GeoUnit, i::Int64, j::Int64) = x.val[i,j]
Base.getindex(x::GeoUnit, i::Int64) = x.val[i]

Base.setindex!(x::GeoUnit, v::Any, i::Int64, j::Int64, k::Int64) = x.val[i,j,k] = v
Base.setindex!(x::GeoUnit, v::Any, i::Int64, j::Int64) = x.val[i,j] = v
Base.setindex!(x::GeoUnit, v::Any, i::Int64) = x.val[i] = v

"""
    GeoUnits

    Structure that holds parameters used for non-dimensionalization
"""
@with_kw_noshow struct GeoUnits{TYPE}
    # Selectable input parameters
    temperature     =   1               #   Characteristic temperature  [C or K]
    length          =   1               #   Characteristic length unit  [km, m or -]
    stress          =   1               #   Characteristic stress unit  [MPa, Pa or -]
    time            =   1               #   Characteristic time unit    [Myrs, s pr -]
    viscosity       =   1

    # primary characteristic units
    K               =   1;                      # temperature in SI units for material parameter scaling
    s               =   1;                      # time in SI units for material parameter scaling
    m               =   1;                      # length in SI units for material parameter scaling
    Pa              =   1;                      # stress in SI units
    kg              =   upreferred(Pa*m*s^2)    # compute mass from pascal. Note: this may result in very large values

    # Main SI units (used later for nondimensionalization)
    Length          =   m
    Mass            =   kg
    Time            =   s
    Temperature     =   K
    Amount          =   1mol
    Second          =   s

    # Not defined, as they are not common in geodynamics:
    #Current
    #Luminosity

    # Derived units
    N               =   kg*m/s^2        # Newton
    J               =   N*m             # Joule
    W               =   J/s             # Watt
    area            =   m^2             # area   in SI units for material parameter scaling
    volume          =   m^3             # volume
    velocity        =   m/s
    density         =   kg/m^3
    acceleration    =   m/s^2
    force           =   kg*m/s^2
    strainrate      =   1/s
    heatcapacity    =   J/kg/K
    conductivity    =   W/m/K

    # Helpful
    SecYear         =   3600*24*365.25
    Myrs            =   1e6
    cmYear          =   SecYear*100     # to transfer m/s -> cm/yr

end

"""
    GEO_units(;length=1000km, temperature=1000C, stress=10MPa, viscosity=1e20Pas)

Creates a non-dimensionalization object using GEO units.

GEO units implies that upon dimensionalization, `time` will be in `Myrs`, `length` in `km`, stress in `MPa`, etc.
which is more convenient for typical geodynamic simulations than SI units
The characteristic values given as input can be in arbitrary units (`km` or `m`), provided the unit is specified.

# Examples:
```julia-repl
julia> CharUnits = GEO_units()
Employing GEO units
Characteristic values:
         length:      1000 km
         time:        0.3169 Myrs
         stress:      10 MPa
         temperature: 1000.0 °C
julia> CharUnits.velocity
1.0e-7 m s⁻¹
```
If we instead have a crustal-scale simulation, it is likely more appropriate to use a different characteristic `length`:
```julia-repl
julia> CharUnits = GEO_units(length=10km)
Employing GEO units
Characteristic values:
         length:      10 km
         time:        0.3169 Myrs
         stress:      10 MPa
         temperature: 1000.0 °C
```
"""
function GEO_units(;length=1000km, temperature=1000C, stress=10MPa, viscosity=1e20Pas)

    if unit(temperature)==NoUnits;  temperature = temperature*C;        end
    if unit(length)==NoUnits;       length      = length*u"km";         end
    if unit(stress)==NoUnits;       stress      = stress*u"MPa";        end
    if unit(viscosity)==NoUnits;    viscosity   = viscosity*u"Pa*s";    end

    T       =   uconvert(C,      temperature)
    Le      =   uconvert(km,     length);
    Sigma   =   uconvert(MPa,    stress)
    Eta     =   uconvert(Pas,    viscosity)

    T_SI    =   uconvert(K,      T);
    Le_SI   =   uconvert(m,      Le);
    Sigma_SI=   uconvert(Pa,     Sigma)
    Time_SI =   Eta/Sigma_SI;
    t       =   uconvert(Myrs,   Time_SI)

    GeoUnits{GEO}(length=Le, temperature=T,      stress=Sigma,  viscosity=Eta, time=t,
                  m=Le_SI,   K=T_SI,   Pa=Sigma_SI,  s=Time_SI)
end

"""
    CharUnits = SI_units(length=1000m, temperature=1000K, stress=10Pa, viscosity=1e20)

Specify the characteristic values using SI units

# Examples:
```julia-repl
julia> CharUnits = SI_units(length=1000m)
Employing SI units
Characteristic values:
         length:      1000 m
         time:        1.0e19 s
         stress:      10 Pa
         temperature: 1000.0 K
```
Note that the same can be achieved if the input is given in `km`:
```julia-repl
julia> CharUnits = SI_units(length=1km)
```
"""
function SI_units(;length=1000m, temperature=1000K, stress=10Pa, viscosity=1e20Pas)

    if unit(temperature)==NoUnits;  temperature = temperature*K;     end
    if unit(length)==NoUnits;       length      = length*u"m";         end
    if unit(stress)==NoUnits;       stress      = stress*u"Pa";        end
    if unit(viscosity)==NoUnits;    viscosity   = viscosity*u"Pa*s";    end

    T       =   uconvert(K,     temperature)
    Le      =   uconvert(m,     length);
    Sigma   =   uconvert(Pa,    stress)
    Eta     =   uconvert(Pas,   viscosity)

    T_SI    =   uconvert(K,      T);
    Le_SI   =   uconvert(m,      Le);
    Sigma_SI=   uconvert(Pa,     Sigma)
    Time_SI =   Eta/Sigma_SI;
    t       =   uconvert(s,      Time_SI)

    GeoUnits{SI}(length=Le, temperature=T,      stress=Sigma,  viscosity=Eta, time=t,
                 m=Le_SI,   K=T_SI,   Pa=Sigma_SI,  s=Time_SI)
end

"""
    CharUnits = NO_units(length=1, temperature=1, stress=1, viscosity=1)

Specify the characteristic values in non-dimensional units

# Examples:
```julia-repl
julia> using GeoParams;
julia> CharUnits = NO_units()
Employing NONE units
Characteristic values:
         length:      1
         time:        1.0
         stress:      1
         temperature: 1.0
```
"""
function NO_units(;length=1, temperature=1, stress=1, viscosity=1)

    if unit(temperature)!=NoUnits;  error("temperature should not have units")    end
    if unit(length)!=NoUnits;       error("length should not have units")    end
    if unit(stress)!=NoUnits;       error("stress should not have units")    end
    if unit(viscosity)!=NoUnits;    error("viscosity should not have units")    end

    T       =   temperature
    Le      =   length;
    Sigma   =   stress
    Eta     =   viscosity
    Time    =   Eta/Sigma;

    GeoUnits{NONE}(length=Le, temperature=T, stress=Sigma, viscosity=Eta, time=Time,
                 m=Le, K=T, Pa=Sigma, s=Time)
end

"""
    Nondimensionalize(param, CharUnits::GeoUnits{TYPE})

Nondimensionalizes `param` using the characteristic values specified in `CharUnits`

# Example 1
```julia-repl
julia> using GeoParams;
julia> CharUnits =   GEO_units();
julia> v         =   3cm/yr
3 cm yr⁻¹
julia> v_ND      =   Nondimensionalize(v, CharUnits)
0.009506426344208684
```
# Example 2
In geodynamics one sometimes encounters more funky units
```julia-repl
julia> CharUnits =   GEO_units();
julia> A         =   6.3e-2MPa^-3.05*s^-1
0.063 MPa⁻³·⁰⁵ s⁻¹
julia> A_ND      =   Nondimensionalize(A, CharUnits)
7.068716262102384e14
```

In case you are interested to see how the units of `A` look like in different units, use this function from the [Unitful](https://github.com/PainterQubits/Unitful.jl) package:
```julia-repl
julia> uconvert(u"Pa^-3.05*s^-1",A)
3.157479571851836e-20 Pa⁻³·⁰⁵
```
and to see it decomposed in the basic `SI` units of length, mass and time:
```julia-repl
julia> upreferred(A)
3.1574795718518295e-20 m³·⁰⁵ s⁵·¹ kg⁻³·⁰⁵
```
"""
function Nondimensionalize(param, g::GeoUnits{TYPE}) where {TYPE}
    if typeof(param) == String
        param_ND = param # The parameter is a string, cannot be nondimensionalized
    elseif unit(param)!=NoUnits
        dim         =   Unitful.dimension(param);                   # Basic SI units
        char_val    =   1.0;
        foreach((typeof(dim).parameters[1])) do y
            val = upreferred(getproperty(g, Unitful.name(y)))       # Retrieve the characteristic value from structure g
            pow = Float64(y.power)                                  # power by which it should be multiplied
            char_val *= val^pow                                     # multiply characteristic value
        end
        param_ND = upreferred.(param)/char_val
    else
        param_ND = param # The parameter has no units, so there is no way to determine how to nondimensionize it
    end
    return param_ND
end

function Nondimensionalize(param::Array, g::GeoUnits{TYPE}) where {TYPE}
    if unit(param[1])!=NoUnits
        dim         =   Unitful.dimension.(param);                   # Basic SI units
        char_val    =   1.0;
        foreach((typeof(dim[1]).parameters[1])) do y
            val = upreferred(getproperty(g, Unitful.name(y)))       # Retrieve the characteristic value from structure g
            pow = Float64(y.power)                                  # power by which it should be multiplied
            char_val *= val^pow                                     # multiply characteristic value
        end
        param_ND = upreferred.(param)/char_val
    else
        param_ND = param # The parameter has no units, so there is no way to determine how to nondimensionize it
    end
    return param_ND
end

"""
    Nondimensionalize!(param::GeoUnit, CharUnits::GeoUnits{TYPE})

Nondimensionalizes `param` (given as GeoUnit) using the characteristic values specified in `CharUnits` in-place

# Example 1
```julia-repl
julia> using GeoParams;
julia> CharUnits =   GEO_units();
julia> v         =   GeoUnit(3cm/yr)
3 cm yr⁻¹
julia> Nondimensionalize!(v, CharUnits)
0.009506426344208684
```
# Example 2
```julia-repl
julia> CharUnits =   GEO_units();
julia> A         =   GeoUnit(6.3e-2MPa^-3.05*s^-1)
0.063 MPa⁻³·⁰⁵ s⁻¹
julia> A_ND      =   Nondimensionalize(A, CharUnits)
7.068716262102384e14
```
"""
function Nondimensionalize!(param::GeoUnit, g::GeoUnits{TYPE}) where {TYPE}
    if unit.(param.val)!=NoUnits
        dim         =   Unitful.dimension(param.val[1]);                   # Basic SI units
        char_val    =   1.0;
        foreach((typeof(dim).parameters[1])) do y
            val = upreferred(getproperty(g, Unitful.name(y)))       # Retrieve the characteristic value from structure g
            pow = Float64(y.power)                                  # power by which it should be multiplied
            char_val *= val^pow                                     # multiply characteristic value
        end
        param.val = upreferred.(param.val)/char_val;
    else
        param = param # The parameter has no units, so there is no way to determine how to nondimensionize it
    end
end

function Nondimensionalize!(param::String, g::GeoUnits{TYPE}) where {TYPE}
    param_ND = param
    return nothing
end

function Nondimensionalize!(param::Array, g::GeoUnits{TYPE}) where {TYPE}

    if unit.(param[1])!=NoUnits
        dim         =   Unitful.dimension(param[1]);                   # Basic SI units
        char_val    =   1.0;
        foreach((typeof(dim).parameters[1])) do y
            val = upreferred(getproperty(g, Unitful.name(y)))       # Retrieve the characteristic value from structure g
            pow = Float64(y.power)                                  # power by which it should be multiplied
            char_val *= val^pow                                     # multiply characteristic value
        end
        param = upreferred.(param)/char_val;
    else
        param = param # The parameter has no units, so there is no way to determine how to nondimensionize it
    end

end

"""
    Nondimensionalize!(MatParam::AbstractMaterialParam, CharUnits::GeoUnits{TYPE})

Non-dimensionalizes a material parameter structure (e.g., Density, CreepLaw)

"""
function Nondimensionalize!(MatParam::AbstractMaterialParam, g::GeoUnits{TYPE}) where {TYPE}
    for param in fieldnames(typeof(MatParam))
        if typeof(getfield(MatParam, param))==GeoUnit
            z=getfield(MatParam, param)
            Nondimensionalize!(z, g)
            setfield!(MatParam, param, z)
        end
    end
end

"""
    Nondimensionalize!(phase_mat::MaterialParams, g::GeoUnits{TYPE})

Nondimensionalizes all fields within the Material Parameters structure that contain material parameters
"""
function Nondimensionalize!(phase_mat::AbstractMaterialParamsStruct, g::GeoUnits{TYPE}) where {TYPE}

    for param in fieldnames(typeof(phase_mat))
        fld = getfield(phase_mat, param)
        if ~isnothing(fld)
            if typeof(fld[1]) <: AbstractPhaseDiagramsStruct

                # in case we employ a phase diagram
                temp = PerpleX_LaMEM_Diagram(fld[1].Name, CharDim = g)

                setfield!(phase_mat, param, (temp,))
            else
                # otherwise non-dimensionalize
                for i=1:length(fld)
                    if typeof(fld[i]) <: AbstractMaterialParam
                        Units.Nondimensionalize!(fld[i],g)
                    end
                end
            end
        end
    end
    phase_mat.Nondimensional = true
end

"""
    Dimensionalize(param, param_dim::Unitful.FreeUnits, CharUnits::GeoUnits{TYPE})

Dimensionalizes `param` into the dimensions `param_dim` using the characteristic values specified in `CharUnits`.

# Example
```julia-repl
julia> CharUnits =   GEO_units();
julia> v_ND      =   Nondimensionalize(3cm/yr, CharUnits)
0.031688087814028945
julia> v_dim     =   Dimensionalize(v_ND, cm/yr, CharUnits)
3.0 cm yr⁻¹
```

"""
function Dimensionalize(param_ND, param_dim::Unitful.FreeUnits, g::GeoUnits{TYPE}) where {TYPE}

    dim         =   Unitful.dimension(param_dim);                   # Basic SI units
    char_val    =   1.0;
    foreach((typeof(dim).parameters[1])) do y
        val = upreferred(getproperty(g, Unitful.name(y)))       # Retrieve the characteristic value from structure g
        pow = Float64(y.power)                                  # power by which it should be multiplied
        char_val *= val^pow                                     # multiply characteristic value
    end
    param = uconvert.(param_dim, param_ND*char_val)

    return param
end

"""
    Dimensionalize!(param::GeoUnit, CharUnits::GeoUnits{TYPE})

Dimensionalizes `param` again to the values that it used to have using the characteristic values specified in `CharUnits`.

# Example
```julia-repl
julia> CharUnits =   GEO_units();
julia> x = GeoUnit(3cm/yr)
julia> Nondimensionalize!(x, CharUnits)
julia> Dimensionalize!(x, CharUnits)
3.0 cm yr⁻¹
```

"""
function Dimensionalize!(param::GeoUnit, g::GeoUnits{TYPE}) where {TYPE}

    dim         =   Unitful.dimension(param.unit);                   # Basic SI units
    char_val    =   1.0;
    foreach((typeof(dim).parameters[1])) do y
        val = upreferred(getproperty(g, Unitful.name(y)))       # Retrieve the characteristic value from structure g
        pow = Float64(y.power)                                  # power by which it should be multiplied
        char_val *= val^pow                                     # multiply characteristic value
    end
    param.val = uconvert.(param.unit, param.val*char_val)

end

"""
    Dimensionalize!(MatParam::AbstractMaterialParam, CharUnits::GeoUnits{TYPE})

Dimensionalizes a material parameter structure (e.g., Density, CreepLaw)

"""
function Dimensionalize!(MatParam::AbstractMaterialParam, g::GeoUnits{TYPE}) where {TYPE}

    for param in fieldnames(typeof(MatParam))
        if typeof(getfield(MatParam, param))==GeoUnit
            z=getfield(MatParam, param)
            Dimensionalize!(z, g)
            setfield!(MatParam, param, z)
        end
    end

end

"""
    isDimensional(MatParam::AbstractMaterialParam)

`true` if MatParam is in dimensional units.
"""
function isDimensional(MatParam::AbstractMaterialParam)
    isDim = false;
    for param in fieldnames(typeof(MatParam))
        if typeof(getfield(MatParam, param))==GeoUnit
            z=getfield(MatParam, param)
            if unit(z.val)!=NoUnits
                isDim=true;
            end
        end
    end
    return isDim
end

"""
    Dimensionalize!(phase_mat::MaterialParams, g::GeoUnits{TYPE})

Dimensionalizes all fields within the Material Parameters structure that contain material parameters
"""
function Dimensionalize!(phase_mat::AbstractMaterialParamsStruct, g::GeoUnits{TYPE}) where {TYPE}

    for param in fieldnames(typeof(phase_mat))
        fld = getfield(phase_mat, param)
        if ~isnothing(fld)
            for i=1:length(fld)
                if typeof(fld[i]) <: AbstractMaterialParam
                    Units.Dimensionalize!(fld[i],g)
                end
            end
        end
    end
    phase_mat.Nondimensional = false

end

# Define a view for the GEO_Units structure
function show(io::IO, g::GeoUnits{TYPE})  where {TYPE}
    print(io, "Employing $TYPE units \n",
              "Characteristic values: \n",
              "         length:      $(g.length)\n",
              "         time:        $(round(ustrip(g.time),digits=4)) $(unit(g.time))\n",
              "         stress:      $(g.stress)\n",
              "         temperature: $(Float64(g.temperature))\n")
end

# This replaces the viewer of the Unitful package, such that the printing of units is done as floats (better)
function Unitful.superscript(i::Rational{Int64}; io=nothing)
    string(superscript(float(i)))
    if io === nothing
        iocontext_value = nothing
    else
        iocontext_value = get(io, :fancy_exponent, nothing)
    end
    if iocontext_value isa Bool
        fancy_exponent = iocontext_value
    else
        v = get(ENV, "UNITFUL_FANCY_EXPONENTS", Sys.isapple() ? "true" : "false")
        t = tryparse(Bool, lowercase(v))
        fancy_exponent = (t === nothing) ? false : t
    end
    if fancy_exponent
        return superscript(float(i))
    else
        return  "^" * string(float(i))
    end

end

Unitful.superscript(i::Float64) = map(repr(i)) do c
    c == '-' ? '\u207b' :
    c == '1' ? '\u00b9' :
    c == '2' ? '\u00b2' :
    c == '3' ? '\u00b3' :
    c == '4' ? '\u2074' :
    c == '5' ? '\u2075' :
    c == '6' ? '\u2076' :
    c == '7' ? '\u2077' :
    c == '8' ? '\u2078' :
    c == '9' ? '\u2079' :
    c == '0' ? '\u2070' :
    c == '.' ? '\u0387' :
    error("unexpected character")
end

end


export OnePlusEvo, evaluate!, populate!

"""

OnePlusEvo implements a 1+λ evolutionary algorithm. To do this, a new subtype of
AbstractEvolution is created, OnePlusEvo, and the following functions are
defined for this type:

evaluate
populate

the function generation can also be defined. Without specification, the default
generation function (which is void) is applied.
"""
mutable struct OnePlusEvo{T} <: AbstractEvolution
    config::NamedTuple
    logger::CambrianLogger
    population::Array{T}
    fitness::Function
    gen::Int
end

function OnePlusEvo{T}(cfg::NamedTuple, fitness::Function;
                      logfile=string("logs/", cfg.id, ".csv")) where T
    logger = CambrianLogger(logfile)
    population = initialize(T, cfg)
    OnePlusEvo(cfg, logger, population, fitness, 0)
end

function oneplus_populate!(e::AbstractEvolution)
    p1 = max_selection(e.population)
    e.population[1] = p1
    for i in 2:prod(e.config.population.size)
        e.population[i] = mutate(p1)
    end
end

evaluate!(e::OnePlusEvo) = fitness_evaluate!(e, e.fitness)
populate!(e::OnePlusEvo) = oneplus_populate!(e)


pluralize(io, n) = n == 1 ? print(io, n) : print(io, n, "s")

function Base.summary{T,S,N}(sig::AbstractSignal{T,S,N})
    b = IOBuffer()
    N == 1 && print(b, length(sig), "-", sig.dims[1] == symbol("") ? "element" : sig.dims[1])
    N > 1 && print(b, join(size(sig), "x"))

    print(b, " ", typeof(sig).name)
    N > 1 && sig.dims[1] != symbol("") && (print(b, " with "); pluralize(b, size(sig, 1)); print(b, " ", sig.dims[1]))
    N > 2 && sig.dims[2] != symbol("") && (print(b, " across "); pluralize(b, size(sig, 2)); print(b, " ", sig.dims[2]))
    # N > 3 && ... # TODO: how do I show other dimension names? Do I at all?
    takebuf_string(b)
end

function Base.writemime(io::IO, m::MIME"text/plain", sig::Signal)
    print(io, summary(sig))
    length(sig) == 0 && return
    println(io, ":")
    print(io, "  Each ", !isempty(sig.dims) && sig.dims[1] != symbol("") ? sig.dims[1] : "element",
              " has ", length(sig.time), " datapoints from ", sig.time[1], " to ", sig.time[end])
    isa(sig, RegularSignal) && print(io, ", at ", samplingfreq(sig))

    show_signal(io, sig)
end

# If the signal contains nested signals display them nicely with sparklines
# I'm not sure I like this since it displays the shape of the data transposed
function show_signal{T,S<:Signal}(io::IO, sig::Signal{T, S}, limit_output::Bool=true)
    length(sig) > 1 && return # TODO: display multiple signals?

    # Determine screen size
    rows, cols = limit_output ? Base.tty_size() : (typemax(Int), typemax(Int))
    rows = min(rows-5, length(sig.time))

    # Gather the times
    ts = Array(ByteString, rows)
    for i=1:min(length(ts), rows)
        ts[i] = string(sig.time[i], ": ")
    end
    twidth = maximum([strwidth(t) for t in ts]) + 2

    # And display the signal for each time
    s = sig[1]
    for i=1:length(ts)
        println(io)
        print(io, lpad(ts[i], twidth))
        if length(s[i]) > cols - twidth
            spark(io, s[i][1:(cols - twidth - 2)])
            print(io, " …")
        else
            spark(io, s[i])
        end
    end

    # And display a continuation mark if we have more data
    if rows < length(sig.time)
        println(io)
        print(io, "  ⋮", lpad(": ", twidth-3), "⋮")
    end
end

function show_signal(io::IO, sig::Signal, limit_output::Bool=true)
    # TODO: How should generic signals of many channels display?
    println(io)
    Base.showarray(io, sig.data; header=false)
end


using DiffEqJump, DiffEqBase, SafeTestsets

@time begin
  @time @safetestset "Spatial utilities" begin include("utils_test.jl") end
  @time @safetestset "Spatial A + B <--> C" begin include("ABC.jl") end
  @time @safetestset "Pure diffusion" begin include("diffusion.jl") end
end


# This file was generated, do not modify it. # hide
fit!(clf)
ŷ = predict(clf, X2)
ŷ[1:3]


"""
  `code_similarity(onestr, twostr[, code=fuzzy_soundex])`

  Produce a similarity estimate based on string comparison of the result of
  coding the two input strings.

  `code` can be any of: `soundex`, `phonex`, `phonix`, `[fuzzy_soundex]`, `nysiis`,
  `metaphone` or `match_rating_encode`. Results from using `double_metaphone` may be
  misleading due to the double encoding of that system.

  This comparison function was originally used by the authors of Fuzzy Soundex. A similar
  but non-normalised comparison scheme exists for Soundex, see `editex`, and a specific
  comparison algorithm exists for `match_rating_encode`, see `match_rating`.
"""
function code_similarity(onestr::String, twostr::String, code=fuzzy_soundex)
  #encode strings
  ionestr = code(onestr)
  itwostr = code(twostr)

  return code_similarity_internal(ionestr, itwostr, code)
end

" Internal code similarity function (not for raw strings). "
function code_similarity_internal(ionestr::String , itwostr::String, code=fuzzy_soundex)

  l1 = length(ionestr)
  l2 = length(itwostr)

  common = 0
  for i in 1:min(l1,l2)
    if ionestr[i] == itwostr[i]
      common += 1
    end
  end

  ð = (2*common)/(l1+l2)
  return ð
end

"""
  `code_match(str, array[, code=fuzzy_soundex, permissive=0.0])`

  Returns those items in `array` which are phonetic matches to `str` according
  to the algorithm `code`.

  Matches are normally on exact codes. The `permissive` option allows you to set
  a lower threshold for comparison for most `code` options, using `code_similarity()`.
  The value should be in the range 0:1.  Permissive coding is not possible for
  `double_metaphone`.

  `match_rating_encode` will be treated like other codes rather than how it is used
  in the overall approach. Those wishing to match items according to that system can
  call

      filter(item -> meets_match_rating(str, item), array)

  instead.
"""
function code_match(str::String, array::Array{T,1}, code=fuzzy_soundex::Function, permissive=0.0::Float64) where {T<:String}

  #catch weird input
  if permissive > 1 || permissive < 0
    error("Permissive value must be in range 0:1")
  end

  #If partial matching
  if permissive > 0 && permissive < 1

    #catch attempt at double_metaphone (makes code messy)
    if code == double_metaphone
      error("Cannot currently code match with double_metaphone")
    end

    #returns any better matches than the set level
    return filter(array) do item
      sim_val = code_similarity(item, str, code)
      return sim_val > permissive
    end
  end

  #encode key
  encoded_key = code(str)

  #filter straight matches
  return filter(array) do item

    encoded_val = code(item)

    #handle double_metaphone
    if code == double_metaphone && isa(encoded_val, Array{String,1})
      return (encoded_key in encoded_val)
    else
      return encoded_key == encoded_val
    end
  end

end

"""
  `code_cluster(array[, code=phonix, lower_threshold=0.7, higher_threshold=0.9, stochastic=true])`

  Uses a phonetic similarity measure to cluster an array of strings.

  This function uses a canopy clustering method to group input strings based
  on their phonetic similarity, as judged by `code_similarity`. The phonetic
  coding method to be used can be specified, as can the upper and lower thresholds
  for similarity and whether or not to use randomly sampled centroids for the clusters.

  - Setting `higher_threshold=1` and `lower_threshold=1` will result in only exact
    code matches, which may be particularly sensible for coding schemes with variable
    length codes.

  - Setting `lower_threshold` to the same value as `higher_threshold` will result
    in cleanly separated clusters -- that is, each element of the original array
    will be assigned to one and only one cluster.

  - Setting `stochastic=false` will cause centroid to be selected from left to right
    rather than randomly, meaning that the output clusters will be deterministic. Behaviour
    without this setting is to select a centroid randomly, which can result in some clusters
    appearing or disappearing between runs on the same data where either threshold < 1.
"""
function code_cluster(array::Array{T, 1}, code=phonix::Function, lower_threshold=0.7, higher_threshold=0.9, stochastic=true) where {T<:String}

  #Sanity check
  if lower_threshold > higher_threshold
    error("Lower threshold cannot be greater than higher threshold.")
  end

  clusters = Vector[]

  #local copy of input array to shrink
  coderef = copy(array)

  #encoded version of input array
  inarray = map(code, array)

  while length(inarray) > 0

    if stochastic
      centroid = rand(inarray)
    else
      centroid = inarray[1]
    end

    #similarity of all items to selected centroid
    similarities = map(x -> code_similarity_internal(x, centroid), inarray)

    #select words which are greater than the lower threshold to add to this cluster
    incluster = filter(x -> similarities[findall(y -> y == x, coderef)[1]] >= lower_threshold, coderef)
    push!(clusters, incluster)

    #remove codes and words which are higher than the higher threshold from consideration
    inarray = filter(x -> similarities[findall(y -> y == x, inarray)[1]] < higher_threshold, inarray)
    coderef = filter(x -> similarities[findall(y -> y == x, coderef)[1]] < higher_threshold, coderef)
  end
  return clusters
end

export code_similarity, code_match, code_cluster


@testset "1025.divisor-game.jl" begin
    @test divisor_game(2) == true
    @test divisor_game(3) == false
end


# ------------------------------------------------------------------- #
#                          FlicsFlow.jl                               #
#                       -----------------                             #
#     A General Three-dimensional Unstructured Finite Volume Solver   #
#                                                                     #
#                    Copyright(c) 2021 NUMFLICS                       #
# ------------------------------------------------------------------- #
using Pkg
Pkg.activate("C:\\Users\\arifa\\.julia\\dev\\FlicsFlow")

using FlicsFlow

# Get the current case directory
casedir = @__DIR__

# Read the OpenFOAM Mesh
mesh = readopenfoammesh(casedir)


macro stridedloops(N, itersym, dimsym, args...)
    _stridedloops(N, itersym, dimsym, args...)
end
function _stridedloops(N::Int, itersym::Symbol, dimsym::Symbol, args...)
    mod(length(args),3)==1 || error("Wrong number of arguments")
    body = args[end]
    ex = Expr(:escape, body)
    for i=1:3:length(args)-1
        ex=Cartesian.sreplace!(ex,args[i],Cartesian.inlineanonymous(args[i], 1))
    end
    for dim = 1:N
        itervar = Cartesian.inlineanonymous(itersym, dim)
        dimvar = Cartesian.inlineanonymous(dimsym, dim)
        preargs = Any[]
        postargs = Any[]
        for i=1:3:length(args)-1
            indnew = Cartesian.inlineanonymous(args[i], dim)
            start = (dim < N ? Cartesian.inlineanonymous(args[i], dim+1) : args[i+1])
            step = Cartesian.inlineanonymous(args[i+2], dim)
            push!(preargs,:($(esc(indnew)) = $(esc(start))))
            push!(postargs,:($(esc(indnew)) += $(esc(step))))
        end
        preexpr=Expr(:block,preargs...)
        postexpr=Expr(:block,postargs...)
        ex = quote
            $preexpr
            for $(esc(itervar)) = 1:$(esc(dimvar))
                $ex
                $postexpr
            end
        end
    end
    if N == 0 # even a zero-dimensional array has one element and thus requires a single run
        preargs = Any[]
        for i=1:3:length(args)-1
            indnew = Cartesian.inlineanonymous(args[i], 1)
            start = args[i+1]
            push!(preargs,:($(esc(indnew)) = $(esc(start))))
        end
        preexpr=Expr(:block,preargs...)
        ex = quote
            $preexpr
            $ex
        end
    end
    ex
end

macro gentracekernel(N1,N2,order,alpha,A,beta,C,startA,startC,odims,cdims,ostridesA,cstridesA,ostridesC)
    _gentracekernel(N1,N2,order,alpha,A,beta,C,startA,startC,odims,cdims,ostridesA,cstridesA,ostridesC)
end
function _gentracekernel(N1::Int,N2::Int,order::Symbol,alpha::Symbol,A::Symbol,beta::Symbol,C::Symbol,
    startA::Symbol,startC::Symbol,odims::Symbol,cdims::Symbol,ostridesA::Symbol,cstridesA::Symbol,ostridesC::Symbol)
    ex=quote
        local indA1, indA2, indC
        local i,j
        if $(esc(order))==0
            local gamma
            gamma=$(esc(beta))
            @stridedloops($N1, i, $(esc(cdims)), indA1, $(esc(startA)), $(esc(cstridesA)), begin
                @stridedloops($N2, j, $(esc(odims)), indA2, indA1, $(esc(ostridesA)), indC, $(esc(startC)), $(esc(ostridesC)), begin
                    @inbounds $(esc(C))[indC]=gamma*$(esc(C))[indC]+$(esc(alpha))*$(esc(A))[indA2]
                end)
                gamma=one($(esc(beta)))
            end)
        else
            @stridedloops($N2, j, $(esc(odims)), indA1, $(esc(startA)), $(esc(ostridesA)), indC, $(esc(startC)), $(esc(ostridesC)), begin
                local localC
                @inbounds localC=$(esc(beta))*$(esc(C))[indC]
                @stridedloops($N1, i, $(esc(cdims)), indA2, indA1, $(esc(cstridesA)), @inbounds localC+=$(esc(alpha))*$(esc(A))[indA2])
                @inbounds $(esc(C))[indC]=localC
            end)
        end
    end
    ex
end

macro gencontractkernel(N1,N2,N3,order,alpha,A,conjA,B,conjB,beta,C,startA,startB,startC,odimsA,odimsB,cdims,ostridesA,cstridesA,ostridesB,cstridesB,ostridesCA,ostridesCB)
    _gencontractkernel(N1,N2,N3,order,alpha,A,conjA,B,conjB,beta,C,startA,startB,startC,odimsA,odimsB,cdims,ostridesA,cstridesA,ostridesB,cstridesB,ostridesCA,ostridesCB)
end
function _gencontractkernel(N1::Int,N2::Int,N3::Int,order::Symbol,
    alpha::Symbol,A::Symbol,conjA::Symbol,B::Symbol,conjB::Symbol,beta::Symbol,C::Symbol,
    startA::Symbol,startB::Symbol,startC::Symbol,odimsA::Symbol,odimsB::Symbol,cdims::Symbol,
    ostridesA::Symbol,cstridesA::Symbol,ostridesB::Symbol,cstridesB::Symbol,ostridesCA::Symbol,ostridesCB::Symbol)
    ex=quote
        local indA1, indA2, indB1, indB2, indC1, indC2, gamma
        # we still have to implement other orders
        if $(esc(order))==3 # k,j,i
            @stridedloops($N3, k, $(esc(cdims)), indA1, $(esc(startA)), $(esc(cstridesA)), indB1, $(esc(startB)), $(esc(cstridesB)), begin
                gamma=$(esc(beta))
                @stridedloops($N2, j, $(esc(odimsB)), indB2, indB1, $(esc(ostridesB)), indC1, $(esc(startC)), $(esc(ostridesCB)), begin
                    @inbounds localB=($(esc(conjB))=='C' ? conj($(esc(B))[indB2]) : $(esc(B))[indB2])
                    @stridedloops($N1, i, $(esc(odimsA)), indA2, indA1, $(esc(ostridesA)), indC2, indC1, $(esc(ostridesCA)), begin
                        @inbounds localA=($(esc(conjA))=='C' ? conj($(esc(A))[indA2]) : $(esc(A))[indA2])
                        @inbounds localC=gamma*$(esc(C))[indC2]
                        @inbounds $(esc(C))[indC2]=localC+$(esc(alpha))*localA*localB
                    end)
                end)
                gamma=one($(esc(beta)))
            end)
        else # j,i,k
            @stridedloops($N2, j, $(esc(odimsB)), indB1, $(esc(startB)), $(esc(ostridesB)), indC1, $(esc(startC)), $(esc(ostridesCB)), begin
                @stridedloops($N1, i, $(esc(odimsA)), indA1, $(esc(startA)), $(esc(ostridesA)), indC2, indC1, $(esc(ostridesCA)), begin
                    @inbounds localC=$(esc(beta))*$(esc(C))[indC2]
                    @stridedloops($N3, k, $(esc(cdims)), indA2, indA1, $(esc(cstridesA)), indB2, indB1, $(esc(cstridesB)), begin
                        @inbounds localA=($(esc(conjA))=='C' ? conj($(esc(A))[indA2]) : $(esc(A))[indA2])
                        @inbounds localB=($(esc(conjB))=='C' ? conj($(esc(B))[indB2]) : $(esc(B))[indB2])
                        localC+=$(esc(alpha))*localA*localB
                    end)
                    @inbounds $(esc(C))[indC2]=localC
                end)
            end)
        end
    end
    ex
end


#===============================================================================

    The main object defined in this file is: SamplingPair.
    It is a pairing of two `SamplingUnit`s and can be used for
    smoothing or inference problems.

===============================================================================#

"""
    struct SamplingPair{TGP,TGPb,TW,TWn,TX}
        u::SamplingUnit{TGP,TGPb,TW,TWn,TX}
        u°::SamplingUnit{TGP,TGPb,TW,TWn,TX}
    end

A pairing of two `SamplingUnit`s, can be used for smoothing or inference
problems.

    SamplingPair(
        aux_laws, recording, tts, args=tuple();
        aux_laws_blocking=aux_laws, artificial_noise=1e-11,
        solver_choice_blocking=args
    )

Base constructor.

# Arguments
---
- `aux_laws`:
- `recording`:
- `tts`:
- `args`:
- `aux_laws_blocking`:
- `artificial_noise`:
- `solver_choice_blocking`:
"""
struct SamplingPair{TGP,TGPb,TW,TWn,TX}
    u::SamplingUnit{TGP,TGPb,TW,TWn,TX}
    u°::SamplingUnit{TGP,TGPb,TW,TWn,TX}

    function SamplingPair(
            aux_laws, recording, tts, args=tuple();
            aux_laws_blocking=aux_laws, artificial_noise=1e-11,
            solver_choice_blocking=args
        )
        u = SamplingUnit(
            aux_laws, recording, tts, args;
            aux_laws_blocking = aux_laws_blocking,
            artificial_noise = artificial_noise,
            solver_choice_blocking = solver_choice_blocking
        )
        u° = deepcopy(u)
        TGP,TGPb,TW,TWn,TX = all_eltypes(u)
        new{TGP,TGPb,TW,TWn,TX}(u, u°)
    end
end


using NQCDynamics
using Test
using OrdinaryDiffEq

@testset "select_reduction" begin
    @test Ensembles.select_reduction(:mean) isa Ensembles.MeanReduction
    @test Ensembles.select_reduction(:sum) isa Ensembles.SumReduction
    @test Ensembles.select_reduction(:append) isa Function
end

@testset "get_u_init" begin

    @testset "AbstractOutput, reduction=$reduction" for reduction in (:mean, :sum)
        reduction = Ensembles.select_reduction(reduction)
        tspan = (0.0, 1.0)
        u0 = 1.0
        saveat = 1
        output = Ensembles.OutputFinal()
        u_init = Ensembles.get_u_init(reduction, saveat, Dict(), tspan, u0, output)
        @test u_init ≈ zero(u0)
    end

    @testset "reduction=:append" begin
        reduction = Ensembles.select_reduction(:append)
        output = Ensembles.OutputFinal()
        u0 = 1
        saveat = 1
        u_init = Ensembles.get_u_init(reduction, saveat, Dict(), (0, 1), u0, output)
        @test u_init == []
    end

    @testset "TimeCorrelationFunction, reduction=$reduction" for reduction in (:mean, :sum)
        reduction = Ensembles.select_reduction(reduction)
        sim = Simulation{FSSH}(Atoms(1), DoubleWell())
        saveat = 0.1
        tspan = (0.0, 1.0)
        u0 = DynamicsVariables(sim, 1, 0, SingleState(1, Adiabatic()))
        output = TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Diabatic())
        u_init = Ensembles.get_u_init(reduction, saveat, Dict(), tspan, u0, output)
        @test u_init ≈ [zeros(2, 2) for _ ∈ tspan[1]:0.1:tspan[2]]
    end

end

@testset "$reduction" for reduction in (Ensembles.MeanReduction(0), Ensembles.SumReduction())
    prob = ODEProblem((u,p,t) -> 1.01u, 0.5, (0.0,1.0))
    prob_func(prob, i, repeat) = remake(prob,u0=rand())
    output = Ensembles.OutputFinal()
    u_init = Ensembles.output_template(output, rand())
    ensemble = EnsembleProblem(prob, prob_func=prob_func, output_func=output, reduction=reduction, u_init=u_init)
    sol = solve(ensemble, Tsit5(), trajectories=1e3, batch_size=20)
    if reduction isa Ensembles.SumReduction
        @test sol.u / 1e3 ≈ 0.5 * exp(1.01) rtol=1e-1
    else
        @test sol.u ≈ 0.5 * exp(1.01) rtol=1e-1
    end
end


export DARTStrain1st!, DARTStrain2nd!, DARTSevaltrain1st!, all_ws_sansbn, all_αs

using Flux
using Flux: onehotbatch, onecold
using Juno
using Base.Iterators
using StatsBase: mean
using Zygote
using Zygote: @nograd
using LinearAlgebra
using CUDA

function accuracy(m, x, y)
    mx = m(x)
    showmx = mx |>cpu
    showy = y|>cpu
    for i in 1:size(showmx,2)
        @show collect(zip(softmax(showmx[:,i]), showy[:,i]))
    end
    mean(onecold(mx, 1:10)|>cpu .== onecold(y|>cpu, 1:10))
end

function DARTStrain1st!(loss, model, train, val, opt_α, opt_w, losses=[0f0,0f0], epoch = 1; cbepoch = () -> (), cbbatch = () -> ())
    local train_loss
    local val_loss
    w = all_ws_sansbn(model)
    α = all_αs(model)
    opt_w.os[2].t = epoch - 1
    for (train_batch, val_batch) in zip(TrainCuIterator(train), TrainCuIterator(val))
        if epoch == 29 || epoch == 30
            @show accuracy(model, train_batch...)
            @show accuracy(model, val_batch...)
        end
        display("grad weights")
        @time begin
            gsw = gradient(w) do
                train_loss = loss(model, train_batch...)
                return train_loss
            end
            losses[1] = train_loss
        end
        foreach(CUDA.unsafe_free!, train_batch)
        Flux.Optimise.update!(opt_w, w, gsw)
        CUDA.reclaim()
        display("grad alpha")
        @time begin
            gsα = gradient(α) do
                val_loss = loss(model, val_batch...)
                return val_loss
            end
            losses[2] = val_loss
            if epoch == 29 || epoch == 30
                @show losses
            end
        end
        foreach(CUDA.unsafe_free!, val_batch)
        Flux.Optimise.update!(opt_α, α, gsα)
        CUDA.reclaim()
        cbbatch()
    end
    cbepoch()
end

all_ws(model::DARTSEvalModel) = Flux.params([model.stem, model.cells..., model.global_pooling, model.classifier])

function DARTSevaltrain1st!(loss, model, train, opt_w, losses=[0f0,0f0], epoch = 1; cbepoch = () -> (), cbbatch = () -> ())
    w = all_ws(model)
    local train_loss
    opt_w.os[2].t = epoch - 1
    for train_batch in EvalCuIterator(train)
        gsw = gradient(w) do
            train_loss = loss(model, train_batch...)
            return train_loss
        end
        losses[1] = train_loss
        CUDA.reclaim()
        GC.gc()
        Flux.Optimise.update!(opt_w, w, gsw)
        cbbatch()
    end
    cbepoch()
end

function update_all!(opt, ps::Params, gs)
    for (p, g) in zip(ps, gs)
        g == nothing && continue
        Flux.Optimise.update!(opt, p, g)
    end
end

function DARTStrain2nd!(loss, model, train, val, opt; cb = () -> ())
    function grad_loss(model, ps, batch, verbose = false)
        gs = gradient(ps) do
            loss(model, batch...)
        end
    end

    ξ = opt.eta

    w = all_ws(model)
    α = all_αs(model)

    for (train_batch, val_batch) in zip(CuIterator(train), CuIterator(val))
        gsα = grad_loss(model, α, val_batch)
        model_prime = deepcopy(model)
        model_prime_α = deepcopy(model_prime)
        w_prime = all_ws(model_prime_α)

        gsw_prime = grad_loss(model_prime, w_prime, val_batch)

        Flux.Optimise.update!(Descent(ξ), w_prime, gsw_prime)

        model_prime_minus = deepcopy(model_prime)

        w_prime_minus = all_ws(model_prime_minus)

        gsw_prime_minus = grad_loss(model_prime_minus, w_prime_minus, val_batch)

        model_minus = deepcopy(model)
        model_plus = deepcopy(model)

        w_minus = all_ws(model_minus)
        w_plus = all_ws(model_plus)

        epsilon = 0.01 ./ norm([gsw_prime_minus[w_] for w_ in w_prime_minus.order])
        update_all!(
            Descent(epsilon),
            w_minus,
            [gsw_prime_minus[w_] for w_ in w_prime_minus.order],
        )
        update_all!(
            Descent(-1 * epsilon),
            w_plus,
            [gsw_prime_minus[w_] for w_ in w_prime_minus.order],
        )

        gsα_prime = grad_loss(model_prime_α, all_αs(model_prime_α), val_batch)

        gsα_plus = grad_loss(model_plus, all_αs(model_plus), train_batch)
        gsα_minus = grad_loss(model_minus, all_αs(model_minus), train_batch)

        update_all!(opt, α, [gsα_prime[a] for a in all_αs(model_prime_α).order])
        update_all!(
            Descent(-1 * ξ^2 / (2 * epsilon)),
            α,
            [gsα_plus[a] for a in all_αs(model_plus).order],
        )
        update_all!(
            Descent(ξ^2 / (2 * epsilon)),
            α,
            [gsα_minus[a] for a in all_αs(model_minus).order],
        )

        gsw = grad_loss(model, w, train_batch)
        Flux.Optimise.update!(opt, w, gsw)

        cb()
    end
end


export NewtonRaphson

abstract type AbstractNonLinearSolver end

"""
    NewtonRaphson <: AbstractNonLinearSolver

Newton-Raphson algorithm. Used to solve the non-linear equation
``g(x, u) = 0``, at a fixed control ``u``.

### Attributes
- `maxiter::Int` (default 20): maximum number of iterations
- `tol::Float64` (default `1e-8`): tolerance of the algorithm
- `verbose::Int` (default `NONE`): verbosity level

"""
struct NewtonRaphson <: AbstractNonLinearSolver
    maxiter::Int
    tol::Float64
    verbose::Int
end
NewtonRaphson(; maxiter=20, tol=1e-8, verbose=VERBOSE_LEVEL_NONE) = NewtonRaphson(maxiter, tol, verbose)

"""
    ConvergenceStatus

Convergence status returned by a non-linear algorithm.

### Attributes
- `has_converged::Bool`: states whether the algorithm has converged.
- `n_iterations::Int`: total number of iterations of the non-linear algorithm.
- `norm_residuals::Float64`: final residual.
- `n_linear_solves::Int`: number of linear systems ``Ax = b`` resolved during the run.

"""
struct ConvergenceStatus
    has_converged::Bool
    n_iterations::Int
    norm_residuals::Float64
    n_linear_solves::Int
end

# Sparse utilities
mutable struct Spmat{VTI<:AbstractVector, VTF<:AbstractVector}
    colptr::VTI
    rowval::VTI
    nzval::VTF

    # create 2 Spmats from complex matrix
    function Spmat{VTI, VTF}(mat::SparseMatrixCSC{Complex{Float64}, Int}) where {VTI, VTF}
        matreal = new(VTI(mat.colptr), VTI(mat.rowval), VTF(real.(mat.nzval)))
        matimag = new(VTI(mat.colptr), VTI(mat.rowval), VTF(imag.(mat.nzval)))
        return matreal, matimag
    end
    # copy constructor
    function Spmat{VTI, VTF}(mat) where {VTI, VTF}
        return new(VTI(mat.colptr), VTI(mat.rowval), VTF(mat.nzval))
    end
end

function _copy_csc!(J_dest, J_src, shift)
    @inbounds for i in 1:size(J_src, 2)
        for j in J_src.colptr[i]:J_src.colptr[i+1]-1
            row = J_src.rowval[j]
            @inbounds J_dest[row+shift, i] = J_src.nzval[j]
        end
    end
end

function _transfer_sparse!(J_dest::SparseMatrixCSC, J_src::SparseMatrixCSC, shift)
    _copy_csc!(J_dest, J_src, shift)
end

function _transfer_sparse!(J_dest::CUSPARSE.CuSparseMatrixCSR, J_src::CUSPARSE.CuSparseMatrixCSR, shift)
    nnz_start = J_dest.rowPtr[shift+1]
    nnz_ = nnz(J_src)
    Jnnz = @view J_dest.nzVal[nnz_start:nnz_start+nnz_-1]
    Jnnz .= J_src.nzVal
end

# projection operator
function project!(w::VT, u::VT, u♭::VT, u♯::VT) where VT<:AbstractArray
    w .= max.(min.(u, u♯), u♭)
end

# Utils function to solve transposed linear system  A' x = y
# Source code taken from:
# https://github.com/JuliaGPU/CUDA.jl/blob/master/lib/cusolver/wrappers.jl#L78L111
function csclsvqr!(A::CUSPARSE.CuSparseMatrixCSC{Float64},
                    b::CUDA.CuVector{Float64},
                    x::CUDA.CuVector{Float64},
                    tol::Float64,
                    reorder::Cint,
                    inda::Char)
    n = size(A,1)
    desca = CUSPARSE.CuMatrixDescriptor(
        CUSPARSE.CUSPARSE_MATRIX_TYPE_GENERAL,
        CUSPARSE.CUSPARSE_FILL_MODE_LOWER,
        CUSPARSE.CUSPARSE_DIAG_TYPE_NON_UNIT, inda)
    singularity = Ref{Cint}(1)
    CUSOLVER.cusolverSpDcsrlsvqr(CUSOLVER.sparse_handle(), n, A.nnz, desca, A.nzVal, A.colPtr, A.rowVal, b, tol, reorder, x, singularity)

    if singularity[] != -1
        throw(SingularException(singularity[]))
    end

    x
end


#hello

#okay, this is a very quick test of a basic blockchain in julia
# for educational purpoess probably

#first we set up our types

abstract type CryptoPrimitive end

abstract type Hash <: CryptoPrimitive end
abstract type MAC <: CryptoPrimitive end
abstract type Signature <: CryptoPrimitive end

abstract type AbstractBlockChain end
abstract type AbstractBlock end

abstract type AbstractGenesisBlock <: AbstractBlock end

abstract type AbstractBlockData end
abstract type AbstractMessage <:AbstractBlockData end

primitive type 256Hash <: Hash 256 end
primitive type 512Hash <: Hash 512 end

primitive type 256MAC <: MAC 256 end
primitive type 512MAC <: MAC 512 end

primitive type 256Signature <: Signature end
primitive type 512Signature <: Signature end

#now for our more complex types

const genesis_message = "This is how the world begins: not with a whimper, but with a bang!"

type Blockchain <: AbstractBlockChain
	blocks:: Array{AbstractBlock}
	timeCreated:: Date
	Description:: AbstractString
	id:: Hash
	sig:: Signature
	MAC:: MAC
end

type Block <: AbstractBlock
	index:: Integer
	hash:: Hash
	previousHash:: Hash
	timestamp:: Date
	data:: Array{Message}
	difficulty:: Integer
	nonce:: Integer

	function Block(index, hash, previousHash, timestamp, data, difficulty, nonce)
		return new(index, hash, previousHash, timestamp, data, difficulty,nonce)
	end

end

type GenesisBlock <: AbstractGenesisBlock
	index:: Integer
	hash:: Hash
	previousHash:: Hash
	timestamp:: Date
	data::Array{Message}
	difficulty:: Integer
	nonce:: Integer

	function GenesisBlock()
		const genesis_message = getGenesisMessage()
		index = 0
		data = Message(0, genesis_message)
		difficulty = 0
		nonce = 0
		previousHash = Hash(0)
		timestamp = 1 # pass
		hash=  caluclate_hash(index, data[0].data, difficulty, nonce, previousHash, timestamp)
		return new(index, hash, previousHash, timestamp, data, difficulty, nonce)
	end
end

type Message <: AbstractMessage
	id:: AbstractString
	data:: AbstractString
	MAC:: MAC

	function Message(id:: Integer, data::AbstactString)
		mac = calculate_MAC(id, data)
		return new(id, data, mac)
	end
	function Message(id:: Integer, data::AbstractString, encrypt_func:: function) {
		encrypt_data = encrypt_func(data)
		mac = calculate_MAC(id, encrypt_data)
		return new(id, encrypt_data, mac)
	end

	function Message(data)
		id = calculate_id()
		mac = calculate_MAC(id, data)
		return new(id, data, mac)
	end
	function Message(data, encrypt_func:: function)
		id = calculate_id()
		encrypt_data = encrypt_func(data)
		mac = calculate_MAC(id, encrypt_data)
		return new(id, encrypt_data, mac)
	end

end

function calculate_hash(inputs)
	# so originally we're going to d something here
	# but I don't know how
	#either use a crypto implementation or roll our own
	# so now we're just going to return
	return inputs
	end

function calculate_MAC(id, data)
	#no idea how to implement this either
	return MAC()
end

function getGenesisMessage()
	return genesis_message
end
# okay, now we do operations on the block chain and functoins thereby

function getLatestBlock(blockchain:: AbstractBlockChain)
	return blockchain.blocks[-1]
end

function getBlocks(blockchain:: AbstractBlockChain)
	return blockchain.blocks
end

function generateNextBlock(blockchain:: AbstractBlockChain, blockData:: Array{Message})
	prevBlock = getLatestBlock(blockchain)
	nextIndex = prevBlock.index +1
	timestamp = getDateTime()
	nextHash = calculate_hash((nextIndex, prevBlock.hash,timestamp, blockData))
	difficulty = getDifficulty()
	nonce = getNonce()
	newBlock = Block(index, hash, previousHash, timestamp, blockData, difficulty,nonce)
	return newBlock
end

function getDateTime()
	Error("not implemented yet")
end

function getDifficulty()
	Error("not implemented yet")
end

function getNonce()
	Error("not implemented yet")
end

function isValidBlock(block:: AbstractBlock, prevBlock::AbstractBlock)
	if prevBlock.index +1 != block.index
		Error('New block index does not match')
		return false
	end
	if block.previousHash != prevBlock.hash
		Error('Block previous hash does not match')
		return false
	end
	if calculateHash(block) != block.hash
		Error('Block hash invalid')
		return false
	end
	if !isValidTimestamp(block.timestamp, prevBlock.timestamp)
		Error('Block timestamp is invalid')
		return false
	end
	return true
end

function isValidGenesis(proposedGenesis:: AbstractGenesisBlock)
	if proposedGenesis.index != 0
		Error('Genesis does not start at 0')
		return false
	end
	if proposedGenesis.hash != calculate_hash(getGenesisMessage())
		Error('Genesis block hash invalid')
		return false
	end
	if proposedGenesis.previousHash != Hash(0)
		Error('Genesis block previous hash invalid')
		return false
	end
	#we might need to validate timestamp, but not sure how
	return true
end

function validateBlockChainBlocks(blocks:: Array{AbstractBlock})
	len = lengh(blocks)
	if !isValidGenesis(blocks[0])
		return false
	end
	for i in 1:len
		if !isValidBlock(blocks[i], blocks[i-1])
			return false
		end
	end
	return true
end

function validateBlockChainProperties(blockchain:: AbstractBlockChain)
	#need to implement properly
	return true
end

function validateBlockChain(blockchain:: AbstractBlockChain)
	return validateBlockChainProperties(blockchain) && validateBlockChainBlocks(blockchain.blocks)
end

function addBlockToChain(blockchain:: AbstractBlockChain, block:: AbstractBlock)
	#do we assume chain is valid here - we might as well
	if !isValidBlock(block, getLatestBlock(blockchain))
		Error('Tried to add invalid block')
		return false
	end
	push!(blockchain.blocks, block)
	# we may need to recalibrate the MAC or something of the chain
	# if so do that here!
	return true
end

function calculateBlockchainSignature(blockchain<: AbstractBlockChain)
	#we're not sure how to implement that AbstractMessage
	# so instead we're just filling this with a placeholder
	return blockchain.Signature
end

# get the cumulative difficulty of a blockchain to see if it works!
function getCumulativeDifficulty = function(blockchain:: AbstractBlockChain)
	blocks = blockchain.blocks
	if length(blocks) <=0
		Error('Blockchain has no length!')
	end
	total = 0
	for block in blocks
		hash = block.hash
		diff = block.difficulty
		if hashMatchesDifficulty(hash, diff)
			total += 2**diff
		end
	else
		Error('not matching hash/difficulty found')
		return -1
	end
	return total
end

function getDifficulty(blockchain:: AbstractBlockChain, difficulty_adjustment_interval, block_generation_interval)
	latestBlock = getLatestBlock(blockchain)
	if latestBlock.index % difficulty_adjustment_interval ==0 && latestBlock.index>=0
		return getAdjustedDifficulty(latestBlock, blockchain,difficulty_adjustment_interval,block_generation_interval)
	end
	return latestBlock.difficulty
end

function getAdjustedDifficulty(latestBlock:: AbstractBlock, blockchain::AbstractBlockChain, difficulty_adjustment_interval, block_generation_interval)
	const blockTimeAdjustmentFactor = 2
	blocks = blockchain.blocks
	prevAdjustmentBlock = blocks[length(blocks)-difficulty_adjustment_interval]
	timeExpected = block_generation_interval* difficulty_adjustment_interval
	if isValidTimestamp(latestBlock) && isValidTimestamp(prevAdjustmentBlock)
		timeTaken = latestBlock.timestamp - prevAdjustmentBlock.timestamp
		if timeTaken < timeExpected /blockTimeAdjustmentFactor
			return prevAdjustmentBlock.difficulty +1
		end
		if timeTaken > timeExpected * blockTimeAdjustmentFactor
			return prevAdjustmentBlock.difficulty -1
		end
		return prevAdjustmentBlock.difficulty
	Error('One of the block timestamps is not valid')
end


⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)
⊗(a::AbstractArray, b::AbstractArray) = a * b'


# Auxiliary functions for location processing
function get_ENU(node::Int, model::Simulation)
	"Get ENU coordinates of a specific node in the Simulation"
    model.m.nodes[m.n[node]]
end

function get_coordinates(agent::Commuter, model::Simulation)
	"Get XY coordinates of a specific node in the Simulation"
    pos = get_ENU(agent.pos, model)
    (getX(pos),
     getY(pos))
end

# Auxiliary functions for colors
agentcolor(agent) = agent.infected ? :red : :black
agentsize(agent) = agent.infected ? 6 : 5

# Plot simulation state
function plot_simstate(model::Simulation, background=Plots.current())
	"""
	Description:
		Plots simulation state on specific moment in time (after run).
	Example:
		plotmap(s.m, width=500, height=500)
		plot_simstate(s)
	"""
	# Plots black/red(if infected) agents routes
    for agent in model.agents
        path = [agent.pos]
        while path[end] in keys(agent.path)
            push!(path, agent.path[path[end]])
        end
        length(path) < 2 && continue
        a_x, a_y = get_coordinates(agent, model)
        enu_coords = get_ENU.(path[2:end], Ref(model))

        a_x2, a_y2 = (getX(enu_coords[1]), getY(enu_coords[1]))

        !(a_x == a_x2 && a_y == a_y2) &&  plot!(
                background, [a_x, a_x2], [a_y, a_y2],
                color=agentcolor(agent), arrow =arrow(:open))

        length(enu_coords) > 2 &&  plot!(
                background, getX.(enu_coords),  getY.(enu_coords),
                color=agentcolor(agent), arrow =arrow(:closed))
    end

	# Plots black wagons routes
    for wagon in model.wagons
        path = [wagon.pos]
        while path[end] in keys(wagon.path)
            push!(path, wagon.path[path[end]])
        end
        length(path) < 2 && continue
        a_x, a_y = get_coordinates(wagon, model)
        enu_coords = get_ENU.(path[2:end], Ref(model))

        a_x2, a_y2 = (getX(enu_coords[1]), getY(enu_coords[1]))

        !(a_x == a_x2 && a_y == a_y2) &&  plot!(
                background, [a_x, a_x2], [a_y, a_y2],
                color=:black, arrow =arrow(:open), linewidth = 3)

        length(enu_coords) > 2 &&  plot!(
                background, getX.(enu_coords),  getY.(enu_coords),
                color=:black, arrow =arrow(:closed), linewidth = 3)
    end

	# Define simulation objects appearance
    colors = vcat([agentcolor(agent) for agent in s.agents], [:yellow for wagon in s.wagons])
    sizes = vcat([agentsize(agent) for agent in model.agents], [10 for wagon in model.wagons])
    markers = vcat([:circle for i in 1:length(model.agents)], [:square for i in 1:length(model.wagons)])
    pos_agents = [get_coordinates(agent, model) for agent in model.agents]
    pos_wagons = [get_coordinates(wagon, model) for wagon in model.wagons]
    pos = vcat(pos_agents, pos_wagons)

	# Plot scatter
    scatter!(background,
        pos;
        markercolor = colors,
        markersize = sizes,
        markershapes = markers,
        label = "",
        markerstrokewidth = 0.5,
        markerstrokecolor = :black,
        markeralpha = 0.9
    )

    # Annotate how many agents are inside a TTC car.
    pos_wagons_x = [w[1] for w in pos_wagons]
    pos_wagons_y = [w[2] for w in pos_wagons]
    annot_text = [text(string(length(wagon.passengers)),8,:black) for wagon in model.wagons]
    annotate!(background, pos_wagons_x, pos_wagons_y, annot_text)

end


using YaoBase
using Test

# mocked registers
struct TestRegister <: AbstractRegister{2} end

YaoBase.nqudits(::TestRegister) = 8
YaoBase.nactive(::TestRegister) = 2

export TestInterfaceRegister
struct TestInterfaceRegister <: AbstractRegister{2} end

@testset "Test general interface" begin
    @test_throws MethodError nactive(TestInterfaceRegister())
    @test_throws MethodError nqubits(TestInterfaceRegister())
    @test_throws MethodError nremain(TestInterfaceRegister())
end

@testset "adjoint register" begin
    @test adjoint(TestRegister()) isa AdjointRegister
    @test adjoint(adjoint(TestRegister())) isa TestRegister
    @test nqubits(adjoint(TestRegister())) == 8
    @test nactive(adjoint(TestRegister())) == 2
end


# Iterator State Placeholders
@inline _dummy_iterate_state(itr) = iterate(itr)[2]
@inline _dummy_iterate_state(itr::Array) = 0
@inline _dummy_iterate_state(itr::OrdinalRange{T}) where {T} = zero(T)
@inline _dummy_iterate_state(itr::StaticArray) = (Base.OneTo(1), 0)
@inline _dummy_iterate_state(itr::Base.Generator) = _dummy_iterate_state(itr.iter)
@inline _dummy_iterate_state(itr::Iterators.Filter) = _dummy_iterate_state(itr.itr)

@inline _dummy_iterate_result(itr) = iterate(itr)
@inline _dummy_iterate_result(itr::AbstractArray{T}) where {T} = (zero(T), _dummy_iterate_state(itr))

# ControlInterval Iteration
@inline Base.iterate(c::ControlInterval) = (c, nothing)
@inline Base.iterate(c::ControlInterval, ::Any) = nothing
@inline Base.eltype(::Type{CI}) where {CI<:ControlInterval} = CI
@inline Base.length(::ControlInterval) = 1

# Propagate
struct Propagate{D,S,C,T}
    f::D
    x0::S
    cs::C
    ss::T
end

Base.length(itr::Propagate) = length(itr.ss)
Base.eltype(::Type{Propagate{D,S,C,T}}) where {D,S,C,T} = S

@inline function _start(itr::Propagate)
    cs_iter_result = iterate(itr.cs)
    if (cs_done = cs_iter_result === nothing)
        cs_iter_result = _dummy_iterate_result(itr.cs)
    end
    ((itr.x0, zero(eltype(itr.ss)), cs_done, cs_iter_result),)
end

@inline function Base.iterate(itr::Propagate, state=_start(itr))
    x, t, cs_done, cs_iter_result = state[1][1], state[1][2], state[1][3], state[1][4]
    ss_state = Base.tail(state)

    ss_iter_result = iterate(itr.ss, ss_state...)
    ss_iter_result === nothing && return nothing
    s, ss_state = ss_iter_result

    c, cs_state = cs_iter_result
    if !cs_done
        while s >= t + duration(c)
            x = propagate(itr.f, x, c)
            t += duration(c)
            cs_iter_result = iterate(itr.cs, cs_state)
            cs_iter_result === nothing && return x, ((x, t, true, (c, cs_state)), ss_state)
            c, cs_state = cs_iter_result
        end
        y = propagate(itr.f, x, c, s - t)
    else
        y = x
    end

    y, ((x, t, cs_done, cs_iter_result), ss_state)
end

# InstantaneousControl
struct InstantaneousControl{C,T}
    cs::C
    ss::T
end

Base.length(itr::InstantaneousControl) = length(itr.ss)
function Base.eltype(::Type{InstantaneousControl{C,T}}) where {C,T}
    typeof(instantaneous_control(zero(eltype(C)), zero(eltype(T))))
end

@inline function _start(itr::InstantaneousControl)
    cs_iter_result = iterate(itr.cs)
    if (cs_done = cs_iter_result === nothing)
        cs_iter_result = _dummy_iterate_result(itr.cs)
    end
    ((zero(eltype(itr.ss)), cs_done, cs_iter_result),)
end

@inline function Base.iterate(itr::InstantaneousControl, state=_start(itr))
    t, cs_done, cs_iter_result = state[1][1], state[1][2], state[1][3]
    ss_state = Base.tail(state)

    ss_iter_result = iterate(itr.ss, ss_state...)
    ss_iter_result === nothing && return nothing
    s, ss_state = ss_iter_result

    c, cs_state = cs_iter_result
    if !cs_done
        while s >= t + duration(c)
            t += duration(c)
            cs_iter_result = iterate(itr.cs, cs_state)
            cs_iter_result === nothing && return (instantaneous_control(c, duration(c)),
                                                  ((t, true, (c, cs_state)), ss_state))
            c, cs_state = cs_iter_result
        end
        ic = instantaneous_control(c, s-t)
    else
        ic = instantaneous_control(c, duration(c))
    end

    ic, ((t, cs_done, cs_iter_result), ss_state)
end


include("util.jl")
using ImageFiltering

function fuck(; debug_level=verbose)
       printd("fuckyou", verbose, debug_level)
       printd("fuckyou", concise, debug_level)
       printd("fuckyou", quiet, debug_level)

end
"""
Replaces by value e (by default machine precision eps())
"""
function zero_protection(a::Array{Float64,2}, e::Float64 = eps())
	(h,w) = size(a)
    #println(" e = $e")
	for j in 1:w
		for i in 1:h
			if @inbounds a[i,j] < e
				@inbounds a[i,j] = e
			end
		end
	end
end

function richardson_lucy(image::Array{Float64,2}, psf::AbstractArray;
                         iterations::Int=10,
                         view_window::Int=1,
                         fom_window::Int=64,
                         debug_level=mute)

    # init
    image_ref  = copy(image)
    psf_mirror = flip2d(psf)
    im_deconv  = 0.5 * ones(size(image))

    printd("view of image center = $(center_view(image_ref, view_window))",
           verbose, debug_level)
    printd("view of image deconv = $(center_view(im_deconv, view_window))",
           verbose, debug_level)
    printd("view of psf center   = $(center_view(collect(psf), view_window))",
           verbose, debug_level)
    printd("view of psf* deconv  = $(center_view(collect(psf_mirror), view_window))",
           verbose, debug_level)

    dff0 = 1.
    DFR = zeros(iterations) * 1.  # Relative differences

    for i in 1:iterations
        # Using imfilter with reflect to produce a convolution (rather than a correlation),
        # psf must be centered so that final image has the same size than original image
        x = imfilter(im_deconv, reflect(centered(psf)), Fill(0) )

        printd(" iteration = $i x ($i) view = $(center_view(x, view_window)) ",
               verbose, debug_level)

        # protect against division by zero
        limit_boundaries(x, eps())

        # relative blur
        relative_blur = image ./ x

        printd(" relBlur($i) view = $(center_view(relative_blur, view_window)) ",
               verbose, debug_level)

        # deconvolution
        im_deconv = im_deconv .* imfilter(relative_blur, reflect(centered(psf_mirror)), Fill(0))

        printd("im_deconv($i) view = $(center_view(relative_blur, view_window)) ",
               verbose, debug_level)

        # Relative differences
        dff = image_fom(collect(center_view(image_ref, 64)), collect(center_view(im_deconv, 64)))
        delta = (dff0 - dff) / dff0

        printd(" iteration $i ", concise, debug_level)
        printd(" diff($i)  = $dff ", concise, debug_level)
        printd(" delta($i)  = $delta ", concise, debug_level)

        DFR[i] = delta

        dff0 = dff
        image_ref = copy(im_deconv)
    end

    # clip image ( 0 < x <= 1)
    limit_boundaries(im_deconv, 0.,1.)
    return im_deconv, DFR

end

function frichardson_lucy(image::Array{Float64,2}, psf::AbstractArray; iterations::Int=10)

    println("Running RL with $iterations iterations")
    psf_mirror = flip2d(psf)
    im_deconv  = 0.5 * ones(size(image))

    for i in 1:iterations
        # Using imfilter with reflect to produce a convolution (rather than a correlation),
        # psf must be centered so that final image has the same size than original image
        x = imfilter(im_deconv, reflect(centered(psf)), Fill(0))

        # protect against division by zero
        #zero_protection(x)

        # relative blur
        relative_blur = image ./ x

        # deconvolution
        im_deconv = im_deconv .* imfilter(relative_blur,
                                          reflect(centered(psf_mirror)), Fill(0))
    end

    # clip image ( 0 < x <= 1)
    limit_boundaries(im_deconv, 0.,1.)
    return im_deconv

end


# ==========================================================================================================
# Backend funcion

formula(trm::TableRegressionModel) = trm.mf.f

function deviances(trm::TableRegressionModel{<: Union{LinearModel, GeneralizedLinearModel}, <: AbstractArray{T}};
                    type::Int = 1, kwargs...) where {T <: BlasReal}
    isa(trm.model.pp, DensePredChol) || throw(
            ArgumentError("Methods for other PredChol types is not implemented; use model with DensesPredChol instead."))

    assign, f = trm.mm.assign, formula(trm).rhs
    # Determine null model by link function
    ## model with 0 factors as null
    ## if 0 factor is not allowed,
    ### eg, ~ 1 + A + B + ...
    ### ~ 1 as null
    ### eg, ~ 0 + A + A & B
    ### ~ A as null
    # start value
    start = isnullable(trm.model) ? 0 : first(assign)
    err1 = ArgumentError("Invalid set of model specification for ANOVA; not enough variables provided.")
    err2 = ArgumentError("Invalid set of model specification for ANOVA; try adding variables without zeros.")
    if type == 1
        todel = union(start, assign)
        # ~ 0 + A, ~ 1
        start > 0 && (length(todel) > 1 || throw(err1))
        devs = zeros(T, length(todel) + 1)
        @inbounds for id in eachindex(todel)
            devs[id] = deviance(trm, @view(todel[id + 1:end]); kwargs...)
        end
        devs = -diff(devs)
    elseif type == 2
        # Problems with Fullrank promotion occur
        ## ~ 0 + A (categorical) + B (categorical)
        ## A is aliased with 1, but 1 is not in formula.
        ## A is promoted to full rank.
        ## B is also aliased with 1, but it is not promoted.
        ## Doing type 2 and 3 anova cause problem
        ## when deleting promoted categorical variables.
        # Some warning or checks?
        todel = unique(assign)
        if start > 0
            # ~ 0 + A + A & B, all terms are related to A, ~ A as null
            selectcoef(f, Val(first(todel))) == Set(todel) && popfirst!(todel)
            # ~ 0 + A, ~ 1
            if isempty(todel)
                throw(err1)
            # ~ 0 + 2 categorical
            # Do not inspect InteractionTerm
            elseif length(todel) == 2
                all(t -> isa(t, CategoricalTerm), f.terms[todel]) && throw(err2)
            end
        end
        devs = zeros(T, length(todel) + 1)
        # cache fitted
        dict_devs = Dict{Set{Int}, T}()
        @inbounds for (id, del) in enumerate(todel)
            delcoef = selectcoef(f, Val(del))
            dev1 = get(dict_devs, delcoef, (push!(dict_devs, delcoef => deviance(trm, delcoef; kwargs...)); dict_devs[delcoef]))
            delete!(delcoef, del)
            dev2 = get(dict_devs, delcoef, (push!(dict_devs, delcoef => deviance(trm, delcoef; kwargs...)); dict_devs[delcoef]))
            devs[id] = dev1 - dev2
        end
        devs[end] = deviance(trm, 0; kwargs...)
    else
        todel = unique(assign)
        if start > 0
            # ~ 0 + A, ~ 1
            if length(todel) <= 1
                throw(err1)
            # ~ 0 + 2 categorical or
            # ~ 1 + 1 categorical
            # Do not inspect InteractionTerm
            elseif length(todel) == 2
                if isa(f.terms[todel][1], Union{CategoricalTerm, InterceptTerm{true}})
                    isa(f.terms[todel][1], CategoricalTerm) && throw(err2)
                end
            end
        end
        devs = zeros(T, length(todel) + 1)
        @inbounds for (id, del) in enumerate(todel)
            devs[id] = deviance(trm, del; kwargs...)
        end
        devr = deviance(trm, 0; kwargs...)
        devs .-= devr
        devs[end] = devr
    end
    # every method end with deviance(trm, 0; kwargs...), ie full model fit.
    installbeta!(trm.model.pp) # ensure model unchanged
    tuple(devs...)
end

submm(trm::TableRegressionModel, exclude) =
    trm.mm.m[:, map(!in(exclude), trm.mm.assign)]
submm(trm::TableRegressionModel, exclude::Int) =
    trm.mm.m[:, map(!=(exclude), trm.mm.assign)]

# Backend for LinearModel
# Only used by type 2 SS
function deviance(trm::TableRegressionModel{<: LinearModel}, exclude)
    p = trm.model.pp

    # reschema to calculate new X, reset beta
    X = submm(trm, exclude)
    p.beta0 = p.delbeta = repeat([0], size(X, 2))
    p.scratchbeta = similar(p.beta0)

    # cholesky
    F = Hermitian(float(X'X))
    p.chol = updatechol(p.chol, F)

    # reset scratch
    p.scratchm1 = similar(X)
    p.scratchm2 = similar(p.chol.factors)

    isempty(trm.model.rr.wts) ? delbeta!(p, X, trm.model.rr.y) : delbeta!(p, X, trm.model.rr.y, trm.model.rr.wts)
    updateμ!(trm.model.rr, linpred(p, X))
    # installbeta is ommited
end

updatechol(::CholeskyPivoted, F::AbstractMatrix{<: BlasReal}) =
    cholesky!(F, Val(true), tol = -one(eltype(F)), check = false)

updatechol(::Cholesky, F::AbstractMatrix{<: BlasReal}) = cholesky!(F)

# Backend for GeneralizedLinearModel
function deviance(trm::TableRegressionModel{<: GeneralizedLinearModel{<: GlmResp{<: GLM.FPVector, <: GLM.UnivariateDistribution, <: Link}}}, exclude;
                    verbose::Bool = false,
                    maxiter::Integer = 30,
                    minstepfac::Real = 0.001,
                    atol::Real = 1e-6,
                    rtol::Real = 1e-6,
                    kwargs...)

    # Check arguments
    maxiter >= 1       || throw(ArgumentError("maxiter must be positive"))
    0 < minstepfac < 1 || throw(ArgumentError("minstepfac must be in (0, 1)"))

    # Extract fields and set convergence flag
    cvg, p, r = false, trm.model.pp, trm.model.rr
    lp = r.mu

    # subset X, reset beta
    X = submm(trm, exclude)
    p.beta0 = p.delbeta = repeat([0], size(X, 2))
    p.scratchbeta = similar(p.beta0)
    # cholesky
    F = Hermitian(float(X'X))
    p.chol = updatechol(p.chol, F)

    # reset scratch
    p.scratchm1 = similar(X)
    p.scratchm2 = similar(p.chol.factors)

    # Initialize β, μ, and compute deviance
    delbeta!(p, X, GLM.wrkresp(r), r.wrkwt)
    linpred!(lp, p, X)
    updateμ!(r, lp)
    installbeta!(p)
    devold = deviance(trm.model)

    for i = 1:maxiter
        f = 1.0 # line search factor
        local dev

        # Compute the change to β, update μ and compute deviance
        try
            delbeta!(p, X, r.wrkresid, r.wrkwt)
            linpred!(lp, p, X)
            updateμ!(r, lp)
            dev = deviance(trm.model)
        catch e
            isa(e, DomainError) ? (dev = Inf) : rethrow(e)
        end

        # Line search
        ## If the deviance isn't declining then half the step size
        ## The rtol*dev term is to avoid failure when deviance
        ## is unchanged except for rouding errors.
        while dev > devold + rtol*dev
            println("ok")
            f /= 2
            f > minstepfac || error("step-halving failed at beta0 = $(p.beta0)")
            try
                updateμ!(r, linpred(p, X, f))
                dev = deviance(trm.model)
            catch e
                isa(e, DomainError) ? (dev = Inf) : rethrow(e)
            end
        end
        installbeta!(p, f)

        # Test for convergence
        verbose && println("Iteration: $i, deviance: $dev, diff.dev.:$(devold - dev)")
        if devold - dev < max(rtol*devold, atol)
            cvg = true
            break
        end
        @assert isfinite(dev)
        devold = dev
    end
    cvg || throw(ConvergenceException(maxiter))
    devold
end

# Linear prediction
function delbeta!(p::DensePredChol{<: BlasReal, <: Cholesky},
                    X::Matrix{<: BlasReal},
                    r::Vector{<: BlasReal})
    ldiv!(p.chol, mul!(p.delbeta, transpose(X), r))
    p
end
# β = (X'X)⁻¹X'y

function delbeta!(p::DensePredChol{<: BlasReal, <: CholeskyPivoted},
                    X::Matrix{<: BlasReal},
                    r::Vector{<: BlasReal})
    ch = p.chol
    beta = mul!(p.delbeta, adjoint(X), r)
    rnk = rank(ch)
    if rnk == length(beta)
        ldiv!(ch, beta)
    elseif size(ch.factors) == (0, 0) # Rank undtermined when the matrix is 0x0
        ldiv!(ch, beta)
    else
        permute!(beta, ch.piv)
        for k = (rnk + 1):length(beta)
            beta[k] = -zero(T)
        end
        LAPACK.potrs!(ch.uplo, view(ch.factors, 1:rnk, 1:rnk), view(beta, 1:rnk))
        invpermute!(beta, ch.piv)
    end
    p
end

# exterior weights
function delbeta!(p::DensePredChol{<: BlasReal, <: Cholesky},
                    X::Matrix{<: BlasReal},
                    r::Vector{<: BlasReal},
                    wt::Vector{<: BlasReal})
    scr = mul!(p.scratchm1, Diagonal(wt), X)
    cholesky!(Hermitian(mul!(cholfactors(p.chol), transpose(scr), X), :U))
    mul!(p.delbeta, transpose(scr), r)
    ldiv!(p.chol, p.delbeta)
    p
end

function delbeta!(p::DensePredChol{<: BlasReal, <: CholeskyPivoted},
                    X::Matrix{<: BlasReal},
                    r::Vector{<: BlasReal},
                    wt::Vector{<: BlasReal})
    cf = cholfactors(p.chol)
    piv = p.chol.piv
    cf .= mul!(p.scratchm2, adjoint(mul!(p.scratchm1, Diagonal(wt), X)), X)[piv, piv]
    cholesky!(Hermitian(cf, Symbol(p.chol.uplo)))
    ldiv!(p.chol, mul!(p.delbeta, transpose(p.scratchm1), r))
    p
end

# experimental
function delbeta!(p::SparsePredChol{<: BlasReal},
                    X::Matrix{<: BlasReal},
                    r::Vector{<: BlasReal},
                    wt::Vector{<: BlasReal})
    scr = mul!(p.scratch, Diagonal(wt), X)
    XtWX = X'*scr
    c = p.chol = cholesky(Symmetric{eltype(XtWX),typeof(XtWX)}(XtWX, 'L'))
    p.delbeta = c \ mul!(p.delbeta, adjoint(scr), r)
end

# skip updating beta0
linpred(p::LinPred, X::Matrix, f::Real = 1.0) =
    linpred!(Vector{eltype(p.X)}(undef, size(p.X, 1)), p, X, f)

function linpred!(out, p::LinPred, X::Matrix, f::Real = 1.0)
    mul!(out, X, iszero(f) ? p.beta0 : broadcast!(muladd, p.scratchbeta, f, p.delbeta, p.beta0))
end

# Create nestedmodels
function nestedmodels(trm::TableRegressionModel{<: LinearModel}; null::Bool = true, kwargs...)
    null = null && isnullable(trm.model.pp.chol)
    assign = unique(trm.mm.assign)
    pop!(assign)
    dropcollinear, range = null ? (false, union(0, assign)) : (true, assign)
    wts = trm.model.rr.wts
    trms = map(range) do id
        # create sub-formula, modify schema, create mf and mm
        mf, mm = subtablemodel(trm, id)
        y = response(mf)
        TableRegressionModel(fit(trm.mf.model, mm.m, y; wts, dropcollinear, kwargs...), mf, mm)
    end
    (trms..., trm)
end

function nestedmodels(trm::TableRegressionModel{<: GeneralizedLinearModel}; null::Bool = true, kwargs...)
    null = null && isnullable(trm.model)
    distr = trm.model.rr.d
    link = typeof(trm.model.rr).parameters[3]()
    wts = trm.model.rr.wts
    offset = trm.model.rr.offset
    assign = unique(trm.mm.assign)
    pop!(assign)
    range = null ? union(0, assign) : assign
    trms = map(range) do id
        # create sub-formula, modify schema, create mf and mm
        mf, mm = subtablemodel(trm, id)
        y = response(mf)
        TableRegressionModel(fit(trm.mf.model, mm.m, y, distr, link; wts, offset, kwargs...), mf, mm)
    end
    (trms..., trm)
end

nestedmodels(::Type{LinearModel}, formula, data; null::Bool = true, kwargs...) =
    nestedmodels(lm(formula, data; kwargs...); null, kwargs...)
nestedmodels(::Type{GeneralizedLinearModel}, formula, data,
                distr::UnivariateDistribution, link::Link = canonicallink(distr);
                null::Bool = true, kwargs...) =
    nestedmodels(glm(formula, data, distr, link; kwargs...); null, kwargs...)

# backend for implementing nestedmodels or dropterms (in future)
function subtablemodel(trm::TableRegressionModel, id; reschema::Bool = false)
    # create sub-formula, modify schema, create mf and mm
    f = formula(trm)
    subf = subformula(f.lhs, f.rhs, id; reschema)
    if reschema
        #=
        terms = setdiff(getterms(f.rhs), getterms(subf.rhs))
        filter!(!=(Symbol("1")), terms)
        pair = collect(pairs(trm.mf.data))
        filter!(x->!in(x.first, terms), pair)
        =#
        mf = ModelFrame(subf, trm.mf.data;
                        model = trm.mf.model, contrasts = extract_contrasts(f))
    else
        #schema = deepcopy(trm.mf.schema)
        #for term in terms
        #    pop!(schema.schema, Term(term))
        #end
        mf = ModelFrame(subf, schema, trm.mf.data, trm.mf.model)
    end
    mm = ModelMatrix(mf)
    mf, mm
end

# Null model for CholeskyPivoted is unstable
## For nestedmodels
isnullable(::CholeskyPivoted) = false
isnullable(::Cholesky) = true
isnullable(::InverseLink) = false
isnullable(::InverseSquareLink) = false
isnullable(::Link) = true

## For deviances
isnullable(::LinearModel) = true
isnullable(model::GeneralizedLinearModel) = isnullable(typeof(model.rr).parameters[3]())


module ITensorVisualizationBase

using AbstractTrees
using Compat
using GeometryBasics
using Graphs
using ITensors
using ITensors.ITensorVisualizationCore
using LinearAlgebra
using MetaGraphs
using NetworkLayout
using SparseArrays
using Statistics

# Avoid conflict between `Graphs.contract` and `ITensors.contract`
using Graphs:
  Graphs,
  AbstractEdge,
  AbstractGraph,
  SimpleGraph,
  SimpleDiGraph,
  add_edge!,
  add_vertex!,
  all_neighbors,
  dst,
  edges,
  ne,
  neighbors,
  nv,
  src,
  vertices

using ITensors: data, QNIndex

import ITensors.ITensorVisualizationCore: visualize, visualize!, visualize_sequence

export @visualize,
  @visualize!,
  @visualize_noeval,
  @visualize_noeval!,
  @visualize_sequence,
  @visualize_sequence_noeval,
  circuit_network,
  itensornetwork,
  layered_layout,
  IndexLabels

# Some general graph functionality
include("graphs.jl")

# Some general layout functionality
include("layered_layout.jl")

# Backends interface
include("backends_interface.jl")
include("defaults.jl")

# Conversion betweens graphs and ITensor networks
include("itensor_graph.jl")

# Visualizing ITensor networks
include("visualize.jl")

end


# This file is a part of TextSearch.jl

import KCenters
export update!, joinmodel
"""
    update!(a::VectorModel, b::VectorModel)

Updates `a` with `b` inplace; returns `a`.
"""
function update!(a::VectorModel, b::VectorModel)
m = a.m
for (token, idfreq) in b.tokens
    x = get(a.tokens, token, nothing)
    if x === nothing
        m += 1
        a.tokens[token] = IdFreq(m, x.freq)
    else
        a.tokens[token] = IdFreq(idfreq.id, idfreq.freq + x.freq)
    end
end

a.maxfreq = max(a.maxfreq, b.maxfreq)
a.n += b.n
a.m = m
a.id2token = Dict(idfreq.id => t for (t, idfreq) in a.tokens)
a
end

"""
    joinmodel(arr::AbstractVector{VectorModel})
    joinmodel(arr::AbstractVector{EntModel})

Joins a list of models into a single one; frequencies and weights are computed as the mean of its non-missing occurrences
"""
function joinmodel(arr::AbstractVector{VectorModel})
    item = first(arr)
    # Vocabulary = Dict{Symbol, IdFreq}
    tokens_ = Dict{Symbol,Int}()
    counts_ = Dict{Symbol,Int32}()

    maxfreq = 0
    for model in arr
        maxfreq = max(maxfreq, model.maxfreq)
        for (k, v) in model.tokens
            tokens_[k] = get(tokens_, k, 0.0) + v.freq # ceil(Int, 0.5 * (
            counts_[k] = get(counts_, k, 0) + 1
        end
    end

    tokens = TextSearch.Vocabulary()
    for (k, v) in tokens_
        tokens[k] = TextSearch.IdFreq(length(tokens) + 1, ceil(Int, v / counts_[k]))
    end

    id2token = Dict(idfreq.id => token for (token, idfreq) in tokens)
    VectorModel(item.config, tokens, id2token, maxfreq, length(tokens_), item.n)
end

function glue(arr::AbstractVector{EntModel})
    item = first(arr)
    tokens_ = Dict{Symbol,Float32}()
    counts_ = Dict{Symbol,Int32}()

    for model in arr
        for (k, v) in model.tokens
            tokens_[k] = get(tokens_, k, 0.0) + v.weight
            counts_[k] = get(counts_, k, 0) + 1
        end
    end

    tokens = TextSearch.WeightedVocabulary()
    for (k, v) in tokens_
        tokens[k] = TextSearch.IdWeight(length(tokens) + 1, v / counts_[k])
    end

    id2token = Dict(w.id => token for (token, w) in tokens)

    EntModel(item.config, tokens, id2token, length(tokens), item.n)
end


using DelimitedFiles
using Flux

# load a specialized model format
function loadPyNNTxt(dest::AbstractString)
  mw = Vector{Array{Float32}}()
  @show files = readdir(dest)
  for f in files
    push!(mw, readdlm(joinpath(dest,f)))
  end
  return mw
end

# convenience function to load specific model format from tensorflow
function loadPose2OdoNNModelIntoFlux(dest::AbstractString; pad::Bool=false)
  weights = loadPyNNTxt(dest::AbstractString)
  buildPose2OdoNN_01_FromWeights(weights, pad=pad)
end

# function buildPose2OdoNN_01_FromWeightsGPU(pywe)
#   buildPose2OdoNN_01_FromElements(collect(pywe[1]) |> gpu, collect(pywe[2][:]) |> gpu, collect(pywe[3]') |> gpu, collect(pywe[4][:]) |> gpu, collect(pywe[5]') |> gpu, collect(pywe[6][:]) |> gpu) |> gpu
# end
#
# function loadPose2OdoNNModelIntoFluxGPU(dest::AbstractString)
#   weights = loadPyNNTxt(dest::AbstractString)
#   buildPose2OdoNN_01_FromWeightsGPU(weights)
# end

#


const DROPOUT       = 0
const BATCH_SIZE    = 32
const SEQ_LENGTH    = 32
const DIM_HIDDEN    = 256
const DIM_EMBED     = 256
const LSTM_N_LAYER  = 2
const N_EPOCH       = 21
const BASE_LR       = 0.01
const WEIGHT_DECAY  = 0.00001
const CLIP_GRADIENT = 1
const NAME          = :ptb
const N_GPU         = 4
const USE_GPU       = true
const DATA_TR_RATIO = 0.9
const CKPOINT_PREFIX = joinpath(dirname(@__FILE__), "checkpoints/$NAME")

const BATCH_SIZE_SMP= 10
const SAMPLE_LENGTH = 100
const SAMPLE_START  = 'a'

const UNKNOWN_CHAR  = Char(0)
const INPUT_FILE    = joinpath(dirname(@__FILE__), "input.txt")
const VOCAB_FILE    = joinpath(dirname(@__FILE__), "vocab.dat")


module AdvancedHMC

const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_AHMC", "0")))

using Statistics: mean, var, middle
using LinearAlgebra: Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky
using StatsFuns: logaddexp
using LazyArrays: BroadcastArray
using Random: GLOBAL_RNG, AbstractRNG
using ProgressMeter
using Parameters: @unpack, reconstruct
using ArgCheck: @argcheck

import StatsBase: sample

# Support of passing a vector of RNGs
Base.rand(rng::AbstractVector{<:AbstractRNG}) = rand.(rng)
Base.randn(rng::AbstractVector{<:AbstractRNG}) = randn.(rng)
function Base.randn(
    rng::AbstractVector{<:AbstractRNG},
    T,
    sz::Int...)
    return cat(randn.(rng, T, first(sz))...; dims=2)
end

# Notations
# ℓπ: log density of the target distribution
# ∂ℓπ∂θ: gradient of the log density of the target distribution
# θ: position variables / model parameters
# r: momentum variables

include("adaptation/Adaptation.jl")
export UnitEuclideanMetric, DiagEuclideanMetric, DenseEuclideanMetric
export NesterovDualAveraging, Preconditioner, NaiveHMCAdaptor, StanHMCAdaptor
using .Adaptation

using .Adaptation: AbstractScalarOrVec
import .Adaptation: adapt!, NesterovDualAveraging

include("hamiltonian.jl")
export Hamiltonian
include("integrator.jl")
export Leapfrog, JitteredLeapfrog, TemperedLeapfrog
include("trajectory.jl")
export StaticTrajectory, HMCDA, NUTS, EndPointTS, SliceTS, MultinomialTS, ClassicNoUTurn, GeneralisedNoUTurn, find_good_eps

include("diagnosis.jl")
include("sampler.jl")
export sample

### Init

using Requires

function __init__()
    include(joinpath(@__DIR__, "contrib/diffeq.jl"))
    include(joinpath(@__DIR__, "contrib/ad.jl"))
end

end # module


### Basic Particle Filter ###
# implements the POMDPs.jl Updater interface
"""
    BasicParticleFilter(predict_model, reweight_model, resampler, n_init::Integer, rng::AbstractRNG)
    BasicParticleFilter(model, resampler, n_init::Integer, rng::AbstractRNG)

Construct a basic particle filter with three steps: predict, reweight, and resample.

In the second constructor, `model` is used for both the prediction and reweighting.
"""
mutable struct BasicParticleFilter{PM,RM,RS,RNG<:AbstractRNG,PMEM} <: Updater
    predict_model::PM
    reweight_model::RM
    resampler::RS
    n_init::Int
    rng::RNG
    _particle_memory::PMEM
    _weight_memory::Vector{Float64}
end

## Constructors ##
function BasicParticleFilter(model, resampler, n::Integer, rng::AbstractRNG=Random.GLOBAL_RNG)
    return BasicParticleFilter(model, model, resampler, n, rng)
end

function BasicParticleFilter(pmodel, rmodel, resampler, n::Integer, rng::AbstractRNG=Random.GLOBAL_RNG)
    return BasicParticleFilter(pmodel,
                               rmodel,
                               resampler,
                               n,
                               rng,
                               particle_memory(pmodel),
                               Float64[]
                              )
end

"""
    particle_memory(m)

Return a suitable container for particles produced by prediction model `m`.

This should usually be an empty `Vector{S}` where `S` is the type of the state for prediction model `m`. Size does not matter because `resize!` will be called appropriately within `update`.
"""
function particle_memory end

function update(up::BasicParticleFilter, b::ParticleCollection, a, o)
    pm = up._particle_memory
    wm = up._weight_memory
    resize!(pm, n_particles(b))
    resize!(wm, n_particles(b))
    # predict_and_reweight!(pm, wm, up.predict_model, up.reweight_model, b, a, o, up.rng)
    predict!(pm, up.predict_model, b, a, o, up.rng)
    reweight!(wm, up.reweight_model, b, a, pm, o, up.rng)

    return resample(up.resampler,
                    WeightedParticleBelief(pm, wm, sum(wm), nothing),
                    up.predict_model,
                    up.reweight_model,
                    b, a, o,
                    up.rng)
end

function Random.seed!(f::BasicParticleFilter, seed)
    Random.seed!(f.rng, seed)
    return f
end

"""
    predict!(pm, m, b, u, rng)
    predict!(pm, m, b, u, y, rng)

Fill `pm` with predicted particles for the next time step.

A method of this function should be implemented by prediction models to be used in a [`BasicParticleFilter`](@ref). `pm` should be a correctly-sized vector created by [`particle_memory`](@ref) to hold a one-step-propagated particle for each particle in `b`.

Normally the observation `y` is not needed, so most prediction models should implement the first version, but the second is available for heuristics that use `y`.

# Arguments
- `pm::Vector`: memory for holding the propagated particles; created by [`particle_memory`](@ref) and resized to `n_particles(b)`.
- `m`: prediction model, the "owner" of this function
- `b::ParticleCollection`: current belief; each particle in this belief should be propagated one step and inserted into `pm`.
- `u`: control or action
- `rng::AbstractRNG`: random number generator; should be used for any randomness in propagation for reproducibility.
- `y`: measuerement/observation (usually not needed)
"""
function predict! end

"""
    reweight!(wm, m, b, a, pm, y)
    reweight!(wm, m, b, a, pm, y, rng)

Fill `wm` likelihood weights for each particle in `pm`.

A method of this function should be implemented by reweighting models to be used in a [`BasicParticleFilter`](@ref). `wm` should be a correctly-sized vector to hold weights for each particle in pm.

Normally `rng` is not needed, so most reweighting models should implement the first version, but the second is available for heuristics that use random numbers.

# Arguments
- `wm::Vector{Float64}`: memory for holding likelihood weights.
- `m`: reweighting model, the "owner" of this function
- `b::ParticleCollection`: previous belief; `pm` should contain a propagated particle for each particle in this belief
- `u`: control or action
- `pm::Vector`: memory for holding current particles; these particle have been propagated by `predict!`.
- `y`: measurement/observation
- `rng::AbstractRNG`: random number generator; should be used for any randomness for reproducibility.
"""
function reweight! end

predict!(pm, m, b, a, o, rng) = predict!(pm, m, b, a, rng)
reweight!(wm, m, b, a, pm, o, rng) = reweight!(wm, m, b, a, pm, o)

"""
    predict(m, b, u, rng)

Simulate each of the particles in `b` forward one time step using model `m` and contol input `u` returning a vector of states. Calls [`predict!`](@ref) internally - see that function for documentation.

This function is provided for convenience only. New models should implement `predict!`.
"""
function predict end

function predict(m, b, args...)
    pm = particle_memory(m)
    resize!(pm, n_particles(b))
    predict!(pm, m, b, args...)
    return pm
end
predict(f::BasicParticleFilter, args...) = predict(f.predict_model, args...)

"""
    reweight(m, b, u, pm, y)

Return a vector of likelihood weights for each particle in `pm` given observation `y`.

`pm` can be generated with `predict(m, b, u, rng)`.

This function is provided for convenience only - new reweighting models should implement `reweight!`.
"""
function reweight end

function reweight(m, b, args...)
    wm = Vector{Float64}(undef, n_particles(b))
    reweight!(wm, m, b, args...)
    return wm
end
reweight(f::BasicParticleFilter, args...) = reweight(f.reweight_model, args...)