#!/bin/bash
# Copyright 2024 SPC-EPFL
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------

# run CHEASE from an input from the local database euitm or IMAS (tree ids or euitm) and write to database or IMAS
#
# database_to_database shot_in run_in shot_out run_out [cocos_in [cocos_out [occurence_nb_in [occurence_nb_out [treename [treeuser_in [ treeuser_out [ipsign_out [b0sign_out [time_in [tokamak_in [tokamak_out [version_in [version_out [backend_in [backend_out [nitmopt]]]]]]]]]]]]]]]]]
#
# defaults: cocos_in=11 (or 13 if not on iter), cocos_out=cocos_in
# Note: for ITER should set default cocos_out=11 in the file, and 13 for EU
# Note: for IMAS, set treename='ids' in this file as default, for CPOs set treename=euitm
#
# ipsign_out: by default (-9) will use sign(Ip) in eqdsk file and perform the transformation from cocos_in to cocs_out
#             If a specific sign is wanted, then set ipsign_out=+1 or -1
#             Note that the standard for ITER is cocos=11, Ip<0 and B0<0, so can use ipsign_out=b0sign_out=-1
#
# b0sign_out: same meaning as ipsign_out but for sign of B0
#
# time_in: to extract the equilibrium nearby a specific time instead of first time by default
#
# Note: if cocos_in=11 or 2 do not work, try 7
#
# New:
#      export chease_exe=chease_imas_hdf5 (to use a specific executable, now the default at IO is chease_imas_hdf5)
# Note:
#      You can use the top part of the CHEASE output, which contains the namelist, and copy it to "chease_namelist" file, then change some parameters and run chease_imas[_hdf5]
#

# Can override/add variables to the namelist file by adding --override_a=b to the *START* and to the *END* of the list of arguments.
############# start of grab_override file ###########
# Grab override variables and add to temp file, which will later be added to the namelist file, and later deleted.
# override statements should be before and/or after the standard list, but not in the middle
non_override_seen=0
override_file=$(mktemp)
for var in "$@"
do
    [ "${var:0:11}" == "--override_" ] && {
        # strip "--override_" from front of var and add to file
        echo ${var#--override_} >> ${override_file}
        # shift, such that the remaining $1/... logic later still works.
        if [ ${1:0:11} == "--override_" ]; then
           shift # note will not shift anymore once all first overrides have been scanned
        fi
    } || {
        # Mark the first time we see an arg not starting with --override_
        (( non_override_seen++ ))
    }
done
############# end of grab_override file ###########

curdir=`pwd`

name_iter=`uname -n | grep iter`
if [ $? -eq 0 ]
then
    # assume on iter machine thus imas database
    default_treename='ids'
    # default cocos_out should be 11 for ITER, 13 for euitm
    cocos_in=11
    chease_exe=${chease_exe:-chease_imas}
else
    # assume on gateway
    xitmenvx=x`echo $ITM_ENVIRONMENT_LOADED`x
    if [ $xitmenvx = 'xx' ]
    then
        # default cocos_in should be 11 for ITER, 13 for euitm
        default_treename='ids'
        cocos_in=11
        chease_exe=${chease_exe:-chease_imas}
    else
        default_treename='euitm'
        # assume on gateway for euitm database
        cocos_in=13
        chease_exe=${chease_exe:-chease_itm}
    fi
fi

cocos_out=$cocos_in
occurence_nb_in=0
occurence_nb_out=0
ipsign_out=-9
b0sign_out=-9
time_in=1.
tree_tokamak_in='test'
tree_tokamak_out='test'
tree_user_in=`echo $USER`
tree_user_out=`echo $USER`
tree_majorversion_in='3'
tree_majorversion_out='3'
tree_backend_in=13 # hdf5=13, mdsplus=12
tree_backend_out=13
nitmopt=11

if [ ${non_override_seen} -lt 4 -o ${non_override_seen} -gt 21 ]
then
   echo "needs 4 to 21 arguments: shot_in run_in shot_out run_out [cocos_in [cocos_out [occurence_nb_in [occurence_nb_out [treename [treeuser_in [treeuser_out [ipsign_out [b0sign_out [time_in [tree_tokamak_in [tree_tokamak_out [tree_majorversion_in [tree_majorversion_out [tree_backend_in [tree_backend_out [nitmopt]]]]]]]]]]]]]]]]]"
   echo " default cocos_in = "$cocos_in
   echo " default cocos_out = "$cocos_out
   echo " default occurence_nb_in = "$occurence_nb_in
   echo " default occurence_nb_out = "$occurence_nb_out
   echo " default treename: "$default_treename
   echo "for some codes might want to impose negative Ip and B0 thus set ipsign_out=-1 and b0sign_out=-1some signs of Ip or B0, note that positive Ip and B0 is automatically provided as extra eqdsk output (in /tmp/username) (default -9, -9) to follow sign from cocos"
   echo "to get the correct Ip and B0 signs from the transformation from cocos_in to cocos_out, do not set ipsign_out nor b0sign_out (hence -9)"
   echo " default ipsign_out: "$ipsign_out
   echo " default b0sign_out: "$b0sign_out
   echo " default time_in: "$time_in " (note if an array is provided then give it in between quotes to be taken as a single input argument)"
   echo " default tree_tokamak_in: "$tree_tokamak_in
   echo " default tree_tokamak_out: "$tree_tokamak_out
   echo " default tree_user_in: "$tree_user_in
   echo " default tree_user_out: "$tree_user_out
   echo " default tree_majorversion_in: "$tree_majorversion_in
   echo " default tree_majorversion_out: "$tree_majorversion_out
   echo " default tree_backend_in (hdf5=13, mdsplus=12): "$tree_backend_in
   echo " default tree_backend_out (hdf5=13, mdsplus=12): "$tree_backend_out
   echo " default nitmopt (11 creates, 21 opens on an existing database case): "$nitmopt
   echo " "
   echo "All output files are in /tmp/username/."
   echo " "
   exit
fi

shot_number_in=$1
run_number_in=$2
shot_number_out=$3
run_number_out=$4

jobname=${shot_number_in}_${run_number_in}

if [ ${non_override_seen} -ge 5 ]
then
  cocos_in=$5
fi

if [ ${non_override_seen} -ge 6 ]
then
  cocos_out=$6
fi

if [ ${non_override_seen} -ge 7 ]
then
  occurence_nb_in=$7
fi
if [ ${non_override_seen} -ge 8 ]
then
  occurence_nb_out=$8
fi

treename=$default_treename
if [ ${non_override_seen} -ge 9 ]
then
  treename=$9
fi

if [ ${non_override_seen} -ge 10 ]
then
  tree_user_in=${10}
fi
if [ ${non_override_seen} -ge 11 ]
then
  tree_user_out=${11}
fi

if [ ${non_override_seen} -ge 12 ]
then
  ipsign_out=${12}
fi
if [ ${non_override_seen} -ge 13 ]
then
  b0sign_out=${13}
fi

if [ ${non_override_seen} -ge 14 ]
then
  time_in=${14}
fi
if [ ${non_override_seen} -ge 15 ]
then
  tree_tokamak_in=${15}
fi
if [ ${non_override_seen} -ge 16 ]
then
  tree_tokamak_out=${16}
fi
if [ ${non_override_seen} -ge 17 ]
then
  tree_majorversion_in=${17}
fi
if [ ${non_override_seen} -ge 18 ]
then
  tree_majorversion_out=${18}
fi
if [ ${non_override_seen} -ge 19 ]
then
  tree_backend_in=${19}
fi
if [ ${non_override_seen} -ge 20 ]
then
  tree_backend_out=${20}
fi
if [ ${non_override_seen} -ge 21 ]
then
  nitmopt=${21}
fi

che_opt=1

filedir_in=`pwd`

filedir_out='/tmp/'$USER
mkdir $filedir_out 2> /dev/null

os=`which $chease_exe`
if [ $? -ne 0 ]
then
    echo ""
    echo "should have $chease_exe in the PATH"
    echo ""
    exit
fi

# run on tmp:
tmpdir=/tmp/`date|sed s/[^0-9]//g`$$
mkdir $tmpdir
cd $tmpdir
mkdir core
pwd

echo " using chease from :"$os > log_database_to_database

echo "***               $eqdskfilename" > chease_namelist
cat >>chease_namelist <<'EOF'
***
***NEGP=-1,NER= 1  TCV case with ECRH
*** NSURF=6,NEQDSK=1,  NEQDXTPO=-1
 &EQDATA
 RELAX= 0.2,  NDIAGOP= 1,  NBSEXPQ=0000,
EOF
echo " NITMOPT= $nitmopt," >> chease_namelist
echo " NITMSHOT= $shot_number_in, $shot_number_out," >> chease_namelist
echo " NITMRUN= $run_number_in, $run_number_out," >> chease_namelist
echo " NITMOCC= $occurence_nb_in, $occurence_nb_out" >> chease_namelist
echo " COCOS_IN=$cocos_in," >> chease_namelist
echo " COCOS_OUT=$cocos_out," >> chease_namelist
# Previously some compilers did not accept the quotes for the array of string, just remove them in the namelist
echo " TREEITM='$treename','$treename'" >> chease_namelist
echo " TREE_TOKAMAK='$tree_tokamak_in','$tree_tokamak_out'" >> chease_namelist
echo " TREE_USER='$tree_user_in','$tree_user_out'" >> chease_namelist
echo " TREE_MAJORVERSION=$tree_majorversion_in,$tree_majorversion_out" >> chease_namelist
echo " TREE_BACKEND=$tree_backend_in,$tree_backend_out" >> chease_namelist
echo " TIME_REF=$time_in," >> chease_namelist
echo " SIGNIPXP=$ipsign_out," >> chease_namelist
echo " SIGNB0XP=$b0sign_out," >> chease_namelist
cat >>chease_namelist <<'EOF'
 NVERBOSE=3,
 NPROPT=-2,
 NIDEAL=6, NPLOT=1,    NTCASE=0,   NSMOOTH=1,
 NS=60,    NT=60,      NPSI=220,    NCHI=220, NISO=220, NTNOVA=12,
 NS=40,    NT=40,      NPSI=180,    NCHI=180, NISO=180, NTNOVA=12,
 CPRESS=1.0, QSPEC=0.70,   CSSPEC=0.0,
 CFNRESS=1.00,
 NRSCAL=0, NCSCAL=2,   NTMF0=0,
 NBAL=0,   NBLOPT=0,     CFBAL=10.00,

 NOPT=0,
 R0EXP=8.79999995E-01,
 TENSPROF= -0.3,
 TENSBND = -10,
 TENSPROF= -0.05,
 TENSBND = -0.05,

 NSURF=6,  ELONG=2.045,  TRIANG=0.7,  BEANS=0., CETA=0.24, SGMA=0.,
           ASPCT= 0.28123462E+00,
           ASPCT= 0.28,

 AT4=0.0, 27.926, -42.339, 16.664, 0.0,
 AT4=28.232, -48.496, 35.964, -14.26, 0.0,
 AT4=29500., -68768., 272720., -1147400., 2798300., -3873600., 2842600., -852840.,
 AT3=0.52719, 0.11069, -0.11517, -0.083952, 0.15572, -0.048391, 0., 0.
 AT3=0.52503, 0.92754, 0.21896, -2.4078, 8.1211, -13.87, 11.653, -3.7942,
 AT2=1.5145, 0.12811, -4.9335, 36.43, -120.88, 201.1, -163.58, 51.79,
 AT2=1.5165, 0.14189, -5.0417, 36.759, -121.11, 200.38, -162.23, 51.152,
 ETAEI= 0.1, RPEOP= 0.5,  RZION=1.5,
 NPPFUN=8, NPP=1,      AP= 0.0,-0.8, 0.0, 0.0,
 NFUNC=4,  NSTTP=2,    AT= 0., -0.30761536E+01,0.72318357, 0.0,
 NSOUR=8,

 NDIFPS=0, NDIFT=1,
 NMESHC=1, NPOIDC=2, SOLPDC=.70,CPLACE=.95,.99,1.0,
                                CWIDTH=.10,.02,.05,
 NMESHA=0, NPOIDA=1, SOLPDA=.60,APLACE=.00,.70,1.0,
                                AWIDTH=.05,.07,.05,
 NMESHA=2, NPOIDQ=6, SOLPDA=.10,QPLACE=0.95,0.95,1.0,2.0,3.00,4.0,4.0,
                                QWIDTH=0.06,.04,.08,.05,0.05,0.05,
 NMESHA=0, NPOIDQ=8, SOLPDA=.10,QPLACE=0.95,0.95,1.0,2.0,3.00,3.0 ,4.0,4.0,
                                QWIDTH=0.06,.04, .03,.06,  .05,0.01,0.04,.01
 NMESHA=2, NPOIDQ=2, SOLPDA=.60,QPLACE=1.00,1.00,1.00,2.00,2.00,2.00,3.00,3.00,4.00,4.00,
                                 QWIDTH=0.13,0.06,0.01,0.09,0.05,0.02,0.07,0.02,0.04,0.01,
 NMESHA=2, NPOIDQ=10, SOLPDA=.10,QPLACE=1.00,1.00,1.00,2.00,2.00,2.00,3.00,3.00,4.00,4.00,
                                 QWIDTH=0.13,0.06,0.01,0.09,0.05,0.02,0.07,0.02,0.04,0.01,
 NMESHA=0, NPOIDQ=10, SOLPDA=.10,QPLACE=1.00,1.00,2.00,2.00,3.00,3.00,4.00,4.00,4.41,4.41,
                                 QWIDTH=0.13,0.04,0.09,0.04,0.07,0.02,0.04,0.01,0.01,0.001,
 NMESHD=1, NPOIDD=2, SOLPDD=.60,DPLACE=-1.80,-1.80,4.0,
                                DWIDTH=.18,.08,.05,
 NMESHD=0, NPOIDD=2, SOLPDD=.60,DPLACE=-1.80,-1.80,4.0,
                                DWIDTH=.18,.08,.05,
 NMESHE=0, NPOIDE=4, SOLPDE=.50,EPLACE=-1.70,-1.70, 1.70, 1.70,
                                EWIDTH=.18,.08,.18,.08,
 EPSLON=1.0E-8, GAMMA=1.6666666667,
 NTURN=20, NBLC0=16,   NPPR=24,
 MSMAX=1,  NINMAP=100,  NINSCA=100,
 NSYM=0,   NEGP= 0,    NER=2, NV=40, NVEXP=1,  REXT=10.0, R0W=1., RZ0W=0.,
 NSYM=0,   NEGP=-2,    NER=2, NV=40, NVEXP=1,  REXT=10.0, R0W=1., RZ0W=0.,
 NSYM=0,   NEGP=-1,    NER=1, NV=40, NVEXP=1,  REXT=10.0, R0W=1., RZ0W=0.,
 NEQDXTPO=-1001,
 NEQDSK=0,  NEQDXTPO=-1,
 NEQDSK=1,  NEQDXTPO=1,
 PSISCL= 1.0,  NRBOX=33, NZBOX=33,
 PSISCL= 1.0,  NRBOX=33, NZBOX=65,
 PSISCL= 1.0,  NRBOX=129, NZBOX=129,
 PSISCL= 1.0,  NRBOX=257, NZBOX=257,
 PSISCL= 1.0,  NRBOX=157, NZBOX=157,
 PSISCL= 1.0,  NRBOX=129, NZBOX=129,
EOF

# Add override vars
cat ${override_file} >> chease_namelist

cat >>chease_namelist <<'EOF'
 /
 &NEWRUN
 AL0=-3.0E-03, NWALL= 1, REXT=10.0,
 NLGREN= .F., WNTORE=1.0, NV=40,
 NLDIAG=11*.T.,
 NAL0AUTO=1,
 /
EOF

cat chease_namelist >> log_database_to_database

#
#     Run CHEASE
#
if [ $che_opt = '1' ]
then
    cp chease_namelist $filedir_out/o.$jobname
    echo "" >> $filedir_out/o.$jobname
    $chease_exe >> $filedir_out/o.$jobname
    ls -al  >> log_database_to_database
    aaaa=`ls -1 | egrep -x 'EQDSK_COCOS_'..'.OUT'`
    cp $aaaa $filedir_out/$aaaa.$jobname
    aaaa=`ls -1 | egrep -x 'EQDSK_COCOS_'..'_POS.OUT'`
    cp $aaaa $filedir_out/$aaaa.$jobname
    cp EXPEQ.OUT $filedir_out/EXPEQ.OUT.$jobname
    cp EXPEQ.OUT.TOR $filedir_out/EXPEQ.OUT.TOR.$jobname
    cp EXPTNZ.OUT $filedir_out/EXPTNZ.OUT.$jobname
    cp NOUT $filedir_out/NOUT.$jobname
    cp EXPEQ_KEPLER.IN $filedir_out/EXPEQ_KEPLER.IN.$jobname
    date  >> log_database_to_database
    cp log_database_to_database $filedir_out/log_database_to_database.$jobname

    echo ""
    echo " Chease done"
    echo ""
fi

cd $curdir
#rm -r $tmpdir
rm ${override_file}
