using Roots

using UUIDs
using Observables
using CSV

using ScottishTaxBenefitModel
using .BCCalcs
using .Definitions
using .ExampleHelpers
using .FRSHouseholdGetter
using .GeneralTaxComponents
using .ModelHousehold
using .Monitor
using .Results
using .Runner
using .RunSettings
using .SimplePovertyCounts: GroupPoverty
using .SingleHouseholdCalculations
using .STBIncomes
using .STBOutput
using .STBParameters
using .Utils

const BASE_UUID = UUID("985c312f-129b-4acd-9e40-cb629d184183")

function initialise_settings()::Settings
    settings = Settings()
	settings.uuid = BASE_UUID
	settings.means_tested_routing = modelled_phase_in
    settings.run_name="run-$(date_string())"
	settings.income_data_source = ds_frs
	settings.dump_frames = false
	settings.do_marginal_rates = false
	settings.poverty_line = 100.0 # doesn't matter so long as +ive
	settings.requested_threads = 4
	settings.dump_frames = true
	return settings
end

const BASE_SETTINGS = initialise_settings()

function load_system()::TaxBenefitSystem
	sys = load_file( joinpath( Definitions.MODEL_PARAMS_DIR, "sys_2021_22.jl" ))
	#
	# Note that as of Budget21 removing these doesn't actually happen till May 2022.
	#
	load_file!( sys, joinpath( Definitions.MODEL_PARAMS_DIR, "sys_2021-uplift-removed.jl"))
	# uc taper to 55
	load_file!( sys, joinpath( Definitions.MODEL_PARAMS_DIR, "budget_2021_uc_changes.jl"))
	weeklyise!( sys )
	return sys
end

settings = initialise_settings()

sys = load_system()
chsys = deepcopy( sys )
chsys.scottish_child_payment.amount = 20.0

params = [sys, chsys]

# results = do_one_run( settings, params, obs )

mutable struct Thing
    a :: Float64
end

mutable struct RunParameters{T<:AbstractFloat}
    params :: TaxBenefitSystem{T}
    settings :: Settings
	base_cost :: T
end

function run( x :: Number, things :: RunParameters )
    obs = Observable(
		Progress(settings.uuid, "",0,0,0,0))
    nsr = deepcopy( things.params.it.non_savings_rates )
    things.params.it.non_savings_rates .+= x
    results = do_one_run(things.settings, [things.params], obs )
    things.params.it.non_savings_rates = nsr
	summary = summarise_frames(results,settings)
	nc = summary.income_summary[1][1,:net_cost]
	return round( nc - things.base_cost, digits=0 )
    # x^2 + things.params.it.non_savings_rates[1]
end

function baserun_cost()
    obs = Observable(
		Progress(settings.uuid, "",0,0,0,0))
	results = do_one_run(settings, [sys], obs )
	summary = summarise_frames(results,settings)
	return summary.income_summary[1][1,:net_cost]
end

nc = baserun_cost()

things = RunParameters( chsys, settings, nc )

zerorun = ZeroProblem( run, 0.0 )

incch = solve( zerorun, things )

println( "incch=$incch")


using BQCESubroutine
using BQCESubroutine.Utilities
using LinearAlgebra
using YaoLocations: YaoLocations, plain, Locations, CtrlLocations, merge_locations
using Test
using BenchmarkTools
using YaoArrayRegister
using LoopVectorization
using PaddedMatrices
using MatchCore
using MLStyle
using BQCESubroutine: BitContext, BitRoutine
using BQCESubroutine: diag_kernel, perm_kernel, generic_kernel, codegen_basic_broutine
using BQCESubroutine: codegen_broutine2x2_multi, codegen_broutine, hoist_entries!

BQCESubroutine.ENABLE_THREADS[]
BQCESubroutine.disable_threads()
U = rand(ComplexF64, 4, 4);
@benchmark broutine!(st, $U, Locations((3, 24))) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark instruct!(st, $U, (3, 24)) setup=(st=rand(ComplexF64, 1 << 25))

BQCESubroutine.disable_threads()
@benchmark broutine!(st, Val(:Z), Locations((3, 24))) setup=(st=rand(ComplexF64, 1 << 25))

BQCESubroutine.enable_threads()
@benchmark broutine!(st, Val(:Y), Locations((3, 24))) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark instruct!(st, Val(:Y), (3, 24)) setup=(st=rand(ComplexF64, 1 << 25))


"""
    create_geometry(setup; recreate = true)

Create cells, surfaces and finite element mesh. If `recreate = false`, previous geometry
will be reused.

This function does the following:
- Check geometry setup consistency
- Create or load cell configuration
- Create or load surface triangulation
- Call TetGen
- Deform domain
- Split mesh into compartments
For custom geometries with more than one compartment, call `split_mesh`
directly instead. This requires facet and element labels.
"""
function create_geometry(setup; recreate = true)
    (; refinement) = setup

    # File name for saving or loading geometry
    filename = joinpath(setup.meshdir, setup.name)

    # Make sure that folder exists
    dir = dirname(filename)
    isdir(dir) || mkpath(dir)

    if hasfield(typeof(setup), :ecs_shape)
        setup.ecs_shape ∈ [:no_ecs, :box, :convex_hull, :tight_wrap] ||
            error("Invalid ECS shape")
    end

    if setup isa Union{CylinderSetup,SphereSetup}
        # Check if cell description file is already available
        cellfilename = filename * "_cells"
        if isfile(cellfilename) && !recreate
            cells = read_cells(cellfilename)
        else
            cells = create_cells(setup)
            save_cells(cells, cellfilename)
        end
    else
        cells = nothing
    end

    # Make directory for storing finite elements mesh
    is_stl = endswith(filename, ".stl")
    if is_stl
        setup.ecs_shape == :no_ecs || error("ECS is only available for surface meshes")
        filename = filename[1:end-4]
    end
    save_meshdir_path = filename * "_dir"
    isdir(save_meshdir_path) || mkpath(save_meshdir_path)

    # Use an existing finite elements mesh or create a new finite elements mesh. The name of
    # the finite elements mesh is stored in the string `fname_tetgen_femesh`
    refinement_str = isinf(refinement) ? "" : "_refinement$refinement"
    fname_tetgen = "$save_meshdir_path/$(split(filename, "/")[end])$(refinement_str)_mesh"

    # Read or create surface triangulation
    if is_stl
        surfaces = nothing
    elseif isfile(fname_tetgen * ".node") && isfile(fname_tetgen * ".poly") && !recreate
        surfaces = read_surfaces(fname_tetgen)
    else
        if isa(setup, NeuronSetup)
            surfaces = create_surfaces(setup, filename)
        else
            surfaces = create_surfaces(setup, cells)
        end
        save_surfaces(fname_tetgen, surfaces)
    end

    # Add ".1" suffix to output file name, since this is what Tetgen does
    fname_tetgen_femesh = fname_tetgen * ".1"

    if isfile(fname_tetgen_femesh * ".node") && !recreate
        # Read global mesh from Tetgen output
        mesh_all = read_tetgen(fname_tetgen_femesh)
    elseif is_stl
        error("Not implemented")
    else
        mesh_all = call_tetgen(surfaces, refinement)
        save_tetgen(mesh_all, fname_tetgen_femesh)
    end

    # Deform domain
    if isa(setup, CylinderSetup) && (setup.bend > 1e-16 || setup.twist > 1e-16)
        @debug "Deforming domain with bend $(setup.bend) and twist $(setup.twist)"
        deform_domain!(mesh_all.points, setup.bend, setup.twist)
    end

    # Split mesh into compartments
    mesh = split_mesh(mesh_all)

    mesh, surfaces, cells
end


# Generate voicings exercises.

import JSON

# This ID is incremented for every problem generated.  Barring changes
# to the generation, this should generate the same ids on subsequent
# runs.
cur_id = 1

function emit(question, answer)
    global cur_id
    println("INSERT INTO \"probs\" VALUES($cur_id, '$question', '$answer');")
    cur_id += 1
end

type Transpose
    name :: String
    amount :: Int
end

# Transpositions for licks.  The order is mixed up a bit.
transposes = [
    Transpose("C", 0),
    Transpose("G", 7),
    Transpose("F", 5),
    Transpose("D", 2),
    Transpose("B♭", 10),
    Transpose("A", 9),
    Transpose("E♭", 3),
    Transpose("E", 4),
    Transpose("B", 11),
    Transpose("A♭", 8),
    Transpose("F♯", 6),
    Transpose("C♯", 1),
    Transpose("G♭", 6) ]

function transpose(notes :: Vector{Vector{Int}}, by :: Transpose)
    map(notes) do chord
        chord .+ by.amount
    end
end

function make_notes(notes :: Vector{Vector{Int}}, by :: Transpose)
    Dict("type" => "lick",
        "notes" => transpose(notes, by))
end

# Generate patterns for a particular voicing.  The name will be
# printed before, and the hint afterwards.  This should be enough for
# the reader to understand how to voice the chord.  The chord start
# can be used to skip part of the progression.
function gen_cycle(name, notes)
    for by in transposes
        emit("$name in $(by.name)", JSON.json(make_notes(notes, by)))
    end
end

println("DELETE FROM probs;")

# Jaki Byard, lick 1
gen_cycle("Byard 1",
    [ [60],[62],[64],[67],[64,65],[67],[63,64],[60],[57],[60] ])
gen_cycle("Byard 2",
    [ [50], [53], [57], [60], [59], [62], [65], [67], [68], [70], [71], [74], [72], [67] ])
gen_cycle("Martin 1",
    [ [63], [64], [69], [67], [64], [60], [65], [63], [64], [60],
      [57], [56], [55], [55] ])
gen_cycle("Martin 2",
    map([ [63],[64],[67],[70],[73],[75],[72] ]) do chord
        chord .- 5
    end)


"""When using distance geometry to generate geometries, lots of bad
distances get produced which messes up the resulting distributions.
Trim a series of bond length distribution files to a given frction of
the data.

This is a quick hack to trim those files. Not robust and will not
be needed when we have proper distributions"""

using ArgMacros
using CSV
using DataFrames
using LinearAlgebra

"""Very simplistic trimming of rhs and lhs based on the `counts` vector.
Args:
  distance: vector of distance values
  counts: vector of counts
  fraction: what fraction of data to retain
  output_fname: where to write the result
"""
function process_file_contents(distance::Vector{Float64},
                      counts::Vector{Int64},
                      fraction::Real,
                      output_fname::String)::Bool
  ntot = sum(counts)
  if ntot < 100
    @warn("Too few points $(ntot), skipping")
    return true
  end
  nkeep = round(Int, fraction * ntot)
  @info("N = $(ntot) fraction $fraction NKeep $(nkeep)")
  lhs = 1
  rhs = length(counts)
  removed = 0
  to_remove = ntot - nkeep
  while removed  < to_remove
    # If equal counts, make an arbitrary decision.
    if counts[lhs] <= counts[rhs]
      removed += counts[lhs]
      lhs += 1
    else
      removed += counts[rhs]
      rhs -= 1
    end
  end
  while counts[rhs] == 0
    rhs -= 1
  end
  while counts[lhs] == 0
    lhs += 1
  end
  x = distance[lhs:rhs]
  y = counts[lhs:rhs]
  to_write = DataFrame(distance=x, count=y)
  CSV.write(output_fname, to_write, writeheader=false)
  return true
end

function process_file(input_fname::String, output_fname::String,
                      fraction::Float64)::Bool
  if ! isfile(input_fname)
    @warn("Skipping $input_fname")
    return true
  end
  @info("Processing $input_fname")
  df = CSV.File(input_fname) |> DataFrame
  return process_file_contents(df[!, 1], df[!, 2], fraction, output_fname)
end

function main()
  @inlinearguments begin
    @argumentrequired String input_stem "-i" "--input_stem"
    @argumentrequired String output_stem "-o" "--output_stem"
    @argumentrequired Float64 fraction "-f" "--fraction"
  end

  bstart = 0
  bond_types = [0, 1, 2, 3]

  atomic_numbers = [1, 6, 7, 8, 9]
  for (t1, t2) in filter(x->x[1] <= x[2], collect(Iterators.product(ntuple(i->atomic_numbers, 2)...))[:])
    for btype in bstart:3
      input_fname = "$input_stem.$(t1).$(btype).$(t2)"
      output_fname = "$output_stem.$(t1).$(btype).$(t2)"
      process_file(input_fname, output_fname, fraction)
    end
  end
end

main()


R_CNFL = Tariff("CNFL",
                "Residential",
                [(0:200, 63.13), (200:300, 96.87), (300:1000, 100.14)],
                [(0:100, 0.0)],
                0.06,
                21.03,
                3.51,
                0.10,
                [0.992832136 0.999770704 1.00427654 0.983489891 1.009385917 1.012387984 0.979508228 0.994153679 0.97686098 1.015651323 1.015252586 1.016430032],
				1893.90);

CI_CNFL = Tariff("CNFL",
                "Commercial - Industrial",
                [(0:3000, 106.68), (3001:10000, 64.22)],
                [(0:8,  80426.32), (8:1000,  10053.29 )],
                0.06,
                21.03,
                3.51,
                0.10,
                [0.997759211, 1.029462815, 0.970880339, 1.004265761, 1.070739906, 1.039263159, 0.993586335, 1.002523373, 0.983852239, 0.950187922, 0.990476922, 0.967002018],
				1893.90);

TMT_CNFL = Tariff("CNFL",
                "Medium Voltage",
                [("punta", 54.42), ("valle", 27.21), ("noche", 19.59)],
                [("punta",  9542.79), ("valle",  6790.00), ("noche",  4310.42)],
                0.06,
                21.03,
                3.51,
                0.10,
                [1.024524936, 0.949492666, 1.025369474, 1.061316875, 1.023024878, 0.996315038, 0.992322444, 0.986442682, 0.947434961, 1.016134986, 1.006473428, 0.971147632],
				1893.90);

R_ICE = Tariff("ICE",
                "Residential",
                [(0:40.99,  3189.60), (41:200.99,  79.74), (201:1000, 143.71)],
                [(0:100, 0.0)],
                0.06,
                21.03,
                4.50,
                0.01,
                [0.992832136 0.999770704 1.00427654 0.983489891 1.009385917 1.012387984 0.979508228 0.994153679 0.97686098 1.015651323 1.015252586 1.016430032],
				3189.60);

CI_ICE = Tariff("ICE",
                "Commercial - Industrial",
                [(0:3000, 120.10), (3001:10000,  71.86)],
                [(0:1000,  11878.17)],
                0.06,
                21.03,
                4.50,
                0.01,
                [0.997759211, 1.029462815, 0.970880339, 1.004265761, 1.070739906, 1.039263159, 0.993586335, 1.002523373, 0.983852239, 0.950187922, 0.990476922, 0.967002018],
				3189.60);

TMT_ICE = Tariff("ICE",
                "Medium Voltage",
                [("punta",  68.90), ("valle",  25.60), ("noche",  15.76)],
                [("punta",  11181.25), ("valle",  7806.89), ("noche",  5000.50)],
                0.06,
                21.03,
                4.50,
                0.01,
                [1.024524936, 0.949492666, 1.025369474, 1.061316875, 1.023024878, 0.996315038, 0.992322444, 0.986442682, 0.947434961, 1.016134986, 1.006473428, 0.971147632],
				3189.60);

SD = Tariff("ICE",
                "Bulk System",
                [("punta",  53.98), ("valle",  44.22), ("noche",  37.54)],
                [("punta",  2863.19), ("valle",  2863.19), ("noche",  2863.19)],
                0.02,
                0.0,
                0.0,
                0.0,
                [0.0],
				3189.60);


module EMISAgentSimulation

#################################################################################
# Exports

# Export Structs
export CaseDefinition

export AgentSimulation

export Investor

export KalmanFilter
export InvestorBelief
export Scenario
export Forecast
export Perfect
export Imperfect

export RiskPreference
export RiskNeutral
export RiskAverse

export MarketPrices

export Finance

export Product
export OperatingProduct
export ReserveDirection
export ReserveUpEMIS
export ReserveDownEMIS
export ReserveUpEMIS
export OperatingReserve
export Energy
export Capacity
export REC

export BuildPhase
export Existing
export Planned
export Queue
export Option
export Retired

export DeviceEMIS
export Project
export GeneratorEMIS
export StorageEMIS

export BatteryTech
export BatteryEMIS

export HydroTech
export HydroGenEMIS

export RenewableTech
export RenewableGenEMIS

export ThermalTech
export ThermalGenEMIS

export ZonalLine

export MarketProject
export EnergyMarket
export ReserveUpMarket
export ReserveDownMarket
export CapacityMarket
export RECMarket
export MarketCollection
export MarketClearingProblem

# Export Functions

# Export Core functions
export calculate_annualized_costs
export calculate_capcost_and_wacc
export calculate_derating_data
export calculate_npv
export capacity_market_clearing
export cem
export cluster
export create_agent_simulation
export create_capacity_demand_curve
export create_expected_marketdata
export create_PSY_generator
export create_realized_marketdata
export economicdispatch
export energy_mkt_clearing
export finish_construction!
export make_investments!
export rec_market_clearing
export retire_old!
export retire_unprofitable!
export run_agent_simulation
export start_construction!

# Export Utility functions
export calculate_operating_profit
export calculate_required_processes
export create_parallel_workers
export dir_exists
export leaftypes
export make_case_data_dir
export make_results_dir
export read_data
export remove_leap_day!
export size_in_MW
export update_derating_factor!
export update_installed_cap!
export update_PSY_timeseries!
export update_simulation_derating_data!
export write_data

# Export Getter Functions
export get_accepted_perc
export get_active_power_limit
export get_active_power_limits
export get_activeprojects
export get_allprojects
export get_all_techs
export get_annual_cashflow
export get_annual_growth
export get_base_dir
export get_belief_update
export get_bus
export get_capacity_forward_years
export get_capacity_price
export get_capacity_bid
export get_capacity_factors
export get_cap_cost_multiplier
export get_capex_years
export get_case
export get_constant
export get_construction_year
export get_data_dir
export get_decidedprojects
export get_decision_year
export get_derating
export get_derating_data
export get_device_size
export get_discount_rate
export get_effective_investment_cost
export get_efficiency
export get_end_life_year
export get_energy_price
export get_error_covariance_estimate
export get_existing
export get_expected_npv
export get_expected_utility
export get_finance_data
export get_fixed_OM_cost
export get_forecast
export get_forecast_type
export get_from_zone
export get_fuel
export get_heterogeneity
export get_hour_weight
export get_info_symmetry
export get_input_active_power_limits
export get_investment_cost
export get_investor_data
export get_investors
export get_iteration_year
export get_kalman_filter
export get_lag_time
export get_life_time
export get_line_rating
export get_lines
export get_marginal_cost
export get_market_prices
export get_markets
export get_max_annual_projects
export get_maxcap
export get_max_limit
export get_maxperc
export get_measurement_covariance
export get_mincap
export get_multiplier
export get_name
export get_num_rep_days
export get_operation_cost
export get_options
export get_output_active_power_limits
export get_ownedby
export get_parallel_investors
export get_parallel_scenarios
export get_parameter_multipliers
export get_parameter_values
export get_peak_load
export get_planned
export get_prices
export get_probability
export get_process_covariance
export get_products
export get_projects
export get_queue
export get_queue_cost
export get_ramp_limits
export get_realized_profit
export get_rec_certificates
export get_rec_price
export get_rec_bid
export get_rep_hour_weight
export get_reservedown_price
export get_reserveup_price
export get_results_dir
export get_retired
export get_retirement_year
export get_risk_aversion
export get_risk_coefficient
export get_risk_preference
export get_rolling_horizon
export get_scenario_npv
export get_scenario_profit
export get_scenario_utility
export get_scenario_data
export get_siip_market_clearing
export get_simulation_years
export get_soc
export get_solver
export get_start_year
export get_state_estimate
export get_state_measurement_matrix
export get_state_transition_matrix
export get_storage_capacity
export get_system_ED
export get_system_UC
export get_system_services
export get_tech
export get_sys_dir
export get_time_limits
export get_to_zone
export get_total_horizon
export get_type
export get_uncertainty
export get_zone
export get_zones
export update_belief!

# Export Find functions
export find_active_invested_projects
export find_energy_product
export find_operating_products
export find_option_projects

# Export Setter Functions
export set_investors!
export set_iteration_year!
export set_peak_load!

#################################################################################
# Imports
import AxisArrays
import Clustering
import ClusterManagers
import CSV
import DataFrames
import Dates
import Distributed
import FileIO
import InteractiveUtils
import JuMP
export optimizer_with_attributes
import JLD2
import LinearAlgebra
import PowerSystems
import PowerSimulations
import InfrastructureSystems

const PSY = PowerSystems
const PSI = PowerSimulations
const IS = InfrastructureSystems

import Random
import UUIDs
import StatsBase
import Statistics
import TimeSeries

const TS = TimeSeries

import Base.convert
import JuMP.value
import MathOptInterface
import MathOptInterface: AbstractOptimizer

const MOI = MathOptInterface

################################################################################
# Includes

# Include all structs:

include("structs/products/Product.jl")
include("structs/products/Energy.jl")
include("structs/products/OperatingReserve.jl")
include("structs/products/Capacity.jl")
include("structs/products/REC.jl")

include("structs/Finance.jl")

include("structs/devices/BuildPhase.jl")
include("structs/devices/Project.jl")
include("structs/devices/BatteryEMIS.jl")
include("structs/devices/HydroGenEMIS.jl")
include("structs/devices/RenewableGenEMIS.jl")
include("structs/devices/ThermalGenEMIS.jl")

include("structs/CaseDefinition.jl")
include("structs/MarketPrices.jl")
include("structs/Scenario.jl")

include("structs/InvestorBelief.jl")

include("structs/Forecast.jl")
include("structs/RiskPreference.jl")
include("structs/ZonalLine.jl")

include("structs/Investor.jl")

include("structs/AgentSimulation.jl")

include("structs/market_structs/MarketProject.jl")
include("structs/market_structs/EnergyMarket.jl")
include("structs/market_structs/ReserveUpMarket.jl")
include("structs/market_structs/ReserveDownMarket.jl")
include("structs/market_structs/CapacityMarket.jl")
include("structs/market_structs/RECMarket.jl")
include("structs/market_structs/MarketCollection.jl")
include("structs/market_structs/MarketClearingProblem.jl")

# Include utilities:
include("utils/general.jl")                 # General utilities.
include("utils/market_utils.jl")           # Utils for actual and expected market clearing modules.
include("utils/product_utils.jl")          # Utils for products provided by the projects
include("utils/siip_psy_utils.jl")         # Utils for SIIP PowerSystems.
include("utils/conversion_utils.jl")       # Define new convert functions for changing buildphase of projects.
include("utils/read_and_write_utils.jl")   # Read and write utils.
include("utils/parallel_utils.jl")         # Utils for parallelizing price prediction runs.
include("utils/finance_utils.jl")          # Functions for calculating adjusted CAPEX and WACC

#Include files containing functions for creating the simulation structs from the given data.
include("struct_creators/simulation_structs/product_creator.jl")
include("struct_creators/simulation_structs/project_creator.jl")
include("struct_creators/simulation_structs/project_vector_creator.jl")
include("struct_creators/simulation_structs/investor_creator.jl")
include("struct_creators/simulation_structs/rts_psy_creator.jl")
include("struct_creators/simulation_structs/siip_psy_device_creator.jl")
include("struct_creators/simulation_structs/agent_simulation_creator.jl")

#Include files containing functions for creating the structs used for market clearing modules.
include("struct_creators/market_structs/market_project_creator.jl")
include("struct_creators/market_structs/cem_creator.jl")
include("struct_creators/market_structs/economic_dispatch_creator.jl")

#Include Test System Parsers
include("test_system_parsers/test_system_reader.jl")
include("test_system_parsers/rts_reader.jl")

#Include Expected and Actual Market Simultion functions.
include("markets_simulation/cem.jl")
include("markets_simulation/expected_market_simulation.jl")
include("markets_simulation/economicdispatch.jl")
include("markets_simulation/siip_simulation_definition.jl")
include("markets_simulation/actual_energy_mkt_clearing.jl")
include("markets_simulation/actual_capacity_mkt_clearing.jl")
include("markets_simulation/actual_rec_mkt_clearing.jl")
include("markets_simulation/actual_market_simulation.jl")

#Include Investor functions
    include("investor_functions/investor_iteration.jl")            # Runs investors annual iteration.

    #### Predictions #####################
    include("investor_functions/prediction/prediction_methodology.jl")  # Functions for running investors' price prediction methodology.
    include("investor_functions/prediction/operating_profit.jl")        # Functions for calculating expected operating market profits.
    include("investor_functions/prediction/capacity_profit.jl")         # Functions for calculating expected capacity market profits.
    include("investor_functions/prediction/REC_profit.jl")              # Functions for calculating expected REC market profits.
    include("investor_functions/prediction/total_profit.jl")            # Functions for updating the expected profits from all markets.

    #### Decisions ######################
    include("investor_functions/decisions/buildphase_conversion.jl")         # Functions for evaluating when to convert the buildphase of projects.
    include("investor_functions/decisions/investment_decision.jl")           # Functions for making investment decisions
    include("investor_functions/decisions/retirement_decision.jl")           # Functions for making retirement decisions

    #### Decision Metrics ######
    include("investor_functions/decision_metrics/npv_functions.jl")          # Functions for calculating NPV.
    include("investor_functions/decision_metrics/utility_functions.jl")      # Functions for calculating expected utility.

    ### Realized Profits and Updates ######
    include("investor_functions/realized_profits_calculator.jl")   # Functions for calculating realized profits from different markets
    include("investor_functions/annual_updates.jl")                # Functions for updating investor revenues and forecasts each year.

#Include derating factor updating methodology
include("derating_factor_updates/derating_factor_calculator.jl")  # Derating factor calculation for renewables

#Include k-medoids clustering for representative days selection.
include("representative_days/clustering.jl")

#Include main investment simulation function.
include("investment_simulation.jl")

end


include("actorcritic.jl")
include("reinforce.jl")
include("sarsa.jl")
include("q.jl")
include("fdpolicysearch.jl")
include("nactd.jl")
include("ppo.jl")

export ActorCritic, REINFORCE, NACTD
export Sarsa, ExpectedSarsa, QLearning
export PPO
export AbstractFiniteDifferencePolicySearch, BasicRandomSearch, AugmentedRandomSearch

function act!(agent::TA, env::TE, rng::AbstractRNG, learn::Bool) where {TA<:AbstractPolicyAgent,TE<:AbstractEnvironment} # do not learn if learn flag is set
    logp = get_action!(agent.π, env.state, rng)
    r = step!(env, agent.π.action, rng)
end


function integrate_scenario_radau(rr::RadauIntegrator{MechanismScenario{Dual{Nothing,Float64,NC}},NX,NC,NR,NSM};
        t_start::Float64=0.0, t_final::Float64=1.0, max_steps::Int64=1000, is_disp_count::Bool=false) where {NX,NC,NR,NSM}

    x = get_state(rr.de_object)
    data_state = zeros(max_steps + 1, NX)
    data_time = zeros(max_steps + 1)
    data_state[1, :] = x
    return integrate_scenario_radau(rr, data_state, data_time, t_final=t_final, is_disp_count=is_disp_count)
end

function integrate_scenario_radau(rr::RadauIntegrator{MechanismScenario{Dual{Nothing,Float64,NC}},NX,NC,NR,NSM},
        data_state::Matrix{Float64}, data_time::Vector{Float64}; t_final::Float64=1.0,
        is_disp_count::Bool=false) where {NX,NC,NR,NSM}

    mech_scen = rr.de_object
    s1_state, s2_state = size(data_state)
    s_time = length(data_time)
    (s2_state == NX) || error("The length of x ($(length(x))) is different from the number of variables in the scenario ($n_dof). Did you forget the state variables for bristle friction?")
    (s1_state == s_time) || error("data_state and data_time are different lengths")
    t_cumulative = data_time[1]
    x = data_state[1, :]
    for k = 2:s1_state
        is_disp_count && println(k)
        if mech_scen.discrete_controller != nothing
            while mech_scen.discrete_controller.t_last <= t_cumulative
                mech_scen.discrete_controller.t_last += mech_scen.discrete_controller.dt
                mech_scen.discrete_controller.control(x, mech_scen, t_cumulative)
            end
        end
        h, x = solveRadau(rr, x, t_cumulative)
        principal_value!(mech_scen, x)
        t_cumulative += h
        data_time[k] = t_cumulative
        data_state[k, :] = x
        if t_final < t_cumulative
            return data_time[1:k], data_state[1:k, :]
        end
    end
    return data_time, data_state
end


function collect_branches( subsys, nodes )  #=

collects the branches in a given circuit. It is assumed that the information of the circuit is contained in the variable "subsys" given as the first argument of this function. The variable "nodes" given as the second argument of this function is an array of the independent nodes in the circuit. The independent nodes may have been collected from the circuit description using the function collect_nodes(). The branches collected by this function will be returned as the array "branches" of the struct Branch. The number "num_branches" of branches (sum of branch phases) and the number "num_dyns" of dynamic branches will also be returned. In the array "branches", the node names have been replaced by their index numbers according to the array "nodes".

Author: Taku Noda
Started on Nov. 11, 2019
Update on June 29, 2021

Copyright 2021 Central Research Institute of Electric Power Industry

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

=#

# -- Internally used function ------------------------------------------

function find( node, nodes )
index = 0
for i = 1:length( nodes )
  if nodes[ i ] == node
    index = i
    break
  end
end
return index
end

# -- Main body starts here ---------------------------------------------

branches = Branch[]
num_branches = 0
num_dyns = 0
num_nls = 0

for elm in eachelement( subsys )
  if elm.name == "branch"
    name = ""; type = ""; nph = 0; val_str = ""
    lhs = Int[]; rhs = Int[]; p_vars = 0
    p_de = 0
    p_ns = 0; plx = []; ply = [];
    #
    type = elm[ "type" ]
    name = elm[ "name" ]
    #
    if !haskey( elm, "nph" )  # if single-phase,
      nph = 1
      if ( type == "R" )||( type == "L" )||( type == "G" )||( type == "C" )
        val_str = elm[ "val" ]
      elseif ( type == "NR" )||( type == "NL" )||( type == "NC" )||( type == "E" )||( type == "J" )||( type == "Sw" )
        val_str = elm.content
      end
    else  # if multiphase,
      nph = parse( Int, elm[ "nph" ] )
      for elm1 in eachelement( elm )
        if elm1.name == "val"
          val_str = elm1.content
          break
        end
      end
    end
    if ( type == "R" )||( type == "L" )||( type == "G" )||( type == "C" )||( type == "NR" )||( type == "NL" )||( type == "NC" )||( type == "Sw" )
      val_str = '['*replace( val_str, r"\t|\n" => "" )*']'
    elseif ( type == "E" )||( type == "J" )
      val_str = replace( val_str, r" |\t|\n" => "" )
    end
    val = Meta.parse( val_str );
    if ( type == "NR" )||( type == "NL" )||( type == "NC" )
      pldat = eval( val )
      plx = pldat[ :, 1 ]
      ply = pldat[ :, 2 ]
    elseif type == "Sw"
      plx = eval( val )
    end
    #
    if ( type == "L" )||( type == "C" )||( type == "NL" )||( type == "NC" )
      p_de = num_dyns + 1
      num_dyns = num_dyns + nph
    else
      p_de = 0
    end
    if ( type[ 1 ] == 'N' )||( type == "E" )||( type == "J" )||( type == "CSw" )||( type == "CE" )||( type == "CJ" )||( type == "VP" )||( type == "CP" )
      p_ns = num_branches + 1
    end
    if ( type[ 1 ] == 'N' )
      num_nls = num_nls + 1
      p_nl = num_nls
    end
    if ( ( type == "R" )||( type == "L" )||( type == "G" )||( type == "C" ) )&&( nph != 1 )
      for elm1 in eachelement( elm )
        nph = parse( Int, elm[ "nph" ] )
        if elm1.name == "lhs"
          node_names = split( elm1.content, ";" )
          if length( node_names ) != nph
            omegasim_error( "For the branch $name, the number of lefthand-side node names specified does not match the number of phases." )
          end
          for node_name in node_names
            push!( lhs, find( node_name, nodes ) )
          end
        elseif elm1.name == "rhs"
          node_names = split( elm1.content, ";" )
          if length( node_names ) != nph
            omega_error( "For the branch $name, the number of righthand-side node names specified does not match the number of phases." )
          end
          for node_name in node_names
            push!( rhs, find( node_name, nodes ) )
          end
        end
      end
    else
      push!( lhs, find( elm[ "lhs" ], nodes ) )
      push!( rhs, find( elm[ "rhs" ], nodes ) )
    end
    #
    p_vars = num_branches + 1
    num_branches = num_branches + nph
    #
    branch = Branch( name, type, nph, val, lhs, rhs, p_vars, p_de, p_ns, plx, ply, [ -1 ] )
    push!( branches, branch )
  end
end

return branches, num_branches, num_dyns

end


adj = [0 1 0 1;
       1 0 1 0;
       0 1 0 1;
       1 0 1 0]

@testset "defaultstaterepresentation.jl" begin

    @testset "DefaultStateRepresentation structure" begin
        g = SeaPearl.CPLayerGraph()
        nodeFeatures = [1.0f0 1.0f0; 2.0f0 2.0f0]
        variableIdx = 1
        dsr = SeaPearl.DefaultStateRepresentation{SeaPearl.DefaultFeaturization,SeaPearl.DefaultTrajectoryState}(g, nodeFeatures, nothing, variableIdx, nothing)

        @test dsr.cplayergraph == g
        @test dsr.nodeFeatures == nodeFeatures
        @test dsr.variableIdx == 1

    end

    @testset "DefaultSR from CPModel" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(3, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        dsr = SeaPearl.DefaultStateRepresentation(model)
        SeaPearl.update_representation!(dsr, model, x)

        @test Matrix(LightGraphs.LinAlg.adjacency_matrix(dsr.cplayergraph)) == [0 0 1 1 0 0
                                                                                0 0 1 1 0 0
                                                                                1 1 0 0 0 1
                                                                                1 1 0 0 1 1
                                                                                0 0 0 1 0 0
                                                                                0 0 1 1 0 0]

        @test dsr.nodeFeatures == Float32[  1 1 0 0 0 0
                                        0 0 1 1 0 0
                                        0 0 0 0 1 1]
        @test dsr.variableIdx == 3
        @test SeaPearl.cpVertexFromIndex(SeaPearl.CPLayerGraph(model), dsr.variableIdx).variable == model.variables["x"]
    end

    @testset "update_representation!()" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        dsr = SeaPearl.DefaultStateRepresentation(model)
        SeaPearl.update_representation!(dsr, model, x)

        @test Matrix(LightGraphs.LinAlg.adjacency_matrix(dsr.cplayergraph)) == [0 0 1 1 0 0
                                                    0 0 1 1 0 0
                                                    1 1 0 0 1 1
                                                    1 1 0 0 1 1
                                                    0 0 1 1 0 0
                                                    0 0 1 1 0 0]

        @test dsr.nodeFeatures == Float32[   1 1 0 0 0 0
                                                        0 0 1 1 0 0
                                                        0 0 0 0 1 1]

        @test dsr.variableIdx == 3
        @test SeaPearl.cpVertexFromIndex(SeaPearl.CPLayerGraph(model), dsr.variableIdx).variable == model.variables["x"]

        SeaPearl.assign!(y, 2)
        g = SeaPearl.CPLayerGraph(model)

        SeaPearl.update_representation!(dsr, model, y)

        @test Matrix(LightGraphs.LinAlg.adjacency_matrix(dsr.cplayergraph)) == [0 0 1 1 0 0
                                                    0 0 1 1 0 0
                                                    1 1 0 0 1 1
                                                    1 1 0 0 1 0
                                                    0 0 1 1 0 0
                                                    0 0 1 0 0 0]

        @test dsr.nodeFeatures == Float32[   1 1 0 0 0 0
                                                        0 0 1 1 0 0
                                                        0 0 0 0 1 1]
        @test dsr.variableIdx == 4
        @test SeaPearl.cpVertexFromIndex(g, dsr.variableIdx).variable == model.variables["y"]

    end

    @testset "trajectoryState constructor" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        dsr = SeaPearl.DefaultStateRepresentation(model)
        SeaPearl.update_representation!(dsr, model, x) #add x as the branching variable
        dts = SeaPearl.DefaultTrajectoryState(dsr)  #creates the DefaultTrajectoryState object

        @test dts.variableIdx == SeaPearl.indexFromCpVertex(dsr.cplayergraph, SeaPearl.VariableVertex(x))

    end

    @testset "BatchedDefaultTrajectoryState constructor" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        dsr = SeaPearl.DefaultStateRepresentation(model)
        SeaPearl.update_representation!(dsr, model, x) #add x as the branching variable
        dts = SeaPearl.DefaultTrajectoryState(dsr)  #creates the DefaultTrajectoryState object

        batchedDtsSingle = dts |> cpu
        @test batchedDtsSingle.fg.graph[:,:,1]== adjacency_matrix(dts.fg)
        @test batchedDtsSingle.fg.nf[:,:,1] == dts.fg.nf
        @test batchedDtsSingle.fg.ef[:,:,1] == dts.fg.ef
        @test batchedDtsSingle.fg.gf[:,1] == dts.fg.gf
        @test batchedDtsSingle.variableIdx[1] == dts.variableIdx

        SeaPearl.assign!(x, 2)
        SeaPearl.fixPoint!(model, SeaPearl.getOnDomainChange(x))
        SeaPearl.update_representation!(dsr, model, x) #add x as the branching variable
        dts2 = SeaPearl.DefaultTrajectoryState(dsr)  #creates the DefaultTrajectoryState object

        batchedDts = [dts,dts2] |> cpu
        @test size(batchedDts.fg.graph,3) == 2
        @test size(batchedDts.fg.nf,3) == 2
        @test size(batchedDts.fg.ef,3) == 2
        @test size(batchedDts.fg.gf,2) == 2
        @test size(batchedDts.variableIdx,1) == 2
    end
end


__precompile__()
module Oanda
using Requests, JSON, SQLite
export price_stream, books, pairs, log_json, log_db, order
export Tick
type Tick
    ask::Float64
    bid::Float64
    sent::DateTime
    recd::DateTime
    pair::String
    #live::Bool
end
function Tick(i::Dict{String,Any})
    Tick(
        parse(Float64,i["asks"][1]["price"]), #ask
        parse(Float64,i["bids"][1]["price"]), #bid
        parse(DateTime,i["time"]),            #sent
        now(),                                #recd
        i["instrument"],                      #pair
        #i["tradeable"],                      #live
    )
end
include("init.jl")  #return values for account and api_key
const HEADERS = Dict("Content-Type" => "application/json", "Authorization" => "Bearer " * API_KEY, "AcceptDatetimeFormat" => "RFC3339")
include("pairs.jl") #return array of tradeable instruments for given account
include("prices.jl")#return buffered stream of price data for instrument(s), #return next price from stream
include("books.jl") #return order book and/or position book for instrument(s) and time period(s)
include("logs.jl")  #write prices or program output to file
include("order.jl") #format and send orders to Oanda for execution
end


# Neighbour finders

export
    NoNeighbourFinder,
    find_neighbours!,
    DistanceNeighbourFinder

"""
    NoNeighbourFinder()

Placeholder neighbour finder that returns no neighbours.
When using this neighbour finder, ensure that `nl_only` for the interactions is
set to `false`.
"""
struct NoNeighbourFinder <: NeighbourFinder end

"""
    find_neighbours!(simulation, neighbour_finder, step_n; parallel=true)

Obtain a list of close atoms in a system.
Custom neighbour finders should implement this function.
"""
function find_neighbours!(s::Simulation,
                            ::NoNeighbourFinder,
                            ::Integer;
                            kwargs...)
    return
end

"""
    DistanceNeighbourFinder(nb_matrix, n_steps, dist_cutoff)
    DistanceNeighbourFinder(nb_matrix, n_steps)

Find close atoms by distance.
"""
struct DistanceNeighbourFinder{T} <: NeighbourFinder
    nb_matrix::BitArray{2}
    n_steps::Int
    dist_cutoff::T
end

function DistanceNeighbourFinder(nb_matrix::BitArray{2},
                                 n_steps::Integer)
    return DistanceNeighbourFinder(nb_matrix, n_steps, 1.2)
end

function find_neighbours!(s::Simulation,
                            nf::DistanceNeighbourFinder,
                            step_n::Integer;
                            parallel::Bool=true)
    !iszero(step_n % nf.n_steps) && return

    neighbours = s.neighbours
    empty!(neighbours)
    sqdist_cutoff = nf.dist_cutoff ^ 2

    if parallel && nthreads() > 1
        nl_threads = [Tuple{Int, Int}[] for i in 1:nthreads()]

        @threads for i in 1:length(s.coords)
            nl = nl_threads[threadid()]
            ci = s.coords[i]
            nbi = nf.nb_matrix[:, i]
            for j in 1:(i - 1)
                r2 = sum(abs2, vector(ci, s.coords[j], s.box_size))
                if r2 <= sqdist_cutoff && nbi[j]
                    push!(nl, (i, j))
                end
            end
        end

        for nl in nl_threads
            append!(neighbours, nl)
        end
    else
        for i in 1:length(s.coords)
            ci = s.coords[i]
            nbi = nf.nb_matrix[:, i]
            for j in 1:(i - 1)
                r2 = sum(abs2, vector(ci, s.coords[j], s.box_size))
                if r2 <= sqdist_cutoff && nbi[j]
                    push!(neighbours, (i, j))
                end
            end
        end
    end
end


using HTTP
using ..Fed: Configuration, serialize_payload, deserialize_payload

"""
    Node(host, port, fit, config)

Orchestrate the training on local data of the clients participating in the
federated training. A `Node` is registerating itself to the server on startup,
and then wait for server requests for training on its local data.

The `fit` function is called "proxied function". It should define the training
loop and respect the following signature:

    fit(weights::Vector{Float32})::Vector{Float32}

Where `weights` is a flatten 1D array contianing the model parameters.
"""
struct Node
    host::String
    port::Int
    fit::Function
    config::Configuration
end

"""
    register_to_server(node::Node)

Register the `node` to the server, letting it knows that a node is available to
take part in the training.

!!! info

    **API Client**: this is an HTTP client request.
"""
function register_to_server(node::Node)::HTTP.Response
    endpoint = node.config.base.serverurl * node.config.base.register_node
    payload = "http://$(node.host):$(node.port)"
    response = HTTP.request("POST", endpoint, [], payload)
    return response
end

"""
    fit_service(node::Node, request::HTTP.Request)

Fit the received weights contained in the `request` with the node's local data,
and return the updated weights to the caller. This function is calling the
"proxied function" `fit` from the `Node`.

!!! info

    **API Endpoint**: this is an HTTP client endpoint.
"""
function fit_service(node::Node, request::HTTP.Request)::HTTP.Response
    weights = deserialize_payload(node.config.payload_serde, request.body, "http://127.0.0.1:9090")
    weights = node.fit(weights)
    payload = serialize_payload(node.config.payload_serde, weights)
    return HTTP.Response(200, payload)
end


###################################################
## projectors/networkprojector.jl
## link geographical and network data
###################################################

"""
    `NetworkProjector`: abstract type that creates a link between NetworkData and GeoData
    is used to project geographical data onto network and to use network results
    must implement attributes:
    - `network::Network`
    - `trips::GeoData` : reference to trip data
    methods:
    - `preloadData!`: takes GeoData and load it
    - `getNetworkTrips`: returns network trips associated with given geo-trips
    - `getTripTiming` : returns timing estimation for a geographical trip (random or known)
"""
abstract type NetworkProjector end

function Base.show(io::IO, p::NetworkProjector)
    typeName = split(string(typeof(p)),".")[end]
    nTrips = length(p.trips)
    println(io,"Network Projector: $(typeName)")
    println(io,"$nTrips trips pre-loaded")
end

"""
    `NetworkData` : convenience function to create NetworkData using Projector
"""
function NetworkData(
        proj::NetworkProjector,
        trips::AbstractArray{Int64,1} = eachindex(proj.trips),
        minTimes::AbstractArray{Float64,2} = maxSpeedTimes(proj.network))
    # Create the "NetworkTrip" array
    return NetworkData(proj.network, getNetworkTrips(proj, trips), minTimes)
end


import Graphs: in_edges, in_degree, in_neighbors, out_edges, out_degree, out_neighbors
import Base: isempty

type BayesianNetwork <: AbstractGraph{BayesianNode, BayesianEdge}
    nodes::Array{BayesianNode,1}
    edges::Array{BayesianEdge,1}

    finclist::Array{Array{BayesianEdge,1},1}   #forward incidence list
    binclist::Array{Array{BayesianEdge,1},1}   #backward incidence list

    cpds::Dict{CPD,Any} ## Add types if needed
    checked_cpds::Dict{CPD,Uint8} # Keep a list of having or not: 0 = not processed, 1 = being processed, 2 = not derivable, 3 = derivable

    function BayesianNetwork{V <: BayesianNode}(_nodes::Array{V,1}, _edges::Array{BayesianEdge,1})
        _nodes = convert(Array{BayesianNode,1}, _nodes)
        l_nodes = length(_nodes)
        _cpds=Dict{CPD, Any}()
        _checked_cpds=Dict{CPD,Uint8}()
        if l_nodes > 0
            map(x -> assign_index(x[2],x[1]), enumerate(_nodes))
            for _n in _nodes
                if has_pd(_n)
                    if length(probabilities(_n.pd)) > 0
                        _cpds[P(_n.label)] = _n.pd
                        _checked_cpds[P(_n.label)] = 0
                    end
                end
            end
        end
        if length(_edges) > 0
            map(x -> assign_index(x[2],x[1]), enumerate(_edges))
        end
        b = new(_nodes, Array(BayesianEdge,0), multivecs(BayesianEdge, l_nodes), multivecs(BayesianEdge, l_nodes), _cpds, _checked_cpds)
        for edge in _edges
            add_edge!(b, edge)
        end
        b
    end
end

BayesianNetwork(n::Array{None,1}, e::Array{None,1}) = BayesianNetwork(Array(BayesianNode,0),Array(BayesianEdge,0))
BayesianNetwork{V <: BayesianNode}(n::Array{V,1}) = BayesianNetwork(n,Array(BayesianEdge,0))
BayesianNetwork{V <: BayesianNode}(n::Array{V,1}, e::Array{None,1}) = BayesianNetwork(n,Array(BayesianEdge,0))
BayesianNetwork() = BayesianNetwork([], [])

nodes(g::BayesianNetwork) = g.nodes
num_nodes(g::BayesianNetwork) = length(nodes(g))

edges(g::BayesianNetwork) = g.edges
num_edges(g::BayesianNetwork) = length(edges(g))

cpds(g::BayesianNetwork) = g.cpds
num_cpds(g::BayesianNetwork) = length(cpds(g))

checked_cpds(g::BayesianNetwork) = g.checked_cpds
num_checked_cpds(g::BayesianNetwork) = length(calculated_cpds(g))

function reset_checked_cpds!(g::BayesianNetwork)
    for i in collect(keys(checked_cpds(g)))
        checked_cpds(g)[i] = 0
    end
end

function set_cpd_derivable!(g::BayesianNetwork, cpd::CPD, value::Bool)
    v = value ? 3 : 2
    checked_cpds(g)[cpd] = v
end

function cpd_derivable(g::BayesianNetwork, cpd::CPD)
    if (check_checked_cpd(g, cpd))
        checked_cpds(g)[cpd] == 3
    else
        false
    end
end

function start_processing_cpd!(g::BayesianNetwork, cpd::CPD)
    checked_cpds(g)[cpd] = 1
end

function cpd_is_processing(g::BayesianNetwork, cpd::CPD)
    if check_checked_cpd(g, cpd)
        checked_cpds(g)[cpd] == 1
    else
        false
    end
end

function cpd_has_been_processed(g::BayesianNetwork, cpd::CPD)
    if check_checked_cpd(g, cpd)
        checked_cpds(g)[cpd] > 1
    else
        false
    end
end

function check_checked_cpd(g::BayesianNetwork, cpd::CPD)
    try
        getindex(checked_cpds(g), cpd)
        true
    catch
        false
    end
end

function assign_index(g_elem, i::Int)
    if g_elem.index != 0
        throw("Attempting to reassign node")
    end
    g_elem.index=i
end

function add_node!{T <: BayesianNode}(g::BayesianNetwork, n::T)
    assign_index(n, num_nodes(g) + 1)

    push!(g.nodes, n)

    # Create entries for vertex in backward and forward incidence lists
    push!(g.finclist, Int[])
    push!(g.binclist, Int[])

    ##Update what is put into the cpds dictionary when decided
    if typeof(n) == DBayesianNode && has_pd(n)
        add_probability!(g, P(n.label), n.pd)
        ##Check just made in case someone construct a graph in an non-obvious manner, will probably not be run
        if length(in_edges(n, g)) > 0
            add_probability!(g, P(n.label|[edge.source.label for edge in in_edges(n, g)]))
        end
    else
        add_probability!(g, P(n.label))
    end

    n
end

add_edge!{S <: BayesianNode, T <: BayesianNode}(g::BayesianNetwork, s::S, t::T) = add_edge!(g, BayesianEdge(s,t))

function add_edge!(g::BayesianNetwork, e::BayesianEdge)
    assign_index(e, num_edges(g) + 1)

    u = source(e)
    v = target(e)

    if !nodes_in_network(g, [u, v])
        throw("Attempting to add an edge to a network where the nodes are not present")
    end

    ui = node_index(u)::Int
    vi = node_index(v)::Int

    # Add edge to network and to forward list for source and to backward list for target
    push!(g.edges, e)
    push!(g.finclist[ui], e)
    push!(g.binclist[vi], e)
    add_cpd_for_edge!(g,u,v)

    e
end

function add_cpd_for_edge!(g::BayesianNetwork, s::BayesianNode,t::BayesianNode)
    if typeof(s) == DBayesianNode && has_pd(s) && has_pd(t)
        # g.cpds[CPD(t.label, [edge.source.label for edge in in_edges(t,g)])] = null
    else
        # g.cpds[CPD(t.label, s.label)] = null
    end
end

function nodes_in_network{V <: BayesianNode}(g::BayesianNetwork, ns::Array{V,1})
    for node in ns
        if !node_in_network(g,node)
            return false
        end
    end
    true
end

function node_in_network{V <: BayesianNode}(g::BayesianNetwork, n::V)
    if node_index(n) > length(nodes(g)) || node_index(n) == 0
        false
    else
        nodes(g)[node_index(n)] == n
    end
end

function symbols_in_network(g::BayesianNetwork, ss::Array{Symbol,1})
    for s in ss
        if find_node_by_symbol(g,s) == false
            return false
        end
    end
    true
end

function find_node_by_symbol(g::BayesianNetwork, s::Symbol)
    for node in g.nodes
        if node.label == s
            return node
        end
    end
    false
end

function in_edges{V <: BayesianNode}(n::V, g::BayesianNetwork)
    if !node_in_network(g, n)
        Array{BayesianEdge,1}[]
    else
        g.binclist[node_index(n)]
    end
end

in_degree{V <: BayesianNode}(n::V, g::BayesianNetwork) = length(in_edges(n, g))
in_neighbors{V <: BayesianNode}(n::V, g::BayesianNetwork) = map(e -> e.source, in_edges(n, g))

function out_edges{V <: BayesianNode}(n::V, g::BayesianNetwork)
    if !node_in_network(g, n)
        Array{BayesianEdge,1}[]
    else
        g.finclist[node_index(n)]
    end
end

out_degree{V <: BayesianNode}(n::V, g::BayesianNetwork) = length(out_edges(n, g))
out_neighbors{V <: BayesianNode}(n::V, g::BayesianNetwork) = map(e -> e.target, out_edges(n, g))

add_probability!{V <: PDistribution}(g::BayesianNetwork, cpd::CPD, pd::V) = _add_probability!(g, cpd, pd)
add_probability!(g::BayesianNetwork, cpd::CPD) = _add_probability!(g, cpd, nothing)

function _add_probability!(g::BayesianNetwork, cpd::CPD, pd)
    if symbols_in_network(g, distribution(cpd)) && symbols_in_network(g, conditionals(cpd))
        g.cpds[cpd] = pd
        checked_cpds(g)[cpd] = 0
        true
    else
        false
    end
end

multivecs{T}(::Type{T}, n::Int) = [T[] for _ =1:n]

#################################################
#
#  query
#
################################################

function legal_configuration(bn::BayesianNetwork, cpd::CPD)
    nodes = Set(cpd.conditionals)
    size = 0
    checkedNodes = Set()
    uncheckedNodes = Set()
    for label in cpd.distribution
        uncheckedNodes = union(Set(gather_nodes_by_edge_out(out_edges(find_node_by_symbol(bn,label), bn))), gather_nodes_by_edge_in(in_edges(find_node_by_symbol(bn,label), bn)))
    end
    while 0 < length(uncheckedNodes)
        for node in uncheckedNodes
            inNodes = gather_nodes_by_edge_in(in_edges(find_node_by_symbol(bn,node.label), bn))
            outNodes = gather_nodes_by_edge_out(out_edges(find_node_by_symbol(bn,node.label), bn))
            union!(uncheckedNodes, setdiff(inNodes,checkedNodes))
            union!(uncheckedNodes, setdiff(outNodes,checkedNodes))
            push!(checkedNodes,node)
            delete!(uncheckedNodes,node)
        end
    end

    for node in checkedNodes
        if node.label in nodes
            delete!(nodes,node.label)
        end
    end
    if isempty(nodes)
        true
    else
        false
    end
end

gather_nodes_by_edge_in(edges::Array{BayesianEdge,1}) = Set(map(edge -> edge.source, edges))

gather_nodes_by_edge_out(edges::Array{BayesianEdge,1}) = Set(map(edge -> edge.target, edges))

function check_requirements(g::BayesianNetwork, cpd::CPD)
    reset_checked_cpds!(g)
    _check_requirements(g, cpd)
end

function _check_requirements(g::BayesianNetwork, cpd::CPD)
    if cpd_has_been_processed(g, cpd)
        cpd_derivable(g, cpd)
    else
        if check_cpd(g, cpd)
            set_cpd_derivable!(g, cpd, true)
            true
        elseif cpd_is_processing(g, cpd)
            false
        else
            start_processing_cpd!(g, cpd)

            val = false
            conds = copy(conditionals(cpd))
            l_conds = length(conds)
            dist = copy(distribution(cpd))
            l_dist = length(dist)
            # If there are conditionals, then try bayes rule
            if l_conds > 0 && check_bayes_rule(g, cpd)
                val = true
            # If there are more than one distribution and there are conditionals, then try to split it
            elseif l_dist > 1 && l_conds > 0 && check_split_distribution(g, cpd)
                val = true
            # If there are more than one distribution but no conditionals
            elseif l_dist > 1 && l_conds == 0 && check_joint_probability_distribution(g, cpd)
                val = true
            end
            set_cpd_derivable!(g, cpd, val)
            val
        end
    end
end

function check_split_distribution(g::BayesianNetwork, cpd::CPD)
    conds = copy(conditionals(cpd))
    dist = copy(distribution(cpd))

    new_dists = 0
    if length(dist) > 1
        new_dists = splice!(dist, 2:length(dist))
    else
        throw("Length of distribution must be larger than 1.")
    end
    _check_requirements(g, P(new_dists|conds)) && _check_requirements(g, P(dist|conds))
end

function check_bayes_rule(g::BayesianNetwork, cpd::CPD)
    conds = copy(conditionals(cpd))
    dist = copy(distribution(cpd))

    opposite = P(conds|dist)
    if _check_requirements(g, opposite)
        to_check = [P(conds), P(dist)]
        for i in to_check
            if !_check_requirements(g, i)
                return false
            end
        end
        true
    else
        false
    end
end

function check_joint_probability_distribution(bn::BayesianNetwork, cpd::CPD)
    for d in distribution(cpd)
        node = find_node_by_symbol(bn, d)
        parent_labels = convert(Array{Symbol,1}, map(p -> p.label, in_neighbors(node, bn)))
        if !_check_requirements(bn, P(d|parent_labels))
            return false
        end
    end
    true
end

function check_cpd(g::BayesianNetwork, cpd::CPD)
    try
        getindex(cpds(g), cpd)
        true
    catch
        false
    end
end

has_cached_result(bn::BayesianNetwork, cpd::CPD) = cpd in collect(keys(cpds(bn))) ? true : false;

function query(bn::BayesianNetwork, cpd::CPD)
    if is_naive_bayes_model(bn, cpd)
        query_naive_bayes(bn,cpd)
    else
        throw("Form of cpd not supported by query at the moment")
    end
end

function is_naive_bayes_model(bn::BayesianNetwork, cpd::CPD)
    if typeof(find_node_by_symbol(bn,distribution(cpd)[1])) == DBayesianNode
        for label in conditionals(cpd)
            if !(typeof(find_node_by_symbol(bn, label)) == CBayesianNode)
                return false
            end
        end
    else
        false
    end
    true
end

function query_naive_bayes(bn::BayesianNetwork, cpd::CPD)
    if !legal_configuration(bn,cpd)
        throw("Invaled configuration")
    elseif has_cached_result(bn,cpd)
        cpds(bn)[cpd]
    elseif check_requirements(bn, cpd)
        calculate_probabilities(bn,cpd)
    end
end

function calculate_probabilities(bn::BayesianNetwork, cpd::CPD)
    res = Dict()
    temp = create_probabilities(cpd)
    for p in temp
        has_cached_result(bn, p) ? calculate_probability!(res, cpds(bn)[p], cpd.parameters[p.distribution[1]]) : nothing;
    end
    for key in keys(res)
        res[key] = get_probability(bn, find_node_by_symbol(bn,cpd.distribution[1]), key)*prod(res[key])
    end
    res
end

get_probability(bn::BayesianNetwork, node::BayesianNode, state) = has_cached_result(bn, P(node.label)) ? cpds(bn)[P(node.label)][state][1] : throw("Node mismatch");

create_probabilities(cpd::CPD) = map(cond -> CPD(cond,cpd.distribution), cpd.conditionals)

function calculate_probability!(dict, dist::ProbabilityDensityDistribution, v)
    for state in states(dist)
        if !(state in collect(keys(dict)))
            dict[state] = Array(Float64,0)
        end
    end

    for i = 1:length(states(dist))
        push!(dict[states(dist)[i]],pdfs(dist)[i](v))
    end
end


using PyCall
using DPMMSubClusters.TorchWrapper: torch_model_2d, torch_model_10d, torch_model_128d

function create_first_local_cluster(group::local_group)
    suff =
        create_sufficient_statistics(group.model_hyperparams.distribution_hyper_params, [])
    post = group.model_hyperparams.distribution_hyper_params
    dist = sample_distribution(post)
    cp = cluster_parameters(
        group.model_hyperparams.distribution_hyper_params,
        dist,
        suff,
        post,
    )
    cpl = deepcopy(cp)
    cpr = deepcopy(cp)
    splittable = splittable_cluster_params(
        cp,
        cpl,
        cpr,
        [0.5, 0.5],
        false,
        ones(burnout_period + 5) * -Inf,
    )
    cp.suff_statistics.N = size(group.points, 2)
    cpl.suff_statistics.N = sum(group.labels_subcluster .== 1)
    cpl.suff_statistics.N = sum(group.labels_subcluster .== 2)
    cluster = local_cluster(
        splittable,
        group.model_hyperparams.total_dim,
        cp.suff_statistics.N,
        cpl.suff_statistics.N,
        cpl.suff_statistics.N,
    )
    @sync for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i split_first_cluster_worker!(group)
    end
    return cluster
end

function create_outlier_local_cluster(group::local_group, outlier_params)
    suff = create_sufficient_statistics(outlier_params, outlier_params, Array(group.points))
    post = calc_posterior(outlier_params, suff)
    dist = sample_distribution(post)
    cp = cluster_parameters(outlier_params, dist, suff, post)
    cpl = deepcopy(cp)
    cpr = deepcopy(cp)
    splittable = splittable_cluster_params(
        cp,
        cpl,
        cpr,
        [0.5, 0.5],
        false,
        ones(burnout_period + 5) * -Inf,
    )
    cp.suff_statistics.N = size(group.points, 2)
    cpl.suff_statistics.N = sum(group.labels_subcluster .== 1)
    cpl.suff_statistics.N = sum(group.labels_subcluster .== 2)
    cluster = local_cluster(
        splittable,
        group.model_hyperparams.total_dim,
        cp.suff_statistics.N,
        cpl.suff_statistics.N,
        cpl.suff_statistics.N,
    )
    # @sync for i in (nworkers()== 0 ? procs() : workers())
    #     @spawnat i split_first_cluster_worker!(group)
    # end
    return cluster
end

function sample_sub_clusters!(group::local_group)
    for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i sample_sub_clusters_worker!(
            group.points,
            group.labels,
            group.labels_subcluster,
        )
    end
end

function sample_sub_clusters_worker!(group_points, group_labels, group_labels_subcluster)

    pts = localpart(group_points)
    labels = localpart(group_labels)
    sub_labels = localpart(group_labels_subcluster)

    for (i, v) in enumerate(clusters_vector)
        create_subclusters_labels!(
            (@view sub_labels[labels.==i]),
            (@view pts[:, labels.==i]),
            v,
        )
    end

end

function create_subclusters_labels!(
    labels::AbstractArray{Int64,1},
    points::AbstractArray{Float32,2},
    cluster_params::thin_cluster_params,
)
    if size(labels, 1) == 0
        return
    end
    parr = zeros(Float32, length(labels), 2)
    log_likelihood!((@view parr[:, 1]), points, cluster_params.l_dist)
    log_likelihood!((@view parr[:, 2]), points, cluster_params.r_dist)
    parr[:, 1] .+= log(cluster_params.lr_weights[1])
    parr[:, 2] .+= log(cluster_params.lr_weights[2])
    sample_log_cat_array!(labels, parr)
end

function sample_labels!(group::local_group, final::Bool, no_more_splits::Bool)
    sample_labels!(group.labels, group.points, final, no_more_splits)
end

function sample_labels!(
    labels::AbstractArray{Int64,1},
    points::AbstractArray{Float32,2},
    final::Bool,
    no_more_splits::Bool,
)
    for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i sample_labels_worker!(labels, points, final, no_more_splits)
    end
end

function sample_labels_worker!(
    labels::AbstractArray{Int64,1},
    points::AbstractArray{Float32,2},
    final::Bool,
    no_more_splits::Bool,
)
    indices = localindices(points)[2]
    lbls = localpart(labels)
    pts = localpart(points)
    log_weights = log.(clusters_weights)
    parr = zeros(Float32, length(indices), length(clusters_vector))
    tic = time()
    for (k, cluster) in enumerate(clusters_vector)
        log_likelihood!(reshape((@view parr[:, k]), :, 1), pts, cluster.cluster_dist)
    end
    for (k, v) in enumerate(clusters_weights)
        parr[:, k] .+= log(v)
    end

    if final
        lbls .= mapslices(argmax, parr, dims = [2])[:]
    else
        sample_log_cat_array!(lbls, parr)
    end
end

function update_splittable_cluster_params!(splittable_cluser::splittable_cluster_params)
    cpl = splittable_cluser.cluster_params_l
    cpr = splittable_cluser.cluster_params_r
    cp = splittable_cluser.cluster_params

    cp.posterior_hyperparams = calc_posterior(cp.hyperparams, cp.suff_statistics)
    cpl.posterior_hyperparams = calc_posterior(cpl.hyperparams, cpl.suff_statistics)
    cpr.posterior_hyperparams = calc_posterior(cpr.hyperparams, cpr.suff_statistics)

end

function create_suff_stats_dict_worker(
    group_pts,
    group_labels,
    group_sublabels,
    hyper_params,
    indices,
)
    suff_stats_dict = Dict()
    if indices == nothing
        indices = collect(1:length(clusters_vector))
    end

    points = localpart(group_pts)
    labels = localpart(group_labels)
    sublabels = localpart(group_sublabels)
    for index in indices
        pts = @view points[:, labels.==index]
        sub_labels = @view sublabels[labels.==index]

        cpl_suff = create_sufficient_statistics(
            hyper_params,
            hyper_params,
            @view pts[:, sub_labels.==1]
        )
        cpr_suff = create_sufficient_statistics(
            hyper_params,
            hyper_params,
            @view pts[:, sub_labels.==2]
        )
        cp_suff = create_sufficient_statistics(hyper_params, hyper_params, pts)
        suff_stats_dict[index] = thin_suff_stats(cp_suff, cpl_suff, cpr_suff)
    end
    return suff_stats_dict
end

function create_suff_stats_dict_node_leader(
    group_pts,
    group_labels,
    group_sublabels,
    hyper_params,
    proc_ids,
    indices,
)
    leader_suff_dict = Dict()
    workers_suff_dict = Dict()
    if indices == nothing
        indices = collect(1:length(clusters_vector))
    end
    for i in proc_ids
        workers_suff_dict[i] = remotecall(
            create_suff_stats_dict_worker,
            i,
            group_pts,
            group_labels,
            group_sublabels,
            hyper_params,
            indices,
        )
    end
    suff_stats_vectors = [[] for i = 1:length(indices)]
    cluster_to_index = Dict([indices[i] => i for i = 1:length(indices)])
    workers_suff_dict_fetched = Dict([k => fetch(v) for (k, v) in workers_suff_dict])
    for (k, v) in workers_suff_dict_fetched
        for (cluster, suff) in v
            push!(suff_stats_vectors[cluster_to_index[cluster]], suff)
        end
    end
    for (k, v) in enumerate(suff_stats_vectors)
        if length(v) > 0
            cp_suff = reduce(aggregate_suff_stats, [x.cluster_suff for x in v])
            cpl_suff = reduce(aggregate_suff_stats, [x.l_suff for x in v])
            cpr_suff = reduce(aggregate_suff_stats, [x.r_suff for x in v])
            cluster_index = indices[k]
            leader_suff_dict[cluster_index] = thin_suff_stats(cp_suff, cpl_suff, cpr_suff)
        end
    end

    return leader_suff_dict
end

function update_suff_stats_posterior!(
    group::local_group,
    indices = nothing,
    use_leader::Bool = true,
)
    workers_suff_dict = Dict()
    if indices == nothing
        indices = collect(1:length(group.local_clusters))
    end
    if use_leader
        for i in collect(keys(leader_dict))
            workers_suff_dict[i] = remotecall(
                create_suff_stats_dict_node_leader,
                i,
                group.points,
                group.labels,
                group.labels_subcluster,
                group.model_hyperparams.distribution_hyper_params,
                leader_dict[i],
                indices,
            )
        end
    else
        for i in (nworkers() == 0 ? procs() : workers())
            workers_suff_dict[i] = @spawnat i create_suff_stats_dict_worker(
                group.points,
                group.labels,
                group.labels_subcluster,
                group.model_hyperparams.distribution_hyper_params,
                indices,
            )
        end
    end
    suff_stats_vectors = [[] for i = 1:length(indices)]
    cluster_to_index = Dict([indices[i] => i for i = 1:length(indices)])
    workers_suff_dict_fetched = Dict([k => fetch(v) for (k, v) in workers_suff_dict])
    for (k, v) in workers_suff_dict_fetched
        for (cluster, suff) in v
            push!(suff_stats_vectors[cluster_to_index[cluster]], suff)
        end
    end
    for (index, v) in enumerate(indices)
        # if outlier_mod > 0 && v == 1
        #     continue
        # end
        if length(suff_stats_vectors[index]) == 0
            continue
        end
        cluster = group.local_clusters[v]
        cp = cluster.cluster_params
        cp.cluster_params.suff_statistics = reduce(
            aggregate_suff_stats,
            [x.cluster_suff for x in suff_stats_vectors[index]],
        )
        cp.cluster_params_l.suff_statistics =
            reduce(aggregate_suff_stats, [x.l_suff for x in suff_stats_vectors[index]])
        cp.cluster_params_r.suff_statistics =
            reduce(aggregate_suff_stats, [x.r_suff for x in suff_stats_vectors[index]])
        cluster.points_count = cp.cluster_params.suff_statistics.N
        update_splittable_cluster_params!(cluster.cluster_params)
    end

end

function split_first_cluster_worker!(group::local_group)
    sub_labels = localpart(group.labels_subcluster)
    pts = localpart(group.points)
    sub_labels .= rand(1:2, length(sub_labels))
end

function split_cluster_local_worker!(
    labels,
    sub_labels,
    points,
    indices::Vector{Int64},
    new_indices::Vector{Int64},
)
    labels = localpart(labels)
    sub_labels = localpart(sub_labels)
    pts = localpart(points)
    for (i, index) in enumerate(indices)
        cluster_sub_labels = @view sub_labels[labels.==index]
        cluster_labels = @view labels[labels.==index]
        cluster_points = @view pts[:, labels.==index]
        cluster_labels[cluster_sub_labels.==2] .= new_indices[i]

        cluster_sub_labels .= rand(1:2, length(cluster_sub_labels))

    end
end

function split_cluster_local!(
    group::local_group,
    cluster::local_cluster,
    index::Int64,
    new_index::Int64,
)

    l_split = copy_local_cluster(cluster)
    l_split.cluster_params = create_splittable_from_params(
        cluster.cluster_params.cluster_params_r,
        group.model_hyperparams.α,
    )
    cluster.cluster_params = create_splittable_from_params(
        cluster.cluster_params.cluster_params_l,
        group.model_hyperparams.α,
    )

    l_split.points_count = l_split.cluster_params.cluster_params.suff_statistics.N
    cluster.points_count = cluster.cluster_params.cluster_params.suff_statistics.N

    group.local_clusters[new_index] = l_split

end

function merge_clusters_worker!(
    group::local_group,
    indices::Vector{Int64},
    new_indices::Vector{Int64},
)
    labels = localpart(group.labels)
    sub_labels = localpart(group.labels_subcluster)

    for (i, index) in enumerate(indices)
        cluster_sub_labels = @view localpart(group.labels_subcluster)[labels.==index]
        cluster_sub_labels .= 1
        cluster_sub_labels =
            @view localpart(group.labels_subcluster)[labels.==new_indices[i]]
        cluster_sub_labels .= 2
        labels[labels.==new_indices[i]] .= index
    end
end

function merge_clusters!(group::local_group, index_l::Int64, index_r::Int64)
    new_splittable_cluster = merge_clusters_to_splittable(
        group.local_clusters[index_l].cluster_params.cluster_params,
        group.local_clusters[index_r].cluster_params.cluster_params,
        group.model_hyperparams.α,
    )
    group.local_clusters[index_l].cluster_params = new_splittable_cluster
    group.local_clusters[index_l].points_count += group.local_clusters[index_r].points_count
    group.local_clusters[index_r].points_count = 0
    group.local_clusters[index_r].cluster_params.cluster_params.suff_statistics.N = 0
    group.local_clusters[index_r].cluster_params.splittable = false
end

function should_split_local!(
    should_split::AbstractArray{Float32,1},
    cluster_params::splittable_cluster_params,
    α::Float32,
    final::Bool,
)
    cpl = cluster_params.cluster_params_l
    cpr = cluster_params.cluster_params_r
    cp = cluster_params.cluster_params
    if final || cpl.suff_statistics.N == 0 || cpr.suff_statistics.N == 0
        should_split .= 0
        return
    end
    post = calc_posterior(cp.hyperparams, cp.suff_statistics)
    lpost = calc_posterior(cp.hyperparams, cpl.suff_statistics)
    rpost = calc_posterior(cp.hyperparams, cpr.suff_statistics)

    log_likihood_l = log_marginal_likelihood(cpl.hyperparams, lpost, cpl.suff_statistics)
    log_likihood_r = log_marginal_likelihood(cpr.hyperparams, rpost, cpr.suff_statistics)
    log_likihood = log_marginal_likelihood(cp.hyperparams, post, cp.suff_statistics)
    log_HR =
        log(α) +
        logabsgamma(cpl.suff_statistics.N)[1] +
        log_likihood_l +
        logabsgamma(cpr.suff_statistics.N)[1] +
        log_likihood_r - (logabsgamma(cp.suff_statistics.N)[1] + log_likihood)
    if log_HR > log(rand())
        should_split .= 1
    end
    # print("LOG HR IS: ")
    # println(log_HR)
end

function check_and_split!(group::local_group, final::Bool, init_type::AbstractString)
    split_arr = zeros(Float32, length(group.local_clusters))
    for (index, cluster) in enumerate(group.local_clusters)
        if outlier_mod > 0 && index == 1
            continue
        end
        if cluster.cluster_params.splittable == true &&
           cluster.cluster_params.cluster_params.suff_statistics.N > 1
            should_split_local!(
                (@view split_arr[index, :]),
                cluster.cluster_params,
                group.model_hyperparams.α,
                final,
            )
        end
    end
    new_index = length(group.local_clusters) + 1
    indices = Vector{Int64}()
    new_indices = Vector{Int64}()
    resize!(group.local_clusters, Int64(length(group.local_clusters) + sum(split_arr)))
    for i = 1:length(split_arr)
        if split_arr[i] == 1

            push!(indices, i)
            push!(new_indices, new_index)
            split_cluster_local!(group, group.local_clusters[i], i, new_index)
            new_index += 1
        end
    end
    all_indices = vcat(indices, new_indices)
    if length(indices) > 0
        for i in (nworkers() == 0 ? procs() : workers())
            @spawnat i split_cluster_local_worker!(
                group.labels,
                group.labels_subcluster,
                group.points,
                indices,
                new_indices,
            )
        end

        # INSERT HERE:
        if init_type == "splitnet_2d"
            for i in all_indices
                println("doing splitnet 2d init")
                splitnet_2d_init!(group,i)
            end
        elseif init_type == "splitnet_10d"
            for i in all_indices
                println("doing splitnet 10d init")
                splitnet_10d_init!(group,i)
            end
        elseif init_type == "splitnet_128d"
            for i in all_indices
                println("doing splitnet 128d init")
                splitnet_128d_init!(group,i)
            end
        elseif init_type == "eigensplit"
            for i in all_indices
                println("doing eigensplit init")
                eigensplit_init!(group,i)
            end
        elseif init_type == "kmeans"
            for i in all_indices
                println("doing kmeans init")
                kmeans_init!(group,i)
            end
        end

    end
    return all_indices
end

function check_and_merge!(group::local_group, final::Bool)
    mergable = zeros(Float32, 1)
    indices = Vector{Int64}()
    new_indices = Vector{Int64}()
    for i = 1:length(group.local_clusters)
        if outlier_mod > 0 && i == 1
            continue
        end
        for j = i+1:length(group.local_clusters)
            if (
                group.local_clusters[i].cluster_params.splittable == true &&
                group.local_clusters[j].cluster_params.splittable == true &&
                group.local_clusters[i].cluster_params.cluster_params.suff_statistics.N >
                0 &&
                group.local_clusters[j].cluster_params.cluster_params.suff_statistics.N > 0
            )
                should_merge!(
                    mergable,
                    group.local_clusters[i].cluster_params.cluster_params,
                    group.local_clusters[j].cluster_params.cluster_params,
                    group.model_hyperparams.α,
                    final,
                )
            end
            if mergable[1] == 1
                merge_clusters!(group, i, j)
                push!(indices, i)
                push!(new_indices, j)
            end
            mergable[1] = 0
        end
    end
    for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i merge_clusters_worker!(group, indices, new_indices)
    end
    return indices
end

function sample_clusters!(group::local_group, first::Bool)
    points_count = Vector{Float32}()
    local_workers = procs(1)[2:end]
    cluster_params_futures = Dict()
    for (i, cluster) in enumerate(group.local_clusters)
        cluster_params_futures[i] =
            sample_cluster_params(cluster.cluster_params, group.model_hyperparams.α, first)
    end
    for (i, cluster) in enumerate(group.local_clusters)
        if outlier_mod > 0 && i == 1
            continue
        end
        cluster.cluster_params = fetch(cluster_params_futures[i])
        cluster.points_count = cluster.cluster_params.cluster_params.suff_statistics.N
        push!(points_count, cluster.points_count)
    end
    push!(points_count, group.model_hyperparams.α)
    group.weights = rand(Dirichlet(Float64.(points_count)))[1:end-1] .* (1 - outlier_mod)
    if outlier_mod > 0
        group.weights = vcat([outlier_mod], group.weights)
    end
end

function create_thin_cluster_params(cluster::local_cluster)
    return thin_cluster_params(
        cluster.cluster_params.cluster_params.distribution,
        cluster.cluster_params.cluster_params_l.distribution,
        cluster.cluster_params.cluster_params_r.distribution,
        cluster.cluster_params.lr_weights,
    )
end

function remove_empty_clusters_worker!(labels, pts_count)
    labels = localpart(labels)
    removed = 0
    for (index, count) in enumerate(pts_count)
        if count == 0
            labels[labels.>index-removed] .-= 1
            removed += 1
        end
    end
end

function remove_empty_clusters!(group::local_group)
    new_vec = Vector{local_cluster}()
    removed = 0
    pts_count = Vector{Int64}()
    for (index, cluster) in enumerate(group.local_clusters)
        push!(pts_count, cluster.points_count)
        if cluster.points_count > 0 ||
           (outlier_mod > 0 && index == 1) ||
           (outlier_mod > 0 && index == 2 && length(group.local_clusters) == 2)
            push!(new_vec, cluster)
        end
    end
    @sync for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i remove_empty_clusters_worker!(group.labels, pts_count)
    end
    group.local_clusters = new_vec
end

function rand_subclusters_labels!(labels::AbstractArray{Int64,1})    #lr_arr = create_array(zeros(Float32,length(labels), 2))
    if size(labels, 1) == 0
        return
    end
    labels .= rand(1:2, size(labels, 1))
end

function reset_bad_clusters_worker!(
    indices::Vector{Int64},
    group_points::AbstractArray{Float32,2},
    group_labels::AbstractArray{Int64,1},
    group_labels_subcluster::AbstractArray{Int64,1},
)
    labels = localpart(group_labels)
    sub_labels = localpart(group_labels_subcluster)
    pts = localpart(group_points)
    for i in indices
        rand_subclusters_labels!((@view sub_labels[labels.==i]))
    end
end

function reset_splitted_clusters!(group::local_group, bad_clusters::Vector{Int64})

    for i in bad_clusters
        group.local_clusters[i].cluster_params.logsublikelihood_hist =
            ones(burnout_period + 5) * -Inf
    end
    for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i reset_bad_clusters_worker!(
            bad_clusters,
            group.points,
            group.labels,
            group.labels_subcluster,
        )
    end
    update_suff_stats_posterior!(group, bad_clusters)
end

function reset_bad_clusters!(group::local_group)
    bad_clusters = Vector{Int64}()
    for (i, c) in enumerate(group.local_clusters)
        cl = c.cluster_params.cluster_params_l
        cr = c.cluster_params.cluster_params_r
        if cl.suff_statistics.N == 0 || cr.suff_statistics.N == 0
            push!(bad_clusters, i)
            c.cluster_params.logsublikelihood_hist = ones(burnout_period + 5) * -Inf
            c.cluster_params.splittable = false
        end
    end
    @sync for i in (nworkers() == 0 ? procs() : workers())
        @spawnat i reset_bad_clusters_worker!(
            bad_clusters,
            group.points,
            group.labels,
            group.labels_subcluster,
        )
    end
    update_suff_stats_posterior!(group, bad_clusters)
end

function broadcast_cluster_params(params_vector, weights_vector)
    @sync for (k, v) in leader_dict
        @spawnat k broadcast_to_node(params_vector, weights_vector, v)
    end
end

#Below code is to prevent some bug with deseraliztion
function broadcast_to_node(params_vector, weights_vector, proc_ids)
    refs = Dict()
    @sync for i in proc_ids
        refs[i] = @spawnat i set_global_data(params_vector, weights_vector)
    end
    alltrue = false
    responses = Dict([k => fetch(v) for (k, v) in refs])
    for (k, v) in responses
        while v == false
            v = remotecall_fetch(set_global_data, k, params_vector, weights_vector)
        end
    end
end

function set_global_data(params_vector, weights_vector)
    succ = true
    try
        global clusters_vector = params_vector
        global clusters_weights = weights_vector
    catch
        succ = false
    end
    return succ
end

function group_step(group::local_group, no_more_splits::Bool, final::Bool, first::Bool, init_type)
    sample_clusters!(group, false)
    broadcast_cluster_params(
        [create_thin_cluster_params(x) for x in group.local_clusters],
        group.weights,
    )
    sample_labels!(group, (hard_clustering ? true : final), no_more_splits)
    sample_sub_clusters!(group)
    update_suff_stats_posterior!(group)
    reset_bad_clusters!(group)
    if no_more_splits == false
        indices = []
        indices = check_and_split!(group, final, init_type)
        update_suff_stats_posterior!(group, indices)
        check_and_merge!(group, final)
    end
    remove_empty_clusters!(group)
    return
end

# PyTORCH MODEL:

function splitnet_2d_init!(group::local_group,cluster_num::Int64)
    pts = Array(group.points)
    labels = Array(group.labels)
    cluster_pts = pts[:,labels .==cluster_num]

    y = torch_model_2d.infer(cluster_pts)
    # preds = argmax(y, dims=2)
    preds = round.(y) .+ 1
    # split_labels =  dropdims([ i for i in preds ], dims=2)
    split_labels =  dropdims(preds, dims=2)

    sub_labels = localpart(group.labels_subcluster)
    sub_labels[labels .== cluster_num] .= split_labels

    # println("SPLITNET COUNT")
    # print(count(split_labels.==1))
    # print(count(split_labels.==2))
end

function splitnet_10d_init!(group::local_group,cluster_num::Int64)
    pts = Array(group.points)
    labels = Array(group.labels)
    cluster_pts = pts[:,labels .==cluster_num]

    y = torch_model_10d.infer(cluster_pts)
    # preds = argmax(y, dims=2)
    preds = round.(y) .+ 1
    # split_labels =  dropdims([ i for i in preds ], dims=2)
    split_labels =  dropdims(preds, dims=2)

    sub_labels = localpart(group.labels_subcluster)
    sub_labels[labels .== cluster_num] .= split_labels
end

function splitnet_128d_init!(group::local_group,cluster_num::Int64)
    pts = Array(group.points)
    labels = Array(group.labels)
    cluster_pts = pts[:,labels .==cluster_num]

    y = torch_model_128d.infer(cluster_pts)
    # preds = argmax(y, dims=2)
    preds = round.(y) .+ 1
    # split_labels =  dropdims([ i for i in preds ], dims=2)
    split_labels =  dropdims(preds, dims=2)

    sub_labels = localpart(group.labels_subcluster)
    sub_labels[labels .== cluster_num] .= split_labels
end

# KMEANS:

function kmeans_init!(group::local_group,cluster_num::Int64)
    pts = Array(group.points)
    labels = Array(group.labels)
    cluster_pts = pts[:,labels .==cluster_num]

    if size(cluster_pts,2) == 1
        return [1]
    end

    split_labels = kmeans(cluster_pts, 2).assignments
    sub_labels = localpart(group.labels_subcluster)
    sub_labels[labels .== cluster_num] .= split_labels

    # println("KMEANS COUNT")
    # print(count(split_labels.==1))
    # print(count(split_labels.==2))
end

# EIGENSPLIT

function eigensplit_init!(group::local_group,cluster_num::Int64)
    N = group.local_clusters[cluster_num].cluster_params.cluster_params.suff_statistics.N
    XXT = group.local_clusters[cluster_num].cluster_params.cluster_params.suff_statistics.S / N
    μ = group.local_clusters[cluster_num].cluster_params.cluster_params.suff_statistics.points_sum / N
    M = XXT - μ*μ'

    F = eigen(M)

    vals1 = copy(F.values)
    vals2 = copy(F.values)

    vals1[end] = 1e-5
    vals2[end-1] = 1e-5

    vals1, vals2

    subcov1 = F.vectors * Diagonal(vals1) * F.vectors'
    subcov2 = F.vectors * Diagonal(vals2) * F.vectors'
    invcov1 = inv(subcov1)
    invcov2 = inv(subcov2)

    @sync for i in workers()
        @spawnat i cov_split_worker!(group, cluster_num, invcov1,invcov2)
    end
end

function cov_split_worker!(group::local_group,cluster_num::Int64, invcov1, invcov2)
    malhan_dist1 = x -> x' * invcov1 * x
    malhan_dist2 = x -> x' * invcov2 * x
    one_or_two = x -> malhan_dist1(x) > malhan_dist2(x) ? 1 : 2
    labels = localpart(group.labels)
    sublabels = @view localpart(group.labels_subcluster)[labels .== cluster_num]
    pts = @view localpart(group.points)[:,labels .== cluster_num]
    for i=1:length(sublabels)
        sublabels[i] = one_or_two(pts[:,i])
    end
end

# function eigensplit_init!(group::local_group,cluster_num::Int64)
#     N = group.local_clusters[cluster_num].cluster_params.cluster_params.suff_statistics.N
#     XXT = group.local_clusters[cluster_num].cluster_params.cluster_params.suff_statistics.S / N
#     μ = group.local_clusters[cluster_num].cluster_params.cluster_params.suff_statistics.points_sum / N
#     M = XXT - μ*μ'
#     # println("vecs:")
#     F = eigen(M)
#     vecs = F.vectors
#     vals = F.values
#     # println("eigsolve:")
#     # v1 = eigsolve(M, 1, maxiter = 10)[2][1]
#     mxindx = findmax(vals)[2]
#     v1=vecs[mxindx,:]
#     means_futures = Vector{Future}(undef,nworkers())
#     for i in workers()
#         means_futures[i-1] =  @spawnat i tranform_points_worker!(group.points,group.labels,(@eval $cluster_num),(@eval $v1),(@eval $μ))
#     end
#     min_pts = [fetch(x) != nothing ? fetch(x)[1] : nothing for x in means_futures if fetch(x) != nothing]
#     max_pts = [fetch(x) != nothing ? fetch(x)[2] : nothing for x in means_futures if fetch(x) != nothing]
#     if length(min_pts) == 0
#         return
#     end

#     min_mean = minimum(min_pts)
#     max_mean = maximum(max_pts)
#     # min_mean = -100
#     # max_mean = +100
#     #Distributed K-Means
#     iter = 0
#     converged = false
#     while(iter < max_split_iter && !converged)
#         means_futures = Vector{Future}(undef,nworkers())
#         @sync for i in workers()
#             means_futures[i-1] = @spawnat i kmeans_iter_worker!(min_mean,max_mean)
#         end
#         min_means = [fetch(x) != nothing ? fetch(x)[1] : nothing for x in means_futures if fetch(x) != nothing]
#         max_means = [fetch(x) != nothing ? fetch(x)[2] : nothing for x in means_futures if fetch(x) != nothing]
#         min_sum = 0
#         min_count = 0
#         for x in min_means
#             min_sum += x[1]
#             min_count += x[2]
#         end
#         max_sum = 0
#         max_count = 0
#         for x in max_means
#             max_sum += x[1]
#             max_count += x[2]
#         end
#         new_min = min_sum / min_count
#         new_max = max_sum / max_count
#         if new_min == min_mean && new_max == max_mean
#             converged = true
#         else
#             min_mean = new_min
#             max_mean = new_max
#         end
#         iter += 1
#     end
#     @sync for i in workers()
#         @spawnat i set_smart_labels_in_worker!(cluster_num,group.labels, group.labels_subcluster,group.points)
#     end
# end

function set_smart_labels_in_worker!(cluster_num,labels,sub_labels,points)
    labels = localpart(labels)
    sub_labels = localpart(sub_labels)
    sub_labels[labels .== cluster_num] .= split_labels
end

function tranform_points_worker!(all_pts,all_labels,cluster_num::Int64,v,μ)

    pts = localpart(all_pts)
    labels = localpart(all_labels)
    rel_points = pts[:,labels .== cluster_num] .- μ
    global transformed_pts = v' * rel_points

    if length(transformed_pts) > 1
        return percentile((@view transformed_pts[:]),0.10), percentile((@view transformed_pts[:]),0.90)
    end
    return nothing
end

function kmeans_iter_worker!(min_mean,max_mean)
    global split_labels = [(abs(x - min_mean) < abs(x-max_mean) ? 1 : 2) for x in (@view transformed_pts[:])]
    if length(split_labels) > 0
        return (sum(transformed_pts[split_labels .== 1]), sum(split_labels .== 1)), (sum(transformed_pts[split_labels .== 2]), sum(split_labels .== 2))
    end
    return nothing
end

function eigens_init!(group::local_group,cluster_num::Int64)
    pts = Array(group.points)
    labels = Array(group.labels)
    cluster_pts = pts[:,labels .==cluster_num]

    if size(points,2) == 1
        return [1]
    end

    m = MultivariateStats.fit(PCA,points; maxoutdim=1)
    v1=projection(m)
    new_pts = transform(m,points)
    a,minpt = findmin(new_pts)
    a,maxpt = findmax(new_pts)
    minpt = minpt[2]
    maxpt = maxpt[2]
    result = kmeans(new_pts, 2, init = [minpt,maxpt])
    return result.assignments
end