text:
# Library Imports

#Based on England data (CHESS and NHS England)
# I want a way to keep this as the "average" disease progression,
# but modify it such that old people have less favorable outcomes (as observed)
# But correspondingly I want people at lower risk to have more favorable outcome on average

# For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr
relativeDeathRisk_given_COVID_by_age = [-0.99742186, -0.99728639, -0.98158438,
                                        -0.9830432 , -0.82983414, -0.84039294,
                                         0.10768979,  0.38432409,  5.13754904]
nTest, nIso, nHS, nAge = 4, 4, 8, 9
stateTensor = ones((nTest, nIso, nHS, nAge))
#ageRelativeDiseaseSeverity = np.array([-0.8, -0.6, -0.3, -0.3, -0.1, 0.1, 0.35, 0.4, 0.5])
# FIXED (above) - this is a guess, find data and fix
#ageRelativeRecoverySpeed = np.array([0.2]*5+[-0.1, -0.2, -0.3, -0.5])
#TODO - this is a guess, find data and fix
ageRelativeRecoverySpeed = zeros(9)# For now we make it same for everyone, makes calculations easier
# For calculations see data_cleaning_py.ipynb, calculations from NHS England dataset as per 05 Apr
caseFatalityRatioHospital_given_COVID_by_age = [0.00856164, 0.03768844, 0.02321319,
            0.04282494, 0.07512237, 0.12550367, 0.167096  , 0.37953452, 0.45757006]

agePopulationTotal = 1000*[8044.056, 7642.473, 8558.707, 9295.024, 8604.251,
                                      9173.465, 7286.777, 5830.635, 3450.616]

function _agePopulationRatio(agePopulationTotal)
    agePopulationTotal *= 55.98/66.27
    return agePopulationTotal/sum(agePopulationTotal)
end

function einsum(str, a, b, T)
    if str=="ijl,j->i"
        return _einsum1(a, b, T)
    elseif str=="ijk,j->ik"
        return _einsum2(a, b, T)
    elseif  str=="ijkl,j->i"
        return _einsum3(a, b, T)
    elseif  str=="ijkl,ijklmnp->imnp"
        return _einsum5(a, b, T)
    end
end

function einsum(str, a, T)
   if str=="ijlml->ijlm"
      return _einsum4(a, T)
  elseif str=="ijkj->ijk"
      return _einsum6(a, T)
  elseif str=="iklkl->ikl"
      return _einsum7(a, T)
  elseif str=="ijljl->ijl"
      return _einsum7(a, T)
  elseif str=="ijkljkm->ijklm"
      return _einsum8(a, T)
  elseif str=="ijkljkl->ijkl"
      return _einsum9(a, T)
  elseif str=="...jkl->..."
      return _einsum10(a, T)
  end
end

function _einsum1(a, b, T) #'ijl,j->i'
    i,_,j = size(a)
    p = zeros(T,i,j)
    for i=1:i, j=1:j
		p[i,j] += dot(a[i,:,j], b)
    end
    return sum(p, dims=1)
end

function _einsum2(a, b, T) #'ijk,j->ik'
    i,_,j = size(a)
    p = zeros(T,i,j)
    for i=1:i, j=1:j
		p[i,j] = dot(a[i,:,j], b)
    end
    return p
end

function _einsum3(a, b, T) #'ijkl,j->i'
    _,j,_,i = size(a)
    p = zeros(T,j,i)
    for i=1:i, j=1:j
		p[j,i] += sum(a[:,j,:,i]*b)
    end
    return sum(p, dims=1)
end

function _einsum4(a, T) #'ijlml->ijlm'
    l,m,l,j,i = size(a)
    p = zeros(T,m,l,j,i)
    for i=1:i, j=1:j, l=1:l
        p[:,l,j,i] = a[l,:,l,j,i]
    end
    return p
end

function _einsum12(a, b, T) # for setting the values in einsum 4
	l,m,l,j,i = size(a)
	for i=1:i, j=1:j, l=1:l
        a[l,:,l,j,i] = b[:,l,j,i]
    end
    return a
end

function _einsum5(a, b, T) #'ijkl,ijklmnp->imnp'
    p, n, m, l, k, j, i = size(b)
    dydt = zeros(T,p,n,m,i)
    for i=1:i, j=1:j, k=1:k, l=1:l, m=1:m, n=1:n, p=1:p
        dydt[p,n,m,i] += a[l,k,j,i] * b[p,n,m,l,k,j,i]
    end
    return dydt
end

function _einsum6(a, T) # 'ijkj->ijk'
    _,k,j,i = size(a)
    p = zeros(T,k,j,i)
    for j=1:j, i=1:i
        p[:,j,i] = a[j,:,j,i]
    end
    return p
end

function _einsum7(a, T) # 'ijljl->ijl' & 'iklkl->ikl'
    _,_,l,k,i = size(a)
    p = zeros(T,l,k,i)
    for i=1:i, k=1:k, l=1:l
        p[l,k,i] = a[l,k,l,k,i]
    end
    return p
end

function _einsum8(a, T) # 'ijkljkm->ijklm'
    m,_,_,l,k,j,i = size(a)
    p = zeros(T,m,l,k,j,i)
    for i=1:i, j=1:j, k=1:k, l=1:l, m=1:m
        p[m,l,k,j,i] = a[m,k,j,l,k,j,i]
    end
    return p
end

function _einsum9(a, T) # 'ijkljkl->ijkl'
    _,_,_,l,k,j,i = size(a)
    p = zeros(T,l,k,j,i)
    for i=1:i, j=1:j, k=1:k, l=1:l
        p[l,k,j,i] = a[l,k,j,l,k,j,i]
    end
    return p
end

function _einsum11(a, b, T) # for setting the values in einsum 10
	_,_,_,l,k,j,i = size(a)
    for i=1:i, j=1:j, k=1:k, l=1:l
        a[l,k,j,l,k,j,i] = b[l,k,j,i]
    end
    return a
end

function _einsum10(a, T) # '...jkl->...'
    return reshape(sum(a, dims=[1,2,3]),size(a)[4:end])
end

agePopulationRatio = _agePopulationRatio(agePopulationTotal)

function trFunc_diseaseProgression(
         ageRelativeRecoverySpeed::Array =
         ageRelativeRecoverySpeed,
         caseFatalityRatioHospital_given_COVID_by_age::Array=
         caseFatalityRatioHospital_given_COVID_by_age,
         nonsymptomatic_ratio::Float64 = 0.86,
                                   # number of days between measurable events
         infect_to_symptoms::Float64 = 5.0,
                                   #symptom_to_death = 16.;
         symptom_to_recovery::Float64= 10.0, # 20.5; #unrealiticly long for old people
         symptom_to_hospitalisation::Float64 = 5.76,
         hospitalisation_to_recovery::Float64 = 14.51,
         IgG_formation::Float64 = 15.0,
                                   # Age related parameters
                                   # for now we'll assume that all hospitalised cases are known (overall 23% of hospitalised COVID patients die. 9% overall case fatality ratio)
                                   # Unknown rates to estimate
         nonsymp_to_recovery::Float64 = 15.0,
         inverse_IS1_IS2::Float64 = 4.0;
         kwargs...)
    # Now we have all the information to build the age-aware multistage SIR model transition matrix
    # The full transition tensor is a sparse map from the Age x HealthState x isolation state to HealthState,
    # and thus is a 4th order tensor itself, representing a linear mapping
    # from "number of people aged A in health state B and isolation state C to health state D.
    #agePopulationRatioByTotal = _agePopulationRatio(agePopulationTotal)
    nAge, nHS, nIso = kwargs[:nAge], kwargs[:nHS], kwargs[:nIso]
    #relativeDeathRisk_given_COVID_by_age = [:relativeDeathRisk_given_COVID_by_age]
    trTensor_diseaseProgression = zeros((nHS, nIso, nHS, nAge))
    # Use basic parameters to regularise inputs
    E_IS1 = 1.0/infect_to_symptoms
    # Numbers nonsymptomatic is assumed to be 86% -> E->IN / E-IS1 = 0.86/0.14
    E_IN = 0.86/0.14 * E_IS1
    # Nonsymptomatic recovery
    IN_R1 = 1.0/nonsymp_to_recovery
    IS1_IS2  = 1.0/inverse_IS1_IS2
    IS2_R1 = 1.0/(symptom_to_recovery - inverse_IS1_IS2)
    R1_R2 = 1.0/IgG_formation

    # Disease progression matrix # TODO - calibrate (together with transmissionInfectionStage)
    # rows: from-state, cols: to-state (non-symmetric!)
    # - this represent excess deaths only, doesn't contain baseline deaths!

    # Calculate all non-serious cases that do not end up in hospitals.
    # Note that we only have reliable death data from hospitals (NHS England),
    # so we do not model people dieing outside hospitals
    diseaseProgBaseline = [
    # to: E,    IN,    IS1,   IS2,    R1,      R2,     D
          0.0  E_IN  E_IS1    0       0        0       0    # from E
          0    0      0       0     IN_R1      0       0    # from IN
          0    0      0    IS1_IS2    0        0       0    # from IS1
          0    0      0       0     IS2_R1     0       0    # from IS2
          0    0      0       0       0       R1_R2    0    # from R1
          0    0      0       0       0        0       0    # from R2
          0    0      0       0       0        0       0    # from D
    ]

    diseaseProgBaseline = transpose(diseaseProgBaseline)
    # TODO can be improved
    # vcat(fill.(x, v)...) ???
    ageAdjusted_diseaseProgBaseline = deepcopy(cat(repeat([diseaseProgBaseline],
                                                              nAge)..., dims=3))
    # Modify all death and R1 rates:
    for ii in range(1, stop = size(ageAdjusted_diseaseProgBaseline)[2])
        # Adjust death rate by age dependent disease severity  ??? check the key args
        ageAdjusted_diseaseProgBaseline[end, ii, :] = adjustRatesByAge_KeepAverageRate(
                                            ageAdjusted_diseaseProgBaseline[end, ii, 1],
                             agePopulationRatio=_agePopulationRatio(agePopulationTotal),
                              ageRelativeAdjustment=relativeDeathRisk_given_COVID_by_age
                              )
        # Adjust recovery rate by age dependent recovery speed
        ageAdjusted_diseaseProgBaseline[end - 2, ii, :] = adjustRatesByAge_KeepAverageRate(
                                            ageAdjusted_diseaseProgBaseline[end - 2, ii, 1],
                                                      agePopulationRatio=agePopulationRatio,
                                             ageRelativeAdjustment=ageRelativeRecoverySpeed
                                             )
    end
    ageAdjusted_diseaseProgBaseline_Hospital = deepcopy(ageAdjusted_diseaseProgBaseline)
    # Calculate hospitalisation based rates, for which we do have data. Hospitalisation can end up with deaths
    # Make sure that the ratio of recoveries in hospital honour the case fatality ratio appropriately
    # IS2 -> death
    ageAdjusted_diseaseProgBaseline_Hospital[end, 4, :] =
                     ageAdjusted_diseaseProgBaseline_Hospital[end - 2, 4, :] .* ( # IS2 -> recovery
                                  caseFatalityRatioHospital_given_COVID_by_age./(  # multiply by cfr / (1-cfr) to get correct rate towards death
                            1 .-  caseFatalityRatioHospital_given_COVID_by_age) )

    #TODO - time to death might be incorrect overall without an extra delay state, especially for young people
    # Non-hospitalised disease progression
    for i1 in [1, 2, 4]
        trTensor_diseaseProgression[2:end, i1, 2:end, :] = ageAdjusted_diseaseProgBaseline
    end
    # hospitalised disease progression
    trTensor_diseaseProgression[2:end, 3, 2:end, :] = ageAdjusted_diseaseProgBaseline_Hospital
    return trTensor_diseaseProgression
end

# Population (data from Imperial #13 ages.csv/UK)
#agePopulationTotal = 1000*[8044.056, 7642.473, 8558.707, 9295.024, 8604.251,
#                                      9173.465, 7286.777, 5830.635, 3450.616]
#agePopulationTotal = 1000.*pd.read_csv("https://raw.githubusercontent.com/ImperialCollegeLondon/covid19model/master/data/ages.csv").iloc[3].values[2:]

# Currently: let's work with england population only instead of full UK, as NHS England + CHESS data is much clearer than other regions
#agePopulationTotal *= 55.98/66.27 # (google england/uk population 2018, assuming age dist is similar)
#agePopulationRatio = agePopulationTotal/sum(agePopulationTotal)

agePopulationRatio = _agePopulationRatio(agePopulationTotal)

function adjustRatesByAge_KeepAverageRate(rate; agePopulationRatio=agePopulationRatio,
                                                ageRelativeAdjustment::Array=nothing,
                                                maxOutRate::Float64=10.0)
    if rate == 0
        return fill(0, size(ageRelativeAdjustment))
    end
    if rate >= maxOutRate
        @warn("covidTesting::adjustRatesByAge_KeepAverageRate Input rate $rate >
                     maxOutRate $maxOutRate, returning input rates")
        return rate*(fill(1, size(ageRelativeAdjustment)))
    end
    out = fill(0, size(ageRelativeAdjustment))
    out[1] = maxOutRate + 1
    while sum(out .>= maxOutRate) > 0
        corrFactor = sum(agePopulationRatio ./ (1 .+ ageRelativeAdjustment))
        out =  rate * (1 .+ ageRelativeAdjustment) * corrFactor
        if sum(out .>= maxOutRate) > 0
            @warn("covidTesting::adjustRatesByAge_KeepAverageRate Adjusted rate
                   larger than $maxOutRate encountered, reducing ageAdjustment
                   variance by 10%")
            tmp_mean = sum(ageRelativeAdjustment)/length(ageRelativeAdjustment)
            ageRelativeAdjustment = tmp_mean .+ sqrt(0.9)*(
                                            ageRelativeAdjustment .- tmp_mean)
        end
    end
    return out
end

# Getting Hospitalised
# -----------------------------------
#ageHospitalisationRateBaseline

# Larger data driver approaches, with age distribution, see data_cleaning_R.ipynb for details

ageHospitalisationRateBaseline = DataFrame(load(joinpath(DATA_DIR,
 "clean_hosp-epis-stat-admi-summ-rep-2015-16-rep_table_6.csv")))[:, end]
ageHospitalisationRecoveryRateBaseline = DataFrame(load(joinpath(DATA_DIR,
 "clean_10641_LoS_age_provider_suppressed.csv")))[:, end]
ageHospitalisationRecoveryRateBaseline = 1.0 ./ ageHospitalisationRecoveryRateBaseline

# Calculate initial hospitalisation (occupancy), that will be used to initialise the model
initBaselineHospitalOccupancyEquilibriumAgeRatio = ageHospitalisationRateBaseline ./
                                                    (ageHospitalisationRateBaseline +
                                                    ageHospitalisationRecoveryRateBaseline) #? check the calculations

# Take into account the NHS work-force in hospitals that for our purposes count
# as "hospitalised S" population, also unaffected by quarantine measures
ageNhsClinicalStaffPopulationRatio = DataFrame(load(joinpath(DATA_DIR,
      "clean_nhsclinicalstaff.csv")))[:,end]

# Extra rate of hospitalisation due to COVID-19 infection stages
# TODO - find / estimate data on this (unfortunately true rates are hard to get due to many unknown cases)
# Symptom to hospitalisation is 5.76 days on average (Imperial #8)

infToHospitalExtra = Array([1e-4, 1e-3, 2e-2, 1e-2])

# For calculations see data_cleaning_py.ipynb, calculations from CHESS dataset as per 05 Apr
relativeAdmissionRisk_given_COVID_by_age = [-0.94886625, -0.96332087, -0.86528671,
                                           -0.79828999, -0.61535305, -0.35214767,
                                            0.12567034,  0.85809052,  3.55950368]

riskOfAEAttandance_by_age = [0.41261361, 0.31560648, 0.3843979 ,
                             0.30475704, 0.26659415,0.25203475,
                             0.24970244, 0.31549102, 0.65181376]

function trFunc_HospitalAdmission(
         ageHospitalisationRateBaseline::Array=
         ageHospitalisationRateBaseline,
         infToHospitalExtra::Array=infToHospitalExtra,
         ageRelativeExtraAdmissionRiskToCovid::Array=
         relativeAdmissionRisk_given_COVID_by_age .*
         riskOfAEAttandance_by_age;
         kwargs...
         )
    nAge, nHS, nI = kwargs[:nAge], kwargs[:nHS], kwargs[:nI]

    trTensor_HospitalAdmission = zeros((nHS, nAge))

    ageAdjusted_infToHospitalExtra = deepcopy(cat(repeat([infToHospitalExtra],
                                                             nAge)..., dims=2))
    for ii in range(1, stop = size(ageAdjusted_infToHospitalExtra)[1])
        ageAdjusted_infToHospitalExtra[ii, :] = adjustRatesByAge_KeepAverageRate(
                     infToHospitalExtra[ii],
                     ageRelativeAdjustment=ageRelativeExtraAdmissionRiskToCovid
                    )
    end
    # Add baseline hospitalisation to all non-dead states
    trTensor_HospitalAdmission[1:end-1, :] .+= reshape(ageHospitalisationRateBaseline,
                                            (1, size(ageHospitalisationRateBaseline)...))
    # Add COVID-caused hospitalisation to all infeted states
    #(TODO: This is a summation fo rates for independent processes, should be correct, but check)
    trTensor_HospitalAdmission[2:(nI+1), :] .+= ageAdjusted_infToHospitalExtra
    return trTensor_HospitalAdmission
end

function trFunc_HospitalDischarge(
    ageHospitalisationRecoveryRateBaseline::Array=
    ageHospitalisationRecoveryRateBaseline,
    dischargeDueToCovidRateMultiplier::Float64=3.0;
    kwargs...
    )
    nAge, nHS = kwargs[:nAge], kwargs[:nHS]
    trTensor_HospitalDischarge = zeros((nHS, nAge))
    # Baseline discharges apply to all non-symptomatic patients (TODO: take into account testing state!)
    trTensor_HospitalDischarge[1:3, :] .+= transpose(
                                     ageHospitalisationRecoveryRateBaseline)

    # No discharges for COVID symptomatic people from the hospital until they recover
    # TODO - check with health experts if this is correct assumption; probably also depends on testing state
    trTensor_HospitalDischarge[4:5, :] .= 0.0
    trTensor_HospitalDischarge[6:7, :] .= dischargeDueToCovidRateMultiplier .*
                            transpose(ageHospitalisationRecoveryRateBaseline)
    return trTensor_HospitalDischarge
end

# Overall new infections include within quarantine and hospital infections
# ------------------------------------------------------------------------

ageSocialMixingBaseline = DataFrame(load(joinpath(DATA_DIR,
"socialcontactdata_UK_Mossong2008_social_contact_matrix.csv")))[:,2:end]
ageSocialMixingBaseline = convert(Matrix, ageSocialMixingBaseline)
ageSocialMixingBaseline = (ageSocialMixingBaseline.+
                           transpose(ageSocialMixingBaseline))/2.0
ageSocialMixingBaseline = ageSocialMixingBaseline
ageSocialMixingDistancing = DataFrame(load(joinpath(DATA_DIR,
"socialcontactdata_UK_Mossong2008_social_contact_matrix_with_distancing.csv")))[:,2:end]
ageSocialMixingDistancing = convert(Matrix, ageSocialMixingDistancing)
ageSocialMixingDistancing = (ageSocialMixingDistancing.+
                             transpose(ageSocialMixingDistancing))/2.0
ageSocialMixingDistancing = ageSocialMixingDistancing

ageSocialMixingIsolation = fill(0.0, size(ageSocialMixingBaseline))

elevatedMixingRatioInHospital = 3.0

withinHospitalSocialMixing = elevatedMixingRatioInHospital *
                           sum(ageSocialMixingBaseline * agePopulationRatio)

transmissionInfectionStage = [0.001, 0.1, 0.6, 0.5]
"""
All new infections, given infected people on all different isolation states (normal, home, hospital)
We use the following assumptions:

- Infectiousness only depends on infection stage, not age or location

- Hospitalised people are assumed to only mix with other hospitalised people (this includes staff!),
in a non-age-dependent manner: withinHospitalSocialMixing

If policySocialDistancing is True
- Non- and home-isolated people mix with non- and home isolated via ageSocialMixingDistancing (averaging interactions)

If policySocialDistancing is False, we assume home-isolation is taken more seriously, but with little effect on non-isolated people
- Non-isolated people mix with each other via ageSocialMixingBaseline, and with home-isolated people via ageSocialMixingIsolation
- Home-isolated people do not mix with each other

This separation will help disentangle the effects of simply a blanket lessening of social distancing
(keeping the policy True but with less effective ageSocialMixingDistancing matrix),
vs case isolation (policy = False, but with serious ageSocialMixingIsolation)
"""
function trFunc_newInfections_Complete(
         stateTensor,
         policySocialDistancing::Bool, # True / False, no default because it's important to know which one we use at any moment!
         policyImmunityPassports::Bool, # True / False, no default because it's important to know which one we use at any moment!
         ageSocialMixingBaseline::Array=
         ageSocialMixingBaseline,
         ageSocialMixingDistancing::Array=
         ageSocialMixingDistancing,
         ageSocialMixingIsolation::Array=
         ageSocialMixingIsolation,
         withinHospitalSocialMixing::Float64=
         withinHospitalSocialMixing,
         transmissionInfectionStage::Array=
         transmissionInfectionStage;
         kwargs...
         )

    nTest, nIso, nHS, nAge, nI = kwargs[:nTest], kwargs[:nIso], kwargs[:nHS],
                                 kwargs[:nAge] , kwargs[:nI]
    ageIsoContractionRate = zeros((nTest, nIso, nAge))
    # Add non-hospital infections
    #--------------------------------
    curNonIsolatedSocialMixing = policySocialDistancing ? ageSocialMixingDistancing : ageSocialMixingBaseline
    # Add baseline interactions only between non-isolated people
    for k1 in [1, 4]
        for k2 in [1, 4]
            ageIsoContractionRate[:,k1,:] .+= reshape(
                  ( einsum("ijl,j->i",
                    stateTensor[:,k2,2:(nI+1),:], transmissionInfectionStage, eltype(transmissionInfectionStage)) * # all infected in non-isolation
                    curNonIsolatedSocialMixing
                  ),
                (1, size(curNonIsolatedSocialMixing)[1]...)
            )
        end
    end
    if policyImmunityPassports
        # If the immunity passports policy is on, everyone who tested antibody positive, can roam freely
        # Therefore replace the interactions between people with testingState = 2 with ageSocialMixingBaseline
        # we do this by using the distributive property of matrix multiplication, and adding extra interactions
        # "ageSocialMixingBaseline"-"curNonIsolatedSocialMixing" with each other (this is zero if no social distancing!)
        # TODO - this is a bit hacky?, but probably correct - double check though!
        for k1 in [1, 4]
            for k2 in [1, 4]
                ageIsoContractionRate[3:end,k1,:] .+=
                         einsum("ijk,j->ik",
                         stateTensor[3:end,k2,2:(nI+1),:], transmissionInfectionStage, eltype(transmissionInfectionStage))* # all infected in non-isolation
                         (ageSocialMixingBaseline.-curNonIsolatedSocialMixing)
            end
        end
    end
    # Add isolation interactions only between isolated and non-isolated people
    # non-isolated contracting it from isolated
    for k1 in [1, 4]
        ageIsoContractionRate[:,k1,:] .+= reshape(
              (einsum("ijl,j->i",
               stateTensor[:,2,2:(nI+1),:], transmissionInfectionStage, eltype(transmissionInfectionStage))* # all infected in isolation
               ageSocialMixingIsolation
              ),
            (1, size(ageSocialMixingIsolation )[1]...)
        )
    end
    # isolated contracting it from non-isolated
    for k1 in [1, 4]
        ageIsoContractionRate[:,1,:] .+= reshape(
               (einsum("ijl,j->i",
                stateTensor[:,k1,2:(nI+1),:], transmissionInfectionStage, eltype(transmissionInfectionStage))* # all infected in non-hospital, non-isolation
                ageSocialMixingIsolation
               ),
            (1, size(ageSocialMixingIsolation )[1]...)
        )
    end
        # isolated cannot contracting it from another isolated
    # Add in-hospital infections (of hospitalised patients, and staff)
    #--------------------------------
    # (TODO - within hospitals we probably want to take into effect the testing state;
    #      tested people are better isolated and there's less mixing)
    ageIsoContractionRate[:,3:end,:] .+= reshape(
                  withinHospitalSocialMixing *
                  einsum("ijkl,j->i",
                 stateTensor[:,3:end,2:(nI+1),:], transmissionInfectionStage, eltype(transmissionInfectionStage)), # all infected in hospital (sick or working)
        (1, 1, size(stateTensor)[end]...)
        )
    return ageIsoContractionRate/sum(stateTensor) # Normalise the rate by total population
end

function trFunc_travelInfectionRate_ageAdjusted(
	     t::Int64, # Time within simulation
	     travelMaxTime::Int64 = 200,
	     travelBaseRate::Float64 = 5e-4, # How many people normally travel back to the country per day # TODO - get data
	     travelDecline_mean::Float64 = 15.0,
	     travelDecline_slope::Float64 = 1.0,
	     travelInfection_peak::Float64 = 1e-1,
	     travelInfection_maxloc::Float64 = 10.0,
	     travelInfection_shape::Float64 = 2.0;
	     kwargs...
         )
	tmpTime = [0:1:travelMaxTime-1;]
	# nAge x T TODO get some realistic data on this
	travelAgeRateByTime = travelBaseRate .* (agePopulationRatio * transpose(1 .- map(logistic,
                        (tmpTime .- travelDecline_mean) ./ travelDecline_slope)))

    # 1 x T TODO get some realistic data on this, maybe make it age weighted
    _scale = travelInfection_maxloc / (travelInfection_shape-1)
    travelContractionRateByTime = map(x -> gammapdf(travelInfection_shape, 1.0,
                                            x/_scale), tmpTime)
    travelContractionRateByTime ./= _scale
    travelContractionRateByTime ./= max(travelContractionRateByTime...)
    travelContractionRateByTime .*= travelInfection_peak

    if t >= size(travelAgeRateByTime)[end]
        return zeros(size(travelAgeRateByTime)[1])
    else
        return travelAgeRateByTime[:,t+1] * travelContractionRateByTime[t+1]
    end
end

# Test parameters
# ---------------

# assumptions about practical (not theoretical, see discrapancy in PCR!) parameters of tests
# TODO - but particular data and references from lit (or estimates based on previous similar tests)

# TODO - MANUAL! - this function is VERY specific to current health state setup, and needs to be manually edited if number of health states change
function inpFunc_testSpecifications(
	PCR_FNR_I1_to_R2::Array = [ 0.9,  0.4, 0.15, 0.35, 0.5, 0.8],
    PCR_FPR::Float64 = 0.01,
    antigen_FNR_I1_to_R2::Array = [ 0.95, 0.6, 0.35, 0.45, 0.6, 0.9],
    antigen_FPR::Float64 = 0.1,
    antibody_FNR_I1_to_R2::Array = [0.99, 0.85, 0.8, 0.65, 0.3, 0.05],
    antibody_FPR_S_to_I4::Array = [0.05, 0.04, 0.03, 0.02, 0.01];
    kwargs...
    )

    nHS, nI, nR = kwargs[:nHS], kwargs[:nI], kwargs[:nR]

    testSpecifications = DataFrame()
    testSpecifications.Name = vcat(
                    ["PCR" for i in 1:nHS],
                    ["Antigen" for i in 1:nHS],
                    ["Antibody" for i in 1:nHS]
                    )
    testSpecifications.OutputTestState = vcat(
                    [1 for i in 1:2*nHS],
                    [2 for i in 1:nHS]
                    )
    testSpecifications.TruePosHealthState = vcat(
                    [[i for i in 1:nI] for j in 1:2*nHS],
                    [[i for i in nI+1:nI+nR] for j in 1:nHS]
                    )
    # In some health states some people are true negatives and some are true positives! (No, makes litte sense to use, just account for it in FPR? Only matters for test makers...)
    # testSpecifications['AmbiguousPosHealthState'] = [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(nI+1, nI+nR+1)]*nHS + [np.arange(1, nI+1)]*nHS # what information state does a pos test transition you to.

    testSpecifications.InputHealthState = vcat(0:nHS-1, 0:nHS-1, 0:nHS-1)

    # These numbers below are "defaults" illustrating the concept, but are modified by the inputs!!!

    testSpecifications.FalseNegativeRate = [ # ratio of positive (infected / immune) people missed by the test
        # For each health stage:
        #  S -> I1 (asymp) -> I2 (mild symp) -> I3 (symp, sick) -> I4 (symp, less sick) -> R1 / R2 (IgM, IgG avail) -> D

        # PCR
            0.,   0.9,            0.4,           0.15,                0.35,              0.5, 0.8,   0.,

        # Antigen
            0.,   0.95,           0.6,           0.35,                0.45,              0.6, 0.9,   0.,

        # Antibody
            0.,   0.99,           0.85,          0.8,                 0.65,              0.3, 0.05,  0.
    ]

    testSpecifications[2:7, :FalseNegativeRate] .= PCR_FNR_I1_to_R2
    testSpecifications[10:15, :FalseNegativeRate] .= antigen_FNR_I1_to_R2
    testSpecifications[18:23, :FalseNegativeRate] .= antibody_FNR_I1_to_R2

    testSpecifications.FalsePositiveRate = [ # ratio of negative (non-infected or not immune) people deemed positive by the test
        # PCR
        0.01, 0.,0.,0.,0., 0.01, 0.01, 0.,

        # Antigen
        0.1, 0.,0.,0.,0., 0.1, 0.1, 0.,

        # Antibody
        0.05, 0.04, 0.03, 0.02, 0.01, 0., 0., 0.
    ]

    testSpecifications[1, :FalsePositiveRate] = PCR_FPR
    testSpecifications[6:7, :FalsePositiveRate] = PCR_FPR
    testSpecifications[9, :FalsePositiveRate] = antigen_FPR
    testSpecifications[14:15, :FalsePositiveRate] = antigen_FPR
    testSpecifications[17:21, :FalsePositiveRate] .= antibody_FPR_S_to_I4

    return testSpecifications
end

function trFunc_testCapacity(
    realTime::Date, # time within simulation (day)
    # PCR capacity - initial
    testCapacity_pcr_phe_total::Float64 = 1e4,
    testCapacity_pcr_phe_inflexday::Date = Date("2020-03-25", "yyyy-mm-dd"),
    testCapacity_pcr_phe_inflexslope::Float64 = 5.0,

    # PCR capacity - increased
    testCapacity_pcr_country_total::Float64 = 1e5,
    testCapacity_pcr_country_inflexday::Date = Date("2020-04-25", "yyyy-mm-dd"),
    testCapacity_pcr_country_inflexslope::Float64 = 10.0,

    # Antibody / antigen capacity
    testCapacity_antibody_country_firstday::Date = Date("2020-04-25", "yyyy-mm-dd"),

    testCapacity_antibody_country_total::Float64 = 5e6,
    testCapacity_antibody_country_inflexday::Date = Date("2020-05-20", "yyyy-mm-dd"),
    testCapacity_antibody_country_inflexslope::Float64 = 20.0,

	testCapacity_antigenratio_country::Float64 = 0.7;

	kwargs...
	)

    # Returns a dictionary with test names and number available at day "t"
    outPCR = (
        #phe phase
        testCapacity_pcr_phe_total * logistic(Dates.value(Day(realTime-testCapacity_pcr_phe_inflexday))/testCapacity_pcr_phe_inflexslope)
        +
        #whole country phase
        testCapacity_pcr_country_total * logistic(Dates.value(Day(realTime-testCapacity_pcr_country_inflexday))/testCapacity_pcr_country_inflexslope)
    )

    if realTime<testCapacity_antibody_country_firstday
        outAntiTotal = 0.0
    else
        outAntiTotal = (
            testCapacity_antibody_country_total * logistic(Dates.value(Day(realTime-testCapacity_antibody_country_inflexday))/testCapacity_antibody_country_inflexslope)
		)
	end

    return Dict([
        ("PCR", outPCR),
        ("Antigen", outAntiTotal*testCapacity_antigenratio_country),
        ("Antibody", outAntiTotal*(1-testCapacity_antigenratio_country))
	]) # Tuples can be used instead (using dictionary to make it identical to python code
end

max_stars_repo_name: vollmersj/coexist-julia

max_forks_repo_name: vollmersj/coexist-julia


text:
export graph_exp_native_jl, graph_exp_native_jl_degopt

function get_expm_pade_coeffs(d, T)
    if d == 13
        return T[
            64764752532480000.0,
            32382376266240000.0,
            7771770303897600.0,
            1187353796428800.0,
            129060195264000.0,
            10559470521600.0,
            670442572800.0,
            33522128640.0,
            1323241920.0,
            40840800.0,
            960960.0,
            16380.0,
            182.0,
            1.0,
        ]
    elseif d == 9
        C = T[
            17643225600.0,
            8821612800.0,
            2075673600.0,
            302702400.0,
            30270240.0,
            2162160.0,
            110880.0,
            3960.0,
            90.0,
            1.0,
        ]
    elseif d == 7
        C = T[
            17297280.0,
            8648640.0,
            1995840.0,
            277200.0,
            25200.0,
            1512.0,
            56.0,
            1.0,
        ]
    elseif d == 5
        C = T[30240.0, 15120.0, 3360.0, 420.0, 30.0, 1.0]
    elseif d == 3
        C = T[120.0, 60.0, 12.0, 1.0]
    else
        error("Unknown degree for expm-Padé coefficients.")
    end
end

"""
     (graph,crefs)=graph_exp_native_jl(A; input=:A)

Creates a graph for the native scaling-and-squaring for the matrix exponential,
as implemented in Julia. The matrix `A` is taken as input to determine the
length of the Padé approximant and the number of squares applied, as well as
determining the type of the coefficients. The kwarg `input` determines the name
of the matrix, in the graph.

**References**

1. N. J. Higham. "Functions of Matrices". SIAM, Philadelphia, PA, 2008.
   DOI: [10.1137/1.9780898717778](https://doi.org/10.1137/1.9780898717778)

2. N. J. Higham. "The Scaling and Squaring Method for the Matrix Exponential
   Revisited". SIAM Journal on Matrix Analysis and Applications,
   26(4):1179-1193, 2005.
   DOI: [10.1137/04061101X](https://doi.org/10.1137/04061101X)

3. "Julia's matrix exponential", [at the time of
   conversion](https://github.com/JuliaLang/julia/blob/697e782ab86bfcdd7fd15550241fe162c51d9f98/stdlib/LinearAlgebra/src/dense.jl#L554).
"""
function graph_exp_native_jl(A; input = :A)
    T = eltype(A)
    nA = opnorm(A, 1)
    d = get_expm_pade_degree_native_lj(nA)
    if d < 13
        C = get_expm_pade_coeffs(d, T)
        return graph_exp_native_jl_low(C, input)
    else
        C = get_expm_pade_coeffs(d, T)
        s = log2(nA / 5.4) # power of 2 later reversed by squaring
        return graph_exp_native_jl_high(C, s, input)
    end
end

function get_expm_pade_degree_native_lj(nA)
    if nA <= 2.1
        if nA > 0.95
            return 9
        elseif nA > 0.25
            return 7
        elseif nA > 0.015
            return 5
        else
            return 3
        end
    else
        return 13
    end
end

# For sufficiently small nA, use lower order Padé-Approximations
function graph_exp_native_jl_low(C, input)
    T = eltype(C)
    graph = Compgraph(T)

    n = length(C)
    s = (div(n, 2) - 1)
    cref = Vector{Tuple{Symbol,Int}}(undef, n)

    add_mult!(graph, :A2, input, input)
    evenmon = Vector{Symbol}(undef, s + 1)
    evenmon[1] = :I
    evenmon[2] = :A2
    for k = 2:s
        sym = Symbol("A$(2*k)")
        add_mult!(graph, sym, :A2, evenmon[k])
        evenmon[k+1] = sym
    end

    a = view(C, 1:2:n-1)
    cref[1:s+1] = add_sum!(graph, :V, a, evenmon, :V)

    a = view(C, 2:2:n)
    cref[s+2:n] = add_sum!(graph, :Ua, a, evenmon, :Ua)
    add_mult!(graph, :U, :Ua, input)

    add_lincomb!(graph, :X, 1.0, :V, 1.0, :U)
    add_lincomb!(graph, :Z, 1.0, :V, -1.0, :U)
    add_ldiv!(graph, :P, :Z, :X)

    add_output!(graph, :P)

    return (graph, cref)
end

# Full scaling and squaring
function graph_exp_native_jl_high(CC, s, input)
    T = eltype(CC)
    graph = Compgraph(T)

    C = input
    if s > 0
        si = ceil(Int, s)
        γ = 1 / convert(T, 2^si)
        C = :C
        add_lincomb!(graph, C, γ, input, 0, :I)
    end

    cref = Vector{Tuple{Symbol,Int}}(undef, 14)

    add_mult!(graph, :A2, C, C)
    add_mult!(graph, :A4, :A2, :A2)
    add_mult!(graph, :A6, :A2, :A4)

    #U  = A * (A6 * (CC[14].*A6 .+ CC[12].*A4 .+ CC[10].*A2) .+
    #          CC[8].*A6 .+ CC[6].*A4 .+ CC[4].*A2 .+ CC[2].*Inn)

    # Ub3= CC[14].*A6 .+ CC[12].*A4 .+ CC[10].*A2
    a = view(CC, 10:2:14)
    cref[1:3] = add_sum!(graph, :Ub3, a, [:A2, :A4, :A6], :Ub)
    add_mult!(graph, :Ub, :Ub3, :A6)

    # Ua = CC[8].*A6 .+ CC[6].*A4 .+ CC[4].*A2 .+ CC[2].*Inn
    a = view(CC, 2:2:8)
    cref[4:7] = add_sum!(graph, :Ua, a, [:I, :A2, :A4, :A6], :Ua)

    add_lincomb!(graph, :Uc, 1.0, :Ub, 1.0, :Ua)
    add_mult!(graph, :U, C, :Uc)

    # V  = A6 * (CC[13].*A6 .+ CC[11].*A4 .+ CC[9].*A2) .+
    #            CC[7].*A6 .+ CC[5].*A4 .+ CC[3].*A2 .+ CC[1].*Inn

    # Vb3= CC[13].*A6 .+ CC[11].*A4 .+ CC[9].*A2
    a = view(CC, 9:2:13)
    cref[8:10] = add_sum!(graph, :Vb3, a, [:A2, :A4, :A6], :Vb)
    add_mult!(graph, :Vb, :Vb3, :A6)

    # Va = CC[7].*A6 .+ CC[5].*A4 .+ CC[3].*A2 .+ CC[1].*Inn
    a = view(CC, 1:2:7)
    cref[11:14] = add_sum!(graph, :Va, a, [:I, :A2, :A4, :A6], :Va)

    add_lincomb!(graph, :V, 1.0, :Vb, 1.0, :Va)

    add_lincomb!(graph, :X, 1.0, :V, 1.0, :U)
    add_lincomb!(graph, :Z, 1.0, :V, -1.0, :U)
    add_ldiv!(graph, :P, :Z, :X)

    Qtm1 = :P
    Qt = :P
    if (s > 0)
        for t = 1:si
            Qt = Symbol("S" * string(t))
            add_mult!(graph, Qt, Qtm1, Qtm1)

            Qtm1 = Qt
        end
    end

    add_output!(graph, Qt)

    return (graph, cref)
end

"""
     (graph,crefs)=graph_exp_native_jl_degopt(A; input=:A)

Same as [`graph_exp_native_jl`](@ref) but with calls to [`graph_degopt`](@ref)
for contruction of numerator and denominator polynomials.
"""
function graph_exp_native_jl_degopt(A; input = :A)
    T = eltype(A)
    nA = opnorm(A, 1)
    d = get_expm_pade_degree_native_lj(nA)
    if d < 13
        C = get_expm_pade_coeffs(d, T)
        return graph_exp_native_jl_low_degopt(C, input)
    else
        C = get_expm_pade_coeffs(d, T)
        s = log2(nA / 5.4) # power of 2 later reversed by squaring
        return graph_exp_native_jl_high_degopt(C, s, input)
    end
end

# For sufficiently small nA, use lower order Padé-Approximations
function graph_exp_native_jl_low_degopt(C, input)
    T = eltype(C)
    graph = Compgraph(T)

    n = length(C)
    s = (div(n, 2) - 1)

    xU = Vector{Tuple{Vector{T},Vector{T}}}(undef, s + 1)
    # Even monomials
    xU[1] = (vcat(zeros(T, 1), one(T)), vcat(zeros(T, 1), one(T)))
    for k = 2:s
        xU[k] =
            (vcat(zeros(T, 2), one(T), zeros(k - 2)), vcat(zeros(T, k), one(T)))
    end

    # U
    a = view(C, 4:2:n)
    xU[s+1] =
        (vcat(zeros(T, 1), one(T), zeros(T, s + 1)), vcat(C[2], zero(T), a))
    zU = vcat(zeros(T, s + 2), one(T))
    (graphU, crefU) = graph_degopt(xU, zU, input = input)
    rename_node!(graphU, Symbol("T2k$(4+s)"), :U, crefU)

    # V
    xV = Vector{Tuple{Vector{T},Vector{T}}}(undef, s)
    xV[1:s] = xU[1:s]
    a = view(C, 3:2:n-1)
    zV = vcat(C[1], zero(T), a)
    (graphV, crefV) = graph_degopt(xV, zV, input = input)
    rename_node!(graphV, Symbol("T2k$(3+s)"), :V, crefV)

    graph = merge_graphs(
        graphU,
        graphV,
        prefix1 = "U",
        prefix2 = "V",
        skip_basic1 = true,
        skip_basic2 = true,
        cref1 = crefU,
        cref2 = crefV,
        input1 = input,
        input2 = input,
    )
    cref = vcat(crefU, crefV)
    empty!(graph.outputs)

    merge_rename_exp_native_jl_degopt!(graph, cref, s + 1)

    add_lincomb!(graph, :X, 1.0, :VV, 1.0, :UU)
    add_lincomb!(graph, :Z, 1.0, :VV, -1.0, :UU)
    add_ldiv!(graph, :P, :Z, :X)

    add_output!(graph, :P)

    return (graph, cref)
end

# Full scaling and squaring
function graph_exp_native_jl_high_degopt(CC, s, input)
    T = eltype(CC)
    graph = Compgraph(T)
    C = input
    if s > 0
        si = ceil(Int, s)
        C = :C
        # Actual scaling below, after graph is created
    end

    #U  = A * (A6 * (CC[14].*A6 .+ CC[12].*A4 .+ CC[10].*A2) .+
    #          CC[8].*A6 .+ CC[6].*A4 .+ CC[4].*A2 .+ CC[2].*Inn)
    xU = Vector{Tuple{Vector{T},Vector{T}}}(undef, 5)

    # Even monomials
    xU[1] = (vcat(zeros(T, 1), one(T)), vcat(zeros(T, 1), one(T)))
    xU[2] = (vcat(zeros(T, 2), one(T)), vcat(zeros(T, 2), one(T)))
    xU[3] = (vcat(zeros(T, 2), one(T), zero(T)), vcat(zeros(T, 3), one(T)))

    # Up = A6 * (CC[14].*A6 .+ CC[12].*A4 .+ CC[10].*A2)
    a = view(CC, 10:2:14)
    xU[4] = (vcat(zeros(T, 4), one(T)), vcat(zeros(T, 2), a))

    # U = A* (CC[8].*A6 .+ CC[6].*A4 .+ CC[4].*A2 .+ CC[2].*Inn + Up)
    a = view(CC, 4:2:8)
    xU[5] =
        (vcat(zero(T), one(T), zeros(T, 4)), vcat(CC[2], zero(T), a, one(T)))

    zU = vcat(zeros(T, 6), one(T))
    (graphU, crefU) = graph_degopt(xU, zU, input = C)
    rename_node!(graphU, :T2k8, :U, crefU)

    # V  = A6 * (CC[13].*A6 .+ CC[11].*A4 .+ CC[9].*A2) .+
    #            CC[7].*A6 .+ CC[5].*A4 .+ CC[3].*A2 .+ CC[1].*Inn
    xV = Vector{Tuple{Vector{T},Vector{T}}}(undef, 4)

    # Even monomials
    xV[1:3] = xU[1:3]

    # Vp = A6 * ( CC[13].*A6 .+ CC[11].*A4 .+ CC[9].*A2 )
    a = view(CC, 9:2:13)
    xV[4] = (vcat(zeros(T, 4), one(T)), vcat(zeros(T, 2), a))

    # V = Vp + CC[7].*A6 .+ CC[5].*A4 .+ CC[3].*A2 .+ CC[1].*Inn
    a = view(CC, 3:2:7)
    zV = vcat(CC[1], zero(T), a, one(T))
    (graphV, crefV) = graph_degopt(xV, zV, input = C)
    rename_node!(graphV, :T2k7, :V, crefV)

    graph = merge_graphs(
        graphU,
        graphV,
        prefix1 = "U",
        prefix2 = "V",
        skip_basic1 = true,
        skip_basic2 = true,
        cref1 = crefU,
        cref2 = crefV,
        input1 = C,
        input2 = C,
    )
    cref = vcat(crefU, crefV)
    empty!(graph.outputs)

    merge_rename_exp_native_jl_degopt!(graph, cref, 4)

    add_lincomb!(graph, :X, 1.0, :VV, 1.0, :UU)
    add_lincomb!(graph, :Z, 1.0, :VV, -1.0, :UU)
    add_ldiv!(graph, :P, :Z, :X)

    if (s > 0)
        γ = 1 / convert(T, 2^si)
        add_lincomb!(graph, C, γ, input, 0, :I)
        xS = Vector{Tuple{Vector{T},Vector{T}}}(undef, si)
        for i = 1:si
            xS[i] = (vcat(zeros(T, i), one(T)), vcat(zeros(T, i), one(T)))
        end
        zS = vcat(zeros(T, si + 1), one(T))
        (graphS, crefS) = graph_degopt(xS, zS, input = :P)

        graph = merge_graphs(
            graph,
            graphS,
            prefix1 = "",
            prefix2 = "S",
            skip_basic1 = true,
            skip_basic2 = true,
            cref1 = cref,
            cref2 = crefS,
            input1 = C,
            input2 = :P,
        )
        cref = vcat(cref, crefS)
    else
        add_output!(graph, :P)
    end

    return (graph, cref)
end

# Remove redundant creation of monomial basis, e.g., VB3 == UB3 and VBa4_3 == UBa4_3
function merge_rename_exp_native_jl_degopt!(graph, cref, M)
    for N in ["VB", "UB"]
        for k in reverse(2:M)
            rename_node!(graph, Symbol(string(N, "$k")), Symbol("B$(k)"), cref)
            for NN in ["a", "b"]
                for kk in reverse(2:k-1)
                    rename_node!(
                        graph,
                        Symbol(string(N, NN, "$(k)_$(kk)")),
                        Symbol(string("B", NN, "$(k)_$(kk)")),
                        cref,
                    )
                end
            end
        end
    end
end

ext: jl


text:
push!(LOAD_PATH,"../src/")
using Documenter, JToric, DocumenterMarkdown

# ensure JToric is loaded for the doc tests
DocMeta.setdocmeta!(JToric, :DocTestSetup, :(using JToric, Random; using JToric.Oscar); recursive = true)

makedocs(sitename="JToric -- Toric Geometry in Julia")

deploydocs(repo = "github.com/oscar-system/JToric.jl.git")


text:
#================
Model
    Validation
        Module
================#
#-------Model Metrics--------____________
"""
      Mean absolute error (MAE) subtracts two arrays and averages the
      difference.\n
      --------------------\n
      actual = [5,10,15,20]\n
      pred = [5,10,15,20]\n
      mae(actual,pred)\n
       """
function mae(actual,pred)
    l = length(actual)
    lp = length(pred)
    if l != lp
        throw(ArgumentError("The array shape does not match!"))
    end
    result = actual-pred
    maeunf = mean(result)
    if maeunf < 0
        maeunf = (maeunf - maeunf) - maeunf
    end
    return(maeunf)
end
# <---- Mean Squared Error ---->
"""
      Mean Square error (MSE) subtracts two arrays, squares the
      difference, and averages the result\n
      --------------------\n
      actual = [5,10,15]\n
      pred = [5,10,15]
      r = mse(actual, pred)\n
       """
function mse(y,ŷ)
    diff = y .- ŷ
    diff = diff .^ 2
    Σdiff = sum(diff)
    return(Σdiff)
end
# <---- R Squared ---->
"""
      R squared is the correlation coefficient of regression, and is found
      by squaring the correlation coefficient.\n
      --------------------\n
      actual = [5,10,15]\n
      pred = [5,10,15]\n
      r2 = r2(actual, pred)\n
       """
function r2(actual,pred)
    l = length(actual)
    lp = length(pred)
    if l != lp
        throw(ArgumentError("The array shape does not match!"))
    end
    r = correlationcoeff(actual,pred)
    rsq = r^2
    rsq = rsq * 100
    return(rsq)
end
"""
Categorical accuracy is used to determine the accuracy of a categorical model
by calculating the percentage of correct predictions.\n
------------------\n
testy = ["hi","hello"]\n
yhat = ["hi","hi"]\n
accuracy = catacc(yhat,testy)\n
.5
    """
function catacc(testy,yhat)
    n = length(yhat)
    if n != length(testy)
        throw(ArgumentError("The array shape does not match!"))
    end
    score = 0
    for (i,w) in zip(yhat,testy)
        if i == w
            score += 1
        end
    end
    acc = score / n
end
macro accuracy(testy, yhat)
    originallen = length(testy)
    sets = Set(testy)
    if originallen * .10 > sets
        println("MAE ", mae(testy, yhat))
        println("MSE ", mse(testy, yhat))
        println("r2 ", r2(testy, yhats), "%")
    else
        println("Accuracy ", catacc(testy, yhat), "%")
    end
end
#--------------------------------------------
# End


text:
# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export ProbabilityWeight
"""
`ProbabilityWeight(x)`

| Set            |``\\oplus``   |  ``\\otimes``  | ``\\bar{0}`` | ``\\bar{1}`` |
|:--------------:|:------------:|:--------------:|:------------:|:------------:|
|``\\mathbb{R}`` |       +      |        *       |   ``0``      |   ``1``      |
"""
struct ProbabilityWeight{T <: AbstractFloat} <: Semiring
  x::T
end

zero(::Type{ProbabilityWeight{T}}) where T = ProbabilityWeight{T}(zero(T))
one(::Type{ProbabilityWeight{T}}) where T = ProbabilityWeight{T}(one(T))

*(a::ProbabilityWeight{T}, b::ProbabilityWeight{T}) where {T <: AbstractFloat} = ProbabilityWeight{T}(a.x * b.x)
+(a::ProbabilityWeight{T}, b::ProbabilityWeight{T}) where {T <: AbstractFloat} = ProbabilityWeight{T}(a.x + b.x)
/(a::ProbabilityWeight{T}, b::ProbabilityWeight{T}) where {T <: AbstractFloat} = ProbabilityWeight{T}(a.x / b.x)

# properties
iscommulative(::Type{W}) where {W<:ProbabilityWeight} = true
isleft(::Type{W}) where {W<:ProbabilityWeight} = true
isright(::Type{W}) where {W<:ProbabilityWeight} = true
isweaklydivisible(::Type{W}) where {W <: ProbabilityWeight}= true
iscomplete(::Type{W}) where {W <: ProbabilityWeight} = true


text:
"""
    Parameters.@with_kw struct Morse{T} <: AdiabaticModel

# References

[J. Chem. Phys. 88, 4535 (1988)](https://doi.org/10.1063/1.453761)
"""
Parameters.@with_kw struct Morse{T} <: AdiabaticModel
    Dₑ::T = 1.0
    x₀::T = 1.0
    a::T  = 1.0
    m::T = 1.0
end

NQCModels.ndofs(harmonic::Morse) = 1

function NQCModels.potential(model::Morse, r::Real)
    (;Dₑ, x₀, a) = model
    return Dₑ * (exp(-a*(r-x₀)) - 1)^2
end

function NQCModels.derivative(model::Morse, r::Real)
    (;Dₑ, x₀, a) = model
    return 2Dₑ * (exp(-a*(r-x₀)) - 1) * -a * exp(-a*(r-x₀))
end

"Eq. 44"
function getω₀(model::Morse)
    (;Dₑ, a, m) = model
    return sqrt(2Dₑ * a^2 / m)
end

"Eq. 36"
function getλ(model::Morse)
    (;Dₑ, a, m) = model
    return sqrt(2m*Dₑ) / a
end

"Eq. 43"
function eigenenergy(model::Morse, n)
    λ = getλ(model)
    n > trunc(Int, λ - 1/2) && throw(DomainError(n, "State $n is unbound!")) # Eq. 40
    ω = getω₀(model)
    harmonic = (n + 1/2)
    Eₙ = harmonic - 1/2λ * harmonic^2
    return Eₙ * ω
end


text:
using Images, GeometryTypes, GLVisualize, Reactive, GLWindow, Colors
using GLAbstraction, GLFW
import GLAbstraction: imagespace, singlepressed
import GLVisualize: moving_average

if !isdefined(:runtests)
    window = glscreen()
end

color_a = Signal(RGBA{Float32}(1,0,0,1))
iconsize = 54
color_v, color_s = widget(color_a, window, area=(iconsize, iconsize))
edit_screen = Screen(window, area = map(window.area) do a
    SimpleRectangle(0, 0, a.w, iconsize)
end)
paint_screen = Screen(window, area = map(window.area) do a
    SimpleRectangle(0, iconsize, a.w, a.h-iconsize)
end)
GLVisualize.add_screen(paint_screen)

using Plots; glvisualize(size = widths(paint_screen))

plt = plot(rand(100));
gui() # tell Plots.jl to plot this in a window

# Plots creates subscreens in which it plots.
# We need to recover that screen
paint_screen = plt.o.children[1]

# create a clear canvas button (unchecked seems vaguely like a good symbol for an empty canvas)
buttobj, button_s = GLVisualize.button(loadasset("unchecked.png"), edit_screen)

# create a slider for the linewidth
slider_w, slider_s = GLVisualize.slider(
    linspace(0.5f0, 20f0, 100), edit_screen,
    thickness = 2f0,
    slider_length = 4 * iconsize,
    icon_size = Signal(iconsize)
)

# view and position them with layout!
_view(
    layout!(SimpleRectangle{Float32}(0, 0, iconsize, iconsize), color_v),
    edit_screen, camera=:fixed_pixel
)
_view(
    layout!(SimpleRectangle{Float32}(iconsize+2, 0, iconsize, iconsize), buttobj),
    edit_screen, camera=:fixed_pixel
)
# translate slider
set_arg!(slider_w, :model, translationmatrix(Vec3f0(iconsize*2+4, 0, 0)))

_view(
    slider_w,
    edit_screen, camera=:fixed_pixel
)

const linebuffer = Signal(fill(Point2f0(NaN), 4))
const colorbuffer = Signal(fill(value(color_s), 4))
const lineobj = visualize(
    linebuffer, :lines, color = colorbuffer,
    thickness = slider_s
)
_view(lineobj, paint_screen, camera=:perspective)

@materialize mouseposition, mouse_buttons_pressed, mouseinside = paint_screen.inputs

camera = paint_screen.cameras[:perspective]

const history = Point2f0[] # preallocate history for moving average

s = map(mouseposition, mouse_buttons_pressed, init=nothing) do mp, mbp
    l0, c0 = map(value, (linebuffer, colorbuffer))
    if singlepressed(mbp, GLFW.MOUSE_BUTTON_LEFT) && value(mouseinside)
        p = imagespace(mp, camera)
        keep, p = moving_average(p, 1.5f0, history)
        if keep
            push!(linebuffer, push!(l0, p))
            push!(colorbuffer, push!(c0, value(color_s)))
        end
    else
        if !isnan(last(l0)) # only push one NaN to seperate
            empty!(history) # reset
            push!(linebuffer, push!(l0, Point2f0(NaN)))
            push!(colorbuffer, push!(c0, value(color_s)))
        end
    end
    nothing
end
# preserve signals, so that it doesn't get garbage collected.
preserve(s)

# reset buffers
s2 = map(button_s) do clicked
    if clicked
        push!(linebuffer, fill(Point2f0(NaN), 4))
        push!(colorbuffer, fill(value(color_s), 4))
    end
    nothing
end
preserve(s2)

if !isdefined(:runtests)
    renderloop(window)
end


text:
using Raylib

function main()
    screenWidth = 800
    screenHeight = 450

    Raylib.InitWindow(screenWidth, screenHeight, "raylib [textures] example - image loading")

    image = Raylib.LoadImage(joinpath(@__DIR__, "resources/raylib_logo.png"))
    texture = Raylib.LoadTextureFromImage(image)
    Raylib.UnloadImage(image)

    while !Raylib.WindowShouldClose()
        Raylib.BeginDrawing()
        Raylib.ClearBackground(Raylib.RAYWHITE)
        Raylib.DrawTexture(texture, fld(screenWidth - texture.width, 2), fld(screenHeight - texture.height, 2), Raylib.WHITE)
        Raylib.DrawText("this IS a texture loaded from an image!", 300, 370, 10, Raylib.GRAY)
        Raylib.EndDrawing()
    end

    Raylib.UnloadTexture(texture)
    Raylib.CloseWindow()
end


text:
function ImGui_ImplOpenGL3_CreateImageTexture(image_width, image_height; format=GL_RGBA, type=GL_UNSIGNED_BYTE)
    id = GLuint(0)
    @c glGenTextures(1, &id)
    glBindTexture(GL_TEXTURE_2D, id)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glPixelStorei(GL_UNPACK_ROW_LENGTH, 0)
    glTexImage2D(GL_TEXTURE_2D, 0, format, GLsizei(image_width), GLsizei(image_height), 0, format, type, C_NULL)
    g_ImageTexture[id] = id
    return Int(id)
end

function ImGui_ImplOpenGL3_UpdateImageTexture(id, image_data, image_width, image_height; format=GL_RGBA, type=GL_UNSIGNED_BYTE)
    glBindTexture(GL_TEXTURE_2D, g_ImageTexture[id])
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, GLsizei(image_width), GLsizei(image_height), format, type, image_data)
end

function ImGui_ImplOpenGL3_DestroyImageTexture(id)
    id = g_ImageTexture[id]
    @c glDeleteTextures(1, &id)
    delete!(g_ImageTexture, id)
    return true
end

function ImGui_ImplOpenGL3_DestroyImages(ctx::Context)
    for (k,v) in ctx.ImageTexture
        ImGui_ImplOpenGL3_DestroyImageTexture(v)
    end
    return true
end


text:
export Role

"""
A [`User`](@ref) role.
More details [here](https://discordapp.com/developers/docs/topics/permissions#role-object).
"""
struct Role
    id::Snowflake
    name::String
    color::Int
    hoist::Bool
    position::Int
    permissions::Int
    managed::Bool
    mentionable::Bool
end
@boilerplate Role :dict :docs :lower :merge


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using LinearAlgebra
using LoopVectorization
using LinearMaps
using ArnoldiMethod

function prepareDiag(L)
    diag = zeros(2^L)
    for state = 1:2^L
        for i = 1:L
            j = i==L ? 1 : i+1
            @inbounds diag[state] -= (((state >> (i-1))&1) == ((state >> (j-1))&1)) ? 1 : -1
        end
    end
    diag
end

function Hfunc!(C, B, diag, L)
    N = length(diag)
    @tturbo for state = 1:N
        C[state] = diag[state] * B[state]
    end
    for i = 1:L
        @tturbo for state = 1:N
            newstate = (state&(~(2^L))) ⊻ (1<<(i-1))
            c = newstate == 0
            newstate = !c*newstate + c*N # remove if statement
            C[newstate] -= B[state]
        end
    end
end
prepareHfunc!(L) = (C, B) -> Hfunc!(C, B, prepareDiag(L), L)

function f_LinearMap(L)
    H_LinearMap = LinearMap(prepareHfunc!(L), 2^L, ismutating=true, issymmetric=true, isposdef=false)
    d_LM, h_LM = partialschur(H_LinearMap, nev=1, which=SR())
    d_LM.eigenvalues[1]
end

f_exact(L) = -2sum(abs(sin((n-1/2) * pi/L)) for n in 1:L)

# %%
# https://juliaphysics.github.io/PhysicsTutorials.jl/tutorials/general/quantum_ising/quantum_ising.html

using LinearAlgebra
using SparseArrays
using ArnoldiMethod
⊗(x,y) = kron(x,y)

function TransverseFieldIsing_sparse(;N,h)
    id = [1 0; 0 1] |> sparse
    σˣ = [0 1; 1 0] |> sparse
    σᶻ = [1 0; 0 -1] |> sparse

    first_term_ops = fill(id, N)
    first_term_ops[1] = σᶻ
    first_term_ops[2] = σᶻ

    second_term_ops = fill(id, N)
    second_term_ops[1] = σˣ

    H = spzeros(Int, 2^N, 2^N) # note the spzeros instead of zeros here
    #for i in 1:N-1
    for i in 1:N # periodic boundary condition
        H -= foldl(⊗, first_term_ops)
        first_term_ops = circshift(first_term_ops,1)
    end

    for i in 1:N
        H -= h*foldl(⊗, second_term_ops)
        second_term_ops = circshift(second_term_ops,1)
    end
    H
end

function f_sparse(L)
    H = TransverseFieldIsing_sparse(N=L, h=1)
    d, h = partialschur(H; nev=1, which=SR())
    d.eigenvalues[1]
end

# %%
f_exact(20)

# %%
@time f_LinearMap(20)
@time f_LinearMap(20)

# %%
@time f_sparse(20)
@time f_sparse(20)

# %%
@time H = TransverseFieldIsing_sparse(N=20, h=1)
@time d, h = partialschur(H; nev=1, which=SR())

# %%
@time f_LinearMap(10)
@time f_LinearMap(10)

# %%
@time f_sparse(10)
@time f_sparse(10)

# %%
L = 10
@time H_LM = LinearMap(prepareHfunc!(L), 2^L, ismutating=true, issymmetric=true, isposdef=false)
@time d_LM, h_LM = partialschur(H_LM, nev=1, which=SR())
@time d_LM, h_LM = partialschur(H_LM, nev=1, which=SR())

# %%
@time H_sparse = TransverseFieldIsing_sparse(N=10, h=1)
@time d, h = partialschur(H_sparse; nev=1, which=SR())
@time d, h = partialschur(H_sparse; nev=1, which=SR())

# %%
using BenchmarkTools

L = 10
@show L
println("LinearMaps:")
H_LM = @btime LinearMap(prepareHfunc!(L), 2^L, ismutating=true, issymmetric=true, isposdef=false)
@btime d_LM, h_LM = partialschur($H_LM, nev=1, which=$(SR()))
println("SparseArrays:")
H_sparse = @btime TransverseFieldIsing_sparse(N=L, h=1)
@btime d, h = partialschur($H_sparse; nev=1, which=$(SR()));

# %%
using BenchmarkTools

L = 20
@show L
println("LinearMaps:")
H_LM = @btime LinearMap(prepareHfunc!(L), 2^L, ismutating=true, issymmetric=true, isposdef=false)
@btime d_LM, h_LM = partialschur($H_LM, nev=1, which=$(SR()))
println("SparseArrays:")
H_sparse = @btime TransverseFieldIsing_sparse(N=L, h=1)
@btime d, h = partialschur($H_sparse; nev=1, which=$(SR()));

# %%

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
module ESS

using Statistics

export ess

function ess(X::Matrix{Float64}, cols::Vector{Int})
    p = length(cols)
    esses = zeros(p)
    for i = 1:p
        esses[i] = ess(X, cols[i])
    end
    return esses
end

function ess(X::Matrix{Float64})
    p = size(X, 2)
    return ess(X, collect(1:p))
end

function ess(X::Matrix{Float64}, col::Int)
    n = size(X, 1)
    ρ = 0.0
    for i = 1:(n - 3)
        p = lag_correlation(X, col, i)
        # @show i
        # @show p
        if p > 0.0
            ρ += p
        else
            break
        end
        # ρ += p
    end
    return n / (1 + 2 * ρ)
end

# function lag_correlation(X::Matrix{Float64}, col::Int, lag::Int)
#     n = size(X, 1)
#     x1 = X[1:(n - lag), col]
#     x2 = X[(lag + 1):n, col]
#     @assert length(x1) == length(x2)
#     ρ = cor(x1, x2)
#     return ρ
# end

function lag_correlation(X::Matrix{Float64}, col::Int, lag::Int)
    si = 0.0
    sj = 0.0
    ssi = 0.0
    ssj = 0.0
    sij = 0.0
    n = size(X, 1) - lag
    for i = 1:n
        j = i + lag
        x = X[i, col]
        y = X[j, col]

        si += x
        ssi += x * x
        sj += y
        ssj += y * y
        sij += x * y
    end

    ei = si / n
    ej = sj / n
    ei2 = ssi / n
    ej2 = ssj / n
    eij = sij / n

    sdi = sqrt(ei2 - ei * ei)
    sdj = sqrt(ej2 - ej * ej)
    covij = eij - ei * ej
    return covij / (sdi * sdj)
end

end


text:
export LoggerActor, logger

"""
    LoggerActor{D}(name::String = "LogActor", io::O) where { D, O }

The `LoggerActor` logs all `next!`/`error!`/`complete!` events that are sent from an Observable.

# Constructor arguments
- `name`: name of the logger. Optional. Default is `LogActor`.
- `io`: io stream to log in, maybe nothing to write to `stdout`

See also: [`Actor`](@ref), [`logger`](@ref)
"""
struct LoggerActor{D, O} <: Actor{D}
    name :: String
    io   :: O

    LoggerActor{D, O}(name::String, io::O) where { D, O } = new(name, io)
end

# Remark: Here nothing for `io` is a workaround for https://github.com/JuliaDocs/Documenter.jl/issues/1245 where println(actor.io, ...) fails on doctest even if actor.io === stdout

on_next!(actor::LoggerActor{D}, data::D) where D = println(actor.io !== nothing ? actor.io : stdout, "[$(actor.name)] Data: $data")
on_error!(actor::LoggerActor, err)               = println(actor.io !== nothing ? actor.io : stdout, "[$(actor.name)] Error: $err")
on_complete!(actor::LoggerActor)                 = println(actor.io !== nothing ? actor.io : stdout, "[$(actor.name)] Completed")

struct LoggerActorFactory{O} <: AbstractActorFactory
    name :: String
    io   :: O
end

create_actor(::Type{L}, factory::LoggerActorFactory{O}) where { L, O } = LoggerActor{L, O}(factory.name, factory.io)

"""
    logger([ io::IO ], name::String = "LogActor")
    logger(::Type{T}, [ io::IO ], name::String = "LogActor") where T

Creation operator for the `LoggerActor` actor.

# Examples

```jldoctest
using Rocket

source = from([ 0, 1, 2 ])
subscribe!(source, logger())
;

# output

[LogActor] Data: 0
[LogActor] Data: 1
[LogActor] Data: 2
[LogActor] Completed
```

```jldoctest
using Rocket

source = from([ 0, 1, 2 ])
subscribe!(source, logger("CustomName"))
;

# output

[CustomName] Data: 0
[CustomName] Data: 1
[CustomName] Data: 2
[CustomName] Completed
```

```jldoctest
using Rocket

buffer = IOBuffer()

source = from([ 0, 1, 2 ])
subscribe!(source, logger(buffer, "CustomBuffer"))
;

print(String(take!(buffer)))
# output

[CustomBuffer] Data: 0
[CustomBuffer] Data: 1
[CustomBuffer] Data: 2
[CustomBuffer] Completed
```

See also: [`LoggerActor`](@ref), [`AbstractActor`](@ref)
"""
logger(name::String = "LogActor")                          = LoggerActorFactory(name, nothing)
logger(io::O, name::String = "LogActor") where { O <: IO } = LoggerActorFactory(name, io)

logger(::Type{T}, name::String = "LogActor")        where { T          } = LoggerActor{T, Nothing}(name, nothing)
logger(::Type{T}, io::O, name::String = "LogActor") where { T, O <: IO } = LoggerActor{T, O}(name, io)


text:
function is_primitive(sig)
    return (Ghost.is_primitive(sig) ||
            is_yota_primitive(sig) ||
            is_chainrules_primitive(sig))
end

########################################################################
#                            GRAD CONTEXT                              #
########################################################################

struct GradCtx
    # map from primal var to its pullback var
    # note: LittleDict is required because vars are mutable
    pullbacks::LittleDict{Variable,Variable}
    # map from primal var to its derivative var
    derivs::LittleDict{Variable,Variable}
end

GradCtx() = GradCtx(LittleDict(), LittleDict())

function rebind_context!(tape::Tape{GradCtx}, st::Dict)
    for (v, dv) in tape.c.derivs
        rebind!(tape, v, st)
        rebind!(tape, dv, st)
    end
    for (v, dv) in tape.c.pullbacks
        rebind!(tape, v, st)
        rebind!(tape, dv, st)
    end
end

########################################################################
#                              GRAD                                    #
########################################################################

const DEBUG_STATE = Ref{Any}()

getderiv(tape::Tape, v::Variable) = get(tape.c.derivs, bound(tape, v), nothing)
setderiv!(tape::Tape, x::Variable, dx::Variable) = (
    tape.c.derivs[bound(tape, x)] = bound(tape, dx)
)
hasderiv(tape::Tape, v::Variable) = getderiv(tape, v) !== nothing

function set_or_add_deriv!(tape::Tape, x::Variable, dx::Variable)
    if !hasderiv(tape, x)
        setderiv!(tape, x, dx)
    else
        old_dx = getderiv(tape, x)
        if tape[dx].val isa Tangent || tape[old_dx].val isa Tangent
            # val = tape[dx].val + tape[old_dx].val
            # new_dx_id = push!(tape, Call, val, (+), [dx.id, old_dx.id])
            new_dx = push!(tape, mkcall(+, dx, old_dx))
        else
            # val = dx.val .+ old_dx.val
            # dot_add_id = record!(tape, Constant, +)
            # new_dx_id = record!(tape, Call, val, broadcast, [dot_add_id, dx.id, old_dx.id])
            new_dx = push!(tape, mkcall(broadcast, +, dx, old_dx))
        end
        setderiv!(tape, x, new_dx)
    end
end

# there's also Core.var"#isa##kw", but I don't understand its semantics
is_kwfunc(f) = endswith(string(f), "##kw")
is_kwfunc(v::Variable) = is_kwfunc(v._op.val)

"""
Transofrm the tape replacing calls `fn(args...)` for which `ChainRule.rrule` is defined
with the following chain or calls:

    * t = rrule(fn, args...)
    * val = _getfield(t, 1)
    * pb = _getfield(t, 2)

where `val = fn(args...)` and `pb` is the pullback function.
"""
function chainrules_transform!(tape::Tape)
    config = YotaRuleConfig()
    i = 1
    while i <= length(tape)
        op = tape[V(i)]
        if (op isa Call
                && is_chainrules_primitive(call_signature(tape, op))
                && !is_yota_primitive(call_signature(tape, op)))
            # replace f(args...) with rrule(f, args...)
            # if op.fn is a kw function, use kw version of rrule
            rr_op = (is_kwfunc(op.fn) ?
                    mkcall(Core.kwfunc(rrule), op.args[1], rrule, op.args[2:end]...) :
                    mkcall(rrule, config, op.fn, op.args...))
            @assert rr_op.val !== nothing "rrule($(op.fn), ...) returned nothing"
            val_op = mkcall(_getfield, V(rr_op), 1)
            pb_op = mkcall(_getfield, V(rr_op), 2)
            tape.c.pullbacks[V(val_op)] = V(pb_op)
            replace!(tape, i => [rr_op, val_op, pb_op]; rebind_to=2)
            i += 3  # rrule + 2 _getfield ops
        else
            i += 1
        end
    end
    return tape
end

"""
Make a single step of backpropagation.
"""
function step_back!(tape::Tape, y::Variable, deriv_todo::Vector{Variable})
    @debug "step_back!() for $(tape[y])"
    df = get_deriv_function(call_signature(tape, tape[y]))
    dy = tape.c.derivs[y]
    if df !== nothing
        # Yota rules
        @debug "Found Yota derivative: $df"
        y_fargs = [tape[y].fn; tape[y].args...]
        dxs = push!(tape, mkcall(df, dy, y_fargs...))
    elseif haskey(tape.c.pullbacks, y)
        # ChainRules
        pb = tape.c.pullbacks[y]
        @debug "Found pullback: $pb"
        dxs = push!(tape, mkcall(pb, dy))
        # propage derivs to rrule variable
        rr = tape[y].args[1]
        y_fargs = is_kwfunc(rr._op.fn) ? tape[rr].args[3:end] : tape[rr].args[2:end]
    else
        sig_str = join(["::$T" for T in Ghost.call_signature(tape, tape[y]).parameters], ", ")
        error("No deriative rule found for op $(tape[y]), " *
              "try defining it using \n\n\tChainRulesCore.rrule($sig_str) = ...\n")
    end
    for (i, x) in enumerate(y_fargs)
        if x isa V
            dx = push!(tape, mkcall(getfield, dxs, i))
            @debug "Updating derivative: $x -> $dx"
            set_or_add_deriv!(tape, x, dx)
            if tape[x] isa Call
                push!(deriv_todo, x)
            end
            @debug "deriv_todo = $(join(deriv_todo, ", "))"
        end
    end
end

const BACKPROP_STATE = Ref{Tuple}()

"""
Backpropagate through the tape, record derivatives as new operations.
"""
function back!(tape::Tape; seed=1)
    # z - final variable (usually a loss)
    # y - resulting variable of current op
    # x - dependencies of y
    # dy - derivative of z w.r.t. y
    z = tape.result
    if (seed == 1) && (ndims(tape[z].val) > 0)
        error("Gradient of a vector-valued function requires a seed")
    end
    dy = push!(tape, Constant(seed))
    # save seed var to use in compilation later
    tape.meta[:seed] = dy
    # set initial derivative value
    tape.c.derivs[z] = dy
    # queue of variables to calculate derivatives for
    deriv_todo = V[z]
    deriv_processed = Set{V}()
    while !isempty(deriv_todo)
        y = popfirst!(deriv_todo)
        ## skip double calculation for the same var in multupath graphs
        if y in deriv_processed
            @debug "Already processed $y, skipping it"
            continue
        end
        try
            step_back!(tape, y, deriv_todo)
        catch e
            BACKPROP_STATE[] = (tape, y, deriv_todo)
            rethrow(e)
        end
        push!(deriv_processed, y)
    end
end

"""
    gradtape(f::Union{Function, DataType}, args...; seed=1)
    gradtape!(tape::Tape; seed=1)

Calculate and record to the tape gradients of `tape[tape.resultid]` w.r.t. `Input` nodes.
See grad() for more high-level API.
"""
function gradtape!(tape::Tape; seed=1)
    # update chainrules if number of rrule or no_rrule methods has changed
    update_chainrules_primitives!(force=false)
    # apply transformations needed for ChainRules
    chainrules_transform!(tape)
    # backpropagate gradients
    back!(tape; seed=seed)
    # add a tuple of (val, (gradients...))
    deriv_vars = [hasderiv(tape, v) ? getderiv(tape, v) : ZeroTangent() for v in inputs(tape)]
    deriv_tuple = push!(tape, mkcall(tuple, deriv_vars...))
    deriv_tuple_unthunked = push!(tape, mkcall(map, ChainRules.unthunk, deriv_tuple))
    new_result = push!(tape, mkcall(tuple, tape.result, deriv_tuple_unthunked))
    tape.result = new_result
    return tape
end

function gradtape(f::Union{Function,DataType}, args...; seed=1)
    _, tape = trace(f, args...; is_primitive=is_primitive, ctx=GradCtx())
    tape = gradtape!(tape; seed=seed)
    return tape
end

"Like Ghost.compile, but adds Yota specific ops"
function grad_compile(tape::Tape)
    ex = Ghost.to_expr(tape)
    seed_var = tape.meta[:seed]
    seed_default_val = tape[seed_var].val
    insert!(ex.args[1].args, 2, Expr(:parameters, Expr(:kw, :seed, seed_default_val)))
    seed_var_name = Ghost.make_name(seed_var.id)
    op_exprs = ex.args[2].args
    for op_ex in op_exprs
        if op_ex.args[1] == seed_var_name
            op_ex.args[2] = :seed
            break
        end
    end
    return Base.eval(@__MODULE__, ex)
end

const GRAD_CACHE = Dict{Any,Any}()

reset!() = empty!(GRAD_CACHE)

"""
    grad(f, args...; seed=1)

Find gradient of a callable `f` w.r.t. its arguments.

`grad()` returns two things: value of `f(args...)` and a tuple of
grafients w.r.t. to its inputs (including the callable itself).

```jldoctest
val, g = grad(x -> sum(x .+ 1), [1.0, 2.0, 3.0])

# output
(9.0, (ZeroTangent(), [1.0, 1.0, 1.0]))
```

By default, `grad()` expects the callable to return a scalar.
Vector-valued functions can be differentiated if a seed (starting value)
is provided. Seed is equivalent to the vector in VJP notation.

```jldoctest
val, g = grad(x -> 2x, [1.0, 2.0, 3.0]; seed=ones(3))

# output
([2.0, 4.0, 6.0], (ZeroTangent(), [2.0, 2.0, 2.0]))
```

All gradients can be applied to original variables using `update!()` function.

See also: [gradtape](@ref)
"""
function grad(f::Union{Function,DataType}, args...; seed=1)
    # key consists of function type and type of argument (for structs) or its size
    cache_key = (f, ([isstruct(arg) ? typeof(arg) : size(arg) for arg in args]...,))
    if haskey(GRAD_CACHE, cache_key)
        gf = GRAD_CACHE[cache_key]
        return gf(f, args...; seed=seed)
    else
        tape = gradtape(f, args...; seed=seed)
        gf = grad_compile(tape)
        GRAD_CACHE[cache_key] = gf
        return tape.retval
    end
end


text:
# Allows UMAP to be called on Experiment objects
function UMAP.umap(e::Experiment, n_components=2; kwargs...)
    umap(convert(Matrix, getdata(e))', n_components; kwargs...)
end


text:
Point(::AnyDomain) = Point(NaN)
Point{T}(x::AnyDomain) where {T} = new(T(NaN))

convert(::Type{Point},::AnyDomain) = Point(NaN)
convert(::Type{Point{T}},::AnyDomain) where T = Point{T}(NaN)

isambiguous(d::Point) = isnan(d.x)

norm(p::Point) = norm(p.x)

getindex(p::Point,k...) = p.x[k...]
first(p::Point) = p.x
last(p::Point) = p.x

Base.isnan(d::Point) = false

issubset(a::Point,d::UnionDomain) = a.x in d

intersect(a::Point,b::Point) = b.x in a ? b : EmptyDomain()
intersect(a::UnionDomain,b::Point) = b.x in a ? b : EmptyDomain()
intersect(a::Domain,b::Point) = b.x in a ? b : EmptyDomain()
intersect(b::Point,a::UnionDomain) = b.x in a ? b : EmptyDomain()
intersect(b::Point,a::Domain) = b.x in a ? b : EmptyDomain()

setdiff(a::Point,b::Point) = a==b ? EmptyDomain() : a
reverseorientation(a::Point) = a

canonicaldomain(a::Point) = Point(0.)
tocanonical(a::Point,x) = x-a.x
fromcanonical(a::Point,x) = x+a.x

points(a::Point,n) = eltype(a)[a.x]
checkpoints(a::Point) = eltype(a)[a.x]


text:
##%% Load the packages

# First, we load the packages.

using ProximalOperators, LinearAlgebra, TSVD, SparseArrays, OSQP, JuMP, MosekTools

##%% All the necessary types

# # ProxRegularSet
#  This is an abstract type that describes the nature of the set.

abstract type ProxRegularSet end

# # problem
# This structure contains the function, the set, value of parameter β, and an intial point.

struct Problem{F <: ProximableFunction, S <: ProxRegularSet, R <: Real, A <: AbstractVecOrMat{ <:Real}}

    f::F # the objective function
    C::S # the constraint set
    β::R # what value of β to pick
    z0::A # one intial state provided by the user

    ## consider having a constructor that will create an intial staate if the user does not give one

end

##%% Settings function

# # settings
#  This struct represents the user settings. It contains information regarding user specified value of different parameters

struct Settings

    # first comes description of different data type

    μ_max::Float64 # starting μ
    μ_min::Float64 # ending μ
    μ_mult_fact::Float64 # multiplicative factor to reduce μ
    n_iter_min::Int64 # minimum number of iterations per fixed μ
    n_iter_max::Int64 # maximum number of iterations per fixed μ
    big_M::Float64 # a big number, it suffices to hard code it, but for now let me keep it
    β::Float64 # β is the strong convexity parameter
    tol::Float64 # tolerance for the fixed point gap
    verbose::Bool # whether to print iteration information
    freq::Int64 # how often to print the iteration solution
    γ_updt_rule::Symbol # for now two rule: :safe and :adaptive, deafult will be :safe

    # constructor function
    function Settings(;
        μ_max = 1.0,
        μ_min = 1e-12,
        μ_mult_fact = 0.85,
        n_iter_min = 1000,
        n_iter_max = 1000,
        big_M = 1e99,
        β = 1e-10,
        tol = 1e-8,
        verbose = true,
        freq = 10,
        γ_updt_rule = :safe
        )
        new(μ_max, μ_min, μ_mult_fact, n_iter_min, n_iter_max, big_M, β, tol, verbose, freq, γ_updt_rule)
    end

end

## structure that describes the state (for the outer loop)

# # state
# This structure contains all the necessary information to describe a state

mutable struct State{T <: AbstractVecOrMat{<: Real}, I <: Integer, R <: Real} # Done
    x::T # the first iterate
    y::T # the second iterate
    z::T # the third iterate
    i::I # iteration number
    fxd_pnt_gap::R # fixed point gap for state
    fsblt_gap::R # how feasible the current point is, this is measured by || x- Π_C(x) ||
    μ::R # current value of μ
    γ::R # current value of γ
end

# Let us initialize the state data structure.

function State(problem::Problem, settings::Settings)

    # this function constructs the very first state for the problem and settings
    z0 = copy(problem.z0)
    x0 = zero(z0)
    y0 = zero(z0)
    i = copy(settings.n_iter_max) # iteration number where best fixed point gap is reached
    fxd_pnt_gap = copy(settings.big_M)
    fsblt_gap = copy(settings.big_M)
    μ = copy(settings.μ_max)
    γ = γ_from_μ(μ, settings) #

    return State(x0, y0, z0, i, fxd_pnt_gap, fsblt_gap, μ, γ)

end

# Time to describe the data structure that we require to initialize our algorithm.

## structure that describes the information to intialize our algorithm:

mutable struct InitInfo{T <: AbstractVecOrMat{<: Real}, R <: Real}

    # this will have z0, γ, μ
    z0::T # initial condition to start an inner iteration
    μ::R # μ required start an inner iteration
    γ::R # γ required to start an inner iteration

end

# Of course, we will need a constructor for the InitInfo type, given the user input. The user inputs are problem::problem, settings::settings.

function InitInfo(problem::Problem, settings::Settings)

    ## constructs the initial NExOS information
    z0 = copy(problem.z0) # need to understand if copy is necessary
    μ = copy(settings.μ_max)
    γ = γ_from_μ(μ, settings)

    return InitInfo(z0, μ, γ)

end

#%%
# Time to define our sets. We have two prox regular sets for now:
# 1. SparseSet: It is of the form: card(x) ≦ r, ||x||_∞ ≦ M
# 2. RankSet: It is of the form: rank(x) ≦ r, ||x||_2 ≦ M

# We also consider the unit hypercube, which is not prox-regular however is defined here to apply our algorithm as a heuristics to the 3-SAT problem.

# define sparse set first, via the `SparseSet` structure.

struct SparseSet{R <: Real, I <: Integer} <: ProxRegularSet
    M::R # M is the upper bound on the elements of x
    r::I # r is the maximum cardinality
    function SparseSet{R, I}(M::R, r::I) where {R <: Real, I <: Integer}
        if r <= 0 || M <= 0
            error("parameter M and r must be a positive integer")
        else
            new(M, r)
        end
    end
end

# Constructor for the sparse set
SparseSet(M::R, r::I) where {R <: Real, I <: Integer} = SparseSet{R,I}(M, r)

# Okay, we have done everything for the sparse set. Now, we do the same for the low rank set. First, we define the structure.

struct RankSet{R<: Real, I <: Integer} <: ProxRegularSet
    M::R
    r::I
    function RankSet{R, I}(M::R, r::I) where {R <: Real, I <: Integer}
        if r <= 0 || M <= 0
            error("parameter M and r must be a positive integer")
        else
            new(M, r)
        end
    end
end

# Let us define the constructor for the low rank set.

RankSet(M::R, r::I) where {R <: Real, I <: Integer} = RankSet{R, I}(M, r)

# Time to define the unit hypercube: {0,1}^n

struct UnitHyperCube{I <: Integer} <: ProxRegularSet
    n::I
    function  UnitHyperCube{I}(n::I) where {I <: Integer}
        if n <= 0
            error("parameter n must be a positive integer")
        else
            new(n)
        end
    end
end

# Constructor for the unit hypercube

UnitHyperCube(n::I) where {I <: Integer} = UnitHyperCube{I}(n)

## Define the Least squarers function for affine rank minimization:

## Start with some helper functions

function Afn_Op_ith_Mat(𝓐, k, i)
    ## gives the i-th matrix of the affine oerator of 𝓐, such that
    ## [𝓐(X)]_i = trace(A_i^T X)
    m, kn = size(𝓐)
    n = convert(Int64, kn/k)
    m, kn = size(𝓐)
    n = convert(Int64, kn/k)
    return 𝓐[:, (i-1)*n+1:i*n]
end

# TODO: Make this function faster
function Affn_Op(𝓐, k, X)
    # Computes 𝓐(X), k is the dimension of the affine operator
    # 𝓐 = [A_1 A_2 ... A_k] ∈ R^ {m × kn}
    m, n = size(X)
    m1, kn = size(𝓐)
    n1 = convert(Int64, kn/k)
    if m1 ≠ m || n1 ≠ n
        @error "sizes of the input matrix and the affine operator matrix do not match"
    end
    z = zeros(k) # output will be stored in this matrix
    for i in 1:k
        z[i] = tr(Afn_Op_ith_Mat(𝓐, k, i)'*X)
    end
    return z
end

function Affn_Op_Mat(𝓐, k, X)
    # Computes 𝓐(X), k is the dimension of the affine operator
    # 𝓐 = [A_1 A_2 ... A_k] ∈ R^ {m × kn}
   A = affine_operator_to_matrix(𝓐, k)
   # output will be stored in this matrix
   return A*vec(X)
end

# TODO: need to make this function faster
function affine_operator_to_matrix(𝓐::AbstractMatrix{R}, k::I) where {R <: Real, I <: Integer}
    ## convert 𝓐 to a matrix A of size k × mn
    ## 𝓐 = [A_1 A_2 ... A_k] ∈ R^ {m × kn}, where each
    ## A_i is an m × n matrix, and our
    ## AffineOp(X) = [ tr(A_1^T X) ; tr(A_2^T X); ... ; tr(A_k^T X) ], where X is an m × n matrix
    ## our goal is coming up with a matrix A such that
    ## AffineOp(X) = A*vec(X)
    m, kn = size(𝓐)
    if mod(kn, k) ≠ 0
        @error "dimension of the matrix [A_1 ... A_k] is not correct"
    end
    n = convert(Int64, kn/k)
    A = zeros(k,m*n)
    for i in 1:k
        A[i,:] = transpose(vec(𝓐[:, (i-1)*n+1:i*n]))
    end
    return A
end

## Let us define the function least squares over matrices

abstract type LeastSquaresOverMatrix <: ProximableFunction end

is_smooth(f::LeastSquaresOverMatrix) = true

is_quadratic(f::LeastSquaresOverMatrix) = true

fun_name(f::LeastSquaresOverMatrix) = "Least squares penalty over matrices"

function LeastSquaresOverMatrix(𝓐::M, b::V, k::I, lam::R=R(1); iterative=false) where {R <: Real, V <: AbstractArray{R}, I <: Integer, M} # BUG: seems that this version of the objective value has a bug, for now use the other one that directly takes a k×mn matrix
    # FIXME: Fix this function later on, for now work with the other version
    ## In this case, 𝓐 is as follows:
    ## 𝓐 = [A_1 A_2 ... A_k] ∈ R^ {m × kn}, where each
    ## A_i is an m × n matrix, and
    # b is the output observation with k entries
    # this function will convert 𝓐 into a matrix Amat of size k × mn that can operate on vec(X) of size mn, which is the vectorized version of the input matrix X of size m × n
    Amat = affine_operator_to_matrix(𝓐, k)
    if iterative == false
        ProximalOperators.LeastSquaresDirect(Amat, b, lam)
    else
        ProximalOperators.LeastSquaresIterative(Amat, b, lam)
    end
end

function LeastSquaresOverMatrix(𝐀::M, b::V, lam::R=R(1); iterative=false) where {R <: Real, V <: AbstractArray{R}, I <: Integer, M}
    # in this case 𝐀 is given as a k × mn matrix, so that it acts directly on the vectorized operator
    if size(𝐀)[1] != length(b)
        @error "number of rows of A must be equal to number of elements of b "
    end
    if iterative == false
        ProximalOperators.LeastSquaresDirect(𝐀, b, lam)
    else
        ProximalOperators.LeastSquaresIterative(𝐀, b, lam)
    end
end

# evaluating the LeastSquaresOverMatrix
# for direct
function (f::ProximalOperators.LeastSquaresDirect)(X::Array{R,2}) where {R <: Real}
    return f(vec(X))
end

# for iterative
function (f::ProximalOperators.LeastSquaresIterative)(X::Array{R,2}) where {R <: Real}
    return f(vec(X))
end

## Function for squared loss for matrix completion problem

# data conversion for the matrix completion problem

# older function: takes much longer to compute
# function matrix_completion_A_b(Zobs::AbstractMatrix{R}) where {R <: Real}
#     # The function will take the matrix Zobs of size mxn such that has it has many missing (NaN) values and few observed value. Let the index of available values of Zobs be Ω. Then the function will create a matrix A and vector b, such that for any matrix X we have
#     #  \[
#     # f(X)=\sum_{(i,j)\in\Omega}(X_{ij}-Z_{ij})^{2}=\|Ax-b\|_{2}^{2}.
#     # \]
#     m, n = size(Zobs)
#     zobs = vec(Zobs)
#     # find the corresponding indices of the missing data again:
#     index_available_vec = findall(x -> isnan(x)== false, zobs)
#     b = zobs[index_available_vec]
#     p = length(index_available_vec)
#     A = spzeros(p, m*n) # m is size of b,
#     for i in 1:p
#         A[i,index_available_vec[i]] = 1.0
#     end
#     return A, b
# end

function matrix_completion_A_b(Zobs::AbstractMatrix{R}) where {R <: Real}
    # The function will take the matrix Zobs of size mxn such that has it has many missing (NaN) values and few observed value. Let the index of available values of Zobs be Ω. Then the function will create a matrix A and vector b, such that for any matrix X we have
    #  \[
    # f(X)=\sum_{(i,j)\in\Omega}(X_{ij}-Z_{ij})^{2}=\|Ax-b\|_{2}^{2}.
    # \]
    m, n = size(Zobs)
    zobs = vec(Zobs)
    # find the corresponding indices of the missing data again:
    index_available_vec = findall(x -> isnan(x)== false, zobs)
    b = zobs[index_available_vec]
    p = length(index_available_vec)
    I_A = 1:p
    J_A = index_available_vec
    V_A = ones(p)
    A = sparse(I_A,J_A,V_A,p,m*n)
    return A, b
end

# time to define the squared loss function

abstract type SquaredLossMatrixCompletion <: ProximableFunction end

is_smooth(f::SquaredLossMatrixCompletion) = true

is_quadratic(f::SquaredLossMatrixCompletion) = true

fun_name(f::SquaredLossMatrixCompletion) = "Squared loss function for matrix completion problem"

function SquaredLossMatrixCompletion(Zobs::AbstractMatrix{R}, lam::R=R(1); iterative = false) where {R <: Real}
    # construct the data from Zobs
    A, b = matrix_completion_A_b(Zobs)

    if iterative == false
        ProximalOperators.LeastSquaresDirect(A, b, lam)
    else
        ProximalOperators.LeastSquaresIterative(A, b, lam)
    end

end


text:
__precompile__()
module TestContinuous
using OpenEphysLoader
using Main.TestUtilities, Main.TestOriginal

using Test

using Random

# Helper functions to test OpenEphysLoader's handling of
# continuous files

export test_OEContArray,
    to_block_contents,
    verify_BlockBuffer,
    damaged_file,
    write_continuous,
    bad_blockhead,
    bad_blocktail,
    rand_block_data,
    to_OE_bytes,
    bad_file

### Test array contents ###
function test_OEContArray(
    io::IOStream,
    ::Type{T},
    testtypes::Vector{DataType},
    D::Vector,
    nblock::Integer,
    recno::Integer,
    startsamp::Integer,
) where T<:OEContArray
    test_OEContArray_interface(T)
    seekstart(io)
    A = T(io)
    test_OEContArray_parts(A, D, nblock, recno, startsamp)
    test_OEContArray(A, D, recno, startsamp)
    for t in testtypes
        seekstart(io)
        A = T(t, io)
        test_OEContArray(A, D, recno, startsamp)
    end
end
function test_OEContArray(
    A::OEContArray,
    D::Vector,
    recno::Integer,
    startsamp::Integer
)
    test_OEArray_interface(A)
    test_OEContArray_contents(A, D, recno, startsamp)
end

function test_OEContArray_parts(
    A::OEContArray,
    D::Vector,
    nblock::Integer,
    recno::Integer,
    startsamp::Integer
)
    nd = length(D)
    verify_ContinuousFile(A.contfile, nd, nblock)
    nblocksamp = OpenEphysLoader.CONT_REC_N_SAMP
    for blockno = 1:fld(nd, nblocksamp)
        block_data, block_body, blockstart =
            to_block_contents(D, blockno, startsamp)
        blockidxes = block_idxes(blockno)
        blockstart = startsamp + blockidxes[1] - 1
        OpenEphysLoader.prepare_block!(A, blockidxes[1])
        @test A.blockno == blockno
        verify_BlockBuffer(A.block,
                           blockstart,
                           recno,
                           block_body,
                           block_data)
    end
end

function block_idxes(blockno::Integer)
    nblocksamp = OpenEphysLoader.CONT_REC_N_SAMP
    return (blockno - 1) * nblocksamp .+ (1:nblocksamp)
end

function to_block_contents(D::Vector, blockno::Integer, startsamp::Integer)
    blockidxes = block_idxes(blockno)
    block_data = D[blockidxes]
    block_body = copy(block_data)
    block_body = reinterpret(UInt8, block_body)
    block_startsamp = startsamp + (blockno - 1) * OpenEphysLoader.CONT_REC_N_SAMP
    return block_data, block_body, block_startsamp
end

function test_OEContArray_contents(
    A::SampleArray{T}, D::Vector, varargs...
) where {T<:Real}
    @test OpenEphysLoader.block_data(A, 1) == A.block.data[1]
    @test OpenEphysLoader.block_data(A, 1024) == A.block.data[1024]
    @test [OpenEphysLoader.convert_data(A, A.block.data[1])] ==
        sample_conversion(T, A.block.data[1:1])
    verify_samples(copy(A), D)
end
function test_OEContArray_contents(
    A::TimeArray{T},
    D::Vector,
    ::Integer,
    startsamp::Integer
) where {T<:Real}
    @test OpenEphysLoader.block_data(A, 1) == A.block.timestamp
    @test OpenEphysLoader.block_data(A, 2) == A.block.timestamp + 1
    @test [OpenEphysLoader.convert_data(A, A.block.timestamp)] ==
        time_conversion(T, A.block.timestamp, 1)
    verify_times(copy(A), startsamp, length(D))
end
function test_OEContArray_contents(
    A::RecNoArray{T},
    D::Vector,
    recno::Integer,
    ::Integer
) where {T<:Integer}
    @test OpenEphysLoader.block_data(A, 1) == recno
    @test OpenEphysLoader.block_data(A, 2) == recno
    @test OpenEphysLoader.convert_data(A, A.block.recordingnumber) ==
        convert(T, recno)
    verify_recnos(copy(A), recno, length(D))
end
function test_OEContArray_contents(
    A::JointArray{Tuple{S,T,R}},
    D::Vector,
    recno::Integer,
    startsamp::Integer
) where {S<:Real,T<:Real,R<:Integer}
    @test OpenEphysLoader.block_data(A, 1) == (A.block.data[1],
                                               A.block.head.timestamp,
                                               recno)
    @test OpenEphysLoader.block_data(A, 2) == (A.block.data[2],
                                               A.block.head.timestamp + 1,
                                               recno)
    datatup = (
        A.block.data[1],
        A.block.head.timestamp,
        A.block.head.recordingnumber
    )
    answertup = (
        sample_conversion(T, A.block.data[1:1])[1],
        time_conversion(T, A.block.head.timestamp, 1)[1],
        convert(T, recno)
    )
    @test OpenEphysLoader.convert_data(A, datatup) == answertup
    samples = S[tup[1] for tup in A]
    times = T[tup[2] for tup in A]
    recnos = R[tup[3] for tup in A]
    nd = length(D)
    verify_samples(samples, D)
    verify_times(times, startsamp, nd)
    verify_recnos(recnos, recno, nd)
end

verify_samples(A::Vector{T}, D::Vector) where {T} = @test sample_conversion(T, D) == A
function verify_times(A::Vector{T}, startsamp::Integer, nsamp::Integer) where T
   @test time_conversion(T, startsamp, nsamp) == A
end
function verify_recnos(A::Vector, recno::Integer, nsamp::Integer)
    @test fill(recno, (nsamp,)) == A
end

### Test interfaces ###
function test_OEArray_interface(A::OEArray)
    @assert !isempty(A) "Tests require a non-empty array"
    @test !isempty(size(A))
    @test !isempty(A[1])
    @test !isempty(A[end])
    @test isa(eltype(A), Type)
    @test !isempty(length(A))
    @test_throws ErrorException A[1] = A[1]
end

function test_OEContArray_interface(::Type{T}) where T<:OEContArray
    fields = fieldnames(T)
    @test :contfile in fields
    @test :block in fields
    @test :blockno in fields
    @test :check in fields
    @test !isempty(methods(T, (IOStream,)))
end

### Verify data types ###
function verify_BlockBuffer(
    block::OpenEphysLoader.DataBlock,
    t::Integer,
    rec::Integer,
    body::AbstractVector{UInt8},
    data::Vector{OpenEphysLoader.CONT_REC_SAMP_BITTYPE},
)
    verify_BlockBuffer(block.head, t, rec)
    @test block.body == body
    @test block.data == data
    @test block.tail == OpenEphysLoader.CONT_REC_END_MARKER
end
function verify_BlockBuffer(
    block::OpenEphysLoader.BlockHeader,
    t::Integer,
    rec::Integer,
    ::AbstractVector{UInt8},
    ::Vector{OpenEphysLoader.CONT_REC_SAMP_BITTYPE},
)
    verify_BlockBuffer(block, t, rec)
end
function verify_BlockBuffer(
    block::OpenEphysLoader.BlockHeader,
    t::Integer,
    rec::Integer
)
    @test block.timestamp == t
    @test block.nsample == OpenEphysLoader.CONT_REC_N_SAMP
    @test block.recordingnumber == rec
end

function verify_ContinuousFile(C::ContinuousFile, nsample::Integer, nblock::Integer)
    @test isopen(C.io) && !eof(C.io)
    @test C.nsample == nsample
    @test C.nblock == nblock
    verify_header(C.header)
end

### Functions to write .continuous files ###
bad_file(io::IOStream) = write(io, "These aren't the droids you're looking for")

function damaged_file(io::IOStream, args...; kwargs...)
    write_continuous(io, args...; kwargs...)
    bad_blockhead(io)
end

function write_continuous(
    io::IOStream,
    d::AbstractArray{T, 1},
    recno::Integer = 0,
    startsamp::Integer = 1,
    recdelay::Integer = 0
) where T<:Integer
    nd = length(d)
    if recdelay >= OpenEphysLoader.CONT_REC_N_SAMP
        error("Delay between file start and recording is too long")
    end
    nstoppad = mod(-(nd + recdelay), OpenEphysLoader.CONT_REC_N_SAMP)
    nblock = cld(nd + recdelay, OpenEphysLoader.CONT_REC_N_SAMP) # ceiling divide
    if nstoppad > 0 || recdelay > 0
        padded = zeros(Int, OpenEphysLoader.CONT_REC_N_SAMP * nblock)
        padded[(1 + recdelay):(end - nstoppad)] = d
    else
        padded = d # no padding needed, renaming for clarity below
    end
    write_fheader_fun()(io)
    tblock = startsamp
    offset = 0
    for blockno in 1:nblock
        writeblock(
            io,
            view(
                padded,
                offset .+ (1:OpenEphysLoader.CONT_REC_N_SAMP)
            ),
            tblock,
            recno
        )
        tblock += OpenEphysLoader.CONT_REC_N_SAMP
        offset += OpenEphysLoader.CONT_REC_N_SAMP
    end
end
function write_continuous(path::String, args...)
    open(path, "w") do io
        write_continuous(io, args...)
    end
end

### Functions to write data blocks ###
good_block(io::IOStream, d::AbstractArray, t::Integer, r::Integer) = writeblock(io, d, t, r)

function bad_blockhead(io::IOStream)
    blockdata = rand_block_data()
    writeblock(io, blockdata; bad_blockhead = true)
end

function bad_blocktail(io::IOStream)
    blockdata = rand_block_data()
    writeblock(io, blockdata; bad_blocktail = true)
end

function writeblock(
    io::IOStream,
    d::AbstractArray,
    t::Integer = 1,
    recno::Integer = 0;
    bad_blockhead::Bool = false,
    bad_blocktail::Bool = false
)
    # write block header
    write(io, OpenEphysLoader.CONT_REC_TIME_BITTYPE(t))
    nsamp_bittype = OpenEphysLoader.CONT_REC_N_SAMP_BITTYPE
    if bad_blockhead
        write(io, zero(nsamp_bittype))
    else
        write(io, nsamp_bittype(OpenEphysLoader.CONT_REC_N_SAMP))
    end
    write(io, OpenEphysLoader.CONT_REC_REC_NO_BITTYPE(recno))

    # Convert data to open ephys format
    oebytes = to_OE_bytes(d)
    # write data
    write(io, oebytes)

    # write tail
    if bad_blocktail
        write(io, b"razzmatazz")
    else
        write(io, OpenEphysLoader.CONT_REC_END_MARKER)
    end
end
function to_OE_bytes(D::AbstractArray{T,1}) where T<:OpenEphysLoader.CONT_REC_SAMP_BITTYPE
    contents = copy(D)
    for i in eachindex(contents)
        @inbounds contents[i] = hton(contents[i])
    end
    oebytes = reinterpret(UInt8, contents)
    return oebytes
end

### Utility functions ###
sample_conversion(::Type{T}, data::AbstractArray{R, 1}) where {T<:Integer, R<:Integer} =
    Vector{T}(copy(data))
sample_conversion(::Type{T}, data::AbstractArray{R, 1}) where {T<:AbstractFloat, R<:Integer} =
    0.195 * Vector{T}(copy(data))
function time_conversion(
    ::Type{T},
    startsamp::Integer,
    nsamp::Integer
) where T<:Integer
    return T[t for t in startsamp:1:(startsamp + nsamp - 1)]
end
function time_conversion(
    ::Type{T},
    startsamp::Integer,
    nsamp::Integer
) where T<:AbstractFloat
    timepoints = (time_conversion(Int, startsamp, nsamp) .- 1) / 30000
    converted_times = similar(timepoints, T)
    copyto!(converted_times, timepoints)
    return converted_times
end

function rand_block_data()
    return rand(
        OpenEphysLoader.CONT_REC_SAMP_BITTYPE,
        OpenEphysLoader.CONT_REC_N_SAMP
    )
end

end


text:
# Example from Improved relaxations for the parametric solutionsof ODEs using differential inequalities
id = :ScottA2013b

pL = [1.0; 1.0; 0.001]
pU = [1200.0; 1200.0; 100.0]
ScottA2013b_x0(p::Vector{T}) where T = [1.53*10^(-4)*one(T); 0.4*one(T); zero(T);
                                        zero(T); 0.0019*one(T); zero(T); zero(T);
                                        zero(T); zero(T); zero(T)]
function ScottA2013b_f!(du::Vector{T}, u::Vector{T}, p::Vector{T}, t) where T

    k = T[53.0*one(T); p[1]; p[2]; p[3]; 1200.0*one(T)]
    K2 = 2081.0
    K3 = 4162.0

    S = zeros(10,5)
    S[1,1] = -1.0; S[2,1] = -1.0; S[3,1] = 1.0;  S[3,2:3] = -1.0; S[3,5] = -2.0;
    S[4,1] = 1.0; S[5,2] = -1.0;  S[5,3] = -1.0; S[6,2] = 1.0;
    S[7,3] = 1.0; S[7,4] = -1.0;  S[8,4] = 1.0;  S[9,4] = 1.0;  S[10,5] = 1.0

    r = zeros(T, 5)
    r[1] = k[1]*u[1]*u[2]
    r[2] = k[2]*(u[3]*u[5] - (1.0/K2)*u[6])
    r[3] = k[3]*(u[3]*u[5] - (1.0/K3)*u[7])
    r[4] = k[4]*u[7]
    r[5] = k[5]*u[3]^2

    du[:] = S*r
    return
end
tspan = (0.0, 2.0)

prob = ODERelaxProb(ScottA2013b_f!, tspan, ScottA2013b_x0, pL, pU)
url = "https://link.springer.com/article/10.1007%2Fs10898-012-9909-0"
source = "ScottA2013b"
desc = "Cyclohexadienyl Radical Oxidation"


text:
println("Running tests on ", gethostname(), ".")
using Test, Random, LinearAlgebra, Statistics
using QMC
using KrylovKit, BlockArrays, MonteCarloObservable, FFTW

@testset "All Tests" begin

    # General tests
    include("tests_linalg.jl")
    include("tests_statistics.jl")

    # O3 model
    include("tests_O3.jl")

    # no interactions
    include("tests_freefermions.jl")

    # Resume
    # include("tests_resume.jl")

end

nothing

max_stars_repo_name: crstnbr/dqmc

max_issues_repo_name: crstnbr/dqmc

max_forks_repo_name: crstnbr/dqmc


text:
# --- User parameters

n = 100
num_samples = 10000

# --- Imports

using BenchmarkTools
using LinearAlgebra
using Statistics

# --- Functions

"""
    Compute QR factorization of `A` using the modified Gram-Schmidt algorithm.
"""
function qr_mgs(A::Matrix)
    # Initialize Q and R
    Q = copy(A)
    R = zero(A)

    for j in 1:size(Q, 2)
        norm = LinearAlgebra.norm(Q[:, j])
        if norm > 0
            Q[:, j] /= norm
            R[j, j] = norm
        else
            continue
        end

        column = Q[:, j]
        for j_inner in (j + 1):size(Q, 2)
            R[j, j_inner] = Q[:, j_inner] ⋅ column
            Q[:, j_inner] -= R[j, j_inner] * column
        end
    end

    Q, R
end;

# --- Test qr_mgs()

n = 10
A = randn(n, n)
Q, R = qr_mgs(A)

println("Relative error Q * R: ", opnorm(A - Q * R)/opnorm(A))
println("Absolute error Q * R: ", opnorm(A - Q * R))
println("opnorm(Q' * Q): ", opnorm(transpose(Q) * Q))

det_A = det(A)
det_Q = det(Q)
det_R = det(R)
println("det(A): ", det_A)
println("det(Q): ", det_Q)
println("det(R): ", det_R)
println("Relative error det(Q * R):", (det_A - det_Q*det_R) / det_A)
println("Absolute error det(Q * R):", det_A - det_Q*det_R)

# Initialize data vectors
mgs_orthogonality_errors = zeros(num_samples)
householder_orthogonality_errors = zeros(num_samples)

mgs_det_errors = zeros(num_samples)
householder_det_errors = zeros(num_samples)

# Collect data
for i in 1:num_samples
    # Generate random matrix
    A = randn(n, n)
    det_A = det(A)

    # Compute QR factorization using modified Gram-Schmidt algorithm
    Q_mgs, R_mgs = qr_mgs(A)

    mgs_orthogonality_errors[i] = opnorm(transpose(Q_mgs) * Q_mgs - LinearAlgebra.I)

    det_R = det(R_mgs)
    mgs_det_errors[i] = abs((abs(det_A) - abs(det_R)) / det_A)

    # Compute QR factorization using Householder triangularization
    F_householder= qr(A)
    Q_householder = F_householder.Q
    R_householder = F_householder.R

    householder_orthogonality_errors[i] =
        opnorm(transpose(Q_householder) * Q_householder - LinearAlgebra.I)

    det_R = det(R_householder)
    householder_det_errors[i] = abs((abs(det_A) - abs(det_R)) / det_A)
end

# --- Orthogonality Error

println("mean(mgs_orthogonality_errors): ", mean(mgs_orthogonality_errors))
println("std(mgs_orthogonality_errors): ", std(mgs_orthogonality_errors))

println("mean(householder_orthogonality_errors): ", mean(householder_orthogonality_errors))
println("std(householder_orthogonality_errors): ", std(householder_orthogonality_errors))

# --- Determinant Error

println("mean(mgs_det_errors): ", mean(mgs_det_errors))
println("std(mgs_det_errors): ", std(mgs_det_errors))

println("mean(householder_det_errors): ", mean(householder_det_errors))
println("std(householder_det_errors): ", std(householder_det_errors))

# --- Special cases that demonstrate large loss of orthogonality with MGS

A = [0.700000 0.70711; 0.70001 0.70711]
det_A = det(A)

# MGS
Q_mgs, R_mgs = qr_mgs(A)
println("MGS orthogonality error: ", opnorm(transpose(Q_mgs) * Q_mgs - LinearAlgebra.I))
println("MGS determinant error: ", abs((abs(det_A) - abs(det(R_mgs))) / det_A))

# Householder triangularization
F_householder = qr(A)
Q_householder = F_householder.Q
R_householder = F_householder.R
println("Householder triangularization orthogonality error: ",
        opnorm(transpose(Q_householder) * Q_householder - LinearAlgebra.I))
println("Householder triangularization determinant error: ",
        abs((abs(det_A) - abs(det(R_householder))) / det_A))

max_stars_repo_name: ktchu/RESEARCH-math-exploration

max_issues_repo_name: ktchu/RESEARCH-math-exploration

max_forks_repo_name: ktchu/RESEARCH-math-exploration