# -*- coding: utf-8 -*-
"""
Feature synthesis utilities for multi-day forecasting

Generates synthetic feature vectors for days beyond the first prediction
using conditional autoregression based on model predictions and historical patterns.
"""

from typing import Dict, Any, Optional


def compute_synthesis_statistics(dataframe) -> Dict[str, float]:
    """
    Analyze training data to learn typical UP/DOWN day patterns for feature synthesis

    Args:
        dataframe: Training dataframe with 'Labels' column (1=UP, 0=DOWN)
                   and financial features (Close, Volume, etc.)

    Returns:
        Dictionary of statistics for feature synthesis:
        - up_day_return_mean/std: Average and std of returns on UP days
        - down_day_return_mean/std: Average and std of returns on DOWN days
        - up_day_volume_ratio: Relative volume on UP vs DOWN days
        - feature means/stds for normalization
    """
    import numpy as np

    if "Labels" not in dataframe.columns:
        raise ValueError("Dataframe must have 'Labels' column for synthesis stats")

    # Separate UP and DOWN days
    up_days = dataframe[dataframe["Labels"] == 1]
    down_days = dataframe[dataframe["Labels"] == 0]

    # Calculate returns if Close column exists
    stats = {}

    if "Close" in dataframe.columns:
        # Calculate daily returns
        up_returns = up_days["Close"].pct_change().dropna()
        down_returns = down_days["Close"].pct_change().dropna()

        stats["up_day_return_mean"] = float(up_returns.mean())
        stats["up_day_return_std"] = float(up_returns.std())
        stats["down_day_return_mean"] = float(down_returns.mean())
        stats["down_day_return_std"] = float(down_returns.std())
    else:
        # Fallback defaults if Close not available
        stats["up_day_return_mean"] = 0.015  # 1.5% typical gain
        stats["up_day_return_std"] = 0.005
        stats["down_day_return_mean"] = -0.012  # 1.2% typical loss
        stats["down_day_return_std"] = 0.005

    # Volume analysis if available
    if "Volume" in dataframe.columns:
        up_volume_mean = up_days["Volume"].mean()
        down_volume_mean = down_days["Volume"].mean()
        stats["up_day_volume_ratio"] = float(
            up_volume_mean / down_volume_mean if down_volume_mean > 0 else 1.2
        )
        stats["down_day_volume_ratio"] = 1.0  # Baseline
    else:
        stats["up_day_volume_ratio"] = 1.2  # UP days typically have 20% more volume
        stats["down_day_volume_ratio"] = 1.0

    # Store overall feature statistics for fallback synthesis
    numeric_cols = dataframe.select_dtypes(include=[np.number]).columns
    feature_cols = [c for c in numeric_cols if c not in ["Labels", "Date"]]

    stats["feature_means"] = {}
    stats["feature_stds"] = {}
    for col in feature_cols:
        stats["feature_means"][col] = float(dataframe[col].mean())
        stats["feature_stds"][col] = float(dataframe[col].std())

    return stats


def synthesize_next_day_features(
    last_window,
    predicted_direction: int,
    training_stats: Dict[str, Any],
    industry_avg: Optional[Any] = None,
):
    """
    Generate synthetic features for the next day using conditional autoregression

    This function creates plausible feature vectors for day +1 based on:
    1. Model's UP/DOWN prediction (50% weight)
    2. Industry momentum if available (30% weight)
    3. Ticker momentum from recent history (20% weight)
    4. Realistic stochastic noise

    Args:
        last_window: Recent feature history, shape (tsize, n_features)
                     where last row is most recent day
        predicted_direction: Model's prediction (1=UP, 0=DOWN)
        training_stats: Statistics from compute_synthesis_statistics()
        industry_avg: Optional industry average features, shape (n_features,)

    Returns:
        Synthetic feature vector for next day, shape (n_features,)
    """
    import numpy as np

    if len(last_window.shape) != 2:
        raise ValueError(
            f"last_window must be 2D (tsize, n_features), got shape {last_window.shape}"
        )

    last_day = last_window[-1]  # Most recent day
    n_features = last_window.shape[1]

    # 1. Select appropriate return statistics based on prediction
    if predicted_direction == 1:  # UP
        typical_return = training_stats["up_day_return_mean"]
        return_std = training_stats["up_day_return_std"]
        volume_ratio = training_stats["up_day_volume_ratio"]
    else:  # DOWN
        typical_return = training_stats["down_day_return_mean"]
        return_std = training_stats["down_day_return_std"]
        volume_ratio = training_stats["down_day_volume_ratio"]

    # 2. Calculate momentum signals
    # Ticker momentum: recent 5-day trend
    if len(last_window) >= 5:
        ticker_momentum = (last_window[-1] - last_window[-5]).mean() / (
            np.abs(last_window[-5]).mean() + 1e-8
        )
    else:
        ticker_momentum = 0.0

    # Industry momentum: if provided
    if industry_avg is not None:
        # Compare last day to industry average
        industry_momentum = (last_day - industry_avg).mean() / (
            np.abs(industry_avg).mean() + 1e-8
        )
    else:
        industry_momentum = 0.0

    # 3. Weighted combination of signals
    # 50% model prediction, 30% industry, 20% ticker momentum
    predicted_return = (
        0.5 * typical_return + 0.3 * industry_momentum + 0.2 * ticker_momentum
    )

    # Add realistic Gaussian noise
    noise = np.random.normal(0, return_std)
    predicted_return += noise

    # 4. Generate synthetic features
    # Simple approach: apply return multiplicatively to all features
    # This assumes features are price-like (OHLC, indicators scale with price)
    synthetic_features = last_day * (1 + predicted_return)

    # Adjust specific feature indices if we know the structure
    # For now, apply return uniformly with some feature-specific variation

    # Add small random variations to make features more realistic
    feature_noise = np.random.normal(0, 0.01, size=n_features)
    synthetic_features = synthetic_features * (1 + feature_noise)

    # Special handling for volume-like features (if we can identify them)
    # Volume typically increases on UP days, decreases on DOWN days
    # This is a heuristic - in production, you'd want to identify which features are volume
    synthetic_features = synthetic_features * volume_ratio

    # Ensure no negative values for features that should be positive
    # (e.g., prices, volume)
    synthetic_features = np.clip(synthetic_features, 0, None)

    return synthetic_features


def synthesize_next_day_features_structured(
    last_window,
    predicted_direction: int,
    training_stats: Dict[str, Any],
    industry_avg: Optional[Any] = None,
    feature_names: Optional[list] = None,
):
    """
    Advanced feature synthesis with structure-aware generation

    This version attempts to identify and synthesize specific feature types
    (OHLC, volume, technical indicators) more accurately.

    Args:
        last_window: Recent feature history, shape (tsize, n_features)
        predicted_direction: Model's prediction (1=UP, 0=DOWN)
        training_stats: Statistics from compute_synthesis_statistics()
        industry_avg: Optional industry average features
        feature_names: List of feature names for structure-aware synthesis

    Returns:
        Synthetic feature vector for next day, shape (n_features,)
    """
    # For now, use simple approach - can be enhanced later
    # with feature name detection (Open, High, Low, Close, Volume, RSI, etc.)
    return synthesize_next_day_features(
        last_window, predicted_direction, training_stats, industry_avg
    )


def validate_synthetic_features(
    synthetic, last_window, predicted_direction: int
) -> bool:
    """
    Validate that synthetic features are reasonable

    Checks:
    - No NaN or Inf values
    - Values within reasonable bounds relative to recent history
    - Direction-consistent price movement (if detectable)

    Args:
        synthetic: Synthetic features to validate
        last_window: Recent history for comparison
        predicted_direction: Expected direction (1=UP, 0=DOWN)

    Returns:
        True if features pass validation, False otherwise
    """
    import numpy as np

    # Check for invalid values
    if np.any(np.isnan(synthetic)) or np.any(np.isinf(synthetic)):
        return False

    # Check that values are within reasonable bounds (e.g., not 1000x last day)
    last_day = last_window[-1]
    ratio = synthetic / (last_day + 1e-8)

    # Features should not change by more than 50% in one day (generous bound)
    if np.any(ratio > 1.5) or np.any(ratio < 0.5):
        return False

    # All checks passed
    return True


def run_forecast(model, device, data_df, ind_mean, cfg, checkpoint):
    """
    Multi-day forecast for neural network models (GRU, Transformer)

    Note: Imports torch locally (following pattern from training_service.py)
    """
    import torch
    import numpy as np
    import pandas as pd
    from automar.shared.services.tuning_service import SCALERS

    forecast_days = cfg.predict.forecast_days or 1
    synthesis_stats = checkpoint.get("synthesis_stats")

    # Extract optimal threshold from checkpoint (stored as fraction 0-1)
    threshold = 0.5  # Default fallback
    if checkpoint and isinstance(checkpoint, dict):
        test_metrics = checkpoint.get("test_metrics")
        if test_metrics and isinstance(test_metrics, dict):
            threshold = test_metrics.get("threshold", 0.5)

    if forecast_days > 1 and not synthesis_stats:
        import warnings

        warnings.warn("Computing synthesis stats from current data", UserWarning)
        synthesis_stats = compute_synthesis_statistics(data_df)

    # Prepare scaler (reuse from tuning_service)
    scaler = SCALERS[cfg.loader.scaler]()
    feature_df = (
        data_df.drop(columns=["Labels"])
        if "Labels" in data_df.columns
        else data_df.iloc[:, :-1]
    )
    all_features = feature_df.values
    scaler.fit(all_features[: int(len(all_features) * 0.8)])

    current_window = scaler.transform(all_features[-cfg.loader.tsize :])
    industry_avg = _extract_industry_features(ind_mean, all_features.shape[1])

    # Forecast loop
    predictions, probabilities, forecast_dates = [], [], []
    last_date = data_df.index[-1]

    for day_ahead in range(forecast_days):
        input_tensor = torch.FloatTensor(current_window).unsqueeze(0).to(device)
        with torch.no_grad():
            output = model(input_tensor)
            prob = torch.sigmoid(output).squeeze().item()
            pred = int(prob > threshold)  # Use optimal threshold from checkpoint

        next_date = pd.bdate_range(start=last_date, periods=day_ahead + 2)[-1]
        predictions.append(pred)
        probabilities.append(prob * 100)
        forecast_dates.append(next_date)

        if day_ahead < forecast_days - 1:
            synthetic = synthesize_next_day_features(
                scaler.inverse_transform(current_window),
                pred,
                synthesis_stats,
                industry_avg,
            )
            current_window = np.vstack(
                [
                    current_window[1:],
                    scaler.transform(synthetic.reshape(1, -1)).flatten(),
                ]
            )

    return {
        "predictions": predictions,
        "probabilities": probabilities,
        "labels": [],
        "metrics": {"threshold": threshold * 100},  # Convert to percentage for frontend
        "confusion_matrix": [],
        "forecast_dates": [d.strftime("%Y-%m-%d") for d in forecast_dates],
        "forecast_days": forecast_days,
        "mode": "forecast",
    }


def run_forecast_logreg(model, data_df, ind_mean, cfg, checkpoint=None):
    """
    Multi-day forecast for logistic regression models
    """
    import numpy as np
    import pandas as pd
    from automar.shared.services.tuning_service import SCALERS

    forecast_days = cfg.predict.forecast_days or 1
    synthesis_stats = None

    # Extract optimal threshold from checkpoint (stored as fraction 0-1)
    threshold = 0.5  # Default fallback
    if checkpoint and isinstance(checkpoint, dict):
        test_metrics = checkpoint.get("test_metrics")
        if test_metrics and isinstance(test_metrics, dict):
            threshold = test_metrics.get("threshold", 0.5)

    if forecast_days > 1:
        import warnings

        warnings.warn("Computing synthesis stats from current data", UserWarning)
        synthesis_stats = compute_synthesis_statistics(data_df)

    # Prepare scaler
    scaler = SCALERS[cfg.loader.scaler]()
    feature_df = (
        data_df.drop(columns=["Labels"])
        if "Labels" in data_df.columns
        else data_df.iloc[:, :-1]
    )
    all_features = feature_df.values
    scaler.fit(all_features[: int(len(all_features) * 0.8)])

    current_window = scaler.transform(all_features[-cfg.loader.tsize :])
    industry_avg = _extract_industry_features(ind_mean, all_features.shape[1])

    # Forecast loop
    predictions, probabilities, forecast_dates = [], [], []
    last_date = data_df.index[-1]

    for day_ahead in range(forecast_days):
        feature_vector = current_window.flatten().reshape(1, -1)
        prediction = model.predict(feature_vector)[0]
        probability = model.predict_proba(feature_vector)[0, 1]

        next_date = pd.bdate_range(start=last_date, periods=day_ahead + 2)[-1]
        predictions.append(int(prediction))
        probabilities.append(float(probability * 100))
        forecast_dates.append(next_date)

        if day_ahead < forecast_days - 1:
            synthetic = synthesize_next_day_features(
                scaler.inverse_transform(current_window),
                int(prediction),
                synthesis_stats,
                industry_avg,
            )
            current_window = np.vstack(
                [
                    current_window[1:],
                    scaler.transform(synthetic.reshape(1, -1)).flatten(),
                ]
            )

    return {
        "predictions": predictions,
        "probabilities": probabilities,
        "labels": [],
        "metrics": {"threshold": threshold * 100},  # Convert to percentage for frontend
        "confusion_matrix": [],
        "forecast_dates": [d.strftime("%Y-%m-%d") for d in forecast_dates],
        "forecast_days": forecast_days,
        "mode": "forecast",
    }


def _extract_industry_features(ind_mean, n_features):
    """Extract industry average features for synthesis"""
    if ind_mean is None:
        return None
    ind_mean_df = (
        ind_mean.drop(columns=["Labels"]) if "Labels" in ind_mean.columns else ind_mean
    )
    features = ind_mean_df.values[-1]
    return features if len(features) == n_features else None
