import pandas as pd
import numpy as np
from itertools import combinations, chain
from collections import Counter
import matplotlib.path
import time
from multiprocessing import Pool

from scipy.sparse import csr_matrix, coo_matrix

import scanpy as sc
from tqdm import tqdm
import networkx as nx

from .hotspot import get_all_hotspot_keys, compute_hotspot_boundary


def coexpression_hotspots(adata, jaccard_matrix=None, hotspot_lookup=None, threshold=0.7,
                          min_size=0, min_genes=5, cutoff=None, verbose=False,
                          processes=1, divisions=10):
    if jaccard_matrix is None:
        if verbose:
            print("Computing hotspot similarity matrix.")
        jaccard_matrix, hotspot_lookup = hotspot_jaccard(adata, processes=processes,
                                                         verbose=verbose, divisions=divisions)

    multigene_hotspots = hotspot_groups(jaccard_matrix, hotspot_lookup,
                                        threshold=threshold, min_genes=min_genes)

    median_sizes = []
    min_size = 5
    #cutoff = 0.5
    for idx, multigene_hotspot in enumerate(multigene_hotspots):
        sizes = []
        for key, v in multigene_hotspot:
            hotspot_size = np.count_nonzero(adata.obs[key] == v)
            sizes.append(hotspot_size)
        median_sizes.append(int(np.median(sizes)))
    multi_arrays = {}
    new_multigene_hotspots = []
    score_dict = {}
    for idx, (multigene_hotspot, median_size) in enumerate(
            sorted(zip(multigene_hotspots, median_sizes), key=lambda x: -x[1])):
        if median_size < min_size:
            break
        hotspots_per_spot = np.zeros(adata.shape[0])

        for idx2, (key, v) in enumerate(multigene_hotspot):
            hotspots_per_spot += (adata.obs[key] == v).astype(np.int_)

        if cutoff is None:
            inds = np.argsort(hotspots_per_spot)[::-1][:median_size]
        else:
            inds = np.where(hotspots_per_spot > cutoff * len(multigene_hotspot))[0]

        score = np.sum(hotspots_per_spot[inds]) / np.sum(hotspots_per_spot)
        score_dict[idx] = score
        # compute essence of overlap

        v = np.full(adata.shape[0], np.nan)
        v[inds] = 0
        multi_arrays[idx] = pd.Categorical(v)
        new_multigene_hotspots.append(multigene_hotspot)

    multi_arrays = {f"hotspots_multi_{idx}": v for idx, (k, v) in enumerate(multi_arrays.items())}
    # adata.obs = pd.concat([adata.obs, pd.DataFrame(multi_arrays, index=adata.obs.index)], axis=1)
    tmp_df = adata.obs.merge(pd.DataFrame(multi_arrays, index=adata.obs.index), how="outer",
                             left_index=True,
                             right_index=True, suffixes=('_old', None))
    tmp_df.drop(tmp_df.filter(regex='_old$').columns.tolist(), axis=1, inplace=True)
    adata.obs = tmp_df
    new_multigene_hotspots = {idx: v for idx, v in enumerate(new_multigene_hotspots)}
    adata.uns["multi_hotspots"] = new_multigene_hotspots

    return len(multi_arrays)


def multi_hotspot_groupings(adata, num_groupings=None):
    multi_hotspots = [v for v in adata.obs if "hotspots_multi" in v]
    multi_hotspots_df = sc.get.obs_df(adata, keys=multi_hotspots)
    # compute to matrix, 0 -> 1, nan -> 0
    multi_hotspots_array = np.nan_to_num(np.array(multi_hotspots_df, dtype=np.float32) + 1).astype(
        np.int_)
    type_counter = Counter()
    for idx in range(adata.shape[0]):
        t = np.where(multi_hotspots_array[idx, :])
        type_counter[t] += 1

    return sorted(type_counter.items(), key=lambda x: -x[1])


def hotspot_decomposition(adata, compute_type_counts=False, neg_weight=1.0):
    # assumes that single-gene hotspots and multi-gene hotspots have already been computed
    all_hotspots = [v for v in adata.obs if "hotspots_" in v and "multi" not in v]
    multi_hotspots = [v for v in adata.obs if "hotspots_multi" in v]

    multi_hotspots_df = sc.get.obs_df(adata, keys=multi_hotspots)
    # compute to matrix, 0 -> 1, nan -> 0
    multi_hotspots_array = np.nan_to_num(np.array(multi_hotspots_df, dtype=np.float32) + 1).astype(
        np.int_)

    weights_dict = hotspot_decomposition_keys(adata, all_hotspots, multi_hotspots_array,
                                              neg_weight=neg_weight)

    type_counter = Counter()
    if compute_type_counts:
        types = [tuple(np.where(weights)[0]) for (weights, _) in weights_dict.values() if
                 np.sum(weights) > 0]
        for t in types:
            for subset in powerset(t):
                type_counter[subset] += 1
        return weights_dict, sorted(type_counter.items(), key=lambda x: -x[1])

    return weights_dict


def hotspot_decomposition_keys(adata, keys, multi_hotspots_array, neg_weight=None):
    weights_dict = {}
    for hotspot in tqdm(keys):
        hotspot_vector = np.array(adata.obs[hotspot], dtype=np.float32)
        hotspot_vector[pd.notnull(hotspot_vector)] = 1
        hotspot_vector = np.nan_to_num(np.array(hotspot_vector, dtype=np.float32)).astype(np.int_)

        weights, score = hotspot_decomposition_sub(hotspot_vector, multi_hotspots_array,
                                                   neg_weight=neg_weight)
        hotspot_name = hotspot[9:]
        weights_dict[hotspot_name] = (weights, score)

    return weights_dict


def hotspot_decomposition_sub(hotspot_vector, multi_hotspots_array, neg_weight=None):
    if neg_weight is None:
        neg_weight = 1.0
    original_size = np.sum(hotspot_vector.flatten()).astype(np.float64)
    weights = np.zeros(multi_hotspots_array.shape[1])
    score = 0.0

    while True:
        overlap = np.sum(np.multiply(hotspot_vector.reshape(-1, 1), multi_hotspots_array), axis=0)
        neg_overlap = np.sum(np.multiply(1 - hotspot_vector.reshape(-1, 1), multi_hotspots_array),
                             axis=0)
        score_change = overlap - neg_weight * neg_overlap
        max_ind = np.argmax(score_change)
        if score_change[max_ind] > 0:
            weights[max_ind] = 1
            intersects = np.logical_and(hotspot_vector, multi_hotspots_array[:, max_ind])
            hotspot_vector[intersects] = 0
            score += score_change[max_ind]
        else:
            break

    score /= original_size
    # check if we met the minimum coverage
    # if np.sum(hotspot_vector.flatten())/original_size < min_covered_frac:
    # we didn't so return nothing
    #    return np.zeros(multi_hotspots_array.shape[1]), 0.0
    return weights, score


def hotspot_jaccard(adata, keys=None, processes=4, verbose=False, divisions=5):
    if keys is None:
        keys = get_all_hotspot_keys(adata)

    num_hotspots = 0
    mean_size = 0
    indptr_counter = 0
    offsets = {}

    indices_list = []
    indptr = [0]

    hotspot_sizes = []

    hotspot_lookup = []

    for idx, hotspot_key in enumerate(keys):
        offsets[hotspot_key] = num_hotspots
        v = pd.Categorical(adata.obs[hotspot_key])
        n_hotspots = len(v.categories)
        for c in range(n_hotspots):
            inds = np.where(v.codes == c)[0]
            indices_list.append(inds)
            indptr_counter += len(inds)
            indptr.append(indptr_counter)
            hotspot_sizes.append(len(inds))
            mean_size += len(inds)
            hotspot_lookup.append((hotspot_key, v.categories[c]))

        num_hotspots += n_hotspots

    # print(f"{counter} total hotspots.")

    indices = np.concatenate(indices_list)
    data = np.ones(indices.shape)
    hotspot_matrix = csr_matrix((data, indices, indptr)).tocsc()

    hotspot_sizes = np.array(hotspot_sizes, dtype=np.float32)

    # Iterate over columns, each representing a spot and the hotspots it belongs to
    num_pairs = 0
    combination_iterators = []
    num_pairs_arr = []
    for col in range(hotspot_matrix.shape[1]):
        hotspots = hotspot_matrix[:, col].nonzero()[0]
        l = len(hotspots)
        num_pairs += l * (l - 1) // 2
        num_pairs_arr.append(l * (l - 1) // 2)
        combination_iterators.append(combinations(hotspots, 2))
    # print(f"{num_pairs} total overlap pairs.")

    # TODO: maybe handle the case we can't import multiprocessing (does that happen?)
    if processes is not None:
        p = Pool(processes)
    # Spatially divide the spots
    coords = adata.obsm['spatial']
    xmin, ymin = np.min(coords, axis=0)
    xmax, ymax = np.max(coords, axis=0)
    xmax += 0.00001
    ymax += 0.00001
    xdiff = xmax - xmin
    ydiff = ymax - ymin
    # Partition up all the spots based on a grid
    data_list = []
    for idx in range(divisions):
        for idy in range(divisions):
            cur_xmin = xmin + (idx / divisions) * xdiff
            cur_xmax = xmin + ((idx + 1) / divisions) * xdiff

            cur_ymin = ymin + (idy / divisions) * ydiff
            cur_ymax = ymin + ((idy + 1) / divisions) * ydiff
            in_current_element = np.all([coords[:, 0] >= cur_xmin, coords[:, 0] < cur_xmax,
                                         coords[:, 1] >= cur_ymin, coords[:, 1] < cur_ymax],
                                        axis=0)
            inds = np.where(in_current_element)[0]
            if len(inds) == 0:
                continue
            cur_iters = []
            cur_num_pairs = 0
            for ind in inds:
                cur_iters.append(combination_iterators[ind])
                cur_num_pairs += num_pairs_arr[ind]
            data_list.append((cur_iters, cur_num_pairs, num_hotspots))

    num_tasks = len(data_list)

    overlap_matrix = csr_matrix((num_hotspots, num_hotspots))

    if verbose:
        list_wrapper = tqdm
    else:
        list_wrapper = lambda x, total: x

    if processes is not None:
        for sub_matrix in list_wrapper(p.imap_unordered(_multi_wrapper, data_list),
                                       total=num_tasks):
            overlap_matrix += sub_matrix
            del sub_matrix
    else:
        for sub_matrix in list_wrapper((_multi_wrapper(data) for data in data_list),
                                       total=num_tasks):
            overlap_matrix += sub_matrix
            del sub_matrix

    if processes is not None:
        p.close()

    del hotspot_matrix
    overlap_matrix = overlap_matrix.tocoo()

    # Compute the jaccard matrix
    if verbose:
        print("Computing jaccard matrix")
    row_hotspot_size = hotspot_sizes[overlap_matrix.row]
    col_hotspot_size = hotspot_sizes[overlap_matrix.col]
    overlap_count = overlap_matrix.data
    jaccard_matrix_data = row_hotspot_size + col_hotspot_size
    jaccard_matrix_data -= overlap_count
    jaccard_matrix_data = overlap_count / jaccard_matrix_data
    if verbose:
        print("Converting jaccard matrix")
    jaccard_matrix = coo_matrix(
        (jaccard_matrix_data, (overlap_matrix.row, overlap_matrix.col)),
        overlap_matrix.shape).tocsr()

    return jaccard_matrix, hotspot_lookup


def _multi_wrapper(data):
    iters, num_pairs, num_hotspots = data
    return compute_overlap_matrix(iters, num_pairs, num_hotspots, verbose=False)


def compute_overlap_matrix(iterator_list, num_pairs, num_hotspots, verbose=True):
    t0 = time.time()
    idx = 0
    if verbose:
        print(f"Constructing overlap matrix with {num_pairs} pairs.")
        #print(f"Current memory usage: {psutil.Process().memory_info().rss/(1024*1024)}")
    row = np.empty(num_pairs, dtype=np.int32)
    col = np.empty(num_pairs, dtype=np.int32)
    if verbose:
        iterator_list = tqdm(iterator_list)
    for it in iterator_list:
        for i, j in it:
            row[idx], col[idx] = i, j
            idx += 1
    data = np.ones(num_pairs, dtype=np.float32)
    if verbose:
        print("Converting overlap matrix")
    overlap_matrix = coo_matrix((data, (row, col)), (num_hotspots, num_hotspots),
                                             dtype=np.float32).tocsr()
    del data, row, col
    if verbose:
        t1 = time.time()
        print(f"Completed in {t1 - t0} seconds.")
    return overlap_matrix


def hotspot_groups(jaccard_matrix, hotspot_lookup, threshold, min_genes=5):
    jaccard_matrix_coo = jaccard_matrix.tocoo()
    above_threshold = jaccard_matrix_coo.data > threshold
    row = jaccard_matrix_coo.row[above_threshold]
    col = jaccard_matrix_coo.col[above_threshold]
    data = jaccard_matrix_coo.data[above_threshold]
    jaccard_matrix_coo.data = data
    jaccard_matrix_coo.row = row
    jaccard_matrix_coo.col = col
    jaccard_matrix_reduced = jaccard_matrix_coo.tocsr()
    del jaccard_matrix_coo

    g = nx.from_scipy_sparse_matrix(jaccard_matrix_reduced)
    g2 = nx.k_core(g, k=min_genes)

    multigene_hotspots = []
    for c in nx.connected_components(g2):
        multigene_hotspot = [hotspot_lookup[i] for i in c]
        multigene_hotspots.append(multigene_hotspot)

    return multigene_hotspots


def compute_multi_boundaries(adata, alpha_max, alpha_min):
    multi_hotspots = [v for v in adata.obs if "hotspots_multi" in v]
    adata.uns["multi_boundaries"] = {}
    for idx, hotspot_key in tqdm(enumerate(multi_hotspots), total=len(multi_hotspots)):
        boundary = compute_hotspot_boundary(adata, hotspot_key, 0, alpha_max=alpha_max,
                                            alpha_min=alpha_min)
        # saving of anndata converts dictionary indices in .uns to strings anyway
        adata.uns["multi_boundaries"][str(idx)] = boundary


def multi_closure(adata):
    hotspot_keys = [v for v in adata.obs if "hotspots_multi" in v]
    points = adata.obsm['spatial']
    for idx, hotspot_key in enumerate(hotspot_keys):
        boundary = adata.uns["multi_boundaries"][idx]
        path = matplotlib.path.Path(boundary)
        v = pd.Categorical(path.contains_points(points, radius=-180), categories=(True,))
        v.categories = [0]
        adata.obs[hotspot_key] = v


def powerset(iterable):
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(1, len(s) + 1))
