from __future__ import annotations

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

import numpy as np

from blazematch._rust import (
    predict_batch,
    predict_forest_batch,
    predict_gbdt_batch,
    predict_svm_batch,
)


def _extract_sklearn_trees(sklearn_model, is_gbdt=False):
    """Extract tree structures into flat numpy arrays with offsets.

    Returns dict with keys: offsets, children_left, children_right, features,
    thresholds, values — all contiguous numpy arrays. Trees are concatenated
    into flat arrays with offsets[i]:offsets[i+1] delimiting each tree.
    This avoids repeated Python list → Rust Vec conversion on every predict call.
    """
    estimators = sklearn_model.estimators_
    # GradientBoosting: estimators_ is 2D (n_estimators, 1) for binary
    if is_gbdt:
        estimators = [e[0] for e in estimators]

    offsets = [0]
    all_cl = []
    all_cr = []
    all_feat = []
    all_thresh = []
    all_val = []

    for est in estimators:
        tree = est.tree_
        all_cl.append(tree.children_left)
        all_cr.append(tree.children_right)
        all_feat.append(tree.feature)
        all_thresh.append(tree.threshold)

        if is_gbdt:
            # GBDT: value shape (n_nodes, 1, 1) -> raw prediction
            values = tree.value.squeeze(-1).squeeze(-1)
        else:
            # RF: value shape (n_nodes, 1, n_classes) -> P(class=1)
            values = tree.value.squeeze(1)  # (n_nodes, n_classes)
            totals = values.sum(axis=1)
            totals[totals == 0] = 1
            values = values[:, 1] / totals  # P(class=1)

        all_val.append(values)
        offsets.append(offsets[-1] + tree.children_left.shape[0])

    return {
        "offsets": np.array(offsets, dtype=np.int64),
        "children_left": np.concatenate(all_cl).astype(np.int64),
        "children_right": np.concatenate(all_cr).astype(np.int64),
        "features": np.concatenate(all_feat).astype(np.int64),
        "thresholds": np.concatenate(all_thresh).astype(np.float64),
        "values": np.concatenate(all_val).astype(np.float64),
    }


def _migrate_tree_data(old_data):
    """Convert old list-of-lists tree_data to flat numpy format."""
    offsets = [0]
    for cl in old_data["children_left"]:
        offsets.append(offsets[-1] + len(cl))

    return {
        "offsets": np.array(offsets, dtype=np.int64),
        "children_left": np.concatenate(
            [np.asarray(x, dtype=np.int64) for x in old_data["children_left"]]
        ),
        "children_right": np.concatenate(
            [np.asarray(x, dtype=np.int64) for x in old_data["children_right"]]
        ),
        "features": np.concatenate(
            [np.asarray(x, dtype=np.int64) for x in old_data["features"]]
        ),
        "thresholds": np.concatenate(
            [np.asarray(x, dtype=np.float64) for x in old_data["thresholds"]]
        ),
        "values": np.concatenate(
            [np.asarray(x, dtype=np.float64) for x in old_data["values"]]
        ),
    }


class MatchModel:
    """Logistic regression model for record linkage scoring.

    Training uses scikit-learn (Python, small labelled set).
    Inference uses Rust-accelerated batch prediction.
    """

    def __init__(self):
        self.weights_: Optional[np.ndarray] = None
        self.bias_: Optional[float] = None
        self._sklearn_model = None

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

    def fit(self, features: np.ndarray, labels: np.ndarray) -> "MatchModel":
        """Train logistic regression on labelled pairs."""
        from sklearn.linear_model import LogisticRegression

        self._sklearn_model = LogisticRegression(max_iter=1000, solver="lbfgs")
        self._sklearn_model.fit(features, labels)

        self.weights_ = self._sklearn_model.coef_[0].astype(np.float64)
        self.bias_ = float(self._sklearn_model.intercept_[0])

        return self

    def set_weights(self, weights: np.ndarray, bias: float) -> "MatchModel":
        """Manually set model weights (e.g., from a saved model)."""
        self.weights_ = np.asarray(weights, dtype=np.float64)
        self.bias_ = float(bias)
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using Rust-accelerated inference."""
        if not self.is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")

        if not features.flags["C_CONTIGUOUS"] or features.dtype != np.float64:
            features = np.ascontiguousarray(features, dtype=np.float64)
        weights = np.ascontiguousarray(self.weights_, dtype=np.float64)
        return np.asarray(predict_batch(features, weights, self.bias_))

    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": "logistic", "weights": self.weights_, "bias": self.bias_}
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "MatchModel":
        """Load model from disk.

        Warning:
            Uses pickle internally. Only load files you trust.
        """
        with open(path, "rb") as f:
            data = pickle.load(f)
        model = cls()
        model.weights_ = data["weights"]
        model.bias_ = data["bias"]
        return model


class RandomForestModel:
    """Random forest model for record linkage scoring.

    Uses sklearn RandomForestClassifier for training.
    Inference is Rust-accelerated via parallel tree traversal with rayon.
    """

    def __init__(self, n_estimators: int = 100, **kwargs):
        self._n_estimators = n_estimators
        self._kwargs = kwargs
        self._model = None
        self._tree_data = None

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

    def fit(self, features: np.ndarray, labels: np.ndarray) -> "RandomForestModel":
        """Train random forest on labelled pairs."""
        from sklearn.ensemble import RandomForestClassifier

        self._model = RandomForestClassifier(
            n_estimators=self._n_estimators,
            random_state=42,
            **self._kwargs,
        )
        self._model.fit(features, labels)
        self._tree_data = _extract_sklearn_trees(self._model, is_gbdt=False)
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using Rust-accelerated tree traversal."""
        if not self.is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        if not features.flags["C_CONTIGUOUS"] or features.dtype != np.float64:
            features = np.ascontiguousarray(features, dtype=np.float64)
        td = self._tree_data
        return np.asarray(predict_forest_batch(
            features,
            td["offsets"],
            td["children_left"],
            td["children_right"],
            td["features"],
            td["thresholds"],
            td["values"],
        ))

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

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

    @classmethod
    def load(cls, path: Union[str, Path]) -> "RandomForestModel":
        with open(path, "rb") as f:
            data = pickle.load(f)
        obj = cls()
        obj._model = data["model"]
        tree_data = data["tree_data"]
        if "offsets" not in tree_data:
            tree_data = _migrate_tree_data(tree_data)
        obj._tree_data = tree_data
        return obj


class GradientBoostingModel:
    """Gradient boosting model for record linkage scoring.

    Uses sklearn GradientBoostingClassifier for training.
    Inference is Rust-accelerated via parallel tree traversal with rayon.
    """

    def __init__(self, n_estimators: int = 100, learning_rate: float = 0.1, **kwargs):
        self._n_estimators = n_estimators
        self._learning_rate = learning_rate
        self._kwargs = kwargs
        self._model = None
        self._tree_data = None
        self._init_value: Optional[float] = None

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

    def fit(self, features: np.ndarray, labels: np.ndarray) -> "GradientBoostingModel":
        """Train gradient boosting on labelled pairs."""
        from sklearn.ensemble import GradientBoostingClassifier

        self._model = GradientBoostingClassifier(
            n_estimators=self._n_estimators,
            learning_rate=self._learning_rate,
            random_state=42,
            **self._kwargs,
        )
        self._model.fit(features, labels)
        self._tree_data = _extract_sklearn_trees(self._model, is_gbdt=True)

        # Extract initial prediction (log-odds of positive class prior).
        # Use staged_predict to get the init value robustly, regardless
        # of what init estimator is used internally.
        try:
            prior = self._model.init_.class_prior_[1]
            self._init_value = float(np.log(prior / (1.0 - prior)))
        except AttributeError:
            # Fallback: compute from label distribution
            pos_rate = np.clip(labels.mean(), 1e-7, 1 - 1e-7)
            self._init_value = float(np.log(pos_rate / (1.0 - pos_rate)))
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using Rust-accelerated tree traversal."""
        if not self.is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        if not features.flags["C_CONTIGUOUS"] or features.dtype != np.float64:
            features = np.ascontiguousarray(features, dtype=np.float64)
        td = self._tree_data
        return np.asarray(predict_gbdt_batch(
            features,
            td["offsets"],
            td["children_left"],
            td["children_right"],
            td["features"],
            td["thresholds"],
            td["values"],
            self._learning_rate,
            self._init_value,
        ))

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

    def save(self, path: Union[str, Path]) -> None:
        data = {
            "type": "gradient_boosting",
            "model": self._model,
            "tree_data": self._tree_data,
            "init_value": self._init_value,
            "learning_rate": self._learning_rate,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "GradientBoostingModel":
        with open(path, "rb") as f:
            data = pickle.load(f)
        obj = cls()
        obj._model = data["model"]
        tree_data = data["tree_data"]
        if "offsets" not in tree_data:
            tree_data = _migrate_tree_data(tree_data)
        obj._tree_data = tree_data
        obj._init_value = data["init_value"]
        obj._learning_rate = data["learning_rate"]
        return obj


class SVMModel:
    """Support Vector Machine model for record linkage scoring.

    Uses sklearn SVC with RBF kernel for training. Features are internally
    standardized (SVM performance is sensitive to feature scale).
    Inference is Rust-accelerated via parallel RBF kernel computation with rayon.
    """

    def __init__(self, **kwargs):
        self._kwargs = kwargs
        self._support_vectors: Optional[np.ndarray] = None
        self._dual_coef: Optional[np.ndarray] = None
        self._intercept: Optional[float] = None
        self._gamma: Optional[float] = None
        self._prob_a: Optional[float] = None
        self._prob_b: Optional[float] = None
        self._mean: Optional[np.ndarray] = None
        self._scale: Optional[np.ndarray] = None
        self._sklearn_model = None

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

    def fit(self, features: np.ndarray, labels: np.ndarray) -> "SVMModel":
        """Train SVM on labelled pairs."""
        from sklearn.svm import SVC

        # Standardize features (critical for RBF kernel performance)
        self._mean = features.mean(axis=0)
        self._scale = features.std(axis=0)
        self._scale[self._scale == 0] = 1.0  # avoid division by zero
        features_scaled = (features - self._mean) / self._scale

        defaults = {"kernel": "rbf", "probability": True, "random_state": 42}
        defaults.update(self._kwargs)
        self._sklearn_model = SVC(**defaults)
        self._sklearn_model.fit(features_scaled, labels)

        self._support_vectors = self._sklearn_model.support_vectors_.astype(np.float64)
        self._dual_coef = self._sklearn_model.dual_coef_[0].astype(np.float64)
        self._intercept = float(self._sklearn_model.intercept_[0])
        # Compute gamma from the public parameter rather than private _gamma
        gamma_param = self._sklearn_model.gamma
        if gamma_param == "scale":
            n_feat = features_scaled.shape[1]
            self._gamma = float(1.0 / (n_feat * features_scaled.var()))
        elif gamma_param == "auto":
            self._gamma = float(1.0 / features_scaled.shape[1])
        else:
            self._gamma = float(gamma_param)
        # Platt scaling parameters for calibrated probabilities
        self._prob_a = float(self._sklearn_model.probA_[0])
        self._prob_b = float(self._sklearn_model.probB_[0])
        return self

    def _scale_features(self, features: np.ndarray) -> np.ndarray:
        """Apply the same standardization used during training."""
        return (features - self._mean) / self._scale

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using Rust-accelerated SVM inference."""
        if not self.is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        features = self._scale_features(features)
        if not features.flags["C_CONTIGUOUS"] or features.dtype != np.float64:
            features = np.ascontiguousarray(features, dtype=np.float64)
        sv = self._support_vectors
        if not sv.flags["C_CONTIGUOUS"] or sv.dtype != np.float64:
            sv = np.ascontiguousarray(sv, dtype=np.float64)
        dc = self._dual_coef
        if not dc.flags["C_CONTIGUOUS"] or dc.dtype != np.float64:
            dc = np.ascontiguousarray(dc, dtype=np.float64)
        return np.asarray(predict_svm_batch(
            features, sv, dc, self._intercept, self._gamma,
            self._prob_a, self._prob_b,
        ))

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

    def save(self, path: Union[str, Path]) -> None:
        data = {
            "type": "svm",
            "support_vectors": self._support_vectors,
            "dual_coef": self._dual_coef,
            "intercept": self._intercept,
            "gamma": self._gamma,
            "prob_a": self._prob_a,
            "prob_b": self._prob_b,
            "mean": self._mean,
            "scale": self._scale,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "SVMModel":
        with open(path, "rb") as f:
            data = pickle.load(f)
        obj = cls()
        obj._support_vectors = data["support_vectors"]
        obj._dual_coef = data["dual_coef"]
        obj._intercept = data["intercept"]
        obj._gamma = data["gamma"]
        obj._prob_a = data["prob_a"]
        obj._prob_b = data["prob_b"]
        obj._mean = data["mean"]
        obj._scale = data["scale"]
        return obj


class LightGBMModel:
    """LightGBM model for record linkage scoring.

    Requires optional lightgbm dependency. Faster training and lower
    memory than XGBoost due to histogram-based binning and leaf-wise
    growth. Uses LightGBM's native C++ inference engine.
    """

    def __init__(self, n_estimators: int = 100, learning_rate: float = 0.1, **kwargs):
        self._n_estimators = n_estimators
        self._learning_rate = learning_rate
        self._kwargs = kwargs
        self._model = None

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

    def fit(self, features: np.ndarray, labels: np.ndarray) -> "LightGBMModel":
        """Train LightGBM on labelled pairs."""
        try:
            from lightgbm import LGBMClassifier
        except ImportError:
            raise ImportError(
                "lightgbm is required for LightGBMModel. "
                "Install it with: pip install lightgbm"
            )

        self._model = LGBMClassifier(
            n_estimators=self._n_estimators,
            learning_rate=self._learning_rate,
            random_state=42,
            verbosity=-1,
            **self._kwargs,
        )
        self._model.fit(features, labels)
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using LightGBM's native inference."""
        if not self.is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        return self._model.predict_proba(features)[:, 1]

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

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

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


class XGBoostModel:
    """XGBoost model for record linkage scoring.

    Requires optional xgboost dependency. Industry standard for
    record linkage competitions. Uses XGBoost's native C++ inference
    engine (already highly optimized — no Rust wrapper needed).
    """

    def __init__(self, n_estimators: int = 100, learning_rate: float = 0.1, **kwargs):
        self._n_estimators = n_estimators
        self._learning_rate = learning_rate
        self._kwargs = kwargs
        self._model = None

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

    def fit(self, features: np.ndarray, labels: np.ndarray) -> "XGBoostModel":
        """Train XGBoost on labelled pairs."""
        try:
            from xgboost import XGBClassifier
        except ImportError:
            raise ImportError(
                "xgboost is required for XGBoostModel. "
                "Install it with: pip install xgboost"
            )

        self._model = XGBClassifier(
            n_estimators=self._n_estimators,
            learning_rate=self._learning_rate,
            random_state=42,
            eval_metric="logloss",
            **self._kwargs,
        )
        self._model.fit(features, labels)
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Score pairs using predict_proba."""
        if not self.is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        return self._model.predict_proba(features)[:, 1]

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

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

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


# Model registry for Linker.fit() convenience
MODEL_REGISTRY = {
    "logistic": MatchModel,
    "random_forest": RandomForestModel,
    "gradient_boosting": GradientBoostingModel,
    "svm": SVMModel,
    "xgboost": XGBoostModel,
    "lightgbm": LightGBMModel,
}
