text:
module TrajectoriesTest

using QPControl.Trajectories
using Test
using LinearAlgebra
using Random
using Rotations
using RigidBodyDynamics
using StaticArrays
using StaticUnivariatePolynomials
using ForwardDiff

using QPControl.Trajectories: subfunctions, breaks
using StaticUnivariatePolynomials: exponential_integral

const SUP = StaticUnivariatePolynomials

@testset "fit_cubic" begin
    rng = MersenneTwister(15)
    for i = 1 : 10
        x0 = rand(rng)
        xf = rand(rng)
        y0 = rand(rng)
        yd0 = rand(rng)
        yf = rand(rng)
        ydf = rand(rng)

        p = fit_cubic(x0=x0, xf=xf, y0=y0, yd0=yd0, yf=yf, ydf=ydf)
        pd = SUP.derivative(p)
        pdd = SUP.derivative(pd)

        atol = 1e-6
        @test p(x0) ≈ y0 atol=atol
        @test pd(x0) ≈ yd0 atol=atol
        @test p(xf) ≈ yf atol=atol
        @test pd(xf) ≈ ydf atol=atol

        @test length(p.coeffs) == 4
    end
end

@testset "fit_quintic" begin
    rng = MersenneTwister(15)
    for i = 1 : 10
        x0 = rand(rng)
        xf = rand(rng)
        y0 = rand(rng)
        yd0 = rand(rng)
        ydd0 = rand(rng)
        yf = rand(rng)
        ydf = rand(rng)
        yddf = rand(rng)

        p = fit_quintic(x0=x0, xf=xf, y0=y0, yd0=yd0, ydd0=ydd0, yf=yf, ydf=ydf, yddf=yddf)
        pd = SUP.derivative(p)
        pdd = SUP.derivative(pd)

        atol = 1e-6
        @test p(x0) ≈ y0 atol=atol
        @test pd(x0) ≈ yd0 atol=atol
        @test pdd(x0) ≈ ydd0 atol=atol
        @test p(xf) ≈ yf atol=atol
        @test pd(xf) ≈ ydf atol=atol
        @test pdd(xf) ≈ yddf atol=atol

        @test length(p.coeffs) == 6
    end
end

@testset "Interpolated, identity interpolator function" begin
    let y0 = 1, yf = 2
        traj = Interpolated(0.0, 1.0, y0, yf)

        y = traj(0.5)
        @test y === 1.5

        y, yd, ydd = traj(0.5, Val(2))
        @test y === 1.5
        @test yd === 1.0
        @test ydd === 0.0
    end

    let f = CartesianFrame3D(), y0 = Point3D(f, 1, 2, 3), yf = Point3D(f, 0, 2, 4)
        traj = Interpolated(0.0, 1.0, y0, yf)
        y, yd, ydd = traj(0.5, Val(2))
        @test y ≈ Point3D(f, 0.5, 2.0, 3.5) atol=1e-15
        @test yd ≈ FreeVector3D(f, -1.0, 0.0, 1.0) atol=1e-15
        @test ydd ≈ FreeVector3D(f, 0.0, 0.0, 0.0) atol=1e-15
    end

    let angle = π / 2, axis = SVector(1.0, 0.0, 0.0), y0 = one(Quat), yf = Quat(AngleAxis(angle, axis...))
        traj = Interpolated(0.0, 1.0, y0, yf)

        y = traj(0.0)
        @test y ≈ y0 atol=1e-15

        y = traj(1.0)
        @test y ≈ yf atol=1e-15

        y, yd, ydd = traj(0.5, Val(2))
        @test y ≈ Quat(AngleAxis(0.5 * angle, axis...)) atol=1e-15
        @test normalize(yd) ≈ axis atol=1e-15
        @test norm(yd) ≈ abs(angle) atol=1e-15
        @test ydd ≈ SVector(0.0, 0.0, 0.0) atol=1e-15
    end
end

@testset "Interpolated, polynomial interpolator function" begin
    interpolator = fit_quintic(x0=0.0, xf=1.0, y0=0.0, yd0=0.0, ydd0=0.0, yf=1.0, ydf=0.0, yddf=0.0)
    x0 = -1.0
    xf = 2.0
    y0 = 2.0
    yf = 3.0
    traj = Interpolated(x0, xf, y0, yf, interpolator, min_num_derivs=Val(2))

    y, yd, ydd = traj(x0, Val(2))
    @test y ≈ y0 atol=1e-10
    @test yd ≈ 0.0 atol=1e-10
    @test ydd ≈ 0.0 atol=1e-10

    y, yd, ydd = traj(xf, Val(2))
    @test y ≈ yf atol=1e-10
    @test yd ≈ 0.0 atol=1e-10
    @test ydd ≈ 0.0 atol=1e-10

    for x in range(x0; stop=xf, length=10)
        _, yd, ydd = traj(x, Val(2))
        yd_forward = ForwardDiff.derivative(traj, x)
        ydd_forward = ForwardDiff.derivative(x -> ForwardDiff.derivative(traj, x), x)
        @test yd ≈ yd_forward atol=1e-4
        @test ydd ≈ ydd_forward atol=1e-4
    end
end

function test_piecewise(traj::Piecewise)
    n = length(subfunctions(traj))
    for (i, t) in enumerate(breaks(traj))
        subtraj = subfunctions(traj)[min(i, n)]
        @test traj(t) == subtraj(t - breaks(traj)[min(i, n)])
        if i < n
            tnext = breaks(traj)[i + 1]
            tmid = (t + tnext) / 2
            @test traj(tmid) == subtraj(tmid - t)
            @test traj(tmid, Val(2)) == subtraj(tmid - t, Val(2))
        end
    end
    t0 = first(breaks(traj))
    tf = last(breaks(traj))
    if traj.clamp
        @test traj(t0 - 1) == traj(t0)
        @test traj(tf + 1) == traj(tf)
    else
        @test_throws DomainError traj(t0 - 1)
        @test_throws DomainError traj(tf + 1)
    end
end

@testset "Piecewise constant" begin
    n = 5
    subtrajectories = [Constant(i) for i = 1 : n]
    breaks = [i^2 - 1 for i = 1 : n + 1]
    test_piecewise(Piecewise(subtrajectories, breaks; clamp=true))
    test_piecewise(Piecewise(subtrajectories, breaks; clamp=false))
end

@testset "Piecewise interpolated" begin
    rng = MersenneTwister(1)
    n = 5
    breaks = [i^2 - 1 for i = 1 : n + 1]
    subtrajectories = map(diff(breaks)) do Δt
        Interpolated(0.0, Float64(Δt), rand(rng), rand(rng))
    end
    test_piecewise(Piecewise(subtrajectories, breaks; clamp=true))
    test_piecewise(Piecewise(subtrajectories, breaks; clamp=false))
end

@testset "PointTrajectory, FreeVectorTrajectory" begin
    rng = MersenneTwister(2)
    interpolator = fit_cubic(x0=0.0, xf=1.0,
        y0=rand(rng), yd0=rand(rng), yf=rand(rng), ydf=rand(rng))
    x0 = -1.0
    xf = 2.0
    y0 = rand(rng, SVector{3})
    yf = rand(rng, SVector{3})
    interpolated = Interpolated(x0, xf, y0, yf, interpolator, min_num_derivs=Val(2))
    frame = CartesianFrame3D()
    t = 1.0
    let
        traj = PointTrajectory(frame, interpolated)
        @test traj(t) === Point3D(frame, interpolated(t))
        point, vel, accel = traj(t, Val(2))
        @test vel === FreeVector3D(frame, interpolated(t, Val(2))[2])
        @test accel === FreeVector3D(frame, interpolated(t, Val(2))[3])
        allocs = let traj=traj, t=t
            @allocated traj(t, Val(2))
        end
        @test allocs == 0
    end
    let
        traj = FreeVectorTrajectory(frame, interpolated)
        @test traj(t) === FreeVector3D(frame, interpolated(t))
        x, xd, xdd = traj(t, Val(2))
        @test xd === FreeVector3D(frame, interpolated(t, Val(2))[2])
        @test xdd === FreeVector3D(frame, interpolated(t, Val(2))[3])
        allocs = let traj=traj, t=t
            @allocated traj(t, Val(2))
        end
        @test allocs == 0
    end
end

@testset "SE(3)" begin
    angular = let angle = π / 2, axis = SVector(1.0, 0.0, 0.0), y0 = one(Quat), yf = Quat(AngleAxis(angle, axis...))
        Interpolated(0.0, 1.0, y0, yf)
    end
    linear = Interpolated(0.0, 1.0, SVector(0.0, 1.0, 2.0), SVector(2.0, 3.0, 4.0))
    body = CartesianFrame3D()
    base = CartesianFrame3D()
    frame = body
    traj = SE3Trajectory(body, base, angular, linear)
    H, T, Td = traj(0.5, Val(2))
    @test H.from == frame
    @test H.to == base
    @test T.body == body
    @test T.base == base
    @test T.frame == frame
    @test Td.body == body
    @test Td.base == base
    @test Td.frame == frame
end

end


text:
## index.jl : unit testing for index.jl
#
# This file is part of MultilevelEstimators.jl - A Julia toolbox for
# Multilevel Monte Carlo Methods (c) Pieterjan Robbe, 2019

@testset "Level                        " begin

level_0 = Level(0)
level_2 = Level(2)
@test typeof(level_0) <: Level
@test typeof(level_2) <: Level
@test level_0 + level_2 == level_2
@test_throws MethodError Level(0.)

end

@testset "Index                        " begin

index_0_0 = Index(0,0)
index_2_3 = Index(2,3)
@test typeof(index_0_0) <: Index{2}
@test typeof(index_2_3) <: Index{2}
@test index_0_0 + index_2_3 == index_2_3
@test_throws MethodError Index(0.,2)

end

ext: jl


text:
module MixedModelsExtras

export icc
export r², r2, adjr², adjr2

using MixedModels
using Statistics
using StatsBase

include("icc.jl")
include("r2.jl")

end


text:
@get_units @with_kw struct FLEXTOPO{T,N}
    # Model time step [s]
    Δt::T | "s"
    # Number of classes
    nclass::Int | "-"
    # Number of cells
    n::Int | "-"
    #dictionary with all possible functions for each store
    dic_function::Dict
    #current class
    kclass::Vector{Int64}
    classes::Vector{String}
    select_snow::Vector{String}
    select_interception::Vector{String}
    select_hortonponding::Vector{String}
    select_hortonrunoff::Vector{String}
    select_rootzone::Vector{String}
    select_fast::Vector{String}
    select_slow::Vector{String}

    #fraction of each class
    hrufrac::Vector{SVector{N,T}} | "-"
    ## PARAMETERS
    ##SNOW
    # Correction factor for precipitation [-]
    pcorr::Vector{T} | "-"
    # Degree-day factor [mm ᵒC⁻¹ Δt⁻¹]
    cfmax::Vector{T} | "mm ᵒC-1 Δt-1"
    # Threshold temperature for snowfall [ᵒC]
    tt::Vector{T} | "ᵒC"
    # Threshold temperature interval length [ᵒC]
    tti::Vector{T} | "ᵒC"
    # Threshold temperature for snowmelt [ᵒC]
    ttm::Vector{T} | "ᵒC"
    # Water holding capacity as fraction of current snow pack [-]
    whc::Vector{T} | "-"
    # Refreezing efficiency constant in refreezing of freewater in snow [-]
    cfr::Vector{T} | "-"
    # Correction factor for precipitation [-]
    rfcf::Vector{T} | "-"
    # Correction factor for snowfall [-]
    sfcf::Vector{T} | "-"
    ## GLACIER
    # Threshold temperature for snowfall above glacier [ᵒC]
    g_tt::Vector{T} | "ᵒC"
    # Degree-day factor [mm ᵒC⁻¹ Δt⁻¹] for glacier
    g_cfmax::Vector{T} | "mm ᵒC-1 Δt-1"
    # Fraction of the snowpack on top of the glacier converted into ice [-]
    g_sifrac::Vector{T} | "-"
    # Water within the glacier [mm]
    glacierstore::Vector{T} | "mm"
    # Fraction covered by a glacier [-]
    glacierfrac::Vector{T} | "-"
    ##INTERCEPTION
    # Maximum interception storage (in forested and non-forested areas) [mm]
    imax::Vector{SVector{N,T}} | "mm"
    # Evap correction [-]
    ecorr::Vector{T} | "-"
    ##HORTON
    # Maximum storage capacity in the hortonian ponding storage [mm]
    shmax::Vector{SVector{N,T}} | "mm"
    #recession coefficient of the hortonian runoff storage [Δt-1]
    khf::Vector{SVector{N,T}} | "Δt-1"
    #maximum modelled accumulated frost resulting in shmin [ᵒC Δt]
    facc0::Vector{SVector{N,T}} | "ᵒC"
    #minimum modelled accumulated frost resulting in shmax [ᵒC Δt]
    facc1::Vector{SVector{N,T}} | "ᵒC"
    #exponent for the decline of infiltration capacity [-]
    fdec::Vector{SVector{N,T}} | "-"
    #maximum infiltration capacity from horton ponding [mm Δt-1]
    fmax::Vector{SVector{N,T}} | "mm Δt-1"
    #minimum storage capacity in horton ponding (relative to shmax) [-]
    shmin::Vector{SVector{N,T}} | "-"
    #melt coefficient for melt of frozen topsoil [-]
    kmf::Vector{SVector{N,T}} | "-"
    ##ROOTZONE
    # maximum root-zone storage capacity [mm]
    srmax::Vector{SVector{N,T}} | "mm"
    # Fraction of root zone storage below which actual evaporation is potential evaporation [-]
    lp::Vector{SVector{N,T}} | "-"
    # Exponent in soil runoff generation equation [-]
    beta::Vector{SVector{N,T}} | "-"
    # maximum percolation rate [mm Δt⁻¹]
    perc::Vector{SVector{N,T}} | "mm Δt-1"
    # maximum capillary rise rate [mm Δt⁻¹]
    cap::Vector{SVector{N,T}} | "mm Δt-1"
    #FAST
    # Exponent for non linear recession [-]
    alfa::Vector{SVector{N,T}} | "-"
    #recession coefficient of fast storage [Δt-1]
    kf::Vector{SVector{N,T}} | "Δt-1"
    # fraction of qrootzone to slowstorage (1-ds to faststorage)
    ds::Vector{SVector{N,T}} | "-"
    # SLOW
    #recession coefficient of slow storage [Δt-1]
    ks::Vector{T} | "Δt-1"

    ## STATES
    ##SNOW
    # Snow water equivalent [mm]
    snow::Vector{T} | "mm"
    # Liquid water content in the snow pack [mm]
    snowwater::Vector{T} | "mm"
    # Interception storage [mm]
    interceptionstorage::Vector{SVector{N,T}} | "mm"
    # Storage in the hortonian ponding reservoir [mm]
    hortonpondingstorage::Vector{SVector{N,T}} | "mm"
    # Storage in the hortonian runoff generation [mm]
    hortonrunoffstorage::Vector{SVector{N,T}} | "mm"
    # Storage in the root-zone [mm]
    rootzonestorage::Vector{SVector{N,T}} | "mm"
    # Storage in the root-zone relative to maximum root-zone storage capacity [mm]
    srootzone_over_srmax::Vector{SVector{N,T}} | "mm"
    # Storage in the fast store [mm]
    faststorage::Vector{SVector{N,T}} | "mm"
    # Storage in the slow reservoir (for qcapillary calc) [mm]
    slowstorage::Vector{T} | "mm"
    #states previous time step to calc water balance [mm]
    states_::Vector{SVector{N,T}} | "mm"
    #states previous time step to calc water balance combined based on perc class. [mm]
    states_m::Vector{T} | "mm"
    #states averaged over classes
    interceptionstorage_m::Vector{T} | "mm"
    hortonpondingstorage_m::Vector{T} | "mm"
    hortonrunoffstorage_m::Vector{T} | "mm"
    srootzone_m::Vector{T} | "mm"
    faststorage_m::Vector{T} | "mm"
    srootzone_over_srmax_m::Vector{T} | "mm"

    ## FLUXES
    #SNOW
    # Precipitation [mm Δt⁻¹]
    precipitation::Vector{T} | "mm Δt-1"
    # Temperature [ᵒC]
    temperature::Vector{T} | "ᵒC"
    # Potential evapotranspiration [mm Δt⁻¹]
    potential_evaporation::Vector{T} | "mm Δt-1"
    # Potential evapotranspiration corrected [mm Δt⁻¹]
    epotcorr::Vector{T} | "mm Δt-1"
    # Precipitation corrected [mm Δt⁻¹]
    precipcorr::Vector{T} | "mm Δt-1"
    # Snow melt + precipitation as rainfall [mm]
    rainfallplusmelt::Vector{T} | "mm Δt-1"
    # Snowfall [mm]
    snowfall::Vector{T} | "mm Δt-1"
    # Snowmelt [mm]
    snowmelt::Vector{T} | "mm Δt-1"
    #INTERCEPTION
    # Potential soil evaporation [mm Δt⁻¹]
    potsoilevap::Vector{SVector{N,T}} | "mm Δt-1"
    # Evaporation from interception storage [mm Δt⁻¹]
    intevap::Vector{SVector{N,T}} | "mm Δt-1"
    #effective precipitation [mm Δt⁻¹]
    precipeffective::Vector{SVector{N,T}} | "mm Δt-1"
    # Evaporation from interception sum classes [mm Δt⁻¹]
    intevap_m::Vector{T} | "mm Δt-1"
    #HORTONPONDING
    # Evaporation from the hortonion ponding storage [-]
    hortonevap::Vector{SVector{N,T}} | "mm Δt-1"
    # Flux from the hortonian ponding storage to the hortonian runoff storage [mm Δt⁻¹]
    qhortonpond::Vector{SVector{N,T}} | "mm Δt-1"
    # Flux from the hortonian ponding storage to the root zone storage [mm Δt⁻¹]
    qhortonrootzone::Vector{SVector{N,T}} | "mm Δt-1"
    # modeled accumulated frost [ᵒC Δt]
    facc::Vector{SVector{N,T}} | "ᵒC Δt"
    # Evaporation from the hortonian sum classes [mm Δt⁻¹]
    hortonevap_m::Vector{T} | "mm Δt-1"
    #HORTONRUNOFF
    # Flux from the hortonian runoff storage [mm Δt⁻¹]
    qhortonrun::Vector{SVector{N,T}} | "mm Δt-1"
    #ROOTZONE
    # Evaporation from the root-zone storage [mm Δt⁻¹]
    rootevap::Vector{SVector{N,T}} | "mm Δt-1"
    # Flux from the root-zone storage [mm Δt⁻¹]
    qrootzone::Vector{SVector{N,T}} | "mm Δt-1"
    # Pref. recharge to fast storage [mm Δt⁻¹]
    qrootzonefast::Vector{SVector{N,T}} | "mm Δt-1"
    # Pref. recharge to slow storage sum classes [mm Δt⁻¹]
    qrootzoneslow_m::Vector{T} | "mm Δt-1"
    # Capillary flux from the slow to the root-zone storage [mm Δt⁻¹]
    qcapillary::Vector{SVector{N,T}} | "mm Δt-1"
    # Capillary flux from the slow to the root-zone storage sum classes [mm Δt⁻¹]
    qcapillary_m::Vector{T} | "mm Δt-1"
    # Percolation flux from the root-zone to the slow storage [mm Δt⁻¹]
    qpercolation::Vector{SVector{N,T}} | "mm Δt-1"
    # Percolation flux from the root-zone to the slow storage sum classes [mm Δt⁻¹]
    qpercolation_m::Vector{T} | "mm Δt-1"
    # Evaporation from the root-zone storage, interception and hortonian [mm Δt⁻¹]
    actevap::Vector{SVector{N,T}} | "mm Δt-1"
    # Evaporation from the root-zone storage, interception and hortonian sum classes [mm Δt⁻¹]
    actevap_m::Vector{T} | "mm Δt-1"
    # Evaporation from the root-zone storage sum classes [mm Δt⁻¹]
    rootevap_m::Vector{T} | "mm Δt-1"
    #FAST
    # runoff from fast reservoir [mm Δt⁻¹]
    qfast::Vector{SVector{N,T}} | "mm Δt-1"
    #SLOW
    # runoff from slow reservoir [mm Δt⁻¹]
    qslow::Vector{T} | "mm Δt-1"
    #Total [mm Δt⁻¹]
    runoff::Vector{T} | "mm Δt-1"
    # fast runoff sum classes [mm Δt⁻¹]
    qfast_tot = Vector{T} | "mm Δt-1"

    ## WATERBALANCES
    #water balance snow store
    wb_snow::Vector{T} | "mm Δt-1"
    #water balance interception storage [mm Δt⁻¹]
    wb_interception::Vector{SVector{N,T}} | "mm Δt-1"
    #water balance hortonian ponding storage [mm Δt⁻¹]
    wb_hortonponding::Vector{SVector{N,T}} | "mm Δt-1"
    #water balance hortonian runoff storage [mm Δt⁻¹]
    wb_hortonrunoff::Vector{SVector{N,T}} | "mm Δt-1"
    #water balance root-zone storage [mm Δt⁻¹]
    wb_rootzone::Vector{SVector{N,T}} | "mm Δt-1"
    #water balance fast storage [mm Δt⁻¹]
    wb_fast::Vector{SVector{N,T}} | "mm Δt-1"
    #water balance slow storage [mm Δt⁻¹]
    wb_slow::Vector{T} | "mm Δt-1"
    #total water balance [mm Δt⁻¹]
    wb_tot::Vector{T} | "mm Δt-1"

end

statevars(::FLEXTOPO) = (:snow, :snowwater, :interceptionstorage, :hortonpondingstorage, :hortonrunoffstorage, :rootzonestorage, :faststorage, :slowstorage)

function common_snow_hbv(flextopo::FLEXTOPO)
    for i = 1:flextopo.n
        # precip correction
        precipcorr = flextopo.precipitation[i] * flextopo.pcorr[i]

        #hbv snow
        snow, snowwater, snowmelt, rainfallplusmelt, snowfall = snowpack_hbv(
            flextopo.snow[i],
            flextopo.snowwater[i],
            precipcorr,
            flextopo.temperature[i],
            flextopo.tti[i],
            flextopo.tt[i],
            flextopo.ttm[i],
            flextopo.cfmax[i],
            flextopo.whc[i],
        )

        #wb
        wb_snow = precipcorr - rainfallplusmelt - snow + flextopo.snow[i] - snowwater + flextopo.snowwater[i]

        #update stores
        # states_ is sum of states of previous time steps to compute WB at the end (per class); states_m is for combined stores
        flextopo.states_m[i] = flextopo.snow[i] + flextopo.snowwater[i]
        flextopo.snowfall[i] = snowfall
        flextopo.snowmelt[i] = snowmelt
        flextopo.snow[i] = snow
        flextopo.snowwater[i] = snowwater
        flextopo.rainfallplusmelt[i] = rainfallplusmelt
        flextopo.precipcorr[i] = precipcorr
        flextopo.wb_snow[i] = wb_snow
    end
end

function common_snow_no_storage(flextopo::FLEXTOPO)
    for i = 1:flextopo.n
        snow = 0.0
        snowwater = 0.0
        snowfall = 0.0
        snowmelt = 0.0
        # precip correction
        precipcorr = flextopo.precipitation[i] * flextopo.pcorr[i]
        rainfallplusmelt = precipcorr

        wb_snow = precipcorr - rainfallplusmelt - snow + flextopo.snow[i]

        #update stores
        # states_ is sum of states of previous time steps to compute WB at the end (per class); states_m is for combined stores
        flextopo.states_m[i] = flextopo.snow[i] + flextopo.snowwater[i]
        flextopo.snowfall[i] = snowfall
        flextopo.snowmelt[i] = snowmelt
        flextopo.snow[i] = snow
        flextopo.snowwater[i] = snowwater
        flextopo.rainfallplusmelt[i] = rainfallplusmelt
        flextopo.wb_snow[i] = wb_snow
        flextopo.precipcorr[i] = precipcorr
    end
end

function common_glaciers(flextopo::FLEXTOPO, config)
    modelglacier = get(config.model, "glacier", false)::Bool
    for i = 1:flextopo.n
        if modelglacier
            # Run Glacier module and add the snowpack on-top of it.
            # Estimate the fraction of snow turned into ice (HBV-light).
            # Estimate glacier melt.

            flextopo.snow[i], snow2glacier, flextopo.glacierstore[i], glaciermelt =
                glacier_hbv(
                    flextopo.glacierfrac[i],
                    flextopo.glacierstore[i],
                    flextopo.snow[i],
                    flextopo.temperature[i],
                    flextopo.g_tt[i],
                    flextopo.g_cfmax[i],
                    flextopo.g_sifrac[i],
                    Second(flextopo.Δt),
                )
            # Convert to mm per grid cell and add to snowmelt
            glaciermelt = glaciermelt * flextopo.glacierfrac[i]
            rainfallplusmelt = flextopo.rainfallplusmelt[i] + glaciermelt
        else
            rainfallplusmelt = flextopo.rainfallplusmelt[i]
        end
    end
end

function interception_no_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n
        intevap = 0.0
        precipeffective = max(flextopo.rainfallplusmelt[i], 0.0)
        interceptionstorage = 0.0

        # correction for potential evaporation
        epotcorr = flextopo.ecorr[i] * flextopo.potential_evaporation[i]
        restevap = max(0.0, epotcorr - intevap)

        #wb interceptionstore
        wb_interception = flextopo.rainfallplusmelt[i] - intevap - precipeffective - interceptionstorage + flextopo.interceptionstorage[i][k]

        #update stores
        flextopo.states_[i] = setindex(flextopo.states_[i], flextopo.interceptionstorage[i][k], k)
        flextopo.interceptionstorage[i] = setindex(flextopo.interceptionstorage[i], interceptionstorage, k)
        flextopo.intevap[i] = setindex(flextopo.intevap[i], intevap, k)
        flextopo.precipeffective[i] = setindex(flextopo.precipeffective[i], precipeffective, k)
        flextopo.potsoilevap[i] = setindex(flextopo.potsoilevap[i], restevap, k)
        flextopo.epotcorr[i] = epotcorr
        flextopo.wb_interception[i] = setindex(flextopo.wb_interception[i], wb_interception, k)

        #average storage over classes
        flextopo.interceptionstorage_m[i] = sum(flextopo.interceptionstorage[i] .* flextopo.hrufrac[i])
        flextopo.intevap_m[i] = sum(flextopo.intevap[i] .* flextopo.hrufrac[i])

    end
end

function interception_overflow(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n
        # rainfall added to interception store interceptionstorage
        interception =
            min(flextopo.rainfallplusmelt[i], flextopo.imax[i][k] - flextopo.interceptionstorage[i][k])
        # current interception storage
        interceptionstorage = flextopo.interceptionstorage[i][k] + interception
        precipeffective = flextopo.rainfallplusmelt[i] - interception

        # correction for potential evaporation
        epotcorr = flextopo.ecorr[i] * flextopo.potential_evaporation[i]

        # evaporation from interception storage
        intevap = min(interceptionstorage, epotcorr)
        interceptionstorage = interceptionstorage - intevap
        restevap = max(0.0, epotcorr - intevap)

        #wb interceptionstore
        wb_interception = flextopo.rainfallplusmelt[i] - intevap - precipeffective - interceptionstorage + flextopo.interceptionstorage[i][k]

        #update stores
        flextopo.states_[i] = setindex(flextopo.states_[i], flextopo.interceptionstorage[i][k], k)
        flextopo.interceptionstorage[i] = setindex(flextopo.interceptionstorage[i], interceptionstorage, k)
        flextopo.intevap[i] = setindex(flextopo.intevap[i], intevap, k)
        flextopo.precipeffective[i] = setindex(flextopo.precipeffective[i], precipeffective, k)
        flextopo.potsoilevap[i] = setindex(flextopo.potsoilevap[i], restevap, k)
        flextopo.epotcorr[i] = epotcorr
        flextopo.wb_interception[i] = setindex(flextopo.wb_interception[i], wb_interception, k)

        #average storage over classes
        flextopo.interceptionstorage_m[i] = sum(flextopo.interceptionstorage[i] .* flextopo.hrufrac[i])
        flextopo.intevap_m[i] = sum(flextopo.intevap[i] .* flextopo.hrufrac[i])

    end
end

function hortonponding_no_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n
        hortonevap = 0.0
        qhortonpond = 0.0
        qhortonrootzone = max(flextopo.precipeffective[i][k], 0.0)
        hortonpondingstorage = 0.0

        wb_hortonponding = flextopo.precipeffective[i][k] - hortonevap - qhortonpond - qhortonrootzone - hortonpondingstorage + flextopo.hortonpondingstorage[i][k]

        #update with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.hortonpondingstorage[i][k], k)
        flextopo.qhortonpond[i] = setindex(flextopo.qhortonpond[i], qhortonpond, k)
        flextopo.qhortonrootzone[i] = setindex(flextopo.qhortonrootzone[i], qhortonrootzone, k)
        flextopo.hortonpondingstorage[i] = setindex(flextopo.hortonpondingstorage[i], hortonpondingstorage, k)
        flextopo.hortonevap[i] = setindex(flextopo.hortonevap[i], hortonevap, k)
        flextopo.wb_hortonponding[i] = setindex(flextopo.wb_hortonponding[i], wb_hortonponding, k)

        #average storage over classes
        flextopo.hortonpondingstorage_m[i] = sum(flextopo.hortonpondingstorage[i] .* flextopo.hrufrac[i])
        flextopo.hortonevap_m[i] = sum(flextopo.hortonevap[i] .* flextopo.hrufrac[i])
    end
end

function hortonponding(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n

        #calculate reduction of shmax due to frost
        facc = min(flextopo.facc[i][k] + flextopo.temperature[i] * flextopo.kmf[i][k], 0.0)
        ft = min(
            max(
                facc / (flextopo.facc1[i][k] - flextopo.facc0[i][k]) -
                flextopo.facc0[i][k] / (flextopo.facc1[i][k] - flextopo.facc0[i][k]),
                flextopo.shmin[i][k],
            ),
            1.0,
        )
        shmax_frost = ft * flextopo.shmax[i][k]

        #add effective precipitation from interception to the horton ponding storage.
        hortonpondingstorage = flextopo.hortonpondingstorage[i][k] + flextopo.precipeffective[i][k]
        # if soil is filled until max capacity, additional water runs of directly
        directrunoff_h = max(hortonpondingstorage - shmax_frost, 0.0)
        #update hortonpondingstorage
        hortonpondingstorage = hortonpondingstorage - directrunoff_h
        #net water which infiltrates in horton (careful, shmax_frost can decrease from one timestep to another)
        netin_hortonpond = max(flextopo.precipeffective[i][k] - directrunoff_h, 0.0)

        #evaporation from the horton ponding storage
        hortonevap =
            hortonpondingstorage > shmax_frost * flextopo.lp[i][k] ? min(hortonpondingstorage, flextopo.potsoilevap[i][k]) :
            min(flextopo.potsoilevap[i][k] * (hortonpondingstorage / (shmax_frost * flextopo.lp[i][k])))
        #update storage
        hortonpondingstorage = hortonpondingstorage - hortonevap
        #update restevap
        restevap = max(0.0, flextopo.potsoilevap[i][k] - hortonevap)

        #excess water from beta function of netin_hortonpond (due to rouding hortonpondingstorage could be slightly larger than shmax_frost, make sure it is always less than 1)
        qhorton_in = netin_hortonpond * (1.0 - pow((1.0 - min(hortonpondingstorage / shmax_frost, 1.0)), flextopo.beta[i][k]))
        #update storage
        hortonpondingstorage = hortonpondingstorage - qhorton_in

        #total water out of horton ponding consists of directrunoff and excess water from beta function
        qhortonpond = directrunoff_h + qhorton_in

        #infiltration to root-zone
        qhortonrootzone =
            min(hortonpondingstorage / shmax_frost, 1.0) > 0.0 ?
            min(
                flextopo.fmax[i][k] *
                exp(-flextopo.fdec[i][k] * (1.0 - min(hortonpondingstorage / shmax_frost, 1.0))),
                hortonpondingstorage,
            ) : 0.0
        hortonpondingstorage = hortonpondingstorage - qhortonrootzone

        #water balance
        wb_hortonponding = flextopo.precipeffective[i][k] - hortonevap - qhortonpond - qhortonrootzone - hortonpondingstorage + flextopo.hortonpondingstorage[i][k]

        #update states
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.hortonpondingstorage[i][k], k)
        flextopo.potsoilevap[i] = setindex(flextopo.potsoilevap[i], restevap, k)
        flextopo.qhortonpond[i] = setindex(flextopo.qhortonpond[i], qhortonpond, k)
        flextopo.qhortonrootzone[i] = setindex(flextopo.qhortonrootzone[i], qhortonrootzone, k)
        flextopo.hortonpondingstorage[i] = setindex(flextopo.hortonpondingstorage[i], hortonpondingstorage, k)
        flextopo.hortonevap[i] = setindex(flextopo.hortonevap[i], hortonevap, k)
        flextopo.wb_hortonponding[i] = setindex(flextopo.wb_hortonponding[i], wb_hortonponding, k)

        #average storage over classes
        flextopo.hortonpondingstorage_m[i] = sum(flextopo.hortonpondingstorage[i] .* flextopo.hrufrac[i])
        flextopo.hortonevap_m[i] = sum(flextopo.hortonevap[i] .* flextopo.hrufrac[i])
    end
end

function hortonrunoff_no_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n
        qhortonrun = 0.0
        hortonrunoffstorage = 0.0

        wb_hortonrunoff = flextopo.qhortonpond[i][k] - qhortonrun - hortonrunoffstorage + flextopo.hortonrunoffstorage[i][k]

        #update with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.hortonrunoffstorage[i][k], k)
        flextopo.qhortonrun[i] = setindex(flextopo.qhortonrun[i], qhortonrun, k)
        flextopo.hortonrunoffstorage[i] = setindex(flextopo.hortonrunoffstorage[i], hortonrunoffstorage, k)
        flextopo.wb_hortonrunoff[i] = setindex(flextopo.wb_hortonrunoff[i], wb_hortonrunoff, k)

        #average storage over classes
        flextopo.hortonrunoffstorage_m[i] = sum(flextopo.hortonrunoffstorage[i] .* flextopo.hrufrac[i])
    end
end

function hortonrunoff(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n

        qhortonrun = min(flextopo.hortonrunoffstorage[i][k], flextopo.hortonrunoffstorage[i][k] * flextopo.khf[i][k])
        hortonrunoffstorage = flextopo.hortonrunoffstorage[i][k] + flextopo.qhortonpond[i][k] - qhortonrun

        wb_hortonrunoff = flextopo.qhortonpond[i][k] - qhortonrun - hortonrunoffstorage + flextopo.hortonrunoffstorage[i][k]

        #update with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.hortonrunoffstorage[i][k], k)
        flextopo.qhortonrun[i] = setindex(flextopo.qhortonrun[i], qhortonrun, k)
        flextopo.hortonrunoffstorage[i] = setindex(flextopo.hortonrunoffstorage[i], hortonrunoffstorage, k)
        flextopo.wb_hortonrunoff[i] = setindex(flextopo.wb_hortonrunoff[i], wb_hortonrunoff, k)

        #average storage over classes
        flextopo.hortonrunoffstorage_m[i] = sum(flextopo.hortonrunoffstorage[i] .* flextopo.hrufrac[i])

    end
end

function rootzone_no_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n
        rootevap = 0.0
        qpercolation = 0.0
        qcapillary = 0.0
        qrootzone = max(flextopo.qhortonrootzone[i][k] + flextopo.rootzonestorage[i][k], 0.0) #if store not empty initial conditions
        rootzonestorage = 0.0

        #compute water balance rootzonestorage storage
        wb_rootzone = flextopo.qhortonrootzone[i][k] - rootevap - qrootzone - qpercolation + qcapillary - rootzonestorage + flextopo.rootzonestorage[i][k]

        #update states and fluxes with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.rootzonestorage[i][k], k)
        flextopo.qrootzone[i] = setindex(flextopo.qrootzone[i], qrootzone, k)
        flextopo.rootzonestorage[i] = setindex(flextopo.rootzonestorage[i], rootzonestorage, k)
        flextopo.srootzone_over_srmax[i] =
            setindex(flextopo.srootzone_over_srmax[i], flextopo.rootzonestorage[i][k] / flextopo.srmax[i][k], k)
        flextopo.qcapillary[i] = setindex(flextopo.qcapillary[i], qcapillary, k)
        flextopo.qpercolation[i] = setindex(flextopo.qpercolation[i], qpercolation, k)
        flextopo.rootevap[i] = setindex(flextopo.rootevap[i], rootevap, k)
        flextopo.actevap[i] = setindex(
            flextopo.actevap[i],
            flextopo.rootevap[i][k] + flextopo.intevap[i][k] + flextopo.hortonevap[i][k],
            k,
        )
        flextopo.wb_rootzone[i] = setindex(flextopo.wb_rootzone[i], wb_rootzone, k)

        #average storage over classes
        flextopo.srootzone_m[i] = sum(flextopo.rootzonestorage[i] .* flextopo.hrufrac[i])
        flextopo.srootzone_over_srmax_m[i] = sum(flextopo.srootzone_over_srmax[i] .* flextopo.hrufrac[i])
        flextopo.rootevap_m[i] = sum(flextopo.rootevap[i] .* flextopo.hrufrac[i])
        flextopo.actevap_m[i] = sum(flextopo.actevap[i] .* flextopo.hrufrac[i])
    end
end

function rootzone_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n

        #added water to the root-zone. NB: if no horton storages: qhortonrootzone is in fact Pe!! (effective precip).
        rootzonestorage = flextopo.rootzonestorage[i][k] + flextopo.qhortonrootzone[i][k]
        # if soil is filled until max capacity, additional water runs of directly
        directrunoff = max(rootzonestorage - flextopo.srmax[i][k], 0.0)
        #update rootzonestorage
        rootzonestorage = rootzonestorage - directrunoff
        #net water which infiltrates in root-zone
        netin_rootzone = flextopo.qhortonrootzone[i][k] - directrunoff

        #evaporation from the rootzone
        rootevap =
            rootzonestorage > flextopo.srmax[i][k] * flextopo.lp[i][k] ?
            min(rootzonestorage, flextopo.potsoilevap[i][k]) :
            min(
                flextopo.potsoilevap[i][k] *
                (rootzonestorage / (flextopo.srmax[i][k] * flextopo.lp[i][k])),
            )
        #update storage
        rootzonestorage = rootzonestorage - rootevap

        #excess water from beta function of netin_rootzone
        qrootzone_in = netin_rootzone * (1.0 - pow(1.0 - rootzonestorage / flextopo.srmax[i][k], flextopo.beta[i][k]))
        #update storage
        rootzonestorage = rootzonestorage - qrootzone_in

        #total water out of root-zone consists of directrunoff and excess water from beta function
        qrootzone = directrunoff + qrootzone_in

        #percolation
        qpercolation = flextopo.perc[i][k] * rootzonestorage / flextopo.srmax[i][k]
        rootzonestorage = rootzonestorage - qpercolation

        #capillary rise
        qcapillary = min(flextopo.cap[i][k] * (1.0 - rootzonestorage / flextopo.srmax[i][k]), flextopo.slowstorage[i])
        rootzonestorage = rootzonestorage + qcapillary

        #compute water balance rootzonestorage storage
        wb_rootzone = flextopo.qhortonrootzone[i][k] - rootevap - qrootzone - qpercolation + qcapillary - rootzonestorage + flextopo.rootzonestorage[i][k]

        #update states and fluxes with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.rootzonestorage[i][k], k)
        flextopo.qrootzone[i] = setindex(flextopo.qrootzone[i], qrootzone, k)
        flextopo.rootzonestorage[i] = setindex(flextopo.rootzonestorage[i], rootzonestorage, k)
        flextopo.srootzone_over_srmax[i] =
            setindex(flextopo.srootzone_over_srmax[i], flextopo.rootzonestorage[i][k] / flextopo.srmax[i][k], k)
        flextopo.qcapillary[i] = setindex(flextopo.qcapillary[i], qcapillary, k)
        flextopo.qpercolation[i] = setindex(flextopo.qpercolation[i], qpercolation, k)
        flextopo.rootevap[i] = setindex(flextopo.rootevap[i], rootevap, k)
        flextopo.actevap[i] = setindex(
            flextopo.actevap[i],
            flextopo.rootevap[i][k] + flextopo.intevap[i][k] + flextopo.hortonevap[i][k],
            k,
        )
        flextopo.wb_rootzone[i] = setindex(flextopo.wb_rootzone[i], wb_rootzone, k)

        #average storage over classes
        flextopo.srootzone_m[i] = sum(flextopo.rootzonestorage[i] .* flextopo.hrufrac[i])
        flextopo.srootzone_over_srmax_m[i] = sum(flextopo.srootzone_over_srmax[i] .* flextopo.hrufrac[i])
        flextopo.rootevap_m[i] = sum(flextopo.rootevap[i] .* flextopo.hrufrac[i])
        flextopo.actevap_m[i] = sum(flextopo.actevap[i] .* flextopo.hrufrac[i])
    end
end

function fast_no_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n
        #make sure ds does not exceed 1
        ds = flextopo.ds[i][k] > 1.0 ? 1.0 : flextopo.ds[i][k]
        #calc inflow to faststorage
        qrootzonefast = flextopo.qrootzone[i][k] * (1.0 - ds)
        qfast = qrootzonefast + flextopo.faststorage[i][k] #if store not empty initial conditions, make sure to empty
        faststorage = 0.0

        #compute wb faststorage
        wb_fast = qrootzonefast - qfast - faststorage + flextopo.faststorage[i][k]

        #update with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.faststorage[i][k], k)
        flextopo.qrootzonefast[i] = setindex(flextopo.qrootzonefast[i], qrootzonefast, k)
        flextopo.qfast[i] = setindex(flextopo.qfast[i], qfast, k)
        flextopo.faststorage[i] = setindex(flextopo.faststorage[i], faststorage, k)
        flextopo.wb_fast[i] = setindex(flextopo.wb_fast[i], wb_fast, k)

        #average storage over classes
        flextopo.faststorage_m[i] = sum(flextopo.faststorage[i] .* flextopo.hrufrac[i])
    end
end

function fast_storage(flextopo::FLEXTOPO)
    k = flextopo.kclass[1]
    for i = 1:flextopo.n

        #make sure ds does not exceed 1
        ds = flextopo.ds[i][k] > 1.0 ? 1.0 : flextopo.ds[i][k]

        #split part of the outflow from the root-zone to the fast runoff (and part as preferential recharge to the slow reservoir)
        qrootzonefast = flextopo.qrootzone[i][k] * (1.0 - ds)

        #fast runoff
        qfast = min(
            flextopo.faststorage[i][k],
            pow(flextopo.faststorage[i][k], flextopo.alfa[i][k]) * flextopo.kf[i][k],
        )
        #update store
        faststorage = flextopo.faststorage[i][k] + qrootzonefast - qfast

        #compute wb faststorage
        wb_fast = qrootzonefast - qfast - faststorage + flextopo.faststorage[i][k]

        #update with setindex
        flextopo.states_[i] =
            setindex(flextopo.states_[i], flextopo.states_[i][k] + flextopo.faststorage[i][k], k)
        flextopo.qrootzonefast[i] = setindex(flextopo.qrootzonefast[i], qrootzonefast, k)
        flextopo.qfast[i] = setindex(flextopo.qfast[i], qfast, k)
        flextopo.faststorage[i] = setindex(flextopo.faststorage[i], faststorage, k)
        flextopo.wb_fast[i] = setindex(flextopo.wb_fast[i], wb_fast, k)

        #average storage over classes
        flextopo.faststorage_m[i] = sum(flextopo.faststorage[i] .* flextopo.hrufrac[i])
    end
end

function slow_no_storage(flextopo::FLEXTOPO)
    for i = 1:flextopo.n
        #make sure ds does not exceed 1
        pref_recharge = flextopo.qrootzone[i] .* min.(flextopo.ds[i], 1.0)
        pref_recharge_sum_classes = sum(pref_recharge .* flextopo.hrufrac[i])

        Qcap_sum_classes = sum(flextopo.qcapillary[i] .* flextopo.hrufrac[i])
        Qperc_sum_classes = sum(flextopo.qpercolation[i] .* flextopo.hrufrac[i])

        Qsin = pref_recharge_sum_classes + Qperc_sum_classes - Qcap_sum_classes
        qslow = Qsin + flextopo.slowstorage[i] # if at start store is not empty
        slowstorage = 0.0

        qfast_tot =
            sum(flextopo.qfast[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.qhortonrun[i] .* flextopo.hrufrac[i])
        runoff = max(0.0, qslow + qfast_tot)

        wb_slow = Qsin - qslow - slowstorage + flextopo.slowstorage[i]

        #update
        flextopo.qfast_tot[i] = qfast_tot
        flextopo.states_m[i] =
            flextopo.states_m[i] +
            sum(flextopo.states_[i] .* flextopo.hrufrac[i]) +
            flextopo.slowstorage[i]
        flextopo.runoff[i] = runoff
        flextopo.qrootzoneslow_m[i] = pref_recharge_sum_classes
        flextopo.qpercolation_m[i] = Qperc_sum_classes
        flextopo.qcapillary_m[i] = Qcap_sum_classes
        flextopo.qslow[i] = qslow
        flextopo.slowstorage[i] = slowstorage
        flextopo.wb_slow[i] = wb_slow
    end
end

function common_slow_storage(flextopo::FLEXTOPO)
    for i = 1:flextopo.n
        #make sure ds does not exceed 1
        pref_recharge = flextopo.qrootzone[i] .* min.(flextopo.ds[i], 1.0)
        pref_recharge_sum_classes = sum(pref_recharge .* flextopo.hrufrac[i])

        Qcap_sum_classes = sum(flextopo.qcapillary[i] .* flextopo.hrufrac[i])
        Qperc_sum_classes = sum(flextopo.qpercolation[i] .* flextopo.hrufrac[i])

        Qsin = pref_recharge_sum_classes + Qperc_sum_classes - Qcap_sum_classes
        slowstorage = flextopo.slowstorage[i] + Qsin

        qslow = min(flextopo.slowstorage[i], flextopo.slowstorage[i] * flextopo.ks[i])
        slowstorage = slowstorage - qslow

        qfast_tot =
            sum(flextopo.qfast[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.qhortonrun[i] .* flextopo.hrufrac[i])
        runoff = max(0.0, qslow + qfast_tot)

        wb_slow = Qsin - qslow - slowstorage + flextopo.slowstorage[i]

        #update
        flextopo.qfast_tot[i] = qfast_tot
        flextopo.states_m[i] =
            flextopo.states_m[i] +
            sum(flextopo.states_[i] .* flextopo.hrufrac[i]) +
            flextopo.slowstorage[i]
        flextopo.runoff[i] = runoff
        flextopo.qrootzoneslow_m[i] = pref_recharge_sum_classes
        flextopo.qpercolation_m[i] = Qperc_sum_classes
        flextopo.qcapillary_m[i] = Qcap_sum_classes
        flextopo.qslow[i] = qslow
        flextopo.slowstorage[i] = slowstorage
        flextopo.wb_slow[i] = wb_slow
    end
end

function watbal(flextopo::FLEXTOPO)
    for i = 1:flextopo.n
        states =
            flextopo.snow[i] +
            flextopo.snowwater[i] +
            sum(flextopo.interceptionstorage[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.hortonpondingstorage[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.hortonrunoffstorage[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.rootzonestorage[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.faststorage[i] .* flextopo.hrufrac[i]) +
            sum(flextopo.slowstorage[i] .* flextopo.hrufrac[i])
        wb_tot =
            flextopo.precipitation[i] - (
                sum(flextopo.intevap[i] .* flextopo.hrufrac[i]) +
                sum(flextopo.hortonevap[i] .* flextopo.hrufrac[i]) +
                sum(flextopo.rootevap[i] .* flextopo.hrufrac[i])
            ) - flextopo.runoff[i] - states + flextopo.states_m[i]
        #update wb
        flextopo.wb_tot[i] = wb_tot
    end
end


text:
using Mmap

function get_laszip_executable_path()
    if Sys.iswindows()
        return joinpath(dirname(@__DIR__), "resources", "laszip.exe")
    elseif Sys.islinux()
        return joinpath(dirname(@__DIR__), "resources", "laszip")
    else
        error("LasIO  with ZIP functionality is only suported for windows and linux!")
    end
end

function get_record_count(header::LasHeader)
    return header.extended_number_of_point_records > 0 ? Int(header.extended_number_of_point_records) : Int(header.records_count)
end

function pointformat(header::LasHeader)
    id = header.data_format_id
    if id == 0x00
        return LasPoint0
    elseif id == 0x01
        return LasPoint1
    elseif id == 0x02
        return LasPoint2
    elseif id == 0x03
        return LasPoint3
    elseif id == 0x04
        return LasPoint4
    elseif id == 0x05
        return LasPoint5
    elseif id == 0x06
        return LasPoint6
    elseif id == 0x07
        return LasPoint7
    elseif id == 0x08
        return LasPoint8
    elseif id == 0x09
        return LasPoint9
    elseif id == 0x0a
        return LasPoint10
    else
        error("unsupported point format $(Int(id))")
    end
end

# skip the LAS file's magic four bytes, "LASF"
skiplasf(s::Union{Stream{format"LAS"}, Stream{format"LAZ"}, IO}) = read(s, UInt32)

function load(f::File{format"LAS"}; mmap=false)
    open(f) do s
        load(s; mmap=mmap)
    end
end

# Load pipe separately since it can't be memory mapped
function load(s::Base.AbstractPipe)
    skiplasf(s)
    header = read(s, LasHeader)
    n = get_record_count(header)
    pointtype = pointformat(header)

    @info "Reading $(n) '$(pointtype)' points"
    pointdata = Vector{pointtype}(undef, n)
    for i=1:n
        i%1000000 == 0 && @info("Read $(i)/$(n) points")
        pointdata[i] = read(s, pointtype)
    end
    header, pointdata
end

function load(s::Stream{format"LAS"}; mmap=false)
    skiplasf(s)
    header = read(s, LasHeader)
    n = get_record_count(header)
    pointtype = pointformat(header)

    @info "Reading $(n) '$(pointtype)' points"
    if mmap
        pointsize = Int(header.data_record_length)
        pointbytes = Mmap.mmap(s.io, Vector{UInt8}, n*pointsize, position(s))
        pointdata = PointVector{pointtype}(pointbytes, pointsize)
    else
        pointdata = Vector{pointtype}(undef, n)
        for i=1:n
            i%1000000 == 0 && @info("Read $(i)/$(n) points")
            pointdata[i] = read(s, pointtype)
        end
    end

    header, pointdata
end

function load(f::File{format"LAZ"})
    # read las from laszip, which decompresses to stdout
    open(`$(get_laszip_executable_path()) -olas -stdout -i $(filename(f))`) do s
        h,p = load(s)
        read(s)
        return h, p
    end
end

function read_header(f::AbstractString)
    open(f) do s
        skiplasf(s)
        read(s, LasHeader)
    end
end

function read_header(s::IO)
    skiplasf(s)
    read(s, LasHeader)
end

function save(f::File{format"LAS"}, header::LasHeader, pointdata::AbstractVector{<:LasPoint})
    open(f, "w") do s
        save(s, header, pointdata)
    end
end

function save(s::Stream{format"LAS"}, header::LasHeader, pointdata::AbstractVector{<:LasPoint})
    # checks
    header_n = get_record_count(header)
    n = length(pointdata)
    msg = "Number of records in header ($header_n) does not match data length ($n)"
    @info "Writing $(n) '$(typeof(pointdata[1]))' points"

    @assert header_n == n msg

    # write header
    write(s, magic(format"LAS"))
    write(s, header)

    # write points
    for p in pointdata
        write(s, p)
    end
end

function save(f::File{format"LAZ"}, header::LasHeader, pointdata::AbstractVector{<:LasPoint})
    # pipes las to laszip to write laz
    open(`$(get_laszip_executable_path()) -olaz -stdin -o $(filename(f))`, "w") do s
        savebuf(s, header, pointdata)
    end
end

# Uses a buffered write to the stream.
# For saving to LAS this does not increase speed,
# but it speeds up a lot when the result is piped to laszip.
function savebuf(s::IO, header::LasHeader, pointdata::AbstractVector{<:LasPoint})
    # checks
    header_n = get_record_count(header)
    n = length(pointdata)
    msg = "number of records in header ($header_n) does not match data length ($n)"
    @assert header_n == n msg
    @info "Writing $(n) '$(typeof(pointdata[1]))' points to LAZ file"

    # write header
    write(s, magic(format"LAS"))
    write(s, header)

    # 2048 points seemed to be an optimum for the libLAS_1.2.las testfile
    npoints_buffered = 2048
    bufsize = header.data_record_length * npoints_buffered
    buf = IOBuffer(sizehint=bufsize)
    # write points
    for (i, p) in enumerate(pointdata)
        write(buf, p)
        if rem(i, npoints_buffered) == 0 || i == n
            write(s, take!(buf))
        end
    end
end


text:
# test all possibilities of readInputData
# Claudia April 2015
# added new name for readInputData function
# Claudia August 2015

include("../src/types.jl")
include("../src/functions.jl")

d=readInputData("1.ms"); #tableCF0.txt
d=readInputData("1.ms",:rand,10); #tableCF3.txt
d=readInputData("1.ms",[1,2,3,4,5]); #tableCF4.txt

d=readInputData("1.ms","allQuartets.txt"); #tableCF1.txt
d=readInputData("1.ms","allQuartets.txt",:rand,10); #tableCF2.txt
d=readInputData("1.ms","allQuartets.txt",true,"try4.txt"); #try4.txt

descData(d)

d = readTrees2CF("1.ms");
d = readTrees2CF("1.ms",filename="try.txt");
d = readTrees2CF("1.ms","allQuartets.txt",filename="hola.txt");
d = readTrees2CF("1.ms","allQuartets.txt",whichQ=:rand,numQ=10,filename="hola.txt");

d= readTableCF("tableCF.txt");
d2= readTableCF("tableCFbad.txt")


text:
using Plots
include(joinpath(@__DIR__, "pendulum_minimum_time.jl"))

# TO solution
X̄, Ū = unpack(Z̄, prob)

@show sum([Ū[t][end] for t = 1:T-1])
t_nominal = [0.0, [sum([Ū[i][end] for i = 1:t]) for t = 1:T-1]...]
plot(t_nominal, hcat(X̄...)', width = 2.0)
plot(t_nominal[1:end-1], hcat(Ū...)[1:1, :]',
	linetype = :steppost, width = 2.0, color = :purple)

# DPO solution
X, U = unpack(Z, prob_dpo.prob.prob.nom)

@show sum([U[t][end] for t = 1:T-1])
t_dpo = [0.0, [sum([U[i][end] for i = 1:t]) for t = 1:T-1]...]
plot(t_dpo, hcat(X...)', width = 2.0)
plot(t_dpo[1:end-1], hcat(U...)[1:1, :]',
	linetype = :steppost, width = 2.0, color = :orange)

max_stars_repo_name: zacmanchester/DirectMotionPlanning

max_issues_repo_name: zacmanchester/DirectMotionPlanning

max_forks_repo_name: zacmanchester/DirectMotionPlanning


text:
using BenchmarkTools, RiskAdjustedLinearizations, MATLAB
include(joinpath(dirname(@__FILE__), "..", "rbc_cc", "rbc_cc.jl"))

# Settings: what do you want to do?
autodiff = false
time_matlab = true

# Set up
autodiff_method = autodiff ? :forward : :central
m_rbc_cc = RBCCampbellCochraneHabits()
m = rbc_cc(m_rbc_cc, 0)
z0 = copy(m.z)
y0 = copy(m.y)
Ψ0 = copy(m.Ψ)

# Use deterministic steady state as guesses
solve!(m, z0, y0; algorithm = :deterministic, autodiff = autodiff_method, verbose = :none)
zdet = copy(m.z)
ydet = copy(m.y)
Ψdet = copy(m.Ψ)

if time_matlab
    println("Relaxation algorithm in MATLAB")
    mat"""
    genaffine_rbc_cc_relaxation;
    """
end

println("Relaxation algorithm in Julia")
@btime begin # called the "iterative" method in the original paper
    solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, autodiff = autodiff_method, verbose = :none)
end

println("Relaxation algorithm with Anderson acceleration")
@btime begin
    solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, use_anderson = true, m = 3, autodiff = autodiff_method, verbose = :none)
end

sparsity, colorvec = compute_sparsity_pattern(m, :relaxation; sparsity_detection = false)
jac_cache = preallocate_jac_cache(m, :relaxation; sparsity_detection = false)

println("Relaxation with matrix coloring of sparse Jacobians")
@btime begin
    solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, autodiff = autodiff_method,
           sparse_jacobian = true, sparsity = sparsity, colorvec = colorvec, verbose = :none)
end

println("Relaxation with matrix coloring of sparse Jacobians and caching")
@btime begin
    solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, autodiff = autodiff_method,
           sparse_jacobian = true, jac_cache = jac_cache, verbose = :none)
end

if time_matlab
    println("Homotopy algorithm in MATLAB")
    mat"""
    genaffine_rbc_cc_homotopy;
    """
end

println("Homotopy algorithm in Julia")
@btime begin # called the "continuation" method in the original paper, but is called homotopy in the original code
    solve!(m, zdet, ydet, Ψdet; algorithm = :homotopy, autodiff = autodiff_method, verbose = :none)
end

sparsity, colorvec = compute_sparsity_pattern(m, :homotopy; sparsity_detection = false)
jac_cache = preallocate_jac_cache(m, :homotopy; sparsity_detection = false)

println("Homotopy with matrix coloring of sparse Jacobians")
@btime begin
    solve!(m, zdet, ydet, Ψdet; algorithm = :homotopy, autodiff = autodiff_method,
           sparse_jacobian = true, sparsity = sparsity, colorvec = colorvec, verbose = :none)
end

println("Homotopy with matrix coloring of sparse Jacobians and caching")
@btime begin
    solve!(m, zdet, ydet, Ψdet; algorithm = :homotopy, autodiff = autodiff_method,
           sparse_jacobian = true, jac_cache = jac_cache, verbose = :none)
end


text:
function initialize_pss!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{M, S, A, TG, PSY.PSSFixed}},
    inner_vars::AbstractVector,
) where {M <: PSY.Machine, S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov} end

function initialize_pss!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{M, S, A, TG, PSY.IEEEST}},
    inner_vars::AbstractVector,
) where {M <: PSY.Machine, S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov}
    #Get Signal Input Integer
    pss = PSY.get_pss(dynamic_device)
    input_code = PSY.get_input_code(pss)
    #Remote bus control not supported

    #Get Input Signal
    u = get_pss_input_signal(
        Base.RefValue{input_code},
        device_states,
        inner_vars,
        1.0,
        dynamic_device,
    )

    #Obtain PSS States
    pss_ix = get_local_state_ix(dynamic_device, typeof(pss))
    pss_states = @view device_states[pss_ix]

    # Get Parameters
    A1 = PSY.get_A1(pss)
    A2 = PSY.get_A2(pss)
    A5 = PSY.get_A5(pss)
    A6 = PSY.get_A6(pss)
    T1 = PSY.get_T1(pss)
    T2 = PSY.get_T2(pss)
    T3 = PSY.get_T3(pss)
    T4 = PSY.get_T4(pss)
    T5 = PSY.get_T5(pss)
    T6 = PSY.get_T6(pss)
    Ks = PSY.get_Ks(pss)
    Ls_min, Ls_max = PSY.get_Ls_lim(pss)
    V_cu = PSY.get_Vcu(pss)
    V_cl = PSY.get_Vcl(pss)

    #Error non-valid parameters
    if A6 > eps() && A2 < eps()
        error(
            "A6 cannot be greater than zero if A2 is zero. Correct data in PSS associated with $(PSY.get_name(dynamic_device))",
        )
    end

    if T1 > eps() && T2 < eps()
        error(
            "T2 cannot be greater than zero if T1 is zero. Correct data in PSS associated with $(PSY.get_name(dynamic_device))",
        )
    end

    if T3 > eps() && T4 < eps()
        error(
            "T3 cannot be greater than zero if T4 is zero. Correct data in PSS associated with $(PSY.get_name(dynamic_device))",
        )
    end

    if T5 < eps()
        error(
            "T5 is not allowed to be zero. Correct data in PSS associated with $(PSY.get_name(dynamic_device))",
        )
    end

    #Compute Parameter Ratios
    A6_A2 = A2 < eps() ? 0.0 : (A6 / A2)
    T1_T2 = T2 < eps() ? 0.0 : (T1 / T2)
    T3_T4 = T4 < eps() ? 0.0 : (T3 / T4)
    KsT5_T6 = T6 < eps() ? 0.0 : (Ks * T5 / T6)

    #Compute steady-state values
    x_p1 = 0.0
    x_p2 = u
    x_p3 = 0.0
    x_p4 = u
    y_f = A6_A2 * x_p2 + (A5 - A1 * A6_A2) * x_p3 + (1.0 - A6_A2) * x_p4
    x_p5 = y_f * (1.0 - T1_T2)
    y_LL1 = x_p5 + T1_T2 * y_f
    x_p6 = y_LL1 * (1.0 - T3_T4)
    y_LL2 = x_p6 + T3_T4 * y_LL1
    x_p7 = y_LL2
    y_out = KsT5_T6 * (y_LL2 - x_p7)

    #Compute and update output signal
    V_ss = clamp(y_out, Ls_min, Ls_max)
    #Compute compensated terminal voltage
    V_R = inner_vars[VR_gen_var]
    V_I = inner_vars[VI_gen_var]
    #To do: Figure out how to compensate terminal voltage
    V_ct = sqrt(V_R^2 + V_I^2)

    #Compute PSS output signal
    V_pss = output_pss_limiter(V_ss, V_ct, V_cl, V_cu)

    #Update Inner Vars
    inner_vars[V_pss_var] = V_pss

    #Update States
    pss_states[1] = x_p1
    pss_states[2] = x_p2
    pss_states[3] = x_p3
    pss_states[4] = x_p4
    pss_states[5] = x_p5
    pss_states[6] = x_p6
    pss_states[7] = x_p7
    return
end

#Currently not working properly.
#=
function initialize_pss!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{M, S, A, TG, PSY.PSSSimple},
) where {M <: PSY.Machine, S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov}

    ω0 = get_ω_ref(dynamic_device)
    P0 = PSY.get_active_power(static)
    τe = get_inner_vars(dynamic_device)[τe_var]

    #Get parameters
    pss = PSY.get_pss(dynamic_device)
    K_ω = PSY.get_K_ω(pss)
    K_p = PSY.get_K_p(pss)

    #It assumes that the system is balanced so:
    get_inner_vars(dynamic_device)[V_pss_var] = K_ω * (ω0 - ω0) + K_p * (ω0 * τe - P0)

end
=#


text:
# TODO
# compile Mod and KMod to elementary gates.
export Mod, KMod

"""
    mod_inverse(x::Int, N::Int) -> Int

Return `y` that `(x*y)%N == 1`, notice the `(x*y)%N` operations in Z* forms a group and this is the definition of inverse.
"""
function mod_inverse(x::Int, N::Int)
    for i=1:N
        (x*i)%N == 1 && return i
    end
    throw(ArgumentError("Can not find the inverse, $x is probably not in Z*($N)!"))
end

"""
    Mod{N} <: PrimitiveBlock{N}

calculates `mod(a*x, L)`, notice `gcd(a, L)` should be 1.
"""
struct Mod{N} <: PrimitiveBlock{N}
    a::Int
    L::Int
    function Mod{N}(a, L) where N
        @assert gcd(a, L) == 1 && L<=1<<N
        new{N}(a, L)
    end
end

function _apply!(reg::ArrayReg{B}, m::Mod{N}) where {B, N}
    nstate = zero(reg.state)
    for i in basis(reg)
        _i = i >= m.L ? i+1 : mod(i*m.a, m.L)+1
        for j in 1:B
            @inbounds nstate[_i,j] = reg.state[i+1,j]
        end
    end
    reg.state = nstate
    reg
end

function Yao.mat(::Type{T}, m::Mod{N}) where {T, N}
    perm = Vector{Int}(undef, 1<<N)
    for i in basis(N)
        @inbounds perm[i >= m.L ? i+1 : mod(i*m.a, m.L)+1] = i+1
    end
    PermMatrix(perm, ones(T, 1<<N))
end

Base.adjoint(m::Mod{N}) where N = Mod{N}(mod_inverse(m.a, m.L), m.L)
Yao.print_block(io::IO, m::Mod{N}) where N = print(io, "Mod{$N}: $(m.a)*x % $(m.L)")

Yao.isunitary(::Mod) = true
# Yao.ishermitian(::Mod) = true  # this is not true for L = 1.
# Yao.isreflexive(::Mod) = true  # this is not true for L = 1.

"""
    KMod{N, K} <: PrimitiveBlock{N}

The first `K` qubits are exponent `k`, and the rest `N-K` are base `a`,
it calculates `mod(a^k*x, L)`, notice `gcd(a, L)` should be 1.
"""
struct KMod{N, K} <: PrimitiveBlock{N}
    a::Int
    L::Int
    function KMod{N, K}(a, L) where {N, K}
        @assert gcd(a, L) == 1 && L<=1<<(N-K)
        new{N, K}(a, L)
    end
end

nqubits_data(m::KMod{N, K}) where {N, K} = N-K
nqubits_control(m::KMod{N, K}) where {N, K} = K

function bint2_reader(T, k::Int)
    mask = bmask(T, 1:k)
    return b -> (b&mask, b>>k)
end

function _apply!(reg::ArrayReg{B}, m::KMod{N, K}) where {B, N, K}
    nstate = zero(reg.state)

    reader = bint2_reader(Int, K)
    for b in basis(reg)
        k, i = reader(b)
        _i = i >= m.L ? i : mod(i*powermod(m.a, k, m.L), m.L)
        _b = k + _i<<K + 1
        for j in 1:B
            @inbounds nstate[_b,j] = reg.state[b+1,j]
        end
    end
    reg.state = nstate
    reg
end

function Yao.mat(::Type{T}, m::KMod{N, K}) where {T, N, K}
    perm = Vector{Int}(undef, 1<<N)
    reader = bint2_reader(Int, K)
    for b in basis(N)
        k, i = reader(b)
        _i = i >= m.L ? i : mod(i*powermod(m.a, k, m.L), m.L)
        _b = k + _i<<K + 1
        @inbounds perm[_b] = b+1
    end
    PermMatrix(perm, ones(T, 1<<N))
end

Base.adjoint(m::KMod{N, K}) where {N, K} = KMod{N, K}(mod_inverse(m.a, m.L), m.L)
Yao.print_block(io::IO, m::KMod{N, K}) where {N, K} = print(io, "Mod{$N, $K}: $(m.a)^k*x % $(m.L)")

Yao.isunitary(::KMod) = true
# Yao.ishermitian(::Mod) = true  # this is not true for L = 1.
# Yao.isreflexive(::Mod) = true  # this is not true for L = 1.

max_stars_repo_name: QuantumBFS/YaoExtensions

max_issues_repo_name: QuantumBFS/YaoExtensions

max_forks_repo_name: QuantumBFS/YaoExtensions


text:
@testset "complete substitution" begin
    @polyvar x y z

    p = x^2 * y + y^2 + x
    @test p(x=>1, y=>2) == 7
    @inferred p(x=>1, y=>2)

    subs1(p, x, y) = @allocated p(x=>1, y=>2)
    subs1(p, x, y)
    @test subs1(p, x, y) == 0

    @test (@inferred x(x=>5)) == 5
    @test (@inferred x(y=>2)) == x
    @test (@inferred x(x=>y)) == y
end

@testset "partial substitution" begin
    @polyvar x y z

    p = x^2 * y + y^2 + x
    p2 = @inferred p(x=>4)
    @test length(terms(p2)) == 3
    @test terms(p2)[1].coefficient == 4
    @test degree(terms(p2)[1].monomial) == 0
    @test terms(p2)[2].coefficient == 16
    @test degree(terms(p2)[2].monomial) == 1
    @test exponents(monomial(terms(p2)[2])) == (1,)
    @test terms(p2)[3].coefficient == 1
    @test degree(terms(p2)[3].monomial) == 2
    @test exponents(monomial(terms(p2)[3])) == (2,)
end


text:
# Simple history container used inside the search.
# Only used to keep track of the actions executed.
struct ActionHistory
    history::Vector{Integer}
    action_space_size::Int32
end

function ActionHistory(g::AbstractEnv)
    ActionHistory(
        map(a->a.idx, gethistory(g)),
        action_space_size(g)
    )
end

function ActionHistory(g::AbstractEnv, idx)
    ActionHistory(
        map(a->a.idx, gethistory(g, idx)),
        action_space_size(g)
    )
end

Base.length(x::ActionHistory) = length(x.history)
clone(x::ActionHistory) = ActionHistory(copy(x.history), x.action_space_size)
last_action(x::ActionHistory) = x.history[end]
to_play(x::ActionHistory) = iseven(length(x.history)) ? PLAYER_ID : OPPONENT_ID
#next_to_play(x::ActionHistory) = to_play(x) == PLAYER_ID ? OPPONENT_ID : PLAYER_ID
add_action!(x::ActionHistory, action::Integer) = push!(x.history, action)

# @todo make this generic
legal_actions(config::MuZeroConfig, x::ActionHistory) = legal_actions(config.game_type, x)
legal_actions(game_type::Type{<:AbstractEnv}, x::ActionHistory) = game_type(x.history) |> getlegalactions


text:
import Base: -,*, convert
using LinearAlgebra

export UnitaryGate, tracenorm, tracedist, average_fidelity, dim, ⊗

abstract type QuantumGate end

struct UnitaryGate{T <: Complex} <: QuantumGate
    U::Matrix{T}
end

function UnitaryGate(U::T) where T<:AbstractMatrix
    if !isunitary(U)
        error("Gate must be unitary!")
    end
    (n,m) = size(U)
    if !(n == m)
        error("Gate matrix must be square!")
    end
    UnitaryGate(complex.(U))
end

Matrix(G::UnitaryGate{T}) where {T} = G.U

*(A::T, B::T) where T <: UnitaryGate = UnitaryGate(A.U*B.U)
*(A::T, B::AbstractMatrix) where T<:UnitaryGate = UnitaryGate(A.U*B)
*(A::AbstractMatrix, B::T) where T<:UnitaryGate = UnitaryGate(A*B.U)
*(A::Tg, x::Tx) where {Tg<:UnitaryGate, Tx<:Number} = UnitaryGate(A.U*x)
*(x::Tx, A::Tg) where {Tg<:UnitaryGate, Tx<:Number} = UnitaryGate(x*A.U)
-(A::T, B::T) where T<: UnitaryGate = UnitaryGate(A.U-B.U)

LinearAlgebra.kron(A::T, B::T) where T<:UnitaryGate = UnitaryGate(kron(A.U, B.U))
⊗(A::T, B::T) where T<:UnitaryGate = kron(A, B)

LinearAlgebra.adjoint(A::UnitaryGate) = UnitaryGate(Matrix(A.U'))
LinearAlgebra.tr(A::UnitaryGate) = tr(A.U)

dim(G::T) where T<:UnitaryGate = size(G.U)[1]

function dim(U::T)::Float64 where T<:AbstractMatrix
    if size(U)[1] == size(U)[2]
        return size(U)[1]
    else
        error("Can only return dimension of square matrix!")
    end
end

tracenorm(G::UnitaryGate) = tracenorm(Matrix(G))
tracedist(A::UnitaryGate, B::UnitaryGate)  = tracedist(Matrix(A), Matrix(B))
frobdist(A::UnitaryGate, B::UnitaryGate) = frobdist(Matrix(A), Matrix(B))

LinearAlgebra.norm(G::UnitaryGate) = norm(G.U, 2)
LinearAlgebra.ishermitian(G::T; tol=ϵ) where T<:UnitaryGate = ishermitian_tol(Matrix(G), tol=ϵ)

"""
Returns the average gate fidelity between two gates.
"""
function average_fidelity(A::T, B::T) where T<:AbstractMatrix
    @assert dim(A) == dim(B)
    d = dim(A)
    return (d + abs(tr(A*B'))^2)/(d*(d+1))
end

average_fidelity(A::T, B::T) where T <: UnitaryGate = average_fidelity(A.U, B.U)


text:
# Example of DandelionWebSocket.jl:
# Send some text and binary frames to ws://echo.websocket.org,
# which echoes them back.

using Requests: URI

using DandelionWebSockets

# Explicitly import the callback functions that we're going to add more methods for.
import DandelionWebSockets: on_text, on_binary,
                            state_connecting, state_open, state_closing, state_closed

# A simple WebSocketHandler which sends a few messages, receives echoes back, and then sends a stop
# signal via a channel when it's done.
mutable struct EchoHandler <: WebSocketHandler
    client::WSClient
    stop_channel::Channel{Any}
    received::Int
    sent::Int
end

# These are called when you get text/binary frames, respectively.
function on_text(h::EchoHandler, s::String)
    h.received += 1
    send_text(h.client, "foo")
    h.sent += 1
end

function on_binary(h::EchoHandler, data::Vector{UInt8})
    #println("Received data: $data")
    h.received += 1
    send_binary(h.client, b"foo")
    h.sent += 1
end

# These are called when the WebSocket state changes.

function state_closing(::EchoHandler)
    #println("State: CLOSING")
end

function state_connecting(::EchoHandler)
    #println("State: CONNECTING")
end

# Called when the connection is open, and ready to send/receive messages.
function state_open(handler::EchoHandler)
    #println("State: OPEN")

    send_text(handler.client, "Ready.")
end

function state_closed(::EchoHandler)
    #println("State: CLOSED")

    # Signal the script that the connection is closed.
    put!(stop_chan, true)
end

stop_chan = Channel{Any}(3)

# Create a WSClient, which we can use to connect and send frames.
client = WSClient()

handler = EchoHandler(client, stop_chan, 0, 0)

uri = URI("ws://server:8080")
println("Connecting to $uri... ")

wsconnect(client, uri, handler)

println("Connected.")

# The first message on `stop_chan` indicates that the connection is closed, so we can exit.
take!(stop_chan)

println("Received:", handler.received)
println("Sent:    ", handler.sent)


text:
# Riemann solver for standard hydro shocktube following Pfrommer et al 2006
# Use left and right state of initial condition to solve shocktube problems

using Roots

"""
    Helper functions to set up Machnumber dependent IC
"""
function solvePr(Pr::Float64,
                 rhol::Float64, rhor::Float64,
                 Pl::Float64, Mach::Float64,
                 γ::Float64)

    γ_1 = γ - 1.0
    γ_pow = γ_1/(2.0*γ)
    η2 = (γ-1.0)/(γ+1.0)

    c_l = sqrt.(γ*Pl/rhol)
    c_r = sqrt.(γ*Pr/rhor)

    f(x) = ( x/Pr - 1.0 ) * sqrt.( (1.0 - η2) / (γ * ( x/Pr + η2 ) )) -
            2.0 / γ_1 * c_l / c_r * ( 1.0 - ( x/Pl )^γ_pow )

    P_m = find_zero(f, (Pr, Pl), Bisection())

    ρ_mr = rhor * ( ( P_m + η2 * Pr ) / ( Pr + η2 * P_m))
    #ρ_ml = rhol * ( P_m / Pl )^(1.0/γ)

    vm = 2.0 * c_l / γ_1 * ( 1.0 - (P_m/Pl)^γ_pow )

    vs = vm / ( 1.0 - rhor/ρ_mr)

    M = vs / c_r

    return M - Mach
end

function solvePrfromMach(rhol::Float64, rhor::Float64,
                         Pl::Float64, Mach::Float64,
                         γ::Float64)

    # first approx without CRs
    findPrHelper(P) = solvePr(P, rhol, rhor,
                              Pl, Mach, γ)

    Pr = find_zero(findPrHelper, (1.e-5, 1.e5), Bisection() )

    return Pr
end

function solvePlfromMach(rhol::Float64, rhor::Float64,
                         Pr::Float64, Mach::Float64,
                         γ::Float64)

    # first approx without CRs
    findPlHelper(P) = solvePr(Pr, rhol, rhor,
                              P, Mach, γ)

    Pl = find_zero(findPlHelper, (1.e-5, 1.e5), Bisection() )

    return Pl
end

function solveMach(Pl::Float64, Pr::Float64,
                   rhol::Float64, rhor::Float64,
                   γ::Float64)

   γ_1 = γ - 1.0
   γ_pow = γ_1/(2.0*γ)
   η2 = (γ-1.0)/(γ+1.0)

   c_l = sqrt.(γ*Pl/rhol)
   c_r = sqrt.(γ*Pr/rhor)

   f(x) = ( x/Pr - 1.0 ) * sqrt.( (1.0 - η2) / (γ * ( x/Pr + η2 ) )) -
           2.0 / γ_1 * c_l / c_r * ( 1.0 - ( x/Pl )^γ_pow )

   P_m = find_zero(f, (Pr, Pl), Bisection())

   ρ_mr = rhor * ( ( P_m + η2 * Pr ) / ( Pr + η2 * P_m))
   ρ_ml = rhol * ( P_m / Pl )^(1.0/γ)

   vm = 2.0 * c_l / γ_1 * ( 1.0 - (P_m/Pl)^γ_pow )

   vs = vm / ( 1.0 - rhor/ρ_mr)

   M = vs / c_r

   return M
end

"""
    Datatypes for IC Parameters and Solution
"""
mutable struct SodParameters

    rhol::Float64           # denisty left
    rhor::Float64           # density right
    Pl::Float64             # pressure left
    Pr::Float64             # pressure right
    Ul::Float64             # internal energy left
    Ur::Float64             # internal energy right
    cl::Float64             # soundspeed left
    cr::Float64             # soundspeed right
    M::Float64              # Mach number
    t::Float64              # time
    x_contact::Float64      # position of the contact discontinuity along the tube
    γ_th::Float64           # adiabatic index of the gas
    γ_exp::Float64          # helper variable
    η2::Float64             # helper variable

    function SodParameters(;rhol::Float64=1.0,  rhor::Float64=0.125,
                            Pl::Float64=0.0,    Pr::Float64=0.0,
                            Ul::Float64=0.0,    Ur::Float64=0.0,
                            Mach::Float64=0.0,  t::Float64,
                            x_contact::Float64=70.0,
                            γ_th::Float64=5.0/3.0)

        γ_exp    = ( γ_th - 1.0 )/( 2.0 * γ_th )
        η2       = (γ_th-1.0)/(γ_th+1.0)

        # calculate Ul and Pl depending on input
        if (Pl == 0.0) & (Ul != 0.0)
            Pl = ( γ_th - 1.0 ) * rhol * Ul
        elseif (Ul == 0.0) & (Pl != 0.0)
            Ul = Pl / ( (γ_th - 1.0) * rhol )
        end

        # calculate Ur and Pr depending on input
        if (Pr == 0.0) & (Ur != 0.0)
            Pr = ( γ_th - 1.0 ) * rhor * Ur
        elseif (Ur == 0.0) & (Pr != 0.0)
            Ur = Pr / ( (γ_th - 1.0) * rhor )
        end

        # error handling
        if (Pr == 0.0) & (Pl == 0.0)
            error("No initial Pressure or energy values given!")
        end

        # solve right or left initial state from target mach number
        if (Pr == 0.0) & (Mach != 0.0)
            Pr = solvePrfromMach(rhol, rhor, Pl, Mach, γ_th)
            Ur = Pr / ( (γ_th - 1.0) * rhor )
        elseif (Pl == 0.0) & (Mach != 0.0)
            Pl = solvePlfromMach(rhol, rhor, Pr, Mach, γ_th)
            Ul = Pl / ( (γ_th - 1.0) * rhol )
        elseif (Ur == 0.0) & (Pr == 0.0) & (Mach == 0.0)
            error("Ur, Pr and Mach are zero! Can't find solution!")
        elseif (Ul == 0.0) & (Pl == 0.0) & (Mach == 0.0)
            error("Ul, Pl and Mach are zero! Can't find solution!")
        end

        if Mach == 0.0
            Mach = solveMach(Pl, Pr, rhol, rhor, γ_th)
        end

        cl = sqrt.(γ_th * Pl / rhol)
        cr = sqrt.(γ_th * Pr / rhor)

        new(rhol, rhor,
            Pl, Pr,
            Ul, Ur,
            cl, cr,
            Mach, t,
            x_contact,
            γ_th, γ_exp, η2)
    end
end

mutable struct SodHydroSolution
    x::Array{Float64,1}         # array of given positions
    rho::Array{Float64,1}       # array of densities along the tube
    rho4::Float64               # density in postshock region
    rho3::Float64               # density between contact disc. and rarefaction wave
    P::Array{Float64,1}         # array of pressures along the tube
    P34::Float64                # pressure between shock and rarefaction wave
    U::Array{Float64,1}         # array of internal energies along the tube
    v::Array{Float64,1}         # array of velocities along the tube
    v34::Float64                # velocity between shock and rarefaction wave
    vt::Float64                 # velocity of rarefaction wave
    vs::Float64                 # shock velocity
    Mach::Float64               # Mach number

    function SodHydroSolution(x::Array{Float64,1})
        N = length(x)
        new(x,
            zeros(N),   # rho
            0.0,        # rho2
            0.0,        # rho3
            zeros(N),   # P
            0.0,        # P34
            zeros(N),   # U
            zeros(N),   # v
            0.0,        # v23
            0.0,        # vt
            0.0,        # vs
            0.0)        # Mach
    end

    # multiple dispatch, in case gadget data is passed directly to the constructor
    # function SodHydroSolution(x::Array{Float32,1})
    #     N = length(x)
    #     new(x,
    #         zeros(N),   # rho
    #         0.0,        # rho2
    #         0.0,        # rho3
    #         zeros(N),   # P
    #         0.0,        # P34
    #         zeros(N),   # U
    #         zeros(N),   # v
    #         0.0,        # v23
    #         0.0,        # vt
    #         0.0,        # vs
    #         0.0)        # Mach
    # end
end

"""
    Functions to solve individual segments of the shocktube
"""

"""
    Pressure
"""
function solveP34(par::SodParameters)
    # Solves the pressure of the middle state in a shocktube

    P_m(P) = ( P/par.Pr - 1.0 ) * sqrt.( ( 1.0 - par.η2 ) / ( par.γ_th *
             ( P/par.Pr + par.η2 ) )) - 2.0 / (par.γ_th - 1.0 ) * par.cl / par.cr *
             ( 1.0 - ( P/par.Pl )^par.γ_exp )

    P34 = find_zero( P_m, (par.Pr, par.Pl), Bisection() )

    return P34
end

function solveP2(x::Float64, par::SodParameters)
    # solves the pressure along the rarefaction wave
    return par.Pl * ( -par.η2 * x / ( par.cl * par.t ) + ( 1.0 - par.η2 ) )^( 1.0/par.γ_exp )
end

function solveP(x::Float64, par::SodParameters, sol::SodHydroSolution)
    # returns P values according to position in shocktube

    if x <= -par.cl * par.t
        return par.Pl
    elseif -par.cl * par.t < x <= -sol.vt * par.t
        P2 = solveP2(x, par)
        return P2
    elseif -sol.vt * par.t < x <= sol.v34 * par.t
        return sol.P34
    elseif sol.v34 * par.t < x <= sol.vs * par.t
        return sol.P34
    elseif sol.vs * par.t < x
        return par.Pr
    else
        println("Error!")
        return 0.0
    end
end

"""
    Density
"""
function solveRho2(x::Float64, par::SodParameters)
    # solves density along the rarefaction wave
    return par.rhol * ( -par.η2 * x / ( par.cl * par.t ) +
                      ( 1.0 - par.η2 ) )^( 2.0 / (par.γ_th - 1.0 ) )
end

function solveRho3(par::SodParameters, sol::SodHydroSolution)
    # solves density left of contact discontiuity and right of refraction wave
    return par.rhol * ( sol.P34/par.Pl )^( 1.0/par.γ_th )
end

function solveRho4(par::SodParameters)
    # solves density left of shockfront
    return ( ( par.γ_th + 1.0 ) * par.M^2 ) /
           ( 2.0 + ( par.γ_th - 1.0 ) * par.M^2 ) * par.rhor
end

function solveRho(x::Float64, par::SodParameters, sol::SodHydroSolution)
    # returns P values according to position in shocktube

    if x <= -par.cl * par.t
        return par.rhol
    elseif -par.cl * par.t < x <= -sol.vt * par.t
        return solveRho2(x, par)
    elseif -sol.vt * par.t < x <= sol.v34 * par.t
        return sol.rho3
    elseif sol.v34 * par.t < x <= sol.vs * par.t
        return sol.rho4
    elseif sol.vs * par.t < x
        return par.rhor
    else
        println("Error!")
        return 0.0
    end
end

"""
    Velocity
"""
function solveV34(par::SodParameters, sol::SodHydroSolution)
    # solves velocity along the isopressure region 2-3
    return 2.0 * par.cl / ( par.γ_th - 1.0 ) * ( 1.0 - (sol.P34 / par.Pl)^par.γ_exp )
end

function solveV2(x::Float64, par::SodParameters)
    # solves the velocity along the rarefaction wave
    return ( 1.0 - par.η2 ) * ( x / par.t + par.cl )
end

function solveVs(par::SodParameters, sol::SodHydroSolution)
    # solves the shock velocity
    return sol.v34 / ( 1.0 - par.rhor / sol.rho4 )
end

function solveVt(par::SodParameters, sol::SodHydroSolution)
    # solves the velocity of the tail of the rarefaction wave
    return par.cl - sol.v34/( 1.0 - par.η2 )
end

function solveV(x::Float64, par::SodParameters, sol::SodHydroSolution)
    # solves the velocity along the shocktube
    if x <= -par.cl * par.t
        return 0.0
    elseif -par.cl * par.t < x <= -sol.vt * par.t
        return solveV2(x, par)
    elseif -sol.vt * par.t < x <= sol.vs * par.t
        return sol.v34
    elseif sol.vs * par.t < x
        return 0.0
    else
        println("Error!")
        return 0.0
    end
end

"""
    Main function for solver
"""
function solveSodShock(x::Array{Float64,1}; par::SodParameters)
    # solves a standard sod shock with zero initial velocity

    N = length(x)

    # set up datatype to store riemann solution
    sol = SodHydroSolution(x)

    # transform into rest-frame of contact discontiuity
    x_in = sol.x .- par.x_contact

    # solve Pressure
    sol.P34 = solveP34(par)

    sol.rho3 = solveRho3(par, sol)
    sol.rho4 = solveRho4(par)

    # solve velocity
    sol.v34 = solveV34(par, sol)
    sol.vs = solveVs(par, sol)
    sol.vt = solveVt(par, sol)
    for i = 1:N
        sol.v[i] = solveV(x_in[i], par, sol)
    end

    for i = 1:N
        sol.P[i] = solveP(x_in[i], par, sol)
    end

    # solve density

    for i = 1:N
        sol.rho[i] = solveRho(x_in[i], par, sol)
    end

    sol.Mach = par.M

    sol.U = sol.P ./ ( (par.γ_th - 1.0) .* sol.rho )

    return sol
end

max_stars_repo_name: LudwigBoess/GadJet

max_issues_repo_name: LudwigBoess/GadJet

max_forks_repo_name: LudwigBoess/GadJet


text:
# -*- coding: utf-8 -*-
module CategoryTheory

using SemanticModels
import Catlab.Theories: dom, codom
import Base: append!, push!, deleteat!, delete!

export ⊔, AbstractMorph, FinSetMorph, dom, codom, verify, func, Decorated, decorations, undecorate, AbstractSpan, leftob, rightob, apexob, Span, left, right, DoublePushout, AbstractCospan, Cospan, pushout

⊔(a::UnitRange, b::UnitRange) = 1:(length(a)+length(b))
⊔(a::AbstractVector{Int}, b::AbstractVector{Int}) = vcat(a,b)

"""    AbstractMorph

an abstract type for representing morphisms. The essential API for subtypes of AbstractMorph are

1. dom(m)::T
2. codom(m)::T
3. func(m)::Function

where T is the type of the objects in the category. See FinSetMorph for an example.
"""
abstract type AbstractMorph end

"""    FinSetMorph{T,F}

morphisms in the category of Finite Sets. The objects are of type UnitRange{Int}.
func(m::FinSetMorph) is a function that takes `Int -> Int`. FinSetMorphs can be constructed
from a list of numbers. For example, `FinSetMorph([1,3,2,3])` is the morphism that takes
`1->1, 2->3, 3->2, 4->3` on domain `1:4` with codomain `1:3`. When you define a morphism from
a list of integers, the codomain is inferred from the largest element of the list. The domain
must always be the `1:l` where `l` is the length of the input list.

```(f::FinSetMorph)(g::G) where G <: AbstractGraph```

lift a finite set morphism (list of integers) to a graph homomorphism by its action on the vertex
set. The graph `h = f(g)` is defined by taking the edges of `g` and relabeling their src and dst
according to the function of `f`.

This method computes a valid graph homomorphism by definition.
"""
struct FinSetMorph{T,F} <: AbstractMorph
    codom::T
    fun::F
end

FinSetMorph(v::AbstractVector{Int}) = FinSetMorph(1:maximum(v), v)

dom(m::FinSetMorph) = 1:length(m.fun)
codom(m::FinSetMorph) = m.codom
func(m::FinSetMorph) = i->m.fun[i]

"""    ⊔(f::FinSetMorph, g::FinSetMorph)

the union of two morphisms in a finite set.
"""
function ⊔(f::FinSetMorph, g::FinSetMorph)
    Y = codom(f) ⊔ codom(g)
    h = f.fun ⊔ g.fun
    FinSetMorph(Y, h)
end

"""    Decorated{M,T}

a decoration applied to the objects of a morphism, where M is a type of morphism and
type T is the category of the decoration
"""
struct Decorated{M}
    f::M
    d::AbstractArray{AbstractModel}
end

# Handle creating a decorated morphism with an array of a single type
Decorated(f, d::AbstractArray{T}) where T<:AbstractModel = Decorated(f, Vector{AbstractModel}(d))
# Handle creating a decorated morphism from a single decoration
Decorated(f, d::T) where T<:AbstractModel = Decorated(f, Vector{AbstractModel}([d]))

# Get the domain or codomain of a decorated morphism
dom(m::Decorated) = dom(m.f)
codom(m::Decorated) = codom(m.f)

# Get the decorations of a decorated morphism
function decorations(m::Decorated)
  return m.d
end
# Get the decorations of AbstractModel T of a decorated morphism
function decorations(m::Decorated, ::Type{T}) where T<:AbstractModel
  filter(x -> isa(x,T), decorations(m))
end

# Remove the decoration of a decorated morphism, and return the original morphism
function undecorate(m::Decorated)
  return m.f
end

# Add a decoration to a decorated morphism
function push!(m::Decorated, decoration::AbstractModel)
  push!(decorations(m), decoration)
end

# Add a collection of decorations to a decorated morphism
function append!(m::Decorated, decorations)
    append!(m.d, decorations)
end

# remove a decoration from a decorated morphism
function deleteat!(m::Decorated, i)
  deleteat!(decorations(m), i)
end

# Remove the decorations of AbstractModel T from a decorated morphism
function delete!(m::Decorated, ::Type{T}) where T<:AbstractModel
  filter!(x -> !isa(x,T), decorations(m))
end

function left(d::Decorated)
  return left(d.f)
end

function right(d::Decorated)
  return right(d.f)
end

"""    AbstractSpan

an abstract type for representing spans. The essential API for subtypes of AbstractSpan are

1. left(s)::M
2. right(s)::M
3. pushout(s)::C

where M is the type of morphisms in the span, and C is the type of cospan that solves
a pushout of span s. See Span for an example.
"""
abstract type AbstractSpan end

leftob(s::AbstractSpan) = codom(left(s))

rightob(s::AbstractSpan) = codom(right(s))

function apexob(s::AbstractSpan)
    a = dom(left(s))
    b = dom(right(s))
    a == b || error("Inconsistent span")
    return a
end

"""    Span{F,G} <: AbstractSpan

a general span type where types F and G are types of morphisms in the span
"""
struct Span{F,G} <: AbstractSpan
    f::F
    g::G
end

function left(s::Span)
    return s.f
end

function right(s::Span)
    return s.g
end

"""    undecorate(s::Span{T,T}) where T <: Decorated

remove decorations of a span of decorated morphisms
"""
function undecorate(s::Span{T,T}) where T <: Decorated
    return Span(undecorate(left(s)), undecorate(right(s)))
end

struct DoublePushout{S<:AbstractSpan, T<:NTuple{3,AbstractMorph}}
    rule::S
    morphs::T
    application::S
end

# TODO DPO CONSTRUCTOR TO SOLVE UNKNOWN DOUBLEPUSHOUT OF FINSET
# take in Span `top` (l, c, r) and finset `l′ `, f: l -> l′
# Solve for c′ using dropdown
#     c′ = setdiff(l′, f(l)) ⊔ c
# Solve pushout for span of c′ ← c → r to get r′

# function pullback

# +
"""    AbstractCospan

an abstract type for representing cospans. The essential API for subtypes of AbstractCospan are

1. left(c)::M
2. right(c)::M
3. pullback(c)::S

where M is the type of morphisms in the cospan, and S is the type of span that solves the
pullback defined by c. See Cospan for an example.
"""
abstract type AbstractCospan end

leftob(c::AbstractCospan) = codom(left(c))

rightob(c::AbstractCospan) = codom(right(c))

function apexob(c::AbstractCospan)
    a = dom(left(c))
    b = dom(right(c))
    a == b || error("Inconsistent cospan")
    return a
end

"""    Cospan{F,G} <: AbstractCospan

a general cospan type where types F and G are types of morphisms in the cospan
"""
struct Cospan{F,G} <: AbstractCospan
    f::F
    g::G
end

function left(c::Cospan)
    return c.f
end

function right(c::Cospan)
    return c.g
end

"""    undecorate(s::Copan{T,T}) where T <: Decorated

remove decorations of a cospan of decorated morphisms
"""
function undecorate(c::Cospan{T,T}) where T <: Decorated
    return Cospan(undecorate(left(c)), undecorate(right(c)))
end

# TODO: This pushout gets the numbers mixed up. It is right up to permutation
# which is fine for when we want to do one pushout. But when we want to chain
# pushouts in order to do composition, "up to permutation" isn't good enough.
"""    pushout(s::Span{T,T}) where T <: FinSetMorph

treat f,g as a span and compute the pushout that is, the cospan of f=(f⊔g) and g=(a⊔b)
"""
function pushout(s::Span{T,U}) where {T <: FinSetMorph, U <: FinSetMorph}
    f_dict = Dict(a=>i for (i, a) in enumerate(left(s).fun))
    g′ = map(n->n in keys(f_dict) ? func(right(s))(f_dict[n]) : n+length(rightob(s)), leftob(s))

    g_dict = Dict(a=>i for (i, a) in enumerate(right(s).fun))
    f′ = map(n->n in keys(g_dict) ? g′[func(left(s))(g_dict[n])] : n, rightob(s))

    u = union(f′, g′)
    u_dict = Dict(a=>i for (i, a) in enumerate(u))
    f′ = FinSetMorph(1:length(u), map(n->u_dict[n], f′))
    g′ = FinSetMorph(1:length(u), map(n->u_dict[n], g′))

    return Cospan(f′, g′)
end

"""    pushout(s::Span{T, T}) where T <: Decorated

treat f,g as a decorated span and compute the pushout that is, the cospan of f=(f⊔g) and g=(a⊔b),
with the decoration of (f⊔g)(d)
"""
function pushout(s::Span{T, T}) where T <: Decorated
    cs = pushout(undecorate(s))
    decorations = map(x->x[1] ⊔ x[2], zip(left(s).d, right(s).d))
    return Decorated(cs, map((right(cs) ⊔ left(cs)), decorations))
end

end


text:
using Main.VoronoiDelaunay
using Main.GeometricalPredicates

struct Scaler
    scale::Any
    transx::Any
    transy::Any

    function Scaler(bbox::Array{Int,2})
        _scale = (abs(1 / (bbox[2] - bbox[1])))
        new(_scale, (1 - bbox[1] * _scale), (1 - bbox[1, 2] * _scale))
    end
end

function scaled_point(unscaled_point::Point2D, scaler::Scaler)
    return Point(
        ((getx(unscaled_point) * scaler.scale) + scaler.transx),
        ((gety(unscaled_point) * scaler.scale) + scaler.transy),
    )
end

function unscaled_point(scaled_point::Point2D, scaler::Scaler)
    return Point(
        (getx(scaled_point) - scaler.transx) / scaler.scale,
        (gety(scaled_point) - scaler.transy) / scaler.scale,
    )
end

function scale_p(p, scaler::Scaler)
    return [(p[1] * scaler.scale) + scaler.transx, (p[2] * scaler.scale) + scaler.transy]
end

function unscale_p(p, scaler::Scaler)
    return [(p[1] - scaler.transx) / scaler.scale, (p[2] - scaler.transy) / scaler.scale]
end

function unscale_y(p, scaler)
    return (p - scaler.transx) / scaler.scale
end

function unscale_x(p, scaler)
    return (p - scaler.transy) / scaler.scale
end

function meshgrid(vx::AbstractVector{T}, vy::AbstractVector{T}) where {T}
    m, n = length(vy), length(vx)
    gx = reshape(repeat(vx, inner = m, outer = 1), m, n)
    gy = reshape(repeat(vy, inner = 1, outer = n), m, n)
    return gx, gy
end

function extractbox(bbox, h0::Real, h1::Real)
    v1 = bbox[1, 1]:h0:bbox[2, 1]
    v2 = bbox[1, 2]:h1:bbox[2, 2]
    return v1, v2
end

function calculateh1(h0::Real)
    return h0 * (sqrt(3) / 2)
end

function shiftevenrows!(x, h0::Real)
    x[2:2:end, :] = x[2:2:end, :] .+ h0 / 2
end

function vectorizededges(triangle)
    a = geta(triangle)
    b = getb(triangle)
    c = getc(triangle)
    ab = Line(Point(getx(a), gety(a)), Point(getx(b), gety(b)))
    ba = Line(Point(getx(b), gety(b)), Point(getx(a), gety(a)))
    bc = Line(Point(getx(b), gety(b)), Point(getx(c), gety(c)))
    cb = Line(Point(getx(c), gety(c)), Point(getx(b), gety(b)))
    ac = Line(Point(getx(a), gety(a)), Point(getx(c), gety(c)))
    ca = Line(Point(getx(c), gety(c)), Point(getx(a), gety(a)))
    return ab, ba, bc, cb, ac, ca
end

function validedges(tess, scaler, fd, geps)
    i = 1
    inside_edges = Array{Line2D,1}()
    outside_edges = Array{Line2D,1}()
    for triangle in tess
        a = geta(triangle)
        b = getb(triangle)
        c = getc(triangle)
        center = unscaled_point(centroid(Primitive(a, b, c)), scaler)
        i += 1
        edges = vectorizededges(triangle)
        if fd([getx(center), gety(center)]) > -geps
            push!(inside_edges, edges...)
        else
            push!(outside_edges, edges...)
        end
    end
    edges = Array{Line2D,1}()
    for edge in delaunayedges(tess)
        push!(edges, Line(geta(edge), getb(edge)))
    end
    inside_edges = filter(edge -> !(edge in outside_edges), inside_edges)
    edges = filter(edge -> !(edge in inside_edges), edges)
    return [edge for edge in edges]
end

function plotedges(edges)
    x = Array{Float64,1}()
    y = Array{Float64,1}()
    for edge in edges
        push!(x, getx(geta(edge)))
        push!(x, getx(getb(edge)))
        push!(x, NaN)
        push!(y, gety(geta(edge)))
        push!(y, gety(getb(edge)))
        push!(y, NaN)
    end
    return x, y
end

function pointstoforces(edges, scaler, Fscale, pfix)
    bars = Array{Point2D,1}()
    barvec = Array{Array{Float64,1},1}()
    points_to_fvces = Dict{Point2D,Array{Float64,1}}()
    for edge in edges
        b = unscaled_point(getb(edge), scaler)
        a = unscaled_point(geta(edge), scaler)
        push!(
            bars,
            Point(getx(a) + ((getx(b) - getx(a)) / 2), gety(a) + ((gety(b) - gety(a)) / 2)),
        )
        push!(
            barvec,
            [
                getx(b) - getx(a)
                gety(b) - gety(a)
            ],
        )
        push!(points_to_fvces, geta(edge) => [0, 0])
        push!(points_to_fvces, getb(edge) => [0, 0])
    end
    L = [sqrt(sum(v_sum .^ 2)) for v_sum in barvec]
    hbars = [fh(getx(p), gety(p)) for p in bars]
    L0 = hbars * Fscale * sqrt(sum(L .^ 2) / sum(hbars .^ 2))
    Fvec = maximum.(L0 - L) ./ L .* barvec
    iterator = 1
    for edge in edges
        a = geta(edge)
        b = getb(edge)
        prev_a = points_to_fvces[a]
        prev_b = points_to_fvces[b]
        push!(points_to_fvces, a => prev_a + (-Fvec[iterator]))
        push!(points_to_fvces, b => prev_b + (Fvec[iterator]))
        iterator = iterator + 1
    end
    for p in pfix
        delete!(points_to_fvces, p)
        push!(points_to_fvces, p => [0, 0])
    end
    return points_to_fvces
end

function finalpositions(points_to_fvces, scaler, deltat, fd, geps, deps, h0)
    new_p = Array{Point2D,1}()
    d_points = Array{Array{Float64,1},1}()
    for (point, force) in points_to_fvces
        p = unscaled_point(point, scaler)
        np = [getx(p), gety(p)] + deltat * force
        push!(new_p, Point2D(np[1], np[2]))
        d = fd(np)
        if d < -geps
            push!(d_points, force)
        end
    end
    final_p = Array{Array{Float64,1},1}()
    for p in new_p
        x = getx(p)
        y = gety(p)
        d = fd([x, y])
        if d > 0
            dgradx = (fd([x + deps, y]) - d) / deps
            dgrady = (fd([x, y + deps]) - d) / deps
            res = [x, y] - [d * dgradx, d * dgrady]
            push!(final_p, res)
        else
            push!(final_p, [x, y])
        end
    end
    final_p = map(p -> scaled_point(Point(p[1], p[2]), scaler), unique(final_p))
    return final_p, moveindex(d_points, deltat, h0)
end

function moveindex(d_points, deltat, h0)
    d = map(row -> sum(deltat * row .^ 2), d_points)
    push!(d, 0)
    return maximum(sqrt.(d) / h0)
end

function movein(p)
    newx = getx(p)
    newy = gety(p)
    if getx(p) <= 1
        newx = round(newx, digits = 8) + 2 *eps(Float64)
    end
    if gety(p) <= 1
        newy = round(newy, digits = 8) + 2 *eps(Float64)
    end
    if getx(p) >= 2
        newx = round(newx, digits = 8) - 3 * eps(Float64)
    end
    if gety(p) >= 2
        newy = round(newy, digits = 8) - 3 * eps(Float64)
    end
    return Point(newx, newy)
end

function distmesh(
    fd,
    fh,
    bbox,
    h0,
    pfix = [],
    ttol = 0.1,
    geps = 0.001 * h0,
    Fscale = 1.2,
    dptol = 0.001,
    deltat = 0.2,
    deps = sqrt(eps(Float64)) * h0,
)
    h1 = calculateh1(h0)
    pold = Inf
    scaler = Scaler(bbox)
    v1, v2 = extractbox(bbox, h0, h1)
    x, y = meshgrid(v1, v2)
    shiftevenrows!(x, h0)
    p = [x[:] y[:]]
    p = [Point(row[1], row[2]) for row in eachrow([x[:] y[:]]) if fd(row) < geps]
    r0 = [1 ./ fh(getx(point), gety(point)) .^ 2 for point in p]
    r0_max = maximum(r0 ./ maximum(r0))
    scaler = Scaler(bbox)
    pfix = [scaled_point(Point(row[1], row[2]), scaler) for row in eachrow(pfix)]
    pfix = map(p -> movein(p), pfix)
    p = [scaled_point(point, scaler) for point in p if (rand(Float64, size(p)))[1] < r0_max]
    p = [p; pfix]
    p = map(p -> movein(p), p)
    p = unique(p)
    while true
        tess = DelaunayTessellation(convex = false)
        push!(tess, p)
        edges = validedges(tess, scaler, fd, geps)
        points_to_fvces = pointstoforces(edges, scaler, Fscale, pfix)
        final_p, move_index =
            finalpositions(points_to_fvces, scaler, deltat, fd, geps, deps, h0)
        if move_index < dptol
            x, y = plotedges(edges)
            break
        else
            p = map(p -> movein(p), final_p)
        end
    end
    x = map(x -> unscale_x(x, scaler), x)
    y = map(y -> unscale_y(y, scaler), y)
    return x, y
end

max_stars_repo_name: jstarczewski/julia-num

max_issues_repo_name: jstarczewski/julia-num

max_forks_repo_name: jstarczewski/julia-num