import random
from typing import Dict, List
import pandas as pd
from datetime import date, timedelta

from app.models import Strategy, Scenario, SimulationRun, RobustnessResult, SimulationStatus

from app.data import data_service

class SimulationEngine:
    def __init__(self) -> None:
        pass

    def run_simulation(self, run: SimulationRun, strategy: Strategy, scenario: Scenario) -> SimulationRun:
        """
        Main entry point for running a simulation.
        """
        print(f"Starting simulation {run.id} for strategy {strategy.name} in scenario {scenario.name}")
        
        # 1. Fetch Data (Real)
        # Extract tickers from strategy
        tickers = [pos.ticker for pos in strategy.allocation]
        start = scenario.start_date or (date.today() - timedelta(days=365*5))
        end = scenario.end_date or date.today()
        
        data = data_service.fetch_data(tickers, start, end)
        
        if data.empty:
            print("No data found for simulation.")
            run.status = SimulationStatus.FAILED
            return run
        
        # 2. Apply Stress/Scenario Logic (Data Level)
        # Apply shocks to the underlying assets before portfolio construction
        # This allows testing correlation breakdowns (diversification failure).
        stressed_data = self._apply_scenario_stress(data, scenario)

        # 3. Apply Strategy
        metrics_returns = self._apply_strategy(stressed_data, strategy)
        
        # 4. Calculate Metrics & Regime Analysis
        metrics = self._calculate_metrics(metrics_returns)
        regime_metrics = self._detect_regimes(metrics_returns)
        metrics.update(regime_metrics)
        
        # 5. Update Run Result
        run.robustness_score = metrics.get("robustness_score", 0.0)
        run.max_drawdown = metrics.get("max_drawdown", 0.0)
        run.sharpe_ratio = metrics.get("sharpe_ratio", 0.0)
        run.status = SimulationStatus.COMPLETED
        
        # Store full result (mock logic for artifact storage)
        print(f"Simulation Metrics: {metrics}")
        
        return run



    def _detect_regimes(self, returns: pd.Series) -> Dict[str, float]:
        """
        Simple Regime Slicing:
        1. Calculate Rolling Volatility (21-day)
        2. Calculate Trend (Price vs 200-day SMA - approximated from returns)
        
        For MVP, we will use a simplified approach on the Portfolio Returns themself
        to classify 'Behavioral Regimes' of the strategy.
        
        Regimes:
        - Low Vol / Bull
        - High Vol / Bull
        - Low Vol / Bear
        - High Vol / Bear
        """
        # Rolling annualized vol
        rolling_vol = returns.rolling(window=21).std() * (252**0.5)
        
        # Cumulative return (Price Proxy)
        price_proxy = (1 + returns).cumprod()
        sma_200 = price_proxy.rolling(window=200).mean()
        
        # Trend: 1 if Price > SMA, -1 if Price < SMA
        trend = (price_proxy > sma_200).astype(int).replace(0, -1)
        
        # Vol Regime: 1 if Vol > Median, -1 if Vol < Median
        median_vol = rolling_vol.median()
        vol_regime = (rolling_vol > median_vol).astype(int).replace(0, -1)
        
        # Combine
        # 1 = Bull/HighVol, 2 = Bull/LowVol, etc.
        # This is a bit complex to return as single metric map.
        # Instead, let's calculate performance IN each regime.
        
        regime_performance = {}
        
        # Masking
        # Bull Market
        bull_mask = trend == 1
        bear_mask = trend == -1
        
        if bull_mask.any():
            regime_performance["Sharpe_Bull"] = self._calculate_sharpe(returns[bull_mask])
        else:
            regime_performance["Sharpe_Bull"] = 0.0
            
        if bear_mask.any():
            regime_performance["Sharpe_Bear"] = self._calculate_sharpe(returns[bear_mask])
        else:
            regime_performance["Sharpe_Bear"] = 0.0

        return regime_performance

    def _calculate_sharpe(self, series: pd.Series) -> float:
        if series.std() == 0:
            return 0.0
        return (series.mean() * 252) / (series.std() * (252**0.5))

    def _apply_scenario_stress(self, data: pd.DataFrame, scenario: Scenario) -> pd.DataFrame:
        """
        Apply scenario transformation to asset returns.
        """
        # Copy to avoid mutating original cache if we had one
        df = data.copy()
        params = scenario.parameters or {}
        
        if scenario.scenario_type == "regime_stress":
            
            # 1. Simple Crash (Legacy support)
            if "Crash" in scenario.name or params.get("shock") == "crash":
                 # Daily drag: -0.1% daily is ~-22% annualized 
                 return df - 0.001 

            # 2. Volatility Spike
            # Param: "vol_mult" (e.g. 1.5x)
            if params.get("shock") == "volatility":
                mult = float(params.get("vol_mult", 1.5))
                return df * mult
                
            # 3. Correlation Breakdown
            # Logic: Collapse all assets towards the cross-sectional mean (Market Mode)
            # Param: "corr_factor" (0.0 to 1.0). 1.0 = Perfect sequence.
            if params.get("shock") == "correlation":
                factor = float(params.get("corr_factor", 0.8))
                
                # Calculate equal-weighted market index of the current assets
                market_index = df.mean(axis=1)
                
                # Blend individual asset with market index
                # New = (1-F)*Old + F*Market
                for col in df.columns:
                    df[col] = (1 - factor) * df[col] + factor * market_index
                
                return df
                
        return df

    def _calculate_metrics(self, returns: pd.Series) -> Dict[str, float]:
        """
        Calculate key robustness metrics.
        """
        if returns.empty:
             return {}
             
        total_return = (1 + returns).prod() - 1
        vol = returns.std() * (252 ** 0.5)
        sharpe = (returns.mean() * 252) / (returns.std() * (252 ** 0.5)) if vol != 0 else 0.0
        
        # Max Drawdown
        cum_returns = (1 + returns).cumprod()
        peak = cum_returns.cummax()
        drawdown = (cum_returns - peak) / peak
        max_dd = drawdown.min() if not drawdown.empty else 0.0
        
        # Robustness Score (0-100)
        # Higher Sharpe & Lower DD = Higher Score
        # Normalize Sharpe (0 to 3) -> 0 to 60
        # Normalize DD (0 to -0.5) -> 40 to 0
        score_sharpe = min(max(sharpe * 20, 0), 60)
        score_dd = min(max((1 + max_dd) * 40, 0), 40)
        score = score_sharpe + score_dd
        
        return {
            "robustness_score": score,
            "max_drawdown": max_dd,
            "sharpe_ratio": sharpe,
            "total_return": total_return
        }

    def _apply_strategy(self, data: pd.DataFrame, strategy: Strategy) -> pd.Series:
        """
        Calculate portfolio returns based on strategy.
        """
        weights = {pos.ticker: pos.weight for pos in strategy.allocation}
        
        # Align data and weights
        portfolio_returns = pd.Series(0.0, index=data.index)
        
        for ticker, weight in weights.items():
            if ticker in data.columns:
                portfolio_returns += data[ticker] * weight
            else:
                pass
                
        return portfolio_returns

engine = SimulationEngine()
