from tqdm import tqdm
import numpy as np
from threadpoolctl import threadpool_limits
from .utils import controlled_compute
from numpy.linalg import LinAlgError
from .logging_utils import logger

__all__ = ['AnnStream']


def fix_knn_query(indices: np.ndarray, distances: np.ndarray, ref_idx: np.ndarray):
    fixed_ind, fixed_dist = indices.copy()[:, 1:], distances.copy()[:, 1:]
    # Identify positions where first index is not a self loop
    mis_idx = indices[:, 0].reshape(1, -1)[0] != ref_idx
    n_mis = mis_idx.sum()
    if n_mis > 0:
        for n, i, j, k in zip(np.where(mis_idx)[0], ref_idx[mis_idx], indices[mis_idx], distances[mis_idx]):
            p = np.where(j == i)[0]
            if len(p) > 0:
                # p is the position of self loop. We exclude this position
                p = p[0]
                j = np.array(list(j[:p]) + list(j[p + 1:]))
                k = np.array(list(k[:p]) + list(k[p + 1:]))
            else:
                # No self found at all. Poor recall? simply remove the last k neighbour
                j = j[:-1]
                k = k[:-1]
            fixed_ind[n] = j
            fixed_dist[n] = k
    return fixed_ind, fixed_dist, n_mis


class AnnStream:
    def __init__(self, data, k: int, n_cluster: int, reduction_method: str,
                 dims: int, loadings: np.ndarray, use_for_pca: np.ndarray,
                 mu: np.ndarray, sigma: np.ndarray,
                 ann_metric: str, ann_efc: int, ann_ef: int, ann_m: int,
                 nthreads: int, ann_parallel: bool, rand_state: int,
                 do_kmeans_fit: bool, disable_scaling: bool, ann_idx):
        self.data = data
        self.k = k
        if self.k >= self.data.shape[0]:
            self.k = self.data.shape[0] - 1
        self.nClusters = max(n_cluster, 2)
        self.dims = dims
        self.loadings = loadings
        if self.dims is None and self.loadings is None:
            raise ValueError("ERROR: Provide either value for atleast one: 'dims' or 'loadings'")
        self.annMetric = ann_metric
        self.annEfc = ann_efc
        self.annEf = ann_ef
        self.annM = ann_m
        self.nthreads = nthreads
        if ann_parallel:
            self.annThreads = self.nthreads
        else:
            self.annThreads = 1
        self.randState = rand_state
        self.batchSize = self._handle_batch_size()
        self.method = reduction_method
        self.nCells, self.nFeats = self.data.shape
        self.clusterLabels: np.ndarray = np.repeat(-1, self.nCells)
        disable_reduction = False
        if self.dims < 1:
            disable_reduction = True
        with threadpool_limits(limits=self.nthreads):
            if self.method == 'pca':
                self.mu, self.sigma = mu, sigma
                if self.loadings is None or len(self.loadings) == 0:
                    if len(use_for_pca) != self.nCells:
                        raise ValueError("ERROR: `use_for_pca` does not have sample length as nCells")
                    if disable_reduction is False:
                        self._fit_pca(disable_scaling, use_for_pca)
                else:
                    # Even though the dims might have been already adjusted according to loadings before calling
                    # AnnStream, it could still be overwritten by _handle_batch_size. Hence need to hard set it here.
                    self.dims = self.loadings.shape[1]
                    # it is okay for dimensions to be larger than batch size here because we will not fit the PCA
                if disable_scaling:
                    if disable_reduction:
                        self.reducer = lambda x: x
                    else:
                        self.reducer = lambda x: x.dot(self.loadings)
                else:
                    if disable_reduction:
                        self.reducer = lambda x: self.transform_z(x)
                    else:
                        self.reducer = lambda x: self.transform_z(x).dot(self.loadings)
            elif self.method == 'lsi':
                if self.loadings is None or len(self.loadings) == 0:
                    if disable_reduction is False:
                        self._fit_lsi()
                else:
                    self.dims = self.loadings.shape[1]
                if disable_reduction:
                    self.reducer = lambda x: x
                else:
                    self.reducer = lambda x: x.dot(self.loadings)
            elif self.method == 'custom':
                if self.loadings is None or len(self.loadings) == 0:
                    logger.warning("No loadings provided for manual dimension reduction")
                else:
                    self.dims = self.loadings.shape[1]
                if disable_reduction:
                    self.reducer = lambda x: x
                else:
                    self.reducer = lambda x: x.dot(self.loadings)
            else:
                raise ValueError(f"ERROR: Unknown reduction method: {self.method}")
            if ann_idx is None:
                self.annIdx = self._fit_ann()
            else:
                self.annIdx = ann_idx
                self.annIdx.set_ef(self.annEf)
                self.annIdx.set_num_threads(1)
            self.kmeans = self._fit_kmeans(do_kmeans_fit)

    def _handle_batch_size(self):
        if self.dims > self.data.shape[0]:
            self.dims = self.data.shape[0]
        batch_size = self.data.chunksize[0]  # Assuming all chunks are same size
        if self.dims >= batch_size:
            self.dims = batch_size - 1  # -1 because we will do PCA +1
            logger.info(f"Number of PCA/LSI components reduced to batch size of {batch_size}")
        if self.nClusters > batch_size:
            self.nClusters = batch_size
            logger.info(f"Cluster number reduced to batch size of {batch_size}")
        return batch_size

    def iter_blocks(self, msg: str = '') -> np.ndarray:
        for i in tqdm(self.data.blocks, desc=msg, total=self.data.numblocks[0]):
            yield controlled_compute(i, self.nthreads)

    def transform_z(self, a: np.ndarray) -> np.ndarray:
        return (a - self.mu) / self.sigma

    def transform_ann(self, a: np.ndarray, k: int = None, self_indices: np.ndarray = None) -> tuple:
        if k is None:
            k = self.k
        # Adding +1 to k because first neighbour will be the query itself
        if self_indices is None:
            i, d = self.annIdx.knn_query(a, k=k)
            return i, d
        else:
            i, d = self.annIdx.knn_query(a, k=k + 1)
            return fix_knn_query(i, d, self_indices)

    def _fit_pca(self, disable_scaling, use_for_pca) -> None:
        from sklearn.decomposition import IncrementalPCA
        # We fit 1 extra PC dim than specified and then ignore the last PC.
        self._pca = IncrementalPCA(n_components=self.dims + 1, batch_size=self.batchSize)
        do_sample_subset = False if use_for_pca.sum() == self.nCells else True
        s, e = 0, 0
        # We store the first block of values here. if such a case arises that we are left with less dims+1 cells to fit
        # then those cells can be added to end_reservoir for fitting. if there are no such cells then end reservoir is
        # just by itself after fitting rest of the cells. If may be the case that the first batch itself has less than
        # dims+1 cells. in that we keep adding cells to carry_over pile until it is big enough.
        end_reservoir = []
        # carry_over store cells that can yet not be added to end_reservoir ot be used for fitting pca directly.
        carry_over = []
        for i in self.iter_blocks(msg='Fitting PCA'):
            if do_sample_subset:
                e = s + i.shape[0]
                i = i[use_for_pca[s:e]]
                s = e
            if disable_scaling is False:
                i = self.transform_z(i)
            if len(carry_over) > 0:
                i = np.vstack((carry_over, i))
                carry_over = []
            if len(i) < (self.dims + 1):
                carry_over = i
                continue
            if len(end_reservoir) == 0:
                end_reservoir = i
                continue
            try:
                self._pca.partial_fit(i, check_input=False)
            except LinAlgError:
                # Add retry counter to make memory consumption doesn't escalate
                carry_over = i
        if len(carry_over) > 0:
            i = np.vstack((end_reservoir, carry_over))
        else:
            i = end_reservoir
        try:
            self._pca.partial_fit(i, check_input=False)
        except LinAlgError:
            logger.warning("{i.shape[0]} samples were not used in PCA fitting due to LinAlgError", flush=True)
        self.loadings = self._pca.components_[:-1, :].T

    def _fit_lsi(self) -> None:
        from gensim.models import LsiModel
        from gensim.matutils import Dense2Corpus

        self._lsiModel = LsiModel(Dense2Corpus(controlled_compute(self.data.blocks[0], self.nthreads).T),
                                  num_topics=self.dims, chunksize=self.data.chunksize[0],
                                  id2word={x: x for x in range(self.data.shape[1])}, extra_samples=0)
        for n, i in enumerate(self.iter_blocks(msg="Fitting LSI model")):
            if n == 0:
                continue
            self._lsiModel.add_documents(Dense2Corpus(i.T))
        self.loadings = self._lsiModel.get_topics().T

    def _fit_ann(self):
        import hnswlib

        dims = self.dims
        if dims < 1:
            dims = self.data.shape[1]
        ann_idx = hnswlib.Index(space=self.annMetric, dim=dims)
        ann_idx.init_index(max_elements=self.nCells, ef_construction=self.annEfc,
                           M=self.annM, random_seed=self.randState)
        ann_idx.set_ef(self.annEf)
        ann_idx.set_num_threads(self.annThreads)
        for i in self.iter_blocks(msg='Fitting ANN'):
            ann_idx.add_items(self.reducer(i))
        return ann_idx

    def _fit_kmeans(self, do_ann_fit):
        from sklearn.cluster import MiniBatchKMeans

        if do_ann_fit is False:
            return None
        kmeans = MiniBatchKMeans(
            n_clusters=self.nClusters, random_state=self.randState,
            batch_size=self.batchSize)
        with threadpool_limits(limits=self.nthreads):
            for i in self.iter_blocks(msg='Fitting kmeans'):
                kmeans.partial_fit(self.reducer(i))
        temp = []
        for i in self.iter_blocks(msg='Estimating seed partitions'):
            temp.extend(kmeans.predict(self.reducer(i)))
        self.clusterLabels = np.array(temp)
        return kmeans
