#!/usr/bin/env python3

"""----------------------------------------------------------------------------
  Script Name: add_info_to_bin_stat.py
  Description: Class bins by quality (High, Medium, Low Quality bins and High Contamination bins) 
               based on contamination and completeness. Parse the total size of the assembly to add
               the size of contigs that have not been binned.  
  Input files: bin stat file and quast report
  Created By:  Jean Mainguy
  Date:        2021-08-02
-------------------------------------------------------------------------------
"""

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


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



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

    parser.add_argument('-s', '--bins_stat', required=True, help="Bins stat file.")
    parser.add_argument('-q', '--quast_report', required=True, help="Quast report.")

    parser.add_argument('-o', '--output_file', type=str,
                        default="bins_stat_and_quality.tsv", 
                        help="Output table with the additional info added.")

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

    args = parser.parse_args()
    return args

def get_bin_cat_quality(df, max_contamination=10, max_high_qual_contamination=5,
                        medium_qual_completeness=50,
                        high_qual_completeness=90):
    conta_filt = df["contamination"] >= max_contamination
    low_filt = (df['completeness'] < medium_qual_completeness ) & (~conta_filt)
    medium_filt = (df['completeness'] >= medium_qual_completeness ) & (~conta_filt)
    high_filt = (df['completeness'] > high_qual_completeness ) & (df["contamination"] < max_high_qual_contamination )

    df.loc[conta_filt, 'Quality'] = 'High-contamination'

    df.loc[low_filt, 'Quality'] = "Low-quality"
    df.loc[medium_filt, 'Quality'] = "Medium-quality"
    df.loc[high_filt, 'Quality'] = "High-quality"
    return 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")

    bins_stat = args.bins_stat
    quast_report = args.quast_report

    # Parse bins stat 
    df = pd.read_csv(bins_stat, sep='\t').set_index('genome', drop=False)

    sum_size_in_bin = df['Size'].sum()

    # Add Bin Quality
    get_bin_cat_quality(df)

    # Get assembly size not in binned
    df_quast = pd.read_csv(quast_report, sep='\t', index_col="Assembly")
    total_assembly_length = int(df_quast.loc["Total length (>= 0 bp)", df_quast.columns[0]])


    df.loc['Not_binned', "Size"] = total_assembly_length - sum_size_in_bin
    df.loc['Not_binned', "Quality"] = 'Not-binned'
    df.loc['Not_binned', "genome"] = 'Not_binned'
    
    # # remove checkm2 specific columns 
    # df = df.drop(['Completeness_Model_Used', 'Additional_Notes'], axis=1)

    # Writting out table
    df.to_csv(args.output_file, sep="\t", index=False)




if __name__ == '__main__':
    main()
