#!/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 an eqdsk file in local directory
#
# database_to_eqdsk eqdskfilename shot run [cocos_in [cocos_out [occurence_nb [treename [ipsign_out [b0sign_out [time_in]]]]]]
#
# 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_HDF5=1 (to run with NIDEAL=9 for ORB5)
#      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.
# source grab_override # could use same source in eqdsk_to_database, but would not be backward compatible since requesting an extra executable in the path
############# 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=0
ipsign_out=-9
b0sign_out=-9
time_in=1.
tree_tokamak_in='test'
tree_user_in=`echo $USER`
tree_majorversion_in='3'

if [ ${non_override_seen} -le 2 -o ${non_override_seen} -ge 14 ]
then
   echo "needs 3 to 13 arguments: eqdsk_filename shot run [cocos_in [cocos_out [occurence_nb [treename [ipsign_out [b0sign_out [time_in] [tree_tokamak [tree_user [tree_majorversion]]]]]]]]]"
   echo "eqdsk_filename is assumed to be in the local directory"
   echo " default cocos_in = "$cocos_in
   echo " default cocos_out = "$cocos_out
   echo " default occurence_nb = "$occurence_nb
   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)"
   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"
   echo " default ipsign_out: "$ipsign_out
   echo " default b0sign_out: "$b0sign_out
   echo " default time_in: "$titime_inme_in
   echo " default tree_tokamak_in: "$tree_tokamak_in
   echo " default tree_user_in: "$tree_user_in
   echo " default tree_majorversion_in: "$tree_majorversion_in
   echo " "
   echo "All output files are in /tmp/username/."
   echo " "
   exit
fi
eqdskfilename=$1
jobname=$eqdskfilename

shot_number=$2
run_number=$3

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

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

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

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

if [ ${non_override_seen} -ge 8 ]
then
  ipsign_out=$8
fi
if [ ${non_override_seen} -ge 9 ]
then
  b0sign_out=$9
fi

if [ ${non_override_seen} -ge 10 ]
then
  time_in=${10}
fi
if [ ${non_override_seen} -ge 11 ]
then
  tree_tokamak_in=${11}
fi
if [ ${non_override_seen} -ge 12 ]
then
  tree_user_in=${12}
fi
if [ ${non_override_seen} -ge 13 ]
then
  tree_majorversion_in=${13}
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_eqdsk

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.,  NDIAGOP= 1,  NBSEXPQ=0000,
 NITMOPT = 1,
EOF
echo " NITMSHOT= $shot_number, 9999," >> chease_namelist
echo " NITMRUN= $run_number, 9999," >> chease_namelist
echo " NITMOCC= $occurence_nb, 9999" >> 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_in'" >> chease_namelist
echo " TREE_USER='$tree_user_in','$tree_user_in'" >> chease_namelist
echo " TREE_MAJORVERSION=$tree_majorversion_in,$tree_majorversion_in" >> 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=4, NPP=1,      AP= 0.0,-0.8, 0.0, 0.0,
 NFUNC=4,  NSTTP=1,    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=40,  NINSCA=40,
 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

if [ 'x'$CHEASE_HDF5'x' = 'x1x' ]
then
  # For generating ogyropsi.h5
  sed -r -i 's/NIDEAL=[ ]*[0-9]+/NIDEAL=9/' chease_namelist
  sed -r -i 's/NER=[ ]*[0-9]+/NER=2/' chease_namelist
  sed -r -i 's/NEGP=[ ]*[+-]?[0-9]+/NEGP=0/' chease_namelist
  echo "CHEASE_HDF5: " $CHEASE_HDF5
fi

cat chease_namelist >> log_database_to_eqdsk

#
#     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_eqdsk
    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
    if [ 'x'$CHEASE_HDF5'x' = 'x1x' ]
    then
      cp ogyropsi.h5 $filedir_out/ogyropsi.$jobname.h5
      cp ogyropsi.dat $filedir_out/ogyropsi.$jobname.dat
    fi
    date  >> log_database_to_eqdsk
    cp log_database_to_eqdsk $filedir_out/log_database_to_eqdsk.$jobname

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

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