from collections.abc import Iterable
from typing import Any, Literal

import numpy as np
from pydantic.v1 import dataclasses

from molflux.modelzoo.info import ModelInfo
from molflux.modelzoo.model import ModelConfig
from molflux.modelzoo.models.sklearn import (
    SKLearnClassificationMixin,
    SKLearnModelBase,
)

try:
    from sklearn.discriminant_analysis import (
        LinearDiscriminantAnalysis as SKLinearDiscriminantAnalysis,
    )
except ImportError as e:
    from molflux.modelzoo.errors import ExtrasDependencyImportError

    raise ExtrasDependencyImportError("sklearn", e) from None


_DESCRIPTION = """
Linear Discriminant Analysis.

A classifier with a linear decision boundary, generated by fitting class
conditional densities to the data and using Bayes' rule.

The model fits a Gaussian density to each class, assuming that all classes
share the same covariance matrix.
"""

_CONFIG_DESCRIPTION = """
Parameters
----------
solver : {'svd', 'lsqr', 'eigen'}, default='svd'
    Solver to use, possible values:
      - 'svd': Singular value decomposition (default).
        Does not compute the covariance matrix, therefore this solver is
        recommended for data with a large number of features.
      - 'lsqr': Least squares solution.
        Can be combined with shrinkage or custom covariance estimator.
      - 'eigen': Eigenvalue decomposition.
        Can be combined with shrinkage or custom covariance estimator.
shrinkage : 'auto' or float, default=None
    Shrinkage parameter, possible values:
      - None: no shrinkage (default).
      - 'auto': automatic shrinkage using the Ledoit-Wolf lemma.
      - float between 0 and 1: fixed shrinkage parameter.
    This should be left to None if `covariance_estimator` is used.
    Note that shrinkage works only with 'lsqr' and 'eigen' solvers.
priors : array-like of shape (n_classes,), default=None
    The class prior probabilities. By default, the class proportions are
    inferred from the training data.
n_components : int, default=None
    Number of components (<= min(n_classes - 1, n_features)) for
    dimensionality reduction. If None, will be set to
    min(n_classes - 1, n_features). This parameter only affects the
    `transform` method.
store_covariance : bool, default=False
    If True, explicitly compute the weighted within-class covariance
    matrix when solver is 'svd'. The matrix is always computed
    and stored for the other solvers.
tol : float, default=1.0e-4
    Absolute threshold for a singular value of X to be considered
    significant, used to estimate the rank of X. Dimensions whose
    singular values are non-significant are discarded. Only used if
    solver is 'svd'.
covariance_estimator : covariance estimator, default=None
    If not None, `covariance_estimator` is used to estimate
    the covariance matrices instead of relying on the empirical
    covariance estimator (with potential shrinkage).
    The object should have a fit method and a ``covariance_`` attribute
    like the estimators in :mod:`sklearn.covariance`.
    if None the shrinkage parameter drives the estimate.
    This should be left to None if `shrinkage` is used.
    Note that `covariance_estimator` works only with 'lsqr' and 'eigen'
    solvers.
"""

SolverT = Literal["svd", "lsqr", "eigen"]


class Config:
    arbitrary_types_allowed = True
    extra = "forbid"


@dataclasses.dataclass(config=Config)
class LinearDiscriminantAnalysisClassifierConfig(ModelConfig):
    solver: SolverT = "svd"
    shrinkage: None | int | Literal["auto"] = None
    priors: Iterable[float] | None = None
    n_components: int | None = None
    store_covariance: bool = False
    tol: float = 1e-4
    covariance_estimator: Any | None = None


class LinearDiscriminantAnalysisClassifier(
    SKLearnClassificationMixin,
    SKLearnModelBase[LinearDiscriminantAnalysisClassifierConfig],
):
    @property
    def _config_builder(self) -> type[LinearDiscriminantAnalysisClassifierConfig]:
        return LinearDiscriminantAnalysisClassifierConfig

    def _info(self) -> ModelInfo:
        return ModelInfo(
            model_description=_DESCRIPTION,
            config_description=_CONFIG_DESCRIPTION,
        )

    def _instantiate_model(self) -> SKLinearDiscriminantAnalysis:
        config = self.model_config
        return SKLinearDiscriminantAnalysis(
            solver=config.solver,
            shrinkage=config.shrinkage,
            priors=None if config.priors is None else np.array(config.priors),
            n_components=config.n_components,
            store_covariance=config.store_covariance,
            tol=config.tol,
            covariance_estimator=config.covariance_estimator,
        )
