using LinearAlgebra
using Distributions
using StatsFuns
using Plots

function individual_loss(Sᵢ, S, D,n=1)
    #= 
        This function simulates how many satellites will be lost 
        from a given constellation.

        TODO: Requires calibration of some sort
    =#
    β₀ = 3e-5 #probability of spontaneous destruction
    β₁ = 1e-6 #proability of collision with satellite in same constellation
    β₂ = 1e-5 #probability of collision with satellite in different constellation
    β₃ = 1e-4 #probability of collision with debris

    #linear probability
    p =  β₀ + β₁ * Sᵢ + β₂ * (S-Sᵢ) + β₃ * D

    dist = Binomial(Sᵢ, p)
    rand(dist)
end


function debris_transition(D, Y⃗, X⃗, δ=0.05, Γ=0.5)
    #δ: 1/δ year decay rate
    #Γ: proportion of the mass of the satellite that is released into orbit during launch

    #new debris released equals the debris lost to deorbit + debris generated by collisions + debris due to launches 
    new_debris = (1-δ) * D .+ sum(Y⃗) .+ Γ*sum(X⃗)
    return new_debris
end


function stocks_transition(S⃗, Y⃗, X⃗)
    return S⃗ .+ X⃗ .- Y⃗
end


function overall_losses(S⃗, D)
    #=
    Simulate the losses for each constellation
    =#
    [ individual_loss(s, sum(S⃗),D) for s in S⃗ ]
end


function inverse_demand_cournot(S⃗,A=100)
    A - sum(S⃗)
end


function profit_function(s,x,C_operation,C_launch,T_launch, T_operation,Price)
    (Price-T_operation-C_operation)*s - (C_launch + T_launch)*x
end





#Get histogram of loss numbers per 10_000
histogram(
    [sum(individual_loss(33,67,500,10_000)) for i in 1:10000]
    ,bins=range(0,200,length=201)
    ,label="Satellites lost per 10,000"
    )



#=
Basic simulations
1. What happens when both participants always launch a single satellite?
2. What happens when a single participant always launches a single satellite?
3. What happens when both participants maintain the cournot equilibrium?
=#

#First try

policy_1(s,S,D) = 1
function policy_2(s,S,D) 
    max(0,33-s)
end
function policy_3(s,S,D) 
    max(0,20-s)
end

#setup state
n=200#00
k=4
state = fill([10.0,1,0,10],n)
profit_path = []

for i in 2:n
    current_state = state[i-1]
    #current_state = state[1]
    S⃗ = current_state[2:end]
    D = current_state[1]

    #Calculate policies
    #X = [policy_2(s,sum(S⃗),D) for s in S⃗]
    X = [2.0,2.0,policy_2(S⃗[3],0,0)]
    #X = [policy_2(S⃗[1],sum(S⃗),D), policy_3(S⃗[2],sum(S⃗),D), ]

    #Information generation
    price = inverse_demand_cournot(S⃗,10000)
    push!(profit_path,[profit_function(s,x,1,2,0, 0,price) for (s,x) in zip(S⃗,X) ])

    #calculate transition to next state
    #loss of satellites
    Y⃗ = overall_losses(S⃗, D)

    #debris
    new_debris = debris_transition(D, Y⃗, X)

    #stocks
    new_stocks = stocks_transition(S⃗,Y⃗,X)

    #update vector of states
    next_state = [new_debris,new_stocks...]

    #update vector
    state[i] = next_state
    println(current_state,next_state)
end


debris = [v[1] for v in state]
stocks1 = [v[2] for v in state]
stocks2 = [v[3] for v in state]
stocks3 = [v[4] for v in state]

plot(
    [debris, stocks1, stocks2, stocks3]
    ,labels=["debris" "stocks 1" "stocks 2" "Stocks 3"]
)

profit_1 = [v[1] for v in profit_path] 
profit_2 = [v[2] for v in profit_path]
profit_3 = [v[3] for v in profit_path]

plot(
    [profit_1,profit_2,profit_3]
    )