"""
AHAD QUANT — Reinforcement Learning Reward Function
Calcule la récompense à chaque step de l'environnement.

Philosophie :
  - Récompense directement le PnL réel (pas l'accuracy)
  - Pénalise les drawdowns (capital preservation)
  - Encourage le Sharpe ratio (qualité des gains)
  - Pénalise légèrement l'inactivité prolongée
"""

import numpy as np
from dataclasses import dataclass, field
from typing import List


# ─── Configuration des poids ─────────────────────────────────────────────────

# PnL
PNL_SCALE          = 1.0     # Multiplicateur du PnL réalisé
UNREALIZED_SCALE   = 0.3     # Poids du PnL non réalisé (moindre que réalisé)

# Drawdown
DRAWDOWN_PENALTY   = 2.0     # Pénalité par unité de drawdown (ex: -0.02 → -0.04)
MAX_DD_THRESHOLD   = 0.05    # Drawdown > 5% : pénalité supplémentaire

# Sharpe
SHARPE_WINDOW      = 20      # Fenêtre des returns pour le Sharpe
SHARPE_BONUS       = 0.1     # Bonus par unité de Sharpe (encouragement)

# Inactivité
IDLE_PENALTY       = -0.0002  # Pénalité par step sans position (encourage le trading)
MAX_IDLE_STEPS     = 48       # Après 48 steps sans trade, la pénalité s'amplifie

# Trades
WIN_BONUS          = 0.05    # Bonus à chaque trade gagnant
LOSS_MULTIPLIER    = 1.5     # Les pertes comptent 1.5x plus que les gains
CONSECUTIVE_WIN    = 0.02    # Bonus supplémentaire pour trades gagnants consécutifs

# Timeout
TIMEOUT_PENALTY    = -0.01   # Pénalité si on tient une position > MAX_HOLD_CANDLES
TIMEOUT_THRESHOLD  = 3       # Steps avant la pénalité de timeout


@dataclass
class RewardTracker:
    """
    État interne du calculateur de récompense.
    Une instance par épisode d'entraînement.
    """
    peak_balance:      float = 1.0       # Balance maximum atteinte (pour drawdown)
    current_balance:   float = 1.0       # Balance courante normalisée
    returns_history:   List[float] = field(default_factory=list)  # Pour Sharpe
    idle_steps:        int = 0           # Steps consécutifs sans position
    consecutive_wins:  int = 0           # Trades gagnants consécutifs
    total_trades:      int = 0
    winning_trades:    int = 0
    candles_in_pos:    int = 0           # Durée de la position courante

    def reset(self):
        self.peak_balance     = 1.0
        self.current_balance  = 1.0
        self.returns_history  = []
        self.idle_steps       = 0
        self.consecutive_wins = 0
        self.total_trades     = 0
        self.winning_trades   = 0
        self.candles_in_pos   = 0


def compute_step_reward(
    tracker:          RewardTracker,
    pnl_realized:     float,   # PnL réalisé ce step (0.0 si pas de close)
    pnl_unrealized:   float,   # PnL non réalisé courant (0.0 si pas en position)
    in_position:      bool,    # L'agent est-il en position ?
    trade_closed:     bool,    # Un trade a-t-il été fermé ce step ?
    fee_paid:         float,   # Frais payés ce step
) -> float:
    """
    Calcule la récompense scalaire pour un step donné.

    Returns:
        reward (float) : récompense à donner à l'agent RL
    """
    reward = 0.0

    # ─── 1. PnL réalisé ──────────────────────────────────────────────────────
    if trade_closed:
        net_pnl = pnl_realized - fee_paid
        if net_pnl >= 0:
            reward += net_pnl * PNL_SCALE + WIN_BONUS
            tracker.consecutive_wins += 1
            reward += tracker.consecutive_wins * CONSECUTIVE_WIN
            tracker.winning_trades += 1
        else:
            # Les pertes pèsent plus que les gains
            reward += net_pnl * PNL_SCALE * LOSS_MULTIPLIER
            tracker.consecutive_wins = 0
        tracker.total_trades += 1
        tracker.candles_in_pos = 0

    # ─── 2. PnL non réalisé (signal continu en position) ─────────────────────
    if in_position:
        reward += pnl_unrealized * UNREALIZED_SCALE
        tracker.candles_in_pos += 1
        tracker.idle_steps = 0

        # Pénalité timeout (position trop longue)
        if tracker.candles_in_pos > TIMEOUT_THRESHOLD:
            reward += TIMEOUT_PENALTY * (tracker.candles_in_pos - TIMEOUT_THRESHOLD)
    else:
        tracker.candles_in_pos = 0
        tracker.idle_steps += 1

    # ─── 3. Inactivité ───────────────────────────────────────────────────────
    if not in_position:
        penalty = IDLE_PENALTY
        if tracker.idle_steps > MAX_IDLE_STEPS:
            penalty *= 2.0   # Double la pénalité après 48h d'inactivité
        reward += penalty

    # ─── 4. Drawdown ─────────────────────────────────────────────────────────
    tracker.current_balance = 1.0 + pnl_realized  # Approx normalisée
    if tracker.current_balance > tracker.peak_balance:
        tracker.peak_balance = tracker.current_balance

    drawdown = (tracker.peak_balance - tracker.current_balance) / tracker.peak_balance
    if drawdown > 0:
        reward -= drawdown * DRAWDOWN_PENALTY
        if drawdown > MAX_DD_THRESHOLD:
            reward -= drawdown * DRAWDOWN_PENALTY  # Double pénalité si > 5%

    # ─── 5. Bonus Sharpe (rétrospectif sur fenêtre glissante) ────────────────
    if trade_closed:
        ret = pnl_realized - fee_paid
        tracker.returns_history.append(ret)
        if len(tracker.returns_history) > SHARPE_WINDOW:
            tracker.returns_history.pop(0)
        sharpe = _rolling_sharpe(tracker.returns_history)
        if sharpe > 0:
            reward += sharpe * SHARPE_BONUS

    return float(np.clip(reward, -1.0, 1.0))


def _rolling_sharpe(returns: List[float]) -> float:
    """Sharpe ratio simplifié sur une liste de returns."""
    if len(returns) < 5:
        return 0.0
    arr = np.array(returns)
    std = arr.std()
    if std < 1e-8:
        return 0.0
    return float(arr.mean() / std)


def compute_episode_metrics(tracker: RewardTracker) -> dict:
    """Métriques de fin d'épisode pour les logs TensorBoard."""
    win_rate = (tracker.winning_trades / tracker.total_trades
                if tracker.total_trades > 0 else 0.0)
    return {
        "ep/total_trades":     tracker.total_trades,
        "ep/win_rate":         win_rate,
        "ep/consecutive_wins": tracker.consecutive_wins,
        "ep/sharpe":           _rolling_sharpe(tracker.returns_history),
    }