"""
Extract best ORF from each orf_group in the output from get_phasing.py and 
computes bio-type and orf-type specific FDRs
By default, all pseudogene biotypes are removed (this is can be avoided by the --keep-pseudo flag)

usage:
python get_results.py --phasing <ribofy_phasing.txt-file> --output <output-file>

"""

import sys
import os
import argparse
import pandas as pd
import numpy as np

from .stats import *


def get_results (phasing, output, keep_pseudo=False, p_methods = ["glm"]):
    """Main function - loads the output from phasing analysis, removes pseudogenes (except if --keep_pseudo is set) and 
    performed multiple testing corrections 

    Parameters
    ----------
    phasing: str
        path/to/phasing file
    output: str
        path/to/output file  
    keep_pseudo: boolean, optional, default = False
        flag to allow pseudogenes in the final output
    p_methods: list, default = ['glm'], possible entries: "glm", "glm_ro", "binom", "wilcox", "taper"
        List of statistical methods to assess phasing
    """

    print ("### get_results ###")

    pd_phasing = pd.read_csv (phasing, sep="\t")

    # remove pseudo genes    
    if not keep_pseudo:
        pd_phasing = pd_phasing[~pd_phasing.bio_type.str.contains("pseudo")]
        
    # group by orf_groups - only keep the most significant orf for each orf_group
    #pd_collapse = pd_phasing.sort_values([pcol], ascending=True) \
    pd_collapse = pd_phasing.sort_values(["total_counts"], ascending=False) \
        .groupby(['orf_group']) \
        .head(1) 

    # calculate group-specific FDR
    for p in p_methods:

        pcol = "p_" + p            
        #pd_collapse = pd_collapse.dropna(subset=[pcol])

        # combine types
        pd_collapse['type'] = pd_collapse[['bio_type', 'orf_type']].agg('_'.join, axis=1)

        # non-filtered padj, i.e. all detected ORFs with p-values are included
        pd_collapse['fdr_' + p] = pd_collapse[pcol].transform (lambda x: p_adjust_bh(x.astype("float"))) 
        pd_collapse['fdr_type_' + p] = pd_collapse.groupby('type')[pcol].transform (lambda x: p_adjust_bh(x.astype("float")))
        pd_collapse['fdr_orftype_' + p] = pd_collapse.groupby('orf_type')[pcol].transform (lambda x: p_adjust_bh(x.astype("float")))
        pd_collapse['fdr_biotype_' + p] = pd_collapse.groupby('bio_type')[pcol].transform (lambda x: p_adjust_bh(x.astype("float")))        
                
        # filtered padj
        pd_collapse = get_filtered_padj (pd_collapse, pcol=pcol, name="filtered_fdr_" + p)    
        pd_collapse = pd_collapse.groupby('type').apply (get_filtered_padj, pcol=pcol, name="filtered_fdr_type_" + p)
        pd_collapse = pd_collapse.groupby('bio_type').apply (get_filtered_padj, pcol=pcol, name="filtered_fdr_biotype_" + p)
        pd_collapse = pd_collapse.groupby('orf_type').apply (get_filtered_padj, pcol=pcol, name="filtered_fdr_orftype_" + p)
        

    print (pd_collapse)
    
    pd_collapse.to_csv (output, sep="\t")

    print ("### Done ###")


def ribofy_results ():

    parser = argparse.ArgumentParser(description='collapse ORFs and calculate FDR')
    parser.add_argument('results', nargs='?', help='') # dummy argument

    # required
    parser.add_argument("--phasing", dest='phasing', required=True, help="phasing - generated by get_phasing.py")
    parser.add_argument("--output", dest='output', default="ribofy.results.txt", help="output; default=ribofy.results.txt")

    # optional    
    parser.add_argument("--p_methods", dest='p_methods', nargs='*', default=["glm"], help="statistics used for enrichment of phased reads; either wilcox or binom")
    parser.add_argument("--keep_pseudo", dest='keep_pseudo', action='store_true', default=False, help="Keep pseudogenes in analysis?")

    args = parser.parse_args()

    get_results (args.phasing, args.output, args.keep_pseudo, args.p_methods)


if __name__ == "__main__":

    ribofy_results ()
