import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
import torch.optim as optim

from .rssm import RSSMState, RSSM
from .networks import Encoder, Decoder, RewardModel, Actor, Critic

def compute_lambda_returns(rewards: torch.Tensor, values: torch.Tensor, gamma: float = 0.99, lambda_: float = 0.95):
    """
    Computes TD(lambda) returns over an imagined trajectory.
    rewards: [horizon, batch]
    values: [horizon, batch]
    """
    horizon = rewards.shape[0]
    returns = []
    last_return = values[-1]
    
    for t in reversed(range(horizon - 1)):
        # TD(lambda) recursive calculation
        ret = rewards[t] + gamma * ((1 - lambda_) * values[t+1] + lambda_ * last_return)
        returns.insert(0, ret)
        last_return = ret
        
    returns.append(values[-1])
    return torch.stack(returns)

class AgenticForecaster(nn.Module):
    def __init__(self, obs_dim: int, action_dim: int, device: torch.device):
        super().__init__()
        self.device = device
        self.action_dim = action_dim
        
        self.encoder = Encoder(obs_dim=obs_dim).to(device)
        self.rssm = RSSM(action_dim=action_dim).to(device)
        
        feature_dim = self.rssm.stoch_dim + self.rssm.deter_dim
        self.decoder = Decoder(feature_dim=feature_dim, obs_dim=obs_dim).to(device)
        self.reward_model = RewardModel(feature_dim=feature_dim).to(device)
        
        self.actor = Actor(feature_dim=feature_dim, action_dim=action_dim).to(device)
        self.critic = Critic(feature_dim=feature_dim).to(device)
        self.target_critic = Critic(feature_dim=feature_dim).to(device)
        self.target_critic.load_state_dict(self.critic.state_dict())
        
        # Optimizers
        self.model_optimizer = optim.Adam(
            list(self.encoder.parameters()) + 
            list(self.rssm.parameters()) + 
            list(self.decoder.parameters()) + 
            list(self.reward_model.parameters()), 
            lr=3e-4
        )
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=1e-4)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-4)

    def train_world_model(self, observations: torch.Tensor, actions: torch.Tensor, rewards: torch.Tensor, free_nats: float = 3.0, kl_scale: float = 1.0):
        """
        observations: [batch_size, seq_len, obs_dim]
        actions: [batch_size, seq_len, action_dim]
        rewards: [batch_size, seq_len]
        """
        batch_size, seq_len, _ = observations.shape
        state = self.rssm.initial_state(batch_size, self.device)
        
        kl_losses = []
        recon_losses = []
        reward_losses = []
        
        for t in range(seq_len):
            obs_embed = self.encoder(observations[:, t, :])
            # Actions align such that action[t-1] led to observation[t]. For t=0, assume zero action
            action_t = actions[:, t-1, :] if t > 0 else torch.zeros(batch_size, self.action_dim, device=self.device)
            
            state, post_stats, prior_stats = self.rssm.observe_step(state, action_t, obs_embed)
            
            features = state.get_features()
            recon_obs = self.decoder(features)
            pred_reward = self.reward_model(features)
            
            recon_losses.append(F.mse_loss(recon_obs, observations[:, t, :], reduction='none').sum(-1).mean())
            reward_losses.append(F.mse_loss(pred_reward.squeeze(-1), rewards[:, t], reduction='mean'))
            
            post_dist = Normal(post_stats[0], post_stats[1])
            prior_dist = Normal(prior_stats[0], prior_stats[1])
            
            # Analytical KL divergence summed over latent dimension, averaged over batch
            kl_loss = torch.mean(torch.sum(torch.distributions.kl.kl_divergence(post_dist, prior_dist), dim=-1))
            # Free bits (nats) implementation
            kl_loss = torch.max(kl_loss, torch.tensor(free_nats).to(self.device))
            kl_losses.append(kl_loss)
            
        model_loss = sum(recon_losses) + sum(reward_losses) + kl_scale * sum(kl_losses)
        
        self.model_optimizer.zero_grad()
        model_loss.backward()
        nn.utils.clip_grad_norm_(self.parameters(), 100.0)
        self.model_optimizer.step()
        
        return {"model_loss": model_loss.item(), "recon_loss": sum(recon_losses).item(), "kl_loss": sum(kl_losses).item()}

    def train_actor_critic(self, start_states: RSSMState, horizon: int = 15, gamma: float = 0.99, lambda_: float = 0.95):
        """
        Trains Actor and Critic purely in latent imagination.
        start_states is sampled from the posterior distribution during world model training.
        """
        # 1. Rollout imagined trajectory
        state = start_states
        features_list = []
        rewards_list = []
        
        for _ in range(horizon):
            features = state.get_features()
            action = self.actor(features.detach(), explore=True)
            
            # Imagine next state driven by the actor's action
            state = self.rssm.imagine_step(state, action)
            
            next_features = state.get_features()
            reward = self.reward_model(next_features)
            
            features_list.append(next_features)
            rewards_list.append(reward.squeeze(-1))
            
        features_stack = torch.stack(features_list)  # [horizon, batch, feature_dim]
        rewards_stack = torch.stack(rewards_list)    # [horizon, batch]
        
        # 2. Compute TD(lambda) values using the target critic
        with torch.no_grad():
            values_stack = self.target_critic(features_stack).squeeze(-1)
            
        returns = compute_lambda_returns(rewards_stack, values_stack, gamma, lambda_)
            
        # 3. Train Actor (Maximize expected lambda return)
        actor_loss = -returns.mean()
        
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        nn.utils.clip_grad_norm_(self.actor.parameters(), 100.0)
        self.actor_optimizer.step()
        
        # 4. Train Critic (Predict lambda returns)
        # Detach features so we don't backprop into the world model
        pred_values = self.critic(features_stack.detach()).squeeze(-1)
        critic_loss = F.mse_loss(pred_values, returns.detach())
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        nn.utils.clip_grad_norm_(self.critic.parameters(), 100.0)
        self.critic_optimizer.step()
        
        # 5. Soft update target critic
        tau = 0.05
        for param, target_param in zip(self.critic.parameters(), self.target_critic.parameters()):
            target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
            
        return {"actor_loss": actor_loss.item(), "critic_loss": critic_loss.item()}

    def act(self, obs: torch.Tensor, prev_state: RSSMState, prev_action: torch.Tensor, explore: bool = False):
        """Used during data collection or evaluation to select an action."""
        with torch.no_grad():
            obs_embed = self.encoder(obs.unsqueeze(0))
            new_state, _, _ = self.rssm.observe_step(prev_state, prev_action.unsqueeze(0), obs_embed)
            features = new_state.get_features()
            action = self.actor(features, explore=explore)
            return action.squeeze(0), new_state

class HomeostaticAgenticForecaster(AgenticForecaster):
    """
    Dreamer agent with homeostatic critic.
    Maintains a stable value baseline while learning deviations.
    """
    
    def __init__(self, obs_dim: int, action_dim: int, device: torch.device):
        super().__init__(obs_dim, action_dim, device)
        
        # Replace critic with homeostatic version
        feature_dim = self.rssm.stoch_dim + self.rssm.deter_dim
        from .networks import HomeostaticCritic
        self.critic = HomeostaticCritic(feature_dim=feature_dim).to(device)
        self.target_critic = HomeostaticCritic(feature_dim=feature_dim).to(device)
        self.target_critic.load_state_dict(self.critic.state_dict())
        
        # Reinitialize optimizers with new parameters
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-4)
        self.training_step = 0
        
    def train_actor_critic(self, start_states: RSSMState, horizon: int = 15, 
                           gamma: float = 0.99, lambda_: float = 0.95):
        """Same as base, but critic updates its setpoint periodically"""
        
        # After computing returns, update homeostasis
        if self.training_step % 100 == 0:  # Every 100 steps
            update_homeostasis = True
        else:
            update_homeostasis = False
            
        self.training_step += 1
            
        # 1. Rollout imagined trajectory
        state = start_states
        features_list = []
        rewards_list = []
        
        for _ in range(horizon):
            features = state.get_features()
            action = self.actor(features.detach(), explore=True)
            
            state = self.rssm.imagine_step(state, action)
            
            next_features = state.get_features()
            reward = self.reward_model(next_features)
            
            features_list.append(next_features)
            rewards_list.append(reward.squeeze(-1))
            
        features_stack = torch.stack(features_list)
        rewards_stack = torch.stack(rewards_list)
        
        # 2. Compute TD(lambda) values using the target critic
        with torch.no_grad():
            values_stack = self.target_critic(features_stack).squeeze(-1)
            
        returns = compute_lambda_returns(rewards_stack, values_stack, gamma, lambda_)
            
        # 3. Train Actor
        actor_loss = -returns.mean()
        
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        nn.utils.clip_grad_norm_(self.actor.parameters(), 100.0)
        self.actor_optimizer.step()
        
        # 4. Train Critic (Predict lambda returns AND update setpoint optionally)
        pred_values = self.critic(features_stack.detach(), update_homeostasis=update_homeostasis).squeeze(-1)
        critic_loss = F.mse_loss(pred_values, returns.detach())
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        nn.utils.clip_grad_norm_(self.critic.parameters(), 100.0)
        self.critic_optimizer.step()
        
        # 5. Soft update target critic
        tau = 0.05
        for param, target_param in zip(self.critic.parameters(), self.target_critic.parameters()):
            target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
            
        return {"actor_loss": actor_loss.item(), "critic_loss": critic_loss.item()}