#!/bin/bash

set -euo pipefail

ln -s ../12.orthofinder/ogs_cds/ .

# parallel MACSE alignment for each orthologous group
mkdir alignments
echo -n >macse.cmds
for og_fasta in ogs_cds/*.fa; do
    og_name=$(basename "$og_fasta" .fa)
    out_stem="alignments/$og_name"
    echo "macse -prog alignSequences -seq $og_fasta -out_AA ${out_stem}.pal -out_NT ${out_stem}.nal > ${out_stem}.log 2>&1" >>macse.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a macse.cmds bash -c "cmd"

# "!" --> "-"
rg -l "!" alignments/*.nal >fs.list
mkdir alignments_clean
for og_fasta in ogs_cds/*.fa; do
    og_name=$(basename "$og_fasta" .fa)
    out_stem="alignments/$og_name"
    echo "Processing $og_name"
    sed 's/!/-/g' "${out_stem}.nal" >"alignments_clean/${og_name}.nal"
done

# classify frameshift loci by independent event count and carrier pattern
python3 "$PROJECTHOME"/99.scripts_v2/utils/classify_frameshifts.py \
    --fasta-list fs.list \
    --output-dir frameshift_summary

# trim alignments with trimAl
mkdir -p alignments_trimmed
echo -n >trimal.cmds
for i in alignments_clean/*.nal; do
    j=$(basename "$i" .nal)
    out_dir="alignments_trimmed"
    echo "trimal -in $i -out $out_dir/${j}.trimed.fa -automated1" >>trimal.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a trimal.cmds bash -c "cmd"

# retain MSAs at least 300 bp long
min_alignment_length=300
mkdir -p alignments_filtered
printf "gene\talignment_length_bp\n" >alignments_shorter_than_300bp.tsv
printf "gene\tremoved_sequence\n" >trimal_removed_sequences.tsv
for i in alignments_trimmed/*.fa; do
    [[ -e "$i" ]] || continue
    j=$(basename "$i" .trimed.fa)
    untrimmed_alignment="alignments_clean/${j}.nal"
    filtered_alignment="alignments_filtered/${j}.fa"
    if [[ ! -s "$untrimmed_alignment" ]]; then
        printf "Missing non-empty pre-trim alignment: %s\n" "$untrimmed_alignment" >&2
        exit 1
    fi

    removed_sequences=()
    mapfile -t removed_sequences < <(
        comm -23 \
            <(awk '/^>/{sub(/^>/, ""); sub(/[[:space:]].*$/, ""); print}' \
                "$untrimmed_alignment" | sort -u) \
            <(awk '/^>/{sub(/^>/, ""); sub(/[[:space:]].*$/, ""); print}' \
                "$i" | sort -u)
    )
    if ((${#removed_sequences[@]} > 0)); then
        rm -f "$filtered_alignment"
        for sequence in "${removed_sequences[@]}"; do
            printf "%s\t%s\n" "$j" "$sequence" >>trimal_removed_sequences.tsv
        done
        continue
    fi

    alignment_length=$(awk '
        /^>/ {
            if (found_sequence)
                exit
            found_sequence = 1
            next
        }
        found_sequence {
            gsub(/[[:space:]]/, "")
            sequence_length += length($0)
        }
        END { print sequence_length + 0 }
    ' "$i")

    if ((alignment_length >= min_alignment_length)); then
        if [[ ! "$i" -ef "$filtered_alignment" ]]; then
            ln -f "$i" "$filtered_alignment"
        fi
    else
        rm -f "$filtered_alignment"
        printf "%s\t%d\n" "$j" "$alignment_length" >>alignments_shorter_than_300bp.tsv
    fi
done

# raxml-ng for each filtered orthologous group
mkdir -p raxml
echo -n >raxml.cmds
for i in alignments_filtered/*.fa; do
    [[ -e "$i" ]] || continue
    j=$(basename "$i" .fa)
    out_prefix="raxml/${j}/${j}"
    mkdir -p "raxml/${j}"
    echo "$SCRIPTS/utils/raxml.sh $i $out_prefix 1 ZJ" >>raxml.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a raxml.cmds bash -c "cmd"

# mrbayes for each filtered orthologous group
mkdir -p mrbayes
echo -n >mrbayes.cmds
for i in alignments_filtered/*.fa; do
    [[ -e "$i" ]] || continue
    j=$(basename "$i" .fa)
    out_prefix="mrbayes/${j}/${j}"
    mkdir -p "mrbayes/${j}"
    echo "bash $SCRIPTS/utils/mrbayes.sh $i ${out_prefix}.nex" >>mrbayes.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a mrbayes.cmds bash -c "cmd"

# mbsum
mkdir -p mbsum
echo -n "" >mbsum.skip.log
echo -n "" >mbsum.cmds
for i in mrbayes/*/mrbayes.log; do
    og_dir=$(dirname "$i")
    base=$(basename "$og_dir")
    ## skip Average standard deviation of split frequencies > 0.01
    dsf=$(awk '/Average standard deviation of split frequencies:/ {out=$7} END{print out+0}' "$i" 2>/dev/null)
    if awk -v d="$dsf" -v min_dsf=0.01 'BEGIN{if (d >= min_dsf) exit 0; exit 1}'; then
        echo "Skipping ${base} due to high DSF: ${dsf}" >>mbsum.skip.log
        continue
    else
        echo "mbsum ${og_dir}/*.t -n 1000 -o mbsum/${base}.in >> $og_dir/mbsum.log 2>&1" >>mbsum.cmds
    fi
done
export PATH=$PATH:$SCRIPTS/bucky/bin
pueue add -- xargs -t -P 12 -I cmd -a mbsum.cmds bash -c "cmd"

#########################################################################
#                                                                       #
#  DEPRECATED SECTION: TreeShrink and FastTree-based filtering          #
#                                                                       #
#########################################################################

# remove long branches with TreeShrink and FastTree
mkdir -p treeshrink
echo -n >fasttree.cmds
for i in alignments_trimmed/*.trimed.fa; do
    j=$(basename "$i" .trimed.fa)
    mkdir -p treeshrink/"${j}"
    absolute_path=$(readlink -f "$i")
    ln -s "$absolute_path" "treeshrink"/"${j}"/input.fasta
    echo "FastTree -nt -gtr -quiet treeshrink/${j}/input.fasta > treeshrink/${j}/input.tree" >>fasttree.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a fasttree.cmds bash -c "cmd"

pueue add -- run_treeshrink.py -f -i treeshrink/ -t input.tree -a input.fasta

# filter orthologous groups after TreeShrink based on minimum taxon number and minimum sequence length
min_taxon=11
min_seq_length=300
mkdir -p final_alignments
echo -n "" >treeshrink_filtered.list
for i in treeshrink/*/output.fasta; do
    j=$(dirname "$i")
    j=$(basename "$j")
    seqlen=$(seqkit fx2tab -C ATCG "$i" | awk '{print $3}' | sort -n | head -n 1)
    seqnum=$(grep -c ">" "$i")
    if [[ $seqnum -eq $min_taxon && $seqlen -ge $min_seq_length ]]; then
        #cp -l "$i" final_alignments/"${j}.fa"
        echo -n ""
    else
        echo "$i" >>treeshrink_filtered.list
    fi
done

# summarize sequences removed by TreeShrink
mkdir -p treeshrink_summary
removed_sequences="treeshrink_summary/treeshrink_removed_sequences.tsv"
removed_by_species="treeshrink_summary/treeshrink_removed_by_species.tsv"
removed_group_summary="treeshrink_summary/treeshrink_removed_group_summary.tsv"
printf "OG\tspecies\n" >"$removed_sequences"
for output_fasta in treeshrink/*/output.fasta; do
    og_name=$(basename "$(dirname "$output_fasta")")
    input_fasta="$(dirname "$output_fasta")/input.fasta"
    comm -23 \
        <(awk '/^>/{sub(/^>/, ""); sub(/[[:space:]].*$/, ""); print}' "$input_fasta" | sort -u) \
        <(awk '/^>/{sub(/^>/, ""); sub(/[[:space:]].*$/, ""); print}' "$output_fasta" | sort -u) |
        while IFS= read -r species; do
            printf "%s\t%s\n" "$og_name" "$species"
        done >>"$removed_sequences"
done
printf "species\tfiltered_OG_count\n" >"$removed_by_species"
awk -F '\t' 'NR > 1 {count[$2]++} END {for (species in count) print species, count[species]}' \
    OFS='\t' "$removed_sequences" | sort -k2,2nr -k1,1 >>"$removed_by_species"
printf "group\tfiltered_OG_count\n" >"$removed_group_summary"
awk -F '\t' '
    NR > 1 {
        if ($2 == "EA" || $2 == "EM" || $2 == "ZJ")
            outgroup[$1] = 1
        else
            ingroup[$1] = 1
    }
    END {
        for (og in outgroup) outgroup_count++
        for (og in ingroup) ingroup_count++
        print "outgroup_EA_EM_ZJ", outgroup_count + 0
        print "ingroup", ingroup_count + 0
    }
' OFS='\t' "$removed_sequences" >>"$removed_group_summary"

# compare OGs containing frameshifts with OGs filtered by TreeShrink
fs_overlap_detail="treeshrink_summary/treeshrink_fs_overlap.tsv"
fs_overlap_summary="treeshrink_summary/treeshrink_fs_overlap_summary.tsv"
printf "OG\tcategory\n" >"$fs_overlap_detail"
awk -F '\t' '
    NR == FNR {
        og = $0
        sub(/^.*\//, "", og)
        sub(/\.nal$/, "", og)
        frameshift[og] = 1
        next
    }
    FNR > 1 {
        treeshrink[$1] = 1
    }
    END {
        for (og in frameshift) {
            if (og in treeshrink)
                print og, "overlap"
            else
                print og, "fs.list_only"
        }
        for (og in treeshrink) {
            if (!(og in frameshift))
                print og, "TreeShrink_only"
        }
    }
' OFS='\t' fs.list "$removed_sequences" | sort -k2,2 -k1,1 >>"$fs_overlap_detail"

printf "category\tOG_count\n" >"$fs_overlap_summary"
awk -F '\t' '
    NR > 1 {
        count[$2]++
    }
    END {
        overlap = count["overlap"] + 0
        fs_only = count["fs.list_only"] + 0
        treeshrink_only = count["TreeShrink_only"] + 0
        print "fs.list_total", overlap + fs_only
        print "TreeShrink_total", overlap + treeshrink_only
        print "overlap", overlap
        print "fs.list_only", fs_only
        print "TreeShrink_only", treeshrink_only
    }
' OFS='\t' "$fs_overlap_detail" >>"$fs_overlap_summary"

# raxml-ng for each filtered orthologous group
mkdir -p raxml
echo -n >raxml.cmds
for i in final_alignments/*.fa; do
    j=$(basename "$i" .fa)
    out_prefix="raxml/${j}/${j}"
    mkdir -p "raxml/${j}"
    echo "$SCRIPTS/utils/raxml.sh $i $out_prefix 1 ZJ" >>raxml.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a raxml.cmds bash -c "cmd"

# mrbayes for each filtered orthologous group
mkdir -p mrbayes
echo -n >mrbayes.cmds
for i in final_alignments/*.fa; do
    j=$(basename "$i" .fa)
    out_prefix="mrbayes/${j}/${j}"
    mkdir -p "mrbayes/${j}"
    echo "bash $SCRIPTS/utils/mrbayes.sh $i ${out_prefix}.nex" >>mrbayes.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a mrbayes.cmds bash -c "cmd"

## treeshrink with per-gene
mkdir -p treeshrink_per_gene
echo -n >treeshrink_per_gene.cmds
for i in final_alignments/*.fa; do
    j=$(basename "$i" .fa)
    out_prefix="treeshrink_per_gene/${j}/${j}"
    mkdir -p "treeshrink_per_gene/${j}"
    echo "bash $SCRIPTS/utils/treeshrink.sh $i $out_prefix" >>treeshrink_per_gene.cmds
done
pueue add -- xargs -t -P 12 -I cmd -a treeshrink_per_gene.cmds bash -c "cmd"