import pandas as pd
import numpy as np


def generate_signals(
    zscore: pd.Series,
    entry_threshold: float = 2.0,
    exit_threshold: float = 0.5,
) -> pd.DataFrame:
    """
    Génère les signaux de trading basés sur le z-score du spread.

    Règles (machine à états) :
    - LONG spread (buy A, short B) quand z-score < -entry_threshold
    - SHORT spread (short A, buy B) quand z-score > +entry_threshold
    - EXIT quand |z-score| < exit_threshold
    - STOP LOSS si |z-score| > 3.5 (le spread diverge trop)

    Parameters
    ----------
    zscore : pd.Series
        Le z-score rolling du spread.
    entry_threshold : float
        Seuil d'entrée en position (défaut : 2.0).
    exit_threshold : float
        Seuil de sortie de position (défaut : 0.5).

    Returns
    -------
    pd.DataFrame
        DataFrame avec colonnes :
        - 'zscore' : le z-score
        - 'signal' : 1 (long spread), -1 (short spread), 0 (flat)
        - 'trade_type' : 'entry_long', 'entry_short', 'exit', 'stop_loss', None
    """
    seuil_stop_loss = 3.5

    # Initialisation des listes pour stocker les résultats
    signaux = np.zeros(len(zscore), dtype=int)
    types_trade = [None] * len(zscore)

    position_actuelle = 0  # 0 = flat, 1 = long, -1 = short

    for i, z in enumerate(zscore.values):
        # Ignorer les NaN (début de la série avant que le rolling soit calculé)
        if np.isnan(z):
            signaux[i] = 0
            continue

        if position_actuelle == 0:
            # Pas de position → chercher une entrée
            if z < -entry_threshold:
                position_actuelle = 1
                types_trade[i] = "entry_long"
            elif z > entry_threshold:
                position_actuelle = -1
                types_trade[i] = "entry_short"
        else:
            # En position → chercher une sortie
            if abs(z) > seuil_stop_loss:
                position_actuelle = 0
                types_trade[i] = "stop_loss"
            elif abs(z) < exit_threshold:
                position_actuelle = 0
                types_trade[i] = "exit"

        signaux[i] = position_actuelle

    return pd.DataFrame(
        {
            "zscore": zscore.values,
            "signal": signaux,
            "trade_type": types_trade,
        },
        index=zscore.index,
    )