"""Unsupervised clustering models for record linkage.

Alternative to Fellegi-Sunter EM — these models cluster pairs into
match/non-match groups without labelled training data.
"""

from __future__ import annotations

import pickle
from pathlib import Path
from typing import Dict, Iterator, Optional, Union

import numpy as np


class KMeansModel:
    """KMeans clustering with k=2 (match/non-match).

    Assigns match probability based on distance to the "match" centroid
    (the centroid with higher mean feature values).

    Args:
        auto_tune: If True, search k=2..5 and pick by silhouette score.
        **kwargs: Passed to sklearn KMeans constructor.
    """

    def __init__(self, auto_tune: bool = False, **kwargs):
        self._auto_tune = auto_tune
        self._kwargs = kwargs
        self._model = None
        self._match_cluster: Optional[int] = None
        self._tuning_results: Optional[Dict] = None

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

    def estimate(self, features: np.ndarray) -> "KMeansModel":
        """Fit KMeans on features."""
        from sklearn.cluster import KMeans
        from sklearn.metrics import silhouette_score

        if self._auto_tune:
            best_k, best_score = 2, -1.0
            results = {}
            for k in range(2, 6):
                model = KMeans(n_clusters=k, random_state=42, n_init=10, **self._kwargs)
                labels = model.fit_predict(features)
                score = silhouette_score(features, labels)
                results[k] = score
                if score > best_score:
                    best_score = score
                    best_k = k
            self._tuning_results = {"silhouette_scores": results, "best_k": best_k}
            self._model = KMeans(n_clusters=best_k, random_state=42, n_init=10, **self._kwargs)
            self._model.fit(features)
        else:
            self._model = KMeans(n_clusters=2, random_state=42, n_init=10, **self._kwargs)
            self._model.fit(features)

        # Identify match cluster: the one with higher mean feature values
        centers = self._model.cluster_centers_
        self._match_cluster = int(np.argmax(centers.mean(axis=1)))
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Return match probability based on distance to centroids."""
        if not self.is_fitted:
            raise RuntimeError("Model is not estimated. Call estimate() first.")

        distances = self._model.transform(features)  # (N, k)
        match_dist = distances[:, self._match_cluster]
        # Select non-match columns without copying the array
        non_match_cols = [i for i in range(distances.shape[1]) if i != self._match_cluster]
        non_match_dist = np.min(distances[:, non_match_cols], axis=1)

        # Sigmoid: P(match) = 1 / (1 + exp(-diff)), clipped to avoid overflow
        diff = np.clip(non_match_dist - match_dist, -500, 500)
        return 1.0 / (1.0 + np.exp(-diff))

    def predict_iter(self, features_iter: Iterator[np.ndarray]) -> Iterator[np.ndarray]:
        for chunk in features_iter:
            yield self.predict(chunk)

    def tuning_summary(self) -> Optional[Dict]:
        """Return auto-tune results, or None if auto_tune was not used."""
        return self._tuning_results

    def save(self, path: Union[str, Path]) -> None:
        data = {
            "type": "kmeans",
            "model": self._model,
            "match_cluster": self._match_cluster,
            "tuning_results": self._tuning_results,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "KMeansModel":
        with open(path, "rb") as f:
            data = pickle.load(f)
        obj = cls()
        obj._model = data["model"]
        obj._match_cluster = data["match_cluster"]
        obj._tuning_results = data.get("tuning_results")
        return obj


class GMMModel:
    """Gaussian Mixture Model with 2 components.

    More flexible than KMeans — models elliptical clusters.
    Returns posterior probability of match component.

    Args:
        auto_tune: If True, search n_components=1..5 and pick by lowest BIC.
        **kwargs: Passed to sklearn GaussianMixture constructor.
    """

    def __init__(self, auto_tune: bool = False, **kwargs):
        self._auto_tune = auto_tune
        self._kwargs = kwargs
        self._model = None
        self._match_component: Optional[int] = None
        self._tuning_results: Optional[Dict] = None

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

    def estimate(self, features: np.ndarray) -> "GMMModel":
        """Fit GMM on features."""
        from sklearn.mixture import GaussianMixture

        if self._auto_tune:
            best_n, best_bic = 2, np.inf
            results = {}
            for n in range(1, 6):
                model = GaussianMixture(
                    n_components=n, random_state=42, **self._kwargs
                )
                model.fit(features)
                bic = model.bic(features)
                results[n] = bic
                if bic < best_bic:
                    best_bic = bic
                    best_n = n
            self._tuning_results = {"bic_scores": results, "best_n_components": best_n}
            self._model = GaussianMixture(
                n_components=best_n, random_state=42, **self._kwargs
            )
            self._model.fit(features)
        else:
            self._model = GaussianMixture(
                n_components=2, random_state=42, **self._kwargs
            )
            self._model.fit(features)

        # Match component: the one with higher mean feature values
        if self._model.n_components == 1:
            self._match_component = 0
        else:
            means = self._model.means_
            self._match_component = int(np.argmax(means.mean(axis=1)))
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Return posterior probability of match component."""
        if not self.is_fitted:
            raise RuntimeError("Model is not estimated. Call estimate() first.")

        proba = self._model.predict_proba(features)
        return proba[:, self._match_component]

    def predict_iter(self, features_iter: Iterator[np.ndarray]) -> Iterator[np.ndarray]:
        for chunk in features_iter:
            yield self.predict(chunk)

    def tuning_summary(self) -> Optional[Dict]:
        """Return auto-tune results, or None if auto_tune was not used."""
        return self._tuning_results

    def save(self, path: Union[str, Path]) -> None:
        data = {
            "type": "gmm",
            "model": self._model,
            "match_component": self._match_component,
            "tuning_results": self._tuning_results,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "GMMModel":
        with open(path, "rb") as f:
            data = pickle.load(f)
        obj = cls()
        obj._model = data["model"]
        obj._match_component = data["match_component"]
        obj._tuning_results = data.get("tuning_results")
        return obj


class DBSCANModel:
    """DBSCAN density-based clustering for record linkage.

    Dense regions in feature space are identified as match clusters.
    Pairs in dense clusters with high mean features get high match scores.
    Noise points (label -1) get low scores.

    Args:
        eps: Maximum distance between samples in a neighborhood. If auto_tune
            is True and eps is not explicitly set, it will be estimated from data.
        min_samples: Minimum samples in a neighborhood for a core point.
        auto_tune: If True, estimate eps from k-distance graph knee detection.
        **kwargs: Passed to sklearn DBSCAN constructor.
    """

    def __init__(self, eps: float = 0.3, min_samples: int = 5,
                 auto_tune: bool = False, **kwargs):
        self._eps = eps
        self._eps_was_set = "eps" in kwargs or eps != 0.3
        self._min_samples = min_samples
        self._auto_tune = auto_tune
        self._kwargs = kwargs
        self._labels: Optional[np.ndarray] = None
        self._match_label: Optional[int] = None
        self._features: Optional[np.ndarray] = None
        self._tuning_results: Optional[Dict] = None

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

    @staticmethod
    def _estimate_eps(features: np.ndarray, min_samples: int) -> float:
        """Estimate eps using k-distance graph knee detection."""
        from sklearn.neighbors import NearestNeighbors

        k = min(min_samples, len(features) - 1)
        nn = NearestNeighbors(n_neighbors=k)
        nn.fit(features)
        distances, _ = nn.kneighbors(features)
        k_distances = np.sort(distances[:, -1])[::-1]  # sorted descending

        # Knee detection via second derivative maximum
        if len(k_distances) < 3:
            return float(np.median(k_distances))

        # Smooth with moving average to reduce noise
        window = max(3, len(k_distances) // 50)
        if window % 2 == 0:
            window += 1
        padded = np.pad(k_distances, window // 2, mode="edge")
        smoothed = np.convolve(padded, np.ones(window) / window, mode="valid")[
            : len(k_distances)
        ]

        # Second derivative
        d2 = np.diff(smoothed, n=2)
        if len(d2) == 0:
            return float(np.median(k_distances))

        knee_idx = int(np.argmax(d2)) + 1  # offset for diff
        return float(k_distances[min(knee_idx, len(k_distances) - 1)])

    def estimate(self, features: np.ndarray) -> "DBSCANModel":
        """Fit DBSCAN on features."""
        from sklearn.cluster import DBSCAN
        from sklearn.neighbors import NearestNeighbors

        # Store training features for NN-based prediction on new data.
        # Note: this doubles memory usage for the feature matrix.
        self._features = features.copy()

        if self._auto_tune and not self._eps_was_set:
            self._eps = self._estimate_eps(features, self._min_samples)
            self._tuning_results = {"estimated_eps": self._eps}

        model = DBSCAN(eps=self._eps, min_samples=self._min_samples, **self._kwargs)
        self._labels = model.fit_predict(features)

        # Find match cluster: cluster with highest mean feature values
        unique_labels = set(self._labels)
        unique_labels.discard(-1)  # remove noise

        if len(unique_labels) == 0:
            # All noise — no clear clusters
            self._match_label = None
        else:
            best_label = -1
            best_mean = -np.inf
            for label in unique_labels:
                mask = self._labels == label
                mean_val = features[mask].mean()
                if mean_val > best_mean:
                    best_mean = mean_val
                    best_label = label
            self._match_label = best_label

        # Build NN model for predict on new data
        self._nn = NearestNeighbors(n_neighbors=min(self._min_samples, len(features)))
        self._nn.fit(features)
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs based on cluster membership.

        Uses nearest-neighbor voting: for each input pair, find its nearest
        neighbors from the training set and compute the fraction that belong
        to the match cluster.
        """
        if not self.is_fitted:
            raise RuntimeError("Model is not estimated. Call estimate() first.")

        if self._match_label is None:
            # No clusters found — fall back to feature magnitude
            return features.mean(axis=1)

        distances, indices = self._nn.kneighbors(features)
        neighbor_labels = self._labels[indices]
        match_votes = (neighbor_labels == self._match_label).mean(axis=1)
        return match_votes.astype(np.float64)

    def predict_iter(self, features_iter: Iterator[np.ndarray]) -> Iterator[np.ndarray]:
        for chunk in features_iter:
            yield self.predict(chunk)

    def tuning_summary(self) -> Optional[Dict]:
        """Return auto-tune results, or None if auto_tune was not used."""
        return self._tuning_results

    def save(self, path: Union[str, Path]) -> None:
        data = {
            "type": "dbscan",
            "labels": self._labels,
            "match_label": self._match_label,
            "features": self._features,
            "nn": self._nn,
            "tuning_results": self._tuning_results,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "DBSCANModel":
        with open(path, "rb") as f:
            data = pickle.load(f)
        obj = cls()
        obj._labels = data["labels"]
        obj._match_label = data["match_label"]
        obj._features = data["features"]
        obj._nn = data["nn"]
        obj._tuning_results = data.get("tuning_results")
        return obj


# Registry for Linker.estimate() convenience
UNSUPERVISED_REGISTRY = {
    "kmeans": KMeansModel,
    "gmm": GMMModel,
    "dbscan": DBSCANModel,
}
