### A Pluto.jl notebook ###
# v0.17.7

using Markdown
using InteractiveUtils

# ╔═╡ 3ce6b18e-deae-44fa-8824-4193e7d6edb3
begin
	using Pkg
	Pkg.activate("../Project.toml")

	using PolyChaos
	using LinearAlgebra
	using DelimitedFiles
	using StatsPlots

	using Plots.PlotMeasures
	using Printf

	using CSV
	using DataFrames

	using NetCDF
	using Dates
	using LaTeXStrings
	using PlutoUI

	using Random
	using JLD
end

# ╔═╡ 893e00b2-0baf-49c2-9c94-884eb0607eb4
begin
	using Revise
	include("./gsa_utils.jl")
end

# ╔═╡ 95e46924-1e6d-11ed-152b-e373176bfe2b
## INCREASE THE CELL WIDTH FOR BIGGER FIGURES
html"""
<style>
	main {
		margin: 0 auto;
		max-width: 2000px;
    	padding-left: max(160px, 10%);
    	padding-right: max(160px, 10%);
	}
  svg {
    width: 100%;
  }
</style>
"""

# ╔═╡ e425f65c-a0a6-4316-b799-b495b73d5a97
md"""
Load required packages (dependencies specified in the `Project.toml` file)
"""

# ╔═╡ ea56d0e1-ef32-4582-908b-e9ba1f576c9d
md"""
Include script with different helper functions. If script is changed, this cell will have to be rerun for changes to reflect.
"""

# ╔═╡ 9415e163-19c0-453c-80fe-90ec1c1f4daf
md"""
Note: Figures are plotted here in the order that they appear in the manuscript.
"""

# ╔═╡ 429c13e8-f91c-41cc-b112-81b6a7f3d41a
md"""
### Preprocessing

"""

# ╔═╡ c3fb873b-2457-4317-83e1-a8dd22db0989
begin
	X_design = CSV.read("./data/design/X_background_CR2152_updated.csv", DataFrame)
	pfss = replace(X_design.PFSS, 1=>"HARMONICS", 2=>"FDIPS")
	surfaceWaveRefl = replace(X_design.UseSurfaceWaveRefl, 1=>"true", 2=>"false")
	select!(X_design, Not([:PFSS, :UseSurfaceWaveRefl]))
	insertcols!(X_design, :PFSS=>pfss, :UseSurfaceWaveRefl=>surfaceWaveRefl)
end

# ╔═╡ 7e81d77e-6aa8-40f2-a900-22f76a996185
md"""
Make selected scatterplots - figure 2 (others can be made in the same fashion by selecting appropriate column names in the plotting arguments)
"""

# ╔═╡ 8b66c931-a842-4011-a06c-93ac0e1445e5
begin
	FactorB0PF = scatter(X_design[:, "FactorB0"], X_design[:, "PoyntingFluxPerBSi"], 
								# zcolor=shiftWLRMSE.PTRMSE, 
								marker=(:black, :circle, 4), 
								xlabel="FactorB0", 
								ylabel="PoyntingFluxPerBSi",
								markerstrokewidth=0,
								label="",
								dpi = 300,
								grid=false
								)
	plot!(sort(X_design.FactorB0), 9e5 ./ (sort(X_design.FactorB0)), line=(:violet, 3.1), label="")
	plot!(xlims=(0.54, 2.7))
	plot!(ylims=(0.3e6, 1.1e6))
	plot!(guidefontsize=30)
	plot!(tickfontsize=25)
	plot!(framestyle=:box)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(right_margin=8mm)
	plot!(yticks=([4e5, 6e5, 8e5, 1e6], ["4e5", "6e5", "8e5", "1e6"]))
	plot!(size=(800, 600))
end

# ╔═╡ 67a24d73-c444-496d-9337-cfbf9cde37c2
md"""
Figure for Supporting Information Document
"""

# ╔═╡ 2b218d46-6a44-4b34-9236-06fb30d4d96a
begin
	LperpStoch = scatter(X_design[:, "LperpTimesSqrtBSi"], X_design[:, "StochasticExponent"], 
									# zcolor=shiftWLRMSE.PTRMSE, 
									marker=(:black, :circle, 4), 
									xlabel="LperpTimesSqrtBSi", 
									ylabel="StochasticExponent",
									markerstrokewidth=0,
									label="",
									dpi = 300,
									grid=false
									)
	plot!(guidefontsize=30)
	plot!(tickfontsize=25)
	plot!(framestyle=:box)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(right_margin=8mm)
	plot!(xticks=([1e5, 2e5, 3e5], ["1e5", "2e5", "3e5"]))
	plot!(size=(800, 600))
end

# ╔═╡ 7b4efa4b-bea0-4517-a7c1-df969a4ca74e
md"""
Figure for supporting information document
"""

# ╔═╡ 82478091-c845-4b26-8092-19f1b9eb7776
begin
	chromoLperp = scatter(X_design[:, "nChromoSi_AWSoM"], X_design[:, "LperpTimesSqrtBSi"], 
									# zcolor=shiftWLRMSE.PTRMSE, 
									marker=(:black, :circle, 4), 
									xlabel="nChromoSi_AWSoM", 
									ylabel="LperpTimesSqrtBSi",
									markerstrokewidth=0,
									label="",
									dpi = 300,
									grid=false
									)
	plot!(guidefontsize=30)
	plot!(tickfontsize=25)
	plot!(framestyle=:box)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(right_margin=8mm)
	plot!(xticks=([1e18, 2e18, 3e18, 4e18, 5e18], ["1e18", "2e18", "3e18", "4e18", "5e18"]))
	plot!(yticks=([1e5, 2e5, 3e5], ["1e5", "2e5", "3e5"]))
	plot!(size=(800, 600))
end

# ╔═╡ 6620dfd1-f147-42ec-8d27-319e7f1cb5aa
begin
	LperpRMin = scatter(X_design[:, "LperpTimesSqrtBSi"], X_design[:, "rMinWaveReflection"], 
									# zcolor=shiftWLRMSE.PTRMSE, 
									marker=(:black, :circle, 4), 
									xlabel="LperpTimesSqrtBSi", 
									ylabel="rMinWaveReflection",
									markerstrokewidth=0,
									label="",
									dpi = 300,
									grid=false
									)
	plot!(guidefontsize=30)
	plot!(tickfontsize=25)
	plot!(framestyle=:box)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(right_margin=8mm)
	plot!(xticks=([1e5, 2e5, 3e5], ["1e5", "2e5", "3e5"]))
	plot!(size=(800, 600))
end

# ╔═╡ 84c005c3-9fa2-4cc4-80b7-105031ef327a
md"""
Scale input parameters to [0-1] range
"""

# ╔═╡ 8eb1989c-5a16-43fb-901e-55ea9663a7df
begin
		lbBg = [0.54, 2e17, 0.3e6, 0.3e5, 0.1, 1]
		ubBg = [2.7, 5e18, 1.1e6, 3e5, 0.34, 1.2]
		
		paramsBgScaled = (X_design[!, 1:6] .- lbBg') ./ (ubBg' - lbBg')
end

# ╔═╡ ba9fe174-0d0e-4396-baf6-088902335076
inputNames = names(paramsBgScaled)

# ╔═╡ fa98a60f-75d7-48d6-81ad-af6d287b1f2f
md"""
Filter out runs that are not physically meaningful based on the exclusion criteria described in section 4.1.
"""

# ╔═╡ dc313671-8a22-4ec1-b74a-36c15398c4b7
excludeRunIdx = [11, 14, 16, 20, 34, 48, 50, 54, 68, 72, 74, 88, 94, 101, 117, 124, 126, 137, 142, 146, 149, 164, 169, 175, 194, 198]

# ╔═╡ c607abf7-e20e-46a1-85ed-7d580a55d9db
begin
	FactorB0PFAll = scatter(X_design[:, "FactorB0"], X_design[:, "PoyntingFluxPerBSi"], 
								# zcolor=shiftWLRMSE.PTRMSE, 
								marker=(:black, :circle, 4), 
								xlabel="FactorB0", 
								ylabel="PoyntingFluxPerBSi",
								markerstrokewidth=0,
								label="",
								dpi = 300,
								grid=false
								)
	scatter!(X_design[excludeRunIdx, "FactorB0"], X_design[excludeRunIdx, "PoyntingFluxPerBSi"], label="Removed / failed")
	plot!(sort(X_design.FactorB0), 9e5 ./ (sort(X_design.FactorB0)), line=(:violet, 3.1), label="")
	plot!(xlims=(0.54, 2.7))
	plot!(ylims=(0.3e6, 1.1e6))
	plot!(guidefontsize=30)
	plot!(tickfontsize=25)
	plot!(framestyle=:box)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(right_margin=8mm)
	plot!(yticks=([4e5, 6e5, 8e5, 1e6], ["4e5", "6e5", "8e5", "1e6"]))
	plot!(size=(800, 600))
end

# ╔═╡ 7139a373-2795-436b-9bb3-6c4596254101
begin
	LperpStoch2 = scatter(X_design[:, "LperpTimesSqrtBSi"], X_design[:, "StochasticExponent"], 
									# zcolor=shiftWLRMSE.PTRMSE, 
									marker=(:black, :circle, 4), 
									xlabel="LperpTimesSqrtBSi", 
									ylabel="StochasticExponent",
									markerstrokewidth=0,
									label="",
									dpi = 300,
									grid=false
									)
	scatter!(X_design[excludeRunIdx, "LperpTimesSqrtBSi"], X_design[excludeRunIdx, "StochasticExponent"], label="Removed / failed")
	plot!(guidefontsize=30)
	plot!(tickfontsize=25)
	plot!(framestyle=:box)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(right_margin=8mm)
	plot!(xticks=([1e5, 2e5, 3e5], ["1e5", "2e5", "3e5"]))
	plot!(size=(800, 600))
end

# ╔═╡ 7ce46450-824e-415c-843b-0990848ab852
# CSV.write("/Users/ajivani/Desktop/Research/SWQUPaper/XDesignRemovedCR2152.csv", X_design[excludeRunIdx, :])

# ╔═╡ e3060049-736d-48da-9767-945d4185153b


# ╔═╡ 14263462-9183-4efc-bd40-957c1443b0ac
finalRuns = setdiff(1:200, excludeRunIdx)

# ╔═╡ 47a8931d-02fc-47ce-950e-f8338d87432c
md"""
Extract QoIs and observations
"""

# ╔═╡ 26379e4c-2fa3-42e6-87b9-29434b61387e
begin
	fn = "./data/background_sims/bg_CR2152.nc"
	UrSim = ncread(fn, "UrSim")
	NpSim = ncread(fn, "NpSim")
	TSim  = ncread(fn, "TSim")
	BSim  = ncread(fn, "BSim")	

	UrObs = ncread(fn, "UrObs")
	NpObs = ncread(fn, "NpObs")
	TObs  = ncread(fn, "TObs")
	BObs  = ncread(fn, "BObs")	
	
	startTime = ncgetatt(fn, "time", "startTime")
	timeElapsed = Dates.Hour.(ncread(fn, "time"))
	times = timeElapsed .+ Dates.DateTime(startTime, "yyyy_mm_ddTHH:MM:SS")
end

# ╔═╡ 678d55ec-afbc-4fb0-9099-48e5fc2a8278
md"""
Plot the QoIs (figure 3). If we want to plot observations, supply `plotObs=true` as a keyword.
"""

# ╔═╡ d0bba806-3b06-498d-b9f7-f6e031e94cbc
md"""
### UQ and GSA
"""

# ╔═╡ 067cce79-6913-4ae3-9f8d-14b128911cd1
md"""
Build surrogate
"""

# ╔═╡ f54baa6e-475f-4265-a692-e6ce06e76e86
begin
	XTrainFinal = Matrix(paramsBgScaled[finalRuns, :])
	YTrainUr = Array{Float64, 2}(UrSim[:, finalRuns]')
	YTrainNp = Array{Float64, 2}(NpSim[:, finalRuns]')
end

# ╔═╡ 753bccd6-66e7-411d-8cba-1a7015c476ea
# build coefficient matrix!
ATrainFinal = buildCoefficientMatrix(XTrainFinal; pceDegree=2)

# ╔═╡ 32d5aef8-67ab-4eeb-b468-aa1e9c68c7d2
begin
	# set regularization coefficients
	lambdaUrFinal = 0.4
	lambdaNpFinal = 5
end

# ╔═╡ 8740834d-2b20-412c-a57a-02a4598e6d4f
begin
	# Perform ridge regression
	betaUrFinal = solveRegPCE(ATrainFinal, YTrainUr; λ=lambdaUrFinal)
	betaNpFinal = solveRegPCE(ATrainFinal, YTrainNp; λ=lambdaNpFinal)
end

# ╔═╡ d58e3a0e-de67-4e28-9b7d-4776f681ca0a


# ╔═╡ 916903b6-8182-4bd6-924d-0424efbc5c29
md"""
Load matrix of test points
"""

# ╔═╡ 52cd0de7-60eb-417c-a942-ade86a4c1dfe
XTestFinal = load("./data/design/CR2152TestFinal.jld", "XTestFinal")[:, 1:6]

# ╔═╡ faef4f04-f921-404e-93b5-f4acbe3aa393
ATest = buildCoefficientMatrix(XTestFinal[:, 1:6]; pceDegree=2);

# ╔═╡ 7ccd891a-2072-4e70-a60c-76a5d61535db
begin
	ATestFinalCoeffs = ATest[:, 2:end]
	ATestFinalFiltered = ATestFinalCoeffs .- mean(ATestFinalCoeffs, dims=1)
	yPredNpFinal = betaNpFinal[1, :] .+ Matrix((ATestFinalFiltered * betaNpFinal[2:end, :])');
	yPredUrFinal = betaUrFinal[1, :] .+ Matrix((ATestFinalFiltered * betaUrFinal[2:end, :])');
end

# ╔═╡ 7af29daf-2618-4bdb-9f77-9d403185e030
md"""
### Show cross validation procedure
"""

# ╔═╡ 9778d9ff-b6d3-43df-892e-a2ea7b795bec
lambdaVals = 10 .^(range(-2, 2, length=40))

# ╔═╡ e98b54e8-b88a-4275-a2c1-e5697a62e7f3
kFoldErrorsUr = kFoldCV(XTrainFinal, YTrainUr, lambdaVals; pceDegree=2, nFolds=10)

# ╔═╡ c8c834aa-b4c3-4969-a6b5-909a9176e5b4
kFoldErrorsNp = kFoldCV(XTrainFinal, YTrainNp, lambdaVals; pceDegree=2, nFolds=10)

# ╔═╡ 025f2e8f-0d30-42ec-8de3-8ea4d8d60fb6
lambdaOptIdxUr = [argmin(kFoldErrorsUr[i, :]) for i in 1:length(times)]

# ╔═╡ 5b2062b8-3c3b-42ea-aabe-56d97b0779c5
kFoldErrorsUr[286, 30]

# ╔═╡ 22f00388-ab63-410f-8a9c-e011416f0393
lambdaOptIdxNp = [argmin(kFoldErrorsNp[i, :]) for i in 1:length(times)]

# ╔═╡ 1b0d8c0e-6001-4ac9-a3bb-f9ba252a31a9
lambdaOptUr = lambdaVals[lambdaOptIdxUr]

# ╔═╡ 6dcc94c9-3fb7-4a4a-b9b7-a5f6a349dc75
lambdaOptNp = lambdaVals[lambdaOptIdxNp]

# ╔═╡ b3578de8-38ea-467e-90ab-a9e8db4db485
begin
kFoldErrorsUrDeg1 = kFoldCV(XTrainFinal, YTrainUr, lambdaVals; pceDegree=1, nFolds=10)
kFoldErrorsUrDeg3 = kFoldCV(XTrainFinal, YTrainUr, lambdaVals; pceDegree=3, nFolds=10)
end

# ╔═╡ 094f9b4e-0ece-416a-a876-3819ee7fedac
kFoldErrorsUrDeg4 = kFoldCV(XTrainFinal, YTrainUr, lambdaVals; pceDegree=4, nFolds=10)

# ╔═╡ 65b600d8-7e98-414c-845b-3503a762e34d
begin
	kFoldErrorsNpDeg1 = kFoldCV(XTrainFinal, YTrainNp, lambdaVals; pceDegree=1, nFolds=10)
	kFoldErrorsNpDeg3 = kFoldCV(XTrainFinal, YTrainNp, lambdaVals; pceDegree=3, nFolds=10)
	kFoldErrorsNpDeg4 = kFoldCV(XTrainFinal, YTrainNp, lambdaVals; pceDegree=4, nFolds=10)
	
end

# ╔═╡ 7328f543-68c4-4d9d-b06d-affda3c71262
avgCVError1 = mean([minimum(kFoldErrorsNpDeg1[i, :]) for i in 1:length(times)])

# ╔═╡ 85a9b0ad-3053-405a-80c8-43444d75da24
avgCVError = mean([minimum(kFoldErrorsNp[i, :]) for i in 1:length(times)])

# ╔═╡ 534cd0e7-fcb2-4753-9f2e-61661607a052
avgCVError4 = mean([minimum(kFoldErrorsNpDeg4[i, :]) for i in 1:length(times)])

# ╔═╡ b6ba45b5-d0d6-485f-9923-94c560906739
avgCVError3 = mean([minimum(kFoldErrorsNpDeg3[i, :]) for i in 1:length(times)])

# ╔═╡ 5fb30df3-92b4-45a6-8cac-2ded08e26e99
begin
	avgCVError1Ur = mean([minimum(kFoldErrorsUrDeg1[i, :]) for i in 1:length(times)])
	avgCVErrorUr = mean([minimum(kFoldErrorsUr[i, :]) for i in 1:length(times)])
	avgCVError3Ur = mean([minimum(kFoldErrorsUrDeg3[i, :]) for i in 1:length(times)])
	avgCVError4Ur = mean([minimum(kFoldErrorsUrDeg4[i, :]) for i in 1:length(times)])
end

# ╔═╡ cec989e4-b995-4718-ae20-dde1cf3345e4
md"""
CV Error Plots (Figure 4)
"""

# ╔═╡ 06fbe236-8681-40d9-9918-d8c47c4a0107
begin
	pAvgUr = plot([1, 2, 3, 4], [avgCVError1Ur, avgCVErrorUr, avgCVError3Ur, avgCVError4Ur], line=(2), marker=(:circle), label="")
	plot!(xlabel="Degree of PCE in total order expansion")
	plot!(ylabel="Minimum CV error")
	plot!(guidefontsize=20)
	plot!(legend=:topright)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(top_margin=2mm)
	plot!(ytickfontsize=17)
	plot!(xtickfontsize=17)
	plot!(size=(800, 500))
end

# ╔═╡ 12dcf64a-d43f-40ee-9921-3027570a52b7
begin
	pAvgNp = plot([1, 2, 3, 4], [avgCVError1, avgCVError, avgCVError3, avgCVError4], line=(2), marker=(:circle), label="")
	plot!(xlabel="Degree of PCE in total order expansion")
	plot!(ylabel="Minimum CV error")
	plot!(guidefontsize=20)
	plot!(legend=:topright)
	plot!(left_margin=5mm)
	plot!(bottom_margin=5mm)
	plot!(top_margin=2mm)
	plot!(ytickfontsize=17)
	plot!(xtickfontsize=17)
	plot!(size=(800, 500))
end

# ╔═╡ d92fc442-23d3-4cb3-bd5f-907d5a2f8f29
begin
	meanEmpiricalNpFinal = mean(yPredNpFinal; dims=2)
	stdEmpiricalNpFinal = std(yPredNpFinal; dims=2)[:]

	meanEmpiricalUrFinal = mean(yPredUrFinal; dims=2)
	stdEmpiricalUrFinal = std(yPredUrFinal; dims=2)[:]
end

# ╔═╡ b57909fa-eba6-462d-85aa-38708abfdc96
md"""
Plot predictive uncertainty (figure 5)
"""

# ╔═╡ 06af5293-d6be-4938-9ce6-5d5b557da1a9
begin
	plotUncertainty(yPredUrFinal, meanEmpiricalUrFinal, stdEmpiricalUrFinal, UrObs, ylabel="Uᵣ [km/s]", times, nSTD=2, trimIndices=(1, 577))
	plot!(guidefontsize=20)
	plot!(tickfontsize=15)
	plot!(legendfontsize=15)
	plot!(labelsize=15)
	plot!(legend=:top)
	plot!(left_margin=3mm)
	plot!(bottom_margin=3mm)	
end

# ╔═╡ e566ed1c-8f0e-4799-976c-c519688bef4f
begin
	plotUncertainty(yPredNpFinal, meanEmpiricalNpFinal, stdEmpiricalNpFinal, NpObs, ylabel="Nₚ [cm⁻³]", ylims=(-20, 120), times, nSTD=2, trimIndices=(1, 577))
	plot!(guidefontsize=20)
	plot!(tickfontsize=15)
	plot!(legend=:topleft)
	plot!(labelsize=15)
	plot!(left_margin=3mm)
	plot!(bottom_margin=3mm)	
end

# ╔═╡ 6d233663-3148-4166-bf5b-d5499ec0d8af
md"""
Sensitivities
"""

# ╔═╡ 5abc9da7-d0b0-4555-b18e-4b7c673f0e08
begin
	gsaUrFinal = gsa(XTrainFinal, YTrainUr; regularize=true, pceDegree=2, lambda=lambdaUrFinal)
	gsaNpFinal = gsa(XTrainFinal, YTrainNp; regularize=true, pceDegree=2, lambda=lambdaNpFinal)
end

# ╔═╡ c9c23639-0567-455a-9a87-ea1e0d6eb3b3
# Extract main effects
begin
	gsaMainUrFinal = processMainEffects(gsaUrFinal)
	gsaMainNpFinal = processMainEffects(gsaNpFinal)
end

# ╔═╡ 6a0fd50a-011d-4f7d-bb76-35ee79375c68
md"""
Plot sensitivities (figures 6 and 7)
"""

# ╔═╡ 4ffe4279-974c-48d7-8887-4f7fc1ea7edd
begin
	pMainUrFinal = plotMainEffects(gsaMainUrFinal, times, inputNames; title = "", dpi=300, actualStartTime=startTime, tickFormat="dd-m")
	plot!(grid=false)
	plot!(ylabel="Sᵢᵗ")
	plot!(titlefontsize=17)
	plot!(xtickfontsize=20)
	plot!(ytickfontsize=20)
	plot!(guidefontsize=23)
	plot!(leftmargin=7mm)
	plot!(rightmargin=5mm)
	plot!(topmargin=2mm)
	plot!(bottommargin=8mm)
	plot!(size=(1200, 600))
end

# ╔═╡ 3eadc4ba-e358-4171-bd03-366d56a89d82
# begin
# 	savefig(pMainUrFinal, "/Users/ajivani/Downloads/pMainUrFinal.png")
# 	savefig(pMainNpFinal, "/Users/ajivani/Downloads/pMainNpFinal.png")

# end

# ╔═╡ be071b08-5fd4-4e8f-a621-9a5a1f9d7414
begin
	pMainNpFinal = plotMainEffects(gsaMainNpFinal, times, inputNames; title="", dpi=300, actualStartTime=startTime, tickFormat="dd-m")
	plot!(grid=false)
	plot!(ylabel="Sᵢᵗ")
	plot!(titlefontsize=17)
	plot!(xtickfontsize=20)
	plot!(ytickfontsize=20)
	plot!(guidefontsize=23)
	plot!(leftmargin=7mm)
	plot!(rightmargin=5mm)
	plot!(topmargin=2mm)
	plot!(bottommargin=8mm)
	plot!(size=(1200, 600))
end

# ╔═╡ 62ff12ba-72a8-428d-9dd2-b3dfab01f7ad
begin
	pUrIntMeanFinal = plotInteractionEffects(gsaUrFinal, times, inputNames, dpi=300) # symmetric matrix (interactions are read from either the upper triangle or the lower triangle)
	plot!(tickfontsize=12)
	plot!(bottom_margin=9mm)
	plot!(right_margin=3mm)
end

# ╔═╡ 917fa984-e031-4556-b152-686737abd3ee
begin
	pNpIntMeanFinal = plotInteractionEffects(gsaNpFinal, times, inputNames, dpi=300)
	plot!(tickfontsize=12)
	plot!(bottom_margin=9mm)
	plot!(right_margin=3mm)
end

# ╔═╡ 3cec0e45-bffe-4905-893a-8fac158fc6d5
md"""
Analysis based on log-transformed quantity
"""

# ╔═╡ 117a60ea-a874-4e3c-8a39-59f0fc57f9b3
YTrainNpLog = Array{Float64, 2}(log.(NpSim[:, finalRuns])');

# ╔═╡ 0f30b6cc-7777-4491-8b86-411fdb0bcd5a
lambdaNpLogFinal = 2

# ╔═╡ 98e3dd8e-8763-4187-a3db-4565a1f5eeeb
betaNpLogFinal = solveRegPCE(ATrainFinal, YTrainNpLog; λ=lambdaNpLogFinal)

# ╔═╡ 94b7993d-0e06-46c4-9735-86c582369fa2
begin
	yPredNpLogFinal = betaNpLogFinal[1, :] .+ Matrix((ATestFinalFiltered * betaNpLogFinal[2:end, :])')
	yPredNpTransformed = exp.(yPredNpLogFinal)
end

# ╔═╡ 91eecee2-26e5-470d-be2d-f00432d7cd3a
md"""
Plot predictions of PCE surrogate trained on Np and log(Np). Also plot time-averaged sensitivities for the latter. (figure 8)
"""

# ╔═╡ 21849735-fe38-4bec-b5e5-a356e56ce80a
gsaNpLogFinal = gsa(XTrainFinal, YTrainNpLog; regularize=true, pceDegree=2, lambda=lambdaNpLogFinal);

# ╔═╡ 936b6a28-a7f3-48ff-b452-6e665f3999fd
begin
	pNpLogIntMeanFinal = plotInteractionEffects(gsaNpLogFinal, times, inputNames, dpi=300)
	plot!(tickfontsize=12)
	plot!(bottom_margin=9mm)
	plot!(right_margin=3mm)
end

# ╔═╡ b4f9217f-231b-416d-8759-39a6650a5c0b
md"""
Bootstrapping
"""

# ╔═╡ ab1a11a9-6670-43f5-8c81-a203153b6af7
md"""
Load saved bootstrap data. Performing the bootstrap can be quite slow, so the actual command for getting the data is commented out.
"""

# ╔═╡ b89e23cb-2e97-48bd-b21c-12e4f361b766
# UrBootstrap = bootstrapGSA(XTrainFinal, YTrainUr; regularize=false, nStart=20, nEnd=140, nStep=20)
# NpBootstrap = bootstrapGSA(XTrainFinal, YTrainNp; regularize=false, nStart=20, nEnd=140, nStep=20)

# ╔═╡ a4fbbd42-47af-47a6-be08-9e248a2902b9
begin
	UrBootstrap = load("./data/bootstrapping/bootstrapUr2152.jld", "UrBootstrap")
	avgBootstrapUr = mean(UrBootstrap; dims=2)[:, 1, :, :]
	avgBootstrapRepsUr = mean(avgBootstrapUr; dims=2)[:, 1, :]
	stdBootstrapRepsUr = std(avgBootstrapUr; dims=2)[:, 1, :]
end

# ╔═╡ f6afc166-dd90-44fb-b3ca-6de6088ebe31
begin
	NpBootstrap = load("./data/bootstrapping/bootstrapNp2152.jld", "NpBootstrap")
	avgBootstrapNp = mean(NpBootstrap; dims=2)[:, 1, :, :]
	avgBootstrapRepsNp = mean(avgBootstrapNp; dims=2)[:, 1, :]
	stdBootstrapRepsNp = std(avgBootstrapNp; dims=2)[:, 1, :]
end

# ╔═╡ 5c8b16a6-3d16-49a9-92d3-7b806802c857
summaryColors = palette(:tab10, rev=true)[6:-1:1]

# ╔═╡ 9ad56db9-adad-425f-897d-4c490757a8ff
md"""
Plot summary of bootstrapping results. (figure 9)
"""

# ╔═╡ 5f623f56-b62c-4997-bef5-c8e82c68283e
begin
	pLineSummaryUr = plot()
	for (idx, eachName) in enumerate(inputNames)
		meanTrend = avgBootstrapRepsUr[idx, :]
		errTrend  = stdBootstrapRepsUr[idx, :] 
		plot!(20:20:140, meanTrend, 
			yerr=errTrend,
			linewidth=2.5,
			linecolor=summaryColors[idx], 
			marker=:circle, 
			markercolor=summaryColors[idx],
			markerstrokecolor=summaryColors[idx],
			markerstrokewidth=1.5,
			label=eachName)
	end
	plot!(legend=:topleft)
	plot!(xticks=(20:20:140, string.(20:20:140)))
	plot!(xlabel="Sample size n")
	plot!(ylabel="Sᵢ")
	plot!(guidefontsize=20)
	plot!(tickfontsize=16)
	plot!(leftmargin=5mm)
	plot!(bottommargin=5mm)
	plot!(framestyle=:box)
	plot!(fg_legend=nothing)
	plot!(bg_legend=nothing)
	plot!(size=(750, 500))
	plot!(ylims=(0, 0.5))
end

# ╔═╡ 92c749d3-a0ae-4031-83d4-27c0966a0a6a
begin
	pLineSummaryNp = plot()
	for (idx, eachName) in enumerate(inputNames)
		meanTrend = avgBootstrapRepsNp[idx, :]
		errTrend  = stdBootstrapRepsNp[idx, :] 
		plot!(20:20:140, meanTrend, 
			yerr=errTrend,
			linewidth=2.5,
			linecolor=summaryColors[idx], 
			marker=:circle, 
			markercolor=summaryColors[idx],
			markerstrokecolor=summaryColors[idx],
			markerstrokewidth=1.5,
			label=eachName)
	end
	plot!(legend=:topleft)
	plot!(xticks=(20:20:140, string.(20:20:140)))
	plot!(xlabel="Sample size n")
	plot!(ylabel="Sᵢ")
	plot!(guidefontsize=20)
	plot!(tickfontsize=16)
	plot!(leftmargin=5mm)
	plot!(bottommargin=5mm)
	plot!(framestyle=:box)
	plot!(fg_legend=nothing)
    plot!(bg_legend=nothing)
	plot!(size=(750, 500))
	plot!(ylims=(0, 0.5))
end

# ╔═╡ 122a2ee3-2561-4940-87ee-5154c219e7c9
md"""
Miscellaneous
"""

# ╔═╡ ec279c30-57e0-4b71-906f-2e059fe5569d
begin
		plotArgsUr = Dict(:palette=>:seaborn_bright,
					   :dateFormat=>"dd-mm HH:MM",
					   :tickInterval=>108,
					   # :simIdx => parse.(Int, EachSimID),
					   :simAlpha=>0.6, 
					   :simWidth=>1.2,
					   :ylabel=>"Uᵣ[km/s]",
					   :title=>"",
					   :startTime=>startTime,
					   :dpi=>200,
					   :plotLabels=>:false,
					   
	                  )
	
		plotArgsNp = Dict(:palette=>:seaborn_bright,
						   :dateFormat=>"dd-mm HH:MM",
						   :tickInterval=>108,
						   # :simIdx => parse.(Int, EachSimID),
						   :simAlpha=>0.6, 
						   :simWidth=>1.2,
						   :ylabel=>"Nₚ[cm⁻³]",
						   # :ylabel=>"Np [cm" * L"^{-3}" * "]",
						   :title=>"",
						   :startTime=>startTime,
						   :dpi=>200,
						   :ylims=>(0, 80),
						   :plotLabels=>:false
						  )
	
		plotArgsB = Dict(:palette=>:seaborn_bright,
						   :dateFormat=>"dd-mm HH:MM",
						   :tickInterval=>108,
						   # :simIdx => parse.(Int, EachSimID),
						   :simAlpha=>0.6, 
						   :simWidth=>1.2,
						   :ylabel=>"B[nT]",
						   # :ylabel=>"Np [cm" * L"^{-3}" * "]",
						   :title=>"",
						   :startTime=>startTime,
						   :dpi=>200,
						   :ylims=>(0, 20),
						   :plotLabels=>:false,
						   :subtractFactor=>0
		                  )
	
		plotArgsT = Dict(:palette=>:seaborn_bright,
					   :dateFormat=>"dd-mm HH:MM",
					   :tickInterval=>108,
					   # :simIdx => parse.(Int, EachSimID),
					   :simAlpha=>0.6, 
					   :simWidth=>1.5,
					   :ylabel=>"T[K]",
					   # :ylabel=>"Np [cm" * L"^{-3}" * "]",
					   :title=>"",
					   :startTime=>startTime,
					   :dpi=>200,
					   :ylims=>(0, 9e5),
					   :plotLabels=>:false,
					   :subtractFactor=>0
					  )
end

# ╔═╡ 8c85fac8-6fcd-4ece-951d-ed905a73f020
begin
	# make plots of the runs that are successful and satisfy the constraints
	pUrSimObsF = plotSimObs(UrSim, UrObs, times, collect(1:200); simIdx=finalRuns, plotArgsUr..., ylims=(100, 800), dateFormat="dd-mm")
	plot!(size=(920, 470))
	plot!(guidefontsize=24)
	plot!(left_margin=7mm)
	plot!(bottom_margin=7mm)
	plot!(top_margin=5.5mm)
	plot!(ytickfontsize=18)
	plot!(xtickfontsize=18)
	plot!(dpi=300)
end

# ╔═╡ e3b927db-8d9f-4061-94e1-20aa70975d3f
begin
	pNpSimObsF = plotSimObs(NpSim, NpObs, times, collect(1:200); simIdx=finalRuns, plotArgsNp..., ylims=(0, 140), dateFormat="dd-mm")
	plot!(size=(920, 470))
	plot!(guidefontsize=24)
	plot!(left_margin=7mm)
	plot!(bottom_margin=7mm)
	plot!(top_margin=5.5mm)
	plot!(ytickfontsize=18)
	plot!(xtickfontsize=18)
	plot!(dpi=300)
	
end

# ╔═╡ fb8fa462-4103-41a0-bde4-cc18929fbdb3
begin
	pBSimObsF  = plotSimObs(BSim, BObs, times, collect(1:200); simIdx=finalRuns, plotArgsB..., dateFormat="dd-mm")
	plot!(size=(920, 470))
	plot!(guidefontsize=24)
	plot!(left_margin=7mm)
	plot!(bottom_margin=7mm)
	plot!(top_margin=5.5mm)
	plot!(ytickfontsize=18)
	plot!(xtickfontsize=18)
	plot!(dpi=300)
end

# ╔═╡ 78b04a27-9e90-49d0-abd3-d5c4d422a1b3
begin
	pTSimObsF = plotSimObs(TSim, TObs, times, collect(1:200); simIdx=finalRuns, plotArgsT..., ylims=(0, 5e5), dateFormat="dd-mm")
	plot!(yticks=(0:1e5:5e5, [0; [@sprintf("%de5", i) for i in 1:5]][:]))
	plot!(size=(920, 470))
	plot!(guidefontsize=24)
	plot!(left_margin=7mm)
	plot!(bottom_margin=7mm)
	plot!(top_margin=5.5mm)
	plot!(ytickfontsize=18)
	plot!(xtickfontsize=18)
	plot!(dpi=300)
end

# ╔═╡ d15af974-8468-4029-a916-18d73f0b258f
begin
	pNpSimObsOrigPred = plotSimObs(yPredNpFinal, NpObs, times, collect(1:200); simIdx=collect(1:400), plotArgsNp..., ylims=(-30, 190), dateFormat="dd-mm")
	plot!(size=(920, 470))
	plot!(guidefontsize=24)
	plot!(left_margin=7mm)
	plot!(bottom_margin=7mm)
	plot!(top_margin=5.5mm)
	plot!(ytickfontsize=18)
	plot!(xtickfontsize=18)
	plot!(dpi=300)
end

# ╔═╡ 84d675a3-e987-4ac5-9c1a-e3775ac82c74
begin
	pNpSimObsLogPred = plotSimObs(yPredNpTransformed, NpObs, times, collect(1:200); simIdx=collect(1:400), plotArgsNp..., ylims=(-30, 190), dateFormat="dd-mm")
	plot!(size=(920, 470))
	plot!(title="Predictions from PCE for log Np")
	plot!(guidefontsize=24)
	plot!(left_margin=7mm)
	plot!(bottom_margin=7mm)
	plot!(top_margin=5.5mm)
	plot!(ytickfontsize=18)
	plot!(xtickfontsize=18)
	plot!(dpi=300)
end

# ╔═╡ Cell order:
# ╠═95e46924-1e6d-11ed-152b-e373176bfe2b
# ╟─e425f65c-a0a6-4316-b799-b495b73d5a97
# ╠═3ce6b18e-deae-44fa-8824-4193e7d6edb3
# ╠═893e00b2-0baf-49c2-9c94-884eb0607eb4
# ╟─ea56d0e1-ef32-4582-908b-e9ba1f576c9d
# ╟─9415e163-19c0-453c-80fe-90ec1c1f4daf
# ╟─429c13e8-f91c-41cc-b112-81b6a7f3d41a
# ╠═c3fb873b-2457-4317-83e1-a8dd22db0989
# ╟─7e81d77e-6aa8-40f2-a900-22f76a996185
# ╠═8b66c931-a842-4011-a06c-93ac0e1445e5
# ╟─67a24d73-c444-496d-9337-cfbf9cde37c2
# ╠═c607abf7-e20e-46a1-85ed-7d580a55d9db
# ╠═2b218d46-6a44-4b34-9236-06fb30d4d96a
# ╟─7b4efa4b-bea0-4517-a7c1-df969a4ca74e
# ╠═7139a373-2795-436b-9bb3-6c4596254101
# ╠═82478091-c845-4b26-8092-19f1b9eb7776
# ╠═6620dfd1-f147-42ec-8d27-319e7f1cb5aa
# ╟─84c005c3-9fa2-4cc4-80b7-105031ef327a
# ╠═8eb1989c-5a16-43fb-901e-55ea9663a7df
# ╠═ba9fe174-0d0e-4396-baf6-088902335076
# ╠═fa98a60f-75d7-48d6-81ad-af6d287b1f2f
# ╠═dc313671-8a22-4ec1-b74a-36c15398c4b7
# ╠═7ce46450-824e-415c-843b-0990848ab852
# ╠═e3060049-736d-48da-9767-945d4185153b
# ╠═14263462-9183-4efc-bd40-957c1443b0ac
# ╟─47a8931d-02fc-47ce-950e-f8338d87432c
# ╠═26379e4c-2fa3-42e6-87b9-29434b61387e
# ╟─678d55ec-afbc-4fb0-9099-48e5fc2a8278
# ╠═8c85fac8-6fcd-4ece-951d-ed905a73f020
# ╠═e3b927db-8d9f-4061-94e1-20aa70975d3f
# ╠═fb8fa462-4103-41a0-bde4-cc18929fbdb3
# ╠═78b04a27-9e90-49d0-abd3-d5c4d422a1b3
# ╟─d0bba806-3b06-498d-b9f7-f6e031e94cbc
# ╟─067cce79-6913-4ae3-9f8d-14b128911cd1
# ╠═f54baa6e-475f-4265-a692-e6ce06e76e86
# ╠═753bccd6-66e7-411d-8cba-1a7015c476ea
# ╠═8740834d-2b20-412c-a57a-02a4598e6d4f
# ╠═32d5aef8-67ab-4eeb-b468-aa1e9c68c7d2
# ╠═d58e3a0e-de67-4e28-9b7d-4776f681ca0a
# ╟─916903b6-8182-4bd6-924d-0424efbc5c29
# ╠═52cd0de7-60eb-417c-a942-ade86a4c1dfe
# ╠═faef4f04-f921-404e-93b5-f4acbe3aa393
# ╠═7ccd891a-2072-4e70-a60c-76a5d61535db
# ╟─7af29daf-2618-4bdb-9f77-9d403185e030
# ╠═9778d9ff-b6d3-43df-892e-a2ea7b795bec
# ╠═e98b54e8-b88a-4275-a2c1-e5697a62e7f3
# ╠═c8c834aa-b4c3-4969-a6b5-909a9176e5b4
# ╠═025f2e8f-0d30-42ec-8de3-8ea4d8d60fb6
# ╠═5b2062b8-3c3b-42ea-aabe-56d97b0779c5
# ╠═22f00388-ab63-410f-8a9c-e011416f0393
# ╠═1b0d8c0e-6001-4ac9-a3bb-f9ba252a31a9
# ╠═6dcc94c9-3fb7-4a4a-b9b7-a5f6a349dc75
# ╠═b3578de8-38ea-467e-90ab-a9e8db4db485
# ╠═094f9b4e-0ece-416a-a876-3819ee7fedac
# ╠═65b600d8-7e98-414c-845b-3503a762e34d
# ╠═7328f543-68c4-4d9d-b06d-affda3c71262
# ╠═85a9b0ad-3053-405a-80c8-43444d75da24
# ╠═534cd0e7-fcb2-4753-9f2e-61661607a052
# ╠═b6ba45b5-d0d6-485f-9923-94c560906739
# ╠═5fb30df3-92b4-45a6-8cac-2ded08e26e99
# ╟─cec989e4-b995-4718-ae20-dde1cf3345e4
# ╠═06fbe236-8681-40d9-9918-d8c47c4a0107
# ╠═12dcf64a-d43f-40ee-9921-3027570a52b7
# ╠═d92fc442-23d3-4cb3-bd5f-907d5a2f8f29
# ╟─b57909fa-eba6-462d-85aa-38708abfdc96
# ╠═06af5293-d6be-4938-9ce6-5d5b557da1a9
# ╠═e566ed1c-8f0e-4799-976c-c519688bef4f
# ╟─6d233663-3148-4166-bf5b-d5499ec0d8af
# ╠═5abc9da7-d0b0-4555-b18e-4b7c673f0e08
# ╠═c9c23639-0567-455a-9a87-ea1e0d6eb3b3
# ╟─6a0fd50a-011d-4f7d-bb76-35ee79375c68
# ╠═4ffe4279-974c-48d7-8887-4f7fc1ea7edd
# ╠═3eadc4ba-e358-4171-bd03-366d56a89d82
# ╠═be071b08-5fd4-4e8f-a621-9a5a1f9d7414
# ╠═62ff12ba-72a8-428d-9dd2-b3dfab01f7ad
# ╠═917fa984-e031-4556-b152-686737abd3ee
# ╟─3cec0e45-bffe-4905-893a-8fac158fc6d5
# ╠═117a60ea-a874-4e3c-8a39-59f0fc57f9b3
# ╠═0f30b6cc-7777-4491-8b86-411fdb0bcd5a
# ╠═98e3dd8e-8763-4187-a3db-4565a1f5eeeb
# ╠═94b7993d-0e06-46c4-9735-86c582369fa2
# ╟─91eecee2-26e5-470d-be2d-f00432d7cd3a
# ╠═d15af974-8468-4029-a916-18d73f0b258f
# ╠═84d675a3-e987-4ac5-9c1a-e3775ac82c74
# ╠═21849735-fe38-4bec-b5e5-a356e56ce80a
# ╠═936b6a28-a7f3-48ff-b452-6e665f3999fd
# ╟─b4f9217f-231b-416d-8759-39a6650a5c0b
# ╟─ab1a11a9-6670-43f5-8c81-a203153b6af7
# ╠═b89e23cb-2e97-48bd-b21c-12e4f361b766
# ╠═a4fbbd42-47af-47a6-be08-9e248a2902b9
# ╠═f6afc166-dd90-44fb-b3ca-6de6088ebe31
# ╠═5c8b16a6-3d16-49a9-92d3-7b806802c857
# ╟─9ad56db9-adad-425f-897d-4c490757a8ff
# ╠═5f623f56-b62c-4997-bef5-c8e82c68283e
# ╠═92c749d3-a0ae-4031-83d4-27c0966a0a6a
# ╟─122a2ee3-2561-4940-87ee-5154c219e7c9
# ╠═ec279c30-57e0-4b71-906f-2e059fe5569d