#!/usr/bin/env python3

## Imports
from Bio import SeqIO
import pandas as pd
import argparse
import os

## Functions

def parse_metaflye(input_fasta, info_file, outdir):
    """
    metaflye fasta contains no information.
    This is contained in the assembly_info.txt file.
    """
    df = pd.read_csv(info_file, sep='\t')
    # this file may have changed symbols for yes, possibilities include a Y or +
    circular_contigs = list(df[df['circ.'].isin(['Y', '+'])]["#seq_name"])

    for count, rec in enumerate(SeqIO.parse(input_fasta, "fasta")):
        if rec.id in circular_contigs:
            outfile = os.path.join(outdir, f"bin_{count}.fa")
            with open(outfile, 'wt') as outfl:
                outfl.write(rec.format("fasta"))

    # with open(circular_fasta, 'wt') as fh_out:
    #     for rec in SeqIO.parse(input_fasta, "fasta"):
    #         if rec.id in circular_contigs:
    #             fh_out.write(rec.format("fasta"))

def parse_hifiasm(input_fasta, outdir):
    """
    hifiasm-meta circular contig names will have an "c" suffix.
    """
    for count, rec in enumerate(SeqIO.parse(input_fasta, "fasta")):
        if rec.description.endswith('c'):
            outfile = os.path.join(outdir, f"bin_{count}.fa")
            with open(outfile, 'wt') as outfl:
                outfl.write(rec.format("fasta"))

    # with open(circular_fasta, 'wt') as fh_out:
    #     for rec in SeqIO.parse(input_fasta, "fasta"):
    #             if rec.id.endswith('c'):
    #                 fh_out.write(rec.format("fasta"))

def main():
    # Manage parameters
    parser = argparse.ArgumentParser( description='Retrieve circular contigs from long-reads assemblies.')
    # Inputs
    group_input = parser.add_argument_group( 'Inputs' )
    group_input.add_argument('-a', '--assembler', required=True, choices=['metaflye', 'hifiasm-meta'], help='Assembler where the assemblies come from.')
    group_input.add_argument('-f', '--input-fasta', required=True, help='The path of assembly fasta file.')
    
    group_input_otu_table = parser.add_argument_group( ' Metaflye ' )
    group_input_otu_table.add_argument('-i','--info-file', default=None, help="The path of metaflye asembly info file.")
    # output
    group_output = parser.add_argument_group( 'Outputs' )
    group_output.add_argument('-o','--outdir', default='circular_contigs', help="The path of circular contigs directory. [Default: %(default)s]" )

    args = parser.parse_args()

    outdir = args.outdir
    os.makedirs(outdir, exist_ok=True)

    # Check for inputs
    if args.assembler == "metaflye":
        if args.info_file is None:
            parser.error("\n\n#ERROR : --info-file is required with metaflye assembler.")
        else:
            parse_metaflye(args.input_fasta, args.info_file, outdir)

    elif args.assembler == "hifiasm-meta":
        parse_hifiasm(args.input_fasta, outdir)


if __name__ == '__main__':
    main()