#!/usr/bin/env	python3

"""
This script is used to identify the minimum partition required to cluster all samples within the same typing category (e.g. ST, serotype...)

By Veronica Mixao
@INSA
"""

import sys
import argparse
import textwrap
import pandas
import glob

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

# running the pipeline	----------

def main():
    
	# argument options
    
	parser = argparse.ArgumentParser(prog="poli_typing.py", formatter_class=argparse.RawDescriptionHelpFormatter, description=textwrap.dedent("""\
									###############################################################################             
									#                                                                             #
									#                                poli_typing.py                               #
									#                                                                             #
									###############################################################################  
									                            
									This script is used to identify the minimum partition required to cluster all 
									samples within the same typing category (e.g. ST, serotype...). It takes as
									input a directory with the *_metadata_w_partitions.tsv table generated by 
									ReporTree for one or multiple pipelines. This table should include a column with
									the typing category for which the analysis is required.
									
									-----------------------------------------------------------------------------"""))
	
	group0 = parser.add_argument_group("Input/Output", "Input/Output specifications")
	group0.add_argument("-d", "--dir", dest="directory", required=True, type=str, help="[MANDATORY] Directory with the *_metadata_w_partitions.tsv file generated by ReporTree.")
	group0.add_argument("-o", "--output", dest="out", required=True, type=str, help="[MANDATORY] Tag for output file name")
	
	group1 = parser.add_argument_group("Analysis", "Analysis specifications")
	group1.add_argument("-c", "--column", dest="column", required=False, type=str, default="MLST_ST", help="[OPTIONAL] Name of the column that should be used for analysis (default: MLST_ST).")
	
	args = parser.parse_args()

	variable_interest = args.column
	partition_code = ["MST-", "single-"]

	final = {}

	print("\n******************************\n")
	print("Running poli_typing.py")
	print(version, "last updated on", last_updated)

	for filename in glob.glob(args.directory + "/*_metadata_w_partitions.tsv"):
		print("Checking " + filename)
		columns_partitions = []
		mx = pandas.read_table(filename,dtype = str)
		if variable_interest not in mx.columns:
			sys.exit("ERROR!! The column name specified was not found in file " + filename)
		if variable_interest not in final.keys():
			final[variable_interest] = []
			final["size"] = []
			list_obs = mx[variable_interest].astype(str).values.tolist()
			val = set(list_obs)
			if 'nan' in val or 'NAN' in val or 'NaN' in val:
				val.remove('nan')
			for v in val:
				final[variable_interest].append(v)
				final["size"].append(list_obs.count(v))
		for col in mx.columns:
			for code in partition_code:
				if code in col:
					columns_partitions.append(col)					
		
		if "/" in filename:
			pip_code = filename.split("/")[-1]
		else:
			pip_code = filename
		pipeline = pip_code.split("_metadata_w_partitions.tsv")[0]
		if pipeline not in final.keys():
			final[pipeline] = []
		for variable in final[variable_interest]:
			check = False
			flt_mx = mx[(mx[variable_interest].astype(str) == variable)]
			for col in columns_partitions:
				clusters = set(flt_mx[col].astype(str).values.tolist())
				if 'nan' in clusters:
					clusters.remove('nan')
				n_clusters = len(clusters)
				if n_clusters == 1:
					check = True
					if "MST-" in col:
						col = col.split("MST-")[1]
					elif "single-" in col:
						col = col.split("single-")[1]
					if "x" in col:
						col = col.split("x")[0]
					final[pipeline].append(col)
					break
			if not check:
				final[pipeline].append("-")
				
	final_matrix = pandas.DataFrame(data = final)
	final_matrix.to_csv(args.out + "_poli.tsv", index = False, header=True, sep ="\t")

if __name__ == "__main__":
    main()