#!/usr/bin/env julia
using CSV, DataFrames
using JuMP, Ipopt
using LinearAlgebra
using SparseArrays
using Mosek, MosekTools
using Distributions
using StatsBase

### load functions
# data scripts
include("scr/fun_data.jl")
include("scr/fun_auxiliary.jl")
# optimization models + supporting functions
include("scr/fun_opt_models.jl")
include("scr/fun_linearization.jl")
# out-of-sample + projection functions
include("scr/fun_post_processing.jl")

### experiment settings
exp_settings = Dict(
# test case
:case => "case_48",
# pressure variance penalty
:Ïˆ_ð�›‘ => 0.1,
# flow variance penalty
:Ïˆ_Ï† => 0,
# prescribed constraint violation probability
:Îµ => 0.01,
# gas extraction standard deviation (% of the nominal level)
:Ïƒ => 0.1,
# det: true - optimize determenistic policies, false - optimize chance-constrained policies
:det => false,
# comp: true - compressor recourse on, false - compressor recourse off
:comp => true,
# valv: true - valve recourse on, false - valve recourse off
:valv => true,
# proj: true - run projection analysis, false - don't run
:proj => false,
# sample compexity
:S => 1000
)

### display input
println()
println("------ Experiment settings ------")
exp_settings[:det] == true ?
println("Policy optimization      :     deterministic") :
println("Policy optimization      :     chance-constrained")
println("Variance penalty         :     Ïˆ_ð�›‘ = $(exp_settings[:Ïˆ_ð�›‘]) ... Ïˆ_Ï† = $(exp_settings[:Ïˆ_Ï†])")
println("Violation probability    :     Îµ = $(exp_settings[:Îµ])")
println("Standard deviation       :     Ïƒ = $(exp_settings[:Ïƒ])")
println("Active pipeline recourse :     comp = $(exp_settings[:comp]) ... valve = $(exp_settings[:valv])")
println("Out of sample setting    :     proj = $(exp_settings[:proj]) ... S = $(exp_settings[:S])")

### run experiment
println()
println("------ Run experiment ------")
# extarct network data
net_data        = load_network_data(exp_settings[:case])
# solve non-convex gas network optimization
sol_non_convex  = non_convex_opt(net_data)
# obtain linearization data
lin_res         = linearization(net_data,sol_non_convex[:model])
# extract forecast data
forecast        = extract_forecast_data(net_data,exp_settings)
# solve chance-constrained gas network optimization
sol_stochastic  = chance_con_opt(net_data,lin_res,forecast,exp_settings)
# get dual solution to the chance-constrained gas network optimization
sol_dual        = stochastic_dual_solution(net_data,sol_stochastic,lin_res,forecast,exp_settings)
# run out of sample analysis
sol_ofs         = out_of_sample(net_data,forecast,sol_stochastic)
# run projection analysis
exp_settings[:proj] == true ? sol_proj        = projection_analysis(net_data,forecast,sol_ofs) : NaN

### display output
println()
println("------ Experiment results ------")
println("Expected cost (\$)              :   $(round(sol_stochastic[:cost],digits=1))")
println("Total pressure variance (MPaÂ²) :   $(round(sum([var(sol_ofs[:Ï�][n,:]) for n in net_data[:N]]),digits=1))")
println("Total flow variance (BMSCFDÂ²)  :   $(round(sum([var(sol_ofs[:Ï†][l,:]) for l in net_data[:E]]),digits=1))")
println("Average comp deployment (kPa)  :   $(round(sum([sqrt(abs(sol_stochastic[:Îº][l])) for l in findall(x->x>0, net_data[:ÎºÌ…])])))")
println("Average valv deployment (kPa)  :   $(round(sum([sqrt(abs(sol_stochastic[:Îº][l])) for l in findall(x->x<0, net_data[:ÎºÌ²])])))")
println("Emp. constraint violation (%)  :   $(round(sol_ofs[:Îµ_stat]*100,digits=5))")
exp_settings[:proj] == true ?
println("Average injection proj (MMSCFD) :   $(round(sol_proj[:Î”Ï‘_mean],digits=3))") : NaN
exp_settings[:proj] == true ?
println("Average regulation proj (kPa)   :   $(round(sol_proj[:Î”Îº_mean],digits=3))") : NaN
println("Total revenue of suppliers (\$) :   $(round(sol_dual[:R_inj],digits=3))")
println("Total revenue of act. pipes (\$):   $(round(sol_dual[:R_act],digits=3))")
println("Total revenue of consumers (\$) :   $(round(sol_dual[:R_con],digits=3))")
