#!/usr/bin/env python3

"""
Format bins stat to multiqc tsv format to generate barplots.

:Example:
python template.py -v
"""

# 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'


from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
import logging
import pandas as pd


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

    parser.add_argument('--bins_stats',  nargs='+', required=True, help="Bins stats with Quality category.")
    parser.add_argument("--out_bins_count", default='bin_count_per_quality.tsv' )
    parser.add_argument("--out_bins_size", default='bin_size_per_quality.tsv' )
    parser.add_argument("-v", "--verbose", help="increase output verbosity",
                        action="store_true")

    args = parser.parse_args()
    return args



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")
    
    
    stat_files = args.bins_stats

    df_list_count = []
    df_list_size = []

    logging.info(f'Parsing {len(stat_files)} bins stats files.')

    for stat_file in stat_files:
        sample_name = stat_file.replace('_stat_and_quality.tsv',"")

        df_ech = pd.read_csv(stat_file, sep='\t')

        #Get table with bin count per quality

        # remove Not_binned category when counting bins in different quality
        df_ech_count = df_ech.loc[~(df_ech['genome'] == "Not_binned")]

        df_ech_gr = df_ech_count.groupby(['Quality']).size().reset_index(name='counts')
        df_ech_gr = df_ech_gr.set_index("Quality")
 
        df_tr = df_ech_gr.transpose()

        df_tr['sample'] = sample_name
        df_tr = df_tr.set_index('sample')

        df_list_count.append(df_tr)

        #Get table with bin size per category (with unbinned size)
        df_ech_gr = df_ech.groupby(['Quality']).agg({"Size":sum}).reset_index()

        df_tr = df_ech_gr.set_index("Quality").transpose()

        df_tr['sample'] = sample_name
        df_tr = df_tr.set_index('sample')
        df_list_size.append(df_tr)

    df_count = pd.concat(df_list_count)
    df_size = pd.concat(df_list_size)



    logging.info(f'Writing {args.out_bins_count}')
    df_count.to_csv(args.out_bins_count, sep='\t')
    
    logging.info(f'Writing {args.out_bins_size}')
    df_size.to_csv(args.out_bins_size, sep='\t')

if __name__ == "__main__":
    main()
