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

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 ReplayBuffer:
    """Standard Experience replay buffer for SAC agent"""
    
    def __init__(self, capacity=100000, state_dim=2, action_dim=1):
        self.capacity = capacity
        self.counter = 0
        
        # Initialize buffer arrays
        self.states = np.zeros((capacity, state_dim), dtype=np.float32)
        self.actions = np.zeros((capacity, action_dim), dtype=np.float32)
        self.rewards = np.zeros((capacity, 1), dtype=np.float32)
        self.next_states = np.zeros((capacity, state_dim), dtype=np.float32)
        self.dones = np.zeros((capacity, 1), dtype=np.float32)
    
    def add(self, state, action, reward, next_state, done):
        """Add experience to buffer"""
        idx = self.counter % self.capacity
        
        # Ensure inputs are correctly shaped numpy arrays
        state = np.array(state, dtype=np.float32).flatten()
        action = np.array(action, dtype=np.float32).flatten()
        reward = np.array([reward], dtype=np.float32)
        next_state = np.array(next_state, dtype=np.float32).flatten()
        done = np.array([done], dtype=np.float32)
        
        if state.shape[0] != self.states.shape[1]:
            sac_logger.error(f"State shape mismatch: {state.shape} vs {self.states.shape[1]}")
            return
        if next_state.shape[0] != self.next_states.shape[1]:
            sac_logger.error(f"Next State shape mismatch: {next_state.shape} vs {self.next_states.shape[1]}")
            return
        if action.shape[0] != self.actions.shape[1]:
            sac_logger.error(f"Action shape mismatch: {action.shape} vs {self.actions.shape[1]}")
            return
        
        self.states[idx] = state
        self.actions[idx] = action
        self.rewards[idx] = reward
        self.next_states[idx] = next_state
        self.dones[idx] = done
        
        self.counter += 1
    
    def sample(self, batch_size):
        """Sample batch of experiences from buffer"""
        max_idx = min(self.counter, self.capacity)
        if max_idx < batch_size:
            print(f"Warning: Trying to sample {batch_size} elements, but buffer only has {max_idx}. Sampling with replacement.")
            indices = np.random.choice(max_idx, batch_size, replace=True)
        else:
             indices = np.random.choice(max_idx, batch_size, replace=False)
        
        states = tf.convert_to_tensor(self.states[indices], dtype=tf.float32)
        actions = tf.convert_to_tensor(self.actions[indices], dtype=tf.float32)
        rewards = tf.convert_to_tensor(self.rewards[indices], dtype=tf.float32)
        next_states = tf.convert_to_tensor(self.next_states[indices], dtype=tf.float32)
        dones = tf.convert_to_tensor(self.dones[indices], dtype=tf.float32)
        
        return states, actions, rewards, next_states, dones
    
    def __len__(self):
        """Get current size of buffer"""
        return min(self.counter, self.capacity)

class SACTradingAgent:
    """V7.3 Enhanced: SAC agent with updated params and architecture fixes."""
    
    def __init__(self, 
                 state_dim=2,         # Standard [pred_ret, uncert]
                 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,  
                 buffer_capacity=100000,
                 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, 
                 min_buffer_size=1000):
        """
        Initialize the SAC agent with enhancements.
        """
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.tau = tau
        self.min_buffer_size = min_buffer_size
        self.target_entropy = tf.constant(target_entropy, dtype=tf.float32)
        self.alpha_auto_tune = alpha_auto_tune

        if self.alpha_auto_tune:
            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=initial_lr) 
        else:
            self.alpha = tf.constant(alpha, dtype=tf.float32)

        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=initial_lr, decay_steps=decay_steps, 
            decay_rate=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() # Outputs [Q_mean, Q_log_std]
        self.critic2 = self._build_critic()
        self.target_critic1 = self._build_critic()
        self.target_critic2 = self._build_critic()

        self.update_target_networks(tau=1.0)
        self.buffer = ReplayBuffer(capacity=buffer_capacity, state_dim=state_dim, action_dim=action_dim)

        sac_logger.info("Enhanced SAC Agent Initialized (V7.3).")
        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" Buffer: {buffer_capacity}, Min Size: {min_buffer_size}, Batch Size: Default 256 (in train)")
        sac_logger.info(f" OU Noise: std={ou_noise_stddev}, theta={ou_noise_theta}, dt={ou_noise_dt}")
        sac_logger.info(f" PER Note: Standard buffer used (PER={False})")
        
    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)
    
    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

        # Explicitly get trainable variables before the tape
        critic1_vars = self.critic1.trainable_variables
        critic2_vars = self.critic2.trainable_variables

        with tf.GradientTape(persistent=True) as tape:
            # Ensure the tape watches the correct variables if needed, though default should be fine
            # tape.watch(critic1_vars)
            # tape.watch(critic2_vars)
            
            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(tf.stop_gradient(target_q_values))
            nll_loss2 = -pred_dist2.log_prob(tf.stop_gradient(target_q_values))
            critic1_loss = tf.reduce_mean(nll_loss1)
            critic2_loss = tf.reduce_mean(nll_loss2)
        
        # Calculate gradients w.r.t the specific variable lists
        critic1_gradients = tape.gradient(critic1_loss, critic1_vars)
        critic2_gradients = tape.gradient(critic2_loss, critic2_vars)
        del tape 
        
        # Apply gradients paired with the specific variable lists using separate optimizers
        self.critic1_optimizer.apply_gradients(zip(critic1_gradients, critic1_vars))
        self.critic2_optimizer.apply_gradients(zip(critic2_gradients, critic2_vars))
        
        return critic1_loss, critic2_loss 
    
    def _update_actor(self, states):
        # Explicitly get trainable variables before the tape
        actor_vars = self.actor.trainable_variables
        
        with tf.GradientTape() as tape:
            # tape.watch(actor_vars)
            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)
        
        # Calculate gradients w.r.t the specific variable list
        actor_gradients = tape.gradient(actor_loss, actor_vars)
        # Apply gradients paired with the specific variable list
        self.actor_optimizer.apply_gradients(zip(actor_gradients, actor_vars))
        return actor_loss, log_probs

    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, batch_size=256):
        """
        Train the enhanced SAC agent.
        Includes alpha auto-tuning.
        Reverted aux tasks and state dim.
        """
        if len(self.buffer) < self.min_buffer_size:
            return {} 
        
        states, actions, rewards, next_states, dones = self.buffer.sample(batch_size)
        
        critic1_loss, critic2_loss = self._update_critics(
            states, actions, rewards, next_states, dones
        )
        
        actor_loss, log_probs = self._update_actor(states)
        
        alpha_loss = None
        if self.alpha_auto_tune:
            alpha_loss = self._update_alpha(log_probs)
        
        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)
        }
        if alpha_loss is not None:
             metrics["alpha_loss"] = float(alpha_loss)
             
        return metrics
    
    def save(self, path):
        try:
            self.actor.save_weights(f"{path}/actor.weights.h5"); self.critic1.save_weights(f"{path}/critic1.weights.h5")
            self.critic2.save_weights(f"{path}/critic2.weights.h5")
            if self.alpha_auto_tune and hasattr(self, 'log_alpha'): np.save(f"{path}/log_alpha.npy", self.log_alpha.numpy())
            sac_logger.info(f"Enhanced SAC Agent weights saved to {path}/")
        except Exception as e: sac_logger.error(f"Error saving SAC weights: {e}")
    
    def load(self, path):
        try:
            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(f"{path}/actor.weights.h5"); self.critic1.load_weights(f"{path}/critic1.weights.h5")
            self.critic2.load_weights(f"{path}/critic2.weights.h5"); self.target_critic1.load_weights(f"{path}/critic1.weights.h5")
            self.target_critic2.load_weights(f"{path}/critic2.weights.h5")
            log_alpha_path = f"{path}/log_alpha.npy"
            if self.alpha_auto_tune and os.path.exists(log_alpha_path): self.log_alpha.assign(np.load(log_alpha_path)); sac_logger.info(f"Loaded log_alpha value")
            sac_logger.info(f"Enhanced SAC Agent weights loaded from {path}/")
        except Exception as e: sac_logger.error(f"Error loading SAC weights from {path}: {e}. Ensure files exist/shapes match.") 