import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense, Concatenate, BatchNormalization, Add
from tensorflow.keras.optimizers import Adam
import tensorflow.keras.backend as K
import tensorflow_probability as tfp # V7.13 Import TFP

tfd = tfp.distributions # V7.13 TFP distribution alias

LOG_STD_MIN = -20 # V7.13 Min log std dev for numerical stability
LOG_STD_MAX = 2   # V7.13 Max log std dev for numerical stability

class SimplifiedSACTradingAgent:
    """
    Simplified SAC Trading Agent optimized for GRU-predicted returns and uncertainty
    with guarantees for performance on M1 chips and smaller datasets.
    V7.13: Updated for 5D state and automatic alpha tuning.
    """
    def __init__(
        self,
        state_dim=5,               # V7.13 Updated state: [pred_ret, unc, z, mom, vol]
        action_dim=1,              # Position size between -1 and 1
        hidden_size=64,            # Reduced network size for faster training
        gamma=0.97,                # Discount factor for faster adaptation
        tau=0.02,                  # Target network update rate
        # alpha=0.1,               # V7.13 Removed: Use automatic alpha tuning
        actor_lr=3e-4,             # V7.13 Default updated
        critic_lr=5e-4,            # V7.13 Default updated
        alpha_lr=3e-4,             # V7.13 Learning rate for alpha tuning
        batch_size=64,             # Smaller batch size for faster updates
        buffer_max_size=20000,     # Smaller buffer for recency bias
        min_buffer_size=1000,      # Start learning after this many experiences
        update_interval=1,         # Update actor every step
        target_update_interval=2,  # Update target networks every 2 steps
        gradient_clip=1.0,         # Clip gradients for stability
        reward_scale=2.0,          # V7.13 Default updated
        use_batch_norm=True,       # Use batch normalization
        use_residual=True,         # Use residual connections
        target_entropy=None,       # V7.13 Target entropy for alpha tuning
        model_dir='models/simplified_sac',
    ):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.hidden_size = hidden_size
        self.gamma = gamma
        self.tau = tau
        # self.alpha = alpha # V7.13 Removed
        self.actor_lr = actor_lr
        self.critic_lr = critic_lr
        self.alpha_lr = alpha_lr # V7.13 Store alpha LR
        self.batch_size = batch_size
        self.buffer_max_size = buffer_max_size
        self.min_buffer_size = min_buffer_size
        self.update_interval = update_interval
        self.target_update_interval = target_update_interval
        self.gradient_clip = gradient_clip
        self.reward_scale = reward_scale
        self.use_batch_norm = use_batch_norm
        self.use_residual = use_residual
        self.model_dir = model_dir
        
        # V7.13 Alpha tuning setup
        if target_entropy is None:
            # Default target entropy heuristic: -dim(A)/2 suggested in instructions
            self.target_entropy = -np.prod(self.action_dim) / 2.0 
        else:
            self.target_entropy = target_entropy
        # Initialize log_alpha (trainable variable) - start near log(0.1) ~ -2.3
        self.log_alpha = tf.Variable(np.log(0.1), dtype=tf.float32, name='log_alpha')
        self.alpha = tfp.util.DeferredTensor(self.log_alpha, tf.exp) # Exponentiated alpha
        self.alpha_optimizer = Adam(learning_rate=self.alpha_lr, name='alpha_optimizer')
        print(f"Initialized SAC with automatic alpha tuning. Target Entropy: {self.target_entropy:.2f}")

        # Experience replay buffer (simple numpy arrays instead of deque for performance)
        self.buffer_counter = 0
        self.buffer_capacity = buffer_max_size
        self.state_buffer = np.zeros((self.buffer_capacity, self.state_dim))
        self.action_buffer = np.zeros((self.buffer_capacity, self.action_dim))
        self.reward_buffer = np.zeros((self.buffer_capacity, 1))
        self.next_state_buffer = np.zeros((self.buffer_capacity, self.state_dim))
        self.done_buffer = np.zeros((self.buffer_capacity, 1))
        
        # Step counter
        self.train_step_counter = 0
        
        # Create actor and critic networks
        self.actor = self._build_actor()
        self.critic_1 = self._build_critic()
        self.critic_2 = self._build_critic()
        self.target_critic_1 = self._build_critic()
        self.target_critic_2 = self._build_critic()
        
        # Initialize target networks with actor and critic's weights
        self.target_critic_1.set_weights(self.critic_1.get_weights())
        self.target_critic_2.set_weights(self.critic_2.get_weights())
        
        # Create optimizers with gradient clipping
        self.actor_optimizer = Adam(learning_rate=self.actor_lr, clipnorm=self.gradient_clip)
        self.critic_optimizer = Adam(learning_rate=self.critic_lr, clipnorm=self.gradient_clip)
        # V7.13 Alpha optimizer already created above
        
        # Loss tracking
        self.actor_loss_history = []
        self.critic_loss_history = []
        
        # Ensure model directory exists
        os.makedirs(self.model_dir, exist_ok=True)

        # V7.9: Explicitly build models to prevent graph mode errors with optimizers
        self._build_models_with_dummy_input()
    
    def _build_models_with_dummy_input(self):
        """Builds actor and critic models with dummy inputs to initialize weights and optimizers."""
        try:
            # Dummy state and action tensors
            dummy_state = tf.zeros((1, self.state_dim), dtype=tf.float32)
            dummy_action = tf.zeros((1, self.action_dim), dtype=tf.float32)

            # Build actor
            self.actor(dummy_state)
            # Build critics
            self.critic_1([dummy_state, dummy_action])
            self.critic_2([dummy_state, dummy_action])
            self.target_critic_1([dummy_state, dummy_action])
            self.target_critic_2([dummy_state, dummy_action])

            # Optionally, build optimizers (Adam usually builds on first apply_gradients)
            # V7.10: Explicitly build optimizers too
            if hasattr(self.actor_optimizer, 'build'): # Check if build method exists (newer TF)
                self.actor_optimizer.build(self.actor.trainable_variables)
            if hasattr(self.critic_optimizer, 'build'):
                self.critic_optimizer.build(self.critic_1.trainable_variables + self.critic_2.trainable_variables)
            # V7.13 Build alpha optimizer
            if hasattr(self.alpha_optimizer, 'build'):
                 self.alpha_optimizer.build([self.log_alpha])
            print("Simplified SAC Agent models and optimizers built explicitly.") # Log success
        except Exception as e:
            print(f"Warning: Failed to explicitly build SAC models/optimizers: {e}")
    
    def _build_actor(self):
        """Build a simplified actor network with optional batch norm and residual connections
           V7.13: Outputs mean and log_std for a Gaussian policy.
        """
        # Input layer
        state_input = Input(shape=(self.state_dim,))
        
        # First hidden layer
        x = Dense(self.hidden_size, activation='relu')(state_input)
        if self.use_batch_norm:
            x = BatchNormalization()(x)
        
        # Second hidden layer
        y = Dense(self.hidden_size, activation='relu')(x)
        if self.use_batch_norm:
            y = BatchNormalization()(y)
        
        # Optional residual connection
        if self.use_residual and self.hidden_size == self.hidden_size:
            z = Add()([x, y])
        else:
            z = y
        
        # V7.13 Output layer(s) for mean and log_std
        mu = Dense(self.action_dim, activation=None, name='mu')(z)
        log_std = Dense(self.action_dim, activation=None, name='log_std')(z)
        
        # V7.13 Clip log_std for stability
        log_std = tf.keras.ops.clip(log_std, LOG_STD_MIN, LOG_STD_MAX)
        
        # V7.13 Create model - outputs mean and log_std
        model = Model(inputs=state_input, outputs=[mu, log_std])
        return model
    
    def _build_critic(self):
        """Build a simplified critic network for Q-value estimation"""
        # Input layers
        state_input = Input(shape=(self.state_dim,))
        action_input = Input(shape=(self.action_dim,))
        
        # Concatenate state and action
        concat = Concatenate()([state_input, action_input])
        
        # First hidden layer
        x = Dense(self.hidden_size, activation='relu')(concat)
        if self.use_batch_norm:
            x = BatchNormalization()(x)
        
        # Second hidden layer
        y = Dense(self.hidden_size, activation='relu')(x)
        if self.use_batch_norm:
            y = BatchNormalization()(y)
        
        # Optional residual connection
        if self.use_residual and self.hidden_size == self.hidden_size:
            z = Add()([x, y])
        else:
            z = y
        
        # Output layer (linear activation for Q-value)
        outputs = Dense(
            1, 
            activation=None,
            kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.003, maxval=0.003)
        )(z)
        
        # Create model
        model = Model(inputs=[state_input, action_input], outputs=outputs)
        return model
    
    def get_action(self, state, deterministic=False):
        """Select trading action (-1 to 1) based on current state
           V7.13: Samples from Gaussian policy, calculates log_prob.
        """
        state_tensor = tf.convert_to_tensor([state], dtype=tf.float32)
        # V7.13 Get mean and log_std from actor
        mu, log_std = self.actor(state_tensor)
        std = tf.exp(log_std)
        
        # Create policy distribution
        policy_dist = tfd.Normal(mu, std)
        
        if deterministic:
            # Use mean for deterministic action
            action = mu
        else:
            # Sample using reparameterization trick
            action = policy_dist.sample()
        
        # Calculate log probability of the sampled/deterministic action
        log_prob = policy_dist.log_prob(action)
        log_prob = tf.reduce_sum(log_prob, axis=1, keepdims=True) # Sum across action dim if > 1
        
        # Apply tanh squashing
        # Action is sampled from Normal, then squashed by tanh
        squashed_action = tf.tanh(action)
        
        # Adjust log_prob for tanh squashing (important!)
        # log π(a|s) = log ρ(u|s) - Σ log(1 - tanh(u)^2)
        # where u is the pre-tanh action sampled from Normal
        # See Appendix C in SAC paper: https://arxiv.org/abs/1801.01290
        log_prob -= tf.reduce_sum(tf.math.log(1.0 - squashed_action**2 + 1e-6), axis=1, keepdims=True)

        # Return the SQUASHED action and its log_prob
        # Return numpy arrays for consistency with buffer etc.
        return squashed_action[0].numpy(), log_prob[0].numpy()
    
    def store_transition(self, state, action, reward, next_state, done):
        """Store experience in replay buffer with efficient circular indexing"""
        # Scale reward
        scaled_reward = reward * self.reward_scale
        
        # Get the index to store experience
        index = self.buffer_counter % self.buffer_capacity
        
        # Store experience
        self.state_buffer[index] = state
        self.action_buffer[index] = action
        self.reward_buffer[index] = scaled_reward
        self.next_state_buffer[index] = next_state
        self.done_buffer[index] = done
        
        # Increment counter
        self.buffer_counter += 1
    
    def sample_batch(self):
        """Sample a batch of experiences from replay buffer"""
        # Get the valid buffer size (min of counter and capacity)
        buffer_size = min(self.buffer_counter, self.buffer_capacity)
        
        # Sample random indices
        batch_indices = np.random.choice(buffer_size, self.batch_size)
        
        # Get batch
        state_batch = self.state_buffer[batch_indices]
        action_batch = self.action_buffer[batch_indices]
        reward_batch = self.reward_buffer[batch_indices]
        next_state_batch = self.next_state_buffer[batch_indices]
        done_batch = self.done_buffer[batch_indices]
        
        return state_batch, action_batch, reward_batch, next_state_batch, done_batch
    
    def update_target_networks(self):
        """Update target critic networks using Polyak averaging"""
        # Update target critic 1
        target_weights = self.target_critic_1.get_weights()
        critic_weights = self.critic_1.get_weights()
        new_weights = []
        for i in range(len(target_weights)):
            new_weights.append(self.tau * critic_weights[i] + (1 - self.tau) * target_weights[i])
        self.target_critic_1.set_weights(new_weights)
        
        # Update target critic 2
        target_weights = self.target_critic_2.get_weights()
        critic_weights = self.critic_2.get_weights()
        new_weights = []
        for i in range(len(target_weights)):
            new_weights.append(self.tau * critic_weights[i] + (1 - self.tau) * target_weights[i])
        self.target_critic_2.set_weights(new_weights)
    
    @tf.function
    def _update_critics(self, states, actions, rewards, next_states, dones):
        """Update critic networks with TensorFlow graph execution"""
        with tf.GradientTape(persistent=True) as tape:
            # V7.18 START: Get next action and log_prob for target Q calculation
            next_mu, next_log_std = self.actor(next_states)
            next_std = tf.exp(next_log_std)
            next_policy_dist = tfd.Normal(next_mu, next_std)
            
            # Sample action for policy evaluation
            next_action_presquash = next_policy_dist.sample()
            next_action = tf.tanh(next_action_presquash) # Apply squashing
            
            # Calculate log_prob, adjusting for tanh squashing
            next_log_prob = next_policy_dist.log_prob(next_action_presquash)
            next_log_prob = tf.reduce_sum(next_log_prob, axis=1, keepdims=True)
            next_log_prob -= tf.reduce_sum(tf.math.log(1.0 - next_action**2 + 1e-6), axis=1, keepdims=True)
            # V7.18 END: Get next action and log_prob
            
            # Get target Q values using the SQUASHED next_action
            target_q1 = self.target_critic_1([next_states, next_action])
            target_q2 = self.target_critic_2([next_states, next_action])
            
            # Use minimum Q value (Double Q learning)
            min_target_q = tf.minimum(target_q1, target_q2)
            
            # V7.18 Calculate target Q including entropy term
            target_q_entropy_adjusted = min_target_q - self.alpha * next_log_prob
            
            # Calculate target values: r + γ(1-done) * (min(Q′(s′,a′)) - α log π(a′|s′))
            target_values = rewards + self.gamma * (1 - dones) * target_q_entropy_adjusted
            
            # Stop gradient flow through target values
            target_values = tf.stop_gradient(target_values)
            
            # Get current Q estimates
            current_q1 = self.critic_1([states, actions])
            current_q2 = self.critic_2([states, actions])
            
            # Calculate critic losses (Huber loss for robustness)
            critic1_loss = tf.reduce_mean(tf.keras.losses.huber(target_values, current_q1, delta=1.0))
            critic2_loss = tf.reduce_mean(tf.keras.losses.huber(target_values, current_q2, delta=1.0))
            critic_loss = critic1_loss + critic2_loss
        
        # Get critic gradients
        critic1_gradients = tape.gradient(critic1_loss, self.critic_1.trainable_variables)
        critic2_gradients = tape.gradient(critic2_loss, self.critic_2.trainable_variables)
        
        # Apply critic gradients
        self.critic_optimizer.apply_gradients(zip(critic1_gradients, self.critic_1.trainable_variables))
        self.critic_optimizer.apply_gradients(zip(critic2_gradients, self.critic_2.trainable_variables))
        
        del tape
        return critic_loss
    
    @tf.function
    def _update_actor_and_alpha(self, states):
        """Update actor network and alpha temperature with TensorFlow graph execution."""
        with tf.GradientTape(persistent=True) as tape:
            # Get policy distribution parameters
            mu, log_std = self.actor(states)
            std = tf.exp(log_std)
            policy_dist = tfd.Normal(mu, std)
            
            # Sample action using reparameterization trick
            actions_presquash = policy_dist.sample()
            log_prob_presquash = policy_dist.log_prob(actions_presquash)
            log_prob_presquash = tf.reduce_sum(log_prob_presquash, axis=1, keepdims=True)
            
            # Apply squashing
            squashed_actions = tf.tanh(actions_presquash)
            
            # Adjust log_prob for squashing
            log_prob = log_prob_presquash - tf.reduce_sum(tf.math.log(1.0 - squashed_actions**2 + 1e-6), axis=1, keepdims=True)
            
            # Get Q values from critics using squashed actions
            q1_values = self.critic_1([states, squashed_actions])
            q2_values = self.critic_2([states, squashed_actions])
            min_q_values = tf.minimum(q1_values, q2_values)
            
            # Calculate actor loss: E[alpha * log_prob - Q]
            actor_loss = tf.reduce_mean(self.alpha * log_prob - min_q_values)
            
            # Calculate alpha loss: E[-alpha * (log_prob + target_entropy)]
            # Note: We use log_alpha and optimize that variable
            alpha_loss = -tf.reduce_mean(self.log_alpha * tf.stop_gradient(log_prob + self.target_entropy))
            
        # --- Actor Gradients ---
        actor_gradients = tape.gradient(actor_loss, self.actor.trainable_variables)
        self.actor_optimizer.apply_gradients(zip(actor_gradients, self.actor.trainable_variables))
        
        # --- Alpha Gradients ---
        alpha_gradients = tape.gradient(alpha_loss, [self.log_alpha])
        self.alpha_optimizer.apply_gradients(zip(alpha_gradients, [self.log_alpha]))
        
        del tape # Release tape resources
        
        return actor_loss, alpha_loss # Return losses for tracking
    
    def train(self, num_iterations=1):
        """Train the agent for a specified number of iterations"""
        # Don't train if buffer doesn't have enough experiences
        if self.buffer_counter < self.min_buffer_size:
            return None # V7.17 Return None if not training
        
        total_actor_loss = 0
        total_critic_loss = 0
        total_alpha_loss = 0 # V7.19 Track alpha loss
        
        iterations_run = 0
        for _ in range(num_iterations):
            # Sample a batch of experiences
            states, actions, rewards, next_states, dones = self.sample_batch()
            
            # Convert to tensors
            states = tf.convert_to_tensor(states, dtype=tf.float32)
            actions = tf.convert_to_tensor(actions, dtype=tf.float32)
            rewards = tf.convert_to_tensor(rewards, dtype=tf.float32)
            next_states = tf.convert_to_tensor(next_states, dtype=tf.float32)
            dones = tf.convert_to_tensor(dones, dtype=tf.float32)
            
            # Update critics
            critic_loss = self._update_critics(states, actions, rewards, next_states, dones)
            total_critic_loss += critic_loss
            
            # Update actor and alpha less frequently for stability?
            # Original implementation updates actor every step, let's stick to that for now
            # if self.train_step_counter % self.update_interval == 0:
            actor_loss, alpha_loss = self._update_actor_and_alpha(states)
            total_actor_loss += actor_loss
            total_alpha_loss += alpha_loss # V7.19 Accumulate alpha loss
                
            # Track losses (consider appending outside the loop if num_iterations > 1)
            self.actor_loss_history.append(actor_loss.numpy())
            self.critic_loss_history.append(critic_loss.numpy())
            
            # Update target networks less frequently
            if self.train_step_counter % self.target_update_interval == 0:
                self.update_target_networks()
            
            # Increment step counter
            self.train_step_counter += 1
            iterations_run += 1
        
        # Return average losses for the iterations run in this call
        avg_actor_loss = total_actor_loss / iterations_run if iterations_run > 0 else tf.constant(0.0)
        avg_critic_loss = total_critic_loss / iterations_run if iterations_run > 0 else tf.constant(0.0)
        avg_alpha_loss = total_alpha_loss / iterations_run if iterations_run > 0 else tf.constant(0.0)
        
        # V7.17 Return tuple (actor_loss, critic_loss) - Needs update for alpha
        # V7.19 Return losses including alpha loss (or maybe just actor/critic for main history)
        return avg_actor_loss, avg_critic_loss # Keep original return for compatibility with main loop plotting for now
    
    def save(self, checkpoint_dir=None):
        """Save the model weights"""
        if checkpoint_dir is None:
            checkpoint_dir = self.model_dir
        
        os.makedirs(checkpoint_dir, exist_ok=True)
        
        # Save actor weights
        self.actor.save_weights(os.path.join(checkpoint_dir, 'actor.weights.h5'))
        
        # Save critic weights
        self.critic_1.save_weights(os.path.join(checkpoint_dir, 'critic_1.weights.h5'))
        self.critic_2.save_weights(os.path.join(checkpoint_dir, 'critic_2.weights.h5'))
        
        # Save target critic weights
        self.target_critic_1.save_weights(os.path.join(checkpoint_dir, 'target_critic_1.weights.h5'))
        self.target_critic_2.save_weights(os.path.join(checkpoint_dir, 'target_critic_2.weights.h5'))
        
        # Save alpha
        np.save(os.path.join(checkpoint_dir, 'alpha.npy'), self.alpha)
        
        print(f"Model saved to {checkpoint_dir}")
    
    def load(self, checkpoint_dir=None):
        """Load the model weights"""
        if checkpoint_dir is None:
            checkpoint_dir = self.model_dir
        
        try:
            # Load actor weights
            self.actor.load_weights(os.path.join(checkpoint_dir, 'actor.weights.h5'))
            
            # Load critic weights
            self.critic_1.load_weights(os.path.join(checkpoint_dir, 'critic_1.weights.h5'))
            self.critic_2.load_weights(os.path.join(checkpoint_dir, 'critic_2.weights.h5'))
            
            # Load target critic weights
            self.target_critic_1.load_weights(os.path.join(checkpoint_dir, 'target_critic_1.weights.h5'))
            self.target_critic_2.load_weights(os.path.join(checkpoint_dir, 'target_critic_2.weights.h5'))
            
            # Load alpha
            if os.path.exists(os.path.join(checkpoint_dir, 'alpha.npy')):
                self.alpha = float(np.load(os.path.join(checkpoint_dir, 'alpha.npy')))
            
            print(f"Model loaded from {checkpoint_dir}")
            return True
        except Exception as e:
            print(f"Error loading model: {e}")
            return False 