# -*- coding: utf-8 -*-
"""
Created on Thu Dec 17 17:45:31 2020

@author: willi
"""

from . import CodonDictionaries
from . import SequenceManipMethods
import numpy as np
import warnings

class PropensityFactory():
    '''
    factory class for the k's
    '''
    def __init__(self):
        self.codon_dicts = CodonDictionaries.CodonDictionaries()
        pass    

    @staticmethod
    def bin_k(k,inds):
        '''
        Generate a k_binned from a given indexing vector

        Parameters
        ----------
        k : array like
            propensity list or ndarray.
        inds : ndarray
            numpy array of index locations, can be generated from even_bin() or intellegent_bin().

        Returns
        -------
        k_binned : list
            propensity rates for the proposed binning strategy.

        '''
        try: 
            k = k.tolist()
        except:
            pass
        k_binned = np.zeros(len(inds)-1)
        binned_ks = []
        for i in range(0,len(inds)-1):
            binned_ks = binned_ks +   [k[inds[i]:inds[i+1]],]  
        
         
        for i in range(0,len(inds)-1):            
            k_binned[i] = 1/ np.sum( 1/np.array(binned_ks[i]))
        return k_binned


    def get_trna_ids(self,nt_seq):
        '''
        Return tRNA ids 0-60 for a given nt_seq

        Parameters
        ----------
        nt_seq : str
            nucleotide sequence.

        Returns
        -------
        list
            list of trna IDs 0-60 (one for each codon).

        '''
        
        if '*' in SequenceManipMethods.SequenceManipMethods().nt2aa(nt_seq)[-1]:
            nt_seq = nt_seq[:-3]
        
        codons = nt_seq.upper()
        seperated_codons = [codons[i:i+3] for i in range(0, len(codons), 3)] 
        return [self.codon_dicts.trna_dict[x] for x in seperated_codons]
        
        

    def get_k(self, nt_seq, k_init, k_elong_mean,k_term):
        '''
        Returns all propensity functions for a given kinit, kelong_mu, and k termination

        Parameters
        ----------
        nt_seq : string
            nucleotide sequence.
        k_init : float
            initation rate.
        k_elong_mean : float
            average kelongation rate for the transcript.
        k_term : float
            termination rate.

        Returns
        -------
        all_k : list
            list of propensity values per location + initiation and termination at respective ends.

        '''
        codons = nt_seq.upper()
        
        genelength = int(len(codons)/3)
        seperated_codons = [codons[i:i+3] for i in range(0, len(codons), 3)] #split codons by 3
        k_elongation = np.zeros((1, genelength))
        tRNA_copynumber = np.zeros((1, genelength))

     
        for i in range(len(seperated_codons)):
            tRNA_copynumber[0, i] = self.codon_dicts.strGeneCopy[seperated_codons[i]]

        mean_tRNA_copynumber = np.mean(list(self.codon_dicts.strGeneCopy_single.values()))

        k_elongation = (tRNA_copynumber / mean_tRNA_copynumber) * k_elong_mean
        all_k = [k_init] + k_elongation.flatten().tolist() + [k_term]
        
        return all_k
    
    
    def get_k_3_frame(self,nt_seq,k_elong_mean):
        '''
        Parameters
        ----------
        nt_seq : str
            nucleotide sequence.
        k_elong_mean : float
            mean elongation rate aa/time.

        Returns
        -------
        kelongs : list
            elongation rates for all 3 frames of a nucleotide sequence stacked in:
                [0+frame L, 1+frame L-1, 2+frame L-1] .

        '''
        kelongs = []
        
        for n in range(3):
            if n !=0:
                codons = nt_seq[n:-(3-n)]
            else:
                codons = nt_seq
            genelength = int(len(codons)/3)
            seperated_codons = [codons[i:i+3] for i in range(0, len(codons), 3)] #split codons by 3
            k_elongation = np.zeros((1, genelength))
            tRNA_copynumber = np.zeros((1, genelength))

     
            for i in range(len(seperated_codons)):
                tRNA_copynumber[0, i] = self.codon_dicts.strGeneCopy[seperated_codons[i]]
    
            mean_tRNA_copynumber = np.mean(list(self.codon_dicts.strGeneCopy_single.values()))
    
            k_elongation = (tRNA_copynumber / mean_tRNA_copynumber) * k_elong_mean
        
           # k_elongation.flatten().tolist()[:-1]
        
            kelongs = kelongs + k_elongation.flatten().tolist()
        
        return kelongs
        
    
    @staticmethod
    def get_binned_k(k,bins):
        '''
        given a propensity vector and an desired binning indices, return pv and pl binned

        Parameters
        ----------
        k : 1 x L  numpy array
            Propensity vector of rates.
        bins : 1xL ind locations sizes
            binning strategy, the sum of this vector should = L.

        Returns
        -------
        propensity_binned: 1 x Nbins numpy array
            the binned propensities.
        bin_sizes: 1 x Nbins numpy array
            list of bin sizes.
        
        '''
        binsize = int(np.floor(len(k)/bins))
        binned_ks = []
        
        k_binned = np.zeros(bins)
        k_lens = np.ones(bins)*binsize
        
        to_redistribute = len(k)%bins

        k_lens[-to_redistribute:] = binsize+1
        
        inds = np.hstack(([0.], np.cumsum(k_lens))).astype(int)     

        
        for i in range(0,bins):
            binned_ks = binned_ks +   [k[inds[i]:inds[i+1]].tolist(),]  
         
        for i in range(0,bins):            
            k_binned[i] = np.mean(binned_ks[i])/len(binned_ks[i])

        return k_binned,k_lens
    
    
    @staticmethod
    def intellegent_bin(pl,nbins ,min_bin = 1):
        '''
        Function to do intellegent binning, focuses resolution on the areas
        defined in the probe location vector
        
        Note if you pass it a minium bin that when min_bin*nbins > length of your sequence
        this function will fail        

        Parameters
        ----------
        pl : numpy array
            numpy array of Ncolor x Length 0's 1's for probe locations.
        nbins : int
            number of desired bins.
        min_bin : int, optional
            DESCRIPTION. The default is 1.

        Returns
        -------
        pl_inds : numpy array
            n bin locations over the vector length given.

        '''
        
        if min_bin*nbins > pl.shape[1]:
            warnings.warn('Desired minimum binsize and desired number of bins is not possible with the length of the probe vector, returning best guess')

        pl_inds = np.where(pl == 1)[1]
       
        if 0 not in pl_inds:
            pl_inds = np.hstack((np.array([0]), pl_inds))
        if pl.shape[1]+1 not in pl_inds:
            pl_inds = np.hstack(( pl_inds,np.array(pl.shape[1]+1) ))
     
        used_bins = len(pl_inds)
        k = len(pl_inds)-1
        j = 0
        to_add = []
        while used_bins < nbins+1:
            if j == k:
                
                j = 0

                prev_pl_inds = pl_inds
                pl_inds = np.hstack(( pl_inds,np.array(to_add)) )
                pl_inds = np.sort(pl_inds)
             
                if np.array_equal(prev_pl_inds,pl_inds):
                    break
                k =  len(pl_inds)-1
                to_add = []

            newbin = int( pl_inds[j]   +  (pl_inds[j+1]-pl_inds[j]  )/2)

            if newbin not in pl_inds:
                if not (np.abs(pl_inds - newbin) <= min_bin).any():
                    to_add.append(newbin)
                    used_bins+=1
            j+=1    
            
        pl_inds = np.hstack(( pl_inds,np.array(to_add)) )
        pl_inds = np.sort(pl_inds) 
        
        return pl_inds

    @staticmethod
    def even_bin(length,nbins):
        '''
        Parameters
        ----------
        length : int
            Length of the vector to bin.
        nbins : int
            Number of desired bins.

        Returns
        -------
        inds : numpy array
            n bin locations over the vector length given.

        
        '''
        binsize = int(np.floor(length/nbins))
        
        k_lens = np.ones(nbins)*binsize
        
        to_redistribute = length%nbins

        k_lens[-to_redistribute:] = binsize+1
        
        inds = np.hstack(([0.], np.cumsum(k_lens))).astype(int)     
        
        return inds