"""Fellegi-Sunter probabilistic record linkage model with EM estimation.

This implements the classic Fellegi-Sunter framework:
- Each comparison is discretized into levels (e.g., exact match, close, no match)
- EM algorithm estimates m-probabilities P(level | match) and u-probabilities P(level | non-match)
- Match weights are log2(m/u) per comparison level
- Total match weight is the sum across all comparisons
- Probability is converted via sigmoid of the total weight
"""

from __future__ import annotations

import pickle
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterator, List, Optional, Union

import numpy as np


@dataclass
class ComparisonLevels:
    """Discretization config for one comparison."""

    thresholds: List[float] = field(default_factory=lambda: [0.7, 1.0])
    # Scores < thresholds[0] → level 0 (no match)
    # thresholds[0] <= score < thresholds[1] → level 1 (partial)
    # score >= thresholds[-1] → level len(thresholds) (exact)


class FellegiSunterModel:
    """Unsupervised probabilistic record linkage model.

    Uses the Fellegi-Sunter framework with Expectation-Maximization to
    estimate match/non-match probabilities without labelled training data.
    """

    def __init__(
        self,
        comparison_levels: Optional[List[ComparisonLevels]] = None,
        n_iterations: int = 25,
        prior_match_rate: float = 0.05,
    ):
        self._comparison_levels = comparison_levels
        self._n_iterations = n_iterations
        self._prior_match_rate = prior_match_rate

        # Estimated parameters
        self.m_probs_: Optional[List[np.ndarray]] = None  # P(level | match) per comparison
        self.u_probs_: Optional[List[np.ndarray]] = None  # P(level | non-match) per comparison
        self.match_weights_: Optional[List[np.ndarray]] = None  # log2(m/u) per level
        self.prior_: Optional[float] = None

    @property
    def is_fitted(self) -> bool:
        return self.m_probs_ is not None

    def _discretize(self, features: np.ndarray) -> np.ndarray:
        """Convert continuous features to discrete levels.

        Args:
            features: (N, F) continuous feature matrix

        Returns:
            (N, F) integer level matrix
        """
        n_pairs, n_features = features.shape
        levels = np.zeros((n_pairs, n_features), dtype=np.int32)

        for j in range(n_features):
            if self._comparison_levels and j < len(self._comparison_levels):
                thresholds = self._comparison_levels[j].thresholds
            else:
                # Default: 3 levels with thresholds at 0.7 and 1.0
                thresholds = [0.7, 1.0]

            for k, t in enumerate(thresholds):
                levels[:, j] += (features[:, j] >= t).astype(np.int32)

        return levels

    def estimate(self, features: np.ndarray) -> "FellegiSunterModel":
        """Run EM algorithm to estimate m and u probabilities.

        Args:
            features: (N, F) feature matrix with similarity scores in [0, 1]
        """
        n_pairs, n_features = features.shape
        levels = self._discretize(features)

        # Number of levels per comparison
        n_levels = []
        for j in range(n_features):
            if self._comparison_levels and j < len(self._comparison_levels):
                n_levels.append(len(self._comparison_levels[j].thresholds) + 1)
            else:
                n_levels.append(3)  # default: 3 levels

        # Initialize m and u probabilities with slight asymmetry
        eps = 1e-6
        m_probs = []
        u_probs = []
        for j in range(n_features):
            nl = n_levels[j]
            # m_probs: higher levels more likely for matches
            m = np.linspace(0.1, 0.5, nl)
            m = m / m.sum()
            m_probs.append(m)
            # u_probs: lower levels more likely for non-matches
            u = np.linspace(0.5, 0.1, nl)
            u = u / u.sum()
            u_probs.append(u)

        lam = self._prior_match_rate  # P(match)

        # EM iterations
        for iteration in range(self._n_iterations):
            # E-step: compute P(match | pair) for each pair
            log_match = np.log(lam + eps)
            log_nonmatch = np.log(1 - lam + eps)

            for j in range(n_features):
                pair_levels = levels[:, j]
                log_match = log_match + np.log(m_probs[j][pair_levels] + eps)
                log_nonmatch = log_nonmatch + np.log(u_probs[j][pair_levels] + eps)

            # Log-sum-exp for numerical stability
            max_log = np.maximum(log_match, log_nonmatch)
            log_total = max_log + np.log(
                np.exp(log_match - max_log) + np.exp(log_nonmatch - max_log)
            )
            gamma = np.exp(log_match - log_total)  # P(match | pair)

            # M-step: update m_probs, u_probs, and lambda
            lam = np.mean(gamma)
            lam = np.clip(lam, eps, 1 - eps)

            for j in range(n_features):
                nl = n_levels[j]
                new_m = np.zeros(nl)
                new_u = np.zeros(nl)

                for k in range(nl):
                    mask = levels[:, j] == k
                    new_m[k] = np.sum(gamma[mask]) + eps
                    new_u[k] = np.sum(1 - gamma[mask]) + eps

                m_probs[j] = new_m / new_m.sum()
                u_probs[j] = new_u / new_u.sum()

        # Store results
        self.m_probs_ = m_probs
        self.u_probs_ = u_probs
        self.prior_ = float(lam)

        # Compute match weights: log2(m/u) per level
        self.match_weights_ = []
        for j in range(n_features):
            weights = np.log2((m_probs[j] + eps) / (u_probs[j] + eps))
            self.match_weights_.append(weights)

        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using estimated Fellegi-Sunter weights.

        Args:
            features: (N, F) feature matrix

        Returns:
            (N,) array of match probabilities
        """
        if not self.is_fitted:
            raise RuntimeError("Model is not estimated. Call estimate() first.")

        levels = self._discretize(features)
        n_pairs, n_features = features.shape

        # Sum match weights across comparisons (log2 scale)
        total_weight = np.zeros(n_pairs)
        for j in range(n_features):
            pair_levels = levels[:, j]
            total_weight += self.match_weights_[j][pair_levels]

        # Convert to probability via sigmoid.
        # Weights are log2(m/u); scale by ln(2) to get natural-log scale
        # for correct probability calibration.
        ln2 = np.log(2.0)
        return 1.0 / (1.0 + np.exp(-total_weight * ln2))

    def predict_iter(self, features_iter: Iterator[np.ndarray]) -> Iterator[np.ndarray]:
        """Streaming prediction over chunked feature matrices."""
        for chunk in features_iter:
            yield self.predict(chunk)

    def save(self, path: Union[str, Path]) -> None:
        """Save model to disk."""
        data = {
            "type": "fellegi_sunter",
            "m_probs": self.m_probs_,
            "u_probs": self.u_probs_,
            "match_weights": self.match_weights_,
            "prior": self.prior_,
            "comparison_levels": self._comparison_levels,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "FellegiSunterModel":
        """Load model from disk."""
        with open(path, "rb") as f:
            data = pickle.load(f)
        model = cls(comparison_levels=data["comparison_levels"])
        model.m_probs_ = data["m_probs"]
        model.u_probs_ = data["u_probs"]
        model.match_weights_ = data["match_weights"]
        model.prior_ = data["prior"]
        return model

    def summary(self) -> dict:
        """Return a summary of estimated parameters."""
        if not self.is_fitted:
            raise RuntimeError("Model is not estimated. Call estimate() first.")

        result = {
            "prior_match_rate": self.prior_,
            "comparisons": [],
        }
        for j in range(len(self.m_probs_)):
            result["comparisons"].append({
                "m_probabilities": self.m_probs_[j].tolist(),
                "u_probabilities": self.u_probs_[j].tolist(),
                "match_weights": self.match_weights_[j].tolist(),
            })
        return result
