# specific to Julia, that is, related to the output file format (language).

#=
F: read the file containing constants
I:
O:
BA
=#
function py2_include_constants_from_literature(src_file_name, pad_string)
  # create src_buffer -
  src_buffer = ""
  # path to distrubtion -
  path_to_src_file = src_file_name
  open(path_to_src_file,"r") do src_file
    for line in eachline(src_file)
      src_buffer *= pad_string*"$line"
    end
  end
  return src_buffer
end

#=
F: generate kinetic buffer
I:
O: return monod affinity constant symbol, mRNA species array,
  protein species array, W_string_array, disassociation_const_string_array.
BA
=#
function py2_build_kinetics_buffer(all_species_dict::Dict{String, Int},
  all_rnx_list::Array, all_txtl_dict::Dict, sys2user::Dict)
  kinetics = "from __future__ import division" *
             "\nimport numpy as np" *
             "\n\ndef calculate_kinetics(X, data_dictionary):" *
             "\n\n\t# load all species dictionary" *
             "\n\tall_species_dict = data_dictionary[\"all_species_dict\"]  # String-->Int"
  # for signaling
  kinetics *= "\n\n\t# generate kinetics for signaling, assume saturation kinetics"
  totalRnxNo = length(all_rnx_list)
  kinetics *= "\n\t# load Monod affinity constants" *
              "\n\tMonodK = data_dictionary[\"MonodAffinityConstantDict\"]  # String-->Float"
  MonodAffinityConstant_String_Array = Array{String,1}()  # collect monod affinity constant symbol
  kinetics *= "\n\t# load reaction kinetic constants" *
              "\n\tkcat = data_dictionary[\"kcat_signaling\"]  # in order of rnx" *
              "\n\n\t# write reaction rate equations" *
              "\n\trnx_rate_vector = np.zeros($totalRnxNo)"
  for (index, rnx) in enumerate(all_rnx_list)  # go thru every rnx
    tmp_rnx_rate = "kcat[$(index - 1)]"
    if isdefined(rnx, :catalysts)  # enzyme
      for token in rnx.catalysts
        tmp_rnx_rate *= "*X[$(all_species_dict[token.oriBioName] - 1)]"
      end
    end
    if isdefined(rnx,:reactants)  # substrate
      for token in rnx.reactants
        tmp_rnx_rate *= "*X[$(all_species_dict[token.oriBioName] - 1)]/(MonodK[\"MonodK~rnx$(index)~$(token.oriBioName)\"]+X[$(all_species_dict[token.oriBioName] - 1)])"
        push!(MonodAffinityConstant_String_Array, "MonodK~rnx$(index)~$(token.oriBioName)")
      end
    end
    kinetics *= "\n\trnx_rate_vector[$(index - 1)] = "*tmp_rnx_rate
  end

  # for TXTL
  W_string_array = Array{String,1}()  # W: weight of protein action on transcription
  disassociation_const_string_array = Array{String,1}()  # disassociation constant in tranfer function
  kinetics *= "\n\n\t# generate kinetics & control terms for TXTL" *
              "\n\t# load data from data dictionary" *
              "\n\tW_value_dict = data_dictionary[\"W_value_dict\"]  #String-->Float" *
              "\n\tcoop = data_dictionary[\"cooperativity\"]" *
              "\n\tdisassociation_constant_dict = data_dictionary[\"transferFunctionDisassociationConstantDict\"]" *
              "\n\tbackground_control_term = data_dictionary[\"backgroundGeneExpressionControlTermDict\"]" *
              "\n\tkcatTX = data_dictionary[\"kcatTranscription\"]" *
              "\n\tRNAPconc = data_dictionary[\"RNAPConcentration\"]" *
              "\n\tgeneConc = data_dictionary[\"avgGeneConcentration\"]" *
              "\n\tsaturationTX = data_dictionary[\"transcriptionSaturationConstant\"]" *
              "\n\tkcatTL = data_dictionary[\"kcatTranslation\"]" *
              "\n\tRIBOconc = data_dictionary[\"RIBOConcentration\"]" *
              "\n\tsaturationTL = data_dictionary[\"translationSaturationConstant\"]"
  totalTXTLNo = length(all_txtl_dict)
  kinetics *= "\n\n\t# write control terms, transcription rate equations, and translation rate equations" *
             "\n\tTX_control_term = {}  # control term: string -> float" *
             "\n\tTX_rate_vector = {}  # transcription rate: string -> float" *
             "\n\tTL_rate_vector = {}  # translation rate: string -> float"
  for (key, txtl) in all_txtl_dict # go thru every txtl
    tmp_protein_string = replace(key, sys2user["MRNA"] => sys2user["PROTEIN"], count=1)
    kinetics *= "\n\t# $key and $tmp_protein_string"
    up_factors_array = Array{String,1}()  # collection of upregulation factors name: targetedmRNA_factor(s)
    if !isempty(txtl.activationProtein)  # upregulation
      for token_array in txtl.activationProtein
        tmp_up = ""
        tmp_up_name = key*"_"
        for token in token_array
          tmp_W = "W~$(token.oriBioName)~$(key)"
          tmp_up *= "W_value_dict[\"$(tmp_W)\"] * X[$(all_species_dict[token.oriBioName] - 1)]**coop / ((disassociation_constant_dict[\"KD~$(token.oriBioName)\"])**coop+X[$(all_species_dict[token.oriBioName] - 1)]**coop)*"
          push!(W_string_array, tmp_W)
          push!(disassociation_const_string_array, "KD~$(token.oriBioName)")
          tmp_up_name *= token.oriBioName*"_"
        end
        tmp_up = chop(tmp_up)
        tmp_up_name = chop(tmp_up_name)
        push!(up_factors_array, tmp_up_name)
        kinetics *= "\n\t$(tmp_up_name) = $(tmp_up)"
      end
    end
    down_factors_array = Array{String,1}()  # collection of downregulation factors name: targetedmRNA_factor(s)
    if !isempty(txtl.inhibitionProtein)  # downregulation
      for token_array in txtl.inhibitionProtein
        tmp_down = ""
        tmp_down_name = key*"_"
        for token in token_array
          tmp_W = "W~$(token.oriBioName)~$(key)"
          tmp_down *= "W_value_dict[\"$(tmp_W)\"] * X[$(all_species_dict[token.oriBioName] - 1)]**coop/((disassociation_constant_dict[\"KD~$(token.oriBioName)\"])**coop+X[$(all_species_dict[token.oriBioName] - 1)]**coop)*"
          push!(W_string_array, tmp_W)
          push!(disassociation_const_string_array, "KD~$(token.oriBioName)")
          tmp_down_name *= token.oriBioName*"_"
        end
        tmp_down = chop(tmp_down)
        tmp_down_name = chop(tmp_down_name)
        push!(down_factors_array, tmp_down_name)
        kinetics *= "\n\t$(tmp_down_name) = $(tmp_down)"
      end
    end
    # combine to form up and down, then merge up and down to form control term
    if !isempty(up_factors_array) && !isempty(down_factors_array)
      up_action = ""
      for up in up_factors_array
        up_action *= up*" +"
      end
      kinetics *= "\n\tup_action_on_$key = $(chop(up_action))"
      down_action = ""
      for down in down_factors_array
        down_action *= down*" +"
      end
      kinetics *= "\n\tdown_action_on_$key = $(chop(down_action))" *
                  "\n\tTX_control_term[\"$key\"] = (background_control_term[\"$key\"] + up_action_on_$key)/(1 + background_control_term[\"$key\"] + up_action_on_$key + down_action_on_$key)"
    elseif !isempty(up_factors_array)
      up_action = ""
      for up in up_factors_array
        up_action *= up*" +"
      end
      kinetics *= "\n\tup_action_on_$key = $(chop(up_action))" *
                  "\n\tTX_control_term[\"$key\"] = (background_control_term[\"$key\"] + up_action_on_$key)/(1 + background_control_term[\"$key\"] + up_action_on_$key)"
    elseif !isempty(down_factors_array)
      down_action = ""
      for down in down_factors_array
        down_action *= down*" +"
      end
      kinetics *= "\n\tdown_action_on_$key = $(chop(down_action))" *
                  "\n\tTX_control_term[\"$key\"] = background_control_term[\"$key\"]/(1 + background_control_term[\"$key\"] + down_action_on_$key)"
    else
      # seems impossible, a TXTL dict is constructed b/c something acts on the key
    end
    kinetics *= "\n\tTX_rate_vector[\"$key\"] = TX_control_term[\"$key\"] * kcatTX * RNAPconc*geneConc/(saturationTX+geneConc)" *
                "\n\tTL_rate_vector[\"$tmp_protein_string\"] = kcatTL * RIBOconc*X[$(all_species_dict[key] - 1)]/(saturationTL+X[$(all_species_dict[key] - 1)])"
  end

  kinetics *= "\n\n\treturn rnx_rate_vector, TX_rate_vector, TL_rate_vector"

  return (kinetics, MonodAffinityConstant_String_Array, W_string_array, disassociation_const_string_array)
end

function py2_build_data_dictionary_buffer(host_type::AbstractString, all_species_array::Array,
  all_species2index_dict::Dict,
  rnx_species_array::Array, all_rnx_list::Array, all_txtl_dict::Dict,
  Monod_affinity_constant_array::Array, W_string_array::Array,
  disassociation_const_array::Array, mRNA_species_array::Array, protein_species_array::Array)
  # data dictionary
  buffer = "from __future__ import division" *
    "\nimport numpy as np" *
    "\nimport math"
  buffer *= "\n\ndef generate_model_parameters_dictionary():"
  # stoichiometry
  buffer *= "\n\t# load stoichiometry" *
    "\n\twith open(\"./stoichiometry.dat\", \"r\") as file:" *
    "\n\t\tmatrix = [[float(s) for s in line.strip().split()] for line in file]" *
    "\n\t\tprint(np.shape(matrix))" *
    "\n\tstoichiometric_matrix = np.matrix(matrix)" *
    "\n\tprint(np.shape(matrix))"

  # initial condition
  No_species = length(all_species_array)
  # initial concentration
  buffer *= "\n\n\t# initial species concentration" *
    "\n\tinitial_condition = [\n\t\t"
  buffer *= join(["0.001,  # $(i-1)  $(all_species_array[i])" for i = 1:(No_species - 1)], "\n\t\t")
  buffer *= "\n\t\t0.001   # $(No_species-1)  $(all_species_array[No_species])"
  buffer *= "\n\t]"

  # all species dictionary
  buffer *= "\n\n\t# all species dictionary"
  buffer *= "\n\tall_species_dict = {}  # String --> Int"
  #all_species_dict = Dict{String, Int64}()  # create all species dict: string --> int64
  for (st, index) in all_species2index_dict
    buffer *= "\n\tall_species_dict[\"$st\"] = $(index - 1)"
    #all_species_dict[st] = index
  end
  # all species dictionary_reversed
  buffer *= "\n\n\t# all species dictionary_reversed"
  buffer *= "\n\tall_species_reversed_dict = {}  # Int --> String"
  buffer *= "\n\tfor key, val in all_species_dict.items():"
  buffer *= "\n\t\tall_species_reversed_dict[val] = key"
  # all rnx species array
  buffer *= "\n\n\t# all rnx species array"
  buffer *= "\n\trnx_species_array = ["
  for st in rnx_species_array
    buffer *= "\n\t\t\"$st\","
  end
  buffer = chop(buffer)*"\n\t]"
  # all mRNA species array
  buffer *= "\n\n\t# all mRNA species array"
  buffer *= "\n\tmRNA_species_array = ["
  for st in mRNA_species_array
    buffer *= "\n\t\t\"$st\","
  end
  buffer = chop(buffer)*"\n\t]"
  # all protein species array
  buffer *= "\n\n\t# all protein species array"
  buffer *= "\n\tprotein_species_array = ["
  for st in protein_species_array
    buffer *= "\n\t\t\"$st\","
  end
  buffer = chop(buffer)*"\n\t]"

  # signaling reaction kinetic constants, reaction name: reactant(s)<rnx type:catalyst(s)>product(s)
  buffer *= "\n\n\t# signaling reaction kinetic constants, reaction name: reactant(s)<rnx type:catalyst(s)>product(s)"
  buffer *= "\n\tkcat_signaling = 1e-3*np.ones($(length(all_rnx_list)))  # kcat[#reaction]: reaction name"
  for (id, rnx) in enumerate(all_rnx_list)
    buffer*= "\n\tkcat_signaling[$(id - 1)] = 1.0  # kcat: $(rnx.rnxName)"
  end
  # Monod affinity constant in signaling reaction, nomenclature: MonodK_#reaction_Substrate
  buffer *= "\n\n\t# Monod affinity constant in signaling reaction, nomenclature: MonodK_#reaction_Substrate"
  buffer *= "\n\tMonodAffinityConstantDict = {}"
  for MK in Monod_affinity_constant_array
    buffer *= "\n\tMonodAffinityConstantDict[\"$MK\"] = 1.0"
  end

  # W_value in transcription control term, nomenclature: W_targetmRNA_actor
  buffer *= "\n\n\n\t# W_value in transcription control term, nomenclature: W_targetmRNA_actor"
  buffer *= "\n\tW_value_dict = {}"
  for WS in W_string_array
    buffer *= "\n\tW_value_dict[\"$WS\"] = 1.0"
  end
  # disassociation constants in transfer function, nomenclature: KD_speciesName
  buffer *= "\n\n\t# disassociation constants in transfer function, nomenclature: KD_speciesName"
  buffer *= "\n\ttransferFunctionDisassociationConstantDict = {}"

  for DC in disassociation_const_array
    buffer *= "\n\ttransferFunctionDisassociationConstantDict[\"$DC\"] = 1.0"
  end
  # transcription specific correction factor
  buffer *= "\n\n\t# transcription specific correction factor"
  buffer *= "\n\ttranscriptionSpecificCorrectionFactorDict = {}"
  for mrna in mRNA_species_array
    buffer *= "\n\ttranscriptionSpecificCorrectionFactorDict[\"$mrna\"] = 1.0  # $mrna"
  end
  # background gene expression control term
  buffer *= "\n\n\t# background gene expression control term"
  buffer *= "\n\tbackgroundGeneExpressionControlTermDict = {}"
  for mrna in mRNA_species_array
    buffer *= "\n\tbackgroundGeneExpressionControlTermDict[\"$mrna\"] = 0.1  # $mrna"
  end
  # translation specific correction factor
  buffer *= "\n\n\t# translation specific correction factor"
  buffer *= "\n\ttranslationSpecificCorrectionFactor = {}"
  for pro in protein_species_array
    buffer *= "\n\ttranslationSpecificCorrectionFactor[\"$pro\"] = 1.0  # $pro"
  end
  # cooperativity number in transfer function
  buffer *= "\n\n\t# cooperativity number in transfer function"
  buffer *= "\n\tcooperativity = 1"
  #buffer *= "\n\n\tbackground_mRNA_synthesis_rate_vector = 0.01*ones($length_TXTL)"
  # load txtl constants buffer
  buffer *= "\n\n \n"
  if host_type == "bacteria"
    buffer *= replace(replace(py2_include_constants_from_literature(
      joinpath(Base.@__DIR__, "txtl_constants_ecoli.jl"),"\n\t"),
      ")^3" => ")**3"), "log(" => "math.log(")
  else
    buffer *= replace(replace(py2_include_constants_from_literature(
      joinpath(Base.@__DIR__, "txtl_constants_hl60.jl"),"\n\t"),
      ")^3" => ")**3"), "log(" => "math.log(")
  end

  #---------------------------------
  # put all stuff in a Dictionary
  buffer *= "\n\n\t####################################"
  buffer *= "\n\t# put all stuff in a Dictionary"
  buffer *= "\n\tdataDictionary = {}  # string -> any"
  buffer *= "\n\tdataDictionary[\"stoichiometric_matrix\"] = stoichiometric_matrix"
  buffer *= "\n\tdataDictionary[\"all_species_dict\"] = all_species_dict"
  buffer *= "\n\tdataDictionary[\"all_species_reversed_dict\"] = all_species_reversed_dict"
  buffer *= "\n\tdataDictionary[\"rnx_species_array\"] = rnx_species_array"
  buffer *= "\n\tdataDictionary[\"mRNA_species_array\"] = mRNA_species_array"
  buffer *= "\n\tdataDictionary[\"protein_species_array\"] = protein_species_array"
  buffer *= "\n\tdataDictionary[\"initial_condition\"] = initial_condition"

  buffer *= "\n\tdataDictionary[\"kcat_signaling\"] = kcat_signaling"
  buffer *= "\n\tdataDictionary[\"MonodAffinityConstantDict\"] = MonodAffinityConstantDict"
  buffer *= "\n\tdataDictionary[\"W_value_dict\"] = W_value_dict"
  buffer *= "\n\tdataDictionary[\"transferFunctionDisassociationConstantDict\"] = transferFunctionDisassociationConstantDict"
  buffer *= "\n\tdataDictionary[\"transcriptionSpecificCorrectionFactorDict\"] = transcriptionSpecificCorrectionFactorDict"
  buffer *= "\n\tdataDictionary[\"backgroundGeneExpressionControlTermDict\"] = backgroundGeneExpressionControlTermDict"
  buffer *= "\n\tdataDictionary[\"translationSpecificCorrectionFactor\"] = translationSpecificCorrectionFactor"
  buffer *= "\n\tdataDictionary[\"cooperativity\"] = cooperativity"
  buffer *= "\n\tdataDictionary[\"RNAPConcentration\"] = rnapII_concentration"
  buffer *= "\n\tdataDictionary[\"RIBOConcentration\"] = ribosome_concentration"
  buffer *= "\n\tdataDictionary[\"degradationConstantmRNA\"] = degradation_constant_mRNA"
  buffer *= "\n\tdataDictionary[\"degradationConstantProtein\"] = degradation_constant_protein"
  buffer *= "\n\tdataDictionary[\"kcatTranscription\"] = kcat_transcription"
  buffer *= "\n\tdataDictionary[\"kcatTranslation\"] = kcat_translation"
  buffer *= "\n\tdataDictionary[\"avgGeneConcentration\"] = avg_gene_concentration"
  buffer *= "\n\tdataDictionary[\"transcriptionSaturationConstant\"] = saturation_transcription"
  buffer *= "\n\tdataDictionary[\"translationSaturationConstant\"] = saturation_translation"
  buffer *= "\n\tdataDictionary[\"specificGrowthRate\"] = maximum_specific_growth_rate - death_rate_constant"

  buffer *= "\n\n\treturn dataDictionary"

  return buffer
end

function py2_build_simulation_buffer(NoExtracellularSpecies::Int64)
  # buffer = build_copyright_header_buffer()
  buffer = "from __future__ import division" *
    "\nfrom Kinetics import calculate_kinetics" *
    "\nimport numpy as np" *
  "\n# set up ODE, get the derivatives
def Balances(y, t, dataDictionary):
  # correct for negatives
  y = [x if x > 0 else 0 for x in y]

  # load data from data dictionary
  rnx_species = dataDictionary[\"rnx_species_array\"]
  mRNA_species = dataDictionary[\"mRNA_species_array\"]
  protein_species = dataDictionary[\"protein_species_array\"]
  stoichiometry = dataDictionary[\"stoichiometric_matrix\"]
  all_species_dict = dataDictionary[\"all_species_dict\"]  # String --> number
  specific_growth_rate = dataDictionary[\"specificGrowthRate\"]
  transcription_correction = dataDictionary[\"transcriptionSpecificCorrectionFactorDict\"]
  translation_correction = dataDictionary[\"translationSpecificCorrectionFactor\"]
  degradation_constant_mRNA = dataDictionary[\"degradationConstantmRNA\"]
  degradation_constant_protein = dataDictionary[\"degradationConstantProtein\"]

  # kinetics
  rnx_rate, transcription_rate, translation_rate = calculate_kinetics(y, dataDictionary)

  # calculate the derivatives
  dydt = np.zeros(len(all_species_dict))

  # signaling reaction network
  X2 = np.zeros(len(rnx_species)-$(NoExtracellularSpecies))  # species inside the cell
  for id, st in enumerate(rnx_species[$(NoExtracellularSpecies):]):
    # get X2 value from y
    X2[id] = y[all_species_dict[st]]

  dX1 = np.dot(stoichiometry[0:$(NoExtracellularSpecies), :], rnx_rate)  # extracellular metabolites
  dX2 = np.dot(stoichiometry[$(NoExtracellularSpecies): , :], rnx_rate) - specific_growth_rate * X2  # INTRA
  # print(np.shape(dX1), np.shape(dX2))
  dX = np.concatenate((dX1, dX2), axis=1)
  # print(np.shape(dX))
  for id, st in enumerate(rnx_species):
    # return derivatives back to dydt
    dydt[all_species_dict[st]] = dX[0, id]

  dydt[0] = specific_growth_rate * y[0]  # for Biomass

  # TXTL network
  for i in range(len(mRNA_species)):
  # seems easy to add additional terms to account for the signaling effect
    mRNA_id = all_species_dict[mRNA_species[i]]
    p_id = all_species_dict[protein_species[i]]
    dydt[mRNA_id] = transcription_rate[mRNA_species[i]] - (specific_growth_rate + transcription_correction[mRNA_species[i]]*degradation_constant_mRNA) * y[mRNA_id]
    dydt[p_id] = translation_rate[protein_species[i]] - (specific_growth_rate + translation_correction[protein_species[i]]*degradation_constant_protein) * y[p_id]

  return dydt"

  return buffer
end

function py2_build_solveODEBalances_buffer(all_species_array::Array,
  all_species_dict::Dict{String, Int}, mRNA_species::Array, protein_species::Array)

  buffer = "from __future__ import division" *
    "\nimport scipy.integrate as integrate" *
    "\nimport matplotlib.pyplot as plt" *
    "\nimport numpy as np" *
    "\nimport math" *
    "\nfrom DataDictionary import generate_model_parameters_dictionary" *
    "\nfrom Kinetics import calculate_kinetics" *
    "\nfrom Balances import Balances" *
    "\n" *
    "\ndef SolveBalances(TStart, TStop, TStep, dataDictionary):" *
    "\n\ttimeNum = int((TStop - TStart)/TStep)" *
    "\n\tt = np.linspace(TStart, TStop, timeNum)"

  # data dataDictionary
  buffer *= "\n\tall_species_reversed_dict = dataDictionary[\"all_species_reversed_dict\"]"
  buffer *= "\n\tinitial_condition = dataDictionary[\"initial_condition\"]"
  buffer *= "\n\ty0 = dataDictionary[\"initial_condition\"]"
  # run simulation
  buffer *= "\n\t# Calling the solver"
  buffer *= "\n\ty = integrate.odeint(Balances, y0, t, args=(dataDictionary,))"
  # plotting
  buffer *= "\n\n\t# plotting results" *
    "\n\tcol = len(y0)" *
    "\n\tsubplt_col = math.ceil(col/4)" *
    "\n\tfor i in range(col):" *
    "\n\t\tplt.subplot(4, subplt_col, i+1)" *
    "\n\t\tplt.plot(t, y[:, i])" *
    "\n\t\tplt.title(all_species_reversed_dict[i])" *
    "\n\tplt.show()"
  buffer *= "\n\treturn t, y"

  # add a test
  buffer *= "\n\n\n"
  buffer *= "## uncomment to test"
  buffer *= "\n# dataDictionary = generate_model_parameters_dictionary()"
  buffer *= "\n# SolveBalances(0, 10, 0.1, dataDictionary)"

  return buffer
end

# FBA data dictionary generation
function py2_generate_FBA_data_dictionary(all_rnx_list::Array,
  rnx_species_array::Array, extra_species_num::Int)
  secrete_id = Set{Int64}()  # For initialize the coefficient array
  # data dictionary
  buffer = "from __future__ import division" *
    "\nimport numpy as np" *
    "\nimport math"
  buffer *= "\n\ndef DataDictionary(time_start,time_stop,time_step):"
  # stoichiometry
  buffer *= "\n\n\t# load stoichiometry" *
    "\n\twith open(\"./stoichiometry.dat\", \"r\") as file:" *
    "\n\t\tmatrix = [[float(s) for s in line.strip().split()] for line in file]" *
    # "\n\t\tprint(np.shape(matrix))" *
    "\n\tstoichiometric_matrix = np.matrix(matrix)"
    # "\n\tprint(np.shape(matrix))"

  # Setup default flux bounds array
  buffer *= "\n\t# Setup default flux bounds array"
  buffer *= "\n\tdefault_flux_bounds_array = ("
  for (id, rnx) in enumerate(all_rnx_list)
    buffer *= "\n\t\t(0, 100.0), # $(id) $(rnx.rnxName)"
  end
  buffer *= "\n\t)"
  # Setup default species bounds array
  buffer *= "\n\n\t# Setup default species bounds array"
  buffer *= "\n\tspecies_bounds_array = ["
  for (id, sp) in enumerate(rnx_species_array)
    if id > extra_species_num
      buffer *= "\n\t\t[0.0, 0.0], # $(id) $(sp)"
    else
      buffer *= "\n\t\t[-1.0, 1.0], # $(id) $(sp)"
    end
  end
  buffer *= "\n\t]"
  # Setup the objective coefficient array
  buffer *= "\n\n\t# Setup the objective coefficient array"
  buffer *= "\n\tobjective_coefficient_array = ["
  for (id, rnx) in enumerate(all_rnx_list)
    buffer *= "\n\t\t0.0, # $(id) $(rnx.rnxName)"
    if rnx.rnxType == "secrete"
        push!(secrete_id, (id-1))
    end
  end
  buffer *= "\n\t]"
  # List of reation strings - used to write flux report
  buffer *= "\n\n\t# List of reation strings - used to write flux report"
  buffer *= "\n\tlist_of_reaction_strings = ["
  for rnx in all_rnx_list
    buffer *= "\n\t\t\"$(rnx.rnxName)\", "
  end
  buffer *= "\n\t]"
  # List of metabolite strings - used to write flux report
  buffer *= "\n\n\t# List of metabolite strings - used to write flux report"
  buffer *= "\n\tlist_of_metabolite_symbols = ["
  for sp in rnx_species_array
    buffer *= "\n\t\t\"$(sp)\","
  end
  buffer *= "\n\t]"
  # return block -
  buffer *= "\n\n"
  buffer *= "\t# =============================== DO NOT EDIT BELOW THIS LINE ============================== #\n"
  buffer *= "\tdata_dictionary = {}\n"
  buffer *= "\tdata_dictionary[\"stoichiometric_matrix\"] = stoichiometric_matrix\n"
  buffer *= "\tdata_dictionary[\"objective_coefficient_array\"] = objective_coefficient_array\n"
  buffer *= "\tdata_dictionary[\"default_flux_bounds_array\"] = default_flux_bounds_array;\n"
  buffer *= "\tdata_dictionary[\"species_bounds_array\"] = species_bounds_array\n"
  buffer *= "\tdata_dictionary[\"list_of_reaction_strings\"] = list_of_reaction_strings\n"
  buffer *= "\tdata_dictionary[\"list_of_metabolite_symbols\"] = list_of_metabolite_symbols\n"
  buffer *= "\tdata_dictionary[\"extra_species_num\"] = $(extra_species_num)\n"
  buffer *= "\t# =============================== DO NOT EDIT ABOVE THIS LINE ============================== #\n"
  buffer *= "\treturn data_dictionary\n"

  # maximize_product_dictionary
  buffer *= "\n\n\n# maximize_product_dictionary"
  buffer *= "\ndef maximize_product_dictionary(time_start,time_stop,time_step):"
  # load the data dictionary -
  buffer *= "\n\t# load the data dictionary -"
  buffer *= "\n\tdata_dictionary = DataDictionary(time_start,time_stop,time_step)"
  # Modify the data dictionary -
  buffer *= "\n\t# Modify the data dictionary -"
  buffer *= "\n\tobjective_coefficient_array = data_dictionary[\"objective_coefficient_array\"]"
  secrete_id = sort(collect(secrete_id))
  for id in secrete_id
      buffer *= "\n\tobjective_coefficient_array[$(id)] = -1"
  end
  buffer *= "\n\t# return -"
  buffer *= "\n\treturn data_dictionary"

  return buffer
end


"""
    SynchronousExecution

Functor object for using synchronous execution in an L-shaped algorithm (assuming multiple Julia cores are available). Create by supplying a [`Synchronous`](@ref) object through `execution` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct SynchronousExecution{H <: AbstractFeasibilityHandler,
                            T <: AbstractFloat,
                            A <: AbstractVector} <: AbstractLShapedExecution
    subworkers::Vector{SubWorker{H,T}}
    decisions::Vector{DecisionChannel}
    subobjectives::A
    model_objectives::A
    metadata::Vector{MetaData}
    cutqueue::CutQueue{T}

    function SynchronousExecution(structure::VerticalStructure{2, 1, <:Tuple{DistributedScenarioProblems}},
                                  ::Type{F}, ::Type{T}, ::Type{A}) where {F <: AbstractFeasibility,
                                                                          T <: AbstractFloat,
                                                                          A <: AbstractVector}
        H = HandlerType(F)
        return new{H,T,A}(Vector{SubWorker{H,T}}(undef, nworkers()),
                          scenarioproblems(structure).decisions,
                          A(),
                          A(),
                          Vector{MetaData}(undef, nworkers()),
                          RemoteChannel(() -> Channel{QCut{T}}(4 * nworkers() * num_scenarios(structure))))
    end
end

function initialize_subproblems!(execution::SynchronousExecution,
                                 scenarioproblems::DistributedScenarioProblems)
    load_subproblems!(execution.subworkers, scenarioproblems, execution.decisions)
    return nothing
end

function finish_initilization!(lshaped::AbstractLShaped, execution::SynchronousExecution)
    append!(execution.subobjectives, fill(1e10, num_thetas(lshaped)))
    append!(execution.model_objectives, fill(-1e10, num_thetas(lshaped)))
    for w in workers()
        execution.metadata[w-1] = RemoteChannel(() -> MetaChannel(), w)
    end
    return lshaped
end

function restore_subproblems!(::AbstractLShaped, execution::SynchronousExecution)
    restore_subproblems!(execution.subworkers)
    return nothing
end

function resolve_subproblems!(lshaped::AbstractLShaped, execution::SynchronousExecution{H,T}) where {H <: AbstractFeasibilityHandler, T <: AbstractFloat}
    # Update metadata
    for w in workers()
        put!(execution.metadata[w-1], timestamp(lshaped), :gap, gap(lshaped))
    end
    @sync begin
        for (i,w) in enumerate(workers())
            worker_aggregator = remote_aggregator(lshaped.aggregation, scenarioproblems(lshaped.structure), w)
            @async remotecall_fetch(resolve_subproblems!,
                                    w,
                                    execution.subworkers[w-1],
                                    execution.decisions[w-1],
                                    lshaped.x,
                                    execution.cutqueue,
                                    worker_aggregator,
                                    timestamp(lshaped),
                                    execution.metadata[w-1])
        end
    end
    # Assume no cuts are added
    added = false
    # Collect incoming cuts
    while isready(execution.cutqueue)
        _, cut::SparseHyperPlane{T} = take!(execution.cutqueue)
        added |= add_cut!(lshaped, cut)
    end
    # Return current objective value and cut_added flag
    return current_objective_value(lshaped), added
end

# API
# ------------------------------------------------------------
function (execution::Synchronous)(structure::VerticalStructure{2, 1, <:Tuple{DistributedScenarioProblems}},
                                  ::Type{F}, ::Type{T}, ::Type{A}) where {F <: AbstractFeasibility,
                                                                          T <: AbstractFloat,
                                                                          A <: AbstractVector}
    return SynchronousExecution(structure, F, T, A)
end

function str(::Synchronous)
    return "Synchronous "
end


using IRTools
using IRTools: IR, slots!, argnames!, varargs!, xcall, SSAValue
using ..Zygote: Argument, argmap, ignored, isexpr

function record!(ir::IR)
  ir = map(ex -> argmap(a -> Argument(a.n+2), ex), ir)
  grads = Dict()
  grad(x::Union{SSAValue,Argument}) = grads[x]
  grad(x) = nothing
  ks = keys(ir)
  for n = 1:length(ir.args)
    grads[Argument(n+2)] = pushfirst!(ir, xcall(Base, :getindex, Argument(2)))
  end
  for k in ks
    ex = ir[k].expr
    if isexpr(ex, :call) && !ignored(ir, ex)
      Δ = insert!(ir, k, xcall(Base, :tuple, grad.(ex.args)...))
      ir[k] = xcall(Forward, :_forward, Δ, ex.args...)
    end
  end
  return ir
end

@generated function _forward(Δ::Tuple, f, args...)
  T = Tuple{f,args...}
  meta = IRTools.meta(T)
  forw = record(IR(meta))
  pushfirst!(forw.args, Any, Any)
  argnames!(meta, Symbol("#self#"), :Δ, :f, :args)
  forw = varargs!(meta, forw, 3)
  forw = slots!(forw)
  return IRTools.update!(meta, forw)
end


module Obfs

using Dates
using Sockets

import Sockets.accept
import Base.read
import Base.write
import Base.readbytes!

using Crypto.Chacha20: Chacha20Encrypt, Chacha20Decrypt
using Crypto.HKDF: hkdf
using ..Common: ioCopy

import ..Common.read
import ..Common.write

const INFO = Array{UInt8}("obfs-tls")
const TIMEOUT = Millisecond(1000)
const ZERO = Millisecond(0)
const MaxSize = 0x5000

mutable struct TLSConfig
    host::Union{String, Sockets.IPAddr}
    port::UInt16
    key::Array{UInt8}
end

mutable struct TLSConnection
    conn::TCPSocket
    cache::IOStream
end

@inline function eof(io::TLSConnection)
    eof(io.conn)
end

@inline function isopen(io::TLSConnection)
    isopen(io.conn)
end

@inline function readbytes!(io::TLSConnection, buff::Array{UInt8}, bytes::Integer)
    readbytes!(io.conn, buff, nbytes)
end

@inline function write(io::TLSConnection, buff::Array{UInt8})
    write(io.conn, buff)
end

@inline function read(io::TLSConnection, buff::Array{UInt8})
    try
        eof(io)
    catch err
        return nothing, err
    end

    isopen(io) ? try
        readbytes(io, buff, 0x05)
    catch err
        return nothing, err
    end : return nothing, Error("Connection Closed")

    nbytes = UInt16(buff[4]) * 0x0100 + buff[5]

    isopen(io) ? try
        readbytes!(io, buff, nbytes)
    catch err
        return nothing, err
    end : return nothing, Error("Connection Closed")

    return nbytes, nothing
end

@inline function read(io::TLSConnection, buff::Array{UInt8}, nbytes::Integer)
    while true
        if nbytes <= bytesavailable(io.cache)
            readbytes!(io.cache, buff, nbytes)
            break
        end

        nread, err = read(io, buff)
        if err != nothing
            return nothing, err
        end

        write(io.cache, unsafe_wrap(Array{UInt8}, pointer(buff), nread))
    end

    return nbytes, nothing
end

@inline function write(io::TLSConnection, buff::Array{UInt8}, nbytes::Integer)
    isopen(io) ? try
        write(io, [[0x17, 0x03, 0x03, ]; UInt8[nbytes >> 8, nbytes & 0xff, ]])
        write(io, unsafe_wrap(Array{UInt8}, pointer(buff), nbytes))
    catch err
        return err
    end : return Error("Connection Closed")

    return nothing
end

mutable struct RandomSet
    lock::ReentrantLock
    set::Set{Array{UInt8}}
end

struct Error <: Exception
    msg::AbstractString
end

function add(random::Array{UInt8}, set::RandomSet)
    lock(random.lock)
    union!(set, [random, ])
    unlock(random.lock)
end

function clear(set::RandomSet)
    lock(random.lock)
    set.set = Set{Array{UInt8}}()
    unlock(random.lock)
end

function hasRandom(random::Array{UInt8}, set::RandomSet)::Bool
    lock(random.lock)
    x = random in set.set
    unlock(random.lock)

    x
end

function genRandom(config::TLSConfig)::Array{UInt8}
    random = Array{UInt8}(undef, 32)
    random[24:32] = rand(UInt8, 9)
    kv, err = hkdf("SHA2-256", random[24:32], config.key, INFO, 44)
    Chacha20Encrypt(random, kv[1:32], kv[33:44], Array{UInt8}(string(Dates.now())))

    random
end

function checkRandom(random::Array{UInt8}, config::TLSConfig)::Bool
    x = Array{UInt8}(undef, 23)
    kv, err = hkdf("SHA2-256", random[24:32], config.key, INFO, 44)
    Chacha20Decrypt(x, kv[1:32], kv[33:44], random[1:23])

    now = try
        Dates.now() - parse(Dates.DateTime, String(x))
    catch err
        return false
    end

    if ZERO < now < TIMEOUT
        true
    else
        false
    end
end

function parseClientHello(buff::Array{UInt8})::Tuple{Union{Array{UInt8}, Nothing}, Union{Array{UInt8}, Nothing}, Union{Exception, Nothing}}
    while true
        pointer = 6;
        if buff[pointer] != 0x01
            break
        end

        pointer += 6;
        random = buff[pointer:pointer+32]

        pointer += 33;
        sessionId = buff[pointer:pointer+32]

        return random, sessionId, nothing
    end

    return nothing, nothing, Error("Not a Valid Client Hello Message")
end

function composeExtension()
    [0x00, ]
end

function composeClientHello(config::TLSConfig)::Array{UInt8}
    [[0x16, ]; [0x03, 0x03, ]; [0x00, 0x00];
        [
            [0x01, ];
            [0x00, 0x01, 0xfc];
            [0x03, 0x03];
            genRandom(config);
            [0x20, ];
            rand(UInt8, 0x20);
            [0x00, 0x1c];
            [0x2a, 0x2a, 0xc0, 0x2b, 0xc0, 0x2f, 0xc0, 0x2c, 0xc0, 0x30, 0xcc, 0xa9, 0xcc, 0xa8, 0xc0, 0x13, 0xc0, 0x14, 0x00, 0x9c, 0x00, 0x9d, 0x00, 0x2f, 0x00, 0x35, 0x00, 0x0a, ];
            [0x01, ];
            [0x00, ];
            [0x01, 0x97, ];
            composeExtension()
        ]
    ]
end

function composeServerHello(sessionId::Array{UInt8})::Array{UInt8}
    [[0x16, ]; [0x03, 0x03, ]; [];
        [
            [0x02, ];
            [0x00, 0x00, 0x4d, ];
            [0x03, 0x03];
            rand(UInt8, 0x20);
            [0x20, ];
            sessionId;
            [0xc0, 0x30, ];
            [0x00, ];
            [0x00, 0x05, ];
            [0xff, 0x01, 0x00, 0x01, 0x00, ]
        ]
    ]
end

function composeChangeCipher()
    [[0x14, ]; [0x03, 0x03, ]; [0x00, 0x01, ]; [0x01, ]]
end

function composeFinished()
    [[0x16, ]; [0x03, 0x03, ]; [0x00, 0x40, ]; rand(UInt8, 0x28); Array{UInt8}(undef, 0x18)]
end

function accept(server::Sockets.TCPServer, config::TLSConfig)::TCPSocket
    if @isdefined set

    else
        global set = RandomSet(ReentrantLock(), Set{Array{UInt8}}())
        @async while true
            sleep(3600)
            clear(set)
        end
    end

    buff = Array{UInt8}(undef, 1024)

    while true
        conn = try
            accept(server)
        catch err
            throw(err)
        end

        nbytes, err = read(conn, buff)
        if err != nothing
            close(conn)
            continue
        end

        random, sessionId, err = parseClientHello(buff)
        if err != nothing
            close(conn)
            continue
        end

        if checkRandom(random, config) && !hasRandom(random, set)
            add(random, set)
            read(conn, buff, 75)
            return conn
        else
            fake = try
                connect(config.host, config.port)
            catch err
            end

            err = write(fake, buff, nbytes)
            if err != nothing
            end

            @async ioCopy(conn, fake)
            @async ioCopy(fake, conn)
        end
    end
end

end # module

if false

using Shadowsocks
using Sockets

config = Shadowsocks.Obfs.TLSConfig(IPv4(13,35,15,164), 0x01bb, Array{UInt8}("julia"))
server = listen(2000)

Shadowsocks.Obfs.accept(server, config)

end


using NQCModels: Free, FrictionModels
using PyCall
using NQCBase
using CubeLDFAModel
using Test
using Unitful

include("test_cube.jl")

aseio = pyimport("ase.io")
ase_atoms = aseio.read("start.in")
atoms, R, cell =  NQCBase.convert_from_ase_atoms(ase_atoms)

model = LDFAModel(Free(), "test.cube", atoms, cell; friction_atoms=[1, 2])

@testset "friction!" begin

    F = zeros(6, 6)
    r = @view R[:,1]
    r .= 0
    FrictionModels.friction!(model, F, R)

    for _=1:10
        r .= 0
        r += rand() * cell.vectors[:,1]
        r += rand() * cell.vectors[:,2]
        r += rand() * cell.vectors[:,3]

        FrictionModels.friction!(model, F, R)
    end

    r = cell.vectors[:,1] + cell.vectors[:,2] + cell.vectors[:,3]
    FrictionModels.friction!(model, F, R)

    r = cell.vectors[:,1] + cell.vectors[:,2] + cell.vectors[:,3] + rand(3)
    FrictionModels.friction!(model, F, R)
end


using BlackBoxOptim: makefitness, NewFitness, SingleObjectiveFitness, SingleNumberFitness
using BlackBoxOptim: MultiObjectiveFitness, VectorFitness
using BlackBoxOptim: fitnessvalues, characteristics, simulations, focusversion, goalversion, evaluator

describe("Fitness types") do
  @repeat test("SingleObjectiveFitness") do
    fval = randn()
    f1 = makefitness(fval)
    @check typeof(f1) <: NewFitness
    @check typeof(f1) <: SingleObjectiveFitness
    @check typeof(f1) == SingleNumberFitness{Float64}

    @check fitnessvalues(f1) == [fval]
    @check characteristics(f1) == [fval]
    @check simulations(f1) == [fval]
    @check evaluator(f1) == nothing
    @check focusversion(f1) == 0
    @check goalversion(f1) == 0
  end

  @repeat test("MultiObjectiveFitness") do
    size = rand(2:117)
    vals = randn(size)
    f2 = makefitness(vals)
    @check typeof(f2) <: NewFitness
    @check typeof(f2) <: MultiObjectiveFitness
    @check typeof(f2) == VectorFitness{Float64}

    @check fitnessvalues(f2) == vals
    @check characteristics(f2) == vals
    @check simulations(f2) == vals
    @check evaluator(f2) == nothing
    @check focusversion(f2) == 0
    @check goalversion(f2) == 0
  end
end


# define a special case of exponential Kernel, because Claudia said so

####
# Population and connection constructors

# the taus are in the trace, now. No unittype or kernel object
# this is pretty much a wrapper around neural traces
struct PopulationStateExpKernel{N}  <: AbstractPopulationState
  label::Symbol
  n::Int64
  traces::NTuple{N,Trace}
end
nneurons(ps::PopulationStateExpKernel) = ps.n

function PopulationStateExpKernel(n::Int64,traces...;
    label::Union{String,Nothing}=nothing)
  label = something(label,rand_label())
  PopulationStateExpKernel(label,n,traces)
end
function reset!(ps::PopulationStateExpKernel)
  for tra in ps.traces
    reset!(tra)
  end
  return nothing
end

struct PopulationExpKernel{N,PS<:PopulationStateExpKernel,
    TC<:NTuple{N,Connection},
    TP<:NTuple{N,PopulationStateExpKernel},
    NL<:AbstractNonlinearity} <: AbstractPopulation
  state::PS
  connections::TC
  pre_states::TP
  input::Vector{Float64} # input only used for external currents
  nonlinearity::NL # nonlinearity if present
  spike_proposals::Vector{Float64} # memory allocation
end
function PopulationExpKernel(state::PopulationStateExpKernel,input::Vector{Float64},
    (conn_pre::Tuple{C,PS} where {C<:Connection,PS<:PopulationStateExpKernel})...;
      nonlinearity::AbstractNonlinearity=NLRelu())
  connections = Tuple(getindex.(conn_pre,1))
  typeassert.(connections,ConnectionExpKernel)
  pre_states = Tuple(getindex.(conn_pre,2))
  spike_proposals = fill(Inf,nneurons(state))
  return PopulationExpKernel(state,connections,pre_states,input,nonlinearity,spike_proposals)
end
# one population only!
function PopulationExpKernel(state::PopulationStateExpKernel,conn::Connection,input::Vector{Float64};
    nonlinearity::AbstractNonlinearity=NLRelu())
  return PopulationExpKernel(state,input,(conn,state);nonlinearity=nonlinearity)
end

struct ConnectionExpKernel{N,PL<:NTuple{N,PlasticityRule}} <: Connection
  weights::Matrix{Float64}
  pre_trace::Trace
  trace_proposal::Vector{Float64}
  plasticities::PL
end
function ConnectionExpKernel(weights::Matrix{Float64},pre_trace::Trace)
  plast = NTuple{0,NoPlasticity}()
  _,npre = size(weights)
  @assert npre == nneurons(pre_trace)
  tra_prop = Vector{Float64}(undef,npre)
  return ConnectionExpKernel(weights,pre_trace,tra_prop,plast)
end
function ConnectionExpKernel(weights::Matrix{Float64},pre_trace::Trace,
     (plasticity_rules::PL where PL<:PlasticityRule)...)
  _,npre = size(weights)
  @assert npre == nneurons(pre_trace)
  tra_prop = Vector{Float64}(undef,npre)
  return ConnectionExpKernel(weights,pre_trace,tra_prop,plasticity_rules)
end

@inline function trace_proposal!(t_now::Real,conn::ConnectionExpKernel)
  return trace_proposal!(conn.trace_proposal,t_now,conn.pre_trace)
end

abstract type Recorder end

struct RecurrentNetworkExpKernel{N,TP<:NTuple{N,AbstractPopulation},TR<:NTuple{N,Recorder}}
  populations::TP
  recorders::TR
end
# one population constructor
function RecurrentNetworkExpKernel(pop::AbstractPopulation,recorders...)
  if isempty(recorders)
    #recorders=NTuple{0,RecNothing}()
    recorders=(RecNothing(),)
    #@show typeof(recorders)
  end
  return RecurrentNetworkExpKernel((pop,),recorders)
end

# reset
function reset!(rn::RecurrentNetworkExpKernel)
  # reset internal variables
  for pop in rn.populations
    reset!(pop.state)
    fill!(pop.spike_proposals,Inf)
    # reset plasticity operators if present
    for conn in pop.connections
      for plast in conn.plasticities
        reset!(plast)
      end
    end
  end
  # reset recorders
  for rec in rn.recorders
    reset!(rec)
  end
  return nothing
end
@inline function npopulations(ntw::RecurrentNetworkExpKernel)
  return length(ntw.populations)
end

function burn_spike!(t_spike::Real,ps::PopulationStateExpKernel,idx_update::Integer)
  # take care of traces
  for tra in ps.traces
    propagate_for_dynamics!(t_spike,tra) # update full trace to t_spike
    update_for_dynamics!(tra,idx_update) # add pulse at index
  end
  return nothing
end

function compute_rate(t_now::Real,external_input::Real,
    pop::PopulationExpKernel, idxneu::Integer)
  ret = external_input
  for conn in pop.connections
    tra_tnow = trace_proposal!(t_now,conn)
    wij_all = view(conn.weights,idxneu,:)
    ret += dot(wij_all,tra_tnow)
  end
  return apply_nonlinearity(ret,pop.nonlinearity)
end

# Thinning algorith, e.g.  Laub,Taimre,Pollet 2015
function compute_next_spike(t_now::Real,pop::PopulationExpKernel,ineu::Integer;Tmax::Real=100.0)
  t_start = t_now
  t = t_now
  ext_inp = pop.input[ineu]
  freq(t) = compute_rate(t,ext_inp,pop,ineu)
  while (t-t_start)<Tmax
    (M::Float64) = freq(t)
    Δt =  -log(rand())/M # rand(Exponential())/M
    t = t+Δt
    u = rand()*M # random between 0 and M
    (u_up::Float64) = freq(t)
    if u <= u_up
      return t
    end
  end
  return Tmax + t_start
end

###
# recorders methods
struct RecNothing <: Recorder end

function reset!(::RecNothing)
  return nothing
end

# general signature: record_stuff!(rec,tfire,popfire,neufire,label_fire,ntw)
# @inline function record_stuff!(::RecNothing,::Real,
#     ::Integer,::Integer,::Symbol,
#     ::RecurrentNetworkExpKernel)
#   return nothing
# end
@inline function record_stuff!(::RecNothing,::Real,
    whatevs...)
  return nothing
end

struct RecFullTrain{N} <: Recorder
  nrec::Int64
  timesneurons::NTuple{N,Tuple{Vector{Float64},Vector{Int64}}}
  k_rec::Vector{Int64}
end
function RecFullTrain(nrec::Integer,(npops::Integer)=1)
  timesneurons=ntuple(_-> (fill(NaN,nrec),fill(-1,nrec)),npops)
  k_rec = fill(0,npops)
  RecFullTrain(nrec,timesneurons,k_rec)
end
function reset!(rec::RecFullTrain)
  fill!(rec.k_rec,0)
  for (spktim,spkneu) in rec.timesneurons
    fill!(spktim,NaN)
    fill!(spkneu,-1)
  end
  return nothing
end

# general signature: record_stuff!(rec,tfire,popfire,neufire,label_fire,ntw)

function record_stuff!(rec::RecFullTrain,tfire::Real,
    popfire::Integer,neufire::Integer,::Symbol,
    ::RecurrentNetworkExpKernel)
  k = rec.k_rec[popfire]+1
  spiketimes,spikeneurons = rec.timesneurons[popfire]
  if !checkbounds(Bool,spiketimes,k)
    @error "recorder full!"
    return nothing
  end
  @inbounds spiketimes[k] = tfire
  @inbounds spikeneurons[k] = neufire
  rec.k_rec[popfire] = k
  return nothing
end

function dynamics_step!(t_now::Real,ntw::RecurrentNetworkExpKernel)
  npops = npopulations(ntw)
  proposals_best = Vector{Float64}(undef,npops)
  neuron_best = Vector{Int64}(undef,npops)
  # for each postsynaptic network, compute spike proposals
  for (kn,pop) in enumerate(ntw.populations)
    # update proposed next spike for each postsynaptic neuron
    nneu = nneurons(pop)
    for ineu in 1:nneu
      pop.spike_proposals[ineu] = compute_next_spike(t_now,pop,ineu)
    end
    # best candidate and neuron that fired it for one input network
    (proposals_best[kn], neuron_best[kn]) = findmin(pop.spike_proposals)
  end
  # select next spike (best across all input_networks)
  (tfire,popfire) = findmin(proposals_best)
  neufire = neuron_best[popfire]
  psfire = ntw.populations[popfire].state
  label_fire = psfire.label
  # update stuff for that specific neuron/population state :
  burn_spike!(tfire,psfire,neufire)
  # apply plasticity rules ! Each rule in each connection in each population
  for pop in ntw.populations
    post = pop.state
    for (conn,pre) in zip(pop.connections,pop.pre_states)
      for plast in conn.plasticities
        plasticity_update!(tfire,label_fire,neufire,post,conn,pre,plast)
      end
    end
  end
  # now trigger recorders.
  # Recorder objects will take care of which population to target
  for rec in ntw.recorders
    record_stuff!(rec,tfire,popfire,neufire,label_fire,ntw)
  end
  # update t_now
  return tfire
end

function dynamics_step_singlepopulation!(t_now::Real,ntw::RecurrentNetworkExpKernel)
  # update proposed next spike for each postsynaptic neuron
  popfire = 1
  pop = ntw.populations[1]
  nneu = nneurons(pop)
  for ineu in 1:nneu
    pop.spike_proposals[ineu] = compute_next_spike(t_now,pop,ineu)
  end
  (tfire, neufire) = findmin(pop.spike_proposals)
  psfire = pop.state
  label_fire = psfire.label
  # update stuff for that specific neuron/population state :
  burn_spike!(tfire,psfire,neufire)
  # apply plasticity rules ! Each rule in each connection in each population
  post = psfire
  for (conn,pre) in zip(pop.connections,pop.pre_states)
    for plast in conn.plasticities
      plasticity_update!(tfire,label_fire,neufire,post,conn,pre,plast)
    end
  end
  # now trigger recorders.
  # Recorder objects will take care of which population to target
  for rec in ntw.recorders
    record_stuff!(rec,tfire,popfire,neufire,label_fire,ntw)
  end
  # update t_now
  return tfire
end


"""
Analytical reports generated during training, for debugging and
hyperparameters tuning.
"""
module Report

# Some imports for docstrings cross-references
import ..MCTS, ..Network, ..MemAnalysisParams

"""
    Report.Loss

Decomposition of the loss in a sum of terms (all have type `Float32`).

- `L` is the total loss: `L == Lp + Lv + Lreg + Linv`
- `Lp` is the policy cross-entropy loss term
- `Lv` is the average value mean square error
- `Lreg` is the L2 regularization loss term
- `Linv` is the loss term penalizing the average weight put by the
  network on invalid actions
"""
struct Loss
  L :: Float32
  Lp :: Float32
  Lv :: Float32
  Lreg :: Float32
  Linv :: Float32
end

"""
    Report.LearningStatus

Statistics about the performance of the neural network on a subset of the
memory buffer.

- `loss`: detailed loss on the samples, as an object of type
    [`Report.Loss`](@ref)
- `Hp`: average entropy of the ``π`` component of samples (MCTS policy);
    this quantity is independent of the network and therefore constant
    during a learning iteration
- `Hpnet`: average entropy of the network's prescribed policy on the samples
"""
struct LearningStatus
  loss :: Loss
  Hp :: Float32 # property of the memory, constant during a learning iteration
  Hpnet :: Float32
end

"""
    Report.Checkpoint

Report generated after a checkpoint evaluation.

- `batch_id`: number of batches after which the checkpoint was computed
- `status_after`: learning status at the checkpoint, as an object of type
   [`Report.LearningStatus`](@ref)
- `reward`: average reward collected by the contender network
- `redundancy`: ratio of duplicate positions encountered during the evaluation,
   not counting the initial position. If this number is too high, you may
   want to increase the move selection temperature.
- `nn_replaced`: true if the current best neural network was updated after
   the checkpoint
"""
struct Checkpoint
  batch_id :: Int
  status_after :: LearningStatus
  reward :: Float64
  redundancy :: Float64
  nn_replaced :: Bool
end

"""
    Report.Learning

Report generated at the end of the learning phase of an iteration.

- `time_convert`, `time_loss`, `time_train` and `time_eval` are the
    amounts of time (in seconds) spent at converting the samples,
    computing losses, performing gradient updates and evaluating checkpoints
    respectively
- `initial_status`: status before the learning phase, as an object of type
    [`Report.LearningStatus`](@ref)
- `losses`: loss value on each minibatch
- `checkpoints`: vector of [`Report.Checkpoint`](@ref) reports
- `nn_replaced`: true if the best neural network was replaced
"""
struct Learning
  time_convert :: Float64
  time_loss :: Float64
  time_train :: Float64
  time_eval :: Float64
  initial_status :: LearningStatus
  losses :: Vector{Float32}
  checkpoints :: Vector{Checkpoint}
  nn_replaced :: Bool
end

"""
    Report.Samples

Statistics about a set of samples, as collected during memory analysis.

- `num_samples`: total number of samples
- `num_boards`: number of distinct board positions
- `Wtot`: total weight of the samples
- `status`: [`Report.LearningStatus`](@ref) statistics of the current network
    on the samples
"""
struct Samples
  num_samples :: Int
  num_boards :: Int
  Wtot :: Float32
  status :: LearningStatus
end

"""
    Report.StageSamples

Statistics for the samples corresponding to a particular game stage,
as collected during memory analysis.

The samples whose statistics are collected in the
[`samples_stats`](@ref Report.Samples) field correspond to historical positions
where the number of remaining moves until the end of the game was in the range
defined by the `min_remaining_length` and `max_remaining_length` fields.
"""
struct StageSamples
  min_remaining_length :: Int
  max_remaining_length :: Int
  samples_stats :: Samples
end

"""
    Report.Memory

Report generated by the memory analysis phase of an iteration. It features
statistics for
  - the whole memory buffer (`all_samples::Report.Samples`)
  - the samples collected during the last self-play iteration
     (`latest_batch::Report.Samples`)
  - the subsets of the memory buffer corresponding to different game stages:
     (`per_game_stage::Vector{Report.StageSamples}`)

See [`MemAnalysisParams`](@ref).
"""
struct Memory
  latest_batch :: Samples
  all_samples :: Samples
  per_game_stage :: Vector{StageSamples}
end

"""
    Report.SelfPlay

Report generated after the self-play phase of an iteration.

- `samples_gen_speed`: average number of samples generated per second
- `average_exploration_depth`: see [`MCTS.average_exploration_depth`](@ref)
- `mcts_memory_footprint`: estimation of the maximal memory footprint of the
    MCTS tree during self-play, as computed by
    [`MCTS.approximate_memory_footprint`](@ref)
- `memory_size`: number of samples in the memory buffer at the end of the
    self-play phase
- `memory_num_distinct_boards`: number of distinct board positions in the
    memory buffer at the end of the self-play phase
"""
struct SelfPlay
  samples_gen_speed :: Float64
  average_exploration_depth :: Float64
  mcts_memory_footprint :: Int
  memory_size :: Int
  memory_num_distinct_boards :: Int
end

"""
    Report.Perfs

Performances report for a subroutine.
- `time`: total time spent, in seconds
- `allocated`: amount of memory allocated, in bytes
- `gc_time`: total amount of time spent in the garbage collector
"""
struct Perfs
  time :: Float64
  allocated :: Int64
  gc_time :: Float64
end

"""
    Report.Iteration

Report generated after each training iteration.
- Fields `self_play`, `memory`, `learning` have types [`Report.SelfPlay`](@ref),
    [`Report.SelfPlay`](@ref) and [`Report.Learning`](@ref) respectively
- Fields `perfs_self_play`, `perfs_memory_analysis` and `perfs_learning` are
    performance reports for the different phases of the iteration,
    with type [`Report.Perfs`](@ref)
"""
struct Iteration
  perfs_self_play :: Perfs
  perfs_memory_analysis :: Perfs
  perfs_learning :: Perfs
  self_play :: SelfPlay
  memory :: Union{Memory, Nothing}
  learning :: Learning
end

"""
    Report.Initial

Report summarizing the configuration of an agent before training starts.
- `num_network_parameters`: see [`Network.num_parameters`](@ref)
- `num_network_regularized_parameters`:
    see [`Network.num_regularized_parameters`](@ref)
- `mcts_footprint_per_node`: see [`MCTS.memory_footprint_per_node`](@ref)
"""
struct Initial
  num_network_parameters :: Int
  num_network_regularized_parameters :: Int
  mcts_footprint_per_node :: Int
end

#####
##### Profiling utilities
#####

macro timed(e)
  quote
    local v, t, mem, gct = Base.@timed $(esc(e))
    v, Perfs(t, mem, gct)
  end
end

end


"""
    graph_from_object(osm_data_object::Union{XMLDocument,Dict};
                      network_type::Symbol=:drive,
                      weight_type::Symbol=:time,
                      graph_type::Symbol=:static,
                      precompute_dijkstra_states::Bool=false,
                      largest_connected_component::Bool=true
                      )::OSMGraph

Creates an `OSMGraph` object from download OpenStreetMap network data, use with `download_osm_network`.

# Arguments
- `osm_data_object::Symbol`: OpenStreetMap network data parsed as either XML or Dictionary object depending on the download method.
- `network_type::Symbol=:drive`: Network type filter, pick from `:drive`, `:drive_service`, `:walk`, `:bike`, `:all`, `:all_private`, `:none`, `:rail`, must match the network type used to download `osm_data_object`.
- `weight_type::Symbol=:time`: Weight type for graph edges, pick from `:distance` (km), `:time` (hours), `:lane_efficiency` (time scaled by number of lanes).
- `graph_type::Symbol=:static`: Type of `LightGraphs.AbstractGraph`, pick from `:static` (StaticDiGraph), `:light` (DiGraph), `:simple_weighted` (SimpleWeightedDiGraph), `:meta` (MetaDiGraph).
- `precompute_dijkstra_states::Bool=false`: Set true to precompute dijkstra parent states for every source node in the graph, *NOTE* this may take a while and may not be possible for graphs with large amount of nodes due to memory limits.
- `largest_connected_component::Bool=true`: Set true to keep only the largest connected components in the network.

# Return
- `OSMGraph`: Container for storing OpenStreetMap node, way, relation and graph related obejcts.
"""
function graph_from_object(osm_data_object::Union{XMLDocument,Dict};
                           network_type::Symbol=:drive,
                           weight_type::Symbol=:time,
                           graph_type::Symbol=:static,
                           precompute_dijkstra_states::Bool=false,
                           largest_connected_component::Bool=true
                           )::OSMGraph
    g = init_graph_from_object(osm_data_object, network_type)
    add_node_and_edge_mappings!(g)
    add_weights!(g, weight_type)
    add_graph!(g, graph_type)
    # Finding connected components can only be done after LightGraph object has been constructed
    largest_connected_component && trim_to_largest_connected_component!(g, g.graph, weight_type, graph_type) # Pass in graph to make type stable
    add_node_tags!(g)
    !(network_type in [:bike, :walk]) && add_indexed_restrictions!(g)

    if precompute_dijkstra_states
        add_dijkstra_states!(g)
    else
        U = DEFAULT_OSM_INDEX_TYPE
        g.dijkstra_states = Vector{Vector{U}}(undef, length(g.nodes))
    end

    add_kdtree!(g)
    @info "Created OSMGraph object with kwargs: network_type=$network_type, weight_type=$weight_type, graph_type=$graph_type, precompute_dijkstra_states=$precompute_dijkstra_states, largest_connected_component=$largest_connected_component"
    return g
end

"""
    graph_from_object(file_path::String;
                      network_type::Symbol=:drive,
                      weight_type::Symbol=:time,
                      graph_type::Symbol=:static,
                      precompute_dijkstra_states::Bool=false,
                      largest_connected_component::Bool=true
                      )::OSMGraph

Creates an `OSMGraph` object from a downloaded OpenStreetMap network data file, the extention must be either `.json`, `.osm` or `.xml`.

# Arguments
- `file_path::String`: OpenStreetMap network data file location.
- `network_type::Symbol=:drive`: Network type filter, pick from `:drive`, `:drive_service`, `:walk`, `:bike`, `:all`, `:all_private`, `:none`, `:rail`, must match the network type used to download `osm_data_object`.
- `weight_type::Symbol=:time`: Weight type for graph edges, pick from `:distance` (km), `:time` (hours), `:lane_efficiency` (time scaled by number of lanes).
- `graph_type::Symbol=:static`: Type of `LightGraphs.AbstractGraph`, pick from `:static` (StaticDiGraph), `:light` (DiGraph), `:simple_weighted` (SimpleWeightedDiGraph), `:meta` (MetaDiGraph).
- `precompute_dijkstra_states::Bool=false`: Set true to precompute dijkstra parent states for every source node in the graph, *NOTE* this may take a while and may not be possible for graphs with large amount of nodes due to memory limits.
- `largest_connected_component::Bool=true`: Set true to keep only the largest connected components in the network.

# Return
- `OSMGraph`: Container for storing OpenStreetMap node, way, relation and graph related obejcts.
"""
function graph_from_file(file_path::String;
                         network_type::Symbol=:drive,
                         weight_type::Symbol=:time,
                         graph_type::Symbol=:static,
                         precompute_dijkstra_states::Bool=false,
                         largest_connected_component::Bool=true
                         )::OSMGraph

    !isfile(file_path) && throw(ArgumentError("File $file_path does not exist"))
    deserializer = file_deserializer(file_path)
    obj = deserializer(file_path)
    return graph_from_object(obj,
                             network_type=network_type,
                             weight_type=weight_type,
                             graph_type=graph_type,
                             precompute_dijkstra_states=precompute_dijkstra_states,
                             largest_connected_component=largest_connected_component)
end

"""
    graph_from_download(download_method::Symbol;
                        network_type::Symbol=:drive,
                        metadata::Bool=false,
                        download_format::Symbol=:json,
                        save_to_file_location::Union{String,Nothing}=nothing,
                        weight_type::Symbol=:time,
                        graph_type::Symbol=:static,
                        precompute_dijkstra_states::Bool=false,
                        largest_connected_component::Bool=true,
                        download_kwargs...
                        )::OSMGraph

Downloads OpenStreetMap network data and creates an `OSMGraph` object.

# Arguments
- `download_method::Symbol`: Download method, choose from `:place_name`, `:bbox` or `:point`.
- `network_type::Symbol=:drive`: Network type filter, pick from `:drive`, `:drive_service`, `:walk`, `:bike`, `:all`, `:all_private`, `:none`, `:rail`.
- `metadata::Bool=false`: Set true to return metadata.
- `download_format::Symbol=:json`: Download format, either `:osm`, `:xml` or `json`.
- `save_to_file_location::Union{String,Nothing}=nothing`: Specify a file location to save downloaded data to disk.
- `weight_type::Symbol=:time`: Weight type for graph edges, pick from `:distance` (km), `:time` (hours), `:lane_efficiency` (time scaled by number of lanes).
- `graph_type::Symbol=:static`: Type of `LightGraphs.AbstractGraph`, pick from `:static` (StaticDiGraph), `:light` (DiGraph), `:simple_weighted` (SimpleWeightedDiGraph), `:meta` (MetaDiGraph).
- `precompute_dijkstra_states::Bool=false`: Set true to precompute dijkstra parent states for every source node in the graph, *NOTE* this may take a while and may not be possible for graphs with large amount of nodes due to memory limits.
- `largest_connected_component::Bool=true`: Set true to keep only the largest connected components in the network.

# Required Kwargs for each Download Method

*`download_method=:place_name`*
- `place_name::String`: Any place name string used as a search argument to the Nominatim API.

*`download_method=:bbox`*
- `minlat::AbstractFloat`: Bottom left bounding box latitude coordinate.
- `minlon::AbstractFloat`: Bottom left bounding box longitude coordinate.
- `maxlat::AbstractFloat`: Top right bounding box latitude coordinate.
- `maxlon::AbstractFloat`: Top right bounding box longitude coordinate.

*`download_method=:point`*
- `point::GeoLocation`: Centroid point to draw the bounding box around.
- `radius::Number`: Distance (km) from centroid point to each bounding box corner.

*`download_method=:polygon`*
- `polygon::AbstractVector`: Vector of longitude-latitude pairs.

# Network Types
- `:drive`: Motorways excluding private and service ways.
- `:drive_service`: Motorways including private and service ways.
- `:walk`: Walkways only.
- `:bike`: Cycleways only.
- `:all`: All motorways, walkways and cycleways excluding private ways.
- `:all_private`: All motorways, walkways and cycleways including private ways.
- `:none`: No network filters.
- `:rail`: Railways excluding proposed and platform.

# Return
- `OSMGraph`: Container for storing OpenStreetMap node, way, relation and graph related obejcts.
"""
function graph_from_download(download_method::Symbol;
                             network_type::Symbol=:drive,
                             metadata::Bool=false,
                             download_format::Symbol=:json,
                             save_to_file_location::Union{String,Nothing}=nothing,
                             weight_type::Symbol=:time,
                             graph_type::Symbol=:static,
                             precompute_dijkstra_states::Bool=false,
                             largest_connected_component::Bool=true,
                             download_kwargs...
                             )::OSMGraph
    obj = download_osm_network(download_method,
                               network_type=network_type,
                               metadata=metadata,
                               download_format=download_format,
                               save_to_file_location=save_to_file_location;
                               download_kwargs...)
    return graph_from_object(obj,
                             network_type=network_type,
                             weight_type=weight_type,
                             graph_type=graph_type,
                             precompute_dijkstra_states=precompute_dijkstra_states,
                             largest_connected_component=largest_connected_component)
end

"""
Adds mappings between nodes, edges and highways to `OSMGraph`.
"""
function add_node_and_edge_mappings!(g::OSMGraph{U,T,W}) where {U <: Integer,T <: Integer,W <: Real}
    for (way_id, way) in g.highways
        @inbounds for (i, node_id) in enumerate(way.nodes)
            if haskey(g.node_to_highway, node_id)
                push!(g.node_to_highway[node_id], way_id)
            else
                g.node_to_highway[node_id] = [way_id]
            end

            if i < length(way.nodes)
                if !way.tags["reverseway"]::Bool
                    o = way.nodes[i] # origin
                    d = way.nodes[i + 1] # destination
                else
                    o = way.nodes[i + 1] # origin and destination reversed
                    d = way.nodes[i]
                end

                g.edge_to_highway[[o, d]] = way_id

                if !way.tags["oneway"]::Bool
                    g.edge_to_highway[[d, o]] = way_id
                end
            end
        end
    end

    @assert(length(g.nodes) == length(g.node_to_highway), "Data quality issue: number of graph nodes ($(length(g.nodes::Dict{T,Node{T}}))) not equal to set of nodes extracted from highways ($(length(g.node_to_highway)))")
    g.node_to_index = OrderedDict{T,U}(n => i for (i, n) in enumerate(collect(keys(g.nodes))))
    g.index_to_node = OrderedDict{U,T}(i => n for (n, i) in g.node_to_index)
end

"""
Adds maxspeed and lanes tags to every `OSMGraph` node.
"""
function add_node_tags!(g::OSMGraph)
    # Custom mean used to minimise allocations
    @inline function _roundedmean(hwys, subkey, T)
        total = zero(T)
        for id in hwys
            total += g.highways[id].tags[subkey]::T
        end
        return round(T, total/length(hwys))
    end

    M = DEFAULT_OSM_MAXSPEED_TYPE
    L = DEFAULT_OSM_LANES_TYPE

    for (id, data) in g.nodes
        highways = g.node_to_highway[id]
        tags_dict = g.nodes[id].tags::Dict{String, Any}
        tags_dict["maxspeed"] = _roundedmean(highways, "maxspeed", M)
        tags_dict["lanes"] = _roundedmean(highways, "lanes", L)
        push!(g.node_coordinates, [data.location.lat, data.location.lon])
    end
end

"""
Finds the adjacent node id on a given way.
"""
function adjacent_node(g::OSMGraph, node::T, way::T)::Union{T,Vector{<:T}} where T <: Integer
    way_nodes = g.highways[way].nodes
    if node == way_nodes[1]
        return way_nodes[2]
    elseif node == way_nodes[end]
        return way_nodes[length(way_nodes) - 1]
    else
        idx = findfirst(isequal(node), way_nodes)
        is_oneway = g.highways[way].tags["oneway"]
        is_reverseway = g.highways[way].tags["reverseway"]
        if is_oneway && !is_reverseway
            return way_nodes[idx + 1]
        elseif is_oneway && is_reverseway
            return way_nodes[idx - 1]
        else
            # via_node is in the middle of a non-oneway highway, this is only the possible when the
            # restriction is "only_straight_on", meaning vehicles can neither turn left nor right.
            return [way_nodes[idx - 1], way_nodes[idx + 1]]
        end
    end
end

"""
Adds restrictions linked lists to `OSMGraph`.

# Example
`[from_way_node_index, ...via_way_node_indices..., to_way_node_index]`
"""
function add_indexed_restrictions!(g::OSMGraph{U,T,W}) where {U <: Integer,T <: Integer,W <: Real}
    g.indexed_restrictions = DefaultDict{U,Vector{MutableLinkedList{U}}}(Vector{MutableLinkedList{U}})

    for (id, r) in g.restrictions
        if r.type == "via_node"
            if r.is_exclusion
                # no_left_turn, no_right_turn, no_u_turn, no_straight_on
                restricted_to_ways = [r.to_way]
            elseif r.is_exclusive
                # only_right_turn, only_left_turn, only_straight_on
                # Multiple to_ways, e.g. if only_right_turn, then left turn and straight on are restricted
                all_ways = g.node_to_highway[r.via_node]
                permitted_to_way = r.to_way
                restricted_to_ways = [w for w in all_ways if w != r.from_way && w != permitted_to_way]
            else
                continue
            end

            from_node = adjacent_node(g, r.via_node, r.from_way)::T
            for to_way in restricted_to_ways
                to_node_temp = adjacent_node(g, r.via_node, to_way)
                to_node = isa(to_node_temp, Integer) ? [to_node_temp] : to_node_temp

                for tn in to_node
                    # only_straight_on restrictions may have multiple to_nodes
                    indices = [g.node_to_index[n] for n in [tn, r.via_node::T, from_node]]
                    push!(g.indexed_restrictions[g.node_to_index[r.via_node]], MutableLinkedList{U}(indices...))
                end
            end

        elseif r.type == "via_way"
            via_way_nodes_list = [g.highways[w].nodes for w in r.via_way::Vector{T}]
            via_way_nodes = join_arrays_on_common_trailing_elements(via_way_nodes_list...)::Vector{T}

            from_way_nodes = g.highways[r.from_way].nodes
            from_via_intersection_node = first_common_trailing_element(from_way_nodes, via_way_nodes)
            from_node = adjacent_node(g, from_via_intersection_node, r.from_way)::T

            to_way_nodes = g.highways[r.to_way].nodes
            to_via_intersection_node = first_common_trailing_element(to_way_nodes, via_way_nodes)
            to_node = adjacent_node(g, to_via_intersection_node, r.to_way)::T

            if to_via_intersection_node == via_way_nodes[end]
                # Ordering matters, see doc string, but
                # we don't want to reorder the array in the Way object
                via_way_nodes = reverse(via_way_nodes)
            end

            indices = [g.node_to_index[n] for n in [to_node, via_way_nodes..., from_node]]
            push!(g.indexed_restrictions[g.node_to_index[to_via_intersection_node]], MutableLinkedList{U}(indices...))
        end
    end
end

"""
Adds edge weights to `OSMGraph`.
"""
function add_weights!(g::OSMGraph, weight_type::Symbol=:distance)
    n_edges = length(g.edge_to_highway)
    o_indices = Vector{Int}(undef, n_edges) # edge origin node indices
    d_indices = Vector{Int}(undef, n_edges) # edge destination node indices
    weights = Vector{Float64}(undef, n_edges)

    W = DEFAULT_OSM_EDGE_WEIGHT_TYPE

    @inbounds for (i, edge) in enumerate(keys(g.edge_to_highway))
        o_loc = g.nodes[edge[1]].location
        d_loc = g.nodes[edge[2]].location

        o_indices[i] = g.node_to_index[edge[1]]
        d_indices[i] = g.node_to_index[edge[2]]

        dist = distance(o_loc, d_loc, :haversine)
        if weight_type == :time || weight_type == :lane_efficiency
            highway = g.edge_to_highway[edge]
            maxspeed = g.highways[highway].tags["maxspeed"]::DEFAULT_OSM_MAXSPEED_TYPE
            if weight_type == :time
                weight = dist / maxspeed
            else
                lanes = g.highways[highway].tags["lanes"]::DEFAULT_OSM_LANES_TYPE
                lane_efficiency = get(LANE_EFFICIENCY, lanes, 1.0)
                weight = dist / (maxspeed * lane_efficiency)
            end
        else
            # Distance
            weight = dist
        end
        weights[i] = max(weight, eps(W))
    end

    n = length(g.nodes)
    g.weights = sparse(o_indices, d_indices, weights, n, n)
    g.weight_type = weight_type
    return g
end

"""
Adds a LightGraphs.AbstractGraph object to `OSMGraph`.
"""
function add_graph!(g::OSMGraph, graph_type::Symbol=:static)
    if graph_type == :light
        g.graph = DiGraph(g.weights)
    elseif graph_type == :static
        g.graph = StaticDiGraph(DiGraph(g.weights))
    elseif graph_type == :simple_weighted
        g.graph = SimpleWeightedDiGraph(g.weights)
    elseif graph_type == :meta
        g.graph = MetaDiGraph(DiGraph(g.weights))
        for (o, d, w) in zip(findnz(transpose(g.weights))...)
            set_prop!(g.graph, o, d, :weight, w)
        end
    else
        throw(ErrorException("Graph type $graph_type not implemented"))
    end
end

"""
Trims graph object to the largest connected component.
"""
function trim_to_largest_connected_component!(g::OSMGraph{U, T, W}, graph, weight_type::Symbol=:time, graph_type::Symbol=:static) where {U, T, W}
    cc = weakly_connected_components(graph)
    sort!(cc, by=x -> length(x), rev=true)
    indices_to_delete = flatten(cc[2:end])
    nodes_to_delete = [g.index_to_node[i] for i in indices_to_delete]
    highways_to_delete = Set(flatten([g.node_to_highway[n] for n in nodes_to_delete]))

    delete_from_dict!(g.nodes, nodes_to_delete, :on_key)
    delete_from_dict!(g.node_to_index, nodes_to_delete, :on_key)
    delete_from_dict!(g.node_to_highway, nodes_to_delete, :on_key)
    delete_from_dict!(g.index_to_node, indices_to_delete, :on_key)
    delete_from_dict!(g.highways, highways_to_delete, :on_key)
    delete_from_dict!(g.edge_to_highway, highways_to_delete, :on_value)

    for (id, r) in g.restrictions
        if r.via_node !== nothing
            !isempty(intersect([r.via_node::T], nodes_to_delete)) && delete!(g.restrictions, id)
        elseif r.via_way !== nothing
            ways = Set([r.from_way, r.to_way, r.via_way::Vector{T}...])
            !isempty(intersect(ways, highways_to_delete)) && delete!(g.restrictions, id)
        end
    end

    # Trim then rebuild graph so indices are in order
    add_node_and_edge_mappings!(g)
    add_weights!(g, weight_type)
    add_graph!(g, graph_type)
    return g
end

"""
Adds precomputed dijkstra states for every source node in `OSMGraph`, not recommended for graphs with greater than 50k nodes.
"""
function add_dijkstra_states!(g::OSMGraph{U,T,W}) where {U <: Integer,T <: Integer,W <: Real}
    srcs = collect(vertices(g.graph)) # source vertices (origin node indices)

    n = length(srcs)
    @info "Precomputing $n Dijkstra States, this might take a while..."
    n > 50000 && @warn "Precomputing and caching all $n Dijkstra States may not be possible due to memory limits"
    g.dijkstra_states = Vector{Vector{U}}(undef, n)

    Threads.@threads for src in srcs
        g.dijkstra_states[src] = dijkstra(g.graph, src, distmx=g.weights, restrictions=g.indexed_restrictions)
    end
end

"""
Adds KDTree to `OSMGraph` for finding nearest neighbours.
"""
function add_kdtree!(g::OSMGraph)
    node_locations = [node.location for (id, node) in g.nodes]  # node locations must have the same order as node indices
    cartesian_locations = to_cartesian(node_locations)
    g.kdtree = KDTree(cartesian_locations)
end


"""
    function disagg!(d::Dict, set::Dict; save_build = true, overwrite = false)

# Arguments
- `d::Dict` of DataFrames containing the model data.
- `set::Dict` of Arrays describing region, sector, final demand, etc.

# Keywords
- `save_build = true`
- `overwrite = false`
See [`SLiDE.build`](@ref) for keyword argument descriptions.

# Returns
- `d::Dict` of DataFrames containing the model data at the disaggregation step
"""
function disagg(
    dataset::String,
    d::Dict,
    set::Dict;
    save_build::Bool = DEFAULT_SAVE_BUILD,
    overwrite::Bool = DEFAULT_OVERWRITE
    )
    CURR_STEP = "disagg"

    d = merge(d, Dict(
        :r => fill_with((r = set[:r],), 1.0),
        (:yr,:r,:g) => fill_with((yr = set[:yr], r = set[:r], g = set[:g]), 1.0)))

    # If there is already disaggregation data, read it and return.
    d_read = read_build(dataset, CURR_STEP; overwrite = overwrite)
    !(isempty(d_read)) && (return (d_read, set))

    # Run all disaggregation calculations.
    _set_gm!(d, set)

    d[:region] = edit_with(d[:region], Rename(:g,:s))
    _disagg_ys0!(d)
    _disagg_id0!(d)
    _disagg_ty0!(d, set)
    _disagg_va0!(d, set)
    _disagg_ld0!(d)
    _disagg_kd0!(d)

    d[:region] = edit_with(d[:region], Rename(:s,:g))
    _disagg_fdcat!(d)
    _disagg_g0!(d)
    _disagg_i0!(d)
    _disagg_cd0!(d)
    _disagg_c0!(d)

    d[:yh0_temp] = _disagg_yh0!(d)
    _disagg_fe0!(d)
    d[:x0_temp] = _disagg_x0!(d, set)
    d[:s0_temp] = _disagg_s0!(d)
    _disagg_a0!(d)

    _disagg_ta0!(d)
    _disagg_tm0!(d)

    _disagg_thetaa!(d)
    _disagg_m0!(d)
    _disagg_md0!(d)
    d[:rx0_temp] = _disagg_rx0!(d)

    _disagg_diff!(d)
    _apply_diff!(d, set)

    _disagg_bop!(d)
    _disagg_pt0!(d)
    _disagg_dc0!(d)

    _disagg_dd0!(d)
    _disagg_nd0!(d)

    _disagg_dm0!(d)
    _disagg_nm0!(d)
    _disagg_xd0!(d)
    _disagg_xn0!(d)
    _disagg_hhadj!(d)

    # Should other
    d[:xn0][d[:xn0][:,:value] .< 1e-8,:value] .= 0;
    d[:xd0][d[:xd0][:,:value] .< 1e-8,:value] .= 0;

    write_build!(dataset, CURR_STEP, d; save_build = save_build)
    write_build!(dataset, SET_DIR, Dict(k => set[k] for k in [:gm]))
    return (d, set)
end

"""
`ys(yr,r,s,g)`, regional sectoral output

```math
\\bar{ys}_{yr,r,s,g} = \\alpha_{yr,r,s}^{gsp} \\tilde{ys}_{yr,s,g}
```
"""
function _disagg_ys0!(d::Dict)
    println("  Disaggregating ys0(yr,r,s,g), regional sectoral output")
    :r in propertynames(d[:ys0]) && (return d[:ys0])

    d[:ys0] = d[:region] * d[:ys0]
    return d[:ys0]
end

"""
`id(yr,r,g,s)`, regional intermediate demand

```math
\\bar{id}_{yr,r,g,s} = \\alpha_{yr,r,s}^{gsp} \\tilde{id}_{yr,g,s}
```
"""
function _disagg_id0!(d::Dict)
    println("  Disaggregating id0(yr,r,g,s), regional intermediate demand")
    :r in propertynames(d[:id0]) && (return d[:id0])

    d[:id0] = d[:region] * d[:id0]
    return d[:id0]
end

"""
`ty(yr,r,s)`, production tax rate

```math
\\begin{aligned}
\\bar{ty}_{yr,r,s}^{rev} &= \\alpha_{yr,r,s}^{gsp} \\tilde{va}_{yr,va,s} \\;\\forall\\; va = othtax \\\\
\\bar{ty}_{yr,r,s} &= \\frac{\\tilde{ty}_{yr,r,s}}{\\sum_{g} \\bar{ys}_{yr,r,s,g}}
\\end{aligned}
```
"""
function _disagg_ty0!(d::Dict, set::Dict)
    println("  Disaggregating ty0(yr,r,s), production tax rate")
    :va in propertynames(d[:va0]) && _unstack_va0!(d, set)
    idx = findindex(d[:va0])

    ty0_rev = d[:region] * d[:va0][:,[idx;:othtax]]
    d[:ty0] = dropnan(ty0_rev / combine_over(d[:ys0], :g))
    return d[:ty0]
end

"""
`va(yr,va,s)`, regional value added

```math
\\bar{va}_{yr,r,s} = \\alpha_{yr,r,s}^{gsp} \\sum_{va = compen,surplus} \\tilde{va}_{yr,va,s}
```
"""
function _disagg_va0!(d::Dict, set::Dict)
    println("  Disaggregating va0(yr,r,s), regional share of value added.")
    :r in propertynames(d[:va0])  && (return d[:va0])
    :va in propertynames(d[:va0]) && (_unstack_va0!(d, set))
    idx = findindex(d[:va0])

    df = d[:va0][:,[idx;:compen]] + d[:va0][:,[idx;:surplus]]
    d[:va0] = d[:region] * df
    return d[:va0]
end

"""
`ld0(yr,r,s)`, labor demand

```math
\\bar{ld}_{yr,r,s} = \\theta_{yr,r,s}^{ls} \\bar{va}_{yr,s,g}
```
"""
function _disagg_ld0!(d::Dict)
    println("  Disaggregating ld0(yr,r,s), labor demand")
    !(:r in propertynames(d[:va0])) && _disagg_va0!(d, set)
    d[:ld0] = d[:labor] * d[:va0]
    return d[:ld0]
end

"""
`kd0(yr,r,s)`, capital demand

```math
\\bar{kd}_{yr,r,s} = \\bar{va}_{yr,r,s} - \\bar{ld}_{yr,r,s}
```
"""
function _disagg_kd0!(d::Dict)
    println("  Disaggregating kd0(yr,r,s), capital demand")
    d[:kd0] = d[:va0] - d[:ld0]
end

"""
    _disagg_fdcat!(d::Dict)
This function aggregates final demand categories into national consumption (`C`), government
(`G`), and investment (`I`) demand.
"""
function _disagg_fdcat!(d::Dict)
    if "pce" in d[:fd0][:,:fd]
        x = [
            Map(joinpath("crosswalk","fd.csv"), [:fd], [:fdcat], [:fd], [:fd], :inner),
            Combine("sum",propertynames(d[:fd0])),
        ]
        d[:fd0] = edit_with(d[:fd0], x)
        # d[:fd0] = combine_over(d[:fd0], :fd)
    end
    return d[:fd0]
end

"""
`g0(yr,r,g)`, national government demand

```math
\\bar{g}_{yr,r,g} = \\alpha_{yr,r,g}^{sgf} \\sum_{G \\in fd} \\tilde{fd}_{yr,g,fd}
```
"""
function _disagg_g0!(d::Dict)
    println("  Disaggregating g0(yr,r,g), national government demand")
    ("pce" in d[:fd0][:,:fd]) && (_disagg_fdcat!(d))

    df = filter_with(d[:fd0], (fd="G",); drop = true)
    d[:g0] = d[:sgf] * df
    return d[:g0]
end

"""
`i0(yr,r,g)`, national investment demand

```math
\\bar{i}_{yr,r,g} = \\alpha_{yr,r,g}^{gsp} \\sum_{I \\in fd} \\tilde{fd}_{yr,g,fd}
```
"""
function _disagg_i0!(d::Dict)
    println("  Disaggregating i0(yr,r,g), national investment demand")
    !("pce" in d[:fd0][:,:fd]) && (_disagg_fdcat!(d))

    df = filter_with(d[:fd0], (fd="I",); drop = true)
    d[:i0] = d[:region] * df
    return d[:i0]
end

"""
`cd0(yr,r,g)`, national final consumption

```math
\\bar{cd}_{yr,r,g} = \\alpha_{yr,r,g}^{pce} \\sum_{C \\in fd} \\tilde{fd}_{yr,g,fd}
```
"""
function _disagg_cd0!(d::Dict)
    println("  Disaggregating cd0(yr,r,g), national final consumption")
    !("pce" in d[:fd0][:,:fd]) && (_disagg_fdcat!(d))

    df = filter_with(d[:fd0], (fd="C",); drop = true)
    d[:cd0] = d[:pce] * df
    return d[:cd0]
end

"""
`c0(yr,r)`, total final household consumption

```math
\\bar{c}_{yr,r} = \\sum_{g} \\bar{cd}_{yr,r,g}
```
"""
function _disagg_c0!(d::Dict)
    println("  Disaggregating c0(yr,r), total final household consumption")
    !(:cd0 in keys(d)) && _disagg_cd0!(d)
    d[:c0] = combine_over(d[:cd0], :g)
    return d[:c0]
end

"""
`yh0(yr,r,g)`, household production

```math
\\bar{yh}_{yr,r,g} = \\alpha_{yr,r,g} \\tilde{fs}_{yr,g}
```
"""
function _disagg_yh0!(d::Dict)
    if !(:diff in keys(d))
        println("  Disaggregating yh0(yr,r,g), household production")
        d[:yh0] = d[:region] * d[:fs0]
    else
        println("  Applying difference to yh0(yr,r,g), household production")
        d[:yh0] = d[:yh0] + d[:diff]
    end
    return dropmissing!(d[:yh0])
end

"""
`fe0(yr,r)`, total factor supply

```math
\\bar{fe}_{yr,r} = \\sum_{s} \\bar{va}_{yr,r,s}
```
"""
function _disagg_fe0!(d::Dict)
    println("  Disaggregating fe0(yr,r), total factor supply")
    d[:fe0] = combine_over(d[:va0], :s)
    return d[:fe0]
end

"""
`x0(yr,r,g)`, foreign exports

```math
\\bar{x}_{yr,r,g} = \\alpha_{yr,r,g}^{utd} \\tilde{x}_{yr,g}
```
"""
function _disagg_x0!(d::Dict, set::Dict)
    if !(:diff in keys(d))
        println("  Disaggregating x0(yr,r,g), foreign exports")
        !(:notrd in keys(set)) && _set_notrd!(d, set)

        df_exports = filter_with(d[:utd], (t = "exports",); drop = true)
        df_region = filter_with(d[:region], (g = set[:notrd],))

        df_trd = dropmissing(df_exports * d[:x0])
        df_notrd = dropmissing(df_region * d[:x0])

        d[:x0] = [df_trd; df_notrd]
    else
        println("  Applying difference to x0(yr,r,g), foreign exports")
        d[:x0] = d[:x0] + d[:diff]
    end

    return dropmissing!(d[:x0])
end

"""
`s0(yr,r,g)`, total supply

```math
\\bar{s}_{yr,r,g} = \\sum_{s} \\bar{ys}_{yr,r,s,g} + \\bar{yh}_{yr,r,g}
```
"""
function _disagg_s0!(d::Dict)
    if !(:diff in keys(d))
        println("  Disaggregating s0(yr,r,g), total supply")
        d[:s0] = combine_over(d[:ys0], :s) + d[:yh0]
    else
        println("  Applying difference to s0(yr,r,g), total supply")
        d[:s0] = d[:s0] + d[:diff]
    end
    return dropmissing!(d[:s0])
end

"""
`a0(yr,r,g)`, domestic absorption

```math
a_{yr,r,g} = \\bar{cd}_{yr,r,g} + \\bar{g}_{yr,r,g} + \\bar{i}_{yr,r,g} + \\sum_{s}\\bar{id}_{yr,r,g}
```
"""
function _disagg_a0!(d::Dict)
    println("  Disaggregating a0(yr,r,g), domestic absorption")

    d[:a0] = dropmissing(d[:cd0] + d[:g0] + d[:i0] + combine_over(d[:id0], :s))
    return d[:a0]
end

"`ta0(yr,r,g)`: Absorption taxes"
function _disagg_ta0!(d::Dict)
    println("  Disaggregating ta0(yr,r,g), absorption taxes")
    d[:ta0] = d[:ta0] * d[:r]
    return d[:ta0]
end

"`tm0(yr,r,g)`: Import taxes"
function _disagg_tm0!(d::Dict)
    println("  Disaggregating tm0(yr,r,g), import taxes")
    d[:tm0] = d[:tm0] * d[:r]
    return d[:tm0]
end

"""
`thetaa(yr,r,g)`, share of regional absorption

```math
\\alpha_{yr,r,g}^{abs} = \\frac{\\bar{a}_{yr,r,g}}{\\sum_{rr}\\bar{a}_{yr,r,g}}
```
"""
function _disagg_thetaa!(d::Dict)
    println("  Disaggregating thetaa(yr,r,g), share of regional absorption")
    d[:thetaa] = dropnan(d[:a0] / transform_over(d[:a0], :r))
    return d[:thetaa]
end

"""
`m0(yr,r,g)`, foreign Imports

```math
\\bar{m}_{yr,r,g} = \\alpha_{yr,r,g}^{abs} \\tilde{m}_{yr,g}
```
"""
function _disagg_m0!(d::Dict)
    println("  Disaggregating m0(yr,r,g), foreign imports")
    d[:m0] = d[:thetaa] * d[:m0]
    return d[:m0]
end

"""
`md0(yr,r,m,g)`, margin demand

```math
\\bar{md}_{yr,r,m,g} = \\alpha_{yr,r,g}^{abs} \\tilde{md}_{yr,m,g}
```
"""
function _disagg_md0!(d::Dict)
    println("  Disaggregating md0(yr,r,m,g), margin demand")
    d[:md0] = dropmissing(d[:thetaa] * d[:md0])
    return d[:md0]
end

"""
`rx0(yr,r,g)`, re-exports

```math
\\bar{rx}_{yr,r,g} = \\bar{x}_{yr,r,g} - \\bar{s}_{yr,r,g}
```
"""
function _disagg_rx0!(d::Dict; round_digits = DEFAULT_ROUND_DIGITS)
    if !(:diff in keys(d))
        println("  Disaggregating rx0(yr,r,g), re-exports")
        d[:rx0] = d[:x0] - d[:s0]

        if round_digits !== false
            d[:rx0][round.(d[:rx0][:,:value]; digits = round_digits) .< 0, :value] .= 0.0
        end
    else
        println("  Applying difference to rx0(yr,r,g), re-exports")
        d[:rx0] = d[:rx0] + d[:diff]
    end

    return dropmissing(d[:rx0])
end

"""
`dc0`,

```math
\\bar{dc}_{yr,r,g} = \\bar{s}_{yr,r,g} - \\bar{x}_{yr,r,g} + \\bar{rx}_{yr,r,g}
```
"""
function _disagg_dc0!(d::Dict; round_digits = DEFAULT_ROUND_DIGITS)
    # (!!!!) name for this?
    d[:dc0] = dropmissing((d[:s0] - d[:x0] + d[:rx0]))

    # if round_digits !== false
    #     d[:dc0][!,:value] .= round.(d[:dc0][:,:value]; digits = round_digits)
    # end
    (round_digits !== false) && (d[:dc0] = round!(d[:dc0], :value; digits = round_digits))
    return d[:dc0]
end

"""
`pt0`,

```math
\\begin{aligned}
\\bar{pt}_{yr,r,g} = &\\left(1 - \\bar{ta}_{yr,r,g} \\right) \\bar{a}_{yr,r,g} + \\bar{rx}_{yr,r,g}
\\\\               - &\\left(1 + \\bar{tm}_{yr,r,g} \\right) \\bar{m}_{yr,r,g} - \\sum_{m} \\bar{md}_{yr,r,m,g}
\\end{aligned}
```
"""
function _disagg_pt0(d::Dict; round_digits = DEFAULT_ROUND_DIGITS)
    df_pta0 = dropmissing(((d[:yr,:r,:g] - d[:ta0]) * d[:a0]) + d[:rx0])
    df_ptm0 = dropmissing(((d[:yr,:r,:g] + d[:tm0]) * d[:m0]) + combine_over(d[:md0], :m))

    df = df_pta0 - df_ptm0
    (round_digits !== false) && round!(df, :value; digits = round_digits)
    return df
end

function _disagg_pt0!(d::Dict; round_digits = DEFAULT_ROUND_DIGITS)
    d[:pt0] = _disagg_pt0(d, round_digits = round_digits)
    return d[:pt0]
end

"`diff`:"
function _disagg_diff!(d::Dict; round_digits = DEFAULT_ROUND_DIGITS)
    # df = _disagg_pt0!(d)
    df = _disagg_pt0(d; round_digits = DEFAULT_ROUND_DIGITS)

    df[!,:value] .= - min.(0, df[:,:value])
    d[:diff] = dropzero!(df)
end

"""
Add an adjustment to `rx_{yr,r,g}`, `s_{yr,r,g}`, `x_{yr,r,g}`, and `yh_{yr,r,g}`.
"""
function _apply_diff!(d::Dict, set::Dict)
    _disagg_rx0!(d)
    _disagg_s0!(d)
    _disagg_x0!(d, set)
    _disagg_yh0!(d)
end

"""
`bopdef0(yr,r)`, balance of payments (closure parameter)

```math
\\bar{bop}_{yr,r} = \\sum_{g} \\left( \\bar{m}_{yr,r,g} - \\bar{x}_{yr,r,g} \\right)
```
"""
function _disagg_bop!(d::Dict)
    println("  Disaggregating bopdef0(yr,r), balance of payments (closure parameter)")
    d[:bopdef0] = combine_over((d[:m0] - d[:x0]), :g)
    return d[:bopdef0]
end

"`gm`: Commodities employed in margin supply"
function _set_gm!(d::Dict, set::Dict)
    ms0_sum = combine_over(d[:ms0], [:yr,:m])
    md0_sum = combine_over(d[:md0], [:yr,:m])

    ms0_sum[!,:value] .= ms0_sum[:,:value] .!= 0.0
    md0_sum[!,:value] .= md0_sum[:,:value] .!= 0.0

    gm = ms0_sum + md0_sum
    set[:gm] = gm[gm[:,:value] .> 0, :g]
    return set[:gm]
end

"""
`dd0max(yr,r,g)`, maximum regional demand from local market

```math
\\hat{dd}_{yr,r,g} = \\min\\left\\{\\bar{pt}_{yr,r,g}, \\bar{dc}_{yr,r,g} \\right\\}
```
"""
function _disagg_dd0max!(d::Dict)
    println("  Disaggregating dd0max(yr,r,g), maximum regional demand from local market")
    cols = propertynames(d[:pt0])
    df = indexjoin(d[:pt0], d[:dc0]; id = [:pt0,:dc0])
    df[!,:value] = min.(df[:,:pt0], df[:,:dc0])

    d[:dd0max] = df[:,cols]
    return d[:dd0max]
end

"`nd0max(yr,r,g)`, maximum regional demand from national market"
function _disagg_nd0max!(d::Dict)
    println("  Disaggregating nd0max(yr,r,g), maximum regional demand from national market")
    cols = propertynames(d[:pt0])
    df = indexjoin(d[:pt0], d[:dc0]; id = [:pt0,:dc0])
    df[!,:value] = min.(df[:,:pt0], df[:,:dc0])

    d[:nd0max] = df[:,cols]
    return d[:nd0max]
end

"`dd0min(yr,r,g)`, minimum regional demand from local market"
_disagg_dd0min(d::Dict) = d[:pt0] - d[:dd0max]

"`nd0min(yr,r,g)`, minimum regional demand from national market"
_disagg_nd0min(d::Dict) = d[:pt0] - d[:nd0max]

"""
`dd0(yr,r,g)`, regional demand from local market

```math
\\bar{dd}_{yr,r,g} = \\rho_{r,g}^{cfs} \\hat{dd}_{yr,r,g}
```
"""
function _disagg_dd0!(d::Dict)
    println("  Disaggregating dd(yr,r,g), regional demand from local market")
    !(:dd0max in keys(d)) && _disagg_dd0max!(d)
    d[:dd0] = d[:dd0max] * d[:rpc]
    return d[:dd0]
end

"""
`nd0_(yr,r,g)`, regional demand from national market

```math
\\bar{nd}_{yr,r,g} = \\bar{pt}_{yr,r,g} - \\bar{dd}_{yr,r,g}
```
"""
function _disagg_nd0!(d::Dict; round_digits = DEFAULT_ROUND_DIGITS)
    println("  Disaggregating nd(yr,r,g), regional demand from national market")
    df_pt0 = _disagg_pt0(d; round_digits = false)

    d[:nd0] = df_pt0 - d[:dd0]
    # (round_digits !== false) && (d[:nd0][!,:value] .= round.(d[:nd0][:,:value]; digits = round_digits))
    (round_digits !== false) && (d[:nd0] = round!(d[:nd0], :value; digits = round_digits))
    return d[:nd0]
end

# ##########################################################################################

"""
`mrgshr(yr,r,m)`, share of margin demand by region

```math
\\alpha_{yr,r,m}^{md} = \\frac{\\sum_{g}\\bar{md}_{yr,r,m,g}}{\\sum_{r,g}\\bar{md}_{yr,r,m,g}}
```
"""
function _disagg_mrgshr(d::Dict)
    # (!!!!) notation alpha^{md} to match thetaa, alpha^a
    df = combine_over(d[:md0], :g)
    df = df / transform_over(df, :r)
    return df
end

"""
`ms0tot(yr,r,m,g)`, designate total supply of margins

```math
\\hat{ms}_{yr,r,m,g} = \\alpha_{yr,r,m}^{md} \\bar{ms}_{yr,g,m}
```
"""
function _disagg_ms0tot!(d::Dict)
    d[:ms0tot] = _disagg_mrgshr(d) * d[:ms0]
    return d[:ms0tot]
end

"""
`shrtrd(yr,r,g,m)`, share of margin total by margin type

```math
\\beta_{yr,r,g,m}^{mar} = \\frac{\\hat{ms}_{yr,r,g,m}}{\\sum_{m}\\hat{ms}_{yr,r,g,m}}
```
"""
function _disagg_shrtrd!(d::Dict)
    !(:ms0tot in keys(d)) && _disagg_ms0tot!(d)

    df = d[:ms0tot]
    d[:shrtrd] = dropnan(df / transform_over(df, :m))
    return d[:shrtrd]
end

"""
`dm0(yr,r,g,m)`, margin supply from the local market

```math
\\bar{dm}_{yr,r,g,m} = \\min\\left\\{ \\rho_{r,g}^{cfs}\\hat{ms}_{yr,r,g,m},
    \\beta_{yr,r,m,g}^{mar} \\left(\\bar{dc}_{yr,r,g} - \\bar{dd}_{yr,r,g}\\right) \\right\\}
```
"""
function _disagg_dm0!(d::Dict)
    println("  Disaggregating dm0(yr,r,g,m), margin supply from the local market")
    !(:ms0tot in keys(d)) && _disagg_ms0tot!(d)
    !(:shrtrd in keys(d)) && _disagg_shrtrd!(d)

    cols = propertynames(d[:ms0tot])
    dm1 = dropmissing(d[:ms0tot] * d[:rpc])
    dm2 = dropmissing((d[:shrtrd] * d[:dc0]) - d[:dd0])

    df = indexjoin(dm1, dm2; id = [:dm1,:dm2])
    df[!,:value] .= min.(df[:,:dm1], df[:,:dm2])
    d[:dm0] = df[:,cols]
    return d[:dm0]
end

"""
`nm0(yr,r,g,m)`, margin demand from the national market

```math
\\bar{nm}_{yr,r,g,m} = \\hat{ms}_{yr,r,g,m} - \\bar{dm}_{yr,r,g,m}
```
"""
function _disagg_nm0!(d::Dict)
    println("  Disaggregating nm0(yr,r,g,m), margin demand from the national market")
    d[:nm0] = d[:ms0tot] - d[:dm0]
    return d[:nm0]
end

"""
`xd0(yr,r,g)`, regional supply to local market

```math
\\bar{xd}_{yr,r,g} = \\sum_{m}\\bar{dm}_{yr,r,g,m} + \\bar{dd}_{yr,r,g}
```
"""
function _disagg_xd0!(d::Dict)
    println("  Disaggregating xd0(yr,r,g), regional supply to local market")
    d[:xd0] = combine_over(d[:dm0], :m) + d[:dd0]
    return d[:xd0]
end

"""
`xn0(yr,r,g)`, regional supply to national market

```math
\\bar{xn}_{yr,r,g} = \\bar{s}_{yr,r,g} + \\bar{rx}_{yr,r,g} - \\bar{xd}_{yr,r,g} - \\bar{x}_{yr,r,g}
```
"""
function _disagg_xn0!(d::Dict)
    println("  Disaggregating xn0(yr,r,g), regional supply to national market")
    d[:xn0] = d[:s0] + d[:rx0] - d[:xd0] - d[:x0]
    return d[:xn0]
end

"""
`hhadj(yr,r)`, household adjustment

```math
\\begin{aligned}
\\bar{adj^{hh}}_{yr,r} = &\\bar{c}_{yr,r}
\\\\ &- \\sum_{s}\\left( \\bar{ld}_{yr,r,s} + \\bar{kd}_{yr,r,s} + \\bar{yh}_{yr,r,s} \\right) - \\bar{bop}_{yr,r}
\\\\ &- \\sum_{s}\\left( \\bar{ta}_{yr,r,s}\\bar{a0}_{yr,r,s} + \\bar{tm}_{yr,r,s}\\bar{m}_{yr,r,s} + \\bar{ty}_{yr,r,s}\\sum_{g}\\bar{ys}_{yr,r,s,g} \\right)
\\\\ &+ \\sum_{s}\\left( \\bar{g}_{yr,r,s} + \\bar{i}_{yr,r,s} \\right)
\\end{aligned}
```
"""
function _disagg_hhadj!(d::Dict)
    println("  Disaggregating hhadj(yr,r), household adjustment")
    dh = Dict(k => edit_with(copy(d[k]), Rename(:g,:s))
        for k in [:c0,:ld0,:kd0,:yh0,:bopdef0,:ta0,:a0,:tm0,:ty0,:g0,:i0,:m0])
    dh[:ys0] = copy(d[:ys0])

    d[:hhadj] = dh[:c0] -
        combine_over(dh[:ld0] + dh[:kd0] + dh[:yh0], :s) - dh[:bopdef0] -
        combine_over(dh[:ta0]*dh[:a0] + dh[:tm0]*dh[:m0] + dh[:ty0]*combine_over(dh[:ys0],:g), :s) +
        combine_over(dh[:g0] + dh[:i0], :s)

    return d[:hhadj]
end

function _unstack_va0!(d::Dict, set::Dict)
    if :va in propertynames(d[:va0])
        d[:va0] = edit_with(unstack(d[:va0], :va, :value),
            Replace.(Symbol.(set[:va]), missing, 0.0))
    end
    return d[:va0]
end