#!/usr/bin/env python

"""
----------------------------------------------------------------------------
  Script Name: write_unbinned_contigs.py
  Description: Write contigs that are not part of any final bins into a separate fasta file
  Input files: assembly file in fasta and final bins files in fasta
  Created By:  Vincent Darbot and Jean Mainguy
  Date:        2022-24-10
-------------------------------------------------------------------------------
"""

# Metadata
__author__ = 'Vincent Darbot and Jean Mainguy'
__copyright__ = 'Copyright (C) 2022 INRAE'
__license__ = 'GNU General Public License'
__version__ = '0.1'
__email__ = 'support.bioinfo.genotoul@inra.fr'
__status__ = 'dev'


import argparse
import pyfastx


def retrieve_contigs_in_bins(bins_files):
    binned_contigs = set()
    for bin_fasta in bins_files:
        binned_contigs |= set(pyfastx.Fasta(bin_fasta).keys())
    return binned_contigs


def write_unbinned_contigs(binned_contigs, assembly_file, unbinned_outfile):

    with open(unbinned_outfile, 'w') as out_fh:
        for contig, seq in pyfastx.Fasta(assembly_file, build_index=False):
            if contig not in binned_contigs:
                out_fh.write(f'>{contig}\n{seq}\n')


def parse_arguments():
    # Manage parameters.
    parser = argparse.ArgumentParser()

    parser.add_argument('-b', '--bins', nargs='+', required = True, 
                        help = 'List of path samples bins to check')

    parser.add_argument('-a', '--assembly', required = True, 
                        help = 'Path to assembly')

    parser.add_argument('-o', '--unbinned', default='unbinned_contigs.fasta',
                        help = 'Output file with unbinned contigs file.')

    args = parser.parse_args()
    return args

def main():

    args = parse_arguments()

    binned_contigs = retrieve_contigs_in_bins( args.bins)
    write_unbinned_contigs(binned_contigs, args.assembly, args.unbinned )
    

if __name__ == '__main__':
    main()