#!/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 eqdsk file in local directory and write to the local database (tree ids or euitm)
#
# eqdsk_to_database eqdsk_filename shot run [cocos_in [cocos_out [treename [ipsign_out [b0sign_out [tree_tokamak [tree_user [tree_majorversion [nitmopt [tree_backend]]]]]]]]]]
#
# defaults: cocos_in=2, cocos_out=11 (or 13)
# 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
#
# Note: if cocos_in=2 does not work, try 7
#
# Typically call for ITM/cpo database:
#         eqdsk_to_database eqdskfilename shot run cocos_in 13
#
# Typically call for ITMIMAS/ids database:
#         eqdsk_to_database eqdskfilename shot run cocos_in 11 ids -1 -1
#
# 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* or at the *END* of the list of arguments.
# source grab_override # could use same source in database_to_eqdsk, 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_out=11
    chease_exe=${chease_exe:-chease_imas}
else
    # assume on gateway
    xitmenvx=x`echo $ITM_ENVIRONMENT_LOADED`x
    if [ $xitmenvx = 'xx' ]
    then
        #for ids database
        default_treename='ids'
        cocos_out=11
        chease_exe=${chease_exe:-chease_imas}
    else
        #for euitm database
        default_treename='euitm'
        cocos_out=13
        chease_exe=${chease_exe:-chease_itm}
    fi
fi
echo $chease_exe
cocos_in=2
tree_tokamak_out='test'
tree_user_out=`echo $USER`
tree_majorversion_out='3'
nitmopt_out=10
tree_backend_out=13

if [ ${non_override_seen} -le 2 -o ${non_override_seen} -ge 14 ]
then
   echo "needs 3 to 12 arguments: eqdsk_filename shot run [cocos_in [cocos_out [treename [ipsign_out [b0sign_out [tree_tokamak [tree_user [tree_majorversion [nitmopt [tree_backend]]]]]]]]]]"
   echo "eqdsk_filename is assumed to be in the local directory"
   echo " default cocos_in = "$cocos_in
   echo "Note: if cocos_in=2 does not work, try 7"
   echo " default cocos_out = "$cocos_out
   echo " default treename: "$default_treename
   echo "for ITER might want to impose negative Ip and B0 thus set ipsign_out=-1 and b0sign_out=-1"
   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 tree_tokamak_out: "$tree_tokamak_out
   echo " default tree_user_out: "$tree_user_out
   echo " default tree_majorversion_out: "$tree_majorversion_out
   echo " default nitmopt (10 creates, 20 opens on an existing database case): "$nitmopt_out
   echo " default tree_backend (13 hdf5, 12 mdsplus): "$tree_backend_out
   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

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

ipsign_out=-9
if [ ${non_override_seen} -ge 7 ]
then
  ipsign_out=$7
fi
b0sign_out=-9
if [ ${non_override_seen} -ge 8 ]
then
  b0sign_out=$8
fi
if [ ${non_override_seen} -ge 9 ]
then
  tree_tokamak_out=$9
fi
if [ ${non_override_seen} -ge 10 ]
then
  tree_user_in=${10}
fi
if [ ${non_override_seen} -ge 11 ]
then
  tree_majorversion_out=${11}
fi
if [ ${non_override_seen} -ge 12 ]
then
  nitmopt_out=${12}
fi
if [ ${non_override_seen} -ge 13 ]
then
  tree_backend_out=${13}
fi

# to restart from previous equilibrium and change B0EXP (B0EXP=-1.4)
NOUTname=$eqdskfilename

expeqfile=$eqdskfilename

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_eqdsk_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_out," >> chease_namelist
echo " NITMSHOT= 9999, $shot_number," >> chease_namelist
echo " NITMRUN= 9999, $run_number," >> 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_out','$tree_tokamak_out'" >> chease_namelist
echo " TREE_USER='$tree_user_out','$tree_user_out'" >> chease_namelist
echo " TREE_MAJORVERSION=$tree_majorversion_out,$tree_majorversion_out" >> chease_namelist
echo " TREE_BACKEND=$tree_backend_out,$tree_backend_out" >> chease_namelist
echo " SIGNIPXP=$ipsign_out," >> chease_namelist
echo " SIGNB0XP=$b0sign_out," >> chease_namelist
cat >>chease_namelist <<'EOF'
 TIME_REF = 1.,
 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=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=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_eqdsk_to_database

#
#     Run CHEASE
#
if [ $che_opt = '1' ]
then
    cp chease_namelist $filedir_out/o.$jobname
    echo "" >> $filedir_out/o.$jobname
    cp $filedir_in/$expeqfile EQDSK
    cp EQDSK EXPEQ

    echo " " >> $filedir_out/o.$jobname
    $chease_exe >> $filedir_out/o.$jobname
    ls -al  >> log_eqdsk_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
    date  >> log_eqdsk_to_database
    cp log_eqdsk_to_database $filedir_out/log_eqdsk_to_database.$jobname
    echo " "
    echo ""
    echo " Chease done"
    echo ""
fi

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