#!/bin/bash

set -euo pipefail
export PATH="$PATH:${SCRIPTS}/bucky/bin"

# treeshrink
mkdir -p treeshrink
for i in ../13.gene_phylogeny/alignments_filtered/*.fa; do
    j=$(basename "$i" .fa)
    mkdir -p treeshrink/"${j}"
    absolute_path=$(readlink -f "$i")
    ln -s "$absolute_path" "treeshrink"/"${j}"/input.fasta
    ln -s "$(readlink -f ../13.gene_phylogeny/raxml)"/"${j}/${j}".raxml.bestTree "treeshrink"/"${j}"/input.tree
done
pueue add -- run_treeshrink.py -f -i treeshrink/ -t input.tree -a input.fasta -m per-gene

# treeshrink summary
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"

# LBS filter
min_taxon=11
echo -n "" >treeshrink_filtered.list
for i in treeshrink/*/output.fasta; do
    j=$(dirname "$i")
    j=$(basename "$j")
    seqnum=$(grep -c ">" "$i")
    if [[ $seqnum -ge $min_taxon ]]; then
        echo -n ""
    else
        echo "$j" >>treeshrink_filtered.list
    fi
done

# data preparation
ln -s ../13.gene_phylogeny/raxml .
ln -s ../13.gene_phylogeny/mbsum .

##
## 第一套，所有基因保留，仅移除treeshrink过滤掉的基因
##

mkdir -p all_genes
cd all_genes

# aster coalescence tree
for file in ../raxml/*/*.support; do
    gene=$(basename "$(dirname "$file")")
    if ! grep -Fqx -- "$gene" ../treeshrink_filtered.list; then
        cat "$file"
    fi
done >all.ml.tre
mkdir -p species_tree
pueue add -- astral4 -t 6 -i all.ml.tre -o species_tree/coal.tre --root ZJ

# bucky
mkdir bucky
for mbfile in ../mbsum/*.in; do
    gene=$(basename "$mbfile" .in)
    if ! grep -Fqx -- "$gene" ../treeshrink_filtered.list; then
        cp "$mbfile" bucky/"$gene".in
    fi
done
cd bucky
tar -cf input.mb.tar -- ./*.in
rm -- ./*.in
pueue add -- "$SCRIPTS"/utils/bucky.pl -a 1 -o bucky_1 -s input.mb.tar

# SNaQ
mkdir snaq
cd snaq
## run snaq.jl with interactive mode
# julia snaq.jl
cd ../..

##
## 第二套，移除treeshrink过滤掉的基因和移码突变基因
mkdir -p exclude_all_fs
cd exclude_all_fs
cut -f 1 ../../13.gene_phylogeny/frameshift_summary/frameshift_locus_classification.tsv | sed 1d > fs_exclude.list

# aster coalescence tree
for file in ../raxml/*/*.support; do
    gene=$(basename "$(dirname "$file")")
    if ! grep -Fqx -- "$gene" ../treeshrink_filtered.list fs_exclude.list; then
        cat "$file"
    fi
done >all.ml.tre
mkdir -p species_tree
pueue add -- astral4 -t 6 -i all.ml.tre -o species_tree/coal.tre --root ZJ

# bucky
mkdir bucky
for mbfile in ../mbsum/*.in; do
    gene=$(basename "$mbfile" .in)
    if ! grep -Fqx -- "$gene" ../treeshrink_filtered.list fs_exclude.list; then
        cp "$mbfile" bucky/"$gene".in
    fi
done
cd bucky
tar -cf input.mb.tar -- ./*.in
rm -- ./*.in
pueue add -- "$SCRIPTS"/utils/bucky.pl -a 1 -o bucky_1 -s input.mb.tar

mkdir snaq
cd snaq
## run snaq.jl with interactive mode
# julia snaq.jl
cd ../..

## 第三套，移除严重移码突变基因和treeshrink过滤掉的基因
mkdir -p exclude_severe_fs
cd exclude_severe_fs
grep -wv single_species_single_event ../../13.gene_phylogeny/frameshift_summary/frameshift_locus_classification.tsv | cut -f 1 | sed 1d > fs_exclude.list

# aster coalescence tree
for file in ../raxml/*/*.support; do
    gene=$(basename "$(dirname "$file")")
    if ! grep -Fqx -- "$gene" ../treeshrink_filtered.list fs_exclude.list; then
        cat "$file"
    fi
done >all.ml.tre
mkdir -p species_tree
pueue add -- astral4 -t 6 -i all.ml.tre -o species_tree/coal.tre --root ZJ

# bucky
mkdir bucky
for mbfile in ../mbsum/*.in; do
    gene=$(basename "$mbfile" .in)
    if ! grep -Fqx -- "$gene" ../treeshrink_filtered.list fs_exclude.list; then
        cp "$mbfile" bucky/"$gene".in
    fi
done
cd bucky
tar -cf input.mb.tar -- ./*.in
rm -- ./*.in
pueue add -- "$SCRIPTS"/utils/bucky.pl -a 1 -o bucky_1 -s input.mb.tar

mkdir snaq
cd snaq
## run snaq.jl with interactive mode
# julia snaq.jl
cd ../..