text:
"""
    flow_val(b::TapTransformer)

Calculates the From - To comp[lex power flow (Flow injected at the bus) of branch of type
TapTransformer

"""
function flow_val(b::TapTransformer)
    Y_t = 1 / (PowerSystems.get_r(b) + PowerSystems.get_x(b) * 1im)
    c = 1 / PowerSystems.get_tap(b)
    arc = PowerSystems.get_arc(b)
    V_from = arc.from.voltage * (cos(arc.from.angle) + sin(arc.from.angle) * 1im)
    V_to = arc.to.voltage * (cos(arc.to.angle) + sin(arc.to.angle) * 1im)
    I = (V_from * Y_t * c^2) - (V_to * Y_t * c)
    flow = V_from * conj(I)
    return flow
end

"""
    flow_val(b::TapTransformer)

Calculates the From - To complex power flow (Flow injected at the bus) of branch of type
Line

"""
function flow_val(b::Line)
    Y_t = (1 / (PowerSystems.get_r(b) + PowerSystems.get_x(b) * 1im))
    arc = PowerSystems.get_arc(b)
    V_from = arc.from.voltage * (cos(arc.from.angle) + sin(arc.from.angle) * 1im)
    V_to = arc.to.voltage * (cos(arc.to.angle) + sin(arc.to.angle) * 1im)
    I = V_from * (Y_t + (1im * PowerSystems.get_b(b).from)) - V_to * Y_t
    flow = V_from * conj(I)
    return flow
end

"""
    flow_val(b::TapTransformer)

Calculates the From - To complex power flow (Flow injected at the bus) of branch of type
Transformer2W

"""
function flow_val(b::Transformer2W)
    Y_t = 1 / (PowerSystems.get_r(b) + PowerSystems.get_x(b) * 1im)
    arc = PowerSystems.get_arc(b)
    V_from = arc.from.voltage * (cos(arc.from.angle) + sin(arc.from.angle) * 1im)
    V_to = arc.to.voltage * (cos(arc.to.angle) + sin(arc.to.angle) * 1im)
    I = V_from * (Y_t + (1im * PowerSystems.get_primaryshunt(b))) - V_to * Y_t
    flow = V_from * conj(I)
    return flow
end

function flow_val(b::PhaseShiftingTransformer)
    error("Systems with PhaseShiftingTransformer not supported yet")
    return
end

function _update_branch_flow!(sys::System)
    for b in get_components(ACBranch, sys)
        S_flow = flow_val(b)
        b.activepower_flow = real(S_flow)
        b.reactivepower_flow = imag(S_flow)
    end
end

function _write_pf_sol!(sys::System, nl_result)
    result = nl_result.zero
    buses = enumerate(sort(collect(get_components(Bus, sys)), by = x -> get_number(x)))

    for (ix, bus) in buses
        if bus.bustype == BusTypes.REF
            P_gen = result[2 * ix - 1]
            Q_gen = result[2 * ix]
            injection_components = get_components(Generator, sys)
            devices = [d for d in injection_components if d.bus == bus]
            generator = devices[1]
            generator.activepower = P_gen
            generator.reactivepower = Q_gen
        elseif bus.bustype == BusTypes.PQ
            Q_gen = result[2 * ix - 1]
            θ = result[2 * ix]
            injection_components = get_components(Generator, sys)
            devices = [d for d in injection_components if d.bus == bus]
            if length(devices) == 1
                generator = devices[1]
                generator.reactivepower = Q_gen
            end
            bus.angle = θ
        elseif bus.bustype == BusTypes.PV
            Vm = result[2 * ix - 1]
            θ = result[2 * ix]
            bus.voltage = Vm
            bus.angle = θ
        end
    end

    _update_branch_flow!(sys)

    return
end

"""
    solve_powerflow!(sys, solve_function, args...)

Solves a the power flow into the system and writes the solution into the relevant structs.
Updates generators active and reactive power setpoints and branches active and reactive
power flows (calculated in the From - To direction) (see
[flow_val](@ref))

Requires loading NLsolve.jl to run. Internally it uses the make_pf (see
[make_pf](@ref)) to create the problem and solve it. As a result it doesn't enforce
reactivepower limits.

Supports passing NLsolve kwargs in the args. By default shows the solver trace.

Arguments available for `nlsolve`:

* `method` : See NLSolve.jl documentation for available solvers
* `xtol`: norm difference in `x` between two successive iterates under which
  convergence is declared. Default: `0.0`.
* `ftol`: infinite norm of residuals under which convergence is declared.
  Default: `1e-8`.
* `iterations`: maximum number of iterations. Default: `1_000`.
* `store_trace`: should a trace of the optimization algorithm's state be
  stored? Default: `false`.
* `show_trace`: should a trace of the optimization algorithm's state be shown
  on `STDOUT`? Default: `false`.
* `extended_trace`: should additifonal algorithm internals be added to the state
  trace? Default: `false`.

## Examples
```julia
using NLsolve
solve_powerflow!(sys, nlsolve)
# Passing NLsolve arguments
solve_powerflow!(sys, nlsolve, method = :Newton)

```

"""
function solve_powerflow!(sys, nlsolve; args...)
    pf!, x0 = PowerSystems.make_pf(sys)
    res = nlsolve(pf!, x0; args...)
    if res.f_converged
        PowerSystems._write_pf_sol!(sys, res)
        @info("PowerFlow solve converged, the results have been stored in the system")
        return res.f_converged
    end
    @error("The powerflow solver returned convergence = $(res.f_converged)")
    return res.f_converged
end


text:
using Revise
using AutomotiveDrivingModels
using AutoViz
using Random

roadway = gen_straight_roadway(1, 500.0)

num_veh = 2
scene = Scene(num_veh)

models = Dict{Int, DriverModel}()

k_spd = 1.0
politeness = 3.0

# 1: first vehicle, moving the fastest
mlane = MOBIL(.1, politeness = politeness)
mlon = IntelligentDriverModel(k_spd = k_spd, σ = 0.0)
models[1] = Tim2DDriver(.1, mlane = mlane, mlon = mlon)
road_idx = RoadIndex(proj(VecSE2(0.0, 0.0, 0.0), roadway))
base_speed = 10.
veh_state = VehicleState(Frenet(road_idx, roadway), roadway, base_speed)
veh_def = VehicleDef(AgentClass.CAR, 5., 2.)
push!(scene, Vehicle(veh_state, veh_def, 1))
# 2: second vehicle, in the middle, moving at intermediate speed
mlane = MOBIL(.1, politeness = politeness)
mlon = IntelligentDriverModel(k_spd = k_spd, σ = 0.0)
models[2] = Tim2DDriver(.1, mlane = mlane, mlon = mlon)
base_speed = 0.
road_pos = 8.
veh_state = VehicleState(Frenet(road_idx, roadway), roadway, base_speed)
veh_state = move_along(veh_state, roadway, road_pos)
veh_def = VehicleDef(AgentClass.CAR, 5., 2.)
push!(scene, Vehicle(veh_state, veh_def, 2))

rec = SceneRecord(500, 0.1, num_veh)
prime_time = .2
rng = MersenneTwister(1)

s = deepcopy(scene)
simulate!(Any, rec, s, roadway, models, 1)
render(s, roadway, cam=CarFollowCamera(1, 10.0))

ext: jl


text:
"""
    testing_models.jl

Written by Carl Whitfield, University of Manchester, 2021

Only to be included by definitions.jl
"""

#===================== Testing vs. VL params =====================#

#====> PCR sens (Smith et al. JCMB - fitted) <====#
const PCR_VL0 = 8.522/4.408
const PCR_VLdisp = 4.408
const VL_LOD_PCR = 1.0      #Ct cutoff value for +ve test -- assumed from UKHSA data

#====> LFD params (Porton Down) <====#
const LFD_VLmean = 10.836/2.680
const LFD_VLstd = 1.0/2.680
const LFD_sens_max = 0.75 #assumed

#====> LFD params (Care Home Data) <====#
const VL_lims_2021 = collect(2.0:1.0:7.0)
const LFD_relsens_2021 = [0.004,0.025,0.087,0.162,0.414,0.651,0.8]
const VL_lims_Oct = [4.0,6.0]
const LFD_relsens_Oct = [0.032, 0.255, 0.477]

const PCRTATdist = Gamma(6.75,20/3)

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

"""
    logistic_function(x::Float64, x0::Float64, k::Float64, ymax::Float64)

Scaled logistic function

## Arguments:
`x`

`x0` = Point where logistic function is at 50% of max value

`k` = Slope of logistic function

`ymax` = Maximum value of logistic function

## Returns:
`Float64` = ymax / (1  + exp(-k*(x-x0)))
"""
function logistic_function(x::Float64, x0::Float64, k::Float64, ymax::Float64)
    return (ymax / (1  + exp(-k*(x-x0))))
end

"""
    probit_function(x::Float64, x0::Float64, xstd::Float64, ymax::Float64)

Scaled probit function

## Arguments:
`x`

`x0` = Mean of normal distribution kernel

`xstd` = Standar deviation of normal distribution kernel

`ymax` = Maximum value of probit function

## Returns:
`Float64` = ymax * cdf(Normal(x0,xstd),x)
"""
function probit_function(x::Float64, x0::Float64, xstd::Float64, ymax::Float64)
    return ymax * cdf(Normal(x0,xstd),x)
end

"""
    PCRtest_positive_prob(VL::Float64)

PCR test positive probability given a viral load

## Arguments:
`VL` = viral load value (log10 copies/ml)

## Returns:
`Float64` = test positive probability
"""
function PCRtest_positive_prob(VL::Float64)
    if VL > VL_LOD_PCR
        return logistic_function(VL, PCR_VL0, PCR_VLdisp, PCR_sens_max)
    else
        return 0
    end
end

"""
    LFDtest_positive_prob(VL::Float64)

LFD test positive probability given a viral load

## Arguments:
`VL` = viral load value

## Returns:
`Float64` = test positive probability
"""
function LFDtest_positive_prob(VL::Float64)
    #test senstivity defined relative to concurrent PCR
    PCR_prob = PCRtest_positive_prob(VL)
    if LFD_model == SC_data_2021
        VLcat = sum(VL .>= VL_lims_2021)
        return LFD_relsens_2021[1+VLcat]*PCR_prob
    elseif LFD_model == SC_data_Oct
        VLcat = sum(VL .>= VL_lims_Oct)
        return LFD_relsens_Oct[1+VLcat]*PCR_prob
    elseif LFD_model == porton_down
        return probit_function(VL, LFD_VLmean, LFD_VLstd, LFD_sens_max)
    end
end

# """ old version
#     LFDtest_positive_prob1(VL::Float64, sens_rel::Float64=1.0)

# LFD test positive probability given a viral load

# ## Arguments:
# `VL` = viral load value
# `sens_rel` = relative sensitivity scaling factor (default 1)

# ## Returns:
# `Float64` = test positive probability
# """
# function LFDtest_positive_prob(VL::Float64, sens_rel::Float64=1.0)
#     #sens rel is 1 in Peto et al, could scale this for self-administering effects (sens rel < 1)
#     return sens_rel*probit_function(VL, LFD_VLmean, LFD_VLstd, LFD_sens_max)
# end

function get_pos_profile(sim::Dict, ip::Int, protocol::Int)
    sens_scale = 1.0
    Do_test_prob = 1.0
    if haskey(sim, "P_do_test")
        Do_test_prob = sim["P_do_test"]
    end
    if haskey(sim,"test_sens_scale")
        sens_scale = sim["test_sens_scale"]
    end
    if protocol == PCR_mass_protocol
        sim["test_pos_profiles"][ip] = sens_scale .* Do_test_prob .* PCRtest_positive_prob.(sim["VL_profiles"][ip])
    elseif protocol == LFD_mass_protocol
        sim["test_pos_profiles"][ip] = sens_scale .* Do_test_prob .* LFDtest_positive_prob.(sim["VL_profiles"][ip])
    end
end

"""
    init_testing!(sim::Dict, testing_params::Dict, i_day::Int, Ndays::Int)

Function to initialise test positive profiles and test isolation probabilities

## Arguments:
`sim` = Dict generated by `init_VL_and_infectiousness(Ntot::Int, Pisol::Float64)`

`testing_params` = Dict containing testing options, must have:

    "new_comply_prob" => `Float` = Probability that people who would not comply
                       with symptom onset isolation, will comply with testing isolation

    "tperiod" => `Int` = Days between repeat tests

    "protocol" => `String` = Current available options are given by the parameters
                            `PCR_mass_protocol` and `LFD_mass_protocol`

`i_day` = Day to start simulation on (counting from 1)

`Ndays` = Total number of days to simulate (counting from 1)

## Returns:
`Array{Int64,1}` = Days when testing will occur

`Int64` = Next testing day index (from `i_day`)

## See also:
`init_VL_and_infectiousness(Ntot::Int, Pisol::Float64)`
"""
function init_testing!(sim::Dict, testing_params::Dict, i_day::Int, Ndays::Int; fill_pos_profiles::Bool=true,
                       different_start::Bool=false)
    #add test positivity profiles to simulation Dict, i_day is start day, Ndays is total length of sim
    #testing params contains "tperiod" (days between test)
    #testing params contains "protocol" (LFD_mass_protocol or PCR_mass_protocol)
    #testing params contains "testing_enforced" (if true all tests are done)
    #test_miss_prob: if "testing_enforced" is false, random probability of missing a test
    #optional: sens_rel: relative scaling of test sensitivity (does not effect spec)
    #optional: policy_adherence: if "testing_enforced" is false, independent probability of adhering to testing
    #(has same effect as sens_rel on test sensitivity, but also affect false positives)
    sim["test_protocol"] = testing_params["protocol"]
    if haskey(testing_params,"sens_rel")
        sim["test_sens_scale"] = testing_params["sens_rel"]
    end
    if (testing_params["testing_enforced"] == false) && haskey(testing_params,"test_miss_prob")
        sim["P_do_test"] = 1.0 - testing_params["test_miss_prob"]
    end

    sim["will_isolate_with_test"] = ones(Bool,sim["Ntot"])
    if testing_params["testing_enforced"] == false  #if not enforced, we assume same group who does not isolate with symps does not isolate with tests (simplification)
        if haskey(testing_params,"policy_adherence")
            non_adherers = randsubseq(1:sim["Ntot"], 1 - testing_params["policy_adherence"])
            sim["will_isolate_with_test"][non_adherers] .= false
        else
            sim["will_isolate_with_test"][sim["non_isolators"]] .= false
        end
    end
    sim["testing_paused"] = zeros(Bool,sim["Ntot"])
    sim["resume_testing"] = -ones(Int64,sim["Ntot"])

    if different_start
        test_days0 = rand(1:testing_params["tperiod"],sim["Ntot"])
        test_days = Array{Array{Int64,1},1}(undef,sim["Ntot"])
        test_day_counter = ones(Int64,sim["Ntot"])
        for i in 1:sim["Ntot"]
            test_days[i] = collect(test_days0[i]:testing_params["tperiod"]:Int64(ceil(Ndays)))
            test_days[i] = push!(test_days[i], test_days[i][end] + testing_params["tperiod"])
            test_day_counter[i] = 1 + sum(test_days[i] .< i_day)
        end
    else
        test_day0 = rand(1:testing_params["tperiod"])
        test_days = collect(test_day0:testing_params["tperiod"]:Int64(ceil(Ndays)))
        test_days = push!(test_days, test_days[end] + testing_params["tperiod"])
        test_day_counter = 1 + sum(test_days .< i_day)
    end
    sim["test_pos_profiles"] = Array{Array{Float64,1},1}(undef,sim["Ntot"])

    if fill_pos_profiles
        sim["test_pos_profiles"] .= get_pos_profile.(Ref(sim), 1:sim["Ntot"],
                        testing_params["protocol"])
    end

    return test_days, test_day_counter
end

function draw_PCR_delays(N::Int)
    #from distribution in social care setting
    delays_hrs = PCR_TaT_scale * rand(PCRTATdist,N)
    delays_days = Int64.(round.(delays_hrs./24))
    #round delay to nearest day
    return delays_days
end

function get_test_probs(VL_profile::Array{Float64,1}, TestDays::Array{Int64,1},
                        TestTypes::Array{Int64,1}; Do_PCR_test_prob::Float64 = 1.0,
                        Do_LFD_test_prob::Float64 = 1.0, all_or_nothing::Bool = false)

    tp = zeros(length(TestDays))
    TTcondPCR = (TestTypes .== 0) .* (TestDays .<= length(VL_profile)) #PCR is test type 0
    TTcondLFD = (TestTypes .== 1) .* (TestDays .<= length(VL_profile))  #LFD is test type 1
    if all_or_nothing
        Do_all_PCR = (rand() < Do_PCR_test_prob)
        Do_all_LFD = (rand() < Do_LFD_test_prob)
        if Do_all_PCR
            tp[TTcondPCR] = PCRtest_positive_prob.(VL_profile[TestDays[TTcondPCR]])
        else
            tp[TTcondPCR] .= 0
        end
        if Do_all_LFD
            tp[TTcondLFD] = LFDtest_positive_prob.(VL_profile[TestDays[TTcondLFD]])
        else
            tp[TTcondLFD] .= 0
        end
    else
        tp[TTcondPCR] = Do_PCR_test_prob*PCRtest_positive_prob.(VL_profile[TestDays[TTcondPCR]])
        tp[TTcondLFD] = Do_LFD_test_prob*LFDtest_positive_prob.(VL_profile[TestDays[TTcondLFD]])
    end

    return tp
end

function run_testing_scenario!(isol_days::Array{Int64,1}, inf_profile::Array{Float64,1},
        test_pos::Array{Float64,1}, test_result_day::Array{Int64,1}, symp_day::Int64, symp_isol::Bool,
        VL_profile::Array{Float64,1}, Conf_PCR::Array{Bool,1},
        Preisolation::Array{Int64,1} = zeros(Int64,0); Day5release::Bool = false,
        Conf_PCR_fixed::Bool = false, Conf_PCR_sens = PCR_sens_max)

    isol_start_day = copy(test_result_day)
    isol_start_prob = copy(test_pos)
    isol_conf_PCR = copy(Conf_PCR)
    if symp_isol && symp_day > 0
        push!(isol_start_day, symp_day)
        push!(isol_start_prob, 1.0)
        push!(isol_conf_PCR, false)   #conf PCR not needed for symptoms
        itd = sortperm(isol_start_day)
        isol_conf_PCR = isol_conf_PCR[itd]
        isol_start_day = isol_start_day[itd]
        isol_start_prob = isol_start_prob[itd]
    end
    go = true
    it = 1
    empty!(isol_days)
    if length(Preisolation) > 0
        push!(isol_days,Preisolation...)
    end
    while go && it <= length(isol_start_day) && isol_start_day[it] <= length(inf_profile)
        if rand() < isol_start_prob[it]  #positive test or symp
            isol_end = isol_start_day[it] + 10
            if Day5release #test daily from day 4, released as soon as 2 consecutive are negative
                DailyLFDResult = zeros(Bool, 6)
                #assume negative once VL profile has ended
                if length(VL_profile) > isol_start_day[it] + 4
                    TestEnd = min(isol_start_day[it] + 9, length(VL_profile))
                    DailyLFDProb = LFDtest_positive_prob.(VL_profile[isol_start_day[it] + 4:TestEnd])
                    DailyLFDResult[1:length(DailyLFDProb)] = rand(length(DailyLFDProb)) .< DailyLFDProb
                end
                EndEarly = false
                k = 5
                while EndEarly == false && k < 10
                    if DailyLFDResult[k - 4] == false && DailyLFDResult[1 + k - 4] == false
                        isol_end = isol_start_day[it] + k
                        EndEarly = true
                    end
                    k = k+1
                end
            end
            if isol_conf_PCR[it]  #assume confirmatory PCR is done immediately
                if isol_start_day[it] <= length(VL_profile)
                    CPCRprob = 0
                    if Conf_PCR_fixed
                        CPCRprob = Conf_PCR_sens
                    else
                        CPCRprob = PCRtest_positive_prob(VL_profile[isol_start_day[it]])
                    end
                    if rand() > CPCRprob
                        #if conf PCR is negative
                        isol_end = isol_start_day[it] + draw_PCR_delays(1)[1]
                    else
                        go = false   #otherwise isolate in full
                    end
                else
                    go = false   #otherwise isolate in full
                end
            else
                go = false
            end
            if isol_end >= isol_start_day[it]
                push!(isol_days, collect(isol_start_day[it]:1:(isol_end-1))...)
            end
        end
        it = it + 1
    end
    Nisol_days = sum(isol_days)   #record actual number of isolation days

    return Nisol_days
end


text:
module MyTestPkg

include("sum.jl")

export my_sum

end # module


text:
cd(@__DIR__); include("setups/path128.jl");
gr(dpi = 200)

## (a)
plt = plot(size = (500, 400))
    heatmap!(Uf, c = :viridis, ratio = 1, yflip = true, xlim = [1, N], clims = (-1, 1))
    xticks!([1; 16:16:128], vcat(string(1), [string(k) for k in 16:16:128]))
    yticks!([1; 16:16:128], vcat(string(1), [string(k) for k in 16:16:128]))
savefig(plt, "../figs/Path128_Uf_j1.png")

## (b)
plt = plot(size = (500, 560), layout = Plots.grid(2, 1))
    scatter!(diag(Uf), ylim = [-1, 1], c = :black, grid = false, legend = false,
    frame = :box, lw = 1.5, ms = 3)
    bar!(diag(Uf); bar_width = 0.02, c = :black)
    xticks!([1; 16:16:128], vcat(string(1), [string(k) for k in 16:16:128]))
    title!("Diagonal entries")

    scatter!(anti_diag(Uf), ylim = [-1, 1], c = :black, grid = false, legend = false,
    frame = :box, lw = 1.5, ms = 3, subplot = 2)
    bar!(anti_diag(Uf); bar_width = 0.02, c = :black, subplot = 2)
    xticks!([1; 16:16:128], vcat(string(1), [string(k) for k in 16:16:128]))
    title!("Antidiagonal entries", subplot = 2)
savefig(plt, "../figs/Path128_Uf_diag_entries.png")


text:
using UnzipUser
using Test

using Unzip

@testset "UnzipUser.jl" begin
    @test unzip([(1,), (2,), (3,), (4,)]) == ([1, 2, 3, 4],)
end


text:
using Zygote
using Random
using DelimitedFiles
using LinearAlgebra
using GalacticOptim,Optim,NLopt
using DifferentialEquations

using Plots
using LaTeXStrings
#using GLOQ
#using ReverseDiff
include("../src/GLOQ.jl")
pyplot()
# Plotting
fnt = Plots.font("Helvetica",16)
lfnt = Plots.font("Helvetica",12)
Plots.default(titlefont=fnt,
			  guidefont=fnt,
			  tickfont=fnt,
			  legendfont=lfnt,
			  linewidth=2);
# BlackBoxOptim somehow downgrade some packages and as a result breaks the auto-differentiation with Zygote
# we should avoid it.
# NLopt seems to be really fast. Their LBFGS with box constrains is faster than Optim's FMinbox(LBFGS)

########################################################################
# Step 1: set up the problem
########################################################################
N_states = 4;
# True values to generate the synthtetic data
freqs_true = [4.0108; 3.8830; 3.6287];
omegas_true = 2.0*pi.*freqs_true;
gamma1_true   = [2.222222e-05; 4.761904e-05; 4.545455e-5];
gamma2_true   = [4.166667e-05; 6.8965517e-5; 2.3255814e-4];
char_noise_true = [2e-6,6e-5,1.3e-4];

########################################################################
# Step 1a: known parameters to set up the experiment
########################################################################
# Konwn driving frequencies
omr_t1_01 = 2.0*pi*4.0108
omr_t1_12 = 2.0*pi*3.8830
omr_t1_23 = 2.0*pi*3.6287

omr_ramsey_01 = 2.0*pi*(4.0108-5e-4)
omr_ramsey_12 = 2.0*pi*(3.8830-2.0e-3)
omr_ramsey_23 = 2.0*pi*(3.6287-2.0e-3)

omr_echo_01 = 2.0*pi*4.0108
omr_echo_12 = 2.0*pi*3.8830
omr_echo_23 = 2.0*pi*3.6287

# width of the control signal
width = 21.25
TC_01 = 2.5*width

width = 17.00
TC_12 = 2.5*width

width = 17.00
TC_23 = 2.5*width

# Convert the dark time from micro-seconds to nano-seconds
T_T1_01 = 125.0*GLOQ.GLOQ_MICRO_SEC
T_T1_12 = 125.0*GLOQ.GLOQ_MICRO_SEC
T_T1_23 = 125.0*GLOQ.GLOQ_MICRO_SEC
T_Ramsey_01 = 20.0*GLOQ.GLOQ_MICRO_SEC
T_Ramsey_12 = 20.0*GLOQ.GLOQ_MICRO_SEC
T_Ramsey_23 = 10.0*GLOQ.GLOQ_MICRO_SEC
T_Echo_01 = 40.0*GLOQ.GLOQ_MICRO_SEC
T_Echo_12 = 20.0*GLOQ.GLOQ_MICRO_SEC
T_Echo_23 = 10.0*GLOQ.GLOQ_MICRO_SEC

# data size used for characterization
n1use = 201
n2use = 201
n3use = 201
n4yse = 201
n4use = 201
n5use = 201
n6use = 201
n7use = 201
n8use = 201
n9use = 201

# time points evaluated
dt_t1_01 = T_T1_01/200
dt_t1_12 = T_T1_12/200
dt_t1_23 = T_T1_23/200
dt_ramsey_01 = T_Ramsey_01/200
dt_ramsey_12 = T_Ramsey_12/200
dt_ramsey_23= T_Ramsey_23/200
dt_echo_01 = T_Echo_01/200
dt_echo_12 = T_Echo_12/200
dt_echo_23 = T_Echo_23/200

t_t1_01 = zeros(Float64,n1use)
for i = 1:n1use
	t_t1_01[i] = (i-1)*dt_t1_01
end
T_evaluated_t1_01 = t_t1_01[end]

t_t1_12 = zeros(Float64,n2use)
for i = 1:n2use
	t_t1_12[i] = (i-1)*dt_t1_12
end
T_evaluated_t1_12 = t_t1_12[end]

t_t1_23 = zeros(Float64,n3use)
for i = 1:n3use
	t_t1_23[i] = (i-1)*dt_t1_23
end
T_evaluated_t1_23 = t_t1_23[end]

t_ramsey_01 = zeros(Float64,n4use)
for i = 1:n4use
	t_ramsey_01[i] = (i-1)*dt_ramsey_01
end
T_evaluated_ramsey_01 = t_ramsey_01[end]

t_ramsey_12 = zeros(Float64,n5use)
for i = 1:n5use
	t_ramsey_12[i] = (i-1)*dt_ramsey_12
end
T_evaluated_ramsey_12 = t_ramsey_12[end]

t_ramsey_23 = zeros(Float64,n6use)
for i = 1:n6use
	t_ramsey_23[i] = (i-1)*dt_ramsey_23
end
T_evaluated_ramsey_23 = t_ramsey_23[end]

t_echo_01 = zeros(Float64,n7use)
for i = 1:n7use
	t_echo_01[i] = (i-1)*dt_echo_01
end
T_evaluated_echo_01 = t_echo_01[end]

t_echo_12 = zeros(Float64,n8use)
for i = 1:n8use
	t_echo_12[i] = (i-1)*dt_echo_12
end
T_evaluated_echo_12 = t_echo_12[end]

t_echo_23 = zeros(Float64,n9use)
for i = 1:n9use
	t_echo_23[i] = (i-1)*dt_echo_23
end
T_evaluated_echo_23 = t_echo_23[end]

# Knwon initial conditions for the quantum system
u0_t1_01 = [1.0;0.0;0.0;0.0]
u0_t1_12 = [0.0;1.0;0.0;0.0]
u0_t1_23 = [0.0;0.0;1.0;0.0]
u0_ramsey_01 = [1.0;0.0;0.0;0.0]
u0_ramsey_12 = [0.0;1.0;0.0;0.0]
u0_ramsey_23 = [0.0;0.0;1.0;0.0]
u0_echo_01 = [1.0;0.0;0.0;0.0]
u0_echo_12 = [0.0;1.0;0.0;0.0]
u0_echo_23 = [0.0;0.0;1.0;0.0]

v0_t1_01 = [0.0;0.0;0.0;0.0]
v0_t1_12 = [0.0;0.0;0.0;0.0]
v0_t1_23 = [0.0;0.0;0.0;0.0]
v0_ramsey_01 = [0.0;0.0;0.0;0.0]
v0_ramsey_12 = [0.0;0.0;0.0;0.0]
v0_ramsey_23 = [0.0;0.0;0.0;0.0]
v0_echo_01 = [0.0;0.0;0.0;0.0]
v0_echo_12 = [0.0;0.0;0.0;0.0]
v0_echo_23 = [0.0;0.0;0.0;0.0]

########################################################################
# Step 1b: read the synthetic data
########################################################################
data_t1_01 = readdlm("synthetic_data/t1_01.txt",',')
data_t1_12 = readdlm("synthetic_data/t1_12.txt",',')
data_t1_23 = readdlm("synthetic_data/t1_23.txt",',')
data_ramsey_01 = readdlm("synthetic_data/ramsey_01.txt",',')
data_ramsey_12 = readdlm("synthetic_data/ramsey_12.txt",',')
data_ramsey_23 = readdlm("synthetic_data/ramsey_23.txt",',')
data_echo_01 = readdlm("synthetic_data/echo_01.txt",',')
data_echo_12 = readdlm("synthetic_data/echo_12.txt",',')
data_echo_23 = readdlm("synthetic_data/echo_23.txt",',')

########################################################################
# add multiplicative and additive noise to the synthetic data
########################################################################
multiplicative_noise = 1.0.+0.025*randn(n1use)
additive_noise = 0.02*randn(n1use)
data_t1_01 .*= multiplicative_noise
data_t1_01 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n2use)
additive_noise = 0.025*randn(n2use)
data_t1_12 .*= multiplicative_noise
data_t1_12 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n3use)
additive_noise = 0.025*randn(n3use)
data_t1_23 .*= multiplicative_noise
data_t1_23 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n4use)
additive_noise = 0.025*randn(n4use)
data_ramsey_01 .*= multiplicative_noise
data_ramsey_01 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n5use)
additive_noise = 0.025*randn(n5use)
data_ramsey_12 .*= multiplicative_noise
data_ramsey_12 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n6use)
additive_noise = 0.025*randn(n6use)
data_ramsey_23 .*= multiplicative_noise
data_ramsey_23 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n7use)
additive_noise = 0.025*randn(n7use)
data_echo_01 .*= multiplicative_noise
data_echo_01 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n8use)
additive_noise = 0.025*randn(n8use)
data_echo_12 .*= multiplicative_noise
data_echo_12 .+= additive_noise

multiplicative_noise = 1.0.+0.025*randn(n9use)
additive_noise = 0.025*randn(n9use)
data_echo_23 .*= multiplicative_noise
data_echo_23 .+= additive_noise
# population (probability) must be in the range of [0,1],
# and the summation of the population for different states must be 1
function correct_noisy_data!(data)
	_n,_m = size(data)
	for i = 1:_n
		for j = 1:_m
			if(data[i,j]<0.0)
				data[i,j] = 0.0
			end
		end
		data[i,:] ./= sum(data[i,:])
	end
end

correct_noisy_data!(data_t1_01)
correct_noisy_data!(data_t1_12)
correct_noisy_data!(data_t1_23)
correct_noisy_data!(data_echo_01)
correct_noisy_data!(data_echo_12)
correct_noisy_data!(data_echo_23)
correct_noisy_data!(data_ramsey_01)
correct_noisy_data!(data_ramsey_12)
correct_noisy_data!(data_ramsey_23)

global _function_call = 0
global _iteration_number = 0

########################################################
# Step 2: define and solve the optimization problem
########################################################
# Step 2a: define the loss (objective) function
function loss_gala(p,p_keywords)
	# Ramsey experiments
	_rho_Ramsey_01_u,_rho_Ramsey_01_v = GLOQ.RamseyParityForwardSolve(u0_ramsey_01,v0_ramsey_01,
					 2.0*pi.*p[1:3],omr_ramsey_01,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 0,
					 TC_01,t_ramsey_01,N_states)
	_population_Ramsey_01 = GLOQ.get_population(_rho_Ramsey_01_u)

	_rho_Ramsey_12_u,_rho_Ramsey_12_v = GLOQ.RamseyParityForwardSolve(u0_ramsey_12,v0_ramsey_12,
					 2.0*pi.*p[1:3],omr_ramsey_12,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 1,
					 TC_12,t_ramsey_12,N_states)
	_population_Ramsey_12 = GLOQ.get_population(_rho_Ramsey_12_u)

	_rho_Ramsey_23_u,_rho_Ramsey_23_v = GLOQ.RamseyParityForwardSolve(u0_ramsey_23,v0_ramsey_23,
					 2.0*pi.*p[1:3],omr_ramsey_23,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 2,
					 TC_23,t_ramsey_23,N_states)
	_population_Ramsey_23 = GLOQ.get_population(_rho_Ramsey_23_u)

	# Echo experiments
	_rho_Echo_01_u,_rho_Echo_01_v = GLOQ.EchoParityForwardSolve(u0_echo_01,v0_echo_01,
					 2.0*pi.*p[1:3],omr_echo_01,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 0,
					 TC_01,t_echo_01,N_states)
	_population_Echo_01 = GLOQ.get_population(_rho_Echo_01_u)

	_rho_Echo_12_u,_rho_Echo_12_v = GLOQ.EchoParityForwardSolve(u0_echo_12,v0_echo_12,
					 2.0*pi.*p[1:3],omr_echo_12,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 1,
					 TC_12,t_echo_12,N_states)
	_population_Echo_12 = GLOQ.get_population(_rho_Echo_12_u)

	_rho_Echo_23_u,_rho_Echo_23_v = GLOQ.EchoParityForwardSolve(u0_echo_23,v0_echo_23,
					 2.0*pi.*p[1:3],omr_echo_23,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 2,
					 TC_23,t_echo_23,N_states)
	_population_Echo_23 = GLOQ.get_population(_rho_Echo_23_u)

	# T1-decay experiments
	_rho_T1_01_u,_rho_T1_01_v = GLOQ.T1ParityForwardSolve(u0_t1_01,v0_t1_01,
					 2.0*pi.*p[1:3],omr_t1_01,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 0,
					 TC_01,t_t1_01,N_states)
	_population_T1_01 = GLOQ.get_population(_rho_T1_01_u)

	_rho_T1_12_u,_rho_T1_12_v = GLOQ.T1ParityForwardSolve(u0_t1_12,v0_t1_12,
					 2.0*pi.*p[1:3],omr_t1_12,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 1,
					 TC_12,t_t1_12,N_states)
	_population_T1_12 = GLOQ.get_population(_rho_T1_12_u)

	_rho_T1_23_u,_rho_T1_23_v = GLOQ.T1ParityForwardSolve(u0_t1_23,v0_t1_23,
					 2.0*pi.*p[1:3],omr_t1_23,
					 2.0*pi.*p[4:6],
					 p[7:9],p[10:12],
					 2,
					 TC_23,t_t1_23,N_states)
	_population_T1_23 = GLOQ.get_population(_rho_T1_23_u)

	#################################################################################
#	_loss = sum(abs2,_population_Echo_23-data_echo_23).*dt_echo_23/T_evaluated_echo_23

	_loss = sum(abs2,_population_Ramsey_01-data_ramsey_01)*dt_ramsey_01/T_evaluated_ramsey_01+
			sum(abs2,_population_Ramsey_12-data_ramsey_12)*dt_ramsey_12/T_evaluated_ramsey_12+
			sum(abs2,_population_Ramsey_23-data_ramsey_23)*dt_ramsey_23/T_evaluated_ramsey_23+
			sum(abs2,_population_Echo_01-data_echo_01).*dt_echo_01/T_evaluated_echo_01+
			sum(abs2,_population_Echo_12-data_echo_12).*dt_echo_12/T_evaluated_echo_12+
			sum(abs2,_population_Echo_23-data_echo_23).*dt_echo_23/T_evaluated_echo_23+
			sum(abs2,_population_T1_01-data_t1_01).*dt_t1_01/T_evaluated_t1_01+
			sum(abs2,_population_T1_12-data_t1_12).*dt_t1_12/T_evaluated_t1_12+
			sum(abs2,_population_T1_23-data_t1_23).*dt_t1_23/T_evaluated_t1_23

	global _function_call
	_function_call += 1
	if(_function_call%1==0)
		println("Function call: ",_function_call," Loss = ",_loss)
	end
	return _loss
end

########################################################################
# Step 2b: initial guess and bounds for target parameters
########################################################################
# transition frequency
freqs = [4.0108; 3.8830; 3.6287];
omegas = 2.0*pi.*freqs # change things from GHz to radians
# magnitude of the charge noise
char_noise = [1.5e-6, 5.0e-5, 1.0e-4]
# Lindblad parameters
gamma1   = [1.0e-05; 4.0e-05; 4.0e-5]
gamma2   = [2.5e-05; 5.0e-05; 2.5e-4]
# initial guess
p_initial = [freqs; # transition frequencies
			 char_noise;
			 gamma1; # parameters for the decoupling
			 gamma2 # parameters for the dephasing
 		 	 ]# charge noise

# upper and lower bound for optimizaiton
initial_loss = loss_gala(p_initial,[])

lower_bound = [ 4.00; 3.85; 3.60;
				1e-6; 5e-5; 1e-4;
				1e-5; 3e-5; 3e-5;
				2e-5; 2.5e-5; 2e-4;
			  ]
upper_bound = [ 4.10; 3.95; 3.65;
				3e-6; 7e-5; 2e-4;
				3e-5; 5e-5; 5e-5;
				5e-5; 1e-4; 3e-4;
			  ]

#####################################################################################
# Step 2c: optimization interface
#####################################################################################
# interface for the GalacticOptim with gradient computed by auto-differentiation
# (with the Zygote package)
optimization_function_az = OptimizationFunction(loss_gala, GalacticOptim.AutoZygote())
prob_az = GalacticOptim.OptimizationProblem(optimization_function_az, p_initial;
										    lb = lower_bound, ub = upper_bound)
_function_call = 0
# solve the optimization problem
@time sol = solve(prob_az , Opt(:LD_LBFGS,length(p_initial)),
				  maxiters=200,
				  ftol_rel=1e-3
				 )
p_optim = sol.u

########################################################
# Step 3: visualization
########################################################
# Plot results
include("ShowResultsOfCharacterizationSynthetic.jl")
# output
println("Transition frequencies: ", "\n Optimzied: ",p_optim[1:3], "\n True: ",freqs_true,"\n Error: ",p_optim[1:3]-freqs_true,"\n")
println("Charge noise:\n ", " Optimzied: ",p_optim[4:6], "\n True: ",char_noise_true,"\n Error: ",p_optim[4:6]-char_noise_true,"\n")
println("Gamma 1:\n ","\n Optimzed: ",p_optim[7:9],"\n True: ",gamma1_true,"\n Error: ",p_optim[7:9]-gamma1_true,"\n")
println("Gamma 2:\n ","\n Optimzed: ",p_optim[10:12],"\n True: ",gamma2_true,"\n Error: ",p_optim[10:12]-gamma1_true,"\n")