"""Interval Type-2 Takagi-Sugeno-Kang fuzzy inference system."""

from __future__ import annotations

import numpy as np

from t1_fis import T1TSK, _FIRE_EPS


class IT2TSK(T1TSK):
    """IT2 first-order TSK system using the Nie-Tan output method.

    The lower and upper Gaussian widths are sigma * (1 - delta) and
    sigma * (1 + delta). Training reuses T1TSK's standardization, clustering,
    and least-squares consequent fitting with the Nie-Tan firing strengths.
    """

    def __init__(self, n_rules: int, delta: float = 0.2, random_state: int = 42) -> None:
        if not 0.0 <= delta < 1.0:
            raise ValueError(f"delta must satisfy 0 <= delta < 1, got {delta}")
        super().__init__(n_rules=n_rules, random_state=random_state)
        self.delta = float(delta)

    def _membership_bounds(
        self, X_scaled: np.ndarray
    ) -> tuple[np.ndarray, np.ndarray]:
        """Return lower and upper memberships with shape (n, rules, features)."""
        assert self.m_ is not None and self.sigma_ is not None
        sigma_lo = self.sigma_ * (1.0 - self.delta)
        sigma_hi = self.sigma_ * (1.0 + self.delta)
        diff = X_scaled[:, None, :] - self.m_[None, :, :]
        lower = np.exp(-(diff**2) / (2.0 * sigma_lo[None, :, :] ** 2))
        upper = np.exp(-(diff**2) / (2.0 * sigma_hi[None, :, :] ** 2))
        return lower, upper

    def _firing_strength_bounds(
        self, X_scaled: np.ndarray
    ) -> tuple[np.ndarray, np.ndarray]:
        """Return lower and upper rule firing strengths."""
        lower, upper = self._membership_bounds(X_scaled)
        return np.prod(lower, axis=2), np.prod(upper, axis=2)

    def _firing_strengths(self, X_scaled: np.ndarray) -> np.ndarray:
        """Return Nie-Tan effective strengths f_lo + f_hi."""
        firing_lo, firing_hi = self._firing_strength_bounds(X_scaled)
        return firing_lo + firing_hi

    def predict(self, X: np.ndarray) -> np.ndarray:
        """Predict using Nie-Tan defuzzification of lower and upper strengths."""
        self._check_fitted()
        X = np.asarray(X, dtype=float)
        if X.ndim != 2:
            raise ValueError(f"X must be 2D, got shape {X.shape}")
        X_scaled = self._standardize(X)
        firing_lo, firing_hi = self._firing_strength_bounds(X_scaled)
        firing = firing_lo + firing_hi
        rule_out = self._rule_outputs(X_scaled)

        numer = np.sum(firing * rule_out, axis=1)
        denom = np.sum(firing, axis=1)
        zero_firing = denom <= 2.0 * _FIRE_EPS
        predictions = numer / np.maximum(denom, 2.0 * _FIRE_EPS)
        if np.any(zero_firing):
            predictions[zero_firing] = np.mean(rule_out[zero_firing], axis=1)
        return predictions