from collections import OrderedDict
import itertools
import logging
from multiprocessing import Pool
import time
import numpy as np
import pandas as pd
from scipy.special import digamma
from scipy.optimize import linear_sum_assignment
from tqdm.auto import tqdm

import spatial_lda.admm as admm
from spatial_lda.online_lda import LatentDirichletAllocation


def _update_xi(counts, diff_matrix, diff_penalty, sample_id, verbosity=0,
               rho=1e-1, mu=2.0):
    if verbosity >= 1:
        logging.info(f'>>> Infering topic weights for sample {sample_id}')
    weight = 1. / diff_penalty
    cs = digamma(counts) - digamma(np.sum(counts, axis=1, keepdims=True))
    s = weight * np.ones(diff_matrix.shape[0])
    result = admm.admm(cs, diff_matrix, s, rho, verbosity=verbosity, mu=mu)
    return result


def _wrap_update_xi(inputs):
    return _update_xi(**inputs)


def _update_xis(sample_features, difference_matrices, difference_penalty, gamma,
                 n_parallel_processes, verbosity, primal_dual_mu=2, admm_rho=0.1):
    sample_idxs = sample_features.index.map(lambda x: x[0])
    new_xis = np.zeros_like(gamma)
    if n_parallel_processes > 1:
        with Pool(n_parallel_processes) as pool:
            unique_idxs = np.unique(sample_idxs)
            sample_masks = [
                sample_idxs == sample_idx for sample_idx in unique_idxs]
            sample_counts = [gamma[sample_mask, :]
                             for sample_mask in sample_masks]
            sample_diff_matrices = [difference_matrices[sample_idx]
                                    for sample_idx in unique_idxs]
            diff_penalties = [difference_penalty for _ in unique_idxs]
            tasks = OrderedDict((('counts', sample_counts),
                                 ('diff_matrix', sample_diff_matrices),
                                 ('diff_penalty', diff_penalties),
                                 ('sample_id', unique_idxs),
                                 ('verbosity', itertools.repeat(verbosity)),
                                 ('rho', itertools.repeat(admm_rho)),
                                 ('mu', itertools.repeat(primal_dual_mu))))
            # convert into a list of keyword dictionaries
            kw_tasks = [{k: v for k, v in zip(tasks.keys(), values)}
                        for values in list(zip(*tasks.values()))]
            results = list(tqdm(pool.imap(_wrap_update_xi, kw_tasks),
                                total=len(unique_idxs),
                                position=1,
                                desc='Spatial LDA inference'))
            new_xis = np.concatenate(results)
    else:
        for sample_idx in np.unique(sample_idxs):
            sample_mask = sample_idxs == sample_idx
            sample_counts = gamma[sample_mask, :]
            sample_diff_matrix = difference_matrices[sample_idx]
            new_xis[sample_mask] = _update_xi(sample_counts,
                                              sample_diff_matrix,
                                              difference_penalty,
                                              sample_idx,
                                              verbosity=verbosity,
                                              rho=admm_rho,
                                              mu=primal_dual_mu)
    return new_xis


def train(sample_features, difference_matrices, n_topics, difference_penalty=1,
          n_iters=3, n_parallel_processes=1, verbosity=0,
          primal_dual_mu=2, admm_rho=1.0):
    """Train a Spatial-LDA model.
    
    Args:
        sample_features: Dataframe that contains neighborhood features of index cells indexed by (sample ID, cell ID).
                         (See featurization.featurize_samples).
        difference_matrices: Difference matrix corresponding to the spatial regularization structure imposed on the
                             samples. (I.e., which cells should be regularized to have similar priors on topics).
                             (See featurization.make_merged_difference_matrices).
        n_topics: Number of topics to fit.
        difference_penalty: Penalty on topic priors of "adjacent" index cells.
        n_iters: Number of outer-loop iterations (LDA + ADMM) to run.
        n_parallel_processes: Number of parallel processes to use.
        verbosity: Amount of debug / info updates to see.
        primal_dual_mu: mu used in primal-dual updates (see paper for more details).
        admm_rho: rho used in ADMM optimization (see paper for more details).

    Returns:
        A Spatial-LDA model.
    """
    start_time = time.time()
    xis = None
    for i in range(n_iters):
        logging.info(f'>>> Starting iteration {i}')
        m_step_start_time = time.time()
        lda = LatentDirichletAllocation(n_components=n_topics, random_state=0,
                                        n_jobs=n_parallel_processes, max_iter=5,
                                        doc_topic_prior=xis)
        lda.fit(sample_features.values)
        gamma = lda._unnormalized_transform(sample_features.values)
        m_duration = time.time() - m_step_start_time
        logging.info(f'>>> Iteration {i}, M-step took {m_duration} seconds.')
        e_step_start_time = time.time()
        xis = _update_xis(sample_features=sample_features,
                          difference_matrices=difference_matrices,
                          difference_penalty=difference_penalty,
                          gamma=gamma,
                          n_parallel_processes=n_parallel_processes,
                          verbosity=verbosity,
                          primal_dual_mu=primal_dual_mu,
                          admm_rho=admm_rho)
        e_duration = time.time() - e_step_start_time
        logging.info(f'>>> Iteration {i}, E-step took {e_duration} seconds.')

    last_m_step_start = time.time()
    columns = ['Topic-%d' % i for i in range(n_topics)]
    lda.topic_weights = pd.DataFrame(lda.fit_transform(sample_features.values),
                                     index=sample_features.index,
                                     columns=columns)
    logging.info(f'>>> Final M-step took {time.time() - last_m_step_start} seconds.')
    logging.info(f'>>> Training took {time.time() - start_time} seconds.')
    return lda


def _topic_name(i):
    return f'Topic-{i}'


def infer(components, sample_features, difference_matrices, difference_penalty=1,
          n_parallel_processes=1):
    """Run inferrence on a Spatial-LDA model.

    This runs only the ADMM updates to get spatially-regularized topic weights on `sample_features`, allowing
    us to infer topic weights on samples after training a model to get topic parameters.
    
    Args:
        components: The components of the Spatial-LDA model (typically `spatial_lda_model.components_`.
        sample_features: Dataframe that contains neighborhood features of index cells indexed by (sample ID, cell ID).
                         (See featurization.featurize_samples).
        difference_matrices: Difference matrix corresponding to the spatial regularization structure imposed on the
                             samples. (I.e., which cells should be regularized to have similar priors on topics).
                             (See featurization.make_merged_difference_matrices).
        difference_penalty: Penalty on topic priors of "adjacent" index cells.
        n_parallel_processes: Number of parallel processes to use.

    Returns:
        Spatial-LDA model with the same topic parameters as the original model but with new topic-weights corresponding
        to the provided sample_features and difference_matrices.
    """
    start_time = time.time()
    logging.info('>>> Starting inference')
    n_topics = components.shape[0]
    complete_lda = LatentDirichletAllocation(n_components=n_topics,
                                             random_state=0,
                                             n_jobs=n_parallel_processes,
                                             max_iter=2,
                                             doc_topic_prior=None)
    complete_lda.set_components(components)
    gamma = complete_lda._unnormalized_transform(sample_features.values)
    xis = _update_xis(sample_features,
                      difference_matrices,
                      difference_penalty,
                      gamma=gamma,
                      n_parallel_processes=2,
                      verbosity=0)
    complete_lda.doc_topic_prior_ = xis
    columns = [_topic_name(i) for i in range(n_topics)]
    topic_weights = pd.DataFrame(complete_lda.fit_transform(sample_features.values),
                                 index=sample_features.index,
                                 columns=columns)
    complete_lda.topic_weights = topic_weights
    logging.info(f'>>> Inference took {time.time() - start_time} seconds.')
    return complete_lda


def get_component_mapping(stats_1, stats_2):
    similarity = stats_1 @ stats_2.T
    assignment = linear_sum_assignment(-similarity)
    mapping = {k: v for k, v in zip(*assignment)}
    return mapping


def get_consistent_orders(stats_list):
    d = stats_list[0].shape[1]
    n_topics = [stats.shape[0] for stats in stats_list]
    assert all([stats.shape[1] == d for stats in stats_list])
    assert all([n1 <= n2 for n1, n2 in zip(n_topics[:-1], n_topics[1:])])
    orders = [list(range(n_topics[0]))]
    for stats_1, stats_2 in zip(stats_list[:-1], stats_list[1:]):
        n_topics_1 = stats_1.shape[0]
        n_topics_2 = stats_2.shape[0]
        mapping = get_component_mapping(stats_1[orders[-1], :], stats_2)
        mapped = mapping.values()
        unmapped = set(range(n_topics_2)).difference(mapped)
        order = [mapping[k] for k in range(n_topics_1)] + list(unmapped)
        orders.append(order)
    return orders


def apply_order_to_model(model, order):
    model.components_ = model.components_[order, :]
    if not np.isscalar(model.doc_topic_prior):
        model.doc_topic_prior = model.doc_topic_prior[:, order]
    model.doc_topic_prior_ = model.doc_topic_prior_[:, order]
    mapper = {_topic_name(j): _topic_name(i) for i, j in enumerate(order)}
    _df = model.topic_weights
    _df.rename(columns=mapper, inplace=True)
    model.topic_weights = _df.reindex(sorted(_df.columns), axis=1)


def order_topics_consistently(models, use_topic_weights=True):
    if use_topic_weights:
        stats_list = [model.topic_weights.values.T for model in models]
    else:
        stats_list = [model.components_ for model in models]

    orders = get_consistent_orders(stats_list)
    for model, order in zip(models, orders):
        apply_order_to_model(model, order)
