#!/usr/bin/env python

import numpy as np
import pandas as pd
import os
import sys
import glob
import networkx as nx
import argparse

def main():

    options=argparse.ArgumentParser(sys.argv[0],
                usage=argparse.SUPPRESS,
                description='virallcClusters: A tool for defining viral sequence clusters based on percent nucleotide identity of aligned samples',
                prefix_chars='-',
                add_help=True,
                epilog='Written by Chrispin Chaguza, St Jude Children\'s Research Hospital, 2025')

    options.add_argument('--input','-i',action='store',required=True,nargs=1,
                        metavar='input',dest='input',
                        help='Input file generated by alnPairDist')
    options.add_argument('--name','-n',action='store',required=True,nargs=1,
                        metavar='seqlabel',dest='seqlabel',
                        help='Viral sequence label')
    options.add_argument('--prefix','-p',action='store',required=True,nargs=1,
                        metavar='prefix',dest='prefix',
                        help='Cluster label prefix')
    options.add_argument('--out','-o',action='store',required=False,nargs=1,
                        metavar='outfile',dest='outfile',default="clusters.out.tsv",
                        help='Output file containing inferred sequence clusters')
    options.add_argument('--sep','-s',action='store',required=False,nargs=1,
                        metavar='separator',dest='separator',default="",
                        help='Lineage label separator')
    options.add_argument('--thresholds','-t',action='store',required=False,nargs="*",
                        metavar='thresholds',dest='thresholds',default=1,
                        help='Number of threads (default=1)')
    options.add_argument('--verbose','-v',action='store_true',default=False,
                        dest='verbose',help='Show progress')

    options=options.parse_args(args=None if sys.argv[2:] else ['--help'])


    cmdValues = {'input': options.input[0:][0],
                 'seqlabel': options.seqlabel[0:][0],
                 'prefix': options.prefix[0:][0],
                 'output': options.outfile[0:][0] if isinstance(options.outfile,list) else options.outfile,
                 'separator': options.separator[0:][0] if isinstance(options.separator,list) else options.separator,
                 'thresholds': [float(i) for i in options.thresholds[0:]] if isinstance(options.thresholds,list) else options.thresholds,
                 'verbose': options.verbose}

    snpDistDF = pd.read_csv(cmdValues["input"],delimiter="\t",header=0)

    Graph = nx.Graph()
    Graph.add_nodes_from(snpDistDF.seq2.to_list())
    Graph.add_nodes_from(snpDistDF.seq1.to_list())

    snpDistDF1 = snpDistDF[snpDistDF.pid >= cmdValues["thresholds"][0]]
    seqPair = [(snpDistDF1.seq1[i],snpDistDF1.seq2[i]) for i in snpDistDF1.index]

    Graph.add_edges_from(seqPair)
    networkClusters = tuple(nx.connected_components(Graph))

    fhandle = open(f"{cmdValues['seqlabel']}.clusters.tsv","w")

    for i,j in enumerate(networkClusters):
        for s in j:
            FirstLevelLin = str(i+1)
            seqLabel = s
            clusterPrefix = cmdValues["prefix"]
            separator = cmdValues["separator"]

            fhandle.write(f"{seqLabel}\t{clusterPrefix}{separator}{FirstLevelLin}\n")

    fhandle.close();

if __name__=="__main__":
    main()



