#!/bin/bash

source $MWSHPATH/colours.sh
source $MWSHPATH/out.sh

SEARCH_RTP=0
REPLACE=0
FIND=0
FORCE=0
EXPERT=0
VERBOSITY=0
PROGRESS=0
TOP_PATH=/opt/proprietary-apps/gromacs/2018/share/gromacs/top/

while test $# -gt 0; do
  case "$1" in
    -h|--help)
      echo -e "Usage for "$colFunc"pdb_mod.sh"$colClear":"
      echo -e $colArg"-s <PDB>"$colClear" summary of PDB contents"
      exit 5
      ;;
    # -s|--summary)
    #   shift
    #   PDB=$1
    #   shift

    #   HEADER=$(grep "HEADER" $PDB | grep -v "REVDAT" | grep -oP "(?<=HEADER    ).*" | awk '{print $1}')
    #   KEYWDS=$(grep "KEYWDS" $PDB | grep -oP "(?<=KEYWDS    ).*")
    -ff|--force-field)
      shift
      FIELD=$1
      shift
      ;;
    -srtp|--search-rtp)
      shift
      RES=$1
      SEARCH_RTP=1
      shift
      ;;
    -matm|--missing-atom)
      shift
      MATM=$1
      shift
      ;;
    -r|--replace)
      shift
      MATM=$1
      shift
      NEWATM=$1
      shift
      REPLACE=1
      ;;
    -fatm|--find-atom)
      shift
      MATM=$1
      shift
      FIND=1
      ;;
    -i|--input)
      shift
      INPUT=$1
      shift
      ;;
    -o|--output)
      shift
      OUTPUT=$1
      shift
      ;;
    -filter|--filter)
      shift
      FILTER=$1
      shift
      ;;
    -v)
      shift
      VERBOSITY=1
      ;;
    -f)
      shift
      FORCE=1
      ;;
    -l)
      shift
      LOGNUM=$1
      shift
      ;;
    -x|--expert)
      shift
      EXPERT=1
      ;;
    # -p)
    #   shift
    #   # PROGRESS=1
    #   ;;
    *)
      warningOut "Unrecognised CLI flag: $colArg$1"
      break
      ;;
  esac
done

LOGFILE=_pdbmod"$LOGNUM".log

if [[ -z $FIELD ]] ; then
  errorOut "No force field specified."
fi

TOP_PATH=/opt/proprietary-apps/gromacs/2018/share/gromacs/top/$FIELD.ff/

if [ $SEARCH_RTP -eq 1 ] ; then

  FOUND=0
  for FILE in "$TOP_PATH*.rtp"; do
    # echo $FILE
    FOUND=$(grep -n "\[ $RES \]" -m1 $FILE | wc -l)
    if [ $FOUND -eq 1 ] ; then break ; fi
  done

  if [ $FOUND -eq 0 ] ; then
    errorOut "No $RES found in $FIELD!" 1
  fi

  # grep -n "\[ $RES \]" -m1 $FILE
  LINE_NUM=$(grep -n "\[ $RES \]" $FILE | awk '{print $1}')
  LINE_NUM=${LINE_NUM::-2}
  # { grep -n -m1 "\[ $RES \]"; grep -n -m1 "\[ "; } << $TOP_PATH/*.rtp

  RTP_LEN=$(cat $FILE | wc -l)
  let "TAIL = $LINE_NUM - $RTP_LEN"
  # tail $TOP_PATH/*.rtp -n $TAIL | grep -n -m1 '^$' $TOP_PATH/*.rtp
  # tail $TOP_PATH/*.rtp -n $TAIL | grep -n -m5 "\[" $TOP_PATH/*.rtp

  tail $FILE -n $TAIL > __temp__

  # echo $RES > __temp__2

  headerOut "$colArg$RES$colClear$colBold entry in $colFunc$FIELD"
  headerOut "$colFile$FILE$colClear"
  headerOut "line $LINE_NUM and onwards:"

  while read -r LINE; do
    # if [[ $LINE != "[ atoms ]" ]] && [[ $LINE == "[ *" ]] ; then
    #   exit
    # fi
    
    if [[ $LINE == "" ]] ; then exit ; fi
    
    # echo "$LINE" >> __temp__2

    if [[ ! -z $MATM ]] ; then
      if [ $(echo "$LINE" | grep "${MATM:0:1}" | wc -l) -eq 1 ] ; then
        echo "$LINE" | grep -E --color "${MATM:0:1}|${MATM:1:1}|${MATM:2:1}"
      else
        echo "$LINE"
      fi
    else
      echo "$LINE"
    fi

  done < __temp__

  rm __temp__*

  # if [[ -z $MATM ]] ; then


fi

if [ $FIND -eq 1 ] ; then

  if [[ -z $INPUT ]] ; then 
    errorOut "No input file specified."
    exit
  fi

  echo -e "$colBold$colArg$MATM$colClear$colBold appears in "$colClear"$colFile$colBold$INPUT$colClear$colBold in$colClear:"
  echo -e "$colUnderline$colVarName""residue$colClear $colUnderline$colVarName""chain$colClear $colResult$colUnderline""instances"$colClear

  grep "ATOM" $INPUT | grep "  $MATM  " | awk '{print $4" "$5}' > __temp__

  # nano __temp__

  touch __temp__2

  while read -r RESCHAIN; do
    if [ $(grep "$RESCHAIN" __temp__2 | wc -l) -ne 0 ] ; then continue ; fi
    echo $RESCHAIN >> __temp__2
  done < __temp__
  
  # nano __temp__2

  while read -r RESCHAIN; do
    echo -e $colVarName$RESCHAIN $colResult$(grep "$RESCHAIN" __temp__ | wc -l)"x" $colClear
  done < __temp__2

  rm __temp__*

fi

if [ $REPLACE -eq 1 ] ; then


  if [[ -z $OUTPUT ]] ; then 
    if [ $EXPERT -eq 0 ] ; then
      warningOut "No output file specified. Defaulted to $colFile""mod.pdb"
    fi
    OUTPUT="mod.pdb"
  fi
  if [[ -z $INPUT ]] ; then 
    errorOut "No input file specified."
    exit
  fi
  if [[ $INPUT == $OUTPUT ]] ; then 
    if [ $FORCE -eq 1 ] ; then
      if [ $EXPERT -eq 0 ] ; then
        warningOut "Overwriting $colFile"$INPUT"$colWarning!"
      fi
    else
      errorOut "Overwriting is not permitted. Select different input and output files."
      exit
    fi
  fi

  if [[ -z $INPUT ]] ; then 
    errorOut "No input file specified."
    exit
  fi

  if [[ "$OUTPUT" == "__temp__" ]] ; then 
    errorOut "Choose a different output name."
    exit
  fi

  if [[ -z $FILTER ]] ; then 
    if [ $EXPERT -eq 0 ] ; then
      warningOut "All instances of $colArg$MATM$colClear$colWarning being replaced."
    fi
  fi
  
  if [ $VERBOSITY -gt 0 ] ; then
    if [[ -z $FILTER ]] ; then
      echo -ne $colFunc"pdb_mod.sh"$colClear": Renaming $colArg$MATM$colClear to $colArg$NEWATM$colClear... "
    else
      echo -ne $colFunc"pdb_mod.sh"$colClear": Renaming $colArg$MATM$colClear to $colArg$NEWATM$colClear in RESCHAINs containing $colArg$FILTER$colClear... "
    fi
  fi

  rm __temp__ 2> /dev/null
  if [[ ! $INPUT == $OUTPUT ]] ; then
    rm $OUTPUT 2> /dev/null
  fi

  COUNT=0

  # IN_LINES=$(cat $INPUT | wc -l)

  while read -r LINE; do

    # if [ $PROGRESS -eq 1 ] ; then echo ; fi

    if [[ $LINE == ATOM*$MATM* ]] ; then
      OLSTRING="     "
      OLSTRING="  $MATM""${OLSTRING:${#MATM}}"
      NEWSTRING="     "
      NEWSTRING="  $NEWATM""${NEWSTRING:${#NEWATM}}"
      # echo "$LINE"
      if [[ -z $FILTER ]] ; then
        # echo "$LINE" | sed "s/  $MATM  /  $NEWATM  /" >> __temp__
        # echo "$LINE" | sed "s/  $MATM  /  $NEWATM  /" >> $LOGFILE
        echo "$LINE" | sed "s/$OLSTRING/$NEWSTRING/" >> __temp__
        echo "$LINE" | sed "s/$OLSTRING/$NEWSTRING/" >> $LOGFILE
        let "COUNT = COUNT + 1"
      else
        RESCHAIN=$(echo "$LINE" | awk '{print $4" "$5}')
        if [[ $RESCHAIN == *$FILTER* ]] ; then
          echo "$LINE" | sed "s/$OLSTRING/$NEWSTRING/" >> __temp__
          echo "$LINE" | sed "s/$OLSTRING/$NEWSTRING/" >> $LOGFILE
          let "COUNT = COUNT + 1"
        else
          echo "$LINE" >> __temp__
        fi
      fi
    else
      echo "$LINE" >> __temp__
    fi
  done < $INPUT

  # if [ $PROGRESS -eq 1 ] ; then echo ; fi

  mv __temp__ $OUTPUT

  if [ $COUNT -eq 0 ] ; then
    if [ $EXPERT -eq 0 ] ; then
      echo ""
      warningOut "No replacements made"
    fi
  else
    successOut "$COUNT replacements made"
  fi

fi