# ── Install any missing packages ──────────────────────────────────────────────
# Gymnasium is the maintained fork of OpenAI Gym.
# yfinance fetches real OHLCV data from Yahoo Finance.
!pip install gymnasium yfinance --quiet

"""
Central import block.
Keeping all imports here makes dependencies explicit and easy to audit.
"""

# ── Standard library ──────────────────────────────────────────────────────────
import os
import random
import warnings
from collections import deque
from typing import Tuple, List, Dict, Optional

# ── Numerical & data ─────────────────────────────────────────────────────────
import numpy as np
import pandas as pd

# ── Deep learning ─────────────────────────────────────────────────────────────
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical

# ── RL environment ────────────────────────────────────────────────────────────
import gymnasium as gym
from gymnasium import spaces

# ── Visualization ─────────────────────────────────────────────────────────────
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import seaborn as sns

# ── Optional: real market data ────────────────────────────────────────────────
try:
    import yfinance as yf
    YFINANCE_AVAILABLE = True
except ImportError:
    YFINANCE_AVAILABLE = False

# ── Global settings ───────────────────────────────────────────────────────────
warnings.filterwarnings('ignore')          # Suppress deprecation noise
sns.set_theme(style='darkgrid')            # Consistent plot style
plt.rcParams['figure.figsize'] = (14, 5)  # Default figure size

# ── Reproducibility ───────────────────────────────────────────────────────────
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

# ── Device selection: GPU if available, else CPU ──────────────────────────────
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print(f" All imports successful")
print(f"  Using device: {DEVICE}")
print(f" PyTorch version: {torch.__version__}")

def generate_synthetic_price_data(
    n_steps: int = 2000,
    initial_price: float = 100.0,
    mu: float = 0.0002,
    sigma: float = 0.015,
    seed: int = SEED
) -> pd.DataFrame:
    """
    Generate synthetic OHLCV price data using Geometric Brownian Motion (GBM).

    GBM is the standard model for stock prices in quantitative finance. It assumes
    log-returns are normally distributed — a simplification, but useful for learning.

    The formula for each step:
        S(t+1) = S(t) * exp((mu - 0.5*sigma^2)*dt + sigma*sqrt(dt)*Z)
    where Z ~ Normal(0, 1) and dt = 1 (daily steps).

    Parameters
    ----------
    n_steps : int
        Number of daily timesteps to generate (default 2000 ≈ 8 years of trading days).
    initial_price : float
        Starting price of the asset.
    mu : float
        Annualised drift (expected daily return). Positive = upward trend.
    sigma : float
        Daily volatility (standard deviation of log-returns). ~1.5% per day is realistic.
    seed : int
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns: ['Open', 'High', 'Low', 'Close', 'Volume'].

    Example
    -------
    >>> df = generate_synthetic_price_data(n_steps=500)
    >>> print(df.shape)  # (500, 5)
    """
    rng = np.random.default_rng(seed)

    # ── Simulate log-returns and cumulative price path ────────────────────────
    # Each daily return is drawn from a normal distribution
    log_returns = (mu - 0.5 * sigma ** 2) + sigma * rng.standard_normal(n_steps)
    prices = initial_price * np.exp(np.cumsum(log_returns))  # Cumulative product via exp(sum)

    # ── Build OHLC from close prices ──────────────────────────────────────────
    # Intraday noise creates High/Low around the close
    intraday_noise = sigma * 0.5  # High/Low spread ≈ half the daily sigma
    opens  = prices * np.exp(rng.normal(0, intraday_noise * 0.3, n_steps))
    highs  = prices * np.exp(np.abs(rng.normal(0, intraday_noise, n_steps)))
    lows   = prices * np.exp(-np.abs(rng.normal(0, intraday_noise, n_steps)))
    closes = prices  # The simulated path IS the close series

    # ── Simulate volume with mean-reverting log-normal process ────────────────
    # Volume spikes on large price moves (stylised fact of real markets)
    base_volume = 1_000_000
    volume_noise = rng.lognormal(mean=0, sigma=0.5, size=n_steps)
    volumes = base_volume * volume_noise * (1 + 5 * np.abs(log_returns))  # Spike on big moves

    # ── Assemble into a DataFrame with a business-day date index ─────────────
    dates = pd.bdate_range(start='2015-01-01', periods=n_steps)
    df = pd.DataFrame({
        'Open':   opens,
        'High':   highs,
        'Low':    lows,
        'Close':  closes,
        'Volume': volumes.astype(int)
    }, index=dates)

    return df


# ── Demonstration ─────────────────────────────────────────────────────────────
raw_df = generate_synthetic_price_data(n_steps=2000)

print(f"Dataset shape: {raw_df.shape}")
print(f"Date range  : {raw_df.index[0].date()} → {raw_df.index[-1].date()}")
print(f"Price range : ${raw_df['Close'].min():.2f} – ${raw_df['Close'].max():.2f}")
print()
print(raw_df.head())

def engineer_features(df: pd.DataFrame, window_short: int = 10, window_long: int = 50) -> pd.DataFrame:
    """
    Transform raw OHLCV price data into normalised technical features.

    Raw prices are non-stationary (they drift over time) and not comparable
    across assets. Technical indicators convert prices into *relative* signals
    that a neural network can interpret meaningfully.

    Features Computed
    -----------------
    log_return       : log(Close_t / Close_{t-1}) — daily log return
    sma_ratio_short  : Close / SMA(window_short) — short-term trend
    sma_ratio_long   : Close / SMA(window_long)  — long-term trend
    rsi              : Relative Strength Index (0–100 normalised to 0–1)
    bb_pct           : Bollinger Band %B — price position within 2σ band
    volatility       : Rolling std of log-returns (normalised)
    volume_ratio     : Volume / SMA(volume, 20) — relative volume

    Parameters
    ----------
    df : pd.DataFrame
        Raw OHLCV data with columns ['Open', 'High', 'Low', 'Close', 'Volume'].
    window_short : int
        Lookback window for the short-term moving average (default 10 days).
    window_long : int
        Lookback window for the long-term moving average (default 50 days).

    Returns
    -------
    pd.DataFrame
        Feature DataFrame with NaN rows dropped. Shape: (n_steps - window_long, n_features).

    Example
    -------
    >>> features = engineer_features(raw_df)
    >>> print(features.columns.tolist())
    """
    feat = pd.DataFrame(index=df.index)
    close  = df['Close']
    volume = df['Volume']

    # ── 1. Log returns ─────────────────────────────────────────────────────────
    # Log returns are additive and approximately normally distributed
    feat['log_return'] = np.log(close / close.shift(1))

    # ── 2. Price / Moving Average ratios ──────────────────────────────────────
    # Values > 1 → price above average (bullish); < 1 → below (bearish)
    feat['sma_ratio_short'] = close / close.rolling(window_short).mean()
    feat['sma_ratio_long']  = close / close.rolling(window_long).mean()

    # ── 3. RSI (Relative Strength Index) ──────────────────────────────────────
    # RSI measures momentum: >0.7 = overbought, <0.3 = oversold
    rsi_window = 14
    delta   = close.diff()
    gain    = delta.clip(lower=0).rolling(rsi_window).mean()  # Average of up-days
    loss    = (-delta.clip(upper=0)).rolling(rsi_window).mean()  # Average of down-days
    rs      = gain / (loss + 1e-8)  # Relative strength; epsilon prevents division by zero
    feat['rsi'] = 1 - (1 / (1 + rs))  # Normalised RSI to [0, 1]

    # ── 4. Bollinger Band %B ──────────────────────────────────────────────────
    # %B = (Price - Lower Band) / (Upper Band - Lower Band)
    # 1.0 = at upper band, 0.0 = at lower band, 0.5 = at middle (SMA)
    bb_window  = 20
    bb_mid     = close.rolling(bb_window).mean()
    bb_std     = close.rolling(bb_window).std()
    bb_upper   = bb_mid + 2 * bb_std
    bb_lower   = bb_mid - 2 * bb_std
    feat['bb_pct'] = (close - bb_lower) / (bb_upper - bb_lower + 1e-8)

    # ── 5. Rolling volatility ─────────────────────────────────────────────────
    # High volatility = high risk; the agent should learn to be cautious
    feat['volatility'] = feat['log_return'].rolling(window_short).std()

    # ── 6. Volume ratio ───────────────────────────────────────────────────────
    # Volume spike often precedes or confirms price movement
    feat['volume_ratio'] = volume / volume.rolling(20).mean()

    # ── Drop NaN rows introduced by rolling windows ────────────────────────────
    feat.dropna(inplace=True)

    # ── Clip extreme outliers to ±5 standard deviations ───────────────────────
    # Outliers can destabilise neural network training
    for col in feat.columns:
        col_mean = feat[col].mean()
        col_std  = feat[col].std()
        feat[col] = feat[col].clip(col_mean - 5 * col_std, col_mean + 5 * col_std)

    return feat


# ── Demonstration ─────────────────────────────────────────────────────────────
features_df = engineer_features(raw_df)

print(f"Feature matrix shape : {features_df.shape}")
print(f"Features             : {features_df.columns.tolist()}")
print()
print(features_df.describe().round(4))

def load_real_market_data(
    ticker: str = 'SPY',
    start: str = '2015-01-01',
    end: str   = '2023-12-31'
) -> pd.DataFrame:
    """
    Download real OHLCV data from Yahoo Finance and engineer features.

    Uses yfinance under the hood. Falls back gracefully to synthetic data
    if yfinance is unavailable or the download fails (e.g., network issues
    in some Colab environments).

    Parameters
    ----------
    ticker : str
        Yahoo Finance ticker symbol (e.g., 'SPY', 'AAPL', 'BTC-USD').
    start : str
        Start date in 'YYYY-MM-DD' format.
    end : str
        End date in 'YYYY-MM-DD' format.

    Returns
    -------
    pd.DataFrame
        Raw OHLCV DataFrame aligned with engineer_features() expectations.

    Example
    -------
    >>> df = load_real_market_data('AAPL', '2018-01-01', '2023-01-01')
    """
    if not YFINANCE_AVAILABLE:
        print(" yfinance not available — using synthetic data instead.")
        return generate_synthetic_price_data()

    try:
        print(f" Downloading {ticker} from Yahoo Finance ...")
        raw = yf.download(ticker, start=start, end=end, progress=False, auto_adjust=True)

        # yfinance returns multi-level columns for some versions — flatten
        if isinstance(raw.columns, pd.MultiIndex):
            raw.columns = raw.columns.get_level_values(0)

        required = ['Open', 'High', 'Low', 'Close', 'Volume']
        missing  = [c for c in required if c not in raw.columns]
        if missing:
            raise ValueError(f"Missing columns after download: {missing}")

        raw.dropna(inplace=True)
        print(f" Downloaded {len(raw)} rows for {ticker}")
        return raw

    except Exception as exc:
        print(f"  Download failed ({exc}) — using synthetic data instead.")
        return generate_synthetic_price_data()


# ── Choose your data source here ──────────────────────────────────────────────
USE_REAL_DATA = False  # Set True to download SPY from Yahoo Finance

if USE_REAL_DATA:
    raw_df = load_real_market_data('SPY', '2015-01-01', '2023-12-31')
else:
    print("Using synthetic data (set USE_REAL_DATA=True for live data)")

# Re-engineer features on the chosen data source
features_df = engineer_features(raw_df)

# ── Train / Test split at 80% ─────────────────────────────────────────────────
split_idx     = int(len(features_df) * 0.80)
train_features = features_df.iloc[:split_idx]
test_features  = features_df.iloc[split_idx:]
train_prices   = raw_df['Close'].loc[train_features.index]
test_prices    = raw_df['Close'].loc[test_features.index]

print(f"\nTrain period: {train_features.index[0].date()} → {train_features.index[-1].date()} ({len(train_features)} rows)")
print(f"Test period : {test_features.index[0].date()} → {test_features.index[-1].date()} ({len(test_features)} rows)")

def plot_feature_overview(raw_df: pd.DataFrame, features_df: pd.DataFrame) -> None:
    """
    Visualise the raw price series alongside the engineered features.

    This plot lets you sanity-check that features look sensible before
    feeding them into the trading environment.

    Parameters
    ----------
    raw_df : pd.DataFrame
        Original OHLCV DataFrame.
    features_df : pd.DataFrame
        Feature DataFrame produced by engineer_features().

    Returns
    -------
    None  (displays matplotlib figure)
    """
    fig, axes = plt.subplots(4, 1, figsize=(14, 14), sharex=True)
    fig.suptitle('Market Data Overview', fontsize=16, fontweight='bold', y=1.01)

    # ── Panel 1: Close price ───────────────────────────────────────────────────
    common_idx = features_df.index  # Use feature index (NaN rows dropped)
    axes[0].plot(raw_df['Close'].loc[common_idx], color='steelblue', linewidth=1)
    axes[0].set_ylabel('Price ($)')
    axes[0].set_title('Close Price')

    # ── Panel 2: RSI ──────────────────────────────────────────────────────────
    axes[1].plot(features_df['rsi'], color='darkorange', linewidth=0.8)
    axes[1].axhline(0.7, color='red',   linestyle='--', linewidth=0.8, label='Overbought 0.7')
    axes[1].axhline(0.3, color='green', linestyle='--', linewidth=0.8, label='Oversold 0.3')
    axes[1].set_ylabel('RSI (normalised)')
    axes[1].set_title('RSI Momentum')
    axes[1].legend(fontsize=8)

    # ── Panel 3: Bollinger %B and SMA ratios ──────────────────────────────────
    axes[2].plot(features_df['bb_pct'],        color='purple',    linewidth=0.8, label='BB %B')
    axes[2].plot(features_df['sma_ratio_long'], color='teal',     linewidth=0.8, label='SMA ratio (50)')
    axes[2].axhline(1.0, color='grey', linestyle=':', linewidth=0.7)
    axes[2].set_ylabel('Ratio')
    axes[2].set_title('Bollinger %B & SMA Ratio')
    axes[2].legend(fontsize=8)

    # ── Panel 4: Volatility ───────────────────────────────────────────────────
    axes[3].fill_between(features_df.index, features_df['volatility'],
                         alpha=0.5, color='crimson', label='Rolling Volatility')
    axes[3].set_ylabel('Volatility')
    axes[3].set_title('Rolling Volatility (10-day std of log-returns)')
    axes[3].legend(fontsize=8)

    plt.tight_layout()
    plt.show()


plot_feature_overview(raw_df, features_df)

class TradingEnvironment(gym.Env):
    """
    Custom single-asset trading environment following the Gymnasium API.

    The agent steps through historical price data day by day, choosing
    to buy, hold, or sell. Profit/loss is tracked and returned as rewards.

    This class inherits from gym.Env, which enforces the standard interface:
        - reset()  → returns initial observation
        - step()   → returns (observation, reward, terminated, truncated, info)
        - render() → optional visualisation

    Parameters
    ----------
    features : pd.DataFrame
        Feature matrix from engineer_features().
    prices : pd.Series
        Close price series aligned with features.index.
    initial_balance : float
        Starting cash balance in dollars (default $10,000).
    transaction_cost : float
        Proportional cost per trade as a fraction (default 0.001 = 0.1%).
    lookback_window : int
        Number of past timesteps stacked into the observation (default 1 = no history).

    Observation Space
    -----------------
    Box(n_features + 3,) — technical features + portfolio_value + position + steps_left

    Action Space
    ------------
    Discrete(3) — 0: Hold, 1: Buy, 2: Sell

    Example
    -------
    >>> env = TradingEnvironment(train_features, train_prices)
    >>> obs, info = env.reset()
    >>> obs, reward, done, truncated, info = env.step(action=1)
    """

    metadata = {'render_modes': ['human']}

    # ── Action constants for readability ──────────────────────────────────────
    HOLD = 0
    BUY  = 1
    SELL = 2

    def __init__(
        self,
        features: pd.DataFrame,
        prices: pd.Series,
        initial_balance: float = 10_000.0,
        transaction_cost: float = 0.001
    ):
        super().__init__()

        # ── Store data ────────────────────────────────────────────────────────
        self.features         = features.values.astype(np.float32)  # Convert to numpy for speed
        self.prices           = prices.values.astype(np.float32)
        self.initial_balance  = initial_balance
        self.transaction_cost = transaction_cost
        self.n_steps          = len(self.features)
        self.n_features       = self.features.shape[1]

        # ── Define action space: 3 discrete actions ───────────────────────────
        self.action_space = spaces.Discrete(3)

        # ── Define observation space ──────────────────────────────────────────
        # Observation = [market_features (n_features), portfolio_val, position, steps_left]
        obs_dim = self.n_features + 3
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32
        )

        # ── Internal state (initialised in reset) ─────────────────────────────
        self.current_step   = 0
        self.balance        = initial_balance  # Cash on hand
        self.shares_held    = 0.0              # Number of shares currently owned
        self.portfolio_value = initial_balance  # Total value = cash + shares * price
        self.position       = 0               # 0 = flat, 1 = long
        self.trade_history  = []              # Log of all trades
        self._prev_portfolio = initial_balance  # Initialize _prev_portfolio here

    # ─────────────────────────────────────────────────────────────────────────
    def _get_observation(self) -> np.ndarray:
        """
        Construct the observation vector for the current timestep.

        Combines raw market features with portfolio state information.
        All values normalised so no single feature dominates.

        Returns
        -------
        np.ndarray of shape (n_features + 3,)
        """
        market_feats = self.features[self.current_step]  # Shape: (n_features,)

        # Normalise portfolio value relative to initial capital
        norm_portfolio = self.portfolio_value / self.initial_balance

        # Position: 0.0 = flat, 1.0 = fully long
        norm_position = float(self.position)

        # Fraction of episode remaining — helps agent reason about time horizon
        steps_remaining = (self.n_steps - self.current_step) / self.n_steps

        return np.concatenate([
            market_feats,
            [norm_portfolio, norm_position, steps_remaining]
        ]).astype(np.float32)

    # ─────────────────────────────────────────────────────────────────────────
    def reset(
        self,
        seed: Optional[int] = None,
        options: Optional[dict] = None
    ) -> Tuple[np.ndarray, dict]:
        """
        Reset the environment to the beginning of the episode.

        Called at the start of every new training episode.

        Returns
        -------
        observation : np.ndarray
            Initial observation.
        info : dict
            Diagnostic information (empty dict for compatibility).
        """
        super().reset(seed=seed)

        self.current_step    = 0
        self.balance         = self.initial_balance
        self.shares_held     = 0.0
        self.portfolio_value = self.initial_balance
        self.position        = 0
        self.trade_history   = []
        self._prev_portfolio = self.initial_balance  # Track previous value for reward calc

        return self._get_observation(), {}

    # ─────────────────────────────────────────────────────────────────────────
    def step(self, action: int) -> Tuple[np.ndarray, float, bool, bool, dict]:
        """
        Execute one timestep (one trading day) in the environment.

        Applies the chosen action (buy/hold/sell), updates portfolio value,
        computes reward, and advances the clock by one day.

        Parameters
        ----------
        action : int
            0 = Hold, 1 = Buy, 2 = Sell.

        Returns
        -------
        observation : np.ndarray
        reward      : float
        terminated  : bool  — True if episode ends naturally (end of data)
        truncated   : bool  — True if episode cut short (not used here)
        info        : dict  — Diagnostic data for logging
        """
        current_price = self.prices[self.current_step]
        trade_cost    = 0.0  # Will be non-zero only when trading

        # ── Execute action ─────────────────────────────────────────────────────
        if action == self.BUY and self.position == 0:
            # Invest all available cash into shares
            cost_basis       = self.balance * (1 - self.transaction_cost)  # Deduct trading fee
            self.shares_held = cost_basis / current_price
            trade_cost       = self.balance * self.transaction_cost
            self.balance     = 0.0
            self.position    = 1
            self.trade_history.append(('BUY', self.current_step, current_price))

        elif action == self.SELL and self.position == 1:
            # Liquidate all shares to cash
            proceeds      = self.shares_held * current_price * (1 - self.transaction_cost)
            trade_cost    = self.shares_held * current_price * self.transaction_cost
            self.balance  = proceeds
            self.shares_held = 0.0
            self.position = 0
            self.trade_history.append(('SELL', self.current_step, current_price))

        # ── Update portfolio value ─────────────────────────────────────────────
        self.portfolio_value = self.balance + self.shares_held * current_price

        # ── Compute reward ─────────────────────────────────────────────────────
        # Log return of portfolio value — proportional P&L, handles compounding
        if self._prev_portfolio > 0:
            reward = np.log(self.portfolio_value / self._prev_portfolio)
        else:
            reward = 0.0

        # Small penalty for holding a losing position (encourages active management)
        # (optional: uncomment to enable)
        # if self.position == 1 and reward < 0:
        #     reward *= 1.1  # Amplify pain of holding through losses

        self._prev_portfolio = self.portfolio_value

        # ── Advance time ───────────────────────────────────────────────────────
        self.current_step += 1
        terminated = self.current_step >= self.n_steps - 1  # End of episode
        truncated  = False

        info = {
            'portfolio_value': self.portfolio_value,
            'position':        self.position,
            'price':           current_price,
            'trade_cost':      trade_cost,
            'step':            self.current_step
        }

        return self._get_observation(), reward, terminated, truncated, info

    # ─────────────────────────────────────────────────────────────────────────
    def render(self, mode: str = 'human') -> None:
        """Print a one-line status update to stdout."""
        print(
            f"Step {self.current_step:4d} | "
            f"Price ${self.prices[self.current_step]:.2f} | "
            f"Portfolio ${self.portfolio_value:,.2f} | "
            f"Position: {'LONG' if self.position else 'FLAT'}"
        )


# ── Sanity check: step through environment with random actions ─────────────
env_test = TradingEnvironment(train_features, train_prices)
obs, info = env_test.reset()

print(f"Observation shape : {obs.shape}")
print(f"Action space      : {env_test.action_space}")
print(f"Obs space         : {env_test.observation_space}")
print()

# Run 5 random steps
for i in range(5):
    action = env_test.action_space.sample()
    obs, reward, done, trunc, info = env_test.step(action)
    print(f"  Step {i+1}: action={action}, reward={reward:.6f}, portfolio=${info['portfolio_value']:,.2f}")

class ActorCriticNetwork(nn.Module):
    """
    Shared-backbone actor-critic neural network for PPO.

    Architecture
    ------------
    Input (obs) → Shared Backbone (2 × FC + Tanh)
                → Actor Head  → action logits → Categorical distribution
                → Critic Head → scalar state value V(s)

    Sharing the backbone means both actor and critic learn a common
    representation of the market, which improves sample efficiency.

    Parameters
    ----------
    obs_dim : int
        Dimension of the observation vector.
    n_actions : int
        Number of discrete actions (3 for buy/hold/sell).
    hidden_dim : int
        Width of each hidden layer (default 256).

    Example
    -------
    >>> net = ActorCriticNetwork(obs_dim=9, n_actions=3)
    >>> dist, value = net(torch.randn(1, 9))
    >>> action = dist.sample()
    >>> log_prob = dist.log_prob(action)
    """

    def __init__(self, obs_dim: int, n_actions: int, hidden_dim: int = 256):
        super(ActorCriticNetwork, self).__init__()

        # ── Shared backbone: learns a rich representation of market state ──────
        self.backbone = nn.Sequential(
            nn.Linear(obs_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh()
        )

        # ── Actor head: outputs unnormalised scores (logits) for each action ───
        # We use logits (not probabilities) because Categorical handles the softmax
        self.actor_head = nn.Linear(hidden_dim, n_actions)

        # ── Critic head: scalar state-value estimate V(s) ──────────────────────
        self.critic_head = nn.Linear(hidden_dim, 1)

        # ── Weight initialisation: orthogonal init is standard for RL ─────────
        self._init_weights()

    def _init_weights(self) -> None:
        """
        Initialise weights with orthogonal initialisation.

        Orthogonal initialisation preserves variance through layers and
        is recommended for policy gradient methods (Schulman et al., 2015).
        Smaller gain for actor output ensures initial policies are near-uniform
        (high entropy = good exploration at the start).
        """
        for module in self.backbone.modules():
            if isinstance(module, nn.Linear):
                nn.init.orthogonal_(module.weight, gain=np.sqrt(2))
                nn.init.constant_(module.bias, 0.0)

        # Actor: small gain → near-uniform initial distribution → maximum exploration
        nn.init.orthogonal_(self.actor_head.weight, gain=0.01)
        nn.init.constant_(self.actor_head.bias, 0.0)

        # Critic: standard gain for value estimation
        nn.init.orthogonal_(self.critic_head.weight, gain=1.0)
        nn.init.constant_(self.critic_head.bias, 0.0)

    def forward(self, obs: torch.Tensor) -> Tuple[Categorical, torch.Tensor]:
        """
        Forward pass: given observations, return action distribution and value.

        Parameters
        ----------
        obs : torch.Tensor of shape (batch_size, obs_dim)
            Batch of observations.

        Returns
        -------
        dist : torch.distributions.Categorical
            Probability distribution over actions. Sample with dist.sample().
        value : torch.Tensor of shape (batch_size, 1)
            Critic's estimate of V(s) for each observation.
        """
        # Extract shared features from the market state
        features = self.backbone(obs)

        # Actor: logits → Categorical distribution
        action_logits = self.actor_head(features)
        dist          = Categorical(logits=action_logits)  # Handles softmax internally

        # Critic: scalar value estimate
        value = self.critic_head(features)

        return dist, value

    def get_action_and_value(self, obs: torch.Tensor, action: Optional[torch.Tensor] = None):
        """
        Convenience method used inside the PPO update loop.

        If action is provided (training), computes log-prob and entropy of that action.
        If not provided (acting), samples a new action.

        Parameters
        ----------
        obs    : torch.Tensor
        action : torch.Tensor or None

        Returns
        -------
        action   : torch.Tensor
        log_prob : torch.Tensor  — log π(action | obs)
        entropy  : torch.Tensor  — H[π(·|obs)] for entropy regularisation
        value    : torch.Tensor  — V(obs)
        """
        dist, value = self.forward(obs)

        if action is None:
            action = dist.sample()  # Sample during rollout collection

        log_prob = dist.log_prob(action)  # log probability of the chosen action
        entropy  = dist.entropy()         # Entropy encourages exploration

        return action, log_prob, entropy, value


# ── Demonstration ─────────────────────────────────────────────────────────────
obs_dim   = env_test.observation_space.shape[0]
n_actions = env_test.action_space.n

net = ActorCriticNetwork(obs_dim=obs_dim, n_actions=n_actions).to(DEVICE)

# Forward pass with a dummy batch of 4 observations
dummy_obs = torch.randn(4, obs_dim).to(DEVICE)
dist, value = net(dummy_obs)

print(f"Network architecture:\n{net}")
print(f"\nDummy forward pass:")
print(f"  Sampled actions  : {dist.sample()}")
print(f"  Action probs     : {dist.probs.detach().cpu().numpy().round(3)}")
print(f"  Value estimates  : {value.detach().cpu().numpy().flatten().round(4)}")
print(f"\nTotal parameters : {sum(p.numel() for p in net.parameters()):,}")

class PPORolloutBuffer:
    """
    Storage buffer for one PPO rollout (on-policy experience collection).

    Stores transitions and computes Generalised Advantage Estimates (GAE)
    at the end of each rollout before the policy update.

    GAE Formula (Schulman et al., 2016)
    ------------------------------------
    δ_t  = r_t + γ * V(s_{t+1}) * (1 - done_t) - V(s_t)    [TD error]
    A_t  = Σ_{l=0}^{∞} (γλ)^l * δ_{t+l}                   [GAE]

    λ=1  → pure Monte Carlo (high variance, low bias)
    λ=0  → pure TD(0) (low variance, high bias)
    λ=0.95 → standard trade-off in practice

    Parameters
    ----------
    rollout_steps : int
        Number of timesteps to collect per rollout (e.g., 2048).
    obs_dim : int
        Dimension of the observation vector.
    gamma : float
        Discount factor γ — how much to value future rewards (default 0.99).
    gae_lambda : float
        GAE smoothing parameter λ (default 0.95).
    device : torch.device
        Compute device for tensors.

    Example
    -------
    >>> buf = PPORolloutBuffer(rollout_steps=512, obs_dim=9)
    >>> buf.store(obs, action, reward, done, value, log_prob)
    >>> buf.compute_advantages(last_value, last_done)
    >>> batches = buf.get_batches(batch_size=64)
    """

    def __init__(
        self,
        rollout_steps: int,
        obs_dim:       int,
        gamma:         float = 0.99,
        gae_lambda:    float = 0.95,
        device:        torch.device = DEVICE
    ):
        self.rollout_steps = rollout_steps
        self.gamma         = gamma
        self.gae_lambda    = gae_lambda
        self.device        = device
        self.ptr           = 0  # Pointer to current storage position

        # ── Pre-allocate tensors for efficiency ───────────────────────────────
        # Pre-allocation avoids repeated memory allocation during rollout
        self.obs       = torch.zeros(rollout_steps, obs_dim,    device=device)
        self.actions   = torch.zeros(rollout_steps,             device=device, dtype=torch.long)
        self.rewards   = torch.zeros(rollout_steps,             device=device)
        self.dones     = torch.zeros(rollout_steps,             device=device)
        self.values    = torch.zeros(rollout_steps,             device=device)
        self.log_probs = torch.zeros(rollout_steps,             device=device)
        self.advantages = torch.zeros(rollout_steps,            device=device)
        self.returns   = torch.zeros(rollout_steps,             device=device)

    def store(
        self,
        obs:      np.ndarray,
        action:   int,
        reward:   float,
        done:     bool,
        value:    float,
        log_prob: float
    ) -> None:
        """
        Store a single (s, a, r, done, V, log π) transition in the buffer.

        Parameters
        ----------
        obs      : np.ndarray — current observation
        action   : int        — action taken
        reward   : float      — reward received
        done     : bool       — whether the episode ended
        value    : float      — critic's value estimate V(s)
        log_prob : float      — log probability of the action log π(a|s)
        """
        idx = self.ptr  # Current write position
        self.obs[idx]       = torch.tensor(obs,      dtype=torch.float32, device=self.device)
        self.actions[idx]   = torch.tensor(action,   dtype=torch.long,    device=self.device)
        self.rewards[idx]   = torch.tensor(reward,   dtype=torch.float32, device=self.device)
        self.dones[idx]     = torch.tensor(done,     dtype=torch.float32, device=self.device)
        self.values[idx]    = torch.tensor(value,    dtype=torch.float32, device=self.device)
        self.log_probs[idx] = torch.tensor(log_prob, dtype=torch.float32, device=self.device)
        self.ptr           += 1

    def compute_advantages(self, last_value: float, last_done: bool) -> None:
        """
        Compute GAE advantages and lambda-returns (targets for the critic).

        This must be called after the rollout is complete but before sampling batches.
        Works backwards through the buffer (standard GAE implementation).

        Parameters
        ----------
        last_value : float — V(s_T+1), the bootstrap value from the next state
        last_done  : bool  — whether the last step ended the episode
        """
        last_gae = 0.0  # GAE accumulator, initialised to zero
        last_val = last_value
        last_done_flag = float(last_done)

        # ── Backward pass through the buffer ──────────────────────────────────
        for t in reversed(range(self.rollout_steps)):
            # Mask future value if episode ended at this step
            next_non_terminal = 1.0 - (self.dones[t].item() if t < self.rollout_steps - 1 else last_done_flag)
            next_value        = (self.values[t + 1].item() if t < self.rollout_steps - 1 else last_val)

            # TD error δ_t = r_t + γ * V(s_{t+1}) * (1 - done) - V(s_t)
            delta    = self.rewards[t] + self.gamma * next_value * next_non_terminal - self.values[t]

            # Accumulate GAE: A_t = δ_t + γλ * A_{t+1}
            last_gae = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae
            self.advantages[t] = last_gae

        # Lambda-returns = advantage + value (used as critic target)
        self.returns = self.advantages + self.values

        # ── Normalise advantages ───────────────────────────────────────────────
        # Zero-mean, unit-variance advantages reduce variance in policy gradient
        adv_mean = self.advantages.mean()
        adv_std  = self.advantages.std() + 1e-8
        self.advantages = (self.advantages - adv_mean) / adv_std

    def get_batches(self, batch_size: int):
        """
        Yield random mini-batches for PPO update epochs.

        Shuffles indices and yields non-overlapping mini-batches.
        Multiple calls to this (for different epochs) all use the same
        stored rollout data.

        Parameters
        ----------
        batch_size : int
            Size of each mini-batch.

        Yields
        ------
        Tuple of (obs, actions, log_probs, advantages, returns) tensors.
        """
        indices = torch.randperm(self.rollout_steps, device=self.device)  # Shuffle

        for start in range(0, self.rollout_steps, batch_size):
            batch_idx = indices[start: start + batch_size]
            yield (
                self.obs[batch_idx],
                self.actions[batch_idx],
                self.log_probs[batch_idx],
                self.advantages[batch_idx],
                self.returns[batch_idx]
            )

    def reset(self) -> None:
        """Reset the write pointer for the next rollout."""
        self.ptr = 0


print("PPORolloutBuffer defined successfully.")
print()
# Quick size check
buf_test = PPORolloutBuffer(rollout_steps=512, obs_dim=obs_dim)
mem_bytes = sum(t.element_size() * t.nelement() for t in [
    buf_test.obs, buf_test.actions, buf_test.rewards,
    buf_test.dones, buf_test.values, buf_test.log_probs
])
print(f"Buffer memory footprint (512 steps): {mem_bytes / 1024:.1f} KB")

def compute_ppo_loss(
    network:     ActorCriticNetwork,
    obs:         torch.Tensor,
    actions:     torch.Tensor,
    old_log_probs: torch.Tensor,
    advantages:  torch.Tensor,
    returns:     torch.Tensor,
    clip_eps:    float = 0.2,
    vf_coef:     float = 0.5,
    entropy_coef: float = 0.01
) -> Tuple[torch.Tensor, dict]:
    """
    Compute the total PPO loss for one mini-batch update.

    This function implements the full clipped PPO objective:
        L_total = -L_clip + vf_coef * L_value - entropy_coef * H[π]

    Parameters
    ----------
    network       : ActorCriticNetwork — the network being updated
    obs           : Tensor (B, obs_dim) — observations
    actions       : Tensor (B,)         — actions taken during rollout
    old_log_probs : Tensor (B,)         — log π_old(a|s), computed during rollout
    advantages    : Tensor (B,)         — GAE advantages (normalised)
    returns       : Tensor (B,)         — lambda-return targets for critic
    clip_eps      : float               — PPO clipping range ε (default 0.2)
    vf_coef       : float               — value function loss coefficient c₁ (default 0.5)
    entropy_coef  : float               — entropy bonus coefficient c₂ (default 0.01)

    Returns
    -------
    total_loss : torch.Tensor — scalar loss to backpropagate
    metrics    : dict         — diagnostic breakdown for logging

    Example
    -------
    >>> loss, metrics = compute_ppo_loss(net, obs_batch, act_batch, ...)
    >>> loss.backward()
    """
    # ── Get NEW log probs and values from the network (after update so far) ───
    _, new_log_probs, entropy, new_values = network.get_action_and_value(obs, actions)
    new_values = new_values.squeeze(-1)  # Shape: (B,) from (B, 1)

    # ── Probability ratio r_t(θ) = π_new(a|s) / π_old(a|s) ─────────────────
    # Computed in log space for numerical stability, then exponentiated
    log_ratio = new_log_probs - old_log_probs  # log(new/old) = log(new) - log(old)
    ratio     = torch.exp(log_ratio)

    # ── Approximate KL divergence (used for diagnostics, not directly in loss) ─
    # If KL > 0.02, the update is too large and training may be unstable
    approx_kl = ((ratio - 1) - log_ratio).mean().item()

    # ── Clipped surrogate loss ─────────────────────────────────────────────────
    # Unclipped: ratio * advantage (standard policy gradient)
    surr1 = ratio * advantages

    # Clipped: cap the ratio in [1-ε, 1+ε] to prevent too-large updates
    surr2 = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * advantages

    # Take the minimum (pessimistic bound) — this is the key PPO innovation
    # The min prevents the loss from benefiting from going beyond the clip range
    policy_loss = -torch.min(surr1, surr2).mean()  # Negative because we maximise

    # ── Value function loss (MSE between predicted and target returns) ─────────
    value_loss = nn.functional.mse_loss(new_values, returns)

    # ── Entropy bonus: H[π] = -Σ π(a) log π(a) ───────────────────────────────
    # Maximising entropy (minimising -entropy) prevents premature convergence
    entropy_loss = -entropy.mean()  # Negative to turn maximisation into minimisation

    # ── Total loss ─────────────────────────────────────────────────────────────
    total_loss = policy_loss + vf_coef * value_loss + entropy_coef * entropy_loss

    # ── Clip fraction: what % of samples exceeded the clip boundary? ───────────
    clip_fraction = ((ratio - 1).abs() > clip_eps).float().mean().item()

    metrics = {
        'total_loss':     total_loss.item(),
        'policy_loss':    policy_loss.item(),
        'value_loss':     value_loss.item(),
        'entropy':        -entropy_loss.item(),  # Return positive entropy for readability
        'approx_kl':      approx_kl,
        'clip_fraction':  clip_fraction
    }

    return total_loss, metrics


print(" PPO loss function defined.")
print()
print("Key hyperparameters explained:")
print(f"  clip_eps     = 0.2  → allows ratio r_t ∈ [0.8, 1.2] before clipping")
print(f"  vf_coef      = 0.5  → weight of value loss relative to policy loss")
print(f"  entropy_coef = 0.01 → encourages ~1% of loss from exploration bonus")

class PPOAgent:
    """
    Full PPO agent for the trading environment.

    Wraps the actor-critic network, rollout buffer, and PPO update logic
    into a single interface. Handles rollout collection, advantage estimation,
    and multi-epoch mini-batch updates.

    Parameters
    ----------
    obs_dim       : int   — observation dimension
    n_actions     : int   — number of discrete actions
    hidden_dim    : int   — hidden layer width (default 256)
    lr            : float — Adam learning rate (default 3e-4)
    gamma         : float — discount factor (default 0.99)
    gae_lambda    : float — GAE lambda (default 0.95)
    rollout_steps : int   — steps per rollout (default 1024)
    n_epochs      : int   — update epochs per rollout (default 10)
    batch_size    : int   — mini-batch size (default 128)
    clip_eps      : float — PPO clipping ε (default 0.2)
    vf_coef       : float — value loss weight (default 0.5)
    entropy_coef  : float — entropy bonus weight (default 0.01)
    max_grad_norm : float — gradient clipping norm (default 0.5)
    device        : torch.device

    Example
    -------
    >>> agent = PPOAgent(obs_dim=9, n_actions=3)
    >>> agent.collect_rollout(env, n_steps=1024)
    >>> metrics = agent.update()
    """

    def __init__(
        self,
        obs_dim:       int,
        n_actions:     int,
        hidden_dim:    int   = 256,
        lr:            float = 3e-4,
        gamma:         float = 0.99,
        gae_lambda:    float = 0.95,
        rollout_steps: int   = 1024,
        n_epochs:      int   = 10,
        batch_size:    int   = 128,
        clip_eps:      float = 0.2,
        vf_coef:       float = 0.5,
        entropy_coef:  float = 0.01,
        max_grad_norm: float = 0.5,
        device:        torch.device = DEVICE
    ):
        self.device        = device
        self.rollout_steps = rollout_steps
        self.n_epochs      = n_epochs
        self.batch_size    = batch_size
        self.clip_eps      = clip_eps
        self.vf_coef       = vf_coef
        self.entropy_coef  = entropy_coef
        self.max_grad_norm = max_grad_norm

        # ── Network ───────────────────────────────────────────────────────────
        self.network = ActorCriticNetwork(obs_dim, n_actions, hidden_dim).to(device)

        # ── Optimiser: Adam with default betas ─────────────────────────────────
        self.optimiser = optim.Adam(self.network.parameters(), lr=lr, eps=1e-5)

        # ── Rollout buffer ────────────────────────────────────────────────────
        self.buffer = PPORolloutBuffer(
            rollout_steps=rollout_steps,
            obs_dim=obs_dim,
            gamma=gamma,
            gae_lambda=gae_lambda,
            device=device
        )

        # ── Internal state ────────────────────────────────────────────────────
        self._current_obs  = None  # Carry observation across rollout boundaries
        self._current_done = False
        self.training_metrics: List[dict] = []  # Track all update metrics

    # ─────────────────────────────────────────────────────────────────────────
    @torch.no_grad()
    def act(self, obs: np.ndarray) -> Tuple[int, float, float]:
        """
        Sample an action from the current policy (no gradient computation).

        Used during rollout collection and evaluation.

        Parameters
        ----------
        obs : np.ndarray of shape (obs_dim,)

        Returns
        -------
        action   : int   — sampled action
        log_prob : float — log π(action | obs)
        value    : float — V(obs) from critic
        """
        obs_tensor = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
        action, log_prob, _, value = self.network.get_action_and_value(obs_tensor)

        return action.item(), log_prob.item(), value.item()

    # ─────────────────────────────────────────────────────────────────────────
    def collect_rollout(self, env: TradingEnvironment) -> dict:
        """
        Run the current policy in the environment for `rollout_steps` steps.

        Stores all transitions in the buffer and computes GAE advantages
        at the end. Handles environment resets transparently.

        Parameters
        ----------
        env : TradingEnvironment

        Returns
        -------
        rollout_info : dict — episode returns and lengths observed this rollout
        """
        self.buffer.reset()

        # Initialise or reuse the current observation
        if self._current_obs is None:
            self._current_obs, _ = env.reset()
            self._current_done   = False

        episode_rewards = []  # Track returns per completed episode
        ep_reward       = 0.0

        for step in range(self.rollout_steps):
            obs = self._current_obs

            # ── Sample action from policy ─────────────────────────────────────
            action, log_prob, value = self.act(obs)

            # ── Step environment ──────────────────────────────────────────────
            next_obs, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            # ── Store transition ──────────────────────────────────────────────
            self.buffer.store(obs, action, reward, done, value, log_prob)
            ep_reward += reward

            # ── Handle episode boundary ───────────────────────────────────────
            if done:
                episode_rewards.append(ep_reward)
                ep_reward = 0.0
                self._current_obs, _ = env.reset()  # Start fresh episode
                self._current_done   = False
            else:
                self._current_obs  = next_obs
                self._current_done = done

        # ── Bootstrap value of the last state for GAE computation ─────────────
        with torch.no_grad():
            _, _, last_value = self.act(self._current_obs)

        self.buffer.compute_advantages(last_value, self._current_done)

        return {'episode_rewards': episode_rewards}

    # ─────────────────────────────────────────────────────────────────────────
    def update(self) -> dict:
        """
        Perform PPO update: multiple epochs of mini-batch gradient descent.

        Runs `n_epochs` passes over the rollout buffer, sampling mini-batches
        of size `batch_size` each time, and applying the clipped PPO loss.

        Returns
        -------
        update_metrics : dict — mean losses, KL, clip fraction over this update
        """
        all_metrics = []

        for epoch in range(self.n_epochs):
            for obs_b, act_b, logp_b, adv_b, ret_b in self.buffer.get_batches(self.batch_size):

                # Compute PPO loss for this mini-batch
                loss, metrics = compute_ppo_loss(
                    self.network, obs_b, act_b, logp_b, adv_b, ret_b,
                    clip_eps=self.clip_eps,
                    vf_coef=self.vf_coef,
                    entropy_coef=self.entropy_coef
                )

                # Gradient update
                self.optimiser.zero_grad()
                loss.backward()

                # Clip gradients to prevent large destabilising steps
                nn.utils.clip_grad_norm_(self.network.parameters(), self.max_grad_norm)

                self.optimiser.step()
                all_metrics.append(metrics)

            # ── Early stopping if KL divergence is too large ─────────────────
            # If the policy changed too much this epoch, stop updating
            mean_kl = np.mean([m['approx_kl'] for m in all_metrics[-10:]])  # Last 10 batches
            if mean_kl > 0.02:
                break  # Prevent policy from drifting too far

        # ── Aggregate metrics across all mini-batch updates ───────────────────
        agg = {key: np.mean([m[key] for m in all_metrics]) for key in all_metrics[0]}
        self.training_metrics.append(agg)
        return agg

    def save(self, path: str) -> None:
        """Persist network weights and optimiser state to disk."""
        torch.save({
            'network_state':   self.network.state_dict(),
            'optimiser_state': self.optimiser.state_dict()
        }, path)
        print(f" Agent saved to {path}")

    def load(self, path: str) -> None:
        """Load network weights and optimiser state from disk."""
        checkpoint = torch.load(path, map_location=self.device)
        self.network.load_state_dict(checkpoint['network_state'])
        self.optimiser.load_state_dict(checkpoint['optimiser_state'])
        print(f" Agent loaded from {path}")


# ── Instantiate agent ──────────────────────────────────────────────────────
agent = PPOAgent(
    obs_dim       = obs_dim,
    n_actions     = n_actions,
    hidden_dim    = 256,
    lr            = 3e-4,
    gamma         = 0.99,
    gae_lambda    = 0.95,
    rollout_steps = 1024,
    n_epochs      = 10,
    batch_size    = 128,
    clip_eps      = 0.2,
    vf_coef       = 0.5,
    entropy_coef  = 0.01,
    max_grad_norm = 0.5,
    device        = DEVICE
)

print(f"PPO Agent ready.")
print(f"Network params   : {sum(p.numel() for p in agent.network.parameters()):,}")
print(f"Rollout steps    : {agent.rollout_steps}")
print(f"Update epochs    : {agent.n_epochs}")
print(f"Mini-batch size  : {agent.batch_size}")

def train_ppo_agent(
    agent:          PPOAgent,
    env:            TradingEnvironment,
    total_steps:    int = 200_000,
    log_interval:   int = 10,
    save_best:      bool = True,
    checkpoint_path: str = '/tmp/ppo_trading_best.pt'
) -> Dict[str, List]:
    """
    Run the full PPO training loop.

    Alternates between rollout collection and policy update until the
    total_steps budget is exhausted. Logs progress and saves the best
    checkpoint based on average episode return.

    Parameters
    ----------
    agent           : PPOAgent — the agent to train
    env             : TradingEnvironment — training environment
    total_steps     : int — total environment steps to train for (default 200k)
    log_interval    : int — print progress every N updates (default 10)
    save_best       : bool — save checkpoint when return improves (default True)
    checkpoint_path : str — path to save best checkpoint

    Returns
    -------
    history : dict — training history with keys:
        'steps'           : cumulative environment steps
        'episode_returns' : mean episode return per update
        'policy_loss'     : mean policy loss per update
        'value_loss'      : mean value loss per update
        'entropy'         : mean policy entropy per update
        'approx_kl'       : mean KL divergence per update

    Example
    -------
    >>> history = train_ppo_agent(agent, env, total_steps=100_000)
    """
    # ── Training history tracking ─────────────────────────────────────────────
    history = {
        'steps':           [],
        'episode_returns': [],
        'policy_loss':     [],
        'value_loss':      [],
        'entropy':         [],
        'approx_kl':       [],
        'clip_fraction':   []
    }

    n_updates  = total_steps // agent.rollout_steps  # Number of update cycles
    best_return = -np.inf
    steps_done  = 0

    print(f"{'─' * 70}")
    print(f"  PPO Training | {total_steps:,} steps | {n_updates} updates")
    print(f"{'─' * 70}")
    print(f"{'Update':>7} {'Steps':>9} {'Ep.Return':>11} {'Policy L':>10} "
          f"{'Value L':>9} {'Entropy':>9} {'KL':>8}")
    print(f"{'─' * 70}")

    recent_returns = deque(maxlen=50)  # Rolling window for best-checkpoint logic

    for update_idx in range(1, n_updates + 1):

        # ── Phase 1: Collect rollout ──────────────────────────────────────────
        rollout_info = agent.collect_rollout(env)
        steps_done  += agent.rollout_steps

        # ── Phase 2: Update policy ────────────────────────────────────────────
        update_metrics = agent.update()

        # ── Track episode returns from this rollout ───────────────────────────
        ep_returns = rollout_info['episode_rewards']
        if ep_returns:  # At least one episode completed
            mean_return = np.mean(ep_returns)
            recent_returns.extend(ep_returns)
        else:
            mean_return = np.nan  # Mid-episode rollout

        # ── Log to history ────────────────────────────────────────────────────
        history['steps'].append(steps_done)
        history['episode_returns'].append(mean_return)
        history['policy_loss'].append(update_metrics['policy_loss'])
        history['value_loss'].append(update_metrics['value_loss'])
        history['entropy'].append(update_metrics['entropy'])
        history['approx_kl'].append(update_metrics['approx_kl'])
        history['clip_fraction'].append(update_metrics['clip_fraction'])

        # ── Save best checkpoint ──────────────────────────────────────────────
        if save_best and len(recent_returns) >= 5:
            rolling_mean = np.mean(list(recent_returns))
            if rolling_mean > best_return:
                best_return = rolling_mean
                agent.save(checkpoint_path)

        # ── Print progress ────────────────────────────────────────────────────
        if update_idx % log_interval == 0 or update_idx == 1:
            ret_str = f"{mean_return:>11.4f}" if not np.isnan(mean_return) else f"{'N/A':>11}"
            print(
                f"{update_idx:>7d} "
                f"{steps_done:>9,} "
                f"{ret_str} "
                f"{update_metrics['policy_loss']:>10.4f} "
                f"{update_metrics['value_loss']:>9.4f} "
                f"{update_metrics['entropy']:>9.4f} "
                f"{update_metrics['approx_kl']:>8.5f}"
            )

    print(f"{'─' * 70}")
    print(f"Training complete. Best rolling return: {best_return:.4f}")
    return history


# ── Create training environment ────────────────────────────────────────────
train_env = TradingEnvironment(
    features         = train_features,
    prices           = train_prices,
    initial_balance  = 10_000.0,
    transaction_cost = 0.001
)

# ── Run training (200k steps ≈ 1–2 min on CPU) ────────────────────────────
print("Starting training...")
history = train_ppo_agent(
    agent           = agent,
    env             = train_env,
    total_steps     = 200_000,
    log_interval    = 20,
    save_best       = True,
    checkpoint_path = '/tmp/ppo_trading_best.pt'
)