import re
import os
import json
import h5py
import anndata
import h5sparse
import scipy.io
import numpy as np
import pandas as pd

from pathlib import Path
from anndata import AnnData
from matplotlib.image import imread
from scipy.sparse import csr_matrix, csc_matrix
from sklearn.preprocessing import normalize as sp_normalize

from ._utilities import _h5_idx_generator, _create_h5_using_adata


def read_10X(sample_dirs,
             sample_names,
             merge=True,
             num_cells=None,
             min_umis=0,
             use_filtered=False,
             reference=None,
             data_type="rna",
             verbose=True):
    """Read 10X alignment data (including V3)

    This function generates a sparse matrix (cells x genes) from the data generated by 10X's
    cellranger count pipeline. It can process V2 and V3 data together, producing either a single
    merged matrix or list of matrices. Also handles multiple data types produced by 10X V3 (Gene
    Expression, Antibody Capture, CRISPR, CUSTOM).

    Parameters
    ----------
    sample_dirs : list
        List of directories containing either matrix.mtx(.gz) file along with genes.tsv,
        (features.tsv), and barcodes.tsv, or outer level 10X output directory (containing outs directory).
    sample_names : list
        List of names to use for samples (corresponding to sample_dirs)
    merge : bool, optional
        Whether to merge all matrices of the same data type across samples or leave as list
        of matrices (the default is True).
    num_cells : int, optional
        Optional limit on number of cells returned for each sample (only for Gene
        Expression data). Retains the cells with the highest numbers of transcripts
        (the default is None).
    min_umis : int, optional
        Minimum UMI threshold for cells (the default is 0).
    use_filtered : bool, optional
        Whether to use 10X's filtered data (as opposed to raw). Only relevant for
        sample.dirs containing 10X outs directory (the default is False).
    reference : str, optional
        For 10X V<3, specify which reference directory to use if sample_dir is outer
        level 10X directory (only necessary if more than one reference used for sequencing)
        (the default is None).
    data_type : str, optional, 'rna' or 'atac'
        Indicates the protocol of the input data. If not specified, input data will be
        considered scRNA-seq data (the default is 'rna').
    verbose : bool, optional
        Print messages (True by default)

    Returns
    -------
    datalist : list
         List of merged matrices stored as AnnData objects across data types
         (returns sparse matrix if only one data type detected), or nested
         list of matrices organized by sample if merge=F.

    Examples
    --------
    >>> sample_dir1 = "path/to/outer/dir1" # 10X output directory V2 -- contains outs/raw_gene_bc_matrices/<reference>/...
    >>> sample_dir2 = "path/to/outer/dir2" # 10X output directory V3 -- for two data types, Gene Expression and CUSTOM
    >>> dges1 = read_10x([sample_dir1, sample_dir2], ["sample1", "sample2"], min_umis = 50)
    >>> ligerex = create_liger(expr = dges1[["Gene Expression"]], custom = dges1[["CUSTOM"]])
    """

    data_list = []
    data_types = np.array(['Gene Expression'])
    num_samples = len(sample_dirs)

    if num_cells is not None:
        num_cells = np.repeat(num_cells, num_samples)

    for i in range(num_samples):
        # start message
        print('Processing sample ' + sample_names[i])

        # Construct sample path
        sample_dir, suffix = _build_path(sample_dirs[i], use_filtered, reference)

        matrix_file = sample_dir + '/matrix.mtx' + suffix
        barcodes_file = sample_dir + "/barcodes.tsv" + suffix

        # Read in raw data (count matrix)
        raw_data = scipy.io.mmread(matrix_file)

        # filter for UMIs first to increase speed
        umi_pass = np.sum(raw_data, axis=0) > min_umis
        umi_pass = np.asarray(umi_pass).flatten()  # convert to np array
        if umi_pass.shape[0] == 0:
            print('No cells pass UMI cutoff. Please lower it.')
        raw_data = raw_data[:, umi_pass]
        raw_data = csr_matrix(raw_data)  # convert to csr matrix

        # Create column names
        barcodes = pd.read_csv(barcodes_file, sep='\t', header=None)
        barcodes = barcodes.to_numpy().flatten()[umi_pass]

        # remove -1 tag from barcodes
        for i in range(barcodes.size):
            barcodes[i] = re.sub('\-1$', '', barcodes[i])

        col_names = pd.DataFrame(barcodes, columns=['barcodes'])

        # Create row names
        if data_type == 'rna':
            features = pd.read_csv(features_file, sep='\t', header=None).to_numpy()  # convert to np array
            row_names = features[:, 1]

            # equal to make.unique function in R
            count_dict = {}
            for i in range(len(row_names)):
                name = row_names[i]
                if name not in count_dict:
                    count_dict[name] = 0
                if name in row_names:
                    count_dict[name] += 1
                    if count_dict[name] > 1:
                        row_names[i] = row_names[i] + '.' + str(count_dict[name] - 1)

            row_names = pd.DataFrame(row_names, columns=['gene_name'])

        elif data_type == 'atac':
            features = pd.read_csv(features_file, sep='\t', header=None).to_numpy()
            features = np.array(
                [str(feature[0]) + ':' + str(feature[1]) + '-' + str(feature[2]) for feature in features])
            row_names = pd.DataFrame(features, columns=['gene_name'])

        # split based on 10X datatype -- V3 has Gene Expression, Antibody Capture, CRISPR, CUSTOM
        # V2 has only Gene Expression by default and just two columns
        # TODO: check atac feature file
        if features.shape[1] == 1:
            sample_datatypes = np.array(['Chromatin Accessibility'])
            adata = AnnData(csr_matrix(raw_data), obs=row_names, var=col_names)
            adata.uns['sample_name'] = sample_names[i]
            adata.uns['data_type'] = 'Chromatin Accessibility'
            data_list.append(adata)
        elif features.shape[1] < 3:
            sample_datatypes = np.array(['Gene Expression'])
            adata = AnnData(csr_matrix(raw_data), obs=row_names, var=col_names)
            adata.uns['sample_name'] = sample_names[i]
            adata.uns['data_type'] = 'Gene Expression'
            data_list.append(adata)
        else:
            sample_datatypes = features[:, 2]
            sample_datatypes_unique = np.unique(sample_datatypes)
            # keep track of all unique datatypes
            data_types = np.union1d(data_types, sample_datatypes_unique)

            for name in sample_datatypes:
                idx = sample_datatypes == name
                subset_row_names = row_names[idx]
                subset_row_names = pd.DataFrame(subset_row_names, columns=['gene_name'])
                subset_data = raw_data[:, sample_datatypes == name]
                adata = AnnData(csr_matrix(subset_data), obs=subset_row_names, var=col_names)
                adata.uns['sample_name'] = sample_names[i]
                adata.uns['data_type'] = name
                data_list.append(adata)

        # num_cells filter only for gene expression data
    #        if num_cells is not None:
    #            if 'Gene Expression' or 'Chromatin Accessibility' in sample_datatypes and sample_datatypes.shape[0] == 1:
    #                data_label = sample_datatypes.item()
    #                cs = np.sum(samplelist[data_label], axis=0)
    #                limit = samplelist[data_label].shape[1]
    #                if num_cells[i] > limit:
    #                    print('You selected more cells than are in matrix {}. Returning all {} cells.'.format(i, limit))
    #                num_cells[i] = limit
    #                samplelist[data_label] = np.flip(np.sort(samplelist[data_label]))[0:num_cells[i]]

    return_dges = {}
    for datatype in data_types:
        for data in data_list:
            if datatype not in return_dges:
                return_dges[datatype] = []
            else:
                return_dges[datatype].append(data[data])

        return_dges.append(adata)
    if merge:
        print('Merging samples')

        # return_dges = MergeSparseDataAll()
        # if only one type of data present
        if len(return_dges) == 1:
            print('Returning {} data matrix'.format(datatypes))

    else:
        return data_list

    return data_list


def read_10X_h5(sample_dir: str,
                sample_name: str,
                file_name = None,
                use_raw = False,
                backed = False,
                chunk_size = 1000
                ) -> AnnData:
    """
    need to be improved a lot
    :param sample_dir:
    :param sample_name:
    :return:
    """
    ### 0. Parameter checking
    if file_name is not None:
        file_path = Path(sample_dir, file_name)
    else:
        if use_raw:
            file_path = Path(sample_dir, 'raw_feature_bc_matrix.h5')
        else:
            file_path = Path(sample_dir, 'filtered_feature_bc_matrix.h5')

    ### 1. Read h5 file
    with h5py.File(file_path, 'r') as f:
        raw_data = csc_matrix((f['matrix']['data'], f['matrix']['indices'], f['matrix']['indptr']),
                              shape=f['matrix']['shape'])
        features = pd.DataFrame({'gene_name': f['matrix']['features']['name'][()].astype(str)})
        barcodes = pd.DataFrame({'barcodes': f['matrix']['barcodes'][()].astype(str)})

    raw_data = csr_matrix(raw_data.transpose())
    barcodes = barcodes.set_index('barcodes')
    features = features.set_index('gene_name')

    ### 2. Save AnnData or hdf5 if backed
    if not backed:
        adata = AnnData(raw_data, obs=barcodes, var=features,
                        uns={'sample_name': sample_name, 'data_type': 'Gene Expression'})
    else:
        if not os.path.isdir('./results'):
            os.mkdir('./results')

        """
        # create h5 file for each individual sample.
        file_name = './results/' + sample_name + '.hdf5'
        with h5sparse.File(file_name, 'w') as f:
            for left, right in _h5_idx_generator(chunk_size, raw_data.shape[0]):
                if 'raw_data' not in f.keys():
                    f.create_dataset('raw_data', data=raw_data[left:right, :], chunks=(chunk_size,), maxshape=(None,))
                else:
                    f['raw_data'].append(raw_data[left:right, :])
        """
        adata = AnnData(raw_data, obs=barcodes, var=features,
                        uns={'sample_name': sample_name, 'data_type': 'Gene Expression', 'backed_path': file_name})
        _create_h5_using_adata(adata, chunk_size)
        adata.write_h5ad(filename='./results/' + sample_name + '.h5ad')
        adata = anndata.read_h5ad(filename='./results/' + sample_name + '.h5ad', backed=True)

    adata.var_names_make_unique()

    return adata




def read_10X_visium(sample_dirs, sample_names):
    """

    :param sample_dirs:
    :param sample_names:
    :return:
    """
    ### 1. read count matrix
    adata = read_10X_h5(sample_dirs, sample_names)

    ### 2. read spatial info
    p_spatial = Path(sample_dirs, 'spatial')

    if not p_spatial.is_dir():
        print('Spatial file is not existed.')
    else:
        positions = pd.read_csv((p_spatial / 'tissue_positions_list.csv'), header=None)
        positions.columns = [
            'barcodes',
            'in_tissue',
            'array_row',
            'array_col',
            'pxl_row_in_fullres',
            'pxl_col_in_fullres'
        ]
        positions = positions.set_index('barcodes')

        adata.obs = adata.obs.join(positions, how='left')
        adata.obsm['pxl_in_fullres'] = adata.obs[['pxl_row_in_fullres', 'pxl_col_in_fullres']].to_numpy()
        adata.obsm['array_spatial'] = adata.obs[['array_row', 'array_col']].to_numpy()
        adata.obs.drop(
            columns=['array_row', 'array_col', 'pxl_row_in_fullres', 'pxl_col_in_fullres'],
            inplace=True,
        )

        adata.uns['image'] = {
            'lowres': imread(p_spatial / 'tissue_lowres_image.png'),
            'hires': imread(p_spatial / 'tissue_hires_image.png'),
            'scale_factors': json.loads((p_spatial / 'scalefactors_json.json').read_text())
        }

    return adata


def read_10X_multiome(sample_dir, sample_name):
    """

    :param sample_dir:
    :param sample_name:
    :return:
    """
    file_name = sample_dir + '/' + sample_name + '.h5'

    with h5py.File(file_name, 'r') as f:
        raw_data = csc_matrix((f['matrix']['data'], f['matrix']['indices'], f['matrix']['indptr']),
                              shape=f['matrix']['shape'])
        features = pd.DataFrame({'gene_name': f['matrix']['features']['name'][()].astype(str),
                                 'id': f['matrix']['features']['id'][()].astype(str),
                                 'feature_type': f['matrix']['features']['feature_type'][()].astype(str)
                                 })
        barcodes = pd.DataFrame({'barcodes': f['matrix']['barcodes'][()].astype(str)})

    raw_data = csr_matrix(raw_data.transpose())
    barcodes = barcodes.set_index('barcodes')
    features = features.set_index('gene_name')

    rna_idx = features['feature_type'] == 'Gene Expression'
    adata_atac = AnnData(raw_data[:, ~rna_idx], obs=barcodes, var=features[~rna_idx],
                         uns={'sample_name': sample_name, 'data_type': 'Gene Expression'})
    adata_rna = AnnData(raw_data[:, rna_idx], obs=barcodes, var=features[rna_idx],
                        uns={'sample_name': sample_name, 'data_type': 'Gene Expression'})
    adata_atac.var_names_make_unique()
    adata_rna.var_names_make_unique()

    return adata_atac, adata_rna


def read_10X_atac():
    pass

def save(dir, overwrite='None'):
    p = Path(dir)
    if p.exists():
        if overwrite is True:
            pass
    pass


def save_cellxgene(liger_object,
                   save_dir,
                   return_adata=False):
    """ """
    concat_adata = anndata.concat(liger_object.adata_list)
    cellxgene_adata = AnnData(csc_matrix(np.log2(sp_normalize(concat_adata.raw.X, axis=1, norm='l1').toarray() * 10000 + 1)),
                              obs=concat_adata.obs, var=concat_adata.raw.var, uns=concat_adata.uns, obsm=concat_adata.obsm)

    cellxgene_adata.obsm['X_umap'] = cellxgene_adata.obsm['umap_coords']
    cellxgene_adata.obs['leiden'] = cellxgene_adata.obs['cluster'].astype("category")


    if return_adata:
        return cellxgene_adata
    else:
        return None


def write_h5(liger_object,
             save_dir):
    """helper function to save a copy of h5 file"""
    pass


def load(dir, backed='None'):
    pass

def _read_mtx(sample_dirs):
    pass
    #return adata

def _read_10x_v3_h5(sample_dirs):
    pass
    #return adata


def _build_path(sample_dir, use_filtered, reference):
    check_inner = sample_dir + '/outs'

    if os.path.exists(check_inner):
        sample_dir = check_inner
        is_v3 = os.path.exists(sample_dir + '/filtered_feature_bc_matrix')
        matrix_prefix = str(np.where(use_filtered, 'filtered', 'raw'))

        if is_v3:
            sample_dir = sample_dir + '/' + matrix_prefix + '_feature_bc_matrix'
        else:
            if reference is None:
                references = os.listdir(sample_dir + '/raw_gene_bc_matrices')
                if len(references) > 1:
                    raise ValueError('Multiple reference genomes found. Please specify a single one.')
                else:
                    reference = references[0]
        if reference is None:
            reference = ''
        sample_dir = sample_dir + '/' + matrix_prefix + '_gene_bc_matrices/' + reference
    else:
        is_v3 = os.path.exists(sample_dir + '/features.tsv.gz')

    suffix = str(np.where(is_v3, '.gz', ''))
    if data_type == 'rna':
        features_file = str(np.where(is_v3, sample_dir + '/features.tsv.gz', sample_dir + '/genes.tsv'))
    elif data_type == 'atac':
        features_file = str(np.where(is_v3, sample_dir + '/peaks.bed.gz', sample_dir + '/peaks.bed'))

    return sample_dir, suffix

class Data:
    def __init__(self, file_path, file_name):
        pass