"""
For each ORF, get_phasing determines P-site positions for all mapped reads and computes statistics for on-frame enrichment

usage:
python get_phasing.py --bam <bam-file> --orfs <ribofy orfs-file> --offsets <ribofy offsets-file> --output <output-file>

"""

import argparse
from os import remove
import pysam
import pandas as pd
import numpy as np
import warnings
from tqdm import tqdm

from .bam_utils import get_tid_info
from .stats import *

    
def get_phasing_stats (input, p_methods = ['glm']):
    """
    Calculates phasing statistics from input (either 1D array with p-site counts or 2D matrix with frame-stratified counts (one column per frame))
    input: 1D or 2D object with p-site distribtion

    Parameters
    -------
    input: list
        List of p-site counts across ORF
    p_methods: list, default = ['glm'], possible entries: "glm", "glm_ro", "binom", "wilcox", "taper"
        List of statistical methods to assess phasing

    
    Returns
    -------
    out : dict
        Dictionary with p-values for the specified tests (p-methods)
    """
    
    if len (input.shape) == 1: #1D input
        psites = input
        mat = get_2D_matrix (psites)
    else:  #2D input
        mat = input
        psites = mat.reshape(-1)

    if isinstance (p_methods, str):
        p_methods = [p_methods]

    # remove codons without any signal
    mat = mat[~np.all(mat == 0, axis=1)]

    output_stats = {'n' : mat.shape[0]}

    for m in p_methods:
        output_stats[m] = np.nan
    
    if mat.shape[0] == 0:
        return (output_stats)


    with warnings.catch_warnings():
        warnings.simplefilter("ignore")

        if 'glm' in p_methods:
            output_stats['glm'] = get_glm (mat)

        if 'glm_ro' in p_methods:
            output_stats['glm_ro'] = get_glm (mat, remove_outliers=True)

        if 'wilcox' in p_methods:
            output_stats['wilcox'] = get_wilcox (mat) 

        if 'binom' in p_methods:
            output_stats['binom'] = get_binom (mat)

        if 'taper' in p_methods:
            output_stats['taper'] = get_taper (psites)

    return (output_stats)




def get_psites (orf_tid, start, stop, bamfiles, pd_offsets, tid2ref_dict = None, bam_dict = None):
    """
    Iterates through bam-files and orfs to collect all p-sites

    Parameters
    -------
    orf_tid: str
        transcript_id 
    start: int
        start of ORF in transcript
    stop: int 
        end of ORF in transtript
    bamfiles: list (str)
        list of bamfiles to analyse
    pd_offsets: pandas dataframe
        dataframe with offset information for the bamfiles
    tid2ref_dict: dict, optional
        dictionary to lookup bamfile reference
    bam_dict: dict, optional
        dictionary to hold pysam objects,         

    Return
    -------
    out: list
        array with p-site distribution
        
    """

    if tid2ref_dict == None or bam_dict == None:
        
        tid2ref_dict, bam_dict = {},{}
        for bamfile in bamfiles:

            bam_dict[bamfile] = pysam.Samfile (bamfile)

            _, tid2ref = get_tid_info (bamfile)
            tid2ref_dict[bamfile] = tid2ref


    psites = np.zeros ((stop-start))

    # fetch read from all bams mapping on tid
    for bamfile in bamfiles:

        dtid2ref = tid2ref_dict[bamfile]

        if not orf_tid in dtid2ref:
            continue

        tid = dtid2ref[orf_tid]

        bam_offsets = pd_offsets[pd_offsets.bam == bamfile]
        doffsets = bam_offsets[["read_length", "offset_key"]].set_index('read_length').to_dict ()['offset_key']

        bam = bam_dict[bamfile]
        

        if not orf_tid in dtid2ref:
            print (f"ERROR: {orf_tid} not found in bam... skipping")
            continue
        
        
        for read in bam.fetch (tid, start, stop):

            if read.is_reverse:
                continue
                    
            read_length = read.infer_read_length () 
            
            if not read_length in doffsets:
                continue

            length_offset = int(doffsets[read_length])
            
            offset_pos = read.pos + length_offset
                    
            if offset_pos >= start and offset_pos < stop: 
                try:
                    psites[offset_pos-start] += 1
                except IndexError:
                    print ("error", offset_pos, start, len(psites))


    return (psites)


  
def get_phasing (bamfiles, orfs, offsets, output, percentile=0.9, alpha = 0.01, p_methods = ["glm"], shuffle = False, multiplier = [1]):
    """
    get_phasing: Main function. Iterates through all ORF
    
    Parameters
    ----------
    bamfiles: list
        path/to/bamfiles
    orfs: str
        path/to/orfs (generated by ribofy orfs)
    offsets: str
        path/to/offsets (generated by ribofy offsets)
    output: str
        path/to/output
    percentile: float, optional, default=0.9
        percentile of offset-based ORFs in agreement
    alpha: float, optional, default=0.01
        Cutoff for when an offsets is significant
    p_methods: list, default = ['glm'], possible entries: "glm", "glm_ro", "binom", "wilcox", "taper"
        List of statistical methods to assess phasing
    shuffle: boolean, default = False
        Shuffle p-sites - Used in development
    multiplier: list, default = [1]
        Extend or truncate ORF - Used in development to determine how signal-lengths affect phasing stats       
   
    """ 
    
    print ("### get_phasing ###")

    if isinstance (p_methods, str):
        p_methods = [p_methods]
    
    
    if not libmtspec and 'taper' in p_methods:
        print("module 'mtspec' is required for taper statistics")
        p_methods.remove ("taper")

        if len (p_methods) == 0:
            print ("no methods left... exiting")
            return ()



    bam_dict = {}
    tid2ref_dict = {}
        
    print ("loading bam...")
    
    for bamfile in bamfiles:

        bam_dict[bamfile] = pysam.Samfile (bamfile)

        _, tid2ref = get_tid_info (bamfile)
        tid2ref_dict[bamfile] = tid2ref


    
    print ("loading offsets...")

    pd_offsets = pd.read_csv (offsets, sep="\t")

    # only use read_length showing significant phasing
    
    pd_offsets = pd_offsets[pd_offsets.offset_pct >= percentile]
    
    for s in p_methods:
        pd_offsets = pd_offsets[pd_offsets["p_" + s] <= alpha]
    
    header = {}

    # output file
    fout = open (output, "w")

    # checking orfs - line by line
    print ("checking orfs...")

    num_lines = sum(1 for line in open(orfs,'r'))

    ntests = 0

    with open(orfs, 'r') as f:

        for i, line in enumerate (tqdm(f, total=num_lines)):

            columns = line.strip ("\n").split ("\t")

            if i == 0:
                for icol, col in enumerate (columns):
                    header[col] = icol

                # printing output header

                header_col = columns + ["total_counts", "frame0", "frame1", "frame2"]                                
                
                p_headers = ["p_"+s for s in p_methods]

                for f in multiplier:                    
                    header_col += [s + "_" + str(f) if f != 1 else s for s in p_headers + ['n']]                    

                print ("\t".join (header_col), file=fout)
                continue

            orf_id = columns[header['orf_id']] 
            orf_tid = columns[header['tid']] 
            start, stop = int(columns[header['start']]), int(columns[header['stop']])+3

            psites = get_psites (orf_tid, start, stop, bamfiles, pd_offsets, tid2ref_dict=tid2ref_dict, bam_dict=bam_dict)

            if shuffle:
                np.random.shuffle (psites)


            counts = get_counts (psites)


            output_columns = columns + [str (counts['total']), str(counts['frame0']), str(counts['frame1']), str(counts['frame2'])]

            # the multiplier is only relevent to show signal-length vs statistics
            for f in multiplier:
                    
                if f < 1:
                    new_l = int(len (psites)*f)
                    new_l -= new_l%3
                    sub_psites = psites[0:new_l]
                else:
                    sub_psites = psites * int(f)

                output_stats = get_phasing_stats (sub_psites, p_methods)
                
                output_columns += [str(output_stats[p]) for p in p_methods + ['n']]


            print ("\t".join (output_columns), file=fout)

            # testing
            if shuffle and output_stats['glm'] < 1e-6:
                pd_out = pd.DataFrame (get_2D_matrix (psites), columns=["frame0", "frame1", "frame2"])
                pd_out['orfid'] = orf_id
                pd_out.to_csv (output + ".shuffle_test.txt", sep="\t", mode='a' if ntests > 0 else 'w')
                print ("FOUND NOISE:", output)
                ntests += 1


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



def ribofy_phasing ():
    """    
    """

    parser = argparse.ArgumentParser(description='get phasing')
    parser.add_argument("--bam", dest='bam', required=True, nargs="+", help="Bam file - sorted and indexed")
    parser.add_argument("--orfs",   dest='orfs', required=True, help="orfs - generated by get_ORFs.py")
    parser.add_argument("--offsets", dest='offsets', required=True, help="offsets - generated by get_offsets.py")
    parser.add_argument("--output", dest='output', default = "ribofy_phasing.txt", help="output")

    parser.add_argument("--percentile", dest='percentile', default = 0.9, help="Percentile of consistent offset-determinants")
    parser.add_argument("--alpha", dest='alpha', default = 0.01, help="cutoff p-value for phase-detection at percentile")
    parser.add_argument("--p_methods", dest='p_methods', nargs="*", default = ["binom", "wilcox", "glm"], help="Statistics: possibilities: binom, wilcox, glm, and taper")
    parser.add_argument("--multiplier", dest='multiplier', nargs="*", default = [1], type=float, help="multipler")


    args = parser.parse_args()

    get_phasing (args.bam, args.orfs, args.offsets, args.output,
                 percentile=args.percentile,
                 alpha=args.alpha,
                 p_methods=args.p_methods,
                 multiplier=args.multiplier)



if __name__ == "__main__":
        
    ribofy_phasing ()