#Python Libraries
import os
import re

#Python Packages
import pandas as pd
from numpy import sum, log10
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.gridspec as gridspec


def load_mapped_data(idx_dir, sample_map=None):
    '''
    Loads in the reads mapped data from a directory containing
    sample specific idxstats file generated using samtools idxstats.

    Returns a pandas Dataframe with indices as sample names based on filename,
    and three columns: ERCC, GENOME, UNMAPPED detailing the number of reads in each category

    :param idx_dir: path to the directory containing .idxstats files
    :type idx_dir: str
    :param sample_map: Sample map for changing names for figure creation
                        and readability
    :type sample_map: dict with key,value as sample_id, sample_name
    :returns: A dataframe with samples as rows, and columns ERCC, GENOME, UNMAPPED
             indicating the number of reads mapped to each.
    :rtype: pandas.DataFrame
    '''
    mapped_stats = {}
    for filename in os.listdir(idx_dir):
        if filename.endswith('.idxstats'):
            sample = filename.split('.')[0]
            filename = os.path.join(idx_dir, filename)
            df = pd.read_table(filename,
                                sep='\t',
                                names=['contig', 'length',
                                      'reads_mapped', 'reads_unmapped'])
            is_ercc = df.contig.str.contains('ERCC')
            ercc_mapped = df[is_ercc]['reads_mapped'].sum()
            other_mapped = df[~is_ercc]['reads_mapped'].sum()
            unmapped = df['reads_unmapped'].sum()
            mapped_stats[sample] = (ercc_mapped, other_mapped, unmapped)
    mapped_df = pd.DataFrame(mapped_stats).T
    mapped_df.columns = ['ERCC', 'GENOME', 'UNMAPPED']
    mapped_df = mapped_df.ix[sample_map]
    if sample_map:
        mapped_df.index = [sample_map[idx] for idx in mapped_df.index]
    return mapped_df

def load_tpms(rsem_dir, sample_map=None):
    '''
    Load in the TPM values generated by RSEM

    :param rsem_dir: Path to the directory with the rsem genes.results files
    :type rsem_dir: str
    :param sample_map: Sample map for changing names for figure creation
                        and readability
    :type sample_map: dict with key,value as sample_id, sample_name
    :returns: dataframe with TPM values; rows are gene ids, columns are samples
    :rtype: pandas.DataFrame
    '''

    def read_sample(filename, sample_name=None):
        '''
        Reads in RSEM generated .genes.results files, extracts
        TPM column and returns TPPM column only dataframe renamed to sample_name

        :param filename: path to the .genes.results file
        :type filename: str
        :param sample_name: name of sample for
        :type sample_name: str
        '''

        if not sample_name:
            sample_name = os.path.basename(filename).split('.')[0]
        df = pd.read_table(filename, index_col=0)
        df[sample_name] = df['TPM']
        return df[[sample_name]]

    sample_dfs = []
    for fn in os.listdir(rsem_dir):
        sample_id = fn.split('.')[0]
        if fn.endswith('.genes.results') and sample_id in sample_map:
            sample_name = sample_map[sample_id]
            sample_df = read_sample(os.path.join(rsem_dir, fn),
                                    sample_name)
            sample_dfs.append(sample_df)

    tpm_df = pd.concat(sample_dfs, axis=1)

    return tpm_df




def filter_df(tpm_df,
              genes_only=True,
              expressed_in_multiple=False,):
    '''
    A helper function to filter the TPM dataframe based on different classifications

    :param tpm_df: A dataframe of TPM values as generated by load_tpms
    :type tpm_df: pandas.DataFrame
    :param genes_only: Filter out all ERCC TPM values
    :type genes_only: boolean
    :param expressed_in_multiple: Filter out all genes only expressed in one nuclei 
    :type expressed_in_multiple: boolean
    '''

    if genes_only:
        tpm_df = tpm_df.ix[[index for index in tpm_df.index
                            if index.find('ENSG') != -1]]

    if expressed_in_multiple:
        tpm_df = tpm_df[sum(tpm_df > 0, axis=1) > 1]

    return tpm_df.sort(axis=1)

def get_low_mid_high_genes(control_tpm, cutoffs=(0.3333333, 0.6666667)):
    '''
    Get low, mid, and highly expressed genes as defined a control sample series

    :param control_tpm: controm TPM series
    :type control_tpm: pandas.Series
    :param cutoffs: a tuple of cut-offs for low and mid expressed.
    :type cutoffs: tuple
    :returns: a list of gene lists [low_expressed, mid_expressed, high_expressed]
    :rtype: list
    '''
    control_expressed = control_tpm[control_tpm > 0]
    first_quantile = control_expressed.quantile(cutoffs[0])
    second_quantile = control_expressed.quantile(cutoffs[1])
    low_cutoff = control_expressed <= first_quantile
    mid_cutoff = ((control_expressed > first_quantile) &
                  (control_expressed < second_quantile))
    high_cutoff = (control_expressed >= second_quantile)
    low_expressed = control_expressed[low_cutoff]
    mid_expressed = control_expressed[ mid_cutoff]
    high_expressed = control_expressed[high_cutoff]
    return low_expressed, mid_expressed, high_expressed

def label_expression(tpm_df, low, mid, high):
    """
    Populates a new column 'Type' in tpm_df by labeling based on membership of gene in
    low, mid, or high.
    
    :param tpm_df: TPM df as derived from load_tpms
    :type tpm_df: pandas.DataFrame
    :param low: list of low expressed genes based on expression in control
    :type low: list
    :param mid: list of mid expressed genes based on expression in control
    :type mid: list
    :param high: list of high expressed genes based on expression in control
    :type high: list
    :returns: TPM df with a "Type" column
    :rtype: pandas.DataFrame
    """
    def label_type(gene_id):
        if gene_id in low:
            return "Low"
        elif gene_id in mid:
            return "Mid"
        elif gene_id in high:
            return "High"
        else:
            return "Novel"

    tpm_df['Type'] = tpm_df.index.map(label_type)
    return tpm_df

def calculate_relative_expression(tpm_df, low, mid, high):
    """
    Calculates the number of low, mid, and high genes expressed in each sample

    :param tpm_df: TPM df as derived from load_tpms
    :type tpm_df: pandas.DataFrame
    :param low: list of low expressed genes based on expression in control
    :type low: list
    :param mid: list of mid expressed genes based on expression in control
    :type mid: list
    :param high: list of high expressed genes based on expression in control
    :type high: list
    :returns: A dataframe with rows as samples, columns ['Low', 'Mid',' High', 'Novel']
    :rtype: pandas.DataFrame

    """
    tpm_df = label_expression(tpm_df, low, mid, high)
    value_counts = {}
    for sample_name in tpm_df.columns:
        if sample_name == 'Type':
            continue
        sample_df = tpm_df[tpm_df[sample_name] > 0]
        value_counts[sample_name] = sample_df['Type'].value_counts()
    return pd.DataFrame(value_counts).T.fillna(0)

def load_gene_body_coverage(coverage_file, sample_map):
    """
    Loads in gene body coverage files generated by RNASeqC geneBody_coverage.py. 
    Normalizes the gene body coverage by taking the relative coverage calculated as:

    normalized_coverage = ( coverage - minimum_coverage ) / (maximum_coverage - minimum_coverage)

    The maximum and minimum coverage are unique to each sample

    :param coverage_file: path to the coverage file
    :type coverage_file: str
    :param sample_map: Sample map for changing names for figure creation
                        and readability
    :type sample_map: dict with key,value as sample_id, sample_name
    :returns: A normalized and nominal coverage df
    :rtype: pandas.DataFrame, pandas.DataFrame
    """
    coverage_df = pd.read_table(coverage_file, index_col=0).T
    sample_name = lambda x: x.split('.')[0][1:]
    coverage_df.columns = coverage_df.columns.map(sample_name)
    coverage_df = coverage_df[[column for column in coverage_df.columns
                                        if column in sample_map]]
    coverage_df.columns = coverage_df.columns.map(lambda x: sample_map[x])
    coverage_range = coverage_df.max()-coverage_df.min()
    normalized_df = (coverage_df - coverage_df.min())/coverage_range
    return normalized_df, coverage_df

def subset_gtf_by_genes(new_gtf, original_gtf, genes):
    '''
    Reads in a GTF file and writes a new gtf file for a passed list of genes

    :param new_gtf: path to the new gtf file
    :type new_gtf: str
    :param original_gtf: path to the original gtf
    :type original_gtf: str
    :param genes: list of genes to subset
    :type genes: list
    '''

    def extract_field(fieldname, information):
        '''
        Helper function to extract field from a gtf file
        '''
        try:
            return re.search(fieldname+' "([\w.-]+)"', information).group(1)
        except:
            return "None"

    with open(original_gtf) as IN:
        with open(new_gtf, 'w') as OUT:
            for line in IN:
                if line.startswith('#'):
                    OUT.write(line)
                else:
                    gene_id = extract_field("gene_id", line)
                    if gene_id in genes:
                        OUT.write(line)

def load_bedtools_coverage(filename, min_reads=1, min_length=100):
    """
    Loads in bedtool coverage files generated using coverageBed. Filters
    the features based on minimum number of reads mapped and length

    :param filename: path to bedtools coverage file
    :type filename: str
    :param min_reads: for filtering, minimum number of reads mapping to feature
    :type min_reads: int 
    :param min_length: for filtering, minimum length of feature
    :type min_length: int
    :returns: Coverage data frame after filtering
    :rtype: pandas.DataFrame

    """
    columns = ['Chromosome', 'Start', 'End', 'GTF-INFO', 'Reads', 'Basesgt0',
                'Length', 'Fractiongt0']
    coverage_df = pd.read_table(filename, names=columns)
    coverage_df = coverage_df[(coverage_df.Reads >= min_reads) &
                              (coverage_df.Length >= min_length)]
    return coverage_df

def plot_bedtools_coverage(intronic_df, exonic_df, figsize=(8,8)):
    '''
    Creates a plot showing the relative intronic and exonic coverage as 
    calculated by bedtools coverage
    
    :param intronic_df: a bedtools coverage df for introns created by load_bedtools_coverage
    :type intronic_df: pandas.DataFrame
    :param exonic_df: a bedtools coverage df for exons created by load_bedtools_coverage
    :type exonic_df: pandas.DataFrame
    :returns: a figure detailing differences between intronic and exonic regions

    '''
    fig = plt.figure(figsize=figsize)
    gs = gridspec.GridSpec(2, 1, height_ratios=[2,1])
    ax = fig.add_subplot(gs[0])
    ax.scatter(exonic_df['Fractiongt0'], 
               exonic_df['Length'].map(log10),
               label='Exons', color='red', alpha=0.4)
    ax.scatter(intronic_df['Fractiongt0'], 
               intronic_df['Length'].map(log10), 
               alpha=1, label='Introns')
    ax.set_ylim([2, 6])
    ax.set_xlim([0, 1])
    ax.set_ylabel('Length of Feature (log10)')
    ax.set_title('Coverage Consistency Across Features')
    ax.legend()
    ax = fig.add_subplot(gs[1])
    intronic_df['Fractiongt0'].hist(ax=ax,)
    exonic_df['Fractiongt0'].hist(color='red', alpha=0.6, ax=ax)
    ax.set_xlabel('Fraction of Feature Covered with >0 Reads')
    ax.set_ylabel('Number of Features')
    fig.tight_layout()
    sns.despine()
    return fig
