#!/usr/bin/env python

"""--------------------------------------------------------------------
  Script Name: quantification_by_functional_annotation.py
  Description: Create files with number of reads in each sample
               for certain types of functional annotation
  Input files: tsv file with counts of reads and
               functional annotation by gene.
  Created By:  Joanna Fourquet
  Date:        2021-01-07
-----------------------------------------------------------------------
"""

# Metadata.
__author__ = 'Joanna Fourquet \
- GenPhySE NED'
__copyright__ = 'Copyright (C) 2021 INRAE'
__license__ = 'GNU General Public License'
__version__ = '0.1'
__email__ = 'support.bioinfo.genotoul@inra.fr'
__status__ = 'dev'

# Status: dev.

# Modules importation.

import argparse
import re
import sys
import pandas as pd


# Manage parameters.
parser = argparse.ArgumentParser(description = 'Create a file with \
number of reads in each sample for each functional annotation')

parser.add_argument('-i', '--input_file', required = True, \
help = 'Name of tsv file with counts of reads and functional annotation by gene.')

parser.add_argument('-v', '--version', action = 'version', \
version = __version__)

args = parser.parse_args()

# Function of deduplication of rows when there are more than one functional annotation by column.
# From https://webdevdesigner.com/q/split-explode-pandas-dataframe-string-entry-to-separate-rows-10694/
def splitDataFrameList(df,target_column,separator):
    ''' df = dataframe to split,
    target_column = the column containing the values to split
    separator = the symbol used to perform the split

    returns: a dataframe with each entry for the target column separated, with each element moved into a new row. 
    The values in the other columns are duplicated across the newly divided rows.
    '''
    def splitListToRows(row,row_accumulator,target_column,separator):
        split_row = str(row[target_column]).split(separator)
        for s in split_row:
            new_row = row.to_dict()
            new_row[target_column] = s
            row_accumulator.append(new_row)
    new_rows = []
    df.apply(splitListToRows,axis=1,args = (new_rows,target_column,separator))
    new_df = pd.DataFrame(new_rows)
    return new_df

# Read input file.
data = pd.read_csv(args.input_file, sep="\t")
# Nb of columns.
nb_cols = data.shape[1]
data.head()
# Abundance by GOs terms.
data_GOs = splitDataFrameList(data,'GOs',',')
res_GOs = data_GOs.groupby('GOs').sum().iloc[:,0:nb_cols-24]
# Create a column to sum row abundances.
res_GOs['sum'] = res_GOs.sum(axis=1)
# Create a column with amigo link.
res_GOs['GOs'] = "http://amigo.geneontology.org/amigo/term/" + res_GOs.index
res_GOs['GOs'] = res_GOs['GOs'].apply(lambda x: '/' if (x=="http://amigo.geneontology.org/amigo/term/-") else x)
res_GOs['GOs'] = res_GOs['GOs'].apply(lambda x: '/' if (x=="http://amigo.geneontology.org/amigo/term/nan") else x)
res_GOs.to_csv('GOs_abundance.tsv', sep = '\t')

# Abundance by KEGG ko.
data_KEGG_ko = splitDataFrameList(data,'KEGG_ko',',')
res_KEGG_ko = data_KEGG_ko.groupby('KEGG_ko').sum().iloc[:,0:nb_cols-24]
# Create a column to sum row abundances.
res_KEGG_ko['sum'] = res_KEGG_ko.sum(axis=1)
res_KEGG_ko.to_csv('KEGG_ko_abundance.tsv', sep = '\t')

# Abundance by KEGG_Pathway.
data_KEGG_Pathway = splitDataFrameList(data,'KEGG_Pathway',',')
res_KEGG_Pathway = data_KEGG_Pathway.groupby('KEGG_Pathway').sum().iloc[:,0:nb_cols-24]
# Create a column to sum row abundances.
res_KEGG_Pathway['sum'] = res_KEGG_Pathway.sum(axis=1)
res_KEGG_Pathway.to_csv('KEGG_Pathway_abundance.tsv', sep = '\t')

# Abundance by KEGG_Module.
data_KEGG_Module = splitDataFrameList(data,'KEGG_Module',',')
res_KEGG_Module = data_KEGG_Module.groupby('KEGG_Module').sum().iloc[:,0:nb_cols-24]
# Create a column to sum row abundances.
res_KEGG_Module['sum'] = res_KEGG_Module.sum(axis=1)
res_KEGG_Module.to_csv('KEGG_Module_abundance.tsv', sep = '\t')

# Abundance by PFAM.
data_PFAM = splitDataFrameList(data,'PFAMs',',')
res_PFAM = data_PFAM.groupby('PFAMs').sum().iloc[:,0:nb_cols-24]
# Create a column to sum row abundances.
res_PFAM['sum'] = res_PFAM.sum(axis=1)
res_PFAM.to_csv('PFAM_abundance.tsv', sep = '\t')
