from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np

from .preprocess import flatten_features, restore_map

try:
    from sklearn.mixture import GaussianMixture  # type: ignore

    SKLEARN_AVAILABLE = True
except Exception:
    GaussianMixture = None
    SKLEARN_AVAILABLE = False


class SimpleGMM:
    def __init__(
        self,
        n_components: int,
        covariance_type: str = "diag",
        random_state: int = 42,
        reg_covar: float = 1e-6,
        max_iter: int = 40,
        tol: float = 1e-3,
    ) -> None:
        self.n_components = n_components
        self.covariance_type = covariance_type
        self.random_state = random_state
        self.reg_covar = reg_covar
        self.max_iter = max_iter
        self.tol = tol
        self.weights_: np.ndarray | None = None
        self.means_: np.ndarray | None = None
        self.covariances_: np.ndarray | None = None

    def fit(self, X: np.ndarray) -> "SimpleGMM":
        rng = np.random.RandomState(self.random_state)
        n_samples, n_features = X.shape
        if n_samples < self.n_components:
            raise ValueError("Number of samples must be >= number of components.")

        indices = rng.choice(n_samples, self.n_components, replace=False)
        means = X[indices].copy()
        cov = np.tile(X.var(axis=0) + self.reg_covar, (self.n_components, 1))
        weights = np.full(self.n_components, 1.0 / self.n_components)

        prev_ll = None
        for _ in range(self.max_iter):
            log_prob = self._estimate_log_prob(X, means, cov)
            log_prob += np.log(weights + 1e-9)
            max_log = log_prob.max(axis=1, keepdims=True)
            exp = np.exp(log_prob - max_log)
            resp = exp / (exp.sum(axis=1, keepdims=True) + 1e-9)

            Nk = resp.sum(axis=0) + 1e-9
            weights = Nk / n_samples
            means = (resp.T @ X) / Nk[:, None]
            cov = (resp.T @ (X ** 2)) / Nk[:, None] - means**2
            cov = np.clip(cov, self.reg_covar, None)

            ll = float(np.sum(max_log + np.log(exp.sum(axis=1, keepdims=True) + 1e-9)))
            if prev_ll is not None and abs(ll - prev_ll) < self.tol * max(abs(prev_ll), 1.0):
                break
            prev_ll = ll

        self.weights_ = weights
        self.means_ = means
        self.covariances_ = cov
        return self

    def _estimate_log_prob(self, X: np.ndarray, means: np.ndarray, cov: np.ndarray) -> np.ndarray:
        # Diagonal covariance only
        diff = X[:, None, :] - means[None, :, :]
        log_det = np.sum(np.log(2 * np.pi * cov), axis=1)
        quad = np.sum((diff**2) / cov[None, :, :], axis=2)
        return -0.5 * (log_det[None, :] + quad)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        if self.weights_ is None or self.means_ is None or self.covariances_ is None:
            raise ValueError("Model is not fitted.")
        log_prob = self._estimate_log_prob(X, self.means_, self.covariances_)
        log_prob += np.log(self.weights_ + 1e-9)
        max_log = log_prob.max(axis=1, keepdims=True)
        exp = np.exp(log_prob - max_log)
        return exp / (exp.sum(axis=1, keepdims=True) + 1e-9)


@dataclass
class SegmentationResult:
    labels: np.ndarray
    probabilities: np.ndarray | None
    uncertainty: np.ndarray
    model: Any
    mask: np.ndarray


def _calibrate_probabilities(probabilities: np.ndarray, temperature: float) -> np.ndarray:
    if temperature <= 0:
        return probabilities
    logp = np.log(probabilities + 1e-9)
    scaled = np.exp(logp / temperature)
    scaled /= scaled.sum(axis=1, keepdims=True)
    return scaled


def _entropy(probabilities: np.ndarray) -> np.ndarray:
    return -np.sum(probabilities * np.log(probabilities + 1e-9), axis=1)


def fit_gmm(
    features: np.ndarray,
    n_components: int,
    sample_size: int = 80000,
    covariance_type: str = "full",
    random_state: int = 42,
    engine: str = "safe",
) -> Any:
    if features.shape[0] > sample_size:
        idx = np.random.RandomState(random_state).choice(features.shape[0], sample_size, replace=False)
        sample = features[idx]
    else:
        sample = features

    if engine == "sklearn":
        if not SKLEARN_AVAILABLE or GaussianMixture is None:
            raise ValueError("Scikit-learn is not available for GMM.")
        model = GaussianMixture(
            n_components=n_components,
            covariance_type=covariance_type,
            random_state=random_state,
            reg_covar=1e-6,
        )
        model.fit(sample)
        return model

    # Safe mode: numpy-only diagonal GMM
    if covariance_type != "diag":
        covariance_type = "diag"
    model = SimpleGMM(
        n_components=n_components,
        covariance_type=covariance_type,
        random_state=random_state,
    )
    model.fit(sample)
    return model


def predict_proba_chunked(model: Any, features: np.ndarray, chunk_size: int = 200000) -> np.ndarray:
    total = features.shape[0]
    probs = np.zeros((total, model.n_components), dtype="float32")
    for start in range(0, total, chunk_size):
        end = min(start + chunk_size, total)
        probs[start:end] = model.predict_proba(features[start:end])
    return probs


def gmm_segment(
    bands: np.ndarray,
    n_components: int = 5,
    sample_size: int = 80000,
    covariance_type: str = "full",
    temperature: float = 1.0,
    return_probabilities: bool = True,
    engine: str = "safe",
) -> SegmentationResult:
    features, mask = flatten_features(bands)
    if features.size == 0:
        raise ValueError("No valid pixels found for segmentation.")

    mean = features.mean(axis=0)
    std = features.std(axis=0)
    std[std == 0] = 1.0
    features = (features - mean) / std

    model = fit_gmm(features, n_components, sample_size, covariance_type, engine=engine)
    probabilities = predict_proba_chunked(model, features)
    probabilities = _calibrate_probabilities(probabilities, temperature)
    labels = probabilities.argmax(axis=1)
    uncertainty = _entropy(probabilities)

    height, width = bands.shape[1], bands.shape[2]
    labels_map = restore_map(labels.astype("int32"), mask, (height, width), fill_value=-1)
    uncertainty_map = restore_map(uncertainty.astype("float32"), mask, (height, width), fill_value=np.nan)

    if return_probabilities:
        probs_maps = np.stack(
            [restore_map(probabilities[:, idx].astype("float32"), mask, (height, width), fill_value=np.nan)
             for idx in range(probabilities.shape[1])],
            axis=0,
        )
    else:
        probs_maps = None

    return SegmentationResult(
        labels=labels_map,
        probabilities=probs_maps,
        uncertainty=uncertainty_map,
        model=model,
        mask=mask,
    )


def weak_label_mask(probabilities: np.ndarray, threshold: float = 0.85) -> np.ndarray:
    if probabilities.ndim != 2:
        raise ValueError("Probabilities must be 2D (pixels, classes).")
    max_prob = probabilities.max(axis=1)
    return max_prob >= threshold


def gmm_segment_table(df, n_components: int = 5, engine: str = "safe"):
    import pandas as pd

    numeric = df.select_dtypes(include=["number"]).copy()
    if numeric.empty:
        raise ValueError("No numeric columns found for tabular segmentation.")
    numeric = numeric.fillna(numeric.median(numeric_only=True))
    features = numeric.values.astype("float32")

    mean = features.mean(axis=0)
    std = features.std(axis=0)
    std[std == 0] = 1.0
    features = (features - mean) / std

    model = fit_gmm(features, n_components=n_components, engine=engine)
    probabilities = model.predict_proba(features)
    labels = probabilities.argmax(axis=1)
    return labels, probabilities, model, list(numeric.columns)
