#module Cabida_nlo_evol

##############################################
# PARTE "1": CARGA CABIDA                    #
##############################################

using LotMassing, .poly2D, .polyShape, Random, CSV

#Random.seed!(1236)
#Random.seed!(1230)

##############################################
# PARTE "2": GENERACIÓN DE PARÁMETROS        #
##############################################

idPredio = 8 #8 predio = 1,2,3,4,5,6,75,8
conjuntoTemplates = [4] #4 [1:L, 2:C, 4:V]

fpe = FlagPlotEdif3D(true,  # predio
                     true,  # volTeor
                     true,  # restSombra
                     true,  # edif
                     true,  # sombraVolTeor_p
                     true,  # sombraVolTeor_o
                     true,  # sombraVolTeor_s
                     true,  # sombraEdif_p
                     true,  # sombraEdif_o
                     true)  # sombraEdif_s

dcn = datosCabidaNormativa(5, # SEPMIN (m): max(4, separación mínima deslindes) OGUC 2.6.3
                     5, # ANTEJARDIN (m)
                     2.75, # RASANTE (m de altura / m de separación)
                     60, # 24, #ALTURAMAX (m)
                     30, # 9, #MAXPISOS (unidades)
                     .7, # COEFOCUPACION (m2 / m2 de terreno)
                     1000, # SUBPREDIALMIN (m2)
                     2500, # DENSIDADMAX (Habitantes / 10000 m2 de terreno bruto)
                     true, # FLAGDENSIDADBRUTA
                     4, # COEFCONSTRUCTIBILIDAD (m2 / m2 de terreno)
                     [.5, 1, 1, 1], # ESTACIONAMIENTOSPORVIV (unidades / departamento)
                     .15, # PORCADICESTACVISITAS (unidades / estacionamiento vendible)
                     34, # SUPPORESTACIONAMIENTO (m2 / Estacionamiento)
                     .25, # ESTBICICLETAPORVIV (unidades / estacionamientos totales)
                     3, # BICICLETASPOREST
                     true, # FLAGCAMBIOESTPORBICICLETA
                     10, # MAXSUBTE (unidades)
                     .8, # COEFOCUPACIONEST (m2 / superficieTerreno)
                     7, # SEPESTMIN (m)
                     0 # REDUCCIONESTPORDISTMETRO
                     );

dca = datosCabidaArquitectura(2.55, # 2.625, # ALTURAPISO (m / piso)
                         .2, # PORCSUPCOMUN (m2 / m2 útil)
                         .05, # PORCTERRAZA (m2 / m2 útil)
                         8,#12, # ANCHOMAX (m)
                         [],
                         []
                         );

dcc = datosCabidaComercial([20, 30, 38, 45], # SUPDEPTOUTIL (m2)
                     [.3, .4, .4, 1], # MAXPORCTIPODEPTO
                     [56, 55, 54, 52], # PRECIOVENTA (UF / m2 vendible)
                     200 # PRECIOVENTAEST (UF / unidad)
                     );

dcu = datosCabidaUnit(30, # duracionProyecto
             20, # CostoTerreno
             2 * 1.19 / 100, # ComisionCorredor
             0.60, # Demolicion
             20, # LosaSNT
             10, # LosaBNT
             5, # EstacVisita
             3 / 100, # ExtrasPostVenta
             1.75 / 100, # Arquitectura
             2.75 / 100, # GestionAdministracion
             30, # ITO
             30, # Contabilidad
             0.85 / 100, # Legales
             0.145, # Calculo
             250, # MecanicaSuelo
             50, # Topografia
             .45 / 100, # ProyectoElectrico
             0, # ProyectoSanitario
             0, # ProyectoSeguridad
             50, # ProyectoPavimentacion
             25, # ProyectoBasura
             0, # ProyectoCalefaccion
             1.2 / 100, # Marketing
             0.5 / 100, # PilotSalaVentas
             0, # NotarialesLegales
             0, # CBR
             0, # EmpalmesAportes
             600, # DerechosPermisosMunicipales
             15, # InspeccionCertificacionTecnica
             0, # Contribuciones
             120, # CopiaPlanosArquitectura
             60, # CopiaPlanosObra
             300, # HallAcceso
             400, # JardinesPaisajismo
             0, # HabilitacionSalaUsosMultiples
             0, # HabilitacionPiscinaEnfermeria
             0, # CuentasServiciosDeptos
             0, # GastosComunes
             0, # ContribucionesDeptos
             0, # Seguros
             1 / 100, # ResponsabilidadCivil
             1 / 100, # PostVentaInmobiliaria
             0.1 / 100, # SeguroVentaEnVerde
             1 / 100 # Imprevistos
             );

dcf = datosCabidaFlujo([0.05, 0.05, 0.05, 0.05, 0.10, 0.70], # IngresosVentas
             [1, 0, 0, 0, 0, 0], # CostoTerreno
             [0.05, 0.10, 0.35, 0.35, .15, 0], # CostoConstruccion
             [0.40, 0.15, 0.15, 0.15, .15, 0], # CostoHonorariosProyectos
             [0, 0.10, 0.20, 0.20, 0.25, 0.25], # CostoVenta
             [0.1, 0.40, 0.25, 0.15, 0.05, 0.05], # CostoInmobiliarioObra
             [0, 0, 0, 0, 0.50, 0.50], # CostosHabilitacion
             [0, 0, 0, 0, 0.25, 0.75], # CostosPuestaEnMarcha
             [0.05, 0.05, 0.05, 0.05, 0.05, 0.75], # CostosAtencionCliente
             [0, 0.20, 0.20, 0.20, 0.20, 0.20], # Imprevistos
             [0.05, 0.10, 0.35, 0.35, .15, 0], # ingresoLineaCredito
             [0, 0, 0.10, 0.35, 0.35, .2], # pagoLineaCredito
             .27, # TasaImpuestoRenta
             .05, # TasaInteresLineaCredito
             0.8 * dcu.duracionProyecto, # DuracionLineaCredito
             15 # Tasaciones
             );

dcr = datosCabidaRentabilidad(1.20) # RetornoExigido

if idPredio == 1
    factorCorreccion = 2;
    x = factorCorreccion * [0 30 40 45 10]';
    y = factorCorreccion * [10 0 15 35 30]';
# Calle: [4]
    dcp = datosCabidaPredio(x, y, [1 4], [20 20], 1, 200);

elseif idPredio == 2
    factorCorreccion = 5;
    x = factorCorreccion * [0 10 15 20 25 13 12 4 9]';
    y = factorCorreccion * [10 0 12 7 20 25 22 28 14]';
# Calle: [1 5]
    dcp = datosCabidaPredio(x, y, [1 5], [20 20], 1, 200);

elseif idPredio == 3
# EJEMPLO 1 (Córdova y Figueroa con Vargas Fontecilla, Quinta Normal)
    factorCorreccion = 0.8345;
    x = factorCorreccion * [-7869557.5 -7869559.7 -7869535.4 -7869535.5 -7869490.4 -7869488.3]';
    y = factorCorreccion * [-3952906.8 -3952965.7 -3952966.5 -3952970.9 -3952972.4 -3952909.6]';
# Calle: [1 6]
    dcp = datosCabidaPredio(x, y, [1 6], [18 20], 0, 200);

elseif idPredio == 4
# EJEMPLO 2 (Catedral llegando a General Velásquez, Quinta Normal)
    factorCorreccion = 0.8547;
    x = factorCorreccion * [-7869442.1 -7869413.9 -7869409.5 -7869438.5]';
    y = factorCorreccion * [-3953939.3 -3953935.5 -3953861.6 -3953864.5]';
# Calle: [1]
    dcp = datosCabidaPredio(x, y, [1], [20], 1, 200);

elseif idPredio == 5
    factorCorreccion = 1;
    x = factorCorreccion * [0 150 150 100 100 50 50 0]';
    y = factorCorreccion * [0 0 100 100 50 50 100 100]';
# Calle: [4]
    dcp = datosCabidaPredio(x, y, [1 8], [15 15], 1, 200);

elseif idPredio == 6
    # EJEMPLO 6 (Augusto Leguía)
    factorCorreccion = 0.4;
    x = factorCorreccion * [568 571 612 613 646 648 683 681]';
    y = factorCorreccion * [-339 -405 -403 -418 -419 -452 -442 -355]';
    # Calle: [1 6]
    dcp = datosCabidaPredio(x, y, [1 6 8], [15 10 20], 0, 200);

elseif idPredio == 7
    # EJEMPLO 7 (Independencia)

    dcn.SEPMIN = 5 # SEPMIN (m): max(4, separación mínima deslindes) OGUC 2.6.3
    dcn.ANTEJARDIN = 5 # ANTEJARDIN (m)
    dcn.ALTURAMAX = 30 # 24, #ALTURAMAX (m)
    dcn.MAXPISOS = 20 # 9, #MAXPISOS (unidades)
    dcn.COEFOCUPACION = .5 # COEFOCUPACION (m2 / m2 de terreno)
    dcn.SUBPREDIALMIN = 300 # SUBPREDIALMIN (m2)
    dcn.DENSIDADMAX = 2000 # DENSIDADMAX (Habitantes / 10000 m2 de terreno bruto)
    dcn.COEFCONSTRUCTIBILIDAD = 3 # COEFCONSTRUCTIBILIDAD (m2 / m2 de terreno)

    dca.ANCHOMAX = 10 # ANCHOMAX (m)

    #dcc.SUPDEPTOUTIL = [20, 35, 50, 60] # SUPDEPTOUTIL (m2)
    #dcc.PRECIOVENTA = [65, 58, 53, 50] # PRECIOVENTA (UF / m2 vendible)

    dcc.SUPDEPTOUTIL = [30, 40, 50, 65] # SUPDEPTOUTIL (m2)
    dcc.PRECIOVENTA = [65, 58, 53, 50] # PRECIOVENTA (UF / m2 vendible)
    dcc.MAXPORCTIPODEPTO = [.50, .10, .33, .18] # MAXPORCTIPODEPTO
    dcr = datosCabidaRentabilidad(1.15) # RetornoExigido

    nombreArchivo = "Independencia.csv"
    loadData = CSV.File(string("C:/Users/rjvia/Downloads/", nombreArchivo); header=false)
    numDatos = length(loadData)
    x = zeros(1,numDatos)
    y = zeros(1,numDatos)
    for i = 1:numDatos
        x[i] = loadData[i].Column1
        y[i] = loadData[i].Column2
    end
    areaSup = x[1]
    x = x[2:end]
    y = y[2:end]
    V = [x y]
    factorCorreccion = ajusteArea(V, areaSup)
    dcp = datosCabidaPredio(factorCorreccion * x, factorCorreccion * y, [1 4], [16.5 13.7], 1, 200);

elseif idPredio == 8
    # EJEMPLO 7 (La Florida)

    dcn.SEPMIN = 5 # SEPMIN (m): max(4, separación mínima deslindes) OGUC 2.6.3
    dcn.ANTEJARDIN = 5 # ANTEJARDIN (m)
    dcn.ALTURAMAX = 60 #47 # 24, #ALTURAMAX (m)
    dcn.MAXPISOS = 40 # 9, #MAXPISOS (unidades)
    dcn.COEFOCUPACION = .4 # COEFOCUPACION (m2 / m2 de terreno)
    dcn.SUBPREDIALMIN = 1000 # SUBPREDIALMIN (m2)
    dcn.DENSIDADMAX = 20000 # DENSIDADMAX (Habitantes / 10000 m2 de terreno bruto)
    dcn.COEFCONSTRUCTIBILIDAD = 3.2 # COEFCONSTRUCTIBILIDAD (m2 / m2 de terreno)

    dcc.SUPDEPTOUTIL = [20, 35, 45, 55] # SUPDEPTOUTIL (m2)
    dcc.PRECIOVENTA = [65, 63, 60, 57] # PRECIOVENTA (UF / m2 vendible)
    dcc.MAXPORCTIPODEPTO = [0, 1, 1, 1];

    dca.ANCHOMAX = 8 # Ancho Crujía (m)

    nombreArchivo = "predioLaFlorida2.csv"
    loadData = CSV.File(string("C:/Users/rjvia/Downloads/", nombreArchivo); header=false)
    numDatos = length(loadData)
    x = zeros(1,numDatos)
    y = zeros(1,numDatos)
    for i = 1:numDatos
        x[i] = loadData[i].Column1
        y[i] = loadData[i].Column2
    end
    areaSup = x[1]
    x = x[2:end]
    y = y[2:end]
    V = [x y]
    V = polyShape.reversePath(V);
    x = V[1:end,1]
    y = V[1:end,2]
    factorCorreccion = ajusteArea(V, areaSup)
    dcp = datosCabidaPredio(factorCorreccion * x, factorCorreccion * y, [1 3], [17.5 55], 0, 200);

elseif idPredio == 9
    # EJEMPLO 9 (El Dante)

    dcn.SEPMIN = 8 # SEPMIN (m): max(4, separación mínima deslindes) OGUC 2.6.3
    dcn.ANTEJARDIN = 7 # ANTEJARDIN (m)
    dcn.ALTURAMAX = 52.5 #47 # 24, #ALTURAMAX (m)
    dcn.MAXPISOS = 15 # 9, #MAXPISOS (unidades)
    dcn.COEFOCUPACION = .3 # COEFOCUPACION (m2 / m2 de terreno)
    dcn.SUBPREDIALMIN = 1500 # SUBPREDIALMIN (m2)
    dcn.DENSIDADMAX = 880 # DENSIDADMAX (Habitantes / 10000 m2 de terreno bruto)
    dcn.FLAGDENSIDADBRUTA = false # FLAGDENSIDADBRUTA
    dcn.COEFCONSTRUCTIBILIDAD = 2.8 # COEFCONSTRUCTIBILIDAD (m2 / m2 de terreno)
    dcn.ESTACIONAMIENTOSPORVIV = [1, 2, 2, 2] # ESTACIONAMIENTOSPORVIV
    dcn.FLAGCAMBIOESTPORBICICLETA = false # FLAGCAMBIOESTPORBICICLETA

    dcc.SUPDEPTOUTIL = [50, 90, 110, 140] # SUPDEPTOUTIL (m2)
    dcc.PRECIOVENTA = [95, 91, 89, 87] # PRECIOVENTA (UF / m2 vendible)
    dcc.MAXPORCTIPODEPTO = [1, 1, 1, 1];
    dcc.PRECIOVENTAEST = 350 # PRECIOVENTAEST (UF / unidad)

    dcu.LosaSNT = 30 # LosaSNT
    dcu.LosaBNT = 10 # LosaBNT

    dca.ANCHOMAX = 8 # Ancho Crujía (m)
    dca.ALTURAPISO = 2.7 # 2.625, # ALTURAPISO (m / piso)
    dca.PORCSUPCOMUN = .2 # PORCSUPCOMUN (m2 / m2 útil)

    #nombreArchivo = "el_dante_2.csv"
    nombreArchivo = "predio_ElDante.csv"
    loadData = CSV.File(string("C:/Users/rjvia/Downloads/", nombreArchivo); header=false)
    numDatos = length(loadData)
    x = zeros(1,numDatos)
    y = zeros(1,numDatos)
    for i = 1:numDatos
        x[i] = loadData[i].Column1
        y[i] = loadData[i].Column2
    end
    areaSup = x[1]
    x = x[2:end]
    y = y[2:end]
    V = [x y]
    factorCorreccion = ajusteArea(V, areaSup)
    x = factorCorreccion * x
    y = factorCorreccion * y

#    dcp = datosCabidaPredio(x, y, [1,2,3], [15,15,15], 0, 200);
    dcp = datosCabidaPredio(x, y, [1,2,3,4,5,6], [15,15,15,15,15,20], 0, 200);

end

resultados, ps_predio, ps_base, xopt_cs, fopt_cs, polyCorte,
ps_SombraVolTeor_p, ps_sombraEdif_p,
ps_SombraVolTeor_o, ps_sombraEdif_o,
ps_SombraVolTeor_s, ps_sombraEdif_s = executaCalculoCabidas(dcp, dcn, dca, dcc, dcu, dcf, dcr, fpe, conjuntoTemplates);


Network = SimpleNet

netparams = SimpleNetHP(
    width = 200,
    depth_common = 8,
    use_batch_norm = true,
    batch_norm_momentum = 1.0,
)

self_play = SelfPlayParams(
    num_games = 2000,
    mcts = MctsParams(
        use_gpu = true,
        num_workers = 1,
        num_iters_per_turn = 200,
        cpuct = 1.0,
        temperature = ConstSchedule(1.0),
        dirichlet_noise_ϵ = 0.2,
        dirichlet_noise_α = 1.0,
    ),
)

arena = ArenaParams(
    num_games = 10,
    reset_mcts_every = 1,
    update_threshold = 0.00,
    flip_probability = 0.5,
    mcts = MctsParams(
        self_play.mcts,
        temperature = ConstSchedule(0.3),
        dirichlet_noise_ϵ = 0.1,
    ),
)

learning = LearningParams(
    samples_weighing_policy = LOG_WEIGHT,
    l2_regularization = 1e-4,
    optimiser = CyclicNesterov(
        lr_base = 1e-3,
        lr_high = 1e-2,
        lr_low = 1e-3,
        momentum_high = 0.9,
        momentum_low = 0.8,
    ),
    batch_size = 24,
    loss_computation_batch_size = 1024,
    nonvalidity_penalty = 1.0,
    min_checkpoints_per_epoch = 0,
    max_batches_per_checkpoint = 5_000,
    num_checkpoints = 1,
)

params = Params(
    arena = arena,
    self_play = self_play,
    learning = learning,
    num_iters = 10,
    memory_analysis = MemAnalysisParams(num_game_stages = 5),
    ternary_rewards = true,
    use_symmetries = false,
    mem_buffer_size = PLSchedule(80_000),
)
benchMcts = MctsParams(
    use_gpu = false,
    num_workers = 1,
    num_iters_per_turn = 100,
    cpuct = 0.9,
    temperature = ConstSchedule(1.0),
    dirichlet_noise_ϵ = 0.2,
    dirichlet_noise_α = 1.0,
)
benchmark = [
    Benchmark.Duel(
        Benchmark.NetworkOnly(),
        Benchmark.MctsRollouts(benchMcts),
        num_games = 20,
        flip_probability = 0.5,
    ),
    Benchmark.Duel(
        Benchmark.Full(self_play.mcts),
        Benchmark.MctsRollouts(self_play.mcts),
        num_games = 20,
        flip_probability = 0.4,
    ),
]


using Statistics

"""calculate number of hydrogenerators"""
function countgenerators!(data::Dict)
    data["hydro"]["nHyd"] = size(data["hydro"]["Hydrogenerators"],1)
end

"""compoute upstream_hydro"""
function upstream_hydro!(data::Dict)
    for i in 1:data["hydro"]["nHyd"]
        data["hydro"]["Hydrogenerators"][i]["upstrem_hydros_turn"] = []
        data["hydro"]["Hydrogenerators"][i]["upstrem_hydros_spill"] = []
    end
    for i in 1:data["hydro"]["nHyd"]
        for idx in data["hydro"]["Hydrogenerators"][i]["downstream_turn"]
            j = findall(x->x["index"]==idx,data["hydro"]["Hydrogenerators"])
            if length(j)!=1
                error("Incoherent downstream_turn list")
            end
            j = j[1]
            data["hydro"]["Hydrogenerators"][j]["upstrem_hydros_turn"] = append!(data["hydro"]["Hydrogenerators"][j]["upstrem_hydros_turn"],i)
        end
        for idx in data["hydro"]["Hydrogenerators"][i]["downstream_spill"]
            j = findall(x->x["index"]==idx,data["hydro"]["Hydrogenerators"])
            if length(j)!=1
                error("Incoherent downstream_spill list")
            end
            j = j[1]
            data["hydro"]["Hydrogenerators"][j]["upstrem_hydros_spill"] = append!(data["hydro"]["Hydrogenerators"][j]["upstrem_hydros_spill"],i)
        end
    end
    for i in 1:data["hydro"]["nHyd"]
        unique!(data["hydro"]["Hydrogenerators"][i]["upstrem_hydros_turn"])
        unique!(data["hydro"]["Hydrogenerators"][i]["upstrem_hydros_spill"])
    end
    return nothing
end

"""create ref for anonimous variables on model"""
function createvarrefs!(sp::JuMP.Model,pm::AbstractPowerModel)
    for listvarref in values(PowerModels.var(pm,pm.cnw,pm.ccnd))
        for variableref in values(listvarref)
            if typeof(variableref) == JuMP.VariableRef
                sp[Symbol(name(variableref))] = variableref
            end
        end
    end
end

"""count available inflow data"""
function countavailableinflow!(data::Dict)
    data["hydro"]["size_inflow"] = size(data["hydro"]["Hydrogenerators"][1]["inflow"])
end

"""circular index"""
function cidx(i::Int,n::Int)
    mod(i,n)==0 ? n : mod(i,n)
end

"""hydrogenerators indexes"""
function idx_hydro(data::Dict)
    [hydro["i_grid"] for hydro in data["hydro"]["Hydrogenerators"] if hydro["index_grid"] != nothing]
end

"""find vector index of hydrogenerators on grid"""
function index2i!(data::Dict)
    for i=1:data["hydro"]["nHyd"]
        if data["hydro"]["Hydrogenerators"][i]["index_grid"] != nothing
            i_grid = findall(x->x["index"] == data["hydro"]["Hydrogenerators"][i]["index_grid"],data["powersystem"]["gen"])
            if length(i_grid) != 1
                error("Generator of hydro with index $(data["hydro"]["Hydrogenerators"][i]["index"]) not identifiable")
            end
            data["hydro"]["Hydrogenerators"][i]["i_grid"] = parse(Int64,i_grid[1])
        end
    end
    return nothing
end

"""gather useful information from data"""
function gatherusefulinfo!(data::Dict)
    # calculate number of hydrogenerators
    countgenerators!(data)

    # compoute upstream_hydro
    upstream_hydro!(data)

    # count available inflow data
    countavailableinflow!(data)

    # find index hydrogen
    index2i!(data)

    # add zeroed loads on buses w no loads for deficit coherence
    add_loads!(data)

    return nothing
end

"""Quantile Scenarios"""
function quantile_scen(scen::Array{Float64,2},quant::Float64)
    return quantile_scen(scen,[quant])[:,1]
end

"""Quantile Scenarios"""
function quantile_scen(scen::Array{Float64,2},quants::Array{Float64};output_dict::Bool=false)
    quantiles = [Statistics.quantile(scen[i, :], quant) for i = 1:size(scen, 1),quant in quants]
    if output_dict
        output = Dict()
        for col = 1:length(quants)
            quant = quants[col]
            output["$(quant*100)%"] = quantiles[:,col]
        end
        return output
    end

    return quantiles
end

"""Multilayer Dict to Onelayer Dict"""
function flat_dict(mlt_dict::Dict{Any,Any})
    if typeof(collect(values(mlt_dict))[1]) != Dict{Any,Any}
        return mlt_dict
    end

    one_dict = Dict()
    kws = collect(keys(mlt_dict))

    recursion_ret = [flat_dict(i) for i in values(mlt_dict)]
    for i = 1:size(recursion_ret,1)
        item = recursion_ret[i]
        for (ikw,val) in item
            one_dict[kws[i]*"_"*ikw] = val
        end
    end

    return one_dict

end

"""truncate values dict"""
function signif_dict(one_dict::Dict, digts::Integer)
    for kw in keys(one_dict)
        if typeof(one_dict[kw]) <: AbstractFloat
            one_dict[kw] = round.(one_dict[kw],digits=digts)
        end
    end
    return one_dict
end

"""add zeroed loads on buses w no loads for deficit coherence"""
function add_loads!(data::Dict)
    locked_buses = unique([load["load_bus"] for load in values(data["powersystem"]["load"])])
    free_buses = setdiff(collect(1:length(data["powersystem"]["bus"])),locked_buses)

    load_idx = length(data["powersystem"]["load"])+1
    for bus in free_buses
        data["powersystem"]["load"]["$load_idx"] = load_dict(   load_bus = bus,
                                                                pd = 0,
                                                                index = load_idx)
        load_idx +=1
    end
end

""" creates a load dictionary """
function load_dict(;
        load_bus= 1, # bus where load is present
        status= 1, # status (ON: 1, OFF:0)
        pd= 1, # reactive power demand
        qd= 0.12*pd, # real power demand
        index= 1) # index

    lod = Dict( "load_bus"=> load_bus,
        "status"=> status,
        "qd"=> qd,
        "pd"=> pd,
        "index"=> index)
end

""" water value """
function water_energy!(data::Dict)
    countgenerators!(data)
    for i in 1:data["hydro"]["nHyd"]
        if !haskey(data["hydro"]["Hydrogenerators"][i],"water_energy")
            water_energy_res!(data,i)
        end
    end
    return nothing
end

""" water value """
function water_energy_res!(data::Dict,res::Int)
    data["hydro"]["Hydrogenerators"][res]["water_energy"] = 0.0
    water_val_turn = data["hydro"]["Hydrogenerators"][res]["production_factor"]
    for idx in data["hydro"]["Hydrogenerators"][res]["downstream_turn"]
        j = findall(x->x["index"]==idx,data["hydro"]["Hydrogenerators"])
        if length(j)!=1
            error("Incoherent downstream_turn list")
        end
        j = j[1]
        water_val_turn += water_energy_res!(data,j)
    end
    water_val_spill = 0.0
    for idx in data["hydro"]["Hydrogenerators"][res]["downstream_spill"]
        j = findall(x->x["index"]==idx,data["hydro"]["Hydrogenerators"])
        if length(j)!=1
            error("Incoherent downstream_spill list")
        end
        j = j[1]
        water_val_spill += water_energy_res!(data,j)
    end

    data["hydro"]["Hydrogenerators"][res]["water_energy"] += water_val_turn >= water_val_spill ? water_val_turn : water_val_spill

    return data["hydro"]["Hydrogenerators"][res]["water_energy"]
end


using TemporalGPs:
    predict,
    step_marginals,
    step_logpdf,
    step_filter,
    invert_dynamics,
    step_posterior,
    storage_type,
    is_of_storage_type

using Zygote, StaticArrays

println("lgssm:")
@testset "lgssm" begin

    rng = MersenneTwister(123456)

    storages = (
        dense=(name="dense storage Float64", val=ArrayStorage(Float64)),
        static=(name="static storage Float64", val=SArrayStorage(Float64)),
    )
    emission_types = (
        small_output=(name="small output", val=SmallOutputLGC),
        large_output=(name="large output", val=LargeOutputLGC),
        scalar_output=(name="scalar output", val=ScalarOutputLGC),
    )
    settings = [
        (tv=:time_varying, N=1, Dlat=3, Dobs=2, storage=storages.dense),
        (tv=:time_varying, N=49, Dlat=3, Dobs=2, storage=storages.dense),
        (tv=:time_invariant, N=49, Dlat=3, Dobs=2, storage=storages.dense),
        (tv=:time_varying, N=49, Dlat=1, Dobs=1, storage=storages.dense),
        (tv=:time_varying, N=1, Dlat=3, Dobs=2, storage=storages.static),
        (tv=:time_invariant, N=49, Dlat=3, Dobs=2, storage=storages.static),
    ]
    orderings = [
        Forward(),
        Reverse(),
    ]
    Qs = [
        Val(:dense),
        # Val(:diag), diag tests don't work because `FiniteDiffernces.to_vec`.
    ]

    @testset "($tv, $N, $Dlat, $Dobs, $(storage.name), $(emission.name), $order, $Q)" for
        (tv, N, Dlat, Dobs, storage) in settings,
        emission in emission_types,
        order in orderings,
        Q in Qs

        # Print current iteration to prevent CI timing out.
        println(
            "(time_varying=$tv, N=$N, Dlat=$Dlat, Dobs=$Dobs, " *
            "storage=$(storage.name), emissions=$(emission.val), ordering=$order)",
        )

        # Build LGSSM.
        model = if emission.val ∈ (SmallOutputLGC, LargeOutputLGC)
            random_lgssm(rng, order, Val(tv), emission.val, Dlat, Dobs, N, Q, storage.val)
        elseif emission.val == ScalarOutputLGC
            random_lgssm(rng, order, Val(tv), emission.val, Dlat, Dobs, N, storage.val)
        else
            throw(error("Unrecognised storage $(emission.val)"))
        end

        # Verify the correct output types has been obtained.
        @test eltype(model.emissions) <: emission.val

        # Verify that model properties are as requested.
        @test storage_type(model) == storage.val

        @test length(model) == N
        @test getindex(model, N) isa TemporalGPs.ElementOfLGSSM

        y = first(rand(model))
        x = TemporalGPs.x0(model)

        @testset "step_marginals" begin
            @inferred step_marginals(x, model[1])
            adjoint_test(step_marginals, (x, model[1]))
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_marginals, (x, model[1]))
            end
        end
        @testset "step_logpdf" begin
            args = (ordering(model[1]), x, (model[1], y))
            @inferred step_logpdf(args...)
            adjoint_test(step_logpdf, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_logpdf, args)
            end
        end
        @testset "step_filter" begin
            args = (ordering(model[1]), x, (model[1], y))
            @inferred step_filter(args...)
            adjoint_test(step_filter, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_filter, args)
            end
        end
        @testset "invert_dynamics" begin
            args = (x, x, model[1].transition)
            @inferred invert_dynamics(args...)
            adjoint_test(invert_dynamics, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(invert_dynamics, args)
            end
        end
        @testset "step_posterior" begin
            args = (ordering(model[1]), x, (model[1], y))
            @inferred step_posterior(args...)
            adjoint_test(step_posterior, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_posterior, args)
            end
        end

        # Run standard battery of LGSSM tests.
        test_interface(
            rng, model;
            rtol=1e-5,
            atol=1e-5,
            context=NoContext(),
            max_primal_allocs=10,
            max_forward_allocs=35,
            max_backward_allocs=50,
            check_allocs=storage.val isa SArrayStorage,
        )
    end
end


using DelimitedFiles
using LinearAlgebra
using Clustering
using StatsBase
using Distributed
using DataFrames
using CSV
import ClusterManagers
using LinearAlgebra
using Distances
using Random
using Base
using ArgParse

function ParseCommands()

    s = ArgParseSettings()

    @add_arg_table s begin
        "--mutation_file", "-f"
            help = "File with the matrix count"
            required = true
        "--min_signatures_to_extract", "-s"
            help = "Minimum number of signatures to extract"
            arg_type = Int
            required = true
        "--max_signatures_to_extract", "-m"
            help = "Maximum number of signatures to extract"
            arg_type = Int
            required = true
        "--outpath", "-o"
            help = "Folder where the results will be saved"
            required = true
        "--iterations", "-i"
            help = "Total NMF iterations performed"
            arg_type = Int
            default = 1024
        "--workers", "-n"
            help = "Number of slurm parallel process or local cores"
            arg_type = Int
            default = 1
        "--weak", "-w"
            help = "Whether to remove weak mutation types. If unsure, use default"
            arg_type = Int
            default = 0
        "--slurm"
            help = "Whether to run the pipeline using SLURM"
            action = :store_true

    end
    return parse_args(s)
end

parsed_args = ParseCommands()
numberWorkers = parsed_args["workers"]
println("Adding workers...")
if parsed_args["slurm"]
    ClusterManagers.addprocs_slurm(numberWorkers)
else
    addprocs(numberWorkers)
end

# get the path where the script is located
@eval @everywhere path_root=@__DIR__

# add all the necessary scripts
include(string(path_root, "/src/removeWeak.jl"))
include(string(path_root,"/src/addWeak.jl"))
include(string(path_root,"/src/filterOutIterations.jl"))
include(string(path_root,"/src/evaluateStability.jl"))
include(string(path_root,"/src/deconvoluteMutCounts.jl"))
include(string(path_root,"/src/extractSignaturesParallelNoArray.jl"))
include(string(path_root,"/src/removeAllSingleSignatures.jl"))
include(string(path_root,"/src/removeOneSig.jl"))
include(string(path_root,"/src/evalSingleSample.jl"))
include(string(path_root,"/src/stabilities.jl"))

# mutation file
mutation_file = parsed_args["mutation_file"]

# number of signatures to extract
min_numberProcessesToExtract = parsed_args["min_signatures_to_extract"]
max_numberProcessesToExtract = parsed_args["max_signatures_to_extract"]

# total number of iterations
totalIterations = parsed_args["iterations"]

# outpath where they will be saved
outpath = parsed_args["outpath"]

# define whether is a Pan whole-genome, so weak types won't be removed
isWeak = parsed_args["weak"]
remove_weak_muttypes=0.001
if isWeak==0
    remove_weak_muttypes=0.00
end

# add seed so that the extraction is reproducible
Random.seed!(123456)

totalReplicates = 100
removeLastPercentage = 0.07

# Load input
genomesOriginal, sample_names = readdlm(mutation_file, '\t', Int, '\n'; header=true)
sample_header = vec(Symbol.(sample_names))

# Deconvolute mutation counts
genomes, sort_idx, n_types_to_remove = removeWeak(genomesOriginal, remove_weak_muttypes)

# start processing
name_ttype = splitext(basename(mutation_file))[1]

@info "Processing $name_ttype"
@info "Outpath $outpath"

if parsed_args["slurm"]
    @info "Slurm mode"
end

# Bootstrap variables
@eval @everywhere parallelGenomes=$genomes
@eval @everywhere originalGenomes=$genomesOriginal
@eval @everywhere parallelTotalIterations=$totalIterations
@eval @everywhere OutPath=$outpath

@eval @everywhere mutfile = $name_ttype
@eval @everywhere initial_mutfile = $mutation_file
@eval @everywhere SavePath = joinpath(OutPath, mutfile)
mkpath(SavePath)

@everywhere parallelGenomesColsSum = sum(parallelGenomes, dims=1)
@everywhere parallelNormGenomes = parallelGenomes ./ repeat(parallelGenomesColsSum, size(parallelGenomes, 1), 1 )

for numberProcessesToExtract = min_numberProcessesToExtract : max_numberProcessesToExtract

    @eval @everywhere parallelTotalProcesses=$numberProcessesToExtract

    @info "Extracting $numberProcessesToExtract mutational signatures for $totalIterations iterations using $numberWorkers cores"

    #signature extraction
    @time processes, exposures, processStab, processStabAvg, exposures_fitting, avgReconstructionError, avgReconstructionErrorPercentage, avgStability, reconError, reconErrorPercentage, simError = deconvoluteMutCounts(numberProcessesToExtract, totalIterations, sort_idx, n_types_to_remove)

    # save outfiles for each signature extracted
    file = string(SavePath, "/exposures_", parallelTotalProcesses)
    CSV.write(file, DataFrame(exposures), header = sample_header, delim = "\t")

    file = string(SavePath, "/processes_",  parallelTotalProcesses)
    CSV.write(file, DataFrame(processes), delim = "\t")

    file = string(SavePath, "/processesStabAvg_", parallelTotalProcesses)
    processes_stab = convert(DataFrame, processStabAvg)
    CSV.write(file, processes_stab, delim = "\t")

    file = string(SavePath, "/exposures_fitting_", parallelTotalProcesses)
    CSV.write(file, DataFrame(exposures_fitting), header = sample_header, delim = "\t")

    file = string(SavePath, "/reconError_", parallelTotalProcesses)
    CSV.write(file, DataFrame(reconError'), header = sample_header, delim="\t")

    file = string(SavePath, "/reconErrorPercentage_", parallelTotalProcesses)
    CSV.write(file, DataFrame(reconErrorPercentage'), header = sample_header, delim="\t")

    file = string(SavePath, "/simError_", parallelTotalProcesses)
    CSV.write(file, DataFrame(simError'), header = sample_header, delim="\t")

    file = string(SavePath, "/avgReconstructionError_", parallelTotalProcesses)
    writedlm(file, avgReconstructionError,  "\t")

    file = string(SavePath, "/avgReconstructionErrorPercentage_", parallelTotalProcesses)
    writedlm(file, avgReconstructionErrorPercentage,  "\t")

    file = string(SavePath, "/avgStability_", parallelTotalProcesses)
    writedlm(file, avgStability,  "\t")
end

println("Removing workers...")
if numberWorkers>1
    for i in workers()
        rmprocs(i)
    end
end


using DSGE, ModelConstructors, Dates, Test

@testset "sys1: no rule-switch, sys2 parameter-switch in 1st forecast quarter but b/c no rule switch, implied system is the same" begin
    m = Model1002("ss10")
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) # Parameter and/or rule switch in 2nd forecast quarter
    ModelConstructors.set_regime_val!(m[:α], 1, m[:α].value)
    ModelConstructors.set_regime_val!(m[:α], 2, m[:α].value)
    ModelConstructors.set_regime_val!(m[:α], 3, m[:α].value / 2.)
    ModelConstructors.set_regime_val!(m[:α], 4, m[:α].value)

    sys1 = compute_system(m)

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)
    m <= Setting(:gensys_regimes, [1:2])
    m <= Setting(:gensys2_regimes, [2:4])
    sys2 = compute_system(m)

    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] != sys2[3, :TTT]
    @test sys1[4, :TTT] == sys2[4, :TTT]
end

@testset "sys1: no rule-switch, no parameter-switch. sys2: rule- and parameter-switch in 1st forecast quarter (so system 3 should be different)" begin
    m = Model1002("ss10")
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:α], 1, m[:α].value) #2.523)
    ModelConstructors.set_regime_val!(m[:α], 2, m[:α].value)
    ModelConstructors.set_regime_val!(m[:α], 3, 0.00001) #
    ModelConstructors.set_regime_val!(m[:α], 4, m[:α].value)

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)
    m <= Setting(:gensys_regimes, [1:2])
    m <= Setting(:gensys2_regimes, [2:4])

    sys2 = compute_system(m)

    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] != sys2[3, :TTT]
    @test sys1[4, :TTT] == sys2[4, :TTT]
end

@testset "sys1: no rule-switch, no parameter-switch. sys2: rule-switch in 1st forecast quarter, parameter-switch in 2nd forecastquarter" begin
    m = Model1002("ss10")
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:α], 1, m[:α].value) #2.523)
    ModelConstructors.set_regime_val!(m[:α], 2, m[:α].value)
    ModelConstructors.set_regime_val!(m[:α], 3, m[:α].value)
    ModelConstructors.set_regime_val!(m[:α], 4, 0.00001) #

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)
    m <= Setting(:gensys_regimes, [1:2])
    m <= Setting(:gensys2_regimes, [2:4])
    sys2 = compute_system(m)

    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] != sys2[3, :TTT]
    @test sys1[4, :TTT] != sys2[4, :TTT]
end

@testset "sys1: no rule-switch, no parameter-switch. sys2: no rule-switch, yes parameter-switch (in 1st forecast quarter)" begin
    m = Model1002("ss10")
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:α], 1, m[:α].value) #2.523)
    ModelConstructors.set_regime_val!(m[:α], 2, m[:α].value)
    ModelConstructors.set_regime_val!(m[:α], 3, 0.00001) #
    ModelConstructors.set_regime_val!(m[:α], 4, m[:α].value)

    sys2 = compute_system(m)

    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] != sys2[3, :TTT]
    @test sys1[4, :TTT] == sys2[4, :TTT]

end

### Now with QQ
@testset "sys1: rule-switch (1st forecast quarter), no parameter-switch. sys2: rule-switch (1st forecast quarter), no parameter-switch (parameter switch with same values)" begin
    m = Model1002("ss10")

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:σ_g], 1, m[:σ_g].value) #2.523)
    ModelConstructors.set_regime_val!(m[:σ_g], 2, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 3, m[:σ_g].value) #
    ModelConstructors.set_regime_val!(m[:σ_g], 4, m[:σ_g].value)
    sys2 = compute_system(m)

    @test sys1[1, :QQ] == sys2[1, :QQ]
    @test sys1[2, :QQ] == sys2[2, :QQ]
    @test sys1[3, :QQ] == sys2[3, :QQ]
    @test sys1[4, :QQ] == sys2[4, :QQ]
end

@testset "sys1: no rule-switch, no parameter-switch. sys2: rule- and parameter-switch in 1st forecast quarter (so system 3 should be different)" begin
    m = Model1002("ss10")

    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:σ_g], 1, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 2, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 3, 0.00001)
    ModelConstructors.set_regime_val!(m[:σ_g], 4, m[:σ_g].value)

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)

    sys2 = compute_system(m)

    @test sys1[1, :QQ] == sys2[1, :QQ]
    @test sys1[2, :QQ] == sys2[2, :QQ]
    @test sys1[3, :QQ] != sys2[3, :QQ]
    @test sys1[4, :QQ] == sys2[4, :QQ]
end

@testset "sys1: no rule-switch, sys2 parameter-switch in shocks in 1st forecast quarter so only QQ changes" begin
    # rule switching doesn't affect QQ here
    m = Model1002("ss10")

    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:σ_g], 1, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 2, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 3, m[:σ_g].value / 2.)
    ModelConstructors.set_regime_val!(m[:σ_g], 4, m[:σ_g].value)

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)
    m <= Setting(:gensys_regimes, [1:2])
    m <= Setting(:gensys2_regimes, [2:4])
    sys2 = compute_system(m)

    @test sys1[1, :QQ] == sys2[1, :QQ]
    @test sys1[2, :QQ] == sys2[2, :QQ]
    @test sys1[3, :QQ] != sys2[3, :QQ]
    @test sys1[4, :QQ] == sys2[4, :QQ]
    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] ≈ sys2[3, :TTT] # populated by gensys2, so slight numerical differences
    @test sys1[4, :TTT] == sys2[4, :TTT]
end

@testset "sys1: no rule-switch, no parameter-switch. sys2: rule-switch in 1st forecast quarter, parameter-switch in 2nd forecastquarter" begin
    # rule switching doesn't affect QQ here
    m = Model1002("ss10")

    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:σ_g], 1, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 2, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 3, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 4, 0.00001)

    # Use gensys2 for anticipated policy rule change
    m <= Setting(:gensys2, true)
    m <= Setting(:gensys_regimes, [1:2])
    m <= Setting(:gensys2_regimes, [2:4])
    sys2 = compute_system(m)

    @test sys1[1, :QQ] == sys2[1, :QQ]
    @test sys1[2, :QQ] == sys2[2, :QQ]
    @test sys1[3, :QQ] == sys2[3, :QQ]
    @test sys1[4, :QQ] != sys2[4, :QQ]
    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] ≈ sys2[3, :TTT] # populated by gensys2, so slight numerical differences
    @test sys1[4, :TTT] == sys2[4, :TTT]
end

@testset "sys1: no rule-switch, no parameter-switch. sys2: no rule-switch, yes parameter-switch (in 1st forecast quarter)" begin
    m = Model1002("ss10")
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)
    # Number of regimes OVERALL
    m <= Setting(:n_regimes, 4)
    # Number of regimes in forcast
    m <= Setting(:n_fcast_regimes, 2)
    # Number of regimes in history
    m <= Setting(:n_hist_regimes, 2)
    # Dates of all regime switches (including both rule and parameter switches)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30))) #Parameter and/or rule switch in 2nd forecast quarter

    sys1 = compute_system(m)

    ModelConstructors.set_regime_val!(m[:σ_g], 1, m[:σ_g].value) #2.523)
    ModelConstructors.set_regime_val!(m[:σ_g], 2, m[:σ_g].value)
    ModelConstructors.set_regime_val!(m[:σ_g], 3, 0.00001) #
    ModelConstructors.set_regime_val!(m[:σ_g], 4, m[:σ_g].value)

    sys2 = compute_system(m)

    @test sys1[1, :QQ] == sys2[1, :QQ]
    @test sys1[2, :QQ] == sys2[2, :QQ]
    @test sys1[3, :QQ] != sys2[3, :QQ]
    @test sys1[4, :QQ] == sys2[4, :QQ]
    @test sys1[1, :TTT] == sys2[1, :TTT]
    @test sys1[2, :TTT] == sys2[2, :TTT]
    @test sys1[3, :TTT] == sys2[3, :TTT]
    @test sys1[4, :TTT] == sys2[4, :TTT]
end

@testset "Adding redundant regimes does not affect gensys2" begin
    # Add additional check for using replace_eqcond "falsely"
    m = Model1002("ss10")
    # Always need regime switching on, even if you're not switching parameters
    m <= Setting(:regime_switching, true)

    m <= Setting(:date_forecast_start, Date(2020, 9, 30))
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30),
                                     5 => Date(2020, 12, 31),
                                     6 => Date(2021, 3, 31))) #Parameter and/or rule switch in 2nd forecast quarter
    setup_regime_switching_inds!(m)

    m <= Setting(:gensys2, true)
    m <= Setting(:replace_eqcond, true)
    m <= Setting(:regime_eqcond_info, Dict(4 => DSGE.EqcondEntry(DSGE.zlb_rule(), [1., 0.])))
    sys1 = compute_system(m)
    m <= Setting(:regime_dates, Dict(1 => date_presample_start(m), #Start
                                     2 => Date(1990, 3, 31), # Parameter switch in history
                                     3 => Date(2020, 6, 30), # Parameter and/or rule switch in 1st forecast quarter
                                     4 => Date(2020, 9, 30),
                                     5 => Date(2020, 12, 31),
                                     6 => Date(2021, 3, 31),
                                     7 => Date(2021, 6, 30))) #Parameter and/or rule switch in 2nd forecast quarter
    setup_regime_switching_inds!(m)
    sys2 = compute_system(m)

    for i in 1:n_regimes(sys1)
        @test sys1[i, :TTT] ≈ sys2[i, :TTT]
        @test sys1[i, :RRR] ≈ sys2[i, :RRR]
        @test sys1[i, :CCC] ≈ sys2[i, :CCC]
    end
    @test sys2[7, :TTT] ≈ sys1[6, :TTT]
    @test sys2[7, :RRR] ≈ sys1[6, :RRR]
    @test sys2[7, :CCC] ≈ sys1[6, :CCC]
end


using Measures

const CategoricalAesthetic =
    Union{Nothing, IndirectArray}

const NumericalAesthetic =
    Union{Nothing, AbstractMatrix, AbstractVector}

const NumericalOrCategoricalAesthetic =
    Union{CategoricalAesthetic, NumericalAesthetic}

@varset Aesthetics begin
    x,            Union{NumericalOrCategoricalAesthetic, Distribution}
    y,            Union{NumericalOrCategoricalAesthetic, Distribution}
    z,            Union{Nothing, Function, NumericalAesthetic}
    xend,         NumericalAesthetic
    yend,         NumericalAesthetic

    size,         Union{CategoricalAesthetic,Vector,Nothing}
    shape,        Union{CategoricalAesthetic,Vector,Nothing}
    color,        Union{CategoricalAesthetic,Vector,Nothing}
    alpha,        NumericalOrCategoricalAesthetic
    linestyle,    Union{CategoricalAesthetic,Vector,Nothing}

    label,        CategoricalAesthetic
    group,        CategoricalAesthetic

    xmin,         NumericalAesthetic
    xmax,         NumericalAesthetic
    ymin,         NumericalAesthetic
    ymax,         NumericalAesthetic

    # hexagon sizes used for hexbin
    xsize,        NumericalAesthetic
    ysize,        NumericalAesthetic

    # fixed lines
    xintercept,   NumericalAesthetic
    yintercept,   NumericalAesthetic
    intercept,    NumericalAesthetic
    slope,        NumericalAesthetic

    # boxplots
    middle,       NumericalAesthetic
    lower_hinge,  NumericalAesthetic
    upper_hinge,  NumericalAesthetic
    lower_fence,  NumericalAesthetic
    upper_fence,  NumericalAesthetic
    outliers,     NumericalAesthetic
    width,        NumericalAesthetic

    # subplots
    xgroup,       CategoricalAesthetic
    ygroup,       CategoricalAesthetic

    # guides
    xtick,        NumericalAesthetic
    ytick,        NumericalAesthetic
    xgrid,        NumericalAesthetic
    ygrid,        NumericalAesthetic
    color_key_colors,     Maybe(AbstractDict)
    color_key_title,      Maybe(AbstractString)
    color_key_continuous, Maybe(Bool)
    color_function,       Maybe(Function)
    titles,               Maybe(Dict{Symbol, AbstractString})
    shape_key_title,      Maybe(AbstractString)
    size_key_title,       Maybe(AbstractString)
    size_key_vals,        Maybe(AbstractDict)

    # mark some ticks as initially invisible
    xtickvisible,         Maybe(Vector{Bool})
    ytickvisible,         Maybe(Vector{Bool})

    # scale at which ticks should become visible
    xtickscale,           Maybe(Vector{Float64})
    ytickscale,           Maybe(Vector{Float64})

    # plot viewport extents
    xviewmin,     Any
    xviewmax,     Any
    yviewmin,     Any
    yviewmax,     Any

    # labeling functions
    x_label,      Function, showoff
    y_label,      Function, showoff
    xtick_label,  Function, showoff
    ytick_label,  Function, showoff
    color_label,  Function, showoff
    xgroup_label, Function, showoff
    ygroup_label, Function, showoff
    shape_label, Function, showoff
    size_label,   Function, showoff

    # pseudo-aesthetics
    pad_categorical_x, Union{Missing,Bool}, missing
    pad_categorical_y, Union{Missing,Bool}, missing
end

# Calculating fieldnames at runtime is expensive
const valid_aesthetics = fieldnames(Aesthetics)

function show(io::IO, data::Aesthetics)
    maxlen = 0
    print(io, "Aesthetics(")
    for name in valid_aesthetics
        val = getfield(data, name)
        if !ismissing(val) && val != nothing
            print(io, "\n  ", string(name), "=")
            show(io, getfield(data, name))
        end
    end
    print(io, "\n)\n")
end

# Alternate aesthetic names
const aesthetic_aliases =
    Dict{Symbol, Symbol}(:colour        => :color,
                                 :x_min         => :xmin,
                                 :x_max         => :xmax,
                                 :y_min         => :ymin,
                                 :y_max         => :ymax,
                                 :x_group       => :xgroup,
                                 :y_group       => :ygroup,
                                 :x_viewmin     => :xviewmin,
                                 :x_viewmax     => :xviewmax,
                                 :y_viewmin     => :yviewmin,
                                 :y_viewmax     => :yviewmax,
                                 :x_group_label => :xgroup_label,
                                 :y_group_label => :ygroup_label,
                                 :x_tick        => :xtick,
                                 :y_tick        => :ytick,
                                 :x_grid        => :xgrid,
                                 :y_grid        => :ygrid)

# Index as if this were a data frame
getindex(aes::Aesthetics, i::Integer, j::AbstractString) = getfield(aes, Symbol(j))[i]

# Return the set of variables that are non-nothing.
function defined_aesthetics(aes::Aesthetics)
    vars = Set{Symbol}()
    for name in valid_aesthetics
        getfield(aes, name) === nothing || push!(vars, name)
    end
    vars
end

# Checking aesthetics and giving reasonable error messages.

# Raise an error if any of the given aesthetics are not defined.
#
# Args:
#   who: A string naming the caller which is printed in the error message.
#   aes: An Aesthetics object.
#   vars: Symbol that must be defined in the aesthetics.
#
undefined_aesthetics(aes::Aesthetics, vars::Symbol...) =
        setdiff(Set(vars), defined_aesthetics(aes))

function assert_aesthetics_defined(who::AbstractString, aes::Aesthetics, vars::Symbol...)
    undefined_vars = undefined_aesthetics(aes, vars...)
    isempty(undefined_vars) || error("The following aesthetics are required by ",who,
            " but are not defined: ", join(undefined_vars,", "),"\n")
end

function assert_aesthetics_undefined(who::AbstractString, aes::Aesthetics, vars::Symbol...)
    defined_vars = intersect(Set(vars), defined_aesthetics(aes))
    isempty(defined_vars) || error("The following aesthetics are defined but incompatible with ",
            who,": ",join(defined_vars,", "),"\n")
end

function assert_aesthetics_equal_length(who::AbstractString, aes::Aesthetics, vars::Symbol...)
    defined_vars = filter(var -> !(getfield(aes, var) === nothing), [vars...])

    if !isempty(defined_vars)
        n = length(getfield(aes, first(defined_vars)))
        for var in defined_vars
            length(getfield(aes, var)) != n && error(
                    "The following aesthetics are required by ",who,
                    " to be of equal length: ",join(defined_vars,", "),"\n")
        end
    end
    nothing
end

# Replace values in a with non-nothing values in b.
#
# Args:
#   a: Destination.
#   b: Source.
#
# Returns: nothing
#
# Modifies: a
#
function update!(a::Aesthetics, b::Aesthetics)
    for name in valid_aesthetics
        issomething(getfield(b, name)) && setfield(a, name, getfield(b, name))
    end
    nothing
end

# Serialize aesthetics to JSON.

# Args:
#  a: aesthetics to serialize.
#
# Returns:
#   JSON data as a string.
#
json(a::Aesthetics) = join([string(a, ":", json(getfield(a, var))) for var in aes_vars], ",\n")

# Concatenate aesthetics.
#
# A new Aesthetics instance is produced with data vectors in each of the given
# Aesthetics concatenated, nothing being treated as an empty vector.
#
# Args:
#   aess: One or more aesthetics.
#
# Returns:
#   A new Aesthetics instance with vectors concatenated.
#
function concat(aess::Aesthetics...)
    cataes = Aesthetics()
    for aes in aess
        for var in valid_aesthetics
            if var in [:xviewmin, :yviewmin]
                mu, mv = getfield(cataes, var), getfield(aes, var)
                setfield!(cataes, var,
                          mu === nothing ? mv :
                             mv == nothing ? mu :
                                 min(mu, mv))
            elseif var in [:xviewmax, :yviewmax]
                mu, mv = getfield(cataes, var), getfield(aes, var)
                setfield!(cataes, var,
                          mu === nothing ? mv :
                             mv == nothing ? mu :
                                 max(mu, mv))
            else
                setfield!(cataes, var,
                          cat_aes_var!(getfield(cataes, var), getfield(aes, var)))
            end
        end
    end
    cataes
end

cat_aes_var!(a::(Nothing), b::(Nothing)) = a
cat_aes_var!(a::(Nothing), b::Union{Function,AbstractString}) = b
cat_aes_var!(a::(Nothing), b) = copy(b)
cat_aes_var!(a, b::(Nothing)) = a
cat_aes_var!(a::Function, b::Function) = a === string || a == showoff ? b : a

function cat_aes_var!(a::Dict, b::Dict)
    merge!(a, b)
    a
end

cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{T}) where {T <: Base.Callable} = append!(a, b)
cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T <: Base.Callable, U <: Base.Callable} =
        a = [a...,b...]

# Let arrays of numbers clobber arrays of functions. This is slightly odd
# behavior, comes up with with function statistics applied on a layer-wise
# basis.
cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T <: Base.Callable, U} = b
cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T, U <: Base.Callable} = a
cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{T}) where {T} = append!(a, b)
cat_aes_var!(a, b) = a

function cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T, U}
    V = promote_type(T, U)
    ab = Array{V}(undef, length(a) + length(b))
    i = 1
    for x in a
        ab[i] = x
        i += 1
    end
    for x in b
        ab[i] = x
        i += 1
    end

    return ab
end

function cat_aes_var!(xs::IndirectArray{T,1}, ys::IndirectArray{S,1}) where {T, S}
    TS = promote_type(T, S)
    return append!(IndirectArray(xs.index, convert(Array{TS},xs.values)),
                   IndirectArray(ys.index, convert(Array{TS},ys.values)))
end

cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T<:Measure, U<:Measure} = [a..., b...]

cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T<:Measure, U} =
    isabsolute(T) ? [a..., b...] : b

cat_aes_var!(a::AbstractVector{T}, b::AbstractVector{U}) where {T, U<:Measure} =
    isabsolute(U) ? a : [a..., b...]

# Summarizing aesthetics

# Produce a matrix of Aesthetic or Data objects partitioning the original
# Aesthetics or Data object by the cartesian product of xgroup and ygroup.
#
# This is useful primarily for drawing facets and subplots.
#
# Args:
#   aes: Aesthetics or Data objects to partition.
#
# Returns:
#   A Array{Aesthetics} of size max(1, length(xgroup)) by
#   max(1, length(ygroup))
#
function by_xy_group(aes::T, xgroup, ygroup,
                     num_xgroups, num_ygroups) where T <: Union{Data, Aesthetics}
    @assert xgroup === nothing || ygroup === nothing || length(xgroup) == length(ygroup)

    n = num_ygroups
    m = num_xgroups

    xrefs = xgroup === nothing ? [1] : xgroup
    yrefs = ygroup === nothing ? [1] : ygroup

    aes_grid = Array{T}(undef, n, m)
    staging = Array{AbstractArray}(undef, n, m)
    for i in 1:n, j in 1:m
        aes_grid[i, j] = T()
    end

    xgroup === nothing && ygroup === nothing && return aes_grid

    function make_pooled_array(::Type{IndirectArray{T,N,A,V}}, arr::AbstractArray) where {T,N,A,V}
        uarr = unique(arr)
        return IndirectArray(A(indexin(arr, uarr)), V(uarr))
    end
    make_pooled_array(::Type{IndirectArray{T,R,N,RA}},
            arr::IndirectArray{T,R,N,RA}) where {T,R,N,RA} = arr

    for var in fieldnames(T)
        # Skipped aesthetics. Don't try to partition aesthetics for which it
        # makes no sense to pass on to subplots.
        if var == :xgroup || var == :ygroup||
           var == :xtick || var == :ytick ||
           var == :xgrid || var == :ygrid ||
           var == :x_viewmin || var == :y_viewmin ||
           var == :x_viewmax || var == :y_viewmax ||
           var == :color_key_colors
            continue
        end

        vals = getfield(aes, var)
        if isa(vals, AbstractArray) && length(vals)>1
            if xgroup !== nothing && length(vals) !== length(xgroup) ||
               ygroup !== nothing && length(vals) !== length(ygroup)
                continue
            end

            for i in 1:n, j in 1:m
                staging[i, j] = similar(vals, 0)
            end

            for (i, j, v) in zip(cycle(yrefs), cycle(xrefs), vals)
                push!(staging[i, j], v)
            end

            for i in 1:n, j in 1:m
                if typeof(vals) <: IndirectArray
                    setfield!(aes_grid[i, j], var,
                              make_pooled_array(typeof(vals), staging[i, j]))
                else
                    if !applicable(convert, typeof(vals), staging[i, j])
                        T2 = eltype(vals)
                        if T2 <: Color T2 = Color end
                        da = Array{Union{Missing,T2}}(undef, length(staging[i, j]))
                        copy!(da, staging[i, j])
                        setfield!(aes_grid[i, j], var, da)
                    else
                        setfield!(aes_grid[i, j], var,
                                  convert(typeof(vals), copy(staging[i, j])))
                    end
                end
            end
        else
            for i in 1:n, j in 1:m
                setfield!(aes_grid[i, j], var, vals)
            end
        end
    end

    aes_grid
end

function inherit(a::Aesthetics, b::Aesthetics;
                 clobber=[])
    acopy = copy(a)
    inherit!(acopy, b, clobber=clobber)
    return acopy
end

function inherit!(a::Aesthetics, b::Aesthetics;
                  clobber=[])
    clobber_set = Set{Symbol}(clobber)
    for field in valid_aesthetics
        aval = getfield(a, field)
        bval = getfield(b, field)
        if field in clobber_set
            setfield!(a, field, bval)
        elseif aval === missing || aval === nothing || aval === string || aval == showoff
            setfield!(a, field, bval)
        elseif field == :xviewmin || field == :yviewmin
            bval != nothing && (aval == nothing || aval > bval) && setfield!(a, field, bval)
        elseif field == :xviewmax || field == :yviewmax
            bval != nothing && (aval == nothing || aval < bval) && setfield!(a, field, bval)
        elseif typeof(aval) <: Dict && typeof(bval) <: Dict
            merge!(aval, getfield(b, field))
        end
    end
    nothing
end


@info "Train baseline lstm language model using enwik8 dataset..."

using Knet

include("../src/data.jl")
include("../src/baseline.jl")
include("../src/train.jl")

datadir = "../data/enwik8"
jld2dir = "../jld2/enwik8.jld2"
BATCHSIZE = 32

if !isfile(jld2dir)
    println("Reading data from directory: $datadir")
    println("Setting batch size to $BATCHSIZE")
    vocab = Vocab("$datadir/train.txt")
    trainfile = TextReader("$datadir/train.txt", vocab)
    validfile = TextReader("$datadir/valid.txt", vocab)
    testfile = TextReader("$datadir/test.txt", vocab)
    dtrn = TextData(trainfile, batchsize=BATCHSIZE)
    ddev = TextData(validfile, batchsize=BATCHSIZE)
    dtst = TextData(testfile, batchsize=BATCHSIZE)
    println("Saving data from $jld2dir")
    Knet.save(jld2dir, "dtrn", dtrn, "dtst", dtst, "ddev", ddev)
else
    println("Loading data from $jld2dir")
    (dtrn, dtst, ddev) = Knet.load(jld2dir, "dtrn", "dtst", "ddev")
    vocab = dtrn.src.vocab
    if dtrn.batchsize != BATCHSIZE
        changebatchsize!(dtrn, BATCHSIZE)
        changebatchsize!(ddev, BATCHSIZE)
        changebatchsize!(dtst, BATCHSIZE)
    end;
end;

println()
@info "Initializing the model and collecting training data..."
epochs, em_size, hidden_size, layers = 30, 256, 256, 2;

println("embedding size: ", em_size);
println("hidden size: ", hidden_size);
println("layers: ", layers);
println("epochs: ", epochs)
ctrn = collect(dtrn);
trn = collect(flatten(shuffle!(ctrn) for i in 1:epochs));
trnmini = ctrn[1:20];
dev = collect(ddev);
model = SimpleLSTMModel(em_size, hidden_size, vocab; layers=layers, dropout=0.2);
# initopt!(model, length(trn))

println()
@info "Starting training, total iteration no: $(length(trn))"
bestmodel = trainadam!(model, trn, dev)

# println()
# @info "Finished training, Starting evaluation ..."
# trnloss = loss(bestmodel, dtrn);
# println("Training set scores:       ", report(trnloss))
devloss = loss(bestmodel, ddev);
println("Development set scores:    ", report(devloss))
testloss = loss(bestmodel, dtst);
println("Test set scores:           ", report(testloss))

@info "Generate text using the trained model"
print(generate(bestmodel, start="United Nations ", maxlength=1024))

# @info "Saving the model as model.jld2"
# Knet.save("model.jld2", "model", model); # uncomment for saving


using Base.Test
using DataFrames
import WCS

const rcf = RunCamcolField(3900, 6, 269)

"""
Get the PSF located at a particular world location in an image.

Args:
 - world_loc: A location in world coordinates.
 - img: An TiledImage

Returns:
 - An array of PsfComponent objects that represents the PSF as a mixture
     of Gaussians.
"""
function get_source_psf(world_loc::Vector{Float64}, img::TiledImage)
    # Some stamps or simulated data have no raw psf information.    In that case,
    # just use the psf from the image.
    if size(img.raw_psf_comp.rrows) == (0, 0)
        # Also return a vector of empty psf params
        return img.psf, fill(fill(NaN, length(PsfParams)), psf_K)
    else
        pixel_loc = WCS.world_to_pix(img.wcs, world_loc)
        psfstamp = img.raw_psf_comp(pixel_loc[1], pixel_loc[2])
        return PSF.fit_raw_psf_for_celeste(psfstamp)
    end
end

function test_blob()
    # A lot of tests are in a single function to avoid having to reload
    # the full image multiple times.
    blob = SDSSIO.load_field_images(rcf, datadir)

    dir = "$datadir/$(rcf.run)/$(rcf.camcol)/$(rcf.field)"
    fname = @sprintf "%s/photoObj-%06d-%d-%04d.fits" dir rcf.run rcf.camcol rcf.field
    cat_entries = SDSSIO.read_photoobj_celeste(fname)

    ea = make_elbo_args(blob, cat_entries,
                       patch_radius_pix=1e-6, fit_psf=false, tile_width=20)

    # Just check some basic facts about the catalog.
    @test length(cat_entries) == 805
    @test 0 < sum([ce.is_star for ce in cat_entries]) < 805

    # Find an object near the middle of the image.
    ctr = WCS.pix_to_world(blob[3].wcs, [blob[3].H / 2, blob[3].W / 2])
    dist = [sqrt((ce.pos[1] - ctr[1])^2 + (ce.pos[2] - ctr[2])^2)
                    for ce in cat_entries]
    obj_index = findmin(dist)[2]    # index of closest object
    obj_loc = cat_entries[obj_index].pos    # location of closest object

    # Test get_source_psf at point while we have the blob loaded.
    test_b = 3
    img = ea.images[test_b]
    ea_obj = make_elbo_args(blob, cat_entries[obj_index:obj_index])
    pixel_loc = WCS.world_to_pix(img.wcs, obj_loc);
    original_psf_val = img.raw_psf_comp(pixel_loc[1], pixel_loc[2])

    original_psf_celeste = PSF.fit_raw_psf_for_celeste(original_psf_val)[1];
    fit_original_psf_val = PSF.get_psf_at_point(original_psf_celeste);

    obj_psf = get_source_psf(ea_obj.vp[1][ids.u], img)[1];
    obj_psf_val = PSF.get_psf_at_point(obj_psf);

    # The fits should match exactly.
    @test_approx_eq_eps(obj_psf_val, fit_original_psf_val, 1e-6)

    # The raw psf will not be as good.
    @test_approx_eq_eps(obj_psf_val, original_psf_val, 1e-2)

    cat_several = [cat_entries[1]; cat_entries[obj_index]]
    ea_several = make_elbo_args(ea_obj.images, cat_several)

    # The second set of vp is the object of interest
    point_patch_psf = PSF.get_psf_at_point(ea_several.patches[2, test_b].psf);
    # The threshold for the test below was formerly 1e-6.
    # Is it a problem I needed to increase it?
    @test_approx_eq_eps(obj_psf_val, point_patch_psf, 1e-4)
end

function test_stamp_get_object_psf()
    stamp_blob, stamp_mp, body = gen_sample_star_dataset();
    img = TiledImage(stamp_blob[3]);
    obj_index =    stamp_mp.vp[1][ids.u]
    pixel_loc = WCS.world_to_pix(img.wcs, obj_index)
    original_psf_val = PSF.get_psf_at_point(img.psf);

    obj_psf_val =
        PSF.get_psf_at_point(get_source_psf(stamp_mp.vp[1][ids.u], img)[1])
    @test_approx_eq_eps(obj_psf_val, original_psf_val, 1e-6)
end

function test_get_tiled_image_source()
    # Test that an object only occurs the appropriate tile's local sources.
    blob, ea, body = gen_sample_star_dataset();

    ea = make_elbo_args(blob, body; patch_radius_pix=1e-6)

    tiled_img = TiledImage(blob[3], tile_width=10);
    for hh in 1:size(tiled_img.tiles, 1), ww in 1:size(tiled_img.tiles, 2)
        tile = tiled_img.tiles[hh, ww]
        loc = Float64[mean(tile.h_range), mean(tile.w_range)]
        for b = 1:5
            ea.vp[1][ids.u] = loc
            pixel_center = WCS.world_to_pix(blob[b].wcs, loc)
            wcs_jacobian = Model.pixel_world_jacobian(blob[b].wcs, pixel_center)
            radius_pix = maxabs(eigvals(wcs_jacobian)) * 1e-6
            ea.patches[1, b] = SkyPatch(loc,
                                        radius_pix,
                                        blob[b].psf,
                                        wcs_jacobian,
                                        pixel_center)
        end
        patches = vec(ea.patches[:, 3])
        local_sources = Model.get_sources_per_tile(tiled_img.tiles,
                                                   Model.patch_ctrs_pix(patches),
                                                   Model.patch_radii_pix(patches))
        @test local_sources[hh, ww] == Int[1]
        for hh2 in 1:size(tiled_img.tiles, 1), ww2 in 1:size(tiled_img.tiles, 2)
            if (hh2 != hh) || (ww2 != ww)
                @test local_sources[hh2, ww2] == Int[]
            end
        end
    end
end

function test_local_source_candidate()
    blob, ea, body = gen_n_body_dataset(100);

    for b=1:ea.N
        # Get the sources by iterating over everything.
        patches = vec(ea.patches[:,b])

        tile_source_map = Model.get_sources_per_tile(ea.images[b].tiles,
                        Model.patch_ctrs_pix(patches),
                        Model.patch_radii_pix(patches))

        # Check that all the actual sources are candidates and that this is the
        # same as what is returned by make_elbo_args
        HH, WW = size(tile_source_map)
        for h=1:HH, w=1:WW
            # Get a set of candidates.
            candidates = Model.local_source_candidates(
                                                ea.images[b].tiles[h, w],
                                                Model.patch_ctrs_pix(patches),
                                                Model.patch_radii_pix(patches))
            @test setdiff(tile_source_map[h, w], candidates) == []
            @test tile_source_map[h, w] == ea.tile_source_map[b][h, w]
        end
    end
end

function test_set_patch_size()
    # Test that the patch size gets most of the light from a variety of
    # galaxy shapes.
    # This shows that the current patch size is actually far too conservative.

    function gal_catalog_from_scale(gal_scale::Float64, flux_scale::Float64)
        CatalogEntry[CatalogEntry(world_location, false,
                                   flux_scale * fluxes, flux_scale * fluxes,
                                   0.1, .01, pi/4, gal_scale, "sample", 0) ]
    end

    srand(1)
    blob0 = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf");
    img_size = 150
    for b in 1:5
            blob0[b].H, blob0[b].W = img_size, img_size
    end
    fluxes = [4.451805E+03,1.491065E+03,2.264545E+03,2.027004E+03,1.846822E+04]

    world_location = WCS.pix_to_world(blob0[3].wcs,
                                      Float64[img_size / 2, img_size / 2])

    for gal_scale in [1.0, 10.0], flux_scale in [0.1, 10.0]
        cat = gal_catalog_from_scale(gal_scale, flux_scale);
        blob = Synthetic.gen_blob(blob0, cat);
        ea = make_elbo_args(blob, cat, tile_width=typemax(Int));

        for b=1:length(blob)
            @assert size(ea.images[b].tiles) == (1, 1)
            tile_image = DeterministicVI.tile_predicted_image(
                ea.images[b].tiles[1,1], ea, ea.tile_source_map[b][1,1]);

            pixel_center = WCS.world_to_pix(blob[b].wcs, cat[1].pos)
            radius = Model.choose_patch_radius(
                pixel_center, cat[1], blob[b].psf, ea.images[b])

            circle_pts = fill(false, blob[b].H, blob[b].W);
            in_circle = 0.0
            for x=1:size(tile_image)[1], y=1:size(tile_image)[2]
                if ((x - pixel_center[1]) ^ 2 + (y - pixel_center[2]) ^ 2) < radius ^ 2
                    in_circle += tile_image[x, y]
                    circle_pts[x, y] = true
                end
            end
            @test in_circle / sum(tile_image) > 0.95

            # Convenient for visualizing:
            # using PyPlot
            # in_circle / sum(tile_image)
            # imshow(tile_image); colorbar()
            # imshow(circle_pts, alpha=0.4)
        end
    end
end

test_blob()
test_stamp_get_object_psf()
test_get_tiled_image_source()
test_local_source_candidate()
test_set_patch_size()


module quickPIV

using FFTW, FileIO, LIBTIFF, Statistics

export PIV

# Shortened type names to keep function signatures short ( when possible one-liners )
include("units.jl")

F      = Union{Float64,Float32} # considering 64 and 32 bit systems
I      = Union{Int64,Int32}
A      = Array
S      = String
C      = Complex
II     = Tuple{I,I}
III    = Tuple{I,I,I}
UnI    = Union{I,II,III}
SySySy = Tuple{Symbol,Symbol,Symbol}
DU     = DistanceUnits
TU     = TimeUnits
UNTS   = Tuple{DU,DU,DU,TU}

Base.convert( ::Type{Union{Int32,Int64}}, a::AbstractFloat ) = round( typeof(0), a )

struct ZNCC   end
struct FFT    end
struct NSQECC end

CORRTYPES = Union{ZNCC,FFT,NSQECC};

include("parameters.jl")
include("pivIO.jl")
include("utils.jl")

include("PIV/subpixel.jl")
include("PIV/signalToNoise.jl")
include("PIV/ccrZNCC.jl")
include("PIV/ccrFFT.jl")
include("PIV/ccrNSQECC.jl")

include("syntheticParticles.jl")

include("postprocessing.jl")

function PIV( data1::Array{T,N}, data2::Array{T,N}, params::PIVParameters;
              checkParams=true, units=nothing ) where {T<:AbstractFloat,N}

	checkParams && checkPIVParameters( size(data1), params )

	return _PIV( data1, data2, params )
end

""" Calling 2D PIV analyses """

function _PIV( img1::A{T,2}, img2::A{T,2}, p::PIVParameters; units=nothing
             ) where {T<:AbstractFloat}

	u, v, sn = PIV_2D( p.corr, img1, img2, p.args... );

	if units !== nothing
	    ratios = changeUnits.(  p.units[1:2], units[1:2] )./changeUnits( p.units, units[4]  )
	    for e in 1:length(u)
	        u[e] = u[e]*ratios[1]
	        v[e] = v[e]*ratios[2]
	    end
	end
	return u, v, sn;
end

""" Calling 3D PIV analyses """

function _PIV( vol1::A{T,3}, vol2::A{T,3}, p::PIVParameters; units=nothing
             ) where {T<:AbstractFloat}

	u, v, w, sn = PIV_3D( p.corr, vol1, vol2, p.args... );

	if units !== nothing
	    ratios = changeUnits.( p.units, units[1:3] )./changeUnits( p.units, units[4] )
	    for e in 1:length(u)
	        u[e] = u[e]*ratios[1]
	        v[e] = v[e]*ratios[2]
	        w[e] = w[e]*ratios[3]
	    end
	end

	return u, v, w, sn;
end

end


module SecureIO

using Nettle

struct SecureSocket <: IO
    socket::IO
    size::Union{Integer,Nothing}
    ch::Channel{UInt8}
    enc::Encryptor
    dec::Decryptor
end

function SecureSocket(socket::IO,key;size=nothing)
    size!=nothing && @assert mod(size,16)==0

    key32 = hexdigest("sha256", "$key")[1:32]
    enc = Encryptor("AES256", key32)
    dec = Decryptor("AES256", key32)

    SecureSocket(socket,size,Channel{UInt8}(Inf),enc,dec)
end

import Base.isopen
isopen(s::SecureSocket) = isopen(s.socket)

import Base.close
close(s::SecureSocket) = close(s.socket)

function addpadding(text::Vector{UInt8},size)
    # Only last two bytes are used to encode the padding boundary. That limits the possible size.
    @assert length(text) + 2 <= size <= 2^16

    endbytes = reinterpret(UInt8, Int16[length(text)])
    paddedtext = [text; UInt8[0 for i in 1:(size - length(text) - 2)]; endbytes]

    return paddedtext
end

function trimpadding(text::Vector{UInt8})
    endbytes = text[end-1:end]
    n = reinterpret(Int16,endbytes)[1]
    return text[1:n]
end

function stack(io::IO,msg::Vector{UInt8})
    frontbytes = reinterpret(UInt8,Int16[length(msg)])
    item = UInt8[frontbytes...,msg...]
    write(io,item)
end

function unstack(io::IO)
    sizebytes = [read(io,UInt8),read(io,UInt8)]
    size = reinterpret(Int16,sizebytes)[1]

    msg = UInt8[]
    for i in 1:size
        push!(msg,read(io,UInt8))
    end
    return msg
end

function unstack(io::IOBuffer)
    bytes = take!(io)
    size = reinterpret(Int16,bytes[1:2])[1]
    msg = bytes[3:size+2]
    if length(bytes)>size+2
        write(io,bytes[size+3:end])
    end
    return msg
end

import Base.write
function _write_fixed(s::SecureSocket,plaintext::Vector{UInt8},size)
    if size - 2 < length(plaintext)
        error("Message with length $(length(plaintext)) does not fit in $size - 2 bytes")
    end

    #paddedtext = add_padding_PKCS5(Vector{UInt8}(plaintext), size)
    paddedtext = addpadding(plaintext, size)

    msgenc = encrypt(s.enc,paddedtext)
    stack(s.socket,msgenc)
end

function _write_elastic(s::SecureSocket,plaintext::Vector{UInt8})
    n = length(plaintext)

    if mod(n+2,16)==0
        size = n+2
    else
        size = (div(n+2,16) + 1)*16
    end

    paddedtext = addpadding(plaintext, size)
    msgenc = encrypt(s.enc,paddedtext)
    stack(s.socket,msgenc)
end

function _write(s::SecureSocket,msg::Vector{UInt8})
    if s.size==nothing
        _write_elastic(s,msg)
    else
        _write_fixed(s,msg,s.size)
    end
end

_write(s::SecureSocket,msg::UInt8) = _write(s,UInt8[msg])
_write(s::SecureSocket,msg::String) = _write(s,Vector{UInt8}(msg))

write(s::SecureSocket,msg::UInt8) = _write(s,msg)
write(s::SecureSocket,msg::String) = _write(s,msg)
write(s::SecureSocket,msg::Vector{UInt8}) = _write(s,msg)

function getbytes(s::SecureSocket)
    ciphertext = unstack(s.socket) ### If s.socket is a buffer.

    deciphertext = decrypt(s.dec,ciphertext)

    #str = trim_padding_PKCS5(deciphertext)
    msg = trimpadding(deciphertext)

    if typeof(msg)==UInt8
        return msg
    else
        return take!(IOBuffer(msg))
    end
end

import Base.read
function read(s::SecureSocket,x::Type{UInt8})
    if isready(s.ch)==false
        bytes = getbytes(s)
        if typeof(bytes)==UInt8
            put!(s.ch,bytes)
        else
            for byte in bytes
                put!(s.ch,byte)
            end
        end
    end

    return take!(s.ch)
end

export SecureSocket, read, write

end # module


"""
    function Base.show(
        io::IO,
        p::Player
    )

Prints out a player.

"""

function Base.show(io::IO, p::Player)
    return print(io, "Player: $(p.id)")
end

"""
    function strength(
        player::Player
    )::Float64

Calculates a player's strength

Arguments
=========
- `player`:
    the player for whose strength to calculate.
"""

function strength(
    player::Player
)::Float64
    return 1 / player.rankings["Ground Truth"]
end

"""
    function participationProbability(;
        player::Player,
        tournament::Tournament,
        ranking_scheme::String,
        players::Vector{Player}
    )::Float64

Calculates the probability a player plays a tournament.

Arguments
=========
- `player`:
    the player under consideration.
- `tournament`:
    the tournament to calculate a player's strength for.
- `ranking_scheme`:
    the player's current ranking to use to determine participation probability.
- `players`:
    all the players in the simulation.
"""

function participationProbability(;
    player::Player,
    tournament::Tournament,
    ranking_scheme::String,
    players::Vector{Player}
)::Float64
    player_rank = player.rankings[ranking_scheme]
    N = length(players)
    n = length(filter(p -> p.rankings[ranking_scheme] < player_rank, players))
    ordinal_rank = N - n
    if ordinal_rank <= 50
        if tournament.type == "Grandslam"
            return 0.37
        elseif tournament.type == "ATPNitto"
            return 0.95
        elseif tournament.type == "ATP1000"
            return 0.62
        elseif tournament.type == "ATP500"
            return 0.53
        elseif tournament.type == "ATP250"
            return 0.29
        else
            return 0.05
        end
    elseif ordinal_rank <= 100
        if tournament.type == "Grandslam"
            return 0.27
        elseif tournament.type == "ATPNitto"
            return 0.90
        elseif tournament.type == "ATP1000"
            return 0.09
        elseif tournament.type == "ATP500"
            return 0.16
        elseif tournament.type == "ATP250"
            return 0.35
        else
            return 0.1
        end
    elseif ordinal_rank <= 200
        if tournament.type == "Grandslam"
            return 0.14
        elseif tournament.type == "ATPNitto"
            return 0.7
        elseif tournament.type == "ATP1000"
            return 0.09
        elseif tournament.type == "ATP500"
            return 0.14
        elseif tournament.type == "ATP250"
            return 0.14
        else
            return 0.15
        end
    elseif ordinal_rank <= 300
        if tournament.type == "Grandslam"
            return 0.08
        elseif tournament.type == "ATPNitto"
            return 0.6
        elseif tournament.type == "ATP1000"
            return 0.07
        elseif tournament.type == "ATP500"
            return 0.06
        elseif tournament.type == "ATP250"
            return 0.06
        else
            return 0.20
        end
    elseif ordinal_rank <= 400
        if tournament.type == "Grandslam"
            return 0.04
        elseif tournament.type == "ATPNitto"
            return 0.5
        elseif tournament.type == "ATP1000"
            return 0.05
        elseif tournament.type == "ATP500"
            return 0.05
        elseif tournament.type == "ATP250"
            return 0.05
        else
            return 0.30
        end
    elseif ordinal_rank <= 500
        if tournament.type == "Grandslam"
            return 0.05
        elseif tournament.type == "ATPNitto"
            return 0.5
        elseif tournament.type == "ATP1000"
            return 0.05
        elseif tournament.type == "ATP500"
            return 0.05
        elseif tournament.type == "ATP250"
            return 0.05
        else
            return 0.40
        end
    elseif ordinal_rank <= 600
        if tournament.type == "Grandslam"
            return 0.05
        elseif tournament.type == "ATPNitto"
            return 0.5
        elseif tournament.type == "ATP1000"
            return 0.05
        elseif tournament.type == "ATP500"
            return 0.05
        elseif tournament.type == "ATP250"
            return 0.05
        else
            return 0.50
        end
    else
        if tournament.type == "Grandslam"
            return 0.05
        elseif tournament.type == "ATPNitto"
            return 0.4
        elseif tournament.type == "ATP1000"
            return 0.05
        elseif tournament.type == "ATP500"
            return 0.05
        elseif tournament.type == "ATP250"
            return 0.05
        else
            return 0.6
        end
    end
end

"""
    function clear_doubles_results(;
        players::Vector{Player}
    )

Clears the doubles results stored within a player struct.

Arguments
=========
- `players`:
    an array of players.
"""

function clear_doubles_results(;
    players::Vector{Player}
)
    for player in players
        player.doubles_results = Dict{String,DoublesTournamentPlayerSummary}()
        player.matches_played = 0
        player.tournaments_played = 0
    end
    return nothing
end

"""
    function initialize_rankings_schemes(;
        players::Vector{Player},
        ranking_schemes::Vector{String},
        verbose::Bool
    )

Initializes the various ranking schemes with the burn in ranking.

Arguments
=========
- `players`:
    an array of players.
- `ranking_schemes`
    an array of the ranking schemes we want to initialize.
"""

function initialize_rankings_schemes(;
    players::Vector{Player},
    ranking_schemes::Vector{String},
    verbose::Bool
)
    for player in players
        for ranking_scheme in ranking_schemes
            if ranking_scheme != "Ground Truth"
                player.rankings[ranking_scheme] = player.rankings["Burn In"]
            end
        end
    end
    if verbose
        println("########")
        println("Ranking schemes initialized.")
        println("########")
        println()
    end
    return nothing
end

"""
    function update_player_ages(;
        players::Vector{Player},
        method::String,
        verbose::Bool
    )

Updates the age for each player in the simulation.

Arguments
=========
- `players`:
    an array of players.
- `method`
    how to initialize the rankings
    options:
    1) Random
    2) Increment
- `verbose`:
    conditional on whether to print results.
"""

function update_player_ages(;
    players::Vector{Player},
    method::String,
    verbose::Bool
)
    if method == "Random"
        for player in players
            player.age = rand(18:40)
        end
    elseif method == "Increment"
        for player in players
            player.age += 1
        end
    end
    if verbose
        println("########")
        if method == "Random"
            println("Player ages randomly initialized.")
        elseif method == "Increment"
            println("Player ages incremented.")
        end
        println("########")
        println()
    end
    return nothing
end


using Luxor
@svg begin
    setopacity(0.85)
    steps = 10
    gap   = 12
    for (n, θ) in enumerate(range(π/4, step=(2π-π/2)/steps, length=steps))
        sethue([Luxor.julia_green,
            Luxor.julia_red,
            Luxor.julia_purple,
            Luxor.julia_blue][mod1(n, 4)])
        sector(Point(-150, 0), 50, 100, θ, θ + 2π/steps - deg2rad(gap), :fill)
    end
     fontsize(90)
     fontface("MyanmarMN-Bold")
     sethue("black")
     text("oefplots.jl", Point(-178, -10) ,halign=:left,valign = :middle)
end 700 200 "logo.svg"


function loadNeuralNet()
    path(s) = normpath("$(@__DIR__)/"*s)
    @load path("../models/agz_128_base.bson") bn
    @load path("../models/agz_128_value.bson") value
    @load path("../models/agz_128_policy.bson") policy
    @load path("../models/weights/agz_128_base.bson") bn_weights
    @load path("../models/weights/agz_128_value.bson") val_weights
    @load path("../models/weights/agz_128_policy.bson") pol_weights
    Flux.loadparams!(bn,bn_weights)
    Flux.loadparams!(value, val_weights)
    Flux.loadparams!(policy, pol_weights)
    # bn = mapleaves(cu, bn)
    # value = mapleaves(cu, value)
    # policy = mapleaves(cu, policy)
    NeuralNet(base_net = bn, value = value, policy = policy)
end

function playMove()
    move = AlphaGo.suggest_move(agz)
    println(go.to_kgs(move))
    play_move!(agz, move)
    return move
end

function playMove(x)
    move = actionToString(x)
    @show move
    play_move!(agz, go.from_kgs(move))
    return move
end

function actionToString(a)
    a["y"] == -1 && a["x"] == -1 && return "pass"
    t = (a["y"], a["x"]) .+ 1
    go.to_kgs(t)
end

function stringToAction(s, c)
    s == "pass" && return Dict("x"=>-1,"y"=>-1,"c"=>c)
    (y, x) = go.from_kgs(s) .- 1
    Dict("x"=>x,"y"=>y,"c"=>c)
end

function isIllegalMove(x)
    AlphaGo.get_position(agz) == nothing && return true
    !AlphaGo.go.is_move_legal(AlphaGo.get_position(agz), go.from_kgs(actionToString(x)))
end

# for rendering board
function updateBoard(move)
    env = Dict()
    pos = AlphaGo.get_position(agz)
    board = []
    if ( pos.board == nothing )
        board = map(i -> map(j-> [], 1:9), 1:9)
    else
        for i=1:go.N
            row = []
            for j=1:go.N
                cont = []
                if pos.board[j, i] == 1 || pos.board[j, i] == -1
                    push!(cont, Dict("x"=>(i - 1),"y"=>(j - 1),"c"=>pos.board[j, i], "type"=>"NORMAL"))
                end
                push!(row, cont)
            end
            push!(board, row)
        end
    end
    env["state"] = Dict("objects"=>board)
    env["action"] = move
    update!(gw, env)
end


"""
    axpy!(a::Number,X::Union{Matrix,RkMatrix,HMatrix},Y::Union{Matrix,RkMatrix,HMatrix})

Perform `Y <-- a*X + Y` in-place. Note that depending on the types of `X` and
`Y`, this may require converting from/to different formats during intermediate
calculations.

In the case where `Y` is an `RkMatrix`, the call `axpy!(a,X,Y)` should
typically be followed by recompression stage to keep the rank of `Y` under
control.

In the case where `Y` is an `HMatrix`, the call `axpy!(a,X,Y)` sums `X` to the
data in the node `Y` (and not on the leaves). In case `Y` has no `data`, it will
simply be assigned `X`. This means that after the call `axpy(a,X,Y)`, the object
`Y` is in a *dirty* state (see [`isclean`][@ref]) and usually a call to
[`flush_to_leaves!`](@ref) or [`flush_to_children!`](@ref) follows.
"""
function axpy!(a,X::Matrix,Y::RkMatrix)
    axpy!(a,RkMatrix(X),Y)
end

# 1.3
function axpy!(a,X::Matrix,Y::HMatrix)
    if hasdata(Y)
        axpy!(a,X,data(Y))
    else
        setdata!(Y,a*X)
    end
    return Y
end

# 2.1
function axpy!(a,X::RkMatrix,Y::Matrix)
    axpy!(a,Matrix(X),Y)
end

#2.2
function axpy!(a,X::RkMatrix,Y::RkMatrix)
    Y.A   = hcat(a*X.A,Y.A)
    Y.B   = hcat(X.B,Y.B)
    return Y
end

# 2.3
function axpy!(a,X::RkMatrix,Y::HMatrix)
    if hasdata(Y)
        axpy!(a,X,data(Y))
    else
        setdata!(Y,a*X)
    end
    return Y
end

#3.1
function axpy!(a,X::HMatrix,Y::Matrix)
    @debug "calling axpy! with `X` and HMatrix and `Y` a Matrix"
    shift = pivot(X) .- 1
    for block in PreOrderDFS(X)
        irange = rowrange(block) .- shift[1]
        jrange = colrange(block) .- shift[2]
        if hasdata(block)
            axpy!(a,data(block),view(Y,irange,jrange))
        end
    end
    return Y
end

# 3.2
function axpy!(a,X::HMatrix,Y::RkMatrix)
    @debug "calling axpby! with `X` and HMatrix and `Y` an RkMatrix"
    # FIXME: inneficient implementation due to conversion from HMatrix to
    # Matrix. Does it really matter? I don't think this function should be
    # called.
    axpy!(a,Matrix(X;global_index=false),Y)
end

# 3.3
function axpy!(a,X::HMatrix,Y::HMatrix)
    # TODO: assumes X and Y have the same structure. How to reinforce this?
    if hasdata(X)
        if hasdata(Y)
            axpy!(a,data(X),Y)
        else
            setdata!(Y,a*data(X))
        end
    end
    @assert size(children(X)) == size(children(Y)) "adding hierarchical matrices requires identical block structure"
    for (bx,by) in zip(children(X),children(Y))
        axpy!(a,bx,by)
    end
    return Y
end

# add a unifor scaling to an HMatrix return an HMatrix
function axpy!(a,X::UniformScaling,Y::HMatrix)
    @assert isclean(Y)
    if hasdata(Y)
        d = data(Y)
        @assert d isa Matrix
        n = min(size(d)...)
        for i=1:n
            d[i,i] += a*X.λ
        end
    else
        n = min(blocksize(Y)...)
        for i=1:n
            axpy!(a,X,children(Y)[i,i])
        end
    end
    return Y
end

Base.:(+)(X::UniformScaling,Y::HMatrix) = axpy!(true,X,deepcopy(Y))
Base.:(+)(X::HMatrix,Y::UniformScaling) = Y+X