"""
Z-Score Engine — Polars Implementation.
Replaces the Pandas engine with high-performance lazy execution.
"""

import yfinance as yf
import polars as pl
from datetime import datetime
from models import ZScoreEvent, ZScoreResult

# Known major historical events for enriched context
KNOWN_EVENTS = {
    (2020, 2): "COVID-19 Market Crash — Global pandemic triggers historic sell-off",
    (2020, 3): "COVID-19 Market Crash — Peak pandemic fear and lockdowns",
    (2008, 8): "2008 Financial Crisis — Lehman Brothers collapse begins",
    (2008, 9): "2008 Financial Crisis — Full-blown credit market freeze",
    (2008, 10): "2008 Financial Crisis — Global contagion and bank bailouts",
    (2008, 11): "2008 Financial Crisis — Continued deleveraging",
    (2001, 8): "Dot-com Bubble Aftermath — Tech sector implosion",
    (2001, 9): "9/11 Attacks — Markets shut down then crash on reopening",
    (2022, 5): "2022 Bear Market — Fed rate hikes crush growth stocks",
    (2022, 6): "2022 Bear Market — Inflation fears peak",
    (2020, 10): "Pre-Election Volatility — Uncertainty ahead of US elections",
    (2011, 7): "US Debt Ceiling Crisis — S&P downgrades US credit rating",
    (2011, 8): "US Debt Ceiling Crisis — Market turmoil continues",
    (2018, 12): "Fed Tightening Scare — December 2018 sell-off",
    (2015, 8): "China Devaluation Shock — Yuan devaluation rattles global markets",
}

def _classify_event(z: float, dt: datetime) -> tuple[str, str]:
    """Classify an event as major/normal and generate context."""
    year = dt.year
    month = dt.month
    abs_z = abs(z)

    known = KNOWN_EVENTS.get((year, month))

    if known or abs_z >= 3.0:
        event_type = "major"
        if known:
            context = known
        elif z < 0:
            context = "Significant market decline — potential crisis event"
        else:
            context = "Significant market rally — unusual positive movement"
    else:
        event_type = "normal"
        if z < 0:
            context = "Notable market decline — day-to-day volatility"
        else:
            context = "Notable market increase — day-to-day volatility"

    return event_type, context

def identify_events(symbol: str, window: int = 100, trigger_min: float = -2.5, trigger_max: float = 2.5) -> ZScoreResult:
    """
    Downloads data with yfinance and calculates rolling Z-scores using Polars.
    Accepts empty symbol to aggregate a basket of top stocks.
    Returns filtered ZScoreResult.
    """
    import random
    
    symbols_to_fetch = [symbol] if symbol else ["SPY", "QQQ", "AAPL", "MSFT", "TSLA", "AMZN", "NVDA", "META"]
    if not symbol:
        # Pick 3 random stocks + SPY to avoid massive token blowout
        symbols_to_fetch = ["SPY"] + random.sample([s for s in symbols_to_fetch if s != "SPY"], 3)
    
    all_events = []
    total_data_points = 0
    
    for current_symbol in symbols_to_fetch:
        try:
            # 1. Fetch data
            data = yf.download(current_symbol, period="5y", progress=False)
            if data.empty:
                continue
                
            # 2. Reset index to get Date as a column, handle multi-layer columns from yf
            df_pd = data.reset_index()
            
            # Flatten columns if yfinance returns multi-index
            if isinstance(df_pd.columns, pl.DataFrame): # safety fallback
                 pass
            
            new_cols = []
            for col in df_pd.columns:
                if isinstance(col, tuple):
                    # Filter out empty strings from the tuple and join
                    parts = [str(c) for c in col if c]
                    new_cols.append('_'.join(parts).strip('_'))
                else:
                    new_cols.append(str(col))
            df_pd.columns = new_cols
            
            # Ensure we have Date and Close
            date_col = next((c for c in df_pd.columns if 'date' in c.lower()), None)
            close_col = next((c for c in df_pd.columns if 'close' in c.lower() and ('_' not in c or current_symbol.lower() in c.lower() or 'close' == c.lower())), None)
            
            if not close_col: # Fallback to just the first column that has 'close'
                close_col = next((c for c in df_pd.columns if 'close' in c.lower()), None)
            
            if not date_col or not close_col:
                print(f"Skipping {current_symbol}: Could not cleanly identify Date/Close columns.")
                continue
                
            df_pd = df_pd[[date_col, close_col]].rename(columns={date_col: "Date", close_col: "Close"})
    
            # Drop NaN
            df_pd = df_pd.dropna()
            if len(df_pd) < window:
                 continue
    
            # 3. Process with Polars
            df = pl.from_pandas(df_pd)
            
            # Ensure correct types
            df = df.with_columns([
                pl.col("Date").cast(pl.Datetime)
            ])
            
            # 4. Calculate Z-Scores Lazy Execution
            q = (
                df.lazy()
                .with_columns([
                    pl.col("Close").rolling_mean(window_size=window).alias("Mean"),
                    pl.col("Close").rolling_std(window_size=window).alias("Std")
                ])
                .with_columns([
                    ((pl.col("Close") - pl.col("Mean")) / pl.col("Std")).alias("Z_Score")
                ])
                .drop_nulls() 
            )
            
            # Force computation
            processed_df = q.collect()
            total_data_points += len(processed_df)
            
            # Filter significant events
            events_df = processed_df.filter(
                (pl.col("Z_Score") <= trigger_min) | (pl.col("Z_Score") >= trigger_max)
            )
            
            # Convert to Pydantic models
            for row in events_df.to_dicts():
                z = round(float(row["Z_Score"]), 3)
                dt = row["Date"]
                event_type, context = _classify_event(z, dt)
                
                # Append stock symbol to context for mixed baskets
                if not symbol:
                    context = f"[{current_symbol}] " + context
                    
                all_events.append(ZScoreEvent(
                    date=dt.strftime("%Y-%m-%d"),
                    price=round(float(row["Close"]), 2),
                    z_score=z,
                    event_type=event_type,
                    context=context,
                    direction="decline" if z < 0 else "rally"
                ))
                
        except Exception as e:
            print(f"Warning: Z-Score calculation failed for {current_symbol}: {str(e)}")
            continue

    if not all_events:
        raise RuntimeError(f"No events found for any requested symbols.")
        
    # Sort events by date descending and limit to top 150 to keep LLM context size healthy
    all_events = sorted(all_events, key=lambda x: abs(x.z_score), reverse=True)[:150]

    return ZScoreResult(
        events=all_events,
        total_events=len(all_events),
        window_size=window,
        data_points=total_data_points
    )
