from __future__ import annotations

from typing import Iterator, Optional

import numpy as np
import pandas as pd

from blazematch.blocking import EmbeddingBlocker, RuleBlocker
from blazematch.clustering import UNSUPERVISED_REGISTRY
from blazematch.config import LinkConfig
from blazematch.features import FeatureComputer
from blazematch.fellegi_sunter import FellegiSunterModel
from blazematch.model import MODEL_REGISTRY, MatchModel


class Linker:
    """End-to-end record linkage pipeline.

    Supervised usage (with labelled data):
        linker = Linker(df_left, df_right, config)
        linker.block().compute_features().fit(labels_df)
        results = linker.predict()

    Unsupervised usage (Fellegi-Sunter, no labels needed):
        linker = Linker(df_left, df_right, config)
        linker.block().compute_features().estimate()
        results = linker.predict()

    All methods return self for chaining.
    """

    def __init__(
        self,
        df_left: pd.DataFrame,
        df_right: pd.DataFrame,
        config: LinkConfig,
        in_memory: bool = True,
    ):
        if len(df_left) == 0 or len(df_right) == 0:
            raise ValueError("Both df_left and df_right must be non-empty.")
        self.df_left = df_left.reset_index(drop=True)
        self.df_right = df_right.reset_index(drop=True)
        self.config = config
        self.config.in_memory = in_memory  # reserved for future streaming mode

        self._feature_computer = FeatureComputer(
            self.df_left, self.df_right, config.comparisons
        )
        self._model: Optional[MatchModel] = None
        self._fs_model: Optional[FellegiSunterModel] = None

        self._pairs: Optional[pd.DataFrame] = None
        self._features: Optional[np.ndarray] = None

    @property
    def pairs(self) -> Optional[pd.DataFrame]:
        """Candidate pairs DataFrame [left_idx, right_idx]."""
        return self._pairs

    @property
    def features(self) -> Optional[np.ndarray]:
        """Feature matrix (N_pairs, N_features)."""
        return self._features

    @property
    def feature_names(self):
        return self._feature_computer.feature_names

    def block(self) -> "Linker":
        """Execute blocking to generate candidate pairs.

        Uses rule-based blocking, embedding-based blocking, or both,
        depending on config. Results are unioned and deduplicated.
        """
        parts = []

        # Rule-based blocking
        if self.config.blocking_rules:
            rule_blocker = RuleBlocker(
                self.df_left,
                self.df_right,
                self.config.blocking_rules,
                self.config.max_candidates_per_record,
            )
            parts.append(rule_blocker.block())

        # Embedding-based blocking
        if self.config.embedding_blocking is not None:
            emb_blocker = EmbeddingBlocker(
                self.df_left,
                self.df_right,
                self.config.embedding_blocking,
                self.config.max_candidates_per_record,
            )
            parts.append(emb_blocker.block())

        if not parts:
            raise ValueError(
                "No blocking rules or embedding config provided. "
                "Set blocking_rules or embedding_blocking in LinkConfig."
            )

        if len(parts) == 1:
            self._pairs = parts[0]
        else:
            # Union and deduplicate
            combined = pd.concat(parts, ignore_index=True)
            self._pairs = combined.drop_duplicates(
                subset=["left_idx", "right_idx"]
            ).reset_index(drop=True)

        return self

    def estimate_ram(
        self,
        model: str = "logistic",
        blocking_reduction_factor: Optional[float] = None,
        system_ram_gb: Optional[float] = None,
        model_kwargs: Optional[dict] = None,
        print_summary: bool = True,
    ) -> "RAMEstimate":
        """Estimate RAM requirements for this linkage job.

        If block() has been called, uses the exact pair count and
        ``blocking_reduction_factor`` is ignored.  Otherwise, falls back
        to heuristic estimation using blocking config and the optional
        ``blocking_reduction_factor``.

        Args:
            model: Model type for training/inference estimation.
            blocking_reduction_factor: Heuristic for pair reduction (ignored if
                block() was already called).
            system_ram_gb: Override detected system RAM (in GB).
            model_kwargs: Optional model parameters (n_estimators, etc.).
            print_summary: If True, print human-readable summary.

        Returns:
            RAMEstimate with per-stage memory estimates.
        """
        import warnings as _warnings

        from blazematch.estimate_ram import estimate_ram as _estimate_ram

        n_pairs = len(self._pairs) if self._pairs is not None else None

        if n_pairs is not None and blocking_reduction_factor is not None:
            _warnings.warn(
                "blocking_reduction_factor is ignored because block() has "
                "already been called and exact pair count is used.",
                stacklevel=2,
            )
            blocking_reduction_factor = None

        return _estimate_ram(
            n_left=len(self.df_left),
            n_right=len(self.df_right),
            comparisons=self.config.comparisons,
            blocking_rules=self.config.blocking_rules or None,
            embedding_blocking=self.config.embedding_blocking,
            model=model,
            n_pairs=n_pairs,
            blocking_reduction_factor=blocking_reduction_factor,
            system_ram_gb=system_ram_gb,
            model_kwargs=model_kwargs,
            print_summary=print_summary,
        )

    def compute_features(self) -> "Linker":
        """Compute similarity features for all candidate pairs."""
        if self._pairs is None:
            raise RuntimeError("Must call block() before compute_features()")
        self._features = self._feature_computer.compute(self._pairs)
        return self

    def fit(self, labels_df: pd.DataFrame, model: str = "logistic", **model_kwargs) -> "Linker":
        """Train supervised model on labelled pairs.

        Args:
            labels_df: DataFrame with columns [left_idx, right_idx, label]
                       where label is 1 (match) or 0 (non-match).
            model: Model type. One of "logistic", "random_forest",
                   "gradient_boosting", "xgboost".
            **model_kwargs: Additional keyword arguments passed to the model constructor.
        """
        if self._features is None:
            raise RuntimeError("Must call compute_features() before fit()")

        if model not in MODEL_REGISTRY:
            raise ValueError(
                f"Unknown model '{model}'. Choose from: {list(MODEL_REGISTRY.keys())}"
            )

        pairs_with_idx = self._pairs.copy()
        pairs_with_idx["_pair_idx"] = range(len(pairs_with_idx))

        merged = labels_df.merge(
            pairs_with_idx, on=["left_idx", "right_idx"], how="inner"
        )

        if len(merged) == 0:
            raise ValueError(
                "No labelled pairs found in candidate pairs. "
                "Ensure labels_df pairs exist in blocking output."
            )

        indices = merged["_pair_idx"].values
        labelled_features = self._features[indices]
        labels = merged["label"].values.astype(np.float64)

        model_cls = MODEL_REGISTRY[model]
        self._model = model_cls(**model_kwargs)
        self._model.fit(labelled_features, labels)
        self._fs_model = None  # clear any previous FS model
        return self

    def estimate(self, model: str = "fellegi_sunter", **kwargs) -> "Linker":
        """Estimate unsupervised model (no labels needed).

        This is the alternative to fit() — it uses unsupervised methods to
        estimate match probabilities from the feature distribution alone.

        Args:
            model: Model type. One of "fellegi_sunter", "kmeans", "gmm", "dbscan".
            **kwargs: Additional keyword arguments passed to the model constructor.
        """
        if self._features is None:
            raise RuntimeError("Must call compute_features() before estimate()")

        if model == "fellegi_sunter":
            self._fs_model = FellegiSunterModel(**kwargs)
            self._fs_model.estimate(self._features)
        elif model in UNSUPERVISED_REGISTRY:
            model_cls = UNSUPERVISED_REGISTRY[model]
            self._fs_model = model_cls(**kwargs)
            self._fs_model.estimate(self._features)
        else:
            valid = ["fellegi_sunter"] + list(UNSUPERVISED_REGISTRY.keys())
            raise ValueError(
                f"Unknown model '{model}'. Choose from: {valid}"
            )

        self._model = None  # clear any previous supervised model
        return self

    def predict(self, threshold: Optional[float] = None) -> pd.DataFrame:
        """Score all candidate pairs.

        Uses whichever model was last configured via fit() or estimate().

        Args:
            threshold: If provided, only return pairs with score >= threshold.

        Returns:
            DataFrame with columns [left_idx, right_idx, score].
        """
        if self._features is None:
            raise RuntimeError("Must call compute_features() before predict()")

        if self._model is not None and self._model.is_fitted:
            scores = self._model.predict(self._features)
        elif self._fs_model is not None and self._fs_model.is_fitted:
            scores = self._fs_model.predict(self._features)
        else:
            raise RuntimeError("Must call fit() or estimate() before predict()")

        result = self._pairs.copy()
        result["score"] = scores

        if threshold is not None:
            result = result[result["score"] >= threshold].reset_index(drop=True)

        return result

    def predict_iter(
        self, chunk_size: Optional[int] = None
    ) -> Iterator[pd.DataFrame]:
        """Streaming prediction for large-scale mode."""
        if self._pairs is None or self._features is None:
            raise RuntimeError(
                "Must call block() and compute_features() before predict_iter()"
            )

        model = self._model or self._fs_model
        if model is None or not model.is_fitted:
            raise RuntimeError("Must call fit() or estimate() before predict_iter()")

        cs = chunk_size or self.config.chunk_size
        n = len(self._pairs)

        for start in range(0, n, cs):
            end = min(start + cs, n)
            chunk_pairs = self._pairs.iloc[start:end].copy()
            chunk_features = self._features[start:end]
            chunk_pairs["score"] = model.predict(chunk_features)
            yield chunk_pairs.reset_index(drop=True)
