import glob
import os
import datetime

from collections import OrderedDict

import pandas as pd
import numpy as np

from sklearn.manifold import TSNE

import scipy.stats as stats
import scipy.sparse as sparse

from sparse_dataframe import SparseDataFrame


def combine_sdf_files(run_folder, folders, verbose=False, **kwargs):
    """function for concatenating SparseDataFrames together"""
    combined = SparseDataFrame()
    combined.rows = []
    columns = set()

    for folder in folders:
        filename = os.path.join(run_folder, folder, f'{folder}.mus.cell-gene.npz')
        if verbose:
            print(f'Reading {filename} ...')
        sdf = SparseDataFrame(filename)
        columns.add(tuple(sdf.columns))
        combined.rows.extend(sdf.rows)
        if combined.matrix is None:
            combined.matrix = sdf.matrix
        else:
            combined.matrix = sparse.vstack((combined.matrix, sdf.matrix),
                                            format='csr')

    assert len(columns) == 1

    combined.columns = columns.pop()

    return combined


def combine_csv_files(folder, globber, verbose=False, **kwargs):
    """generic function for concatentating a bunch of csv files into a single
    pandas Dataframe"""
    dfs = []

    for filename in glob.iglob(os.path.join(folder, globber)):
        if verbose:
            print(f'Reading {filename} ...')
        df = pd.read_csv(filename, **kwargs)
        dfs.append(df)
    combined = pd.concat(dfs)
    return combined


def maybe_to_numeric(series):
    try:
        return pd.to_numeric(series)
    except ValueError:
        return series


def clean_mapping_stats(mapping_stats_original, convert_to_percentage=None):
    """Remove whitespace from all values and convert to numbers"""

    if convert_to_percentage is None:
        convert_to_percentage = set()

    mapping_stats_original = mapping_stats_original.applymap(
        lambda x: (x.replace(',', '').strip().strip('%')
                   if isinstance(x, str) else x))

    numeric = mapping_stats_original.apply(maybe_to_numeric)

    numeric.columns = numeric.columns.map(str.strip)

    # for 10X mapping stats
    numeric.columns = numeric.columns.map(
            lambda x: ('Percent {}'.format(x.replace('Fraction ', ''))
                       if x in convert_to_percentage else x)
    )

    return numeric



def diff_exp(matrix, group1, group2, index):
    """Computes differential expression between group 1 and group 2
    for each column in the dataframe counts.

    Returns a dataframe of Z-scores and p-values."""

    g1 = matrix[group1, :]
    g2 = matrix[group2, :]

    g1mu = g1.mean(0)
    g2mu = g2.mean(0)

    mean_diff = np.asarray(g1mu - g2mu).flatten()
    # E[X^2] - (E[X])^2
    pooled_sd = np.sqrt(
        ((g1.power(2)).mean(0) - np.power(g1mu, 2)) / len(group1)
        + ((g2.power(2)).mean(0) - np.power(g2mu, 2)) / len(group2))
    pooled_sd = np.asarray(pooled_sd).flatten()

    z_scores = np.zeros_like(pooled_sd)
    nz = pooled_sd > 0
    z_scores[nz] = np.nan_to_num(mean_diff[nz] / pooled_sd[nz])

    # t-test
    p_vals = (1 - stats.norm.cdf(np.abs(z_scores))) * 2

    df = pd.DataFrame(OrderedDict([('z', z_scores), ('p', p_vals)]),
                      index=index)

    return df


class Plates(object):

    # Names of commonly accessed columns
    MEAN_READS_PER_CELL = 'Mean reads per well'
    MEDIAN_GENES_PER_CELL = 'Median genes per well'
    PERCENT_ERCC = 'Percent ERCC'
    PERCENT_MAPPED_READS = 'Percent mapped to genome'

    # maybe we should change this to the right thing
    SAMPLE_MAPPING = 'WELL_MAPPING'

    def __init__(self, data_folder, metadata, genes_to_drop='Rn45s',
                 verbose=False, nrows=None):

        plates_folder = os.path.join(data_folder, 'plates')

        counts = combine_csv_files(
            plates_folder, '*.htseq-count-by-cell.csv',
            index_col=[0, 1, 2, 3], verbose=verbose, nrows=nrows)
        mapping_stats = combine_csv_files(
            plates_folder, '*.log-by-cell.csv',
            index_col=[0, 1, 2, 3], verbose=verbose)
        self.genes, self.cell_metadata, self.mapping_stats = \
            self.clean_and_reformat(counts, mapping_stats)

        self.plate_summaries = self.calculate_plate_summaries()

        original_metadata = pd.read_csv(metadata, index_col=0)
        self.plate_metadata = self.clean_plate_metadata(original_metadata)
        self.plate_metadata = self.plate_metadata.loc[
            self.plate_summaries.index]

        if not os.path.exists(os.path.join(data_folder, 'coords')):
            os.mkdir(os.path.join(data_folder, 'coords'))

        self.bulk_smushed_cache_file = os.path.join(data_folder, 'coords',
                                                    'bulk_smushed.csv')
        self.cell_smushed_cache_file = os.path.join(data_folder, 'coords',
                                                    'cell_smushed.pickle')

        self.bulk_smushed = self.compute_bulk_smushing()
        self.cell_smushed = self.compute_cell_smushing()

        self.gene_names = sorted(self.genes.columns)
        self.plate_metadata_features = sorted(self.plate_metadata.columns)

        # Remove pesky genes
        self.genes = self.genes.drop(genes_to_drop, axis=1)

        # Get a counts per million rescaling of the genes
        self.counts_per_million = self.genes.divide(self.genes.sum(axis=1),
                                                    axis=0) * 1e6
        self.top_genes = self.compute_top_genes_per_cell()

        self.data = {'genes': self.genes,
                     'mapping_stats': self.mapping_stats,
                     'cell_metadata': self.cell_metadata,
                     'plate_metadata': self.plate_metadata,
                     'plate_summaries': self.plate_summaries}

    def __repr__(self):
        n_plates = self.plate_summaries.shape[0]