include("../src/Orbits.jl")
using .Orbits
using LinearAlgebra

#Set key dimensions
const N_constellations = 4
const N_debris = 1
const N_states = N_constellations + N_debris

#Setup Economic Models
em2_a = LinearModel(0.95, [1 -0.02 -0.02 0], [5.0 0 0 0])
em2_b = LinearModel(0.95, [-0.02 1 -0.02 0], [0.0 5 0 0])
em2_c = LinearModel(0.95, [0 -0.02 1 -0.02], [0.0 0 5 0])
em2_d = LinearModel(0.95, [0 -0.02 -0.02 1], [0.0 0 0 5])

#Setup Physics

basic_model = BasicPhysics(
    0.002
    ,0.002*(ones(N_constellations,N_constellations) - LinearAlgebra.I)
    ,0.01
    ,0.001
    ,5.0
    ,0.05
)

# Setup NN
bg = BranchGenerator(N_constellations)
operators_policy = bg(operator_policy_function_generator(N_constellations,N_debris),vcat)
planners_policy = bg(operator_policy_function_generator(N_constellations,N_debris),vcat)
planners_value = value_function_generator(64)


# Setup Operators
const operator_array = [ 
    #first operator
    OperatorLoss(		
        em2_a
        ,value_function_generator()
        ,operators_policy #this is held by all operators
        ,params(operators_policy[2][1]) #first operator gets first branch of params
        ,basic_model
    )
    ,OperatorLoss(		
        em2_b
        ,value_function_generator()
        ,operators_policy #this is held by all operators
        ,params(operators_policy[2][2]) #first operator gets first branch of params
        ,basic_model
    )
    ,OperatorLoss(		
        em2_c
        ,value_function_generator()
        ,operators_policy #this is held by all operators
        ,params(operators_policy[2][3]) #first operator gets first branch of params
        ,basic_model
    )
    ,OperatorLoss(		
        em2_d
        ,value_function_generator()
        ,operators_policy #this is held by all operators
        ,params(operators_policy[2][4]) #first operator gets first branch of params
        ,basic_model
    )
]

#sanity check time
@assert length(operator_array) == N_constellations "Mismatch in predetermined number of constellations and the number of operators initialized"

# Setup Planner
pl = PlannerLoss(
		0.95
		,operator_array
		,planners_policy
		,params(planners_policy)
		,planners_value
		,params(planners_value)
        ,basic_model
)

# Export Planner