import pandas as pd
import numpy as np
from statsmodels.tsa.stattools import coint
from statsmodels.regression.linear_model import OLS
from statsmodels.tools import add_constant


def test_cointegration(prices_a: pd.Series, prices_b: pd.Series) -> dict:
    """
    Effectue le test de cointégration d'Engle-Granger sur deux séries de prix.

    Étapes :
    1. Régression linéaire : prices_a = alpha + beta * prices_b + résidus
    2. Test ADF (Augmented Dickey-Fuller) sur les résidus
    3. Si p-value < 0.05 → la paire est cointégrée

    Parameters
    ----------
    prices_a : pd.Series
        Prix du premier actif.
    prices_b : pd.Series
        Prix du second actif.

    Returns
    -------
    dict
        Résultats du test avec les clés :
        - is_cointegrated (bool)
        - p_value (float)
        - beta (float) : coefficient de hedge
        - alpha (float) : intercept de la régression
        - test_statistic (float)
        - critical_values (dict) : seuils à 1%, 5%, 10%
    """
    # Régression OLS : prices_a = alpha + beta * prices_b
    prices_b_avec_constante = add_constant(prices_b)
    modele = OLS(prices_a, prices_b_avec_constante).fit()
    alpha = modele.params.iloc[0]
    beta = modele.params.iloc[1]

    # Test de cointégration d'Engle-Granger via statsmodels
    statistique_test, p_value, valeurs_critiques = coint(prices_a, prices_b)

    return {
        "is_cointegrated": p_value < 0.05,
        "p_value": round(float(p_value), 6),
        "beta": round(float(beta), 6),
        "alpha": round(float(alpha), 6),
        "test_statistic": round(float(statistique_test), 4),
        "critical_values": {
            "1%": round(float(valeurs_critiques[0]), 4),
            "5%": round(float(valeurs_critiques[1]), 4),
            "10%": round(float(valeurs_critiques[2]), 4),
        },
    }


def compute_spread(prices_a: pd.Series, prices_b: pd.Series, beta: float) -> pd.Series:
    """
    Calcule le spread entre deux actifs.

    spread = prices_a - beta * prices_b

    Parameters
    ----------
    prices_a : pd.Series
        Prix du premier actif.
    prices_b : pd.Series
        Prix du second actif.
    beta : float
        Coefficient de hedge issu de la régression.

    Returns
    -------
    pd.Series
        Le spread entre les deux actifs.
    """
    return prices_a - beta * prices_b


def compute_zscore(spread: pd.Series, window: int = 60) -> pd.Series:
    """
    Calcule le z-score rolling du spread.

    z = (spread - rolling_mean) / rolling_std

    Parameters
    ----------
    spread : pd.Series
        Le spread entre les deux actifs.
    window : int
        Fenêtre glissante en jours (défaut : 60, ~3 mois de trading).

    Returns
    -------
    pd.Series
        Le z-score rolling du spread.
    """
    moyenne_glissante = spread.rolling(window=window).mean()
    ecart_type_glissant = spread.rolling(window=window).std()

    zscore = (spread - moyenne_glissante) / ecart_type_glissant
    return zscore