"""
Probability Calibration Module.

Provides calibration methods for probability estimates and reliability diagrams.
Supports both multi-output models (Long/Short) and custom fitting.
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.isotonic import IsotonicRegression
from sklearn.linear_model import LogisticRegression
from utils.logger import logger


class MultiOutputCalibrator:
    """
    Обертка для калибровки вероятностей multi-output моделей (Long/Short).

    Supports:
    - Isotonic Regression (non-parametric)
    - Platt Scaling (Logistic Regression)

    Args:
        method: Calibration method ('isotonic' or 'platt')

    Example:
        >>> calibrator = MultiOutputCalibrator(method='isotonic')
        >>> calibrator.fit(raw_probs, y_true)
        >>> calibrated_probs = calibrator.predict_proba(raw_probs)
    """

    def __init__(self, method: str = "isotonic"):
        """
        Инициализация калибратора.

        Args:
            method: Метод калибровки ('isotonic' или 'platt')
        """
        self.method = method
        self.calibrators = []
        self.is_fitted = False

    def fit(self, raw_probs: list, y_true: np.ndarray):
        """
        Обучает калибраторы для каждого выхода.

        Args:
            raw_probs: list of arrays (вероятности классов для каждого выхода)
                       shape: [(n_samples, 2), (n_samples, 2), ...]
            y_true: массив (n_samples, n_targets)

        Returns:
            self: Возвращает себя для удобства (fluent interface)
        """
        if not isinstance(raw_probs, list):
            raise ValueError("raw_probs must be a list of arrays")

        n_targets = len(raw_probs)
        if n_targets != y_true.shape[1]:
            raise ValueError(
                f"Mismatch: raw_probs has {n_targets} arrays, y_true has {y_true.shape[1]} columns"
            )

        self.calibrators = []

        for i in range(n_targets):
            p_raw = raw_probs[i][:, 1]  # Берем вероятность класса "1"
            y_t = y_true[:, i]

            if self.method == "isotonic":
                cal = IsotonicRegression(out_of_bounds="clip", y_min=0, y_max=1)
                cal.fit(p_raw, y_t)
            elif self.method == "platt":
                cal = LogisticRegression(C=1e10, solver="lbfgs", max_iter=1000)
                # Platt требует 2D массив
                cal.fit(p_raw.reshape(-1, 1), y_t)
            else:
                raise ValueError(f"Method must be 'isotonic' or 'platt', got '{self.method}'")

            self.calibrators.append(cal)

        self.is_fitted = True
        logger.info(f"Calibrated {n_targets} targets using {self.method}.")
        return self

    def predict_proba(self, raw_probs: list) -> list:
        """
        Применяет калибровку к сырым вероятностям.

        Args:
            raw_probs: list of arrays (вероятности классов)

        Returns:
            list of arrays: Калиброванные вероятности (n_samples, 2)
        """
        if not self.is_fitted:
            raise RuntimeError("Calibrator not fitted yet. Call fit() first.")

        calibrated = []
        for i, cal in enumerate(self.calibrators):
            p_raw = raw_probs[i][:, 1]

            if self.method == "isotonic":
                p1 = cal.predict(p_raw)
            else:  # platt
                p1 = cal.predict_proba(p_raw.reshape(-1, 1))[:, 1]

            p1 = np.clip(p1, 0, 1)

            # Возвращаем в формате sklearn (N, 2)
            calibrated.append(np.c_[1 - p1, p1])

        return calibrated

    def predict(self, raw_probs: list) -> np.ndarray:
        """
        Предсказывает классы с калиброванными вероятностями.

        Args:
            raw_probs: list of arrays (вероятности классов)

        Returns:
            np.ndarray: Предсказанные классы (n_samples, n_targets)
        """
        calibrated_probs = self.predict_proba(raw_probs)
        return np.column_stack([p[:, 1] >= 0.5 for p in calibrated_probs])

    def get_calibrator(self, target_idx: int):
        """
        Получает конкретный калибратор по индексу.

        Args:
            target_idx: Индекс таргета (0 = Long, 1 = Short)

        Returns:
            Калибратор sklearn
        """
        if target_idx >= len(self.calibrators):
            raise IndexError(f"Target index {target_idx} out of range")
        return self.calibrators[target_idx]


class CalibratedModel:
    """
    Обертка для модели с применением калибратора.

    Проблема:
        Модель.predict_proba() возвращает "сырые" вероятности, которые могут
        быть смещены (недооценка или переоценка вероятностей).

    Решение:
        Используем Calibrator для приведения вероятностей в корректный диапазон.

    Args:
        model: Обученная модель (RF, XGBoost, LightGBM)
        calibrator: MultiOutputCalibrator или None

    Example:
        >>> calibrator = MultiOutputCalibrator(method='isotonic')
        >>> calibrator.fit(raw_probs, y_true)
        >>> model = CalibratedModel(model, calibrator)
        >>> probs = model.predict_proba(X)
    """

    def __init__(self, model, calibrator=None):
        """
        Инициализация.

        Args:
            model: Обученная модель
            calibrator: MultiOutputCalibrator (по умолчанию None)
        """
        self.model = model
        self.calibrator = calibrator

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Предсказывает классы.

        Args:
            X: Feature matrix (n_samples, n_features)

        Returns:
            np.ndarray: Предсказанные классы (n_samples, n_targets)
        """
        if self.calibrator is not None:
            raw_probs = self.model.predict_proba(X)
            return self.calibrator.predict(raw_probs)
        else:
            return self.model.predict(X)

    def predict_proba(self, X: np.ndarray) -> list:
        """
        Предсказывает калиброванные вероятности.

        Args:
            X: Feature matrix (n_samples, n_features)

        Returns:
            list of arrays: Калиброванные вероятности (n_samples, 2)
        """
        if self.calibrator is not None:
            raw_probs = self.model.predict_proba(X)
            if not isinstance(raw_probs, list):
                raw_probs = [raw_probs]
            return self.calibrator.predict_proba(raw_probs)
        else:
            raw_probs = self.model.predict_proba(X)
            if not isinstance(raw_probs, list):
                return [raw_probs]
            return raw_probs


def plot_reliability_diagram(
    raw_probs: list, cal_probs: list, y_true: np.ndarray, save_path: str = None
):
    """
    Строит график надежности (Reliability Diagram) для сравнения сырых и откалиброванных вероятностей.

    График показывает:
    - Как хорошо модель предсказывает вероятность (x-axis)
    - Какая доля фактических позитивов (y-axis)
    - Линия безошибочной калибровки (x=y)
    - Красная зона: 0.55 - 0.60 (высокая уверенность)

    Args:
        raw_probs: Сырые вероятности (list of arrays, shape [(n_samples, 2), (n_samples, 2)])
        cal_probs: Калиброванные вероятности (list of arrays)
        y_true: True labels (n_samples, n_targets)
        save_path: Путь для сохранения графика (необязательно)

    Example:
        >>> from models.calibration import plot_reliability_diagram
        >>> plot_reliability_diagram(raw_probs, cal_probs, y_test, save_path="plot.png")
    """
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    target_names = ["Long", "Short"]

    for i in range(2):
        ax = axes[i]
        y_t = y_true[:, i]

        # Данные для сырых и калиброванных вероятностей
        prob_true_raw, prob_pred_raw = _calibration_curve(y_t, raw_probs[i][:, 1])
        prob_true_cal, prob_pred_cal = _calibration_curve(y_t, cal_probs[i][:, 1])

        ax.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated")
        ax.plot(prob_pred_raw, prob_true_raw, "s-", label="Raw model")
        ax.plot(prob_pred_cal, prob_true_cal, "o-", label="Calibrated")

        # Выделяем зону 0.55 - 0.60
        ax.axvspan(0.55, 0.60, color="red", alpha=0.1, label="High-conf zone (0.55-0.60)")

        ax.set_title(f"Reliability Diagram: {target_names[i]}")
        ax.set_xlabel("Mean predicted probability")
        ax.set_ylabel("Fraction of positives")
        ax.legend(loc="upper left")
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        logger.info(f"Reliability diagram saved to {save_path}")
    else:
        plt.show()

    return fig, axes


def plot_calibration_comparison(
    raw_probs: list,
    cal_probs: list,
    y_true: np.ndarray,
    target_names: list[str] = ["Long", "Short"],
    save_path: str = None,
):
    """
    Альтернативная визуализация: гистограмма распределения вероятностей.

    Сравнивает сырые и калиброванные вероятности, показывая
    как Isotonic Regression сдвигает распределение.

    Args:
        raw_probs: Сырые вероятности
        cal_probs: Калиброванные вероятности
        y_true: True labels
        target_names: Названия таргетов
        save_path: Путь для сохранения

    Example:
        >>> plot_calibration_comparison(raw_probs, cal_probs, y_test)
    """
    fig, axes = plt.subplots(2, 1, figsize=(12, 8))

    for i, name in enumerate(target_names):
        ax = axes[i]
        y_t = y_true[:, i]

        # Сырые вероятности
        ax.hist(
            raw_probs[i][:, 1][y_t == 0],
            bins=20,
            alpha=0.5,
            label="Raw (negatives)",
            color="blue",
            edgecolor="black",
        )
        ax.hist(
            raw_probs[i][:, 1][y_t == 1],
            bins=20,
            alpha=0.5,
            label="Raw (positives)",
            color="red",
            edgecolor="black",
        )

        # Калиброванные вероятности
        ax.hist(
            cal_probs[i][:, 1][y_t == 0],
            bins=20,
            alpha=0.5,
            label="Calibrated (negatives)",
            color="cyan",
            edgecolor="black",
        )
        ax.hist(
            cal_probs[i][:, 1][y_t == 1],
            bins=20,
            alpha=0.5,
            label="Calibrated (positives)",
            color="orange",
            edgecolor="black",
        )

        ax.set_title(f"Probability Distribution: {name}")
        ax.set_xlabel("Probability")
        ax.set_ylabel("Count")
        ax.legend()
        ax.grid(alpha=0.3)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        logger.info(f"Calibration comparison saved to {save_path}")
    else:
        plt.show()

    return fig, axes


def _calibration_curve(y_true: np.ndarray, y_prob: np.ndarray, n_bins: int = 10):
    """
    Кастомная реализация калибровочной кривой.

    Args:
        y_true: True labels
        y_prob: Predicted probabilities
        n_bins: Количество бинов

    Returns:
        (prob_true, prob_pred): Данные для построения графика
    """
    bins = np.linspace(0, 1, n_bins + 1)
    bin_centers = (bins[:-1] + bins[1:]) / 2

    prob_true = np.zeros(n_bins)
    prob_pred = np.zeros(n_bins)

    for i, (low, high) in enumerate(zip(bins[:-1], bins[1:])):
        mask = (y_prob >= low) & (y_prob < high)
        if mask.sum() > 0:
            prob_true[i] = y_true[mask].mean()
            prob_pred[i] = y_prob[mask].mean()
        else:
            prob_true[i] = np.nan
            prob_pred[i] = np.nan

    # Убираем пустые бины
    valid = ~np.isnan(prob_true)
    return prob_true[valid], prob_pred[valid]


def calibration_score(raw_probs: list, y_true: np.ndarray) -> float:
    """
    Вычисляет Brier Score (MSE между предсказанными и истинными вероятностями).

    Args:
        raw_probs: Сырые вероятности
        y_true: True labels

    Returns:
        Brier Score (чем меньше, тем лучше)
    """
    brier_score = 0
    total = 0

    for i, (p_raw, y_t) in enumerate(zip(raw_probs, y_true.T)):
        brier_score += np.mean((p_raw[:, 1] - y_t) ** 2)
        total += 1

    return brier_score / total if total > 0 else 0