importall LegHemodynamics

function main()

rstflag = "yes" # restarting from scratch or previous simulation
hemoflag = "no" # 10% hemorrhage from left femoral artery
saveflag = "yes" # save solution to .mat file
coupleflag = "no" # coupling to 3D liver tissue model
timeflag = "yes" # solver timing
assimflag = "yes" # patient data assimilation via EnKF

# build solution struct or generate ensemble
if assimflag == "no"
    if rstflag == "no"
        loadfile = "arterylist_2.txt"; # default artery data file for new sim
    elseif rstflag == "yes"
        loadfile = "test.mat"; # restart file
    end
    system = LegHemodynamics.buildall(loadfile;numbeatstotal=1,restart=rstflag);
    savefile = "test.mat" # filename for saving (only used if saveflag == "yes")
elseif assimflag == "yes"
    ensemblesize = 10;
    if rstflag == "no"
        loadfiles = ["arterylist_2_thick_branches.txt" for i=1:ensemblesize];
    elseif rstflag == "yes"
        loadfiles = ["small_dev_7_$i.mat" for i=1:ensemblesize];
    end
    systems = pmap((a1)->LegHemodynamics.buildall(a1;numbeatstotal=1,restart=rstflag),loadfiles);
    savefiles = ["converged_small_dev_1_$i.mat" for i=1:ensemblesize];
    if rstflag == "no" # randomize ICs on new ensemble
        soln = pmap((a1)->LegHemodynamics.applycustomics!(a1),systems);
        systems = [soln[i] for i=1:ensemblesize];
    elseif rstflag == "yes" # ensure uniform time step
        h = [];
        for i = 1:length(systems)
            push!(h,systems[i].solverparams.h)
        end
        for i = 1:length(systems)
            systems[i].solverparams.h = minimum(h);
        end
        systems = pmap((a1,a2)->LegHemodynamics.rediscretizet!(a1,a2),systems,zeros(Int64,length(systems)));
    end
end

# create interpolation object for upstream flow velocity data (needed for proximal BC)
patientdata = MAT.matread("flowrates_mL.mat");
tdata = reshape(patientdata["t"],101);
udata = reshape(patientdata["u_pop_cms"],101)*LegHemodynamics.cmTom;
if assimflag == "no"
    spl = Dierckx.Spline1D(tdata,udata;k=1);
elseif assimflag == "yes"
    spl = [Dierckx.Spline1D(tdata,udata;k=1) for i=1:ensemblesize];
    udata = reshape(patientdata["u_ant_cms"],101)*LegHemodynamics.cmTom;
    splant = Dierckx.Spline1D(tdata,udata;k=1);
    udata = reshape(patientdata["u_post_cms"],101)*LegHemodynamics.cmTom;
    splpost = Dierckx.Spline1D(tdata,udata;k=1);
    udata = reshape(patientdata["u_per_cms"],101)*LegHemodynamics.cmTom;
    splper = Dierckx.Spline1D(tdata,udata;k=1);
end

# timers
if assimflag == "no"
    times = LegHemodynamics.CVTimer();
elseif assimflag == "yes"
    times = [LegHemodynamics.CVTimer() for i=1:ensemblesize];
end

# collect all IDs of terminal/non-terminal branches
terms = Int64[];
splits = Int64[];
if assimflag == "no"
    for i = 1:length(system.branches.ID)
        if isempty(system.branches.children[i])
            push!(terms,i);
        else
            push!(splits,i);
        end
    end
elseif assimflag == "yes"
    for i = 1:length(systems[1].branches.ID)
        if isempty(systems[1].branches.children[i])
            push!(terms,i);
        else
            push!(splits,i);
        end
    end
    term_itr = [terms for i=1:ensemblesize];
    split_itr = [splits for i=1:ensemblesize];
end

# ensemble distributions, allocators, and scalings
if assimflag == "yes"
    nparams = 15;
    nstates = 3;
    nmeas = 3;
    errors = LegHemodynamics.Errors(nparams);

    X = [zeros(nstates) for i=1:ensemblesize];
    θ = [zeros(nparams) for i=1:ensemblesize];
    Y = [zeros(nmeas) for i=1:ensemblesize];

    xhat = zeros(nstates);
    θhat = zeros(nparams);
    yhat = zeros(nmeas);

    yi = [zeros(nmeas) for i=1:ensemblesize];

    # parameter scalings
    θs = zeros(nparams);
    for i=1:ensemblesize
        θs[1] += systems[i].branches.term[14].R[2]; # anterior
        θs[2] += systems[i].branches.term[24].R[2]; # peroneal
        θs[3] += systems[i].branches.term[32].R[2]; # posterior
        θs[4] += systems[i].branches.term[14].C[1];
        θs[5] += systems[i].branches.term[24].C[1];
        θs[6] += systems[i].branches.term[32].C[1];
        θs[7] += systems[i].branches.termscalings[1];
        θs[8] += systems[i].branches.termscalings[2];
        θs[9] += systems[i].branches.termscalings[3];
        θs[10] += systems[i].branches.termscalings[4];
        θs[11] += systems[i].branches.termscalings[5];
        θs[12] += systems[i].branches.termscalings[6];
        # θs[9] += systems[i].branches.A0[14];
        # θs[10] += systems[i].branches.A0[24];
        # θs[11] += systems[i].branches.A0[31];
        θs[13] += systems[i].branches.beta[14];
        θs[14] += systems[i].branches.beta[24];
        θs[15] += systems[i].branches.beta[32];
        # θs[9] += systems[i].branches.YModinPa[14];
        # θs[10] += systems[i].branches.YModinPa[24];
        # θs[11] += systems[i].branches.YModinPa[31];
    end
    θs /= ensemblesize;

    # RTPS allocators
    p = 0.5; # RTPS relaxation amount, ∃ [0.5,0.95]
    c = zeros(ensemblesize);
    σxb = zeros(nstates);
    σxa = zeros(nstates);
    σtb = zeros(nparams);
    σta = ones(nparams);

    tout = Float64[];
    xout = Vector{Float64}[];
    xoutv = Vector{Float64}[];
    yout = Vector{Float64}[];
    youtv = Vector{Float64}[];
    innov = Vector{Float64}[];
    ioutv = Vector{Float64}[];
    θout = Vector{Float64}[];
    θoutv = Vector{Float64}[];
    Pθout = Vector{Float64}[];
    Pθoutv = Vector{Float64}[];
    Pxout = Vector{Float64}[];
    Pxoutv = Vector{Float64}[];
    lbx = Vector{Float64}[];
    ubx = Vector{Float64}[];
    lbxv = Vector{Float64}[];
    ubxv = Vector{Float64}[];

    # default parameter values for non-measured terminal constraints
    Rdefault = 0.75*[77.6909,141.4980,114.7901,102.6618,176.7149,257.4914,18.6383,
        107.4268,291.8191,191.9240,149.3493,35.9008,53.0083,150.1083,259.9696,
        142.5643,95.1046]*LegHemodynamics.gTokg/(LegHemodynamics.mmTom^4);
    Cdefault = [0.08972235,0.04926296,0.06072485,0.06789874,0.03944552,0.02707124,
        0.39707822,0.06488707,0.02388675,0.03631963,0.04667319,0.19569034,0.13150028,
        0.04643721,0.02681318,0.04830696,0.00723676]*(LegHemodynamics.mmTom^4)/LegHemodynamics.gTokg;
end

# time step counter
if assimflag == "no"
    n = system.solverparams.nstart;
elseif assimflag == "yes"
    n = [systems[1].solverparams.nstart for i=1:ensemblesize];
end

# solver loop
tic();
if assimflag == "no"
    while system.solverparams.numbeats < system.solverparams.numbeatstotal
        # time within heart cycle
        tp = system.t[n+1] - sum(system.solverparams.th*system.solverparams.numbeats);
        if mod(n,1000) == 0
            println("Reached time step $n. t = $tp s.")
        end
        # interpolate proximal flow velocity from patient data
        uprox = spl(tp);
        # TVD RK3 time integration
        LegHemodynamics.tvdrk3!(system,times,n,splits,terms,-uprox);
        tic();
        # arterial pressure update
        LegHemodynamics.arterialpressure!(system,n);
        # check for start of new cardiac cycle
        oldnumbeats = system.solverparams.numbeats;
        LegHemodynamics.setnumbeats!(system,n);
        if oldnumbeats < system.solverparams.numbeats
            # update time discretization
            if system.solverparams.numbeats < system.solverparams.numbeatstotal
                LegHemodynamics.updatediscretization!(system);
            end
        end
        n+=1;
        times.tr += toq();
    end
    times.tt += toc();
elseif assimflag == "yes"
    nsamp = [300 for i=1:ensemblesize];
    while systems[1].solverparams.numbeats < systems[1].solverparams.numbeatstotal
        # normal time integration between measurements
        rflag = "predict";
        soln = pmap((a1,a2,a3,a4,a5,a6,a7)->LegHemodynamics.advancetime!(a1,a2,a3,a4,a5,a6,a7;
            runtype=rflag),systems,spl,times,n,nsamp,term_itr,split_itr);
        systems = [soln[i][1] for i=1:ensemblesize];
        n = [soln[i][2] for i=1:ensemblesize];
        times = [soln[i][3] for i=1:ensemblesize];
        # data assimilation
        if assimflag == "yes" && mod((n[1]-1)/nsamp[1],1) == 0
            println("Current time step: $(n[1])")
            println("Current time: $(systems[1].t[n[1]])")

            # time within heart cycle
            tp = systems[1].t[n[1]] - sum(systems[1].solverparams.th*systems[1].solverparams.numbeats);
            # non-dimensional measurement (velocity)
            y = -[splant(tp),splpost(tp),splper(tp)]/LegHemodynamics.us;
            # y = -[splant(tp)]/LegHemodynamics.us;
            println("Scaled velocity measurements at t = $tp s: $y")

            # # generate measurement replicates
            # for i = 1:length(systems)
            #     yi[i][1] = y[1] + rand(Distributions.Normal(0,errors.odev[1]/LegHemodynamics.us));
            #     yi[i][2] = y[2] + rand(Distributions.Normal(0,errors.odev[2]/LegHemodynamics.us));
            #     yi[i][3] = y[3] + rand(Distributions.Normal(0,errors.odev[3]/LegHemodynamics.us));
            # end
            #
            # # forecast parameters and means
            # for i = 1:length(systems)
            #     θ[i][1] = systems[i].branches.term[14].R[2];
            #     θ[i][2] = systems[i].branches.term[24].R[2];
            #     θ[i][3] = systems[i].branches.term[32].R[2];
            #     θ[i][4] = systems[i].branches.term[14].C[1];
            #     θ[i][5] = systems[i].branches.term[24].C[1];
            #     θ[i][6] = systems[i].branches.term[32].C[1];
            #     θ[i][7] = systems[i].branches.termscalings[1];
            #     θ[i][8] = systems[i].branches.termscalings[2];
            #     θ[i][9] = systems[i].branches.termscalings[3];
            #     θ[i][10] = systems[i].branches.termscalings[4];
            #     θ[i][11] = systems[i].branches.termscalings[5];
            #     θ[i][12] = systems[i].branches.termscalings[6];
            #     # θ[i][9] = systems[i].branches.A0[14];
            #     # θ[i][10] = systems[i].branches.A0[24];
            #     # θ[i][11] = systems[i].branches.A0[31];
            #     θ[i][13] = systems[i].branches.beta[14];
            #     θ[i][14] = systems[i].branches.beta[24];
            #     θ[i][15] = systems[i].branches.beta[32];
            #     # θ[i][9] = systems[i].branches.YModinPa[14];
            #     # θ[i][10] = systems[i].branches.YModinPa[24];
            #     # θ[i][11] = systems[i].branches.YModinPa[31];
            # end
            #
            # # forecast parameters (dimensional)
            # θ = [LegHemodynamics.paramwalk!(errors,θs,θ[i]) for i in 1:length(systems)];
            #
            # # non-dimensionalize parameters
            # for i = 1:length(systems)
            #     θ[i] = θ[i]./θs;
            # end
            #
            # # RTPS prior standard deviation
            # for i in 1:length(θ[1])
            #     for j in 1:length(systems)
            #         c[j] = θ[j][i];
            #     end
            #     σtb[i] = std(c;corrected=true);
            # end
            #
            # # parameters back into model
            # for i = 1:length(soln)
            #     systems[i].branches.term[14].R[2] = θ[i][1]*θs[1];
            #     systems[i].branches.term[24].R[2] = θ[i][2]*θs[2];
            #     systems[i].branches.term[32].R[2] = θ[i][3]*θs[3];
            #     systems[i].branches.term[14].C[1] = θ[i][4]*θs[4];
            #     systems[i].branches.term[24].C[1] = θ[i][5]*θs[5];
            #     # systems[i].branches.term[24].C[1] = 2.5e-11;
            #     systems[i].branches.term[32].C[1] = θ[i][6]*θs[6];
            #     systems[i].branches.termscalings[1] = θ[i][7]*θs[7];
            #     systems[i].branches.termscalings[2] = θ[i][8]*θs[8];
            #     systems[i].branches.termscalings[3] = θ[i][9]*θs[9];
            #     systems[i].branches.termscalings[4] = θ[i][10]*θs[10];
            #     systems[i].branches.termscalings[5] = θ[i][11]*θs[11];
            #     systems[i].branches.termscalings[6] = θ[i][12]*θs[12];
            #     # k = 1;
            #     for j in [2,4;] # popliteal
            #         systems[i].branches.term[j].C[1] *= 1/3*(systems[i].branches.termscalings[1]+
            #             systems[i].branches.termscalings[3]+systems[i].branches.termscalings[5]);
            #         systems[i].branches.term[j].R[2] *= 1/3*(systems[i].branches.termscalings[2]+
            #             systems[i].branches.termscalings[4]+systems[i].branches.termscalings[6]);
            #         # systems[i].branches.term[j].C[1] = 1/3*(systems[i].branches.termscalings[1]+
            #             # systems[i].branches.termscalings[3]+systems[i].branches.termscalings[5])*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = 1/3*(systems[i].branches.termscalings[2]+
            #             # systems[i].branches.termscalings[4]+systems[i].branches.termscalings[6])*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [7,9,11,13;] # anterior
            #         systems[i].branches.term[j].C[1] *= systems[i].branches.termscalings[1];
            #         systems[i].branches.term[j].R[2] *= systems[i].branches.termscalings[2];
            #         # systems[i].branches.term[j].C[1] = systems[i].branches.termscalings[1]*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = systems[i].branches.termscalings[2]*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [17,19,21,23;] # posterior
            #         systems[i].branches.term[j].C[1] *= systems[i].branches.termscalings[3];
            #         systems[i].branches.term[j].R[2] *= systems[i].branches.termscalings[4];
            #         # systems[i].branches.term[j].C[1] = systems[i].branches.termscalings[3]*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = systems[i].branches.termscalings[4]*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [25,27,29,31;] # peroneal
            #         systems[i].branches.term[j].C[1] *= systems[i].branches.termscalings[5];
            #         systems[i].branches.term[j].R[2] *= systems[i].branches.termscalings[6];
            #         # systems[i].branches.term[j].C[1] = systems[i].branches.termscalings[5]*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = systems[i].branches.termscalings[6]*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [1:4;] # popliteal
            #         systems[i].branches.beta[j] *= 1/3*(θ[i][13]*θs[13]/systems[i].branches.beta[14]+
            #             θ[i][14]*θs[14]/systems[i].branches.beta[24]+θ[i][15]*θs[15]/systems[i].branches.beta[32]);
            #     end
            #     for j in [5,7:13;] # anteriortibial
            #         systems[i].branches.beta[j] *= θ[i][13]*θs[13]/systems[i].branches.beta[14];
            #     end
            #     for j in [6,15,17:23;] # posteriortibial
            #         systems[i].branches.beta[j] *= θ[i][14]*θs[14]/systems[i].branches.beta[24];
            #     end
            #     for j in [16,25:31;] # peroneal
            #         systems[i].branches.beta[j] *= θ[i][15]*θs[15]/systems[i].branches.beta[32];
            #     end
            #     systems[i].branches.beta[14] = θ[i][13]*θs[13];
            #     systems[i].branches.beta[24] = θ[i][14]*θs[14];
            #     systems[i].branches.beta[32] = θ[i][15]*θs[15];
            #     # systems[i].branches.A0[14] = θ[i][9]*θs[9];
            #     # systems[i].branches.A0[24] = θ[i][10]*θs[10];
            #     # systems[i].branches.A0[31] = θ[i][11]*θs[11];
            #     # for j in [14,24,32]
            #     #     systems[i].branches.term[j].R[1] = systems[i].solverparams.rho*
            #     #         systems[i].branches.c0[j][end]/systems[i].branches.A0[j];
            #     # end
            #     for j = 1:length(systems[i].branches.ID)
            #         systems[i].branches.c0[j][end] = sqrt(0.5*systems[i].branches.beta[j]/
            #             systems[i].solverparams.rho)*systems[i].branches.A0[j]^0.25
            #         if isempty(systems[i].branches.children[j])
            #             systems[i].branches.term[j].R[1] = systems[i].solverparams.rho*
            #                 systems[i].branches.c0[j][end]/systems[i].branches.A0[j];
            #         end
            #         # systems[i].branches.beta[j] = sqrt(pi)*systems[i].branches.thicknessinmm[j]*LegHemodynamics.mmTom*
            #             # systems[i].branches.YModinPa[j]/((1-0.5^2)*systems[i].branches.A0[j]);
            #     end
            # end
            #
            # # find new time step to satisfy CFL condition for all members
            # h = [];
            # for i = 1:length(soln)
            #     for j in 1:length(systems[i].branches.ID)
            #         push!(h,systems[i].solverparams.CFL*systems[i].branches.k[j]/
            #             systems[i].branches.c0[j][end]);
            #     end
            # end
            # if minimum(h) != systems[1].solverparams.h
            #     # println("Old time step size: $(systems[1].solverparams.h) s")
            #     for i = 1:length(soln)
            #         systems[i].solverparams.h = minimum(h);
            #     end
            #     # println("New time step size: $(systems[1].solverparams.h) s")
            #     # re-discretize
            #     if systems[1].solverparams.numbeats < systems[1].solverparams.numbeatstotal
            #         systems = pmap((a1,a2)->LegHemodynamics.rediscretizet!(a1,a2),systems,n);
            #     end
            # end
            #
            # # forecast step
            # rflag = "assim";
            # soln = pmap((a1,a2,a3,a4,a5,a6,a7)->LegHemodynamics.advancetime!(a1,a2,a3,a4,a5,a6,a7;
            #     runtype=rflag),systems,spl,times,n,nsamp,term_itr,split_itr);
            # systems = [soln[i][1] for i=1:ensemblesize];
            # n = [soln[i][2] for i=1:ensemblesize];
            # times = [soln[i][3] for i=1:ensemblesize];
            #
            # # forecast state, measurements
            # X = [zeros(nstates) for i=1:ensemblesize];
            # Y = [zeros(nmeas) for i=1:ensemblesize];
            #
            # for i = 1:ensemblesize
            #     X[i][1] = systems[i].branches.Q[14][1,n[1]]/
            #         (systems[i].branches.A[14][1,n[1]]*LegHemodynamics.us);
            #     X[i][2] = systems[i].branches.Q[24][1,n[1]]/
            #         (systems[i].branches.A[24][1,n[1]]*LegHemodynamics.us);
            #     X[i][3] = systems[i].branches.Q[32][1,n[1]]/
            #         (systems[i].branches.A[32][1,n[1]]*LegHemodynamics.us);
            #     Y[i][1] = systems[i].branches.Q[14][1,n[1]]/
            #         (systems[i].branches.A[14][1,n[1]]*LegHemodynamics.us);
            #     Y[i][2] = systems[i].branches.Q[24][1,n[1]]/
            #         (systems[i].branches.A[24][1,n[1]]*LegHemodynamics.us);
            #     Y[i][3] = systems[i].branches.Q[32][1,n[1]]/
            #         (systems[i].branches.A[32][1,n[1]]*LegHemodynamics.us);
            # end
            #
            # # forecast mean state, parameters, measurement
            # xhat = mean(X);
            # yhat = mean(Y);
            # θhat = mean(θ);
            # # println("Normalized ̂x, first forecast: $xhat")
            # println("Normalized ̂y, first forecast: $yhat")
            # println("Normalized ̂θ, first forecast: $θhat")
            #
            # # forecast params./meas. cross covariance, meas. covariance
            # Pty = zeros(length(θhat),length(yhat))
            # Pyy = zeros(length(yhat),length(yhat))
            # for i = 1:length(soln)
            #     Pty += *((θ[i] .- θhat),(Y[i] .- yhat)');
            #     Pyy += *((Y[i] .- yhat),(Y[i] .- yhat)');
            # end
            # Pty ./= (ensemblesize);
            # Pyy ./= (ensemblesize);
            #
            # # println("Parameter-measurement cross-covariance:")
            # # display(Pty)
            #
            # # add meas. noise to meas. covariance (allows invertibility)
            # Pyy += diagm([(errors.odev[1]/LegHemodynamics.us)^2,
            #     (errors.odev[2]/LegHemodynamics.us)^2,
            #     (errors.odev[3]/LegHemodynamics.us)^2],0);
            # # Pyy += diagm([(errors.odev[1]/LegHemodynamics.us)^2],0);
            #
            # # println("Measurement covariance:")
            # # display(Pyy)
            #
            # # parameter Kalman gain
            # K = Pty*inv(Pyy);
            # # println("Parameter Kalman gain:")
            # # display(K)
            #
            # # parameter analysis step
            # for i = 1:length(soln)
            #     # println("ith measurement replicate: $(yi[i])")
            #     # println("ith forecast measurement: $(Y[i])")
            #     θ[i][:] += K*(yi[i] .- Y[i]);
            # end
            #
            # # RTPS parameter covariance inflation
            # for i in 1:length(θ[1])
            #     for j in 1:length(systems)
            #         c[j] = θ[j][i];
            #     end
            #     σta[i] = std(c;corrected=true);
            # end
            # θhat = mean(θ);
            # for i = 1:length(soln)
            #     θ[i] .= θ[i] .+ p.*((σtb.-σta)./σta).*(θ[i].-θhat);
            # end
            #
            # # analysis parameters back into ensemble members
            # for i = 1:length(soln)
            #     for j = 1:nparams
            #         if θ[i][j]*θs[j] <= errors.lb[j]
            #             # println("Warning: analysis θ_$j below lower bound for member $i.
            #                 # Setting to lower bound of $(errors.lb[j]).")
            #             θ[i][j] = errors.lb[j]/θs[j];
            #         end
            #         if θ[i][j]*θs[j] >= errors.ub[j]
            #             # println("Warning: analysis θ_$j above upper bound for member $i.
            #                 # Setting to upper bound of $(errors.ub[j]).")
            #             θ[i][j] = errors.ub[j]/θs[j];
            #         end
            #     end
            #     systems[i].branches.term[14].R[2] = θ[i][1]*θs[1];
            #     systems[i].branches.term[24].R[2] = θ[i][2]*θs[2];
            #     systems[i].branches.term[32].R[2] = θ[i][3]*θs[3];
            #     systems[i].branches.term[14].C[1] = θ[i][4]*θs[4];
            #     systems[i].branches.term[24].C[1] = θ[i][5]*θs[5];
            #     # systems[i].branches.term[24].C[1] = 2.5e-11;
            #     systems[i].branches.term[32].C[1] = θ[i][6]*θs[6];
            #     systems[i].branches.termscalings[1] = θ[i][7]*θs[7];
            #     systems[i].branches.termscalings[2] = θ[i][8]*θs[8];
            #     systems[i].branches.termscalings[3] = θ[i][9]*θs[9];
            #     systems[i].branches.termscalings[4] = θ[i][10]*θs[10];
            #     systems[i].branches.termscalings[5] = θ[i][11]*θs[11];
            #     systems[i].branches.termscalings[6] = θ[i][12]*θs[12];
            #     # k = 1;
            #     for j in [2,4;] # popliteal
            #         systems[i].branches.term[j].C[1] *= 1/3*(systems[i].branches.termscalings[1]+
            #             systems[i].branches.termscalings[3]+systems[i].branches.termscalings[5]);
            #         systems[i].branches.term[j].R[2] *= 1/3*(systems[i].branches.termscalings[2]+
            #             systems[i].branches.termscalings[4]+systems[i].branches.termscalings[6]);
            #         # systems[i].branches.term[j].C[1] = 1/3*(systems[i].branches.termscalings[1]+
            #             # systems[i].branches.termscalings[3]+systems[i].branches.termscalings[5])*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = 1/3*(systems[i].branches.termscalings[2]+
            #             # systems[i].branches.termscalings[4]+systems[i].branches.termscalings[6])*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [7,9,11,13;] # anterior
            #         systems[i].branches.term[j].C[1] *= systems[i].branches.termscalings[1];
            #         systems[i].branches.term[j].R[2] *= systems[i].branches.termscalings[2];
            #         # systems[i].branches.term[j].C[1] = systems[i].branches.termscalings[1]*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = systems[i].branches.termscalings[2]*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [17,19,21,23;] # posterior
            #         systems[i].branches.term[j].C[1] *= systems[i].branches.termscalings[3];
            #         systems[i].branches.term[j].R[2] *= systems[i].branches.termscalings[4];
            #         # systems[i].branches.term[j].C[1] = systems[i].branches.termscalings[3]*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = systems[i].branches.termscalings[4]*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j in [25,27,29,31;] # peroneal
            #         systems[i].branches.term[j].C[1] *= systems[i].branches.termscalings[5];
            #         systems[i].branches.term[j].R[2] *= systems[i].branches.termscalings[6];
            #         # systems[i].branches.term[j].C[1] = systems[i].branches.termscalings[5]*Cdefault[k];
            #         # systems[i].branches.term[j].R[2] = systems[i].branches.termscalings[6]*Rdefault[k];
            #         # k+=1;
            #     end
            #     for j = 1:length(term_itr[1])
            #         if term_itr[i][j] != 14 && term_itr[i][j] != 24 && term_itr[i][j] != 32
            #             # systems[i].branches.term[term_itr[i][j]].C[1] *= systems[i].branches.termscalings[1];
            #             # if systems[i].branches.term[term_itr[i][j]].C[1] < 1e-11
            #             #     systems[i].branches.term[term_itr[i][j]].C[1] = 1e-11;
            #             # end
            #             if systems[i].branches.term[term_itr[i][j]].C[1] < 0.5*Cdefault[j]
            #                 # println("Warning: Analysis terminal C < lower bound for member $i,
            #                     # artery $j. Setting to 1e-11.")
            #                 systems[i].branches.term[term_itr[i][j]].C[1] = 0.5*Cdefault[j];
            #             elseif systems[i].branches.term[term_itr[i][j]].C[1] < 1e-12
            #                 systems[i].branches.term[term_itr[i][j]].C[1] = 1e-12;
            #             elseif systems[i].branches.term[term_itr[i][j]].C[1] > 2*Cdefault[j];
            #                 systems[i].branches.term[term_itr[i][j]].C[1] = 2*Cdefault[j];
            #             end
            #             # systems[i].branches.term[term_itr[i][j]].R[2] *= systems[i].branches.termscalings[2];
            #             if systems[i].branches.term[term_itr[i][j]].R[2] < 0.5*Rdefault[j]
            #                 systems[i].branches.term[term_itr[i][j]].R[2] = 0.5*Rdefault[j];
            #             elseif systems[i].branches.term[term_itr[i][j]].R[2] > 2*Rdefault[j]
            #                 systems[i].branches.term[term_itr[i][j]].R[2] = 2*Rdefault[j];
            #             end
            #         end
            #     end
            #     for j in [1:4;] # popliteal
            #         systems[i].branches.beta[j] *= 1/3*(θ[i][13]*θs[13]/systems[i].branches.beta[14]+
            #             θ[i][14]*θs[14]/systems[i].branches.beta[24]+θ[i][15]*θs[15]/systems[i].branches.beta[32]);
            #     end
            #     for j in [5,7:13;] # anteriortibial
            #         systems[i].branches.beta[j] *= θ[i][13]*θs[13]/systems[i].branches.beta[14];
            #     end
            #     for j in [6,15,17:23;] # posteriortibial
            #         systems[i].branches.beta[j] *= θ[i][14]*θs[14]/systems[i].branches.beta[24];
            #     end
            #     for j in [16,25:31;] # peroneal
            #         systems[i].branches.beta[j] *= θ[i][15]*θs[15]/systems[i].branches.beta[32];
            #     end
            #     systems[i].branches.beta[14] = θ[i][13]*θs[13];
            #     systems[i].branches.beta[24] = θ[i][14]*θs[14];
            #     systems[i].branches.beta[32] = θ[i][15]*θs[15];
            #     # systems[i].branches.A0[14] = θ[i][9]*θs[9];
            #     # systems[i].branches.A0[24] = θ[i][10]*θs[10];
            #     # systems[i].branches.A0[31] = θ[i][11]*θs[11];
            #     for j = 1:length(systems[i].branches.ID)
            #         # systems[i].branches.beta[j] = sqrt(pi)*systems[i].branches.thicknessinmm[j]*LegHemodynamics.mmTom*
            #             # systems[i].branches.YModinPa[j]/((1-0.5^2)*systems[i].branches.A0[j]);
            #         systems[i].branches.c0[j][end] = sqrt(0.5*systems[i].branches.beta[j]/
            #             systems[i].solverparams.rho)*systems[i].branches.A0[j]^0.25
            #         if isempty(systems[i].branches.children[j])
            #             systems[i].branches.term[j].R[1] = systems[i].solverparams.rho*
            #                     systems[i].branches.c0[j][end]/systems[i].branches.A0[j];
            #         end
            #     end
            #     # for j in [14,24,32]
            #     #     systems[i].branches.term[j].R[1] = systems[i].solverparams.rho*
            #     #             systems[i].branches.c0[j][end]/systems[i].branches.A0[j];
            #     # end
            # end

            # recalculate and output parameter averages (non-dimensional)
            θhat = mean(θ);
            append!(θout,[θhat])
            # println("Normalized analysis parameter averages: $θhat")

            # analysis parameter variance (for post-processing)
            Ptt = zeros(length(θhat))
            for i = 1:length(soln)
                Ptt += (θ[i] .- θhat).^2;
            end
            Ptt ./= ensemblesize;
            append!(Pθout,[Ptt])

            # # find new time step to satisfy CFL condition for all members
            # h = [];
            # for i = 1:length(soln)
            #     for j in 1:length(systems[i].branches.ID)
            #         push!(h,systems[i].solverparams.CFL*systems[i].branches.k[j]/
            #             systems[i].branches.c0[j][end]);
            #     end
            # end
            # if minimum(h) != systems[1].solverparams.h
            #     # println("Old time step size: $(systems[1].solverparams.h) s")
            #     for i = 1:length(soln)
            #         systems[i].solverparams.h = minimum(h);
            #     end
            #     println("New time step size: $(systems[1].solverparams.h) s")
            #     # re-discretize
            #     if systems[1].solverparams.numbeats < systems[1].solverparams.numbeatstotal
            #         systems = pmap((a1,a2)->LegHemodynamics.rediscretizet!(a1,a2),systems,n);
            #     end
            # end
            #
            # # corrected forecast w/ analysis parameters
            # rflag = "assim";
            # soln = pmap((a1,a2,a3,a4,a5,a6,a7)->LegHemodynamics.advancetime!(a1,a2,a3,a4,a5,a6,a7;
            #     runtype=rflag),systems,spl,times,n,nsamp,term_itr,split_itr);
            # systems = [soln[i][1] for i=1:ensemblesize];
            # n = [soln[i][2] for i=1:ensemblesize];
            # times = [soln[i][3] for i=1:ensemblesize];

            # second forecast state, measurements
            X = [zeros(nstates) for i=1:ensemblesize];
            Y = [zeros(nmeas) for i=1:ensemblesize];

            for i = 1:ensemblesize
                X[i][1] = systems[i].branches.Q[14][1,n[1]]/
                    (systems[i].branches.A[14][1,n[1]]*LegHemodynamics.us);
                X[i][2] = systems[i].branches.Q[24][1,n[1]]/
                    (systems[i].branches.A[24][1,n[1]]*LegHemodynamics.us);
                X[i][3] = systems[i].branches.Q[32][1,n[1]]/
                    (systems[i].branches.A[32][1,n[1]]*LegHemodynamics.us);
                Y[i][1] = systems[i].branches.Q[14][1,n[1]]/
                    (systems[i].branches.A[14][1,n[1]]*LegHemodynamics.us);
                Y[i][2] = systems[i].branches.Q[24][1,n[1]]/
                    (systems[i].branches.A[24][1,n[1]]*LegHemodynamics.us);
                Y[i][3] = systems[i].branches.Q[32][1,n[1]]/
                    (systems[i].branches.A[32][1,n[1]]*LegHemodynamics.us);
            end

            # second forecast mean state, measurement
            xhat = mean(X);
            yhat = mean(Y);
            # println("Normalized ̂x, second forecast: $xhat")
            # println("Normalized ̂y, second forecast: $yhat")

            # output second forecast measurement
            append!(yout,[yhat])
            append!(innov,[yhat-y])

            # output ensemble average state
            append!(xout,[xhat])

            # state variance (for post-processing)
            Pxx = zeros(length(xhat))
            for i = 1:length(soln)
                Pxx += (X[i] .- xhat).^2;
            end
            Pxx ./= ensemblesize;
            append!(Pxout,[Pxx])

            # state 2-sigma quantiles (for post-processing)
            for i = 1:length(xhat)
                xq = Float64[];
                for j = 1:length(soln)
                    push!(xq,X[j][i])
                end
                q = quantile(xq,[0.025,0.975]);
                if i == 1
                    append!(lbx,[ones(1)*q[1]])
                    append!(ubx,[ones(1)*q[2]])
                else
                    push!(lbx[end],q[1])
                    push!(ubx[end],q[2])
                end
            end

            # output measurement time
            push!(tout,systems[1].t[n[1]])
        end
    end
    tt = toc();
    for i = 1:ensemblesize
        times[i].tt += tt;
    end
end

# reshape ensemble averages into vectors of individual time series
if assimflag == "yes"
    for i in 1:length(xout[1]) # i indexes state variables
        xo = Float64[];
        Pxo = Float64[];
        lb = Float64[];
        ub = Float64[];
        for j in 1:length(xout) # j indexes time steps
            push!(xo,xout[j][i])
            push!(Pxo,Pxout[j][i])
            push!(lb,lbx[j][i])
            push!(ub,ubx[j][i])
        end
        append!(xoutv,[xo])
        append!(Pxoutv,[Pxo])
        append!(lbxv,[lb])
        append!(ubxv,[ub])
    end

    for i in 1:length(θout[1]) # i indexes state variables
        θo = Float64[];
        Pθo = Float64[];
        for j in 1:length(θout) # j indexes time steps
            push!(θo,θout[j][i])
            push!(Pθo,Pθout[j][i])
        end
        append!(θoutv,[θo])
        append!(Pθoutv,[Pθo])
    end

    for i in 1:length(yout[1])
        yo = Float64[];
        for j in 1:length(yout)
            push!(yo,yout[j][i])
        end
        append!(youtv,[yo])
    end
end

if saveflag == "yes"
    if assimflag == "no"
        file = MAT.matopen(savefile, "w")
        write(file,"system",system)
        close(file)
    elseif assimflag == "yes"
        for i=1:ensemblesize
            file = MAT.matopen(savefiles[i],"w");
            write(file,"system",systems[i])
            close(file)
        end
        vnames = ["t" "x" "Px" "theta" "thetascale" "Pt" "lb" "ub"];
        for i in 1:length(vnames)
            file = MAT.matopen("$(vnames[i])_converged_small_dev_1.mat","w");
            if vnames[i] == "t"
                write(file,"t",tout)
            elseif vnames[i] == "x"
                write(file,"x",xoutv)
            elseif vnames[i] == "Px"
                write(file,"Px",Pxoutv)
            elseif vnames[i] == "theta"
                write(file,"theta",θoutv)
            elseif vnames[i] == "thetascale"
                write(file,"thetascale",θs)
            elseif vnames[i] == "Pt"
                write(file,"Pt",Pθoutv)
            elseif vnames[i] == "lb"
                write(file,"lb",lbxv)
            elseif vnames[i] == "ub"
                write(file,"ub",ubxv)
            end
            close(file)
        end
    end
end

if assimflag == "no"
    if timeflag == "yes"
        return system, n, times
    elseif timeflag == "no"
        return system, n
    end
elseif assimflag == "yes"
    return n,tout,xoutv,youtv,ioutv,θoutv,Pθoutv,Pxoutv,lbxv,ubxv,θs
end

end


export build_momentrelaxation

"""
    momentrelaxation = build_momentrelaxation(relax_ctx, problem, moment_param::Dict{String, Tuple{Set{String}, Int}}, max_cliques::Dict{String, Set{Variable}})

    Compute the `momentrelaxation` of `problem` corresponding to the clique decomposition `max_cliques` and parameters `moment_param`.
"""
function build_momentrelaxation(relax_ctx::RelaxationContext,
                                problem::Problem,
                                momentmat_param::Dict{String, Int},
                                localizingmat_param::Dict{String, Tuple{Set{String}, Int}},
                                max_cliques::Dict{String, Set{Variable}};
                                T=Float64)

    var_to_cliques = Dict{Variable, Set{String}}()
    for (clique, vars) in max_cliques
        for var in vars
            haskey(var_to_cliques, var) || (var_to_cliques[var] = Set{String}())
            push!(var_to_cliques[var], clique)
        end
    end

    ## Building linear-in-moments objective
    objective = Dict{Moment, T}()
    for (expo, val) in problem.objective
        clique = get_exponentclique(expo, var_to_cliques)
        objective[Moment(expo, clique)] = val
    end

    ## Building linear matrix inequalities
    momentmatrices = Dict{Tuple{String, String}, MomentMatrix{T}}()

    ## Build moment matrix
    for (cliquename, vars) in max_cliques
        dcl = momentmat_param[cliquename]
        momentmatrices[(get_momentcstrname(), cliquename)] = MomentMatrix{T}(relax_ctx, vars, dcl, relax_ctx.symmetries,
                                                                                relax_ctx.cstrtypes[get_momentcstrname()],
                                                                                default_clique = cliquename)
    end

    ## Build localizing matrices
    for (cstrname, cstr) in problem.constraints

        cstrtype = get_cstrtype(cstr)
        if cstrtype == :ineqdouble
            cstrname_lo, cstrname_up = get_cstrname_lower(cstrname), get_cstrname_upper(cstrname)

            # Deal with lower inequality
            clique_keys, order = localizingmat_param[cstrname_lo]
            vars, cliquename = collect_cliquesvars(clique_keys, max_cliques)
            # length(clique_keys) == 1 || error(LOGGER, "build_momentrelaxation(): constraint $cstrname spans several cliques ($clique_keys).\nNot supported yet.")

            mmt = MomentMatrix{T}(relax_ctx, vars, order, relax_ctx.symmetries,
                                                          relax_ctx.cstrtypes[cstrname_lo],
                                                          var_to_cliques = var_to_cliques)
            # print_with_color(:green, "$cstrname, :Lo\n") ##NOTE: find better logging system.
            product!(mmt, cstr.p - cstr.lb, var_to_cliques)
            momentmatrices[(cstrname_lo, cliquename)] = mmt

            # Deal with upper inequality
            clique_keys, order = localizingmat_param[cstrname_up]
            vars, cliquename = collect_cliquesvars(clique_keys, max_cliques)
            # length(clique_keys) == 1 || error(LOGGER, "build_momentrelaxation(): constraint $cstrname spans several cliques ($clique_keys).\nNot supported yet.")

            mmt = MomentMatrix{T}(relax_ctx, vars, order, relax_ctx.symmetries,
                                                          relax_ctx.cstrtypes[cstrname_up],
                                                          var_to_cliques = var_to_cliques)
            # print_with_color(:green, "$cstrname, :Up\n") ##NOTE: find better logging system.
            product!(mmt, cstr.ub - cstr.p, var_to_cliques)
            momentmatrices[(cstrname_up, cliquename)] = mmt

            # # Deal with upper inequality, no recomputing of variables or moment matrix if possible
            # clique_keys_up, order_up = localizingmat_param[cstrname_up]
            # if collect(clique_keys) != collect(clique_keys_up)
            #     warn(LOGGER, "clique keys different from lower and upper side of double constraint")
            #     length(clique_keys_up) == 1 || error(LOGGER, "build_momentrelaxation(): constraint $cstrname spans several cliques ($clique_keys).\nNot supported yet.")
            #     vars, cliquename = collect_cliquesvars(clique_keys_up, max_cliques)

            #     mmt = MomentMatrix(relax_ctx, vars, order_up, relax_ctx.symmetries,
            #                                                   relax_ctx.cstrtypes[cstrname_hi],
            #                                                   var_to_cliques = var_to_cliques)
            # elseif order_up != order
            #     warn(LOGGER, "order different from lower and upper side of double constraint")
            #     mmt = MomentMatrix(relax_ctx, vars, order_up, relax_ctx.symmetries,
            #                                                   relax_ctx.cstrtypes[cstrname_hi],
            #                                                   var_to_cliques = var_to_cliques)
            # end

        else
            # either cstrtype == :ineqlo, :ineqhi, :eq
            clique_keys, order = localizingmat_param[get_cstrname(cstrname, cstrtype)]
            vars, cliquename = collect_cliquesvars(clique_keys, max_cliques)
            # length(clique_keys) == 1 || error(LOGGER, "build_momentrelaxation(): constraint $cstrname spans several cliques ($clique_keys).\nNot supported yet.")

            mmt = MomentMatrix{T}(relax_ctx, vars, order, relax_ctx.symmetries,
                                                          relax_ctx.cstrtypes[get_cstrname(cstrname, cstrtype)],
                                                          var_to_cliques = var_to_cliques)

            # print_with_color(:green, "$cstrname, :Up\n") ##NOTE: find better logging system.
            product!(mmt, get_normalizedpoly(cstr, cstrtype), var_to_cliques)

            momentmatrices[(get_cstrname(cstrname, cstrtype), cliquename)] = mmt
        end
    end

    ## Locate clique overlapping moments
    expo_to_cliques = Dict{Exponent, Set{String}}()

    # Collect Exponents per clique (moment matrix)
    for ((ctrobj, clique), mmtmat) in momentmatrices
        if ctrobj == get_momentcstrname()

            for (key, momentpoly) in mmtmat.mm
                for (moment, coeff) in momentpoly
                    expo = product(moment.conj_part, moment.expl_part)
                    haskey(expo_to_cliques, expo) || (expo_to_cliques[expo] = Set{String}())
                    push!(expo_to_cliques[expo], moment.clique)
                end
            end
        end
    end

    nb_expos = length(expo_to_cliques)
    for (expo, cliques) in expo_to_cliques
        length(cliques) > 1 || delete!(expo_to_cliques, expo)
    end

    momentrelaxation = SDPDual{T}(objective, momentmatrices, expo_to_cliques)

    print_build_momentrelax(relax_ctx, momentrelaxation, nb_expos)

    return momentrelaxation
end


# Convenience wrapper
save_subdomain_temperature!(m, state, aux) =
    save_subdomain_temperature!(m, m.moisture, state, aux)

using KernelAbstractions: @print

"""
    save_subdomain_temperature!(
        m::AtmosModel,
        moist::EquilMoist,
        state::Vars,
        aux::Vars,
    )

Updates the subdomain sensible temperature, given:
 - `m`, an `AtmosModel`
 - `moist`, an `EquilMoist` model
 - `state`, state variables
 - `aux`, auxiliary variables
"""
function save_subdomain_temperature!(
    m::AtmosModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
)
    N_up = n_updrafts(m.turbconv)
    ts = recover_thermo_state(m, state, aux)
    ts_up = new_thermo_state_up(m, state, aux, ts)
    ts_en = new_thermo_state_en(m, state, aux, ts)

    @unroll_map(N_up) do i
        aux.turbconv.updraft[i].T = air_temperature(ts_up[i])
    end
    aux.turbconv.environment.T = air_temperature(ts_en)
    return nothing
end

# No need to save temperature for DryModel.
function save_subdomain_temperature!(
    m::AtmosModel,
    moist::DryModel,
    state::Vars,
    aux::Vars,
) end


using WebIO
using JSExpr

function hue_app()
    scope = Scope()
    import!(scope, ["//cdnjs.cloudflare.com/ajax/libs/p5.js/0.5.7/p5.js"])

    # Note: we explicitly qualify @js, @var, and @new with the JSExpr module
    # because Blink.jl exports macros with the same name.
    sketch = JSExpr.@js function (p5)
        JSExpr.@var s = function(p)
            JSExpr.@var barWidth = 20
            JSExpr.@var lastBar = -1

            p.setup = function ()
               p.createCanvas(720, 400)
               p.colorMode(p.HSB, p.height, p.height, p.height)
               p.noStroke()
               p.background(0)
            end

            p.draw = function ()
                JSExpr.@var whichBar = p.mouseX / barWidth
                JSExpr.@var barX = whichBar * barWidth;
                if whichBar != lastBar
                   p.fill(p.mouseY, p.height, p.height)
                   p.rect(barX, 0, barWidth, p.height)
                   lastBar = whichBar
               end
            end
        end
        this.dom.querySelector("#p5container").innerText = "";
        JSExpr.@new p5(s, "p5container");
    end
    onimport(scope, sketch)

    scope.dom = dom"div#p5container"("Loading p5...")

    scope
end

# Display in whatever frontend is avalaible
function main()
    if @isdefined(IJulia) || @isdefined(Juno)
        return hue_app()
    elseif @isdefined(Blink)
        win = Window()
        body!(win, hue_app())
        win
    elseif @isdefined(Mux)
        @sync webio_serve(page("/", req -> hue_app()), rand(8000:9000))
    else
        error("do one of using Mux, using Blink before running the
               example, or run it from within IJulia or Juno")
    end
end

main()


#----------------------------------------------------------------------------
# load data from Imperial College
struct Data
    turing_data
    cases
    deaths
    hospit
    num_observations
    num_total
    dates
    pop
    predictors
    seroprev
end

dan2eng = Dict(
    "Dato"          => "date",
    "Total"         => "total",
    "Hovedstaden"   => "capital",
    "Sj\xe6lland"   => "zealand",
    "Sjælland"      => "zealand",
    "Syddanmark"    => "south",
    "Midtjylland"   => "central",
    "Nordjylland"   => "north",
    "Ukendt Region" => "unknown"
)

function missing2zero(xs)
    i = findfirst(!ismissing, xs)
    xs[1:i] .= 0
    convert(Vector{Int64}, xs)
end
# weekly_average(xs) = [mean(xs[i:(i+6)]) for i in 1:div(length(xs),7)]
# weekly_average(deaths)

function padlast(xs, l)
    n  = length(xs)
    i  = findlast(x -> typeof(x) <: Real, xs) #avoid catching a missing value
    ys = Vector{Union{Float64, Missing}}(undef, l)
    ys[1:n] = xs
    ys[n+1:end] .= xs[i]
    ys
end

function padmissing(xs, l)
    n  = length(xs)
    ys = Vector{Union{eltype(xs), Missing}}(missing, l)
    ys[1:n] = xs
    ys
end

function cum2counts(x)
    x = x |> missing2zero |> diff
    pushfirst!(x, 0)
    x[x.<0] .= 0 # ignore data inconsistencies
    x
end

# plotnpi(x) = plot(dk.date, x; ticks=:native)
# dk.cancel_events|>plotnpi
# dk.internal_movement_restrictions|>plotnpi
# dk.gatherings_restrictions|>plotnpi
# dk.stay_home_restrictions|>plotnpi
# dk.transport_closing|>plotnpi
# dk.workplace_closing|>plotnpi
# dk.school_closing|>plotnpi
# # dk.stringency_index|>plotnpi

const dformat1 = DateFormat("d u Y")
const dformat2 = DateFormat("d-u-Y")
function parsetime(ds)
    ds = replace(ds, "Mrz"=>"Mar")
    ds = replace(ds, "Dez"=>"Dec")
    ds = replace(ds, "."=>"")
    d = try
        Date(ds, dformat1)
    catch
        try
            Date(ds, dformat2)
        catch
            @error "Not a valid time format: $ds"
        end
    end
    d + Year(2000)
end

function znorm!(xs)
    xs .-= mean(skipmissing(xs))
    xs ./= std(skipmissing(xs))
    nothing
end

function minmaxnorm!(xs)
    xs .-= minimum(skipmissing(xs))
    xs ./= maximum(skipmissing(xs))
    nothing
end

function surveynorm!(xs)
    i = findfirst(x->!ismissing(x), xs)
    xs ./= abs(xs[i])
    xs .-= xs[i]
    nothing
end

function preproc(xs::AbstractVector, num_obs, Ntot; normalize=nothing)
    mask = 1:num_obs
    xs = padlast(xs[mask], Ntot)
    !isnothing(normalize) && normalize(xs)
    xs
end

function select_predictors(dk, predictors, num_obs, Ntot, surveyvars)
    predictors = String.(predictors)

    #preprocess mobility data
    dk[!,:nonresidential_mobility] = preproc(
        (
            dk.retail_and_recreation_percent_change_from_baseline +
            dk.grocery_and_pharmacy_percent_change_from_baseline +
            dk.workplaces_percent_change_from_baseline
        ) / 3 / 100 , num_obs, Ntot; normalize=nothing)
    # begin
    #     Mavg = (
    #     preproc(dk.retail_and_recreation_percent_change_from_baseline, num_obs, Ntot) +
    #     preproc(dk.grocery_and_pharmacy_percent_change_from_baseline, num_obs, Ntot) +
    #     preproc(dk.workplaces_percent_change_from_baseline, num_obs, Ntot)
    #     ) ./ -3
    # end
    dk[!,:residential_mobility] = preproc(dk.residential_percent_change_from_baseline, num_obs, Ntot; normalize=nothing) ./ 100

    #preprocess survey data
    for k in surveyvars
        dk[!,k] = preproc(dk[:,k], num_obs, Ntot; normalize=surveynorm!)
    end

    dk[!, :lockdown] .= 0.
    t = findfirst(==(Date("2020-03-14")), dk.date) # lockdown date in DK: 11. (partial), 17. (full)
    dk[t:end, :lockdown] .+= 1
    # select predictors
    miss = .!( in.(predictors, Ref(names(dk)) ) )
    any(miss) && (@error "predictor $(predictors[miss]) not found. Select either nonresidential_mobility, residential_mobility or from survey: $(surveyvars)")
    Array( dk[!,predictors] )
end

rolling(xs) = begin
    n  = length(xs)
    ys = Vector{Union{Missing, Float64}}(missing, n)
    ys[1:n-6] = [maximum(xs[i:i+6]) for i in 1:n-6]
    ys
end

function gamma(μ, cv)
    α = 1/cv^2
    θ = μ/α
    Gamma(α, θ)
end

function inf2hosp(nbins=50, ihr = 1)
    N = 10_000_000
    # SSI: ihr ∈ 0.05%-0.5%
    # derived from Verity et al. See parameters.jl for details
    # incubatine time ~ gamma(5.1, .86) see Flaxman et al
    # onset to hospitalization see Zhang et al. Lancet (2020)
    x = rand(gamma(5.1, .86), N) .+ rand(gamma(2.6, .92), N)
    h = fit(Histogram, x, vcat([0], 1.5:1:(nbins + .5)))
    h.weights / N * ihr
end

function inf2death(nbins = 75, ifr = 1)
    N = 10_000_000
    # ifr = 89 / 100_000 for age < 70 see Erikstrup et al. in Clin. Inf. Dis.
    # report34: onset to death = 19.8, Cv = .85
    r = vcat([0], 1.5:1:(nbins + .5))
    x = rand(gamma(5.1, .86), N) .+ rand(gamma(19.8, .85), N) #rand(gamma(17.8, .45), N)
    h = fit(Histogram, x, r)
    h.weights / N * ifr
end

function seropositive(nbins)
    N = 10_000_000
    # For details see Imperial Report34
    rbins = 1.5:1:(nbins + .5)
    inf2symp = rand(gamma(5.1, .86), N)
    inf2conv = inf2symp .+ rand(Exponential(13.3), N)
    inf2rev = inf2conv .+ rand(Weibull(3.67, 132.88), N)
    seroconverted = map( ecdf(symp2conv), rbins)
    notseroreverted = 1 .- map( ecdf(symp2rev), rbins)
    serconverted .* notseroreverted
end

function inf2case(nbins = 30)
    # ifr = 89 / 100_000 for age < 70 see Erikstrup et al. in Clin. Inf. Dis.
    r = vcat([0], 1.5:1:(nbins + .5))
    diff(cdf.(gamma(5.1, .86), r))
end

function serialinterval(nbins = 30)
    r = vcat([0], 1.5:1:(nbins + .5))
    diff(cdf.(gamma(6.5, .62), r))
end

# function inf2serocon(nbins)
#     N = 10_000_000
#     # onset to seroconversion time from Imperial Report 34
#     x = rand(gamma(5.1, .86), N) .+ rand(Exponential(18.3), N)
#     h = fit(Histogram, x, vcat([0], 1.5:1:150.5))
#     h.weights / N * ihr
# end

function aggregateregions!(df, mun2reg)
  regions = values(mun2reg) |> unique
  for r in regions
    df[!, r] .= 0
  end
  for m in keys(mun2reg)
    df[!, mun2reg[m]] += df[!,m]
  end
end

# function getcases()
#     municipalities = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/municipalities.csv"); delim="\t", )|>DataFrame
#     mun2reg = Dict(municipalities.municipality[i] => municipalities.region[i]
#                     for i in 1:size(municipalities,1))
#     cases    = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/Municipality_cases_time_series.csv") )|>DataFrame
#     cases = select(cases, municipalities.municipality, :date_sample => :date)
#     cases[!, :total] = sum(select(cases, municipalities.municipality)|>eachcol)
#     aggregateregions!(cases, mun2reg)
#     cases
# end

function insert_missing_dates!(df)
    ts = df.date[1] : Day(1) : df.date[end]
    missing_date = .! in.(ts, Ref(df.date))
    if any(missing_date)
        for t in ts[missing_date]
            insertrow!(df, t)
        end
    end
end

function replace_missing!(df, v=0)
    @assert !any( ismissing.( df.date ) )
    for c in eachcol(df)
        replace!(c, missing=>v)
    end
    disallowmissing!(df)
end

function process_ssi(df)
    rename!(df, 1=>:region, 2=>:date, 3=>:counts)
    df.region = replace(df.region, dan2eng...)
    dropmissing!(df, :region)
    df = unstack(df, :date, :region, :counts)
    insert_missing_dates!(df)
    replace_missing!(df)
    DataFrames.transform(df, AsTable(Not(:date)) => sum => :country)
end

# function getdeaths()
#     deaths   = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/Deaths_over_time.csv"))|>DataFrame
#     deaths = deaths[1:end-1,:]
#     rename!(deaths, "Dato"=>"date", "Antal_døde"=>"total")
#     deaths.date = Date.(deaths.date, Ref(dateformat"y-m-d"))
#     deaths
# end
#
# function gethospit()
#     hosp     = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/Newly_admitted_over_time.csv"))|>DataFrame
#     dan2eng = Dict(
#         "Dato"          => "date",
#         "Total"         => "total",
#         "Hovedstaden"   => "capital",
#         "Sjælland"      => "zealand",
#         "Syddanmark"    => "south",
#         "Midtjylland"   => "central",
#         "Nordjylland"   => "north",
#         "Ukendt Region" => "unknown"
#     )
#     rename!(hosp, dan2eng)
#     hosp
# end

function startdate(deaths)
    i = findfirst(>=(10), skipmissing(deaths.country))
    deaths.date[i] - Day(30)
end

enddate(df1, df2) = min(df1.date[end], df2.date[end])
enddate(df, args...) = min(df.date[end], enddate(args...))

function padzeros(xs::AbstractVector, padfirst=0, total=length(xs)+padfirst)
    n  = length(xs)
    ys = zeros(eltype(xs), total)
    ys[padfirst+1:n+padfirst] = xs
    ys
end

function insertrow!(df, d::Date, v=0)
    i = findfirst(==(d-Day(1)), df.date)
    v = [eltype(c) == Date ? d : v for c in eachcol(df)]
    insert!.(eachcol(df), i+1, v)
    nothing
end

function padzeros(df::DataFrame, s::Date)
    npad = (df.date[1] - s).value
    dfpad = DataFrame( :date => s:Day(1):df.date[end] )
    for n in names(df)
        n == "date" && continue
        dfpad[!,n] = padzeros(df[!,n], npad)
    end
    dfpad
end

function limit(df::DataFrame, s::Date, e::Date, pad=true)
    pad && df.date[1] > s && ( df = padzeros(df, s) )
    is = findfirst(==(s), df.date)
    ie = findlast(==(e), df.date)
    df[is:ie, :]
end

function parseweek(s)
    _,w = split(s)
    w = parse(Int, w)
    y = w > 40 ? Date("2020") : Date("2021")
    y + Week(w) - Day(3) #center at day 4
end

# ≈ 15days delay from infection to antibody production
# see Ferguson Report and Zhao 2020: "Antibody responses to SARS-CoV-2 in patients of novel coronavirus disease 2019"
function readseroprev(dates, delay=15)
    df0 = CSV.File(projectdir("data/seroprevalence/", "seroprev_early.csv"), delim="\t")|>DataFrame
    # df1 = CSV.File(projectdir("data/seroprevalence/", "seroprev.csv"), delim="\t")|>DataFrame
    # df1[!,:date] = df1.week .|> parseweek
    # select!(df1, Not(:week))
    # df = vcat(df0, df1)
    df = df0
    df = df[ dates[1] .<= df.date .<= dates[end], :]
    df.date .-= Day(delay)
    parsepercent!(df)
    df.CI = df.CI .|> parseCI
    # df[!,:std_u] = [(CIs[i][2] - df.country[i]) / 1.96 for i in 1:length(CIs)]
    # estimates of std are not symmetric with respect to lower / upper CI bound
    # we choose the larger estimate, i.e. the lower bound
    turing_data = (
        std  = [(df.country[i] - df.CI[i][1]) / 1.96 for i in 1:length(df.CI)],
        mean = df.country,
        idx  = [findfirst(==(d), dates) for d in df.date]
    )
    df, turing_data
end

function parsepercent(x::String)
    x = strip(x, '%')
    x = replace(x, "," => ".")
    parse(Float64, x) / 100
end

function parseCI(x)
    x = strip(x, '%')
    x = replace(x, "," => ".")
    lq, uq = split(x, '-')
    parse(Float64, lq) / 100, parse(Float64, uq) / 100
end

function parsepercent!(df)
    regions = ["capital", "zealand", "south", "central", "north"]
    df.country = parsepercent.(df.country)
    for r in regions
        df[!,r] = parsepercent.(df[:,r])
    end
end

function findindices(dates, ds)
    is = Int[]
    for d in ds
        i = findfirst(==(d), dates)
        push!(is, i)
    end
    is
end

function readsurvey(fname)
    s = CSV.File(fname)|>DataFrame
    s.date = s.date .|> parsetime
    dropmissing!(s)
    for n in names(s)
        c = s[:,n]
        if eltype(c) == String
            c = replace.(c, ","=>".")
            s[!,n] = parse.(Ref(Float64), c)
        end
    end
    s
end

function readsurveys_raw()
    s1 = normpath( homedir(), "data/covidsurvey/contacts.csv" )  |> readsurvey #projectdir("data", "survey", "timeseries_colleagues_strangers.csv") |> readsurvey
    s1 = select(s1, Not([:Rt, :uncertainty_lower, :uncertainty_upper]))
    s2 = normpath( homedir(), "data/covidsurvey/predictors.csv" )  |> readsurvey #projectdir("data", "survey", "predictors.csv") |> readsurvey
    s = innerjoin(s1, s2, on=:date)
    s[3:end,:] # skip first for data quality
end
function readsurveys()
    fname = normpath( homedir(), "data/covidsurvey/smoothed_contacts.csv" )
    df = CSV.File(fname)|>DataFrame
    return df
end

country(df) = df.country

function get_iar_idx(num_tot, num_obs, start_idx, iar_step)
    iar_idx = zeros(Int64, num_tot)
    iar_idx[1:start_idx-1] .= 1
    idx = 2
    for (j,i) in enumerate(start_idx:num_obs)
         iar_idx[i] = idx
         j % iar_step == 0 && (idx += 1)
    end
    if num_obs < num_tot
        iar_idx[num_obs+1:num_tot] .= iar_idx[num_obs]
    end
    iar_idx
end

function load_data(
    observations_end=nothing,
    predictors=nothing,
    cases_start=nothing,
    addhospital=false,
    addseroprev=false,
    addtests=false;
    update=false
)
    iar_step = 7

    cases = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/08_bekraeftede_tilfaelde_pr_dag_pr_regions.csv" ) ) |>
        DataFrame |> process_ssi

    deaths = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/07_antal_doede_pr_dag_pr_region.csv" ) ) |>
        DataFrame |> process_ssi

    hospit = CSV.File(normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/06_nye_indlaeggelser_pr_region_pr_dag.csv" ) ) |>
        DataFrame |> process_ssi

    s      = startdate(deaths)
    e      = enddate(deaths, hospit, cases)
    dates  = collect( s : Day(1) : e )
    deaths = limit(deaths, s, e) |> country
    hospit = limit(hospit, s, e) |> country
    cases  = limit(cases, s, e)  |> country
    num_tot= length(dates)

    dk = begin
        update && run(`Rscript $(datadir())/COVID19.R`)
        dk = load(datadir("COVID19", "DK.rds"))
        limit(dk, s, e, false)
    end

    survey = readsurveys()
    surveyvars = filter(x->x!=="date", names(survey))
    dk = leftjoin(dk, survey, on = :date)

    num_obs = if isnothing(observations_end)
        num_tot
    else
        findfirst(==(Date.(observations_end)), dates)
    end

    lockdown = findfirst(==(Date("2020-03-18")), dates)

    num_case_obs, cases_start_idx = if isnothing(cases_start)
        0, nothing
    else
        cases_start_idx = findfirst(==(Date.(cases_start)), dates)
        num_obs - cases_start_idx + 1, cases_start_idx
    end

    iar_idx = get_iar_idx(num_tot, num_obs, cases_start_idx, iar_step)
    num_iar_steps = length(unique(iar_idx))

    turing_data = OrderedDict(
        :num_impute                  => 6,
        :num_total_days              => num_tot,
        :cases                       => cases[1:num_obs], #padmissing(cases[mask], num_tot),
        :deaths                      => deaths[1:num_obs], #padmissing(deaths[mask], num_tot),
        :π                           => padzeros(inf2death(90), 0, num_tot),#[1:100], #length = 100
        :π2                          => padzeros(inf2case(30), 0, num_tot),#[1:30], #WARNING US-NY DATA
        :epidemic_start              => 31,
        :population                  => dk.population[1],
        :serial_intervals            => padzeros(serialinterval(30), 0, num_tot), #data_usa.turing_data.serial_intervals[1:50],
        :num_iar_steps               => num_iar_steps,
        :iar_idx                     => iar_idx,
        :lockdown                    => lockdown,
        :num_case_obs                => num_case_obs,
        )

    if !isnothing(predictors)
        preds = select_predictors(dk, predictors, num_obs, num_tot, surveyvars)
        covariates_start = findfirst(x->!ismissing(x), preds[:,end])
        preds = convert(Array{Float64,2}, preds[covariates_start:end,:] )
        turing_data[:covariates] = preds
        turing_data[:covariates_start] = covariates_start
    end

    if addhospital
        turing_data[:πh] = padzeros(inf2hosp(50, 1), 0,  num_tot) #length 50
        turing_data[:hospit] = hospit[1:num_obs]
    end

    seroprev, turing_seroprev = readseroprev(dates[1:num_obs])
    if addseroprev
        #turing_data[:πs] = padzeros(inf2serocon(50, 1), 0,  num_tot)
        turing_data[:seroprev_mean] = turing_seroprev.mean * dk.population[1]
        turing_data[:seroprev_std]  = turing_seroprev.std * dk.population[1]
        turing_data[:seroprev_idx]  = turing_seroprev.idx
    end

    turing_data = (;turing_data...)
    Data(
        turing_data,
        cases,
        deaths,
        hospit,
        num_obs, #num_observations
        num_tot,
        dates,
        dk.population,
        predictors,
        seroprev,
    )
end


# LSL.jl: Julia interface for Lab Streaming Layer
# Copyright (C) 2019 Samuel Powell

# test/streaminfo.jl: test stream information functions

@testset "StreamInfo" begin

info = StreamInfo(name="testname",
                  type="testtype",
                  channel_count=2,
                  nominal_srate=10,
                  channel_format=Float64,
                  source_id="testid")

@test name(info) == "testname"
@test type(info) == "testtype"
@test channel_count(info) == 2
@test nominal_srate(info) == 10
@test channel_format(info) == Float64
@test source_id(info) == "testid"

@test version(info) == protocol_version()

channels = append_child(desc(info), "channels")
for label in ["C3", "C4", "Cz", "FPz", "POz", "CPz", "O1", "O2"]
  ch = append_child(channels, "channel")
  append_child_value(ch, "label", label)
  append_child_value(ch, "unit", "microvolts")
  append_child_value(ch, "type", "EEG")
end
append_child_value(desc(info), "manufacturer", "SCCN")
cap = append_child(desc(info), "cap")
append_child_value(cap, "name", "EasyCap")
append_child_value(cap, "size", "54")
append_child_value(cap, "labelscheme", "10-20")

end

# Force cleanup
GC.gc()


import DataFrames
import PredictMD
import Test

a = DataFrames.DataFrame()
a[:x] = ["foo", "foo", "foo", "bar", "bar", "bar"]
contrasts = PredictMD.DataFrameFeatureContrasts(a, [:x])
transformer = PredictMD.MutableDataFrame2DecisionTreeTransformer([:x], :y)
PredictMD.set_feature_contrasts!(transformer, contrasts)

Test.@test length(PredictMD.transform(transformer, a)) == 6

b = DataFrames.DataFrame()
b[:x] = ["bar", "foo"]
Test.@test length(PredictMD.transform(transformer, b)) == 2

c = DataFrames.DataFrame()
c[:x] = ["foo", "bar", "baz"]
Test.@test_throws KeyError PredictMD.transform(transformer, c)

d = DataFrames.DataFrame()
Test.@test_throws ErrorException PredictMD.DataFrameFeatureContrasts(d, [:x, :x])


"""
    find_scored_idx(collection, what::String; check::Function=_->true, score::Function=_->0)

A helper function that runs through a `collection` and select an item that
satisfies the predicate `check` with best score computed by `score`. If no
suitable item is found, an error is thrown referring to the item type as
`what`.
"""
function find_scored_idx(
    collection,
    what::String;
    check::Function = _ -> true,
    score::Function = _ -> 0,
)
    found = 0
    found_score = -Inf
    for (idx, obj) in enumerate(collection)
        if check(obj)
            s = score(obj)
            if s > found_score
                found = idx
                found_score = s
            end
        end
    end
    found > 0 || error("No suitable $what found")
    return found
end

"""
    find_physical_device(instance; args...)

Find a physical device within the instance that satisfies some given
properties; arguments are forwarded to [`find_scored_idx`](@ref).

# Example

    physical_device = find_physical_device(instance, check=props->contains(props.device_name, "MyFavVendor"))
"""
function find_physical_device(instance; args...)
    devs = unwrap(enumerate_physical_devices(instance))
    return devs[find_scored_idx(
        get_physical_device_properties.(devs),
        "physical device";
        args...,
    )]
end

"""
    find_queue_family_idx(physical_device; args...)

Find a queue family index (numbered from 0) for a physical device that
satisfies some given properties; arguments are forwarded to
[`find_scored_idx`](@ref).

# Example

    # Find a queue family with most queues that supports compute and transfer operations.
    find_queue_family_idx(physical_device,
        check=qfp->hasbits(qfp.queue_flags, QUEUE_TRANSFER_BIT | QUEUE_COMPUTE_BIT),
        score=qfp->qfp.queue_count,
    )
"""
find_queue_family_idx(physical_device; args...) =
    find_scored_idx(
        get_physical_device_queue_family_properties(physical_device),
        "queue family";
        args...,
    ) - 1

"""
    find_memory_type_idx(physical_device; check::Function=(_,_)->true, score::Function=(_,_)->0)

Find a memory type index (numbered from 0) for a physical device that satisfies
some given properties. The arguments `check` and `score` behave as with
[`find_scored_idx`](@ref), but take 2 arguments with types `MemoryType` and
`MemoryHeap`.

# Example

    # Find a host visible and device local memory type with largest heap
    find_memory_type_idx(physical_device,
        check=(mt, heap) -> hasbits(mt.property_flags, MEMORY_PROPERTY_HOST_VISIBLE_BIT|MEMORY_PROPERTY_DEVICE_LOCAL_BIT),
        score=(mt, heap) -> heap.size,
    )

"""
function find_memory_type_idx(
    physical_device;
    check::Function = (_, _) -> true,
    score::Function = (_, _) -> 0,
)
    props = get_physical_device_memory_properties(physical_device)
    return find_scored_idx(
        props.memory_types[1:props.memory_type_count],
        "memory type";
        check = mt -> check(mt, props.memory_heaps[mt.heap_index+1]),
        score = mt -> score(mt, props.memory_heaps[mt.heap_index+1]),
    ) - 1
end


struct SearchProblem <: AbstractProblem
    patternset::PatternSet
    trace::Vector{Float64}
    simulator::AbstractSimulator
    solver::AbstractSolver
end


const labels = ["A1", "A2", "E", "B1", "B2", "A"]
const label_map = Dict(key => "$num" for (num, key) ∈ enumerate(labels))
# add the extra character
label_map["-"] = ""
const column_types = [Int64, Int64, Int64, Float64, Int64, Float64, Float64]

"""Function to parse in a CSV formatted line list output from PGopher.
Kwargs specify what to skip and where to read the header in: by default,
the options expect to skip the first line, read the header from the second,
startin parsing from the third, and skips the last line.
"""
function read_pgopher_linelist(filepath; header=2, skipto=3, footerskip=1)
    csv_data = CSV.File(filepath, header=header, skipto=skipto, footerskip=footerskip)
    ν = csv_data[:Position]
    A = csv_data[:A]
    encodings = clean_linelist_encoding(csv_data)
    transitions = StructArray([Transition(ν=ν[i], I=A[i], encoding=encodings[:,i]) for i ∈ eachindex(ν, A)])
    return transitions
end

"""Parse through an energy level list output of PGopher.
"""
function read_pgopher_levels(filepath)
    open(filepath) do file
        lines = readlines(file)[3:end-3]
        return StructArray(parse_level_row.(lines))
    end
end

"""Not intended for public consumption, but encapsulating the work done
on a row actually helps the eventual `StructArray` creation infer
the correct types, which should give downstream performance improvements
(instead of just `Vector{Any}`).
"""
function parse_level_row(row)
    split_line = replace_sym_labels(row) |> split
    # for remaining columns, parse as integers
    remaining = abs(length(column_types) - length(split_line))
    expected_types = vcat(copy(column_types), [Int64 for i in 1:remaining])
    data = [parse(expec_type, value) for (expec_type, value) in zip(expected_types, split_line)]
    Level(e=data[4], g=data[5], encoding=[val for val in data[8:end]])
end

function replace_sym_labels(string_encoding)
    for (key, value) ∈ label_map
        string_encoding = replace(string_encoding, key=>value)
    end
    string_encoding
end

"""Function to clean up the linelist quantum number encoding. The
first action replaces symmetry labels with a numeric mapping, followed
by splitting each row into a vector of strings. The result is combined
to yield a `M x N` matrix for `M` quantum numbers, and `N` entries.

Finally, the matrix values are converted into integers.
"""
function clean_linelist_encoding(csv_data)
    mapreduce(x->replace_sym_labels(x) |> split, hcat, csv_data[:Label]) |> x -> parse.(Int, x)
end