import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model
import tensorflow_probability as tfp
from tensorflow.keras.optimizers.schedules import ExponentialDecay
import logging
import os
import json

sac_logger = logging.getLogger(__name__)
sac_logger.setLevel(logging.INFO)

class OrnsteinUhlenbeckActionNoise:
    def __init__(self, mean, std_deviation, theta=0.15, dt=0.01, x_initial=None):
        self.theta = theta
        self.mean = mean
        self.std_dev = std_deviation
        self.dt = dt
        self.x_initial = x_initial
        self.x_prev = self.x_initial if self.x_initial is not None else np.zeros_like(self.mean)

    def __call__(self):
        x = (
            self.x_prev
            + self.theta * (self.mean - self.x_prev) * self.dt
            + self.std_dev * np.sqrt(self.dt) * np.random.normal(size=self.mean.shape)
        )
        self.x_prev = x
        return x

    def reset(self):
        self.x_prev = self.x_initial if self.x_initial is not None else np.zeros_like(self.mean)

class SACTradingAgent:
    """V7.3 Enhanced: SAC agent with updated params and architecture fixes."""
    
    def __init__(self, 
                 state_dim=5,         # [mu, sigma, edge, |mu|/sigma, position]
                 action_dim=1, 
                 gamma=0.99,          
                 tau=0.005,           
                 initial_lr=3e-4,     
                 decay_steps=100000,  
                 end_lr=5e-6, # Note: End LR not directly used by ExponentialDecay
                 lr_decay_rate=0.96,  
                 ou_noise_stddev=0.2, 
                 ou_noise_theta=0.15, 
                 ou_noise_dt=0.01,    
                 alpha=0.2,           
                 alpha_auto_tune=True,
                 target_entropy=-1.0, 
                 edge_threshold_config: float | None = None,
                 reward_scale_config: float | None = None, 
                 action_penalty_lambda_config: float | None = None):
        """
        Initialize the SAC agent with enhancements.
        Args:
            state_dim (int): The dimension of the state space.
            action_dim (int): The dimension of the action space.
            gamma (float): The discount factor.
            tau (float): The target network update rate.
            initial_lr (float): The initial learning rate.
            decay_steps (int): The number of steps over which to decay the learning rate.
            end_lr (float): The final learning rate. (Not used by ExponentialDecay)
            lr_decay_rate (float): The rate at which to decay the learning rate.
            ou_noise_stddev (float): The standard deviation of the Ornstein-Uhlenbeck noise.
            ou_noise_theta (float): The theta parameter of the Ornstein-Uhlenbeck noise.
            ou_noise_dt (float): The dt parameter of the Ornstein-Uhlenbeck noise.
            alpha (float): The initial alpha value.
            alpha_auto_tune (bool): Whether to automatically tune the alpha value.
            target_entropy (float): The target entropy for the alpha auto-tuning.
            edge_threshold_config (float | None): The edge threshold value from the config.
            reward_scale_config (float | None): Environment reward scale used during training.
            action_penalty_lambda_config (float | None): Action penalty lambda used during training.
        """
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.tau = tau
        self.alpha_auto_tune = alpha_auto_tune
        self.edge_threshold_config = edge_threshold_config
        self.reward_scale_config = reward_scale_config
        self.action_penalty_lambda_config = action_penalty_lambda_config

        # --- Target Entropy Calculation --- #
        effective_target_entropy = target_entropy
        default_target_entropy_value = -1.0 * float(self.action_dim)
        if self.alpha_auto_tune:
            if abs(target_entropy - default_target_entropy_value) < 1e-6:
                 effective_target_entropy = -0.5 * np.log(4.0) 
                 sac_logger.info(f"alpha_auto_tune=True and default target_entropy detected. Setting target_entropy to -0.5*log(4) = {effective_target_entropy:.4f}")
            else:
                 effective_target_entropy = target_entropy
                 sac_logger.info(f"alpha_auto_tune=True, using explicitly provided target_entropy: {effective_target_entropy:.4f}")
            self.log_alpha = tf.Variable(tf.math.log(alpha), trainable=True, name='log_alpha')
            self.alpha = tfp.util.DeferredTensor(self.log_alpha, tf.exp)
            self.alpha_optimizer = tf.keras.optimizers.Adam(learning_rate=float(initial_lr))
        else:
            effective_target_entropy = target_entropy
            self.alpha = tf.constant(alpha, dtype=tf.float32)
            sac_logger.info(f"alpha_auto_tune=False. Using fixed alpha={self.alpha:.4f}")
        self.target_entropy = tf.constant(effective_target_entropy, dtype=tf.float32)
        # --- End Target Entropy --- #

        self.ou_noise = OrnsteinUhlenbeckActionNoise(
            mean=np.zeros(action_dim), 
            std_deviation=float(ou_noise_stddev) * np.ones(action_dim),
            theta=ou_noise_theta, dt=ou_noise_dt)

        self.lr_schedule = ExponentialDecay(
            initial_learning_rate=float(initial_lr),
            decay_steps=int(decay_steps),
            decay_rate=float(lr_decay_rate),
            staircase=False)
        sac_logger.info(f"Using ExponentialDecay LR: init={initial_lr}, steps={decay_steps}, rate={lr_decay_rate}")
        self.actor_optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr_schedule)
        self.critic1_optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr_schedule)
        self.critic2_optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr_schedule)

        # Initialize networks
        self.actor = self._build_actor()
        self.critic1 = self._build_critic()
        self.critic2 = self._build_critic()
        self.target_critic1 = self._build_critic()
        self.target_critic2 = self._build_critic()

        self.update_target_networks(tau=1.0)

        sac_logger.info("Enhanced SAC Agent Initialized (V7.3 - Consolidated).")
        sac_logger.info(f" State Dim: {state_dim}, Action Dim: {action_dim}")
        sac_logger.info(f" Hyperparams: gamma={gamma}, tau={tau}, alpha={'auto' if alpha_auto_tune else alpha}, target_entropy={target_entropy}")
        sac_logger.info(f" LR Schedule: Exponential {initial_lr} -> ? (decay_rate={lr_decay_rate})")
        sac_logger.info(f" OU Noise: std={ou_noise_stddev}, theta={ou_noise_theta}, dt={ou_noise_dt}")

    def _build_actor(self):
        inputs = layers.Input(shape=(self.state_dim,))
        x1 = layers.Dense(128, activation='relu')(inputs); x1_norm = layers.BatchNormalization()(x1)
        x2 = layers.Dense(64, activation='relu')(x1_norm); x2_norm = layers.BatchNormalization()(x2)
        x1_proj = layers.Dense(64)(x1_norm); x_res = layers.add([x1_proj, x2_norm])
        mean = layers.Dense(self.action_dim, activation='tanh')(x_res)
        log_std = layers.Dense(self.action_dim, activation='linear')(x_res)
        # V7.3 Fix: Use Lambda layer for clipping
        log_std_clipped = layers.Lambda(lambda x: tf.clip_by_value(x, -5, 2), name='log_std_clip')(log_std)
        model = Model(inputs=inputs, outputs=[mean, log_std_clipped]); return model
    
    def _build_critic(self):
        """
        Build distributional critic network.
        Outputs: [Q_mean, Q_log_std]
        """
        state_inputs = layers.Input(shape=(self.state_dim,))
        action_inputs = layers.Input(shape=(self.action_dim,))
        concat_inputs = layers.Concatenate()([state_inputs, action_inputs])
        
        x1 = layers.Dense(128, activation='relu')(concat_inputs); x1_norm = layers.BatchNormalization()(x1)
        x2 = layers.Dense(64, activation='relu')(x1_norm); x2_norm = layers.BatchNormalization()(x2)
        x1_proj = layers.Dense(64)(x1_norm); x_res = layers.add([x1_proj, x2_norm])

        q_mean = layers.Dense(1, name='q_mean')(x_res)
        q_log_std = layers.Dense(1, name='q_log_std')(x_res)
        # V7.3 Fix: Use Lambda layer for clipping
        q_log_std_clipped = layers.Lambda(lambda x: tf.clip_by_value(x, -5, 2), name='q_log_std_clip')(q_log_std)
        
        # V7.3 Revert: Remove auxiliary outputs
        model = Model(inputs=[state_inputs, action_inputs], 
                      outputs=[q_mean, q_log_std_clipped]) 
        return model
    
    @tf.function
    def _get_action_distribution(self, state):
        mean, log_std = self.actor(state); std = tf.exp(log_std)
        distribution = tfp.distributions.Normal(loc=mean, scale=std)
        squashed_distribution = tfp.distributions.TransformedDistribution(distribution=distribution, bijector=tfp.bijectors.Tanh())
        return squashed_distribution

    def get_action(self, state, deterministic=False):
        state_tensor = tf.convert_to_tensor([state], dtype=tf.float32)
        mean, log_std = self.actor(state_tensor)
        if deterministic:
            action = tf.tanh(mean)
        else:
            std = tf.exp(log_std); distribution = tfp.distributions.Normal(loc=mean, scale=std)
            raw_action = distribution.sample(); action = tf.tanh(raw_action)
            noise = self.ou_noise(); action = action + noise
            action = tf.clip_by_value(action, -1.0, 1.0)
        return action.numpy()[0]

    def update_target_networks(self, tau=None):
        if tau is None: tau = self.tau
        for target_weight, weight in zip(self.target_critic1.weights, self.critic1.weights): target_weight.assign(tau * weight + (1.0 - tau) * target_weight)
        for target_weight, weight in zip(self.target_critic2.weights, self.critic2.weights): target_weight.assign(tau * weight + (1.0 - tau) * target_weight)
    
    @tf.function
    def _update_critics(self, states, actions, rewards, next_states, dones):
        next_means, next_log_stds = self.actor(next_states); next_stds = tf.exp(next_log_stds)
        next_distributions = tfp.distributions.Normal(loc=next_means, scale=next_stds)
        next_actions_raw = next_distributions.sample(); next_actions_tanh = tf.tanh(next_actions_raw)
        next_log_probs = next_distributions.log_prob(next_actions_raw)
        log_correction = tf.math.log(1.0 - tf.square(next_actions_tanh) + 1e-6)
        next_log_probs = tf.reduce_sum(next_log_probs - log_correction, axis=1, keepdims=True)
        
        target_q1_mean, target_q1_log_std = self.target_critic1([next_states, next_actions_tanh])
        target_q2_mean, target_q2_log_std = self.target_critic2([next_states, next_actions_tanh])
        target_q_min_mean = tf.minimum(target_q1_mean, target_q2_mean)
        target_q = target_q_min_mean - self.alpha * next_log_probs 
        target_q_values = rewards + (1.0 - dones) * self.gamma * target_q
        # Ensure target_q_values is correctly shaped [batch_size, 1]
        target_q_values_stopped = tf.stop_gradient(target_q_values)

        critic1_vars = self.critic1.trainable_variables
        critic2_vars = self.critic2.trainable_variables

        with tf.GradientTape(persistent=True) as tape:
            current_q1_mean, current_q1_log_std = self.critic1([states, actions])
            current_q2_mean, current_q2_log_std = self.critic2([states, actions])

            pred_dist1 = tfp.distributions.Normal(loc=current_q1_mean, scale=tf.exp(current_q1_log_std))
            pred_dist2 = tfp.distributions.Normal(loc=current_q2_mean, scale=tf.exp(current_q2_log_std))
            nll_loss1 = -pred_dist1.log_prob(target_q_values_stopped) # Use stopped gradient here
            nll_loss2 = -pred_dist2.log_prob(target_q_values_stopped) # Use stopped gradient here
            critic1_loss = tf.reduce_mean(nll_loss1)
            critic2_loss = tf.reduce_mean(nll_loss2)
            
            # --- Calculate TD Errors for PER --- #
            # Use the mean of the prediction as the current Q estimate for error calculation
            current_q_min = tf.minimum(current_q1_mean, current_q2_mean)
            td_errors = tf.abs(target_q_values - current_q_min) # Target Q already has stopped gradient from above
            # Ensure td_errors have shape [batch_size, 1] or [batch_size]
            td_errors = tf.squeeze(td_errors) # Squeeze to [batch_size] if needed by buffer
            # --- End TD Error Calculation --- #
        
        critic1_gradients = tape.gradient(critic1_loss, critic1_vars)
        critic2_gradients = tape.gradient(critic2_loss, critic2_vars)
        del tape 
        
        self.critic1_optimizer.apply_gradients(zip(critic1_gradients, critic1_vars))
        self.critic2_optimizer.apply_gradients(zip(critic2_gradients, critic2_vars))
        
        # Return losses and TD errors
        return critic1_loss, critic2_loss, td_errors 
    
    @tf.function
    def _update_actor(self, states):
        actor_vars = self.actor.trainable_variables
        with tf.GradientTape() as tape:
            means, log_stds = self.actor(states); stds = tf.exp(log_stds)
            distributions = tfp.distributions.Normal(loc=means, scale=stds)
            actions_raw = distributions.sample(); actions_tanh = tf.tanh(actions_raw)
            log_probs = distributions.log_prob(actions_raw)
            log_correction = tf.math.log(1.0 - tf.square(actions_tanh) + 1e-6)
            log_probs = tf.reduce_sum(log_probs - log_correction, axis=1, keepdims=True)
            
            q1_mean, _ = self.critic1([states, actions_tanh])
            q2_mean, _ = self.critic2([states, actions_tanh])
            q_min_mean = tf.minimum(q1_mean, q2_mean)
            
            actor_loss = tf.reduce_mean(self.alpha * log_probs - q_min_mean)
        
        actor_gradients = tape.gradient(actor_loss, actor_vars)
        self.actor_optimizer.apply_gradients(zip(actor_gradients, actor_vars))
        return actor_loss, log_probs

    @tf.function
    def _update_alpha(self, log_probs):
        with tf.GradientTape() as tape:
            alpha_loss = -tf.reduce_mean(self.log_alpha * tf.stop_gradient(log_probs + self.target_entropy))
        
        alpha_gradients = tape.gradient(alpha_loss, [self.log_alpha])
        self.alpha_optimizer.apply_gradients(zip(alpha_gradients, [self.log_alpha]))
        return alpha_loss

    def train(self, states, actions, rewards, next_states, dones, importance_weights=None):
        """
        Perform a single training update step using a batch of experience.
        Now accepts and returns values needed for PER.

        Args:
            states (tf.Tensor): Batch of states.
            actions (tf.Tensor): Batch of actions.
            rewards (tf.Tensor): Batch of rewards.
            next_states (tf.Tensor): Batch of next states.
            dones (tf.Tensor): Batch of done flags.
            importance_weights (tf.Tensor, optional): Importance sampling weights for PER. Defaults to None.
            
        Returns:
            dict: Dictionary containing loss metrics and TD errors.
        """
        # Apply importance weights if provided (for PER)
        sample_weights = importance_weights if importance_weights is not None else tf.ones_like(rewards)
        
        # Critic updates now return TD errors
        critic1_loss, critic2_loss, td_errors = self._update_critics(
            states, actions, rewards, next_states, dones
        )
        # Apply importance weights to critic losses (if needed by loss definition, often handled internally)
        # Example: critic1_loss = tf.reduce_mean(critic1_loss * sample_weights)
        # Note: If using NLL loss, applying weights might need careful consideration.
        # For simplicity here, we assume the mean reduction inside _update_critics is sufficient,
        # but a true PER implementation might weight the NLL terms *before* reduction.
        
        actor_loss, log_probs = self._update_actor(states)
        # Actor loss already uses Q values which implicitly account for importance via critic updates? 
        # Typically, actor loss is not directly weighted by IS weights.
        
        alpha_loss = None
        if self.alpha_auto_tune:
            alpha_loss = self._update_alpha(log_probs)
            # Alpha loss typically not weighted by IS weights.
        
        self.update_target_networks()
        
        metrics = {
            "critic1_loss": float(critic1_loss),
            "critic2_loss": float(critic2_loss),
            "actor_loss": float(actor_loss),
            "learning_rate": float(self.lr_schedule(self.actor_optimizer.iterations)),
            "alpha": float(tf.exp(self.log_alpha)) if self.alpha_auto_tune else float(self.alpha),
            "td_errors": td_errors.numpy() # Return TD errors for priority updates
        }
        if alpha_loss is not None:
             metrics["alpha_loss"] = float(alpha_loss)
             
        return metrics
    
    def save(self, path):
        """Saves agent weights and potentially metadata."""
        try:
            os.makedirs(path, exist_ok=True)
            self.actor.save_weights(os.path.join(path, "actor.weights.h5"))
            self.critic1.save_weights(os.path.join(path, "critic1.weights.h5"))
            self.critic2.save_weights(os.path.join(path, "critic2.weights.h5"))

            metadata = {}
            if self.alpha_auto_tune and hasattr(self, 'log_alpha'):
                # Save log_alpha directly
                np.save(os.path.join(path, "log_alpha.npy"), self.log_alpha.numpy())
                metadata['log_alpha_saved'] = True
            else:
                metadata['log_alpha_saved'] = False
                metadata['fixed_alpha'] = float(self.alpha) # Save fixed alpha value

            # Add other relevant metadata if needed (like state_dim, action_dim used during training)
            metadata['state_dim'] = self.state_dim
            metadata['action_dim'] = self.action_dim
            # Add edge threshold to metadata (Step 4-A)
            if self.edge_threshold_config is not None:
                 metadata['edge_threshold'] = self.edge_threshold_config # Use the key specified in revisions.txt
            else:
                 metadata['edge_threshold'] = None # Or a default value like 0.55?
            # --- Add Env Params to Metadata (Task 5.6) --- #
            if self.reward_scale_config is not None:
                 metadata['reward_scale'] = self.reward_scale_config
            else:
                 metadata['reward_scale'] = None # Indicate if not set
            if self.action_penalty_lambda_config is not None:
                 metadata['lambda'] = self.action_penalty_lambda_config # Use lambda key from revisions.txt
            else:
                 metadata['lambda'] = None
            # --- End Add Env Params --- #

            meta_path = os.path.join(path, 'agent_metadata.json')
            with open(meta_path, 'w') as f:
                json.dump(metadata, f, indent=4)

            sac_logger.info(f"SAC Agent weights and metadata saved to {path}/")
        except Exception as e:
            sac_logger.error(f"Error saving SAC weights/metadata: {e}", exc_info=True)

    def load(self, path):
        """Loads agent weights and potentially metadata."""
        try:
            # Load weights (existing logic seems ok, ensures models are built)
            if not self.actor.built: self.actor.build((None, self.state_dim))
            if not self.critic1.built: self.critic1.build([(None, self.state_dim), (None, self.action_dim)])
            if not self.critic2.built: self.critic2.build([(None, self.state_dim), (None, self.action_dim)])
            if not self.target_critic1.built: self.target_critic1.build([(None, self.state_dim), (None, self.action_dim)])
            if not self.target_critic2.built: self.target_critic2.build([(None, self.state_dim), (None, self.action_dim)])
            self.actor.load_weights(os.path.join(path, "actor.weights.h5")); self.critic1.load_weights(os.path.join(path, "critic1.weights.h5"))
            self.critic2.load_weights(os.path.join(path, "critic2.weights.h5")); self.target_critic1.load_weights(os.path.join(path, "critic1.weights.h5"))
            self.target_critic2.load_weights(os.path.join(path, "critic2.weights.h5"))

            # Load metadata
            meta_path = os.path.join(path, 'agent_metadata.json')
            metadata = {}
            if os.path.exists(meta_path):
                with open(meta_path, 'r') as f:
                    metadata = json.load(f)
                sac_logger.info(f"Loaded agent metadata: {metadata}")

                # Load log_alpha if saved and auto-tuning
                log_alpha_path = os.path.join(path, "log_alpha.npy")
                if self.alpha_auto_tune and metadata.get('log_alpha_saved', False) and os.path.exists(log_alpha_path):
                    self.log_alpha.assign(np.load(log_alpha_path))
                    sac_logger.info(f"Restored log_alpha value from saved state.")
                elif not self.alpha_auto_tune and 'fixed_alpha' in metadata:
                    # Restore fixed alpha if not auto-tuning
                    self.alpha = tf.constant(metadata['fixed_alpha'], dtype=tf.float32)
                    sac_logger.info(f"Restored fixed alpha value: {self.alpha:.4f}")

            else:
                sac_logger.warning(f"Agent metadata file not found at {meta_path}. Cannot verify parameters or load log_alpha.")

            sac_logger.info(f"SAC Agent weights loaded from {path}/")
            return metadata # Return metadata for potential checks

        except Exception as e:
            sac_logger.error(f"Error loading SAC weights/metadata from {path}: {e}. Ensure files exist/shapes match.", exc_info=True)
            return {} # Return empty dict on failure

    def clear_buffer(self):
        """Clears the agent's replay buffer."""
        sac_logger.error("Agent or buffer does not have a clear_buffer method.") 