"""
Simplified Trading Environment for SAC Training.

Uses pre-calculated GRU predictions (mu, sigma, p_cal) and actual returns.
"""
import numpy as np
import pandas as pd
import logging

env_logger = logging.getLogger(__name__)

class TradingEnv:
    def __init__(self,
                 mu_predictions: np.ndarray,
                 sigma_predictions: np.ndarray,
                 p_cal_predictions: np.ndarray,
                 actual_returns: np.ndarray,
                 initial_capital: float = 10000.0,
                 transaction_cost: float = 0.0005):
        """
        Initialize the environment.

        Args:
            mu_predictions: Predicted log returns (μ̂).
            sigma_predictions: Predicted volatility (σ̂ = exp(log σ̂)).
            p_cal_predictions: Calibrated probability of price increase (p_cal).
            actual_returns: Actual log returns (y_ret).
            initial_capital: Starting capital for simulation (used notionally in reward).
            transaction_cost: Fractional cost per trade.
        """
        assert len(mu_predictions) == len(sigma_predictions) == len(p_cal_predictions) == len(actual_returns), \
            "All input arrays must have the same length"

        self.mu = mu_predictions
        self.sigma = sigma_predictions
        self.p_cal = p_cal_predictions
        self.actual_returns = actual_returns

        self.initial_capital = initial_capital
        self.transaction_cost = transaction_cost

        self.n_steps = len(actual_returns)
        self.current_step = 0
        self.current_position = 0.0 # Fraction of capital (-1 to 1)
        self.current_capital = initial_capital # Track for info, not used in reward directly

        # State dimension: [mu, sigma, edge, |mu|/sigma, position]
        self.state_dim = 5
        self.action_dim = 1

        env_logger.info(f"TradingEnv initialized with {self.n_steps} steps.")

    def _get_state(self) -> np.ndarray:
        """Construct the state vector for the current step."""
        if self.current_step >= self.n_steps:
             # Handle episode end - return a dummy state or zeros
             return np.zeros(self.state_dim, dtype=np.float32)

        mu_t = self.mu[self.current_step]
        sigma_t = self.sigma[self.current_step]
        p_cal_t = self.p_cal[self.current_step]

        edge_t = 2 * p_cal_t - 1
        z_score_t = np.abs(mu_t) / (sigma_t + 1e-9)

        # State uses position *before* the action for this step is taken
        state = np.array([
            mu_t,
            sigma_t,
            edge_t,
            z_score_t,
            self.current_position
        ], dtype=np.float32)
        return state

    def reset(self) -> np.ndarray:
        """Reset the environment to the beginning."""
        self.current_step = 0
        self.current_position = 0.0
        self.current_capital = self.initial_capital
        env_logger.debug("Environment reset.")
        return self._get_state()

    def step(self, action: float) -> tuple[np.ndarray, float, bool, dict]:
        """
        Execute one time step.

        Args:
            action (float): Agent's desired position size (-1 to 1).

        Returns:
            tuple: (next_state, reward, done, info_dict)
        """
        if self.current_step >= self.n_steps:
            # Should not happen if 'done' is handled correctly, but as safeguard
            env_logger.warning("Step called after environment finished.")
            return self._get_state(), 0.0, True, {}

        # Action is the TARGET position for the *end* of this step
        target_position = np.clip(action, -1.0, 1.0)
        trade_size = target_position - self.current_position

        # Calculate PnL based on position held *during* this step
        step_actual_return = self.actual_returns[self.current_step]
        # Use simple return for PnL calculation: exp(log_ret) - 1
        pnl_fraction = self.current_position * (np.exp(step_actual_return) - 1)

        # Calculate transaction costs for the trade executed now
        cost_fraction = abs(trade_size) * self.transaction_cost

        # Reward is net PnL fraction (doesn't scale with capital directly)
        reward = pnl_fraction - cost_fraction

        # Update internal state for the *next* step
        self.current_position = target_position
        self.current_capital *= (1 + pnl_fraction - cost_fraction) # Update tracked capital
        self.current_step += 1

        # Check if done
        done = self.current_step >= self.n_steps or self.current_capital <= 0

        next_state = self._get_state()
        info = {'capital': self.current_capital, 'position': self.current_position}

        # Log step details periodically
        # if self.current_step % 1000 == 0:
        #     env_logger.debug(f"Step {self.current_step}: Action={action:.2f}, Pos={self.current_position:.2f}, Ret={step_actual_return:.5f}, Rew={reward:.5f}, Cap={self.current_capital:.2f}")

        if done:
             env_logger.info(f"Environment finished at step {self.current_step}. Final Capital: {self.current_capital:.2f}")

        return next_state, reward, done, info

    def close(self):
        """Clean up any resources (if needed)."""
        env_logger.info("TradingEnv closed.")
        pass 