#!/usr/bin/env python

"""----------------------------------------------------------------------------------------------------------------------------------------------------------
  Script Name: plot_kaiju_stat.py
  Description: Generates density plot distribution of kaiju match length
  Input files:
               Verbose files generated by Kaiju for each sample.
  Created By:  Jean Mainguy
  Date:        2022-06-10
----------------------------------------------------------------------------------------------------------------------------------------------------------
"""

# Metadata
__author__ = 'Jean Mainguy - 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, FileType
import gzip
import os
import pandas as pd
from collections import defaultdict
import plotly.express as px
from scipy.stats import gaussian_kde
import numpy as np 
import logging
import sys
from plotly.express.colors import qualitative


def parse_arguments():
    """Parse script arguments."""
    parser = ArgumentParser(description="Generates density plot distribution of kiaju match length",
                            formatter_class=ArgumentDefaultsHelpFormatter)

    parser.add_argument('kaiju_results', nargs='+', help='Verbose files generated by Kaiju')

    parser.add_argument('-o', '--output', default='match_length_kaiju_distribution.html')

    parser.add_argument("--smoothing_val", type=float, default=.3,
                        help="smoothing parameter for the density plot")

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

    args = parser.parse_args()
    return args

def parse_matchlen_from_kaiju_output(kaiju_output):
    proper_open = gzip.open if kaiju_output.endswith('.gz') else open

    matchlen2count = defaultdict(int)

    with proper_open(kaiju_output, "rt") as fl:
        for l in fl:
            match_len = int(l.split('\t')[3])
            matchlen2count[match_len] += 1

    df = pd.DataFrame(matchlen2count.items(), columns=['match_length', "reads"])
    df['reads %'] = 100 * df['reads']/df['reads'].sum()
    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")


    smoothing_parameter = args.smoothing_val
    kaiju_results = args.kaiju_results
    output = args.output

    colors = qualitative.Bold + qualitative.G10 + qualitative.Pastel + qualitative.Antique + qualitative.Safe + qualitative.Vivid + qualitative.Dark24   


    density_df_samples = []
    max_match_length = 0
    sample2df_matchlen = {}

    logging.info(f'Parsing kaiju results')
    for kaiju_result in kaiju_results:
        
        sample_name = kaiju_result.replace('_kaiju.out.gz', "")
        
        df_matchlen =  parse_matchlen_from_kaiju_output(kaiju_result)
        
        max_match_length = max((max_match_length, df_matchlen['match_length'].max()))
        sample2df_matchlen[sample_name] = df_matchlen

    x_vals = np.linspace(0,max_match_length,max_match_length) # Specifying the limits of our data
    logging.info(f'Maximum match length is {max_match_length}')
    
    logging.info(f'Computing density for each sample')

    for sample_name, df_matchlen in sample2df_matchlen.items():
        density = gaussian_kde(df_matchlen['match_length'], weights = df_matchlen['reads'])
        density.covariance_factor = lambda : smoothing_parameter #Smoothing parameter
        density._compute_covariance()
        
        density_info = {"Match length":x_vals, "Density":density(x_vals), 'sample':sample_name}
        
        df_density_sample = pd.DataFrame(data=density_info)
        density_df_samples.append(df_density_sample)

    df_density = pd.concat(density_df_samples)


    logging.info(f'Plot creation')
    fig = px.line(df_density, x='Match length', y='Density', color='sample', 
                title="Kaiju match length distribution", color_discrete_sequence=colors)
    fig.update_traces(line=dict(width=2), opacity=.7)

    logging.info(f'Writing output html in {output}')
    fig.write_html(output)
   

if __name__ == '__main__':
    main()
