#!/usr/bin/env  python3

"""
This script was designed to compare the sample distances when a cluster zoom-in
is performed with a wgMLST schema. 

By Veronica Mixao
@INSA
"""

import os
import argparse
from doctest import ELLIPSIS_MARKER
import textwrap
import pandas
import glob

version = "1.0.0"
last_updated = "2024-07-26"

def check_dist(new_dist_file):
	""" Get info from the new allele distance matrix
	input: filename
	output: max_AD, samples  """

	mx = pandas.read_table(new_dist_file)
	samples = mx[mx.columns[0]].values.tolist()
	mx = mx.set_index(mx.columns[0], drop = True)
	max_value = mx.max().max()
	cluster_len = len(samples)
	if "SRR3973058" in samples:
		print(max_value)
	
	return max_value, samples, cluster_len

def check_old_dist(original_distances, samples):
	""" Get info from the old allele distance matrix
	input: filename
	output: max_AD """

	mx = pandas.read_table(original_distances)
	mx = mx.loc[mx[mx.columns[0]].isin(samples)]
	mx = mx[samples]
	max_value = mx.max().max()
	
	return max_value

def main():
    
	# argument options	----------
    
	parser = argparse.ArgumentParser(prog="wgmlst_exercise.py", formatter_class=argparse.RawDescriptionHelpFormatter, description=textwrap.dedent("""\
									###############################################################################             
									#                                                                             #
									#                              wgmlst_exercise.py                             #
									#                                                                             #
									############################################################################### 
									                            
									This script was designed to compare  the sample distances when a cluster zoom-in
									is performed with a wgMLST schema. 


									                  
									-------------------------------------------------------------------------------"""))
	
	## parameters
	
	parser.add_argument("-d", "--directory", dest="directory", required=True, type=str, help="PATH to directory with the different cluster zoom-in directories (generated by ReporTree)")
	parser.add_argument("-dist", "--original-distances", dest="original_distances", required=True, type=str, help="Distance matrix of the original analysis in which the clusters were determined")
	parser.add_argument("-o", "--output", dest="output", required=True, type=str, help="Tag for output filename")
	
	args = parser.parse_args()
	
	print("\n******************************\n")
	print("Running wgmlst_exercise.py")
	print(version, "last updated on", last_updated)

	## pipeline
	with open(args.output + "_dist_wgMLST.tsv", "w+") as out:
		print("cluster\tcluster_len\tdist_all\tdist_zoom\tdifference\tsamples", file = out)
		for dir in glob.glob(args.directory + "/*"):
			if "cluster_" in dir:
				new_dist_file = dir + "/cluster_dist_hamming.tsv"
				if os.path.exists(new_dist_file):
					max_value, samples, cluster_len = check_dist(new_dist_file)
					old_max = check_old_dist(args.original_distances, samples)
					dif = int(max_value) - int(old_max)
					if "/" in dir:
						cluster = dir.split("/")[-1]
					else:
						cluster = dir
					print("\t", cluster)
					print(str(cluster) + "\t" + str(cluster_len) + "\t" + str(old_max) + "\t" + str(max_value) + "\t" + str(dif) + "\t" + ",".join(samples), file = out)

if __name__ == "__main__":
    main()