#!/usr/bin/env python

"""----------------------------------------------------------------------------
  Script Name: Filter_contig_per_cpm.py
  Description: Calculates the CPM normalization of mapped reads for each 
               contig and returns contigs which have a CPM > cutoff in .fa.
  Input files: Samtools idxstats output file, .fasta file of contigs.
  Created By:  Jean Mainguy
  Date:        2022-24-10
-------------------------------------------------------------------------------
"""

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

# Status: dev

# Modules importation


from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
import pandas as pd
import numpy as np
import logging
import pyfastx

################################################
# Function
################################################

def parse_arguments():
    """Parse script arguments."""
    parser = ArgumentParser(description="...",
                            formatter_class=ArgumentDefaultsHelpFormatter)

    parser.add_argument("-i", "--samtools_idxstats", nargs='+',  required = True, 
    help = "samtools idxstats file containing contig id, \
    sequence length, number of mapped reads or fragments, \
    number of unmapped reads or fragments")

    parser.add_argument('-f', '--fasta_file',  required = True, 
                        help = 'fasta file containing sequences of contigs.')
    parser.add_argument("-c", "--cutoff_cpm", required = True,
                       help = "Minimum number of reads in a contig")
    parser.add_argument("-s", "--select", 
                       help = "Name of outpout .fa file containing contigs which passed cpm cutoff")
    parser.add_argument("-d", "--discard", 
                       help = "Name of outpout .fa file containing contigs which don't passed cpm cutoff")


    parser.add_argument("-v", "--verbose", help="increase output verbosity",
                        action="store_true")

    args = parser.parse_args()
    return args

def combine_idxstat_files(idxstat_files):
    """
    Combine multiple idxstat files that have the same contigs.

    Sum the #_mapped_read_segments column over multiple idxstat files that have the same reference sequences.
    """
    columns_names = ['reference_sequence_name', 
                    'sequence_length',
                    '#_mapped_read_segments',
                    '#_unmapped_read-segments']

    idxstat_df = pd.read_csv(idxstat_files[0],
                            sep ='\t',
                            names = columns_names, 
                            usecols = ['reference_sequence_name', 
                                        'sequence_length',
                                        '#_mapped_read_segments',],
                            comment="*").set_index('reference_sequence_name')
                            
    for idxstat_file in idxstat_files[1:]:
        other_idxstat_df = pd.read_csv(idxstat_file,
                            sep ='\t',
                            names = columns_names, 
                            usecols = ['reference_sequence_name',
                                        '#_mapped_read_segments',],
                            comment="*").set_index('reference_sequence_name')
            
        idxstat_df['#_mapped_read_segments'] += other_idxstat_df['#_mapped_read_segments']

    return idxstat_df

def main():
    args = parse_arguments()

    if args.verbose:
        logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.DEBUG)
        logging.info('Mode verbose ON')

    else:
        logging.basicConfig(format="%(levelname)s: %(message)s")

    cpm_cutoff = float(args.cutoff_cpm)

    # Read input tables 
    idxstat_df = combine_idxstat_files(args.samtools_idxstats)

    # Calculates cpm for each contig
    sum_reads = idxstat_df['#_mapped_read_segments'].sum()
    logging.info(f'Total number of mapped reads {sum_reads}')

    logging.info(f'With a cpm cutoff of {args.cutoff_cpm}, contigs with less than {(sum_reads*cpm_cutoff)/1e6} reads are removed.')
    idxstat_df['cpm_count'] = 1e6 * idxstat_df['#_mapped_read_segments']/sum_reads

    
    # Contigs with nb reads > cutoff
    kept_contigs = idxstat_df.loc[idxstat_df["cpm_count"] >= cpm_cutoff].index

    logging.info(f'{len(kept_contigs)}/{len(idxstat_df)} contigs are kept with a cpm cutoff of {cpm_cutoff}.')
    # Write new fasta files with kept and unkept contigs
    with open(args.select, "w") as out_select_handle, open(args.discard, "w") as out_discard_handle:
        
        for contig, seq in pyfastx.Fasta(args.fasta_file, build_index=False):
            if contig in kept_contigs:
                out_select_handle.write(f'>{contig}\n{seq}\n')
            else:
                out_discard_handle.write(f'>{contig}\n{seq}\n')

if __name__ == "__main__":
    main()
