"""
Adapter для ExpertEnsemble, чтобы интегрировать его в BenchmarkFramework.

ExpertEnsemble требует OHLCV данные, а BenchmarkFramework работает только с фичами.
Адаптер "оборачивает" ExpertEnsemble и имитирует sklearn-интерфейс fit/predict_proba.
"""

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.multioutput import MultiOutputClassifier
from config import MODEL_CONFIG, RANDOM_SEED
from utils.logger import logger
from models.experts import ExpertEnsemble
from models.moe import _build_dataset, _get_feature_cols


class ExpertEnsembleAdapter:
    """
    Адаптер для интеграции ExpertEnsemble в BenchmarkFramework.

    Оборачивает ExpertEnsemble + Fusion RF в sklearn-подобный интерфейс.

    Args:
        feature_cols: Список базовых фичей для ExpertEnsemble

    Example:
        >>> adapter = ExpertEnsembleAdapter(feature_cols=base_features)
        >>> adapter.fit(df, y)
        >>> probs = adapter.predict_proba(df)
    """

    def __init__(self, feature_cols: list[str]):
        """
        Инициализация адаптера.

        Args:
            feature_cols: Список базовых фичей (без OHLCV)
        """
        self.feature_cols = feature_cols
        self.ensemble = None
        self.fusion_model = None
        self.fusion_cols = None
        self.is_fitted = False

    def fit(self, df: pd.DataFrame, y: np.ndarray):
        """
        Обучает LSTM экспертов и Fusion RF.

        Args:
            df: DataFrame с OHLCV данными и фичами (все колонки)
            y: Target vector/matrix (n_samples, n_targets) или (n_samples,)

        Returns:
            self: Для удобства (fluent interface)
        """
        logger.info(f"  [Adapter] Training {len(self.feature_cols)} base features...")

        # 1. Инициализация и обучение LSTM экспертов
        # Устанавливаем seed для воспроизводимости
        np.random.seed(RANDOM_SEED)

        self.ensemble = ExpertEnsemble(len(self.feature_cols))
        self.ensemble.train_all(df, self.feature_cols, verbose=False)

        # 2. Получаем сигналы экспертов
        logger.info(f"  [Adapter] Getting expert signals...")
        signals = self.ensemble.predict_all(df, self.feature_cols)

        # 3. Создаем фичи для Fusion RF
        self.fusion_cols = [c for c in _get_feature_cols() if c in df.columns]
        logger.info(f"  [Adapter] Building fusion features from {len(self.fusion_cols)} columns...")

        X_fusion = _build_dataset(df, signals, self.ensemble, self.fusion_cols)

        # 4. Выравнивание длины (из-за WINDOW=20 длина может быть меньше y)
        min_len = min(len(X_fusion), len(y))
        X_fusion = X_fusion[-min_len:]
        y_aligned = y[-min_len:]

        logger.info(f"  [Adapter] Training Fusion RF on {len(X_fusion)} samples...")
        logger.info(f"    Fusion features: {len(self.fusion_cols)}")
        logger.info(f"    X_fusion shape: {X_fusion.shape}, y shape: {y_aligned.shape}")

        # 5. Обучаем RF на вершинах экспертов
        self.fusion_model = MultiOutputClassifier(
            RandomForestClassifier(**MODEL_CONFIG['rf'])
        )
        self.fusion_model.fit(X_fusion, y_aligned)

        self.is_fitted = True
        logger.info(f"  [Adapter] Fusion RF trained successfully!")
        return self

    def predict_proba(self, df: pd.DataFrame):
        """
        Предсказывает вероятности через экспертов и Fusion RF.

        Args:
            df: DataFrame с OHLCV данными и фичами

        Returns:
            list of arrays: Вероятности для каждого таргета (shape: [(n_samples, 2), ...])
        """
        if not self.is_fitted:
            raise RuntimeError(
                "Model not fitted yet. Call fit() first on DataFrame with OHLCV data."
            )

        # 1. Сигналы экспертов
        signals = self.ensemble.predict_all(df, self.feature_cols)

        # 2. Создаем фичи для Fusion RF
        X_fusion = _build_dataset(df, signals, self.ensemble, self.fusion_cols)

        # 3. Предсказание RF
        probs = self.fusion_model.predict_proba(X_fusion)

        # 4. Выравнивание длины
        # Из-за WINDOW=20 первые n строк будут без предсказаний
        # Заполняем их нулями, чтобы длина совпадала с исходным df
        pad_len = len(df) - len(probs[0])

        if pad_len > 0:
            logger.debug(f"  [Adapter] Padding {pad_len} predictions with neutral probs (0.5)")
            padded_probs = []
            for p in probs:
                # p имеет форму (len, 2). Делаем паддинг нулями в начале
                pad = np.full((pad_len, 2), 0.5)  # Нейтральная вероятность 0.5
                padded_probs.append(np.vstack([pad, p]))
            return padded_probs

        return probs

    def predict(self, X: np.ndarray):
        """
        Предсказывает классы (доступно для совместимости).

        Args:
            X: Feature matrix (n_samples, n_features)

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

    def is_fitted(self) -> bool:
        """
        Проверяет, обучена ли модель.

        Returns:
            bool: True если модель обучена
        """
        return self.is_fitted

    def get_fusion_feature_names(self) -> list[str]:
        """
        Получает список имен фичей для Fusion RF.

        Returns:
            list[str]: Имена фичей
        """
        return self.fusion_cols


# Для backward compatibility с существующим кодом
class ExpertEnsembleCalibrator(ExpertEnsembleAdapter):
    """
    Вариация адаптера с калибровкой.

    Дополнительно применяет калибровку к вероятностям.
    """

    def fit(self, df: pd.DataFrame, y: np.ndarray):
        # Сначала обучаем базовый адаптер
        super().fit(df, y)

        # Применяем калибровку
        from models.calibration import MultiOutputCalibrator

        # Получаем вероятности и обучаем калибратор
        raw_probs = []
        y_true = []

        # Генерируем вероятности для оценки калибратора
        temp_df = df.copy()
        temp_df[self.feature_cols] = df[self.feature_cols]

        try:
            probs_list = self.predict_proba(df)
            for probs in probs_list:
                raw_probs.append(probs[:, 1])
                y_true.append(y[:, 0] if len(y.shape) > 1 else y)
        except:
            pass

        if raw_probs:
            logger.info(f"  [Adapter] Calibrating expert probabilities...")
            calibrator = MultiOutputCalibrator(method="isotonic")
            calibrator.fit(raw_probs, y_true)
            self.calibrator = calibrator
            logger.info(f"  [Adapter] Probabilities calibrated!")

        return self