#!/usr/bin/env bash

[[ $# -ge 3 ]] || {
    echo "usage $0 <srcdir> <trgdir> <node> [<nmax>]"
    cat <<EOF

This script subsamples fasta files for the given <node>.
From each cluster a maximum of <nmax> genes are selected from a list with
increasing score (or decreasing if last parameter starts with 'd').
Each subsampled fasta is saved by taxid in <trgdir>/g<node>_<taxid>.fs_selected .
Also the list of all genes per taxid is save at <trgdir>/g<node>_<taxid>.txt .

It prints to stdout the new fasta file names, including path.

EOF
    exit 0
}


export DIR_SCRIPT="$(cd -P "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
[[ -x $DIR_SCRIPT/subsample_cluster_genes.py ]] || { echo "error: missing script $DIR_SCRIPT/subsample_cluster_genes.py"; exit 255; }
[[ -x $DIR_SCRIPT/group_by_species.py        ]] || { echo "error: missing script $DIR_SCRIPT/group_by_species.py";        exit 255; }
[[ -x $DIR_SCRIPT/get_subsample_nmax.py      ]] || { echo "error: missing script $DIR_SCRIPT/get_subsample_nmax.py";      exit 255; }


calc(){ awk "BEGIN{ print $* }" ;}
error_msg() {
     echo "#ERROR: ${@}" 1>&2
}
warn_msg() {
    echo "#WARNING: ${@}"
}
info_msg() {
    echo "#INFO: ${@}"
}

#===============================================================
# NOTE CODE COPIED FROM mapping_func.sh
#---------------------------------------------------------------
# get score scan direction
# * des     sort genes by ascending score
# * asc     sort genes by ascending score
# * sco     sort genes first by overall species coverage and then by score
# * sde     des, select stepwise
# * sas     asc, select stepwise
# * ssc     sco, select stepwise
map_get_scdir_value() {
    [[ $# -eq 1 ]] || return 255
    [[ -n "$1" ]] || { echo "ssc"; return 0; }
    local v="$1"
    v=${v,,}
    if [[ ${v:0:3} == "des" ]]; then
        # sort genes by decreasing score
        echo "des"
    elif [[ ${v:0:3} == "asc" ]]; then
        # sort genes by ascending score
        echo "asc"
    elif [[ ${v:0:3} == "sde" ]]; then
        # sort genes by decreasing score, select stepwise
        echo "sde"
    elif [[ ${v:0:3} == "sas" ]]; then
        # sort genes by ascending score, select stepwise
        echo "sas"
    elif [[ ${v:0:3} == "sco" ]]; then
        # sort genes first by overall species coverage and then by score
        echo "sco"
    elif [[ ${v:0:3} == "ssc" ]]; then
        # like sco, select stepwise
        echo "ssc"
    else
        return 255
    fi
    return 0
}

map_get_subsample_value() {
    [[ $# -eq 1 ]] || { echo "0"; return 0; }
    [[ -n "$1"  ]] || { echo "0"; return 0; }
    echo "$1"
}

subsample_label() {
    [[ $# -eq 3 ]] || { error_msg "subsample_label: missing args"; return 255; }
    local nodeid="$1"
    local subsample=$(map_get_subsample_value "$2")
    [[ "$subsample" == "0" ]] && return 0
    local scdir=$(map_get_scdir_value "$3")
    [[ -n $scdir ]] || { error_msg "subsample_label: illegal scdir argument $3"; return 255; }
    echo "g${nodeid}_n${subsample}${scdir}"
}

#===============================================================

#
# Subsample fast for given node.
# <srcdir>      root dir of pipeline project
# <trgdir>      where to output subsampled fasta
# <node>        which node to consider
# <scdir>       subsampling startegy - 'ssc' by default
#
subsample_fasta() {
    local srcdir="$1"
    local trgdir="$2"
    local node="$3"
    local subsample="$4"
    local scdirOpt=""
    local scdir="ssc"
    [[ $# -eq 5 ]] && scdir="$5"
    scdir=$(map_get_scdir_value "$scdir")
    
    local Nmax=$subsample
    
    # make sure args are ok
    [[ -n $scdir ]] || {
        map_error_msg "subsample_fasta: illegal scdir argument $5"
        return 255
    }
    (( subsample >0 )) || {
        map_error_msg "subsample_fasta: subsample must be non-zero positive number ($subsamble)"
        return 255
    }
    
    # input cluster
    local fcluster=$srcdir/Cluster/${node}.og
    
    # input spcov, if used
    local fspcov=$srcdir/Cluster/${node}.spcov

    # todo file, for stats
    local ftodo=$srcdir/todo/${node}.todo
    
    # start from best score
    [[ "$scdir" == "des" ]] && scdirOpt='-r'
    
    # start from best score, subsample by steps
    [[ "$scdir" == "sde" ]] && scdirOpt='-r -d'

    # start from worst score, subsample by steps
    [[ "$scdir" == "sas" ]] && scdirOpt='-d'

    # select by sorting genes in cluster first by species coverage and then by score
    local checkSpcov=0
    [[ "$scdir" == "sco" ]] && { scdirOpt="-c $fspcov"; checkSpcov=1; }

    # sco, but stepwise selection
    [[ "$scdir" == "ssc" ]] && { scdirOpt="-c $fspcov -d"; checkSpcov=1; }

    [[ -f $fcluster ]] || { error_msg "missing cluster file $fcluster"; return 255; }
    local facc=0.97
    if (( checkSpcov == 1 )); then
	    [[ -f $fspcov   ]] || { map_error_msg "missing coverage file $fspcov";  return 255; }
        Nmax=$($DIR_SCRIPT/get_subsample_nmax.py -c $fspcov -m $subsample -f $facc)
    fi
    
    info_msg "subsample from cluster $fcluster"
    
    [[ -d $srcdir           ]] || { error_msg "missing source dir at $srcdir"; return 255; }
    [[ -d $srcdir/Sequences ]] || { error_msg "missing Sequences dir at $srcdir/Sequences"; return 255; } 
    [[ -d $srcdir/Cluster   ]] || { error_msg "missing Cluster dir at $srcdir/Cluster"; return 255; }
    [[ -d $trgdir           ]] || { error_msg "missing target dir at $trgdir"; return 255; }
    [[ -f $fcluster         ]] || { error_msg "missing cluster file at $fcluster"; return 255; }
    
    # use the input N=subsample for label and the deduced Nmax for calculation
    local label="$(subsample_label $node $subsample $scdir)"
    info_msg "subsample from node $node , label $label with options -n $Nmax $scdirOpt"
            
    local fgenes=( $($DIR_SCRIPT/subsample_cluster_genes.py -o $fcluster -n "$Nmax" $scdirOpt | sort | $DIR_SCRIPT/group_by_species.py -d $trgdir --label $label ) )
    [[ $? -eq 0 ]] || { error_msg "failed creating fgenes"; return 255; }

    local fractions="$trgdir/${label}.log"
    local subsampledTodo="$trgdir/../todo/${label}.todo"
    local data=
    local saveIFS=$IFS
    {
        echo "# facc  = $facc"
        echo "# Nmax  = $Nmax ($subsample)"
        echo "# scdir = $scdirOpts"
        for data in ${fgenes[@]}; do
            IFS=','
            elems=( $data )
            IFS=$saveIFS
            local taxid=${elems[0]}
            local fgeneList=${elems[1]}
            b=$(basename $fgeneList .txt)
            ffasta=$trgdir/${b}.fs_selected
            fsrc=$srcdir/Sequences/${taxid}.fs_selected
            Nfsrc=$(grep -c '>' $fsrc)
            [[ -f $fgeneList ]] || { error_msg "missing genes file $fgeneList" 1>&2;   continue; }
            [[ -f $fsrc      ]] || { error_msg "missing source file $fsrc" 1>&2;  continue; }

            # output label for subsampled todo
            echo $b

            [[ -f $ffasta    ]] && continue #{ echo "warning: skip existing file $ffasta" 1>&2; continue; }
            {
                echo "# N(src)    $Nfsrc"
                grep -f $fgeneList --no-group-separator -A 1 $fsrc
            }> $ffasta
            Nfasta=$(grep -c '>' $ffasta)
            {
                echo "# fgenes    $fgenes"
                echo "# fsrc      $fsrc"
                echo "# N(fasta)  $Nfasta"
                echo "#END"
            }>> $ffasta
            
        done
    }> $subsampledTodo

    # output stats
    local Nremoved=0
    local Ntotal=0
    local ftmp=$(mktemp /tmp/fsample_XXXXX)
    while read line; do
        [[ -n $line ]] || continue
        [[ "${line:0:1}" == "#" ]] && continue
        forg=$srcdir/Sequences/${line}.fs_selected
        fsub=$trgdir/${label}_${line}.fs_selected
        Ntotal=$((Ntotal+1))
        if [[ -f $fsub ]]; then
            Norg=$(grep -c '>' $forg)
            Nsub=$(grep -c '>' $fsub)
            printf "%s %.4f\n" "$line" $(calc "$Nsub/$Norg")
        else
            Nremoved=$((Nremoved+1))
            echo "$line  0"
        fi
    done < $ftodo > $ftmp

    # fraction removed taxids
    local frem=0
    if (( Ntotal > 0 )); then
        if (( Nremoved > 0 )); then
            frem=$(calc "1.0*$Nremoved/$Ntotal")
        fi
    fi
    
    # stats: ndata min max mean median sum variance
    local stats=($($DIR_SCRIPT/stats.py -f $ftmp  -c 1))
    {
        cat $ftmp
        printf "# Fraction per file remaining min/max/mean : %.2f/%.2f/%.2f\n" "${stats[1]}" "${stats[2]}" "${stats[3]}"
        printf "# Total fraction remaining : %.2f\n" $(calc "${stats[3]}*(1.0-$frem)")
        if (( Ntotal > 0 )); then
            if (( Nremoved == 0 )); then
                printf "# Fraction taxids removed NONE\n"
            else
                printf "# Fraction taxids removed : %.2f%%\n" $(calc "100.0*$frem")
            fi
        fi
    } > $fractions

    rm -f $ftmp
}

SRCDIR="$1"
TRGDIR="$2"
NODE="$3"

NMAX="8"
[[ $# -eq 4 ]] && NMAX="$4"

subsample_fasta "$SRCDIR" "$TRGDIR" "$NODE" "$NMAX"
