import numpy as np
import numba

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils import check_array, check_random_state
from sklearn.utils.validation import _check_sample_weight
from scipy.sparse import issparse, csr_matrix, coo_matrix

from enstop.utils import (
    normalize,
    coherence,
    mean_coherence,
    log_lift,
    mean_log_lift,
    standardize_input,
)
from enstop.plsa import plsa_init


@numba.njit(
    [
        "f4[:,::1](i4[::1],i4[::1],f4[:,::1],f4[:,::1],f4[:,::1],f4)",
        "f4[:,::1](i4[::1],i4[::1],f4[:,:],f4[:,::1],f4[:,::1],f4)",
    ],
    locals={
        "k": numba.types.intp,
        "w": numba.types.uint32,
        "d": numba.types.uint32,
        "z": numba.types.uint16,
        "v": numba.types.float32,
        "nz_idx": numba.types.uint32,
        "norm": numba.types.float32,
    },
    fastmath=True,
    nogil=True,
)
def plsa_e_step_on_a_block(
    block_rows,
    block_cols,
    p_w_given_z_block,
    p_z_given_d_block,
    p_z_given_wd_block,
    probability_threshold=1e-32,
):
    k = p_w_given_z_block.shape[0]

    for nz_idx in range(block_rows.shape[0]):
        if block_rows[nz_idx] < 0:
            break

        d = block_rows[nz_idx]
        w = block_cols[nz_idx]

        norm = 0.0
        for z in range(k):
            v = p_w_given_z_block[z, w] * p_z_given_d_block[d, z]
            if v > probability_threshold:
                p_z_given_wd_block[nz_idx, z] = v
                norm += v
            else:
                p_z_given_wd_block[nz_idx, z] = 0.0
        for z in range(k):
            if norm > 0:
                p_z_given_wd_block[nz_idx, z] /= norm

    return p_z_given_wd_block


@numba.njit(
    [
        "void(i4[::1],i4[::1],f4[::1],f4[:,::1],f4[:,::1],f4[:,::1],f4[::1],f4[::1])",
        "void(i4[::1],i4[::1],f4[::1],f4[:,:],f4[:,:],f4[:,::1],f4[::1],f4[::1])",
    ],
    locals={
        "k": numba.types.intp,
        "w": numba.types.uint32,
        "d": numba.types.uint32,
        "x": numba.types.float32,
        "z": numba.types.uint16,
        "nz_idx": numba.types.uint32,
        "s": numba.types.float32,
    },
    fastmath=True,
    nogil=True,
)
def plsa_partial_m_step_on_a_block(
    block_rows,
    block_cols,
    block_vals,
    p_w_given_z_block,
    p_z_given_d_block,
    p_z_given_wd_block,
    norm_pwz,
    norm_pdz_block,
):
    k = p_w_given_z_block.shape[0]

    for nz_idx in range(block_rows.shape[0]):
        if block_rows[nz_idx] < 0:
            break

        d = block_rows[nz_idx]
        w = block_cols[nz_idx]
        x = block_vals[nz_idx]

        for z in range(k):
            s = x * p_z_given_wd_block[nz_idx, z]

            p_w_given_z_block[z, w] += s
            p_z_given_d_block[d, z] += s

            norm_pwz[z] += s
            norm_pdz_block[d] += s


@numba.njit(
    "void(i4[:,:,::1],i4[:,:,::1],f4[:,:,::1],f4[:,:,::1],f4[:,:,::1],f4[:,:,:,::1],"
    "f4[:,:,:,::1],f4[:,:,:,::1],f4[:,::1],f4[:,:,::1],f4)",
    locals={
        "k": numba.types.intp,
        "z": numba.types.uint16,
        "d": numba.types.uint32,
        "i": numba.types.uint16,
        "j": numba.types.uint16,
        "n_w_blocks": numba.types.intp,
        "n_d_blocks": numba.types.intp,
    },
    parallel=True,
    fastmath=True,
    nogil=True,
)
def plsa_em_step_by_blocks(
    block_rows_ndarray,
    block_cols_ndarray,
    block_vals_ndarray,
    prev_p_w_given_z,
    prev_p_z_given_d,
    blocked_next_p_w_given_z,
    blocked_next_p_z_given_d,
    p_z_given_wd_block,
    blocked_norm_pwz,
    blocked_norm_pdz,
    e_step_thresh=1e-32,
):
    n_d_blocks = block_rows_ndarray.shape[0]
    n_w_blocks = block_rows_ndarray.shape[1]

    # n = prev_p_z_given_d.shape[0]
    # m = prev_p_w_given_z.shape[1]
    k = prev_p_z_given_d.shape[2]

    # zero out the norms for recomputation
    blocked_norm_pdz[:] = 0.0
    blocked_norm_pwz[:] = 0.0

    for i in numba.prange(n_d_blocks):

        for j in numba.prange(n_w_blocks):
            block_rows = block_rows_ndarray[i, j]
            block_cols = block_cols_ndarray[i, j]
            block_vals = block_vals_ndarray[i, j]

            plsa_e_step_on_a_block(
                block_rows,
                block_cols,
                prev_p_w_given_z[j],
                prev_p_z_given_d[i],
                p_z_given_wd_block[i, j],
                np.float32(e_step_thresh),
            )
            plsa_partial_m_step_on_a_block(
                block_rows,
                block_cols,
                block_vals,
                blocked_next_p_w_given_z[i, j],
                blocked_next_p_z_given_d[j, i],
                p_z_given_wd_block[i, j],
                blocked_norm_pwz[i],
                blocked_norm_pdz[j, i],
            )

    prev_p_z_given_d[:] = blocked_next_p_z_given_d.sum(axis=0)
    norm_pdz = blocked_norm_pdz.sum(axis=0)
    prev_p_w_given_z[:] = blocked_next_p_w_given_z.sum(axis=0)
    norm_pwz = blocked_norm_pwz.sum(axis=0)

    # Once complete we can normalize to complete the M step
    for z in numba.prange(k):
        if norm_pwz[z] > 0:
            for w_block in range(prev_p_w_given_z.shape[0]):
                for w_offset in range(prev_p_w_given_z.shape[2]):
                    prev_p_w_given_z[w_block, z, w_offset] /= norm_pwz[z]
        for d_block in range(prev_p_z_given_d.shape[0]):
            for d_offset in range(prev_p_z_given_d.shape[1]):
                if norm_pdz[d_block, d_offset] > 0:
                    prev_p_z_given_d[d_block, d_offset, z] /= norm_pdz[
                        d_block, d_offset
                    ]

    # Zero out the old matrices these matrices for next time
    blocked_next_p_z_given_d[:] = 0.0
    blocked_next_p_w_given_z[:] = 0.0


@numba.njit(
    locals={
        "i": numba.types.uint16,
        "j": numba.types.uint16,
        "k": numba.types.intp,
        "w": numba.types.uint32,
        "d": numba.types.uint32,
        "z": numba.types.uint16,
        "nz_idx": numba.types.uint32,
        "x": numba.types.float32,
        "result": numba.types.float32,
        "p_w_given_d": numba.types.float32,
    },
    fastmath=True,
    nogil=True,
    parallel=True,
)
def log_likelihood_by_blocks(
    block_rows_ndarray,
    block_cols_ndarray,
    block_vals_ndarray,
    p_w_given_z,
    p_z_given_d,
):
    result = 0.0
    k = p_z_given_d.shape[2]

    for i in numba.prange(block_rows_ndarray.shape[0]):
        for j in range(block_rows_ndarray.shape[1]):
            for nz_idx in range(block_rows_ndarray.shape[2]):
                if block_rows_ndarray[i, j, nz_idx] < 0:
                    break

                d = block_rows_ndarray[i, j, nz_idx]
                w = block_cols_ndarray[i, j, nz_idx]
                x = block_vals_ndarray[i, j, nz_idx]

                p_w_given_d = 0.0
                for z in range(k):
                    p_w_given_d += p_w_given_z[j, z, w] * p_z_given_d[i, d, z]

                result += x * np.log(p_w_given_d)

    return result


@numba.njit(fastmath=True, nogil=True)
def plsa_fit_inner_blockwise(
    block_rows_ndarray,
    block_cols_ndarray,
    block_vals_ndarray,
    p_w_given_z,
    p_z_given_d,
    block_row_size,
    block_col_size,
    n_iter=100,
    n_iter_per_test=10,
    tolerance=0.001,
    e_step_thresh=1e-32,
):
    k = p_z_given_d.shape[2]

    n_d_blocks = block_rows_ndarray.shape[0]
    n_w_blocks = block_rows_ndarray.shape[1]
    block_size = block_rows_ndarray.shape[2]

    p_z_given_wd_block = np.zeros(
        (n_d_blocks, n_w_blocks, block_size, k), dtype=np.float32
    )

    blocked_next_p_w_given_z = np.zeros(
        (
            np.int64(n_d_blocks),
            np.int64(n_w_blocks),
            np.int64(k),
            np.int64(block_col_size),
        ),
        dtype=np.float32,
    )
    blocked_norm_pwz = np.zeros((n_d_blocks, k), dtype=np.float32)
    blocked_next_p_z_given_d = np.zeros(
        (
            np.int64(n_w_blocks),
            np.int64(n_d_blocks),
            np.int64(block_row_size),
            np.int64(k),
        ),
        dtype=np.float32,
    )
    blocked_norm_pdz = np.zeros(
        (np.int64(n_w_blocks), np.int64(n_d_blocks), np.int64(block_row_size)),
        dtype=np.float32,
    )

    previous_log_likelihood = log_likelihood_by_blocks(
        block_rows_ndarray,
        block_cols_ndarray,
        block_vals_ndarray,
        p_w_given_z,
        p_z_given_d,
    )

    for i in range(n_iter):
        plsa_em_step_by_blocks(
            block_rows_ndarray,
            block_cols_ndarray,
            block_vals_ndarray,
            p_w_given_z,
            p_z_given_d,
            blocked_next_p_w_given_z,
            blocked_next_p_z_given_d,
            p_z_given_wd_block,
            blocked_norm_pwz,
            blocked_norm_pdz,
            e_step_thresh,
        )

        if i % n_iter_per_test == 0:
            current_log_likelihood = log_likelihood_by_blocks(
                block_rows_ndarray,
                block_cols_ndarray,
                block_vals_ndarray,
                p_w_given_z,
                p_z_given_d,
            )
            change = np.abs(current_log_likelihood - previous_log_likelihood)
            if change / np.abs(current_log_likelihood) < tolerance:
                break
            else:
                previous_log_likelihood = current_log_likelihood

    return p_z_given_d, p_w_given_z


def plsa_fit(
    X,
    k,
    n_row_blocks=8,
    n_col_blocks=8,
    init="random",
    n_iter=100,
    n_iter_per_test=10,
    tolerance=0.001,
    e_step_thresh=1e-32,
    random_state=None,
):
    rng = check_random_state(random_state)
    p_z_given_d_init, p_w_given_z_init = plsa_init(X, k, init=init, rng=rng)

    A = X.tocsr().astype(np.float32)

    n = A.shape[0]
    m = A.shape[1]

    block_row_size = np.uint16(np.ceil(A.shape[0] / n_row_blocks))
    block_col_size = np.uint16(np.ceil(A.shape[1] / n_col_blocks))

    p_z_given_d = np.zeros((block_row_size * n_row_blocks, k), dtype=np.float32)
    p_z_given_d[: p_z_given_d_init.shape[0]] = p_z_given_d_init
    p_z_given_d = p_z_given_d.reshape(n_row_blocks, block_row_size, k)

    p_w_given_z = np.zeros((k, block_col_size * n_col_blocks), dtype=np.float32)
    p_w_given_z[:, : p_w_given_z_init.shape[1]] = p_w_given_z_init
    # p_w_given_z = np.transpose(
    #     p_w_given_z.T.reshape(n_col_blocks, block_col_size, k), axes=[0, 2, 1]
    # ).astype(np.float32, order="C")
    p_w_given_z = np.stack(np.hsplit(p_w_given_z, n_col_blocks))

    A_blocks = [[0] * n_col_blocks for i in range(n_row_blocks)]
    max_nnz_per_block = 0
    for i in range(n_row_blocks):

        row_start = block_row_size * i
        row_end = min(row_start + block_row_size, n)

        for j in range(n_col_blocks):

            col_start = block_col_size * j
            col_end = min(col_start + block_col_size, m)

            A_blocks[i][j] = A[row_start:row_end, col_start:col_end].tocoo()
            if A_blocks[i][j].nnz > max_nnz_per_block:
                max_nnz_per_block = A_blocks[i][j].nnz

    block_rows_ndarray = np.full(
        (n_row_blocks, n_col_blocks, max_nnz_per_block), -1, dtype=np.int32
    )
    block_cols_ndarray = np.full(
        (n_row_blocks, n_col_blocks, max_nnz_per_block), -1, dtype=np.int32
    )
    block_vals_ndarray = np.zeros(
        (n_row_blocks, n_col_blocks, max_nnz_per_block), dtype=np.float32
    )
    for i in range(n_row_blocks):
        for j in range(n_col_blocks):
            nnz = A_blocks[i][j].nnz
            block_rows_ndarray[i, j, :nnz] = A_blocks[i][j].row
            block_cols_ndarray[i, j, :nnz] = A_blocks[i][j].col
            block_vals_ndarray[i, j, :nnz] = A_blocks[i][j].data

    p_z_given_d, p_w_given_z = plsa_fit_inner_blockwise(
        block_rows_ndarray,
        block_cols_ndarray,
        block_vals_ndarray,
        p_w_given_z,
        p_z_given_d,
        block_row_size,
        block_col_size,
        n_iter=n_iter,
        n_iter_per_test=n_iter_per_test,
        tolerance=tolerance,
        e_step_thresh=e_step_thresh,
    )
    p_z_given_d = np.vstack(p_z_given_d)[:n, :]
    p_w_given_z = np.hstack(p_w_given_z)[:, :m]

    return p_z_given_d, p_w_given_z


class BlockParallelPLSA(BaseEstimator, TransformerMixin):
    def __init__(
        self,
        n_components=10,
        init="random",
        n_row_blocks=8,
        n_col_blocks=8,
        n_iter=100,
        n_iter_per_test=10,
        tolerance=0.001,
        e_step_thresh=1e-32,
        transform_random_seed=42,
        random_state=None,
    ):

        self.n_components = n_components
        self.init = init
        self.n_row_blocks = n_row_blocks
        self.n_col_blocks = n_col_blocks
        self.n_iter = n_iter
        self.n_iter_per_test = n_iter_per_test
        self.tolerance = tolerance
        self.e_step_thresh = e_step_thresh
        self.transform_random_seed = transform_random_seed
        self.random_state = random_state

    def fit(self, X, y=None, sample_weight=None):
        """Learn the pLSA model for the data X and return the document vectors.

        This is more efficient than calling fit followed by transform.

        Parameters
        ----------
        X: array or sparse matrix of shape (n_docs, n_words)
            The data matrix pLSA is attempting to fit to.

        y: Ignored

        sample_weight: array of shape (n_docs,)
            Input document weights.

        Returns
        -------
        self
        """
        self.fit_transform(X, sample_weight=sample_weight)
        return self

    def fit_transform(self, X, y=None, sample_weight=None):
        """Learn the pLSA model for the data X and return the document vectors.

        This is more efficient than calling fit followed by transform.

        Parameters
        ----------
        X: array or sparse matrix of shape (n_docs, n_words)
            The data matrix pLSA is attempting to fit to.

        y: Ignored

        sample_weight: array of shape (n_docs,)
            Input document weights.

        Returns
        -------
        embedding: array of shape (n_docs, n_topics)
            An embedding of the documents into a topic space.
        """

        X = check_array(X, accept_sparse="csr")
        X = standardize_input(X)

        if not issparse(X):
            X = csr_matrix(X)

        sample_weight = _check_sample_weight(sample_weight, X, dtype=np.float32)

        if np.any(X.data < 0):
            raise ValueError(
                "PLSA is only valid for matrices with non-negative " "entries"
            )

        row_sums = np.array(X.sum(axis=1).T)[0]
        good_rows = row_sums != 0

        if not np.all(good_rows):
            zero_rows_found = True
            data_for_fitting = X[good_rows]
        else:
            zero_rows_found = False
            data_for_fitting = X

        U, V = plsa_fit(
            data_for_fitting,
            self.n_components,
            n_row_blocks=self.n_row_blocks,
            n_col_blocks=self.n_col_blocks,
            init=self.init,
            n_iter=self.n_iter,
            n_iter_per_test=self.n_iter_per_test,
            tolerance=self.tolerance,
            e_step_thresh=self.e_step_thresh,
            random_state=self.random_state,
        )

        if zero_rows_found:
            self.embedding_ = np.zeros((X.shape[0], self.n_components))
            self.embedding_[good_rows] = U
        else:
            self.embedding_ = U

        self.components_ = V
        self.training_data_ = X

        return self.embedding_
