#!/bin/bash

# Note: this program needs EMBOSS 6.60 to translate the cds into the protein,
# and diamond to compute the Ks tree. So please ensure these software have been
# installed before running this program. You can quickly check it by using `transeq -h`
# and `diamond help` in your command line.

# For our cluster, we use module to manage the software.
# So we load the software by using the command line below.
module load EMBOSS/x86_64/6.6.0


# Function to run Diamond for two species
run_diamond_2sp() {
	species1=$1
	fasta_file1=$2
	species2=$3
	fasta_file2=$4
	cds_file=$5
	pep_file=$6
	out_file=$7
	num_thread=$8

	seqs1=$(cat $fasta_file1)
	seqs1=$(echo "$seqs1" | sed -E "s/>(.*)/\>${species1}_\1/")
	temp_file1=$(mktemp)
	echo "$seqs1" | awk '{print $1}' > $temp_file1

	seqs2=$(cat $fasta_file2)
	seqs2=$(echo "$seqs2" | sed -E "s/>(.*)/\>${species2}_\1/")
	temp_file2=$(mktemp)
	echo "$seqs2" | awk '{print $1}' > $temp_file2

	cat $temp_file1 $temp_file2 > $cds_file

	#module load EMBOSS/x86_64/6.6.0
	transeq -sequence $cds_file -outseq $pep_file -trim
	sed -i 's/_/|/2' $pep_file
	sed -i 's/_1$//' $pep_file

	diamond blastp -d $pep_file -q $pep_file --outfmt 6 -p $num_thread --out $out_file -e 1E-05 -k 500 --quiet

	rm $temp_file1 $temp_file2
}

# Function to run Diamond for a single species
run_diamond() {
	species=$1
	fasta_file=$2
	cds_file=$3
	pep_file=$4
	out_file=$5
	num_thread=$6

	seqs=$(cat $fasta_file)
	seqs=$(echo "$seqs" | sed -E "s/>(.*)/\>${species}_\1/")
	echo "$seqs" | awk '{print $1}' > $cds_file

	transeq -sequence $cds_file -outseq $pep_file -trim
	sed -i 's/_/|/2' $pep_file
	sed -i 's/_1$//' $pep_file

	diamond blastp -d $pep_file -q $pep_file --outfmt 6 -p $num_thread --out $out_file -e 1E-05 -k 500 --quiet
}

# Function to filter the Blast output
run_filter_blast() {
	blastpoutput=$1
	filtered_out=$2

	# Define the awk script for processing the input file
	awk_script='
	BEGIN {
		FS=OFS="\t"
	}

	# Filter out rows where query_id and subject_id are not equal
	$1 != $2 {
		# Split query_id and subject_id using "|"
		split($1, query_id_arr, "|")
		split($2, subject_id_arr, "|")
		species1 = query_id_arr[1]
		gene_id1 = query_id_arr[2]
		species2 = subject_id_arr[1]
		gene_id2 = subject_id_arr[2]

		# Find max bitscores for gene_id1 and gene_id2
		max_bitscore_x1[gene_id1, species2] = ($12 > max_bitscore_x1[gene_id1, species2]) ? $12 : max_bitscore_x1[gene_id1, species2]
		max_bitscore_y2[gene_id2, species1] = ($12 > max_bitscore_y2[gene_id2, species1]) ? $12 : max_bitscore_y2[gene_id2, species1]

		# Store the current line in blast_output array
		if (!(gene_id1 SUBSEP gene_id2 in blast_output) && !(gene_id2 SUBSEP gene_id1 in blast_output)) {
			blast_output[gene_id1, gene_id2] = $0
			species1_gene[gene_id1] = species1
			species2_gene[gene_id2] = species2
		}
	}

	END {
		# Print the filtered and processed rows
		for (gene_pair in blast_output) {
			split(gene_pair, genes, SUBSEP)
			gene1 = (genes[1] < genes[2]) ? genes[1] : genes[2]
			gene2 = (genes[1] < genes[2]) ? genes[2] : genes[1]

			# Calculate c_value and filter based on c_value > 0.5
			split(blast_output[gene_pair], fields, "\t")
			max_bitscore_x1_val = max_bitscore_x1[gene1, species2_gene[gene2]]
			max_bitscore_y2_val = max_bitscore_y2[gene2, species1_gene[gene1]]
			max_bitscore_val = (max_bitscore_x1_val > max_bitscore_y2_val) ? max_bitscore_x1_val : max_bitscore_y2_val
			c_value = (max_bitscore_val != 0) ? (fields[12] / max_bitscore_val) : 0
			if (c_value > 0.5) {
				print gene1, gene2
			}
		}
	}
	'

	awk -e "$awk_script" "$blastpoutput" > "$filtered_out"
}

run_diamond_and_filter() {
	local input_data_file="$1"
	local work_dir="$2"
	local num_core="$3"

	mapfile -t input_data_array < "$input_data_file"

	for (( i = 0; i < ${#input_data_array[@]}; i++ )); do
		line1="${input_data_array[$i]}"
		IFS=$'\t' read -r species fasta_file gff_file <<< "$line1"

		if [[ ! -e "../$gff_file" ]]; then
			echo "Skipping $species: gff_file does not exist"
			continue
		fi

		species=$(echo "$species" | tr ' ' '_')

		outdir="${work_dir}/i-adhore.${species}_vs_${species}"
		out_file="${outdir}/${species}_vs_${species}.blastp.tsv"
		cdsFile="${outdir}/${species}.cds.fasta"
		pepFile="${outdir}/${species}.pep.fasta"
		blastOutFile="${outdir}/${species}_vs_${species}.blastp.tsv"
		out_filter_file="${outdir}/${species}_vs_${species}.blastp.pairs"

		if [[ ! -e "$out_filter_file" ]]; then
			if [[ ! -e "$outdir" ]]; then
				mkdir -p "$outdir"
			fi
			run_diamond "$species" "../$fasta_file" "$cdsFile" "$pepFile" "$blastOutFile" "$num_core"
			run_filter_blast "$blastOutFile" "$out_filter_file"
		fi

		for (( j = i + 1; j < ${#input_data_array[@]}; j++ )); do
			line2="${input_data_array[$j]}"
			IFS=$'\t' read -r other_species other_fasta_file other_gff_file <<< "$line2"

			if [[ ! -e "../$other_gff_file" ]]; then
				echo "Skipping $other_species: gff_file does not exist"
				continue
			fi
			other_species=$(echo "$other_species" | tr ' ' '_')

			outdir="${work_dir}/i-adhore.${species}_vs_$other_species"
			cdsFile="$outdir/${species}_vs_${other_species}.cds.fasta"
			pepFile="$outdir/${species}_vs_${other_species}.pep.fasta"
			blastOutFile="$outdir/${species}_vs_${other_species}.blastp.tsv"
			filteredFile="$outdir/${species}_vs_${other_species}.blastp.pairs"
			# Run Diamond and filter blast results if output files don't exist
			if [[ "$species" != "$other_species" && ! -e "$filteredFile" ]]; then
				run_diamond_2sp "$species" "$fasta_file" "$other_species" "../$other_fasta_file" "$cdsFile" "$pepFile" "$blastOutFile" "$num_core"
				run_filter_blast "$blastOutFile" "$filteredFile"
			fi
			gzip "$blastOutFile"
		done
	done
}

if [ $# -lt 3 ]; then
	echo "Usage: $0 <input_file> <output_dir> <thread_num>"
else
	run_diamond_and_filter $1 $2 $3
fi
