# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Realized-variance forecasting utilities (RQ2/RQ5): OLS forecasts on raw
feature levels, rolling out-of-sample evaluation, and loss functions."""

import numpy as np
import pandas as pd
from numpy.linalg import lstsq


def ols_forecast(train: pd.DataFrame, test: pd.DataFrame,
                 features: list, target: str) -> dict:
    """Fit OLS on ``train`` and evaluate on ``test`` (non-negative forecasts).

    Features enter in levels (no standardization), matching the original
    forecasting design. Returns MSE, MAE, out-of-sample R², QLIKE and the
    fitted coefficients.
    """
    X_train = np.column_stack([np.ones(len(train)), train[features].values])
    y_train = train[target].values
    coefs, _, _, _ = lstsq(X_train, y_train, rcond=None)

    X_test = np.column_stack([np.ones(len(test)), test[features].values])
    y_test = test[target].values
    y_pred = np.maximum(X_test @ coefs, 0)

    mse = np.mean((y_test - y_pred) ** 2)
    mae = np.mean(np.abs(y_test - y_pred))
    ss_res = np.sum((y_test - y_pred) ** 2)
    ss_tot = np.sum((y_test - y_test.mean()) ** 2)
    r2_oos = 1 - ss_res / ss_tot if ss_tot > 0 else 0

    y_pred_safe = np.maximum(y_pred, 1e-10)
    qlike = np.mean(y_test / y_pred_safe - np.log(y_test / y_pred_safe) - 1)

    return {'mse': mse, 'mae': mae, 'r2_oos': r2_oos, 'qlike': qlike, 'coefs': coefs}


def rolling_evaluation(data: pd.DataFrame, models: dict, target_col: str,
                       target_label: str, window: int = 500,
                       step: int = 250) -> pd.DataFrame:
    """Per-ticker rolling-window out-of-sample evaluation of competing models.

    For each ticker and model, fits on ``window`` observations and evaluates
    on the following ``step`` observations, advancing by ``step``. Returns
    one row per (ticker, model) with average losses across windows.
    """
    all_results = []
    for ticker in data['ticker'].unique():
        td = data[data['ticker'] == ticker].reset_index(drop=True)
        if len(td) < window + 100:
            continue

        for model_name, features in models.items():
            sub = td[features + [target_col]].dropna()
            if len(sub) < window + 50:
                continue

            losses = {'mse': [], 'mae': [], 'r2': [], 'qlike': []}
            for start in range(0, len(sub) - window - 50, step):
                train = sub.iloc[start:start + window]
                test = sub.iloc[start + window:start + window + step]
                if len(test) < 10:
                    continue
                try:
                    res = ols_forecast(train, test, features, target_col)
                except np.linalg.LinAlgError:
                    continue
                losses['mse'].append(res['mse'])
                losses['mae'].append(res['mae'])
                losses['r2'].append(res['r2_oos'])
                losses['qlike'].append(res['qlike'])

            if losses['mse']:
                all_results.append({
                    'ticker': ticker,
                    'model': model_name,
                    'target': target_label,
                    'avg_mse': np.mean(losses['mse']),
                    'avg_mae': np.mean(losses['mae']),
                    'avg_r2_oos': np.mean(losses['r2']),
                    'avg_qlike': np.mean(losses['qlike']),
                    'n_windows': len(losses['mse']),
                })
    return pd.DataFrame(all_results)


def diebold_mariano(e1: np.ndarray, e2: np.ndarray) -> float:
    """DM statistic on squared-error differentials (positive favors model 2)."""
    d = e1 ** 2 - e2 ** 2
    return np.mean(d) / (np.std(d) / np.sqrt(len(d))) if np.std(d) > 0 else 0