trading_env_eth.py

# trading_env.py

from __future__ import annotations

import numpy as np

# Prefer gymnasium if available (SB3 supports it), fallback to gym
try:
    import gymnasium as gym
    from gymnasium import spaces
    _GYMNASIUM = True
except ImportError:
    import gym
    from gym import spaces
    _GYMNASIUM = False


class EthereumTradingEnv(gym.Env):
    """
    RL Ethereum Trading Environment (Position-Persistent)

    Key properties:
      - Observation: rolling window of features + 3 state features (position, time_in_trade, unrealized_pnl_dollars)
      - Actions:
          0: HOLD (do nothing)
          1: CLOSE (close position if any)
          2..: OPEN (direction + SL + TP), only effective when flat
      - Position persistence: once open, position remains until:
          - agent sends CLOSE, or
          - SL/TP hit intrabar
      - Friction: spread (% of price) + commission + optional slippage
      - Reward:
          - realized PnL (dollars) minus costs on closes
          - optional shaping via delta unrealized PnL while holding
      - Random episode start to reduce memorization / overfit
      
    Note: For crypto, we use dollars instead of pips, and spread is percentage-based
    """

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

    def __init__(
        self,
        df,
        window_size: int = 30,
        sl_options=None,  # Stop loss in dollars, e.g., [20, 50, 100]
        tp_options=None,  # Take profit in dollars, e.g., [30, 75, 150]
        feature_columns = None,
        spread_pct: float = 0.1,               # Spread as % of price (0.1% = typical for ETH)
        commission_dollars: float = 0.5,       # Flat commission per trade in dollars
        max_slippage_pct: float = 0.05,        # Random slippage as % (0.05% max)
        position_size_eth: float = 0.1,        # How much ETH to trade (0.1 ETH per trade)
        reward_scale: float = 0.01,            # Scale rewards (dollars are larger than pips)
        unrealized_delta_weight: float = 0.001, # Shaping weight on delta-unrealized
        random_start: bool = True,
        min_episode_steps: int = 300,
        episode_max_steps: int | None = None,
        feature_mean: np.ndarray | None = None,
        feature_std: np.ndarray | None = None,
        allow_flip: bool = False,
        hold_reward_weight: float = 0.001,     # Smaller for dollar-based rewards
        open_penalty_dollars: float = 1.0,     # Penalty per open in dollars
        time_penalty_dollars: float = 0.05,    # Cost per hour in trade
    ):
        super().__init__()

        self.df = df.reset_index(drop=True)
        self.n_steps = len(self.df)

        if feature_columns is None:
            self.feature_columns = list(self.df.columns)  # fallback: everything
        else:
            self.feature_columns = list(feature_columns)

        if sl_options is None or tp_options is None:
            raise ValueError("sl_options and tp_options must be provided (e.g. [20, 50, 100] for dollars).")
        self.sl_options = list(sl_options)
        self.tp_options = list(tp_options)

        if self.n_steps <= window_size + 2:
            raise ValueError("Dataframe is too short for the given window_size.")

        self.window_size = int(window_size)

        # Crypto-specific: spread as percentage, not pips
        self.spread_pct = float(spread_pct) / 100.0  # Convert to decimal (0.1% -> 0.001)
        self.commission_dollars = float(commission_dollars)
        self.max_slippage_pct = float(max_slippage_pct) / 100.0  # Convert to decimal
        
        # Position sizing for crypto (in ETH units)
        self.position_size_eth = float(position_size_eth)

        # Reward handling (dollar-based)
        self.reward_scale = float(reward_scale)
        self.unrealized_delta_weight = float(unrealized_delta_weight)
        self.hold_reward_weight = float(hold_reward_weight)
        self.open_penalty_dollars = float(open_penalty_dollars)
        self.time_penalty_dollars = float(time_penalty_dollars)

        # Episode handling
        self.random_start = bool(random_start)
        self.min_episode_steps = int(min_episode_steps)
        self.episode_max_steps = episode_max_steps if episode_max_steps is None else int(episode_max_steps)

        # Optional normalization (fit on train only, pass arrays here)
        self.feature_mean = feature_mean
        self.feature_std = feature_std

        self.allow_flip = bool(allow_flip)

        # --- Actions ---
        # 0: HOLD
        # 1: CLOSE
        # 2..: OPEN(direction, sl_dollars, tp_dollars)
        self.action_map = [("HOLD", None, None, None), ("CLOSE", None, None, None)]
        for direction in [0, 1]:  # 0=short, 1=long
            for sl in self.sl_options:
                for tp in self.tp_options:
                    self.action_map.append(("OPEN", direction, float(sl), float(tp)))

        self.action_space = spaces.Discrete(len(self.action_map))

        # Observation features: df columns + 3 state features
        self.base_num_features = len(self.feature_columns)
        self.state_num_features = 3
        self.num_features = self.base_num_features + self.state_num_features

        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(self.window_size, self.num_features),
            dtype=np.float32
        )

        # Internal state
        self._reset_state()

    # ----------------------------
    # Core Helpers
    # ----------------------------

    def _reset_state(self):
        self.current_step = 0
        self.steps_in_episode = 0
        self.terminated = False
        self.truncated = False

        # Position state
        self.position = 0              # 0=flat, +1=long, -1=short
        self.entry_price = None
        self.sl_price = None
        self.tp_price = None
        self.time_in_trade = 0
        self.prev_unrealized_pips = 0.0

        # Accounting
        self.initial_equity_usd = 10000.0
        self.equity_usd = self.initial_equity_usd

        # Logging
        self.equity_curve = []
        self.last_trade_info = None

    def _get_state_features(self):
        # position in [-1,0,1], time normalized, unrealized in dollars (scaled)
        pos = float(self.position)
        t_norm = float(self.time_in_trade) / 1000.0
        unreal_dollars = float(self._compute_unrealized_dollars()) if self.position != 0 else 0.0
        unreal_scaled = unreal_dollars / 100.0  # Scale to reasonable magnitude
        return np.array([pos, t_norm, unreal_scaled], dtype=np.float32)

    def _compute_unrealized_dollars(self):
        """Calculate unrealized P&L in dollars for current position"""
        if self.position == 0 or self.entry_price is None:
            return 0.0
        close_price = float(self.df.loc[self.current_step, "close"])
        if self.position == 1:  # Long position
            price_diff = close_price - self.entry_price
        else:  # Short position
            price_diff = self.entry_price - close_price
        return price_diff * self.position_size_eth

    def _apply_optional_normalization(self, obs: np.ndarray) -> np.ndarray:
        if self.feature_mean is None or self.feature_std is None:
            return obs
        mean = self.feature_mean.reshape(1, 1, -1)
        std = self.feature_std.reshape(1, 1, -1)
        std = np.where(std == 0, 1.0, std)
        return (obs - mean) / std

    def _get_observation(self):
        start = self.current_step - self.window_size
        if start < 0:
            start = 0

        obs_df = self.df.iloc[start:self.current_step].copy()
        # use only selected feature columns for the agent
        obs_df = obs_df[self.feature_columns]

        # If empty (safety), use the first row repeated
        if len(obs_df) == 0:
            base = np.tile(self.df.iloc[0].values.astype(np.float32), (self.window_size, 1))
        else:
            base = obs_df.values.astype(np.float32)
            if base.shape[0] < self.window_size:
                pad_rows = self.window_size - base.shape[0]
                pad = np.tile(base[0], (pad_rows, 1))
                base = np.vstack([pad, base])

        # Append state features (same for each row)
        state_feat = self._get_state_features()
        state_block = np.tile(state_feat, (self.window_size, 1))
        obs = np.hstack([base, state_block]).astype(np.float32)

        # Optional normalization (only if user passes train-fitted mean/std matching obs dims)
        obs = self._apply_optional_normalization(obs)

        return obs

    def _sample_slippage_pct(self) -> float:
        """Sample random slippage as percentage of price"""
        if self.max_slippage_pct <= 0:
            return 0.0
        return float(np.random.uniform(0.0, self.max_slippage_pct))

    def _cost_dollars_round_trip(self, entry_price: float) -> float:
        """Calculate trading costs in dollars for a round-trip trade"""
        # Spread cost (percentage of price * position size)
        spread_cost = entry_price * self.spread_pct * self.position_size_eth
        # Flat commission
        total_cost = spread_cost + self.commission_dollars
        return total_cost

    def _open_position(self, direction: int, sl_dollars: float, tp_dollars: float):
        """
        Open a new position
        - direction: 1=long, 0=short
        - sl_dollars: stop loss in dollars from entry
        - tp_dollars: take profit in dollars from entry
        """
        close_price = float(self.df.loc[self.current_step, "close"])
        slip_pct = self._sample_slippage_pct()
        slip_price = close_price * slip_pct

        if direction == 1:  # Long position
            entry = close_price + slip_price
            # For long: SL is below entry, TP is above
            # Convert dollars to price levels
            sl_price = entry - (sl_dollars / self.position_size_eth)
            tp_price = entry + (tp_dollars / self.position_size_eth)
            self.position = 1
        else:  # Short position
            entry = close_price - slip_price
            # For short: SL is above entry, TP is below
            sl_price = entry + (sl_dollars / self.position_size_eth)
            tp_price = entry - (tp_dollars / self.position_size_eth)
            self.position = -1

        self.entry_price = entry
        self.sl_price = sl_price
        self.tp_price = tp_price
        self.time_in_trade = 0
        self.prev_unrealized_dollars = 0.0

        self.last_trade_info = {
            "event": "OPEN",
            "step": self.current_step,
            "position": self.position,
            "entry_price": float(self.entry_price),
            "sl_price": float(self.sl_price),
            "tp_price": float(self.tp_price),
            "sl_dollars": float(sl_dollars),
            "tp_dollars": float(tp_dollars)
        }

    def _close_position(self, reason: str, exit_price: float):
        """
        Close current position and calculate P&L in dollars
        """
        # Calculate realized P&L in price difference
        if self.position == 1:  # Long
            price_diff = exit_price - self.entry_price
        else:  # Short
            price_diff = self.entry_price - exit_price
        
        # Convert to dollars
        realized_dollars = price_diff * self.position_size_eth

        # Subtract costs (spread + commission)
        cost_dollars = self._cost_dollars_round_trip(self.entry_price)
        net_dollars = realized_dollars - cost_dollars

        # Update equity in USD
        self.equity_usd += net_dollars

        trade_info = {
            "event": "CLOSE",
            "reason": reason,
            "step": self.current_step,
            "position": self.position,
            "entry_price": float(self.entry_price),
            "exit_price": float(exit_price),
            "realized_dollars": float(realized_dollars),
            "cost_dollars": float(cost_dollars),
            "net_dollars": float(net_dollars),
            "equity_usd": float(self.equity_usd),
            "time_in_trade": int(self.time_in_trade),
        }

        # Reset position state
        self.position = 0
        self.entry_price = None
        self.sl_price = None
        self.tp_price = None
        self.time_in_trade = 0
        self.prev_unrealized_dollars = 0.0

        self.last_trade_info = trade_info
        return net_dollars

    def _check_sl_tp_intrabar_and_maybe_close(self) -> float:
        """
        Checks SL/TP on the *next bar* range [Low, High].
        Conservative rule if both touched: assume SL hits first (worst case).
        Returns realized net pips if closed; otherwise None.
        """
        if self.position == 0:
            return None

        # If last bar, close on close
        if self.current_step >= self.n_steps - 2:
            exit_price = float(self.df.loc[self.current_step, "close"])
            net_pips = self._close_position("END_OF_DATA", exit_price)
            return net_pips

        next_high = float(self.df.loc[self.current_step + 1, "high"])
        next_low = float(self.df.loc[self.current_step + 1, "low"])

        if self.position == 1:
            sl_hit = next_low <= self.sl_price
            tp_hit = next_high >= self.tp_price
            if sl_hit and tp_hit:
                # conservative: SL first
                return self._close_position("SL_AND_TP_SAME_BAR_SL_FIRST", self.sl_price)
            elif sl_hit:
                return self._close_position("SL_HIT", self.sl_price)
            elif tp_hit:
                return self._close_position("TP_HIT", self.tp_price)
        else:
            sl_hit = next_high >= self.sl_price
            tp_hit = next_low <= self.tp_price
            if sl_hit and tp_hit:
                return self._close_position("SL_AND_TP_SAME_BAR_SL_FIRST", self.sl_price)
            elif sl_hit:
                return self._close_position("SL_HIT", self.sl_price)
            elif tp_hit:
                return self._close_position("TP_HIT", self.tp_price)

        return None

    # ----------------------------
    # Gym API
    # ----------------------------

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)

        self._reset_state()

        # Choose start
        if self.random_start:
            max_start = self.n_steps - max(self.min_episode_steps, self.window_size) - 2
            if max_start <= self.window_size:
                self.current_step = self.window_size
            else:
                self.current_step = int(np.random.randint(self.window_size, max_start))
        else:
            self.current_step = self.window_size

        self.steps_in_episode = 0
        self.terminated = False
        self.truncated = False

        obs = self._get_observation()

        if _GYMNASIUM:
            return obs, {}
        return obs

    def step(self, action: int):
        if self.terminated or self.truncated:
            # If someone steps after done, just return current obs with 0 reward
            obs = self._get_observation()
            if _GYMNASIUM:
                return obs, 0.0, True, False, {}
            return obs, 0.0, True, {}

        self.steps_in_episode += 1

        # Reward components (in dollars)
        reward_dollars = 0.0
        info = {}

        act_type, direction, sl_dollars, tp_dollars = self.action_map[int(action)]

        # 1) Apply action logic
        if act_type == "HOLD":
            pass

        elif act_type == "CLOSE":
            if self.position != 0:
                # Close at current close (with slippage)
                close_price = float(self.df.loc[self.current_step, "close"])
                slip_pct = self._sample_slippage_pct()
                slip_price = close_price * slip_pct
                exit_price = close_price - slip_price if self.position == 1 else close_price + slip_price
                reward_dollars += self._close_position("MANUAL_CLOSE", exit_price)

        elif act_type == "OPEN":
            if self.position == 0:
                self._open_position(direction=direction, sl_dollars=sl_dollars, tp_dollars=tp_dollars)
                # penalty for opening a trade to discourage overtrading
                reward_dollars -= self.open_penalty_dollars
            else:
                if self.allow_flip:
                    close_price = float(self.df.loc[self.current_step, "close"])
                    reward_dollars += self._close_position("FLIP_CLOSE", close_price)
                    self._open_position(direction=direction, sl_dollars=sl_dollars, tp_dollars=tp_dollars)
                    reward_dollars -= self.open_penalty_dollars


        # 2) If position is open, check SL/TP on next bar intrabar
        realized_now = self._check_sl_tp_intrabar_and_maybe_close()
        if realized_now is not None:
            reward_dollars += realized_now

        # 3) If still open, apply reward shaping
        if self.position != 0:
            self.time_in_trade += 1

            unreal_now = self._compute_unrealized_dollars()
            delta_unreal = unreal_now - self.prev_unrealized_dollars

            # (a) small bonus for holding a winning trade
            #     proportional to current unrealized profit
            if unreal_now > 0:
                reward_dollars += self.hold_reward_weight * unreal_now

            # (b) optional shaping on change in unrealized (can keep small or zero)
            if self.unrealized_delta_weight != 0.0:
                reward_dollars += self.unrealized_delta_weight * delta_unreal

            # (c) small time cost per bar in a trade to avoid infinite holding
            reward_dollars -= self.time_penalty_dollars

            self.prev_unrealized_dollars = unreal_now


        # 4) Advance time
        self.current_step += 1

        # 5) Termination / truncation
        if self.current_step >= self.n_steps - 1:
            self.terminated = True

        if self.episode_max_steps is not None and self.steps_in_episode >= self.episode_max_steps:
            self.truncated = True

        # 6) Log equity
        self.equity_curve.append(float(self.equity_usd))

        # 7) Build observation
        obs = self._get_observation()

        # 8) Final reward scaling
        reward = float(reward_dollars) * self.reward_scale

        # 9) Info
        info.update({
            "equity_usd": float(self.equity_usd),
            "position": int(self.position),
            "time_in_trade": int(self.time_in_trade),
            "reward_dollars": float(reward_dollars),
            "last_trade_info": self.last_trade_info
        })

        if _GYMNASIUM:
            return obs, reward, self.terminated, self.truncated, info
        else:
            done = bool(self.terminated or self.truncated)
            return obs, reward, done, info

    def render(self):
        print(
            f"Step={self.current_step} | Equity=${self.equity_usd:,.2f} | "
            f"Pos={self.position} | Entry={self.entry_price} | SL={self.sl_price} | TP={self.tp_price}"
        )