#!/bin/bash
# Note: this program needs mafft to align the pep, use trimal to backtranlate the pep to cds, and paml to compute the Ks tree
# For our cluster, we use module to manage the software. So we load the software by using the command line below
module load mafft trimal paml

# Check if both anchorpoints file and fasta file are provided as command-line arguments
if [ $# -lt 4 ]; then
	echo "Usage: $0 <anchorpoints_file> <cds.fasta> <pep.fasta> <thread_num>"
    exit 1
fi

# Anchorpoints file
anchorpoints_file="$1"

# Fasta file
fasta_file="$2"

# Pep file
pep_file="$3"

# Number of threads
num_threads="$4"


start_time=$(date +%s)  # Get the start time of the script

tmp_fifofile="/tmp/$$.fifo"
mkfifo "$tmp_fifofile"   # Create a FIFO file
exec 6<> "$tmp_fifofile"  # Associate FD6 with the FIFO file
rm "$tmp_fifofile"  # Remove the FIFO file (optional)

# Fill FD6 with tokens (newlines) based on the num_threads
for ((i=0; i<${num_threads}; i++)); do
    echo >&6
done

# Output file
output_file="anchorpoints.ks.txt"
echo -e "geneX\tgeneY\tOmega\tKa\tKs" >$output_file

# Function to perform alignment using prank
perform_alignment() {
    gene_x=$1
    gene_y=$2
    fasta_file=$3

    # Create temporary fasta file with sequences for gene_x and gene_y
    tmp_fasta="${gene_x}_${gene_y}.nuc.fa"
    echo ">${gene_x}" > "$tmp_fasta"
    grep -A 1 "$gene_x" "$fasta_file" | tail -n 1 | sed 's/TAA$//i; s/TGA$//i; s/TAG$//i' >> "$tmp_fasta"
    echo ">${gene_y}" >> "$tmp_fasta"
    grep -A 1 "$gene_y" "$fasta_file" | tail -n 1 | sed 's/TAA$//i; s/TGA$//i; s/TAG$//i' >> "$tmp_fasta"

	tmp_pep_fasta="${gene_x}_${gene_y}.pep"
	echo ">${gene_x}" > "$tmp_pep_fasta"
	less "$pep_file" | awk '/^>/ {if (seq) print seq; print; seq=""; next} {seq = seq $0} END {print seq}' | grep -A 1 "$gene_x" | tail -n 1 >> "$tmp_pep_fasta"
	echo ">${gene_y}" >> "$tmp_pep_fasta"
	less "$pep_file" | awk '/^>/ {if (seq) print seq; print; seq=""; next} {seq = seq $0} END {print seq}' | grep -A 1 "$gene_y" | tail -n 1 >> "$tmp_pep_fasta"

    # Perform alignment using mafft and trimal to do the alignment
    aligned_file="${gene_x}_${gene_y}_nuc.aligned.best.fas"
	mafft --quiet "$tmp_pep_fasta" > "$tmp_pep_fasta.mafft.fa"
	trimal -in "$tmp_pep_fasta.mafft.fa" -backtrans "$tmp_fasta" -automated1 -out "$aligned_file" -ignorestopcodon
}

# Function to convert FASTA to PHYLIP format
convert_to_phylip() {
    fasta_file="${gene_x}_${gene_y}_nuc.aligned.best.fas"
    phylip_file="${gene_x}_${gene_y}_nuc.aligned.best.phylip"

	# Use trimal to convert fasta format into phylip format
	trimal -in "$fasta_file" -out "$phylip_file" -phylip_paml 

    # Read the FASTA file and extract the sequence names and sequences
#    sequence_names=()
#    sequences=()
#    current_sequence=""
#
#    while read -r line; do
#        if [[ "$line" =~ ^>(.*)$ ]]; then
#            if [[ -n "$current_sequence" ]]; then
#                sequences+=("$current_sequence")
#                current_sequence=""
#            fi
#            sequence_names+=("${BASH_REMATCH[1]}")
#        else
#            current_sequence+="$line"
#        fi
#    done < "$fasta_file"
#
#    # Add the last sequence
#    if [[ -n "$current_sequence" ]]; then
#        sequences+=("$current_sequence")
#    fi
#
#    # Get the number of sequences and sequence length
#    num_sequences=${#sequence_names[@]}
#    sequence_length=${#sequences[0]}
#
#    # Write the PHYLIP header line
#    echo "$num_sequences $sequence_length" > "$phylip_file"
#
#    # Write the sequence names and sequences in PHYLIP format
#    for ((i = 0; i < num_sequences; i++)); do
#        printf "%-50s%s\n" "${sequence_names[i]}" "${sequences[i]}" >> "$phylip_file"
#    done
}

# Function to compute ks value using PAML
compute_ks() {
    # Create PAML control file
    control_file="${gene_x}_${gene_y}_nuc.codeml.ctrl"
	aligned_file="${gene_x}_${gene_y}_nuc.aligned.best.phylip"
    cat > "$control_file" <<EOF
      seqfile = ${aligned_file}
      outfile = ${gene_x}_${gene_y}_nuc.codeml.out
      noisy = 0
      verbose = 0
      runmode = -2
      seqtype = 1
      CodonFreq = 2
      clock = 0
      aaDist = 0
      model = 0
      NSsites = 0
       icode = 0
      Mgene = 0
      fix_kappa = 0
      kappa = 2
      fix_omega = 0
      omega = 0.4
      fix_alpha = 1
      alpha = 0
      Malpha = 0
      ncatG = 8
      getSE = 0
      RateAncestor = 0
      method = 0
EOF

    # Run PAML to compute ks value
    codeml "$control_file" >/dev/null

    # Extract ks value from PAML output
	omega=$(grep -oP 'dN/dS=\s*\K[0-9.]*' "${gene_x}_${gene_y}_nuc.codeml.out")
	dn=$(grep -oP 'dN=\s*\K[0-9.]*' "${gene_x}_${gene_y}_nuc.codeml.out")
	ds=$(grep -oP 'dN=\s*[0-9.]+\s+dS=\s*\K[0-9.]*' "${gene_x}_${gene_y}_nuc.codeml.out")

	echo -e "$omega\t$dn\t$ds" | grep -v "CODONML"
}

#Read anchorpoints file line by line
while IFS=$'\t' read -r id multiplicon basecluster gene_x gene_y coord_x coord_y is_real_anchorpoint; do
	read -u6
    {
		# Check if gene_x and gene_y are valid anchorpoints
   		if [[ "$is_real_anchorpoint" == "-1" ]]; then
   			# Perform alignment using mafft and trimal
       		aliged_file=$(perform_alignment "$gene_x" "$gene_y" "$fasta_file")

			# Convert fasta into phylip
			phylip_file=$(convert_to_phylip)

   	    	# Compute ks value using PAML
   	    	ks_value=$(compute_ks "$phylip_file")

			omega=$(echo "$ks_value" | awk '{print $1}')
			dn=$(echo "$ks_value" | awk '{print $2}')
			ds=$(echo "$ks_value" | awk '{print $3}')

			# Check if dn and ds are numeric
			if [[ $omega =~ ^[0-9]+(\.[0-9]+)?$ && $dn =~ ^[0-9]+(\.[0-9]+)?$ && $ds =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
				# Append result to output file
				echo -e "$gene_x\t$gene_y\t$omega\t$dn\t$ds" >> "$output_file"
			fi

    	  	# Clean up temporary files
			if [ -f "${gene_x}_${gene_y}.nuc.fa" ]; then
				rm "${gene_x}_${gene_y}.nuc.fa"
			fi
			if [ -f "${gene_x}_${gene_y}.pep" ]; then
				rm "${gene_x}_${gene_y}.pep"
			fi
			if [ -f "${gene_x}_${gene_y}_nuc.aligned.best.fas" ]; then
				rm "${gene_x}_${gene_y}_nuc.aligned.best.fas"
			fi
			if [ -f "${gene_x}_${gene_y}.pep.mafft.fa" ]; then
				rm "${gene_x}_${gene_y}.pep.mafft.fa"
			fi
			if [ -f "${gene_x}_${gene_y}_nuc.aligned.best.phylip" ]; then
				rm "${gene_x}_${gene_y}_nuc.aligned.best.phylip"
			fi
			if [ -f "${gene_x}_${gene_y}_nuc.codeml.ctrl" ]; then
				rm "${gene_x}_${gene_y}_nuc.codeml.ctrl"
			fi
			if [ -f "${gene_x}_${gene_y}_nuc.codeml.out" ]; then
				rm "${gene_x}_${gene_y}_nuc.codeml.out"
			fi
   		fi

        # Return a token to FD6, indicating the completion of the process
        echo >&6
    } &
done < "$anchorpoints_file"

# Wait for all processes to complete
wait

# Get the end time of the script
end_time=$(date +%s)
runtime=$((end_time - start_time))
rm 2ML.dN 2ML.dS 2ML.t 2NG.dN 2NG.dS 2NG.t rst rst1 rub

echo "Script execution time: $runtime seconds"

