# src/core/walk_forward.py
import pandas as pd
from typing import List, Dict, Optional, Callable, Any
from src.core.optimizer import ParameterOptimizer
from src.core.engine import Engine
from src.risk.sizing.base import PositionSizer
from src.risk.stops.base import StopLoss
from ..utils.logger import log

class WalkForwardValidator:
    """
    Orquestador de walk-forward validation.

    Su responsabilidad es:
    - dividir el tiempo en ventanas TRAIN / TEST
    - coordinar optimización y backtest (más adelante)
    """

    def __init__(
        self,
        strategy_class,
        param_grid: Optional[dict],
        data: pd.DataFrame,
        train_window: pd.Timedelta,
        test_window: pd.Timedelta,
        step_size: Optional[pd.Timedelta] = None,
        initial_capital: float = 10_000,
        commission: float = 0.001,
        slippage: float = 0.0005,
        position_size: float = 0.95,
        optimizer_metric: str = "sharpe_ratio",
        verbose: bool = True,
        position_sizer: Optional[PositionSizer] = None,
        stop_loss: Optional[StopLoss] = None,
        max_combinations: Optional[int] = None,
        progress_callback: Optional[callable] = None,
        fixed_params: Optional[dict] = None,
    ):
        self.strategy_class = strategy_class
        self.param_grid = param_grid
        self.fixed_params = fixed_params

        self.data = data.sort_index()

        self.train_window = train_window
        self.test_window = test_window
        self.step_size = step_size or test_window

        self.initial_capital = initial_capital
        self.commission = commission
        self.slippage = slippage
        self.position_size = position_size
        self.optimizer_metric = optimizer_metric
        self.verbose = verbose
        self.position_sizer = position_sizer
        self.stop_loss = stop_loss
        self.max_combinations = max_combinations

        self.progress_callback = progress_callback

        # Validaciones básicas
        if not isinstance(self.data.index, pd.DatetimeIndex):
            raise ValueError("data debe tener un DatetimeIndex")

        if not self.data.index.is_monotonic_increasing:
            raise ValueError("data.index debe estar ordenado cronológicamente")
        
        # ✅ Validación de modo (NUEVO, mínimo y claro)
        if self.param_grid is not None and self.fixed_params is not None:
            raise ValueError("WalkForwardValidator: usa param_grid (optimization) o fixed_params (validation), no ambos.")

        if self.param_grid is None and self.fixed_params is None:
            raise ValueError("WalkForwardValidator: debes pasar param_grid o fixed_params.")

    # ------------------------------------------------------------------
    # 🔹 GENERACIÓN DE VENTANAS TEMPORALES
    # ------------------------------------------------------------------

    def _generate_windows(self) -> List[Dict]:
        """
        Genera ventanas temporales TRAIN / TEST para walk-forward validation.

        Returns:
            Lista de diccionarios con:
            - train_start
            - train_end
            - test_start
            - test_end
        """

        windows: List[Dict] = []

        data_start = self.data.index.min()
        data_end = self.data.index.max()

        train_start = data_start

        window_id = 1

        while True:
            train_end = train_start + self.train_window
            test_start = train_end
            test_end = test_start + self.test_window

            # 🚨 Condición de parada: no hay datos suficientes para TEST
            if test_end > data_end:
                if self.verbose:
                    log.info(
                        f"⛔ Walk-forward detenido: "
                        f"test_end ({test_end}) > data_end ({data_end})"
                    )
                break

            # Extraer subconjuntos (solo para validación de tamaño)
            train_data = self.data.loc[train_start:train_end]
            test_data = self.data.loc[test_start:test_end]

            # Validaciones mínimas
            if len(train_data) < 10:
                log.warning(
                    f"⚠️ Ventana {window_id} ignorada: "
                    f"muy pocos datos en TRAIN ({len(train_data)})"
                )
                train_start += self.step_size
                continue

            if len(test_data) < 5:
                log.warning(
                    f"⚠️ Ventana {window_id} ignorada: "
                    f"muy pocos datos en TEST ({len(test_data)})"
                )
                train_start += self.step_size
                continue

            window = {
                "window_id": window_id,
                "train_start": train_start,
                "train_end": train_end,
                "test_start": test_start,
                "test_end": test_end,
            }

            windows.append(window)

            if self.verbose:
                log.info(
                    f"🪟 WF #{window_id} | "
                    f"TRAIN: {train_start.date()} → {train_end.date()} | "
                    f"TEST: {test_start.date()} → {test_end.date()}"
                )

            # Avanzar ventana
            train_start += self.step_size
            window_id += 1

        if not windows:
            raise RuntimeError("No se generó ninguna ventana válida de walk-forward")

        return windows
    
    # ------------------------------------------------------------------
    # 🔹 EJECUCIÓN WF
    # ------------------------------------------------------------------

    def run(self):
        """
        Ejecuta el walk-forward completo:
        - Optimiza en TRAIN
        - Evalúa en TEST
        - Devuelve un objeto estructurado para análisis y visualización
        """
        self.windows = self._generate_windows()

        rows = []          # filas planas (para DataFrame)
        raw_results = []   # resultados completos (debug / drill-down)

        log.info(f"🚀 Iniciando Walk-Forward con {len(self.windows)} ventanas")

        for w in self.windows:
            wid = w["window_id"]
            
            if self.progress_callback:
                self.progress_callback(
                    window_id=wid,
                    total_windows=len(self.windows)
                )

            wid = w["window_id"]
            log.info(f"▶️ WF #{wid} en ejecución")

            train_data = self.data.loc[w["train_start"]:w["train_end"]]
            test_data = self.data.loc[w["test_start"]:w["test_end"]]

            if train_data.empty or test_data.empty:
                log.warning(f"WF #{wid} ignorado (datos insuficientes)")
                continue

            # 1️⃣ Optimización TRAIN
            best_train_metric = None

            if self.fixed_params is not None:
                # ✅ VALIDATION MODE: sin optimización
                best_params = self.fixed_params
            else:
                # ✅ OPTIMIZATION MODE
                optimizer = ParameterOptimizer(
                    strategy_class=self.strategy_class,
                    data=train_data,
                    initial_capital=self.initial_capital,
                    commission=self.commission,
                    slippage=self.slippage,
                    position_size=self.position_size,
                    position_sizer=self.position_sizer,
                    stop_loss=self.stop_loss,
                    max_combinations=self.max_combinations,
                )

                opt_df = optimizer.optimize(self.param_grid)

                if opt_df.empty:
                    log.warning(f"WF #{wid} sin resultados de optimización")
                    continue

                best_params = optimizer.get_best_params(metric=self.optimizer_metric)
                best_train_metric = opt_df[self.optimizer_metric].max()

            # 2️⃣ Backtest TEST (OOS)
            strategy = self.strategy_class(**best_params)

            engine = Engine(
                strategy=strategy,
                initial_capital=self.initial_capital,
                commission=self.commission,
                slippage=self.slippage,
                position_size=self.position_size,
                position_sizer=self.position_sizer,
                stop_loss=self.stop_loss,
            )

            test_results = engine.run(test_data)

            # 3️⃣ Fila plana (TEST only)
            rows.append({
                "window": wid,
                "train_start": w["train_start"],
                "train_end": w["train_end"],
                "test_start": w["test_start"],
                "test_end": w["test_end"],
                "return_pct": test_results["total_return_pct"],
                "sharpe": test_results["sharpe_ratio"],
                "max_dd_pct": test_results["max_drawdown_pct"],
                "trades": test_results["total_trades"],
                "params": best_params,
            })

            # 4️⃣ Resultado completo (opcional)
            raw_results.append({
                "window_id": wid,
                "best_params": best_params,
                "train_metric": best_train_metric,
                "test_results": test_results,
            })

            log.success(
                f"✅ WF #{wid} | "
                f"Return: {test_results['total_return_pct']:.2f}% | "
                f"Sharpe: {test_results['sharpe_ratio']:.2f}"
            )

        # 5️⃣ Resultado final
        wf_result = {
            "meta": {
                "strategy": self.strategy_class.__name__,
                "train_window": self.train_window,
                "test_window": self.test_window,
                "step_size": self.step_size,
                "optimizer_metric": self.optimizer_metric,
                "n_windows": len(rows),
                "data_start": self.data.index.min(),
                "data_end": self.data.index.max(),
                "mode": "validation" if self.fixed_params is not None else "optimization"
            },
            "windows": pd.DataFrame(rows),
            "raw_results": raw_results,
        }

        log.success("🏁 Walk-Forward completado")

        # Validaciones de integridad
        assert isinstance(wf_result, dict), "wf_result debe ser un dict"
        assert "windows" in wf_result, "wf_result debe contener 'windows'"
        assert "raw_results" in wf_result, "wf_result debe contener 'raw_results'"
        assert isinstance(wf_result["windows"], pd.DataFrame), "'windows' debe ser DataFrame"
        assert isinstance(wf_result["raw_results"], list), "'raw_results' debe ser list"
        assert all(isinstance(r, dict) for r in wf_result["raw_results"]), "raw_results corrupto"

        return wf_result
