#!/bin/bash

##### for target domain in Central America 
latS=-5.
latN=33.0
lonW=242.0  #-118 + 360
lonE=302.0    # -58 +360

#### SST Domain for Western Africa climate (Atlantic, Pacific)  
lonWeof=130. 
lonEeof=340.
latSeof=-20.
latNeof=20.

##### Number of SST and Precip EOFs
numeof=6

#### Model intitial date and forecast period
startmon=5
sznlen=3
ilead=1

obsstartmon=$((ilead + startmon))
obsendmon=$((obsstartmon + sznlen - 1))

currentyr=2026
seasonname="JJA"

#### For R
testLength=5
trainLength=30
leave3Out=TRUE
runserial=FALSE
obst="MSWEP"
#obst="LOCAL"

#### For ploting -- Change country
mycountry="CRI"  # Costa Rica
yrsstsc=1991
yrsstec=2020

### For obsevations
firstyr=1991
endyr=2025
monthly=True
isstate=False
obsres=0.5
obsprecip="prec${seasonname}_${firstyr}-${endyr}${obst}.nc"

### Diagnostics
printProb=False
printProbRaw=False
#Diagnostic Domain for Senegal
tlats=9.
tlatn=10.
tlonw=275.0
tlone=276.0

#### For both R and NCL
###########
szn4mon=False
leadi=$ilead
leadf=$((ilead + sznlen - 1))

allmodelNames=(canesm5 gem5nemo nasa ccsm4 cfsv2 cesm1 ensmean)

allprdctors=(MGI ATL3 Nino34 WVG IOD SIOD PCS1 PCS2 PCS3 PCS4 PCS5 PCS6 PCR1 PCR2 PCR3 PCR4 PCR5 PCR6 ALLP_INT)

slctprdctors=(MGI ATL3 Nino34 WVG IOD SIOD PCS1 PCS2 PCS3 PCS4 PCS5 PCS6 PCR1 PCR2 PCR3 PCR4 PCR5 PCR6 ALLP_INT)

nprdFirst=$(( ${#slctprdctors[@]} - 1 ))
#######################
join_by_comma() {
  local IFS=","
  echo "$*"
}

quote_items_for_r() {
  local out=()
  for x in "$@"; do
    out+=("\"$x\"")
  done
  local IFS=","
  echo "c(${out[*]})"
}

quote_items_for_ncl() {
  local out=()
  for x in "$@"; do
    out+=("\"$x\"")
  done
  local IFS=","
  echo "(/${out[*]}/)"
}

allmodelNamesR=$(quote_items_for_r "${allmodelNames[@]}")
allmodelNamesNCL=$(quote_items_for_ncl "${allmodelNames[@]}")
slctPrdctrR=$(quote_items_for_r "${slctprdctors[@]}")
allprdctorNCL=$(quote_items_for_ncl "${allprdctors[@]}")
slctprdctrNCL=$(quote_items_for_ncl "${slctprdctors[@]}")

##########
#1 Get observational validation data
cat > paramsObs.ncl <<EOF
startday = 1
iMonth   = $startmon
startmon = $obsstartmon
endmon   = $obsendmon
SEASON   = "$seasonname"
YearS    = $firstyr
YearE    = $endyr
latS     = $latS
latN     = $latN
lonW     = $lonW
lonE     = $lonE
obsT     = "$obst"
monthly  = $monthly
Res      = $obsres
EOF

if [ $obst = "LOCAL" ]; then
 cat paramsObs.ncl templateLocalObs.ncl > obsPrecip.ncl
else
 cat paramsObs.ncl templateObs.ncl > obsPrecip.ncl
fi
ncl obsPrecip.ncl
##################################

#######################
#### 2. SST 
cat > paramsSST.ncl <<EOF
iMon           = $startmon
YYYY           = $currentyr
SEASON         = "$seasonname"
yrS            = $firstyr
yrE            = $endyr
Leadi          = $leadi
Leadf          = $leadf
latSeof        = $latSeof
latNeof        = $latNeof
lonWeof        = $lonWeof
lonEeof        = $lonEeof
numeof         = $numeof
yrSc           = $yrsstsc
yrEc           = $yrsstec
modelFileNames = $allmodelNamesNCL
EOF

#2 Get NMME SST data
cat paramsSST.ncl templateSSTDownload.ncl > sstDownloadStdz.ncl

ncl sstDownloadStdz.ncl

#######################
#### 3. Get Precip 
cat > paramsPrec.ncl <<EOF
iMon           = $startmon
YYYY           = $currentyr
SEASON         = "$seasonname"
yrS            = $firstyr
yrE            = $endyr
Leadi          = $leadi
Leadf          = $leadf
latSeof        = $latS
latNeof        = $latN
lonWeof        = $lonW
lonEeof        = $lonE
numeof         = $numeof
yrSc           = $yrsstsc
yrEc           = $yrsstec
modelFileNames = $allmodelNamesNCL
EOF

#3 Get NMME precip data
cat paramsPrec.ncl templatePrecipDownload.ncl > precDownloadStdz.ncl

ncl precDownloadStdz.ncl

#############################
### 4. R script
####################

cat > params.R <<EOF
iMonth <- $startmon
SznLen <- $sznlen
ThisYr <- $currentyr
SEASON <- "$seasonname"
fileprecnc <- "$obsprecip"
FirstYr <- $firstyr
LastYr <- $endyr
iLead <- $ilead
trainLength <- $trainLength
testLength <- $testLength
Leave3Out <- $leave3Out
RUN_SERIAL <- $runserial
modelFileNames <- $allmodelNamesR
slctPrdctr <- $slctPrdctrR
nprdFirst <- $nprdFirst
EOF


####params.R is sourced inside multiNominalRegression.R
cp ./templateDriver.R  multiNominalRegression.R
Rscript multiNominalRegression.R

############################
##### 5. Plot individual single model probabilities
constant=".ncl"

for v in "${slctprdctors[@]}"; do
    outfile="plotSingleModel${v}${constant}"

    cat > paramsSingle.ncl <<EOF
PRDCTR        = "${v}"
iLead          = $ilead
Szn4Month      = $szn4mon
iMonth         = $startmon
STATE          = $isstate
currentYr      = $currentyr
SEASON         = "$seasonname"
latS           = $latS
latN           = $latN
lonW           = $lonW
lonE           = $lonE
COUNTRY        = "$mycountry"
Res            = $obsres
printProb      = $printProb
printProbRaw   = $printProbRaw
tlats          = $tlats
tlatn          = $tlatn
tlonw          = $tlonw
tlone          = $tlone
modelFileNames = $allmodelNamesNCL
EOF
    cat paramsSingle.ncl templateIndividualProb.ncl > "$outfile"
    ncl "./$outfile"
    echo $outfile
done

############################
##### 6. Plot individual average model probabilities
constant=".ncl"

for v in "${slctprdctors[@]}"; do
    outfile="plotAvgModel${v}${constant}"

    cat > paramsAvg.ncl <<EOF
PRDCTR        = "${v}"
iLead          = $ilead
Szn4Month      = $szn4mon
iMonth         = $startmon
STATE          = $isstate
currentYr      = $currentyr
startmon       = $obsstartmon
SEASON         = "$seasonname"
latS           = $latS
latN           = $latN
lonW           = $lonW
lonE           = $lonE
COUNTRY        = "$mycountry"
Res            = $obsres
printProb      = $printProb
printProbRaw   = $printProbRaw
tlats          = $tlats
tlatn          = $tlatn
tlonw          = $tlonw
tlone          = $tlone
modelFileNames = $allmodelNamesNCL
EOF

    cat paramsAvg.ncl ./templateEnsProb.ncl > "$outfile"
    ncl "./$outfile"
    echo $outfile
done

############################
##### 7. Plot Average Predictors Skills for indiviual models and ensemble

cat > paramsAllSkill.ncl <<EOF
PRDCTRall      = $slctprdctrNCL
iLead          = $ilead
Szn4Month      = $szn4mon
iMonth         = $startmon
STATE          = $isstate
currentYr      = $currentyr
startmon       = $obsstartmon
SEASON         = "$seasonname"
latS           = $latS
latN           = $latN
lonW           = $lonW
lonE           = $lonE
COUNTRY        = "$mycountry"
Res            = $obsres
printProb      = $printProb
printProbRaw   = $printProbRaw
tlats          = $tlats
tlatn          = $tlatn
tlonw          = $tlonw
tlone          = $tlone
modelFileNames = $allmodelNamesNCL
EOF

cat paramsAllSkill.ncl templateSingleAvgPrdctors.ncl > plotAvgPrdctorsSkill.ncl

ncl ./plotAvgPrdctorsSkill.ncl
    
############################
##### 8. Plot Average Predictors Skills for indiviual models and ensemble

cat > paramsSkillWtd.ncl <<EOF
PRDCTRall      = $slctprdctrNCL
iLead          = $ilead
Szn4Month      = $szn4mon
iMonth         = $startmon
STATE          = $isstate
currentYr      = $currentyr
startmon       = $obsstartmon
SEASON         = "$seasonname"
latS           = $latS
latN           = $latN
lonW           = $lonW
lonE           = $lonE
COUNTRY        = "$mycountry"
Res            = $obsres
printProb      = $printProb
printProbRaw   = $printProbRaw
tlats          = $tlats
tlatn          = $tlatn
tlonw          = $tlonw
tlone          = $tlone
nprdFirst      = $nprdFirst
modelFileNames = $allmodelNamesNCL
EOF

cat paramsSkillWtd.ncl templateEnsSkill.ncl > plotSkillWtdAvg.ncl

ncl ./plotSkillWtdAvg.ncl


############################


echo `date`
