#!/bin/bash

usage="
Script to tabulate data obtained from a BayesTrait analysis.

Usage:

BTrait_analyzer.sh [M/B] <BayesTrait.log.file> <Output_file_name>

where:
	-h show this help text
	[M/B] type of analysis output:
		\"M\" --> Maximum Likelihood
		\"B\" --> Bayesian"

while getopts 'ht' option
do
  case "$option" in
    h) echo "$usage"
       exit
       ;;
  esac
done

analysis_type=$1
infile=$2
outfile=$3
first_data_line=$(grep -n Tree $infile | cut -d":" -f1 | tail -n1)
last_data_line=$(cat $infile | wc -l)

summarizeML ()
{
  sed -n "$first_data_line","$last_data_line"p $infile | sed 's/ /_/g' | tr [:blank:] "\n" | sed '/^$/d' > extraction.tmp1
  row_qty=$(( $(cat extraction.tmp1 | wc -l)/2 ))
  head -n"$row_qty" extraction.tmp1 > extraction.tmp2
  tail -n"$row_qty" extraction.tmp1 > extraction.tmp3
  paste extraction.tmp2 extraction.tmp3 -d"," | grep -v ^Lh | grep -v ^q | grep -v ^Tree > extraction.tmp1
  rm extraction.tmp2 extraction.tmp3

  if [[ -z $(grep "_S" extraction.tmp1) ]]
  then
    echo "Single Character"

#    for i in $(cat Node_equiv.csv )
#    do
#      family=$(echo $i | cut -d"," -f1)
#      node=$(echo $i | cut -d"," -f2)
#      sed -i "s/^$family"_"/$node"_"/g" extraction.tmp1
#    done

    charNums=$(cut -d"(" -f2 extraction.tmp1 | cut -d")" -f1 | sort -u )

    sed -i -e 's/_P(/,0,/g;s/)//g' extraction.tmp1
    echo "Node,Character" > $outfile.probBase
    cut -d"," -f1,2 extraction.tmp1 | sort -u | sort -t"," -k1 -h >> $outfile.probBase

    for i in $charNums
    do
      intlines=$(cut -d"," -f3 extraction.tmp1 | grep -n $i | cut -d":" -f1 | tr "\n" ";" | sed 's/;$/p/g;s/;/p;/g')

      echo "Prob."$i"" > "$outfile".prob"$i"
      sed -n $intlines extraction.tmp1 | sort -t"," -k1 -h | cut -d"," -f4 >> "$outfile".prob"$i"

    done
    paste $outfile.probBase $(ls "$outfile".prob[0-9] | sort -h) -d"," > $outfile

    rm *.prob*

  else


    echo "Multiple Characters"
#    for i in $(cat Node_equiv.csv )
#    do
#      family=$(echo $i | cut -d"," -f1)
#      node=$(echo $i | cut -d"," -f2)
#      sed -i "s/^$family,/$node,/g" extraction.tmp1
#    done
    sed -i -e 's/_-_S(/,/g;s/)_-_P(/,/g;s/)//g' extraction.tmp1


    zero_lines=$(cut -d"," -f3 extraction.tmp1 | grep -n "0" | cut -d":" -f1 | tr "\n" ";" | sed 's/;$/p/g;s/;/p;/g')
    one_lines=$(cut -d"," -f3 extraction.tmp1 | grep -n "1" | cut -d":" -f1 | tr "\n" ";" | sed 's/;$/p/g;s/;/p;/g')

    echo "Node,Character,Prob.0" > "$outfile".tmp1
    sed -n $zero_lines extraction.tmp1 | sort -t"," -k2 -h | cut -d"," -f1,2,4 >> "$outfile".tmp1
    echo "Prob.1" > "$outfile".tmp2
    sed -n $one_lines extraction.tmp1 | sort -t"," -k2 -h | cut -d"," -f4 >> "$outfile".tmp2

    paste "$outfile".tmp1 "$outfile".tmp2 -d"," > $outfile
  fi

}

summarizeBI ()
{

  if [[ -n $(grep "S([0-9])" $infile) ]]
  then
    echo "Multiple Characters"

    sed -n "$first_data_line","$last_data_line"p $infile | tr "\t" "," | sed -e 's/ /_/g;s/_-_S(/_/g;s/_-_P(/_/g;s/)//g' | sed '/^$/d' > extraction.tmp1
    Rscript -e 'inData<-read.csv("extraction.tmp1",header=T); inData<-tail(na.exclude(as.data.frame(colMeans(inData))),n=-6); colnames(inData)<-c("Node,Character,State,Prob"); write.csv(inData,"extraction.tmp1",quote=F)'
    sed -i 's/^,//g;s/_/,/g' extraction.tmp1

#    for i in $(cat Node_equiv.csv )
#    do
#      family=$(echo $i | cut -d"," -f1)
#      node=$(echo $i | cut -d"," -f2)
#      sed -i "s/^$family,/$node,/g" extraction.tmp1
#    done

    zero_lines=$(cut -d"," -f3 extraction.tmp1 | grep -n "0" | cut -d":" -f1 | tr "\n" ";" | sed 's/;$/p/g;s/;/p;/g')
    one_lines=$(cut -d"," -f3 extraction.tmp1 | grep -n "1" | cut -d":" -f1 | tr "\n" ";" | sed 's/;$/p/g;s/;/p;/g')
    echo "Node,Character,Prob.0" > "$outfile".tmp1
    sed -n $zero_lines extraction.tmp1 | sort -t"," -k2 -h | cut -d"," -f1,2,4 >> "$outfile".tmp1
    echo "Prob.1" > "$outfile".tmp2
    sed -n $one_lines extraction.tmp1 | sort -t"," -k2 -h | cut -d"," -f4 >> "$outfile".tmp2

    paste "$outfile".tmp1 "$outfile".tmp2 -d"," > $outfile

  else
    echo "Single Character"

    sed -n "$first_data_line","$last_data_line"p $infile | tr "\t" "," | sed -e 's/ /_/g;s/_P(/_/g;s/)//g'  > extraction.tmp1
    dataCol1=$(head -n1 extraction.tmp1 | tr "," "\n" | grep -n "^Root" | cut -d":" -f1 | head -n1)
    stateNums=$(head -n1 extraction.tmp1 | tr "," "\n" | grep "_" | grep -v Mean | grep -v ^Tree | cut -d"_" -f2 | sort -u)
    stateQty=$(echo $stateNums | wc -w)
    
    mv extraction.tmp1 extraction.tmp2
    cut -d"," -f$dataCol1- extraction.tmp2 | sed 's/,$//g' > extraction.tmp1
    cut -d"," -f1-$(($dataCol1 -1)) extraction.tmp2 | sed 's/,$//g' > means.tmp1
    rm extraction.tmp2

    Rscript -e 'inData<-read.csv("extraction.tmp1",header=T); inData<-na.exclude(as.data.frame(colMeans(inData))); colnames(inData)<-c("Node,State,Prob"); write.csv(inData,"extraction.tmp1",quote=F); inData<-read.csv("means.tmp1",header=T); inData<-na.exclude(as.data.frame(colMeans(inData))); colnames(inData)<-c("Parameter,Mean"); write.csv(inData,"means.tmp1",quote=F)'
    sed 's/^,//g' means.tmp1 > "$outfile"_Means.csv
    sed -i 's/^,//g;s/_/,/g' extraction.tmp1

#    for i in $(cat Node_equiv.csv )
#    do
#      family=$(echo $i | cut -d"," -f1)
#      node=$(echo $i | cut -d"," -f2)
#      sed -i "s/^$family,/$node,0,/g" extraction.tmp1
#    done

    echo "Node,Character" > $outfile.probBase
    cut -d"," -f1,2 extraction.tmp1 | tail -n+2 | sort -u | sort -t"," -k1 -h >> $outfile.probBase

    for i in $stateNums
    do
      intlines=$(cut -d"," -f3 extraction.tmp1 | grep -n $i | cut -d":" -f1 | tr "\n" ";" | sed 's/;$/p/g;s/;/p;/g')

      echo "Prob."$i"" > "$outfile".prob"$i"
      sed -n $intlines extraction.tmp1 | sort -t"," -k1 -h | cut -d"," -f4 >> "$outfile".prob"$i"

    done

    paste $outfile.probBase $(ls "$outfile".prob[0-9] | sort -h) -d"," > $outfile
    rm *.prob*

  fi

}


if [[ $analysis_type == "M" ]]
then
  summarizeML
elif [[ $analysis_type == "B" ]]
then
  summarizeBI
else
  echo "Option $analysis_type is not valid. Try using \"M\" or \"B\""
fi

rm *tmp*
