"""
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
import gymnasium as gym
from omegaconf import DictConfig # Added for config typing

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,
                 bar_imputed_flags: np.ndarray, # Added imputed flags
                 config: DictConfig, # Added config
                 initial_capital: float = 10000.0,
                 transaction_cost: float = 0.0005,
                 reward_scale: float = 100.0,
                 action_penalty_lambda: float = 0.0):
        """
        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).
            bar_imputed_flags: Boolean array indicating if a bar was imputed.
            config: OmegaConf configuration object.
            initial_capital: Starting capital for simulation (used notionally in reward).
            transaction_cost: Fractional cost per trade.
            reward_scale: Multiplier for the reward signal.
            action_penalty_lambda: Coefficient for the action magnitude penalty (λ).
        """
        assert len(mu_predictions) == len(sigma_predictions) == len(p_cal_predictions) == len(actual_returns) == len(bar_imputed_flags), \
            "All input arrays (predictions, returns, imputed_flags) 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.bar_imputed = bar_imputed_flags.astype(bool) # Store imputed flags
        self.config = config # Store config

        self.initial_capital = initial_capital
        self.transaction_cost = transaction_cost
        self.reward_scale = reward_scale
        self.action_penalty_lambda = action_penalty_lambda

        # --- Revision 5: Calculate action penalty based on transaction cost --- #
        if self.transaction_cost > 0:
            # Override passed lambda if cost is positive
            self._internal_action_penalty_lambda = 0.01 / self.transaction_cost
            env_logger.info(f"Using calculated action penalty lambda: {self._internal_action_penalty_lambda:.4f} (0.01 / {self.transaction_cost})")
        else:
            # Use passed value if cost is zero (or default to 0)
            self._internal_action_penalty_lambda = self.action_penalty_lambda 
            env_logger.info(f"Using provided action penalty lambda: {self._internal_action_penalty_lambda:.4f}")
        # --- End Revision 5 --- #

        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

        # --- Define Gym Spaces ---
        self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(self.action_dim,), dtype=np.float32)
        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.state_dim,), dtype=np.float32)
        # --- End Define Gym Spaces ---

        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:
             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]

        # Calculate edge based on p_cal shape (binary vs ternary)
        if isinstance(p_cal_t, (np.ndarray, list)) and len(p_cal_t) == 3:
            # Ternary: edge = max(P(up), P(down)) - P(flat)
            # Assuming order [Down, Flat, Up] for p_cal_t
            edge_t = max(p_cal_t[2], p_cal_t[0]) - p_cal_t[1]
        elif isinstance(p_cal_t, (float, np.number)):
            # Binary: edge = 2 * P(up) - 1
            edge_t = 2 * p_cal_t - 1
        else:
            env_logger.error(f"Unexpected type/shape for p_cal_t at step {self.current_step}: {p_cal_t}. Using edge=0.")
            edge_t = 0.0

        _EPS = 1e-9 # Define epsilon locally
        z_score_t = np.abs(mu_t) / (sigma_t + _EPS)

        # 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)
        """
        info = {'capital': self.current_capital, 'position': self.current_position, 'is_imputed_step_skipped': False}

        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, info

        # --- Handle Imputed Bar --- #
        imputed = self.bar_imputed[self.current_step]
        if imputed:
            mode = self.config.sac.imputed_handling
            env_logger.debug(f"SAC step {self.current_step} on imputed bar: handling={mode}")
            if mode == "skip":
                self.current_step += 1
                next_state = self._get_state() # Get state for the *next* actual step
                # Return 0 reward, not done, but indicate skip for buffer handling
                info['is_imputed_step_skipped'] = True 
                return next_state, 0.0, False, info
            elif mode == "hold":
                # Action is forced to maintain current position
                action = self.current_position 
            elif mode == "penalty":
                # Calculate reward penalty based on config
                target_position_penalty = np.clip(action, -1.0, 1.0)
                reward = -self.config.sac.action_penalty * (target_position_penalty - self.current_position)**2
                # Update position based on agent's intended action (clipped)
                self.current_position = target_position_penalty
                # Update capital notionally (no actual return, only cost if implemented)
                # Cost is implicitly 0 here as there's no trade size if pos doesn't change
                # If penalty mode allowed position change, cost would apply.
                # For simplicity, we don't add cost here for the penalty step.
                self.current_step += 1
                next_state = self._get_state()
                scaled_reward = reward * self.reward_scale # Scale the penalty
                done = self.current_step >= self.n_steps
                info['capital'] = self.current_capital
                info['position'] = self.current_position
                return next_state, scaled_reward, done, info
            # else: default behavior (treat as normal bar) - implicitly handled by falling through
        # --- End Handle Imputed Bar --- #

        # --- Normal Step Logic (if not imputed or handling mode allows fallthrough like 'hold') --- #
        # 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

        # --- Apply Action Penalty (Revision 5) --- #
        # Penalty is applied to the raw reward *before* scaling
        # Uses the *trade size* for the penalty, with internal lambda
        if self._internal_action_penalty_lambda > 0:
             action_penalty = self._internal_action_penalty_lambda * abs(trade_size)
             reward -= action_penalty
             env_logger.debug(f"Step {self.current_step}: Action penalty applied: {action_penalty:.5f} (lambda={self._internal_action_penalty_lambda:.4f}, trade_size={trade_size:.3f})")
        # --- End Action Penalty --- #

        # --- Apply Reward Scaling (Task 5.1) --- #
        scaled_reward = reward * self.reward_scale
        # --- End Reward Scaling --- #

        # 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()
        # Update info dict (capital/position might have changed in normal step)
        info['capital'] = self.current_capital
        info['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, scaled_reward, done, info

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