from typing import Literal

import numpy as np
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import MinMaxScaler

from .clle import clle
from .init_pca import init_pca
from .learning_l import learning_l
from .learning_s import learning_s
from .opt_scale import opt_scale
from .pps import pps


def sude(
    X: np.ndarray,
    no_dims: int = 2,
    k1: int = 20,
    normalize: bool = True,
    large: bool = False,
    initialize: Literal["le", "pca", "mds"] = "le",
    agg_coef: float = 1.2,
    T_epoch: int = 50,
):
    """
    This function returns representation of the N by D matrix X in the
    lower-dimensional space. Each row in X represents an observation.

    Parameters are:

    'no_dims'      - A positive integer specifying the number of dimension of
                     the representation Y.
                     Default: 2
    'k1'           - A non-negative integer specifying the number of nearest
                     neighbors for PPS to sample landmarks. It must be
                     smaller than N.
                     Default: adaptive
    'normalize'    - Logical scalar. If true, normalize X using min-max
                     normalization. If features in X are on different scales,
                     'Normalize' should be set to true because the learning
                     process is based on nearest neighbors and features with
                     large scales can override the contribution of features
                     with small scales.
                     Default: True
    'large'        - Logical scalar. If true, the data can be split into
                     multiple blocks to avoid the problem of memory overflow,
                     and the gradient can be computed block by block using
                     'learning_l' function.
                     Default: False
    'initialize'   - A string specifying the method for initializing Y before
                     manifold learning.
        'le'       - Laplacian eigenmaps.
        'pca'      - Principal component analysis.
        'mds'      - Multidimensional scaling.
                     Default: 'le'
    'agg_coef'     - A positive scalar specifying the aggregation
                     coefficient.
                     Default: 1.2
    'T_epoch'      - Maximum number of epochs to take.
                     Default: 50

    """
    assert isinstance(no_dims, int) and no_dims > 0 and no_dims < X.shape[1]
    assert isinstance(k1, int) and k1 >= 0 and no_dims < X.shape[0]
    assert normalize in [True, False]
    assert large in [True, False]
    assert initialize in ["le", "pca", "mds"]
    assert agg_coef >= 0
    assert isinstance(T_epoch, int) and T_epoch > 0

    # Remove duplicate observations
    X, orig_id = np.unique(X, axis=0, return_inverse=True)

    # Obtain size and dimension of data
    n, dim = X.shape

    # Normalize the data
    if normalize:
        X = MinMaxScaler().fit(X).transform(X)

    # Perform PPS to obtain the landmarks
    if k1 > 0:
        if n > 5000 and dim > 50:
            xx = init_pca(X, no_dims, 0.8)
            get_knn = NearestNeighbors(n_neighbors=k1 + 1).fit(xx).kneighbors(
                xx,
                return_distance=False,
            )
        else:
            get_knn = NearestNeighbors(n_neighbors=k1 + 1).fit(X).kneighbors(
                X,
                return_distance=False,
            )
        _, rnn = np.unique(get_knn, return_counts=True)
        id_samp = pps(get_knn, rnn, 1)
    else:
        get_knn = []
        rnn = []
        id_samp = list(range(n))

    X_samp = X[id_samp]

    # Compute embedding of landmarks
    if not large:
        Y_samp, k2 = learning_s(
            X_samp,
            k1,
            get_knn,
            rnn,
            id_samp,
            no_dims,
            initialize,
            agg_coef,
            T_epoch,
        )
    else:
        Y_samp, k2 = learning_l(
            X_samp,
            k1,
            get_knn,
            rnn,
            id_samp,
            no_dims,
            initialize,
            agg_coef,
            T_epoch,
        )

    # Compute embedding of non-landmarks
    if k1 > 0:
        id_rest = np.setdiff1d(range(n), id_samp)
        X_rest = X[id_rest]
        Y_rest = np.zeros((len(id_rest), no_dims))
        # Compute the optimal scale for each landmark
        scale = opt_scale(X_samp, Y_samp, k2)
        top_k = no_dims + 1
        near_dis, near_samp = NearestNeighbors(n_neighbors=top_k).fit(
            X_samp
        ).kneighbors(X_rest)
        for i in range(len(id_rest)):
            near_top_k = near_samp[i]
            top_X = X_samp[near_top_k]
            top_Y = Y_samp[near_top_k]
            N_dis = near_dis[i, 0] * scale[near_top_k[0]]
            # Perform CLLE
            Y_rest[i] = clle(top_X, top_Y, X_rest[i], N_dis)
        YY = np.zeros((n, no_dims))
        YY[id_rest] = Y_rest
        YY[id_samp] = Y_samp
    else:
        YY = Y_samp
    Y = YY[orig_id]

    return Y
