import argparse
import math
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
import copy
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.amp import autocast, GradScaler
from datasets import load_dataset
from transformers import AutoTokenizer


# Set the device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def parse_args():
    parser = argparse.ArgumentParser(description='Train World Model with Transformer outputs.')
    parser.add_argument('--model_name', type=str, default='gpt2', help='Pretrained model name or path')
    parser.add_argument('--dataset_name', type=str, default='wikitext', help='Dataset name from HuggingFace Datasets')
    parser.add_argument('--dataset_config', type=str, default='wikitext-2-raw-v1', help='Dataset configuration name')
    parser.add_argument('--batch_size', type=int, default=2, help='Batch size')
    parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs')
    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length')
    parser.add_argument('--mcts_iterations', type=int, default=5, help='Number of MCTS Iterations')
    parser.add_argument('--mcts_exploration_constant', type=float, default=1.414, help='Learning rate')
    parser.add_argument('--accumulation_steps', type=int, default=4, help='Gradient accumulation steps')
    parser.add_argument('--learning_rate', type=float, default=1e-4, help='Learning rate')
    parser.add_argument('--weight_decay', type=float, default=1e-2, help='Weight decay')
    parser.add_argument('--alpha', type=float, default=0.1, help='Entropy regularization weight')
    parser.add_argument('--beta', type=float, default=0.1, help='Variance regularization weight')
    parser.add_argument('--max_grad_norm', type=float, default=1.0, help='Max gradient norm for clipping')
    parser.add_argument('--save_dir', type=str, default='./models', help='Directory to save the models')
    parser.add_argument('--temperature', type=float, default=1.0, help='Temperature parameter for entropy and variance')
    parser.add_argument('--transformer_model_path', type=str, required=True, help='Path to the saved Transformer model')
    args = parser.parse_args()
    return args


def load_data(args, tokenizer):
    # Load the dataset
    dataset = load_dataset(args.dataset_name, args.dataset_config)
    
    # Ensure the tokenizer has a padding token
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def tokenize_function(examples):
        return tokenizer(examples['text'], truncation=True, max_length=args.max_length)

    tokenized_datasets = dataset.map(
        tokenize_function,
        batched=True,
        num_proc=4,
        remove_columns=dataset['train'].column_names,
    )

    # Build inputs and labels for language modeling
    block_size = args.max_length

    def group_texts(examples):
        # Concatenate all texts
        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
        total_length = len(concatenated_examples['input_ids'])
        # We drop the small remainder
        total_length = (total_length // block_size) * block_size
        # Split by chunks of block_size
        result = {
            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result['labels'] = result['input_ids'].copy()
        return result

    lm_datasets = tokenized_datasets.map(
        group_texts,
        batched=True,
        num_proc=4,
    )

    # Create DataLoader
    train_dataset = lm_datasets['train']
    eval_dataset = lm_datasets['validation'] if 'validation' in lm_datasets else lm_datasets['test']

    data_collator = lambda data: {
        'input_ids': torch.tensor([f['input_ids'] for f in data], dtype=torch.long),
        'labels': torch.tensor([f['labels'] for f in data], dtype=torch.long)
    }

    train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size, collate_fn=data_collator)
    eval_loader = DataLoader(eval_dataset, shuffle=False, batch_size=args.batch_size, collate_fn=data_collator)

    return train_loader, eval_loader


def save_all_models(transformer_model, representation_network, dynamics_network, prediction_network, action_encoder, save_dir, epoch):
    """
    Save all models to the specified directory.

    Args:
        transformer_model (nn.Module): Transformer model.
        representation_network (nn.Module): Representation network.
        dynamics_network (nn.Module): Dynamics network.
        prediction_network (nn.Module): Prediction network.
        action_encoder (nn.Module): Action encoder.
        save_dir (str): Directory to save the models.
        epoch (int): Current epoch number.
    """
    os.makedirs(save_dir, exist_ok=True)

    torch.save(transformer_model.state_dict(), os.path.join(save_dir, f'transformer_model_epoch_{epoch}.pt'))
    torch.save(representation_network.state_dict(), os.path.join(save_dir, f'representation_network_epoch_{epoch}.pt'))
    torch.save(dynamics_network.state_dict(), os.path.join(save_dir, f'dynamics_network_epoch_{epoch}.pt'))
    torch.save(prediction_network.state_dict(), os.path.join(save_dir, f'prediction_network_epoch_{epoch}.pt'))
    torch.save(action_encoder.state_dict(), os.path.join(save_dir, f'action_encoder_epoch_{epoch}.pt'))

    print(f"All models saved for epoch {epoch}.")


class RotaryPositionalEncoding(nn.Module):
    def __init__(self, d_model):
        super(RotaryPositionalEncoding, self).__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x):
        seq_len, batch_size, _ = x.size()
        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
        sinusoid_inp = torch.einsum("i,j->ij", t, self.inv_freq)
        sin = sinusoid_inp.sin().unsqueeze(1)  # (seq_len, 1, d_model/2)
        cos = sinusoid_inp.cos().unsqueeze(1)  # (seq_len, 1, d_model/2)

        x1 = x[..., 0::2]
        x2 = x[..., 1::2]

        # Apply rotation
        x_rotated = torch.zeros_like(x)
        x_rotated[..., 0::2] = x1 * cos - x2 * sin
        x_rotated[..., 1::2] = x1 * sin + x2 * cos

        return x_rotated


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        self.linear_q = nn.Linear(d_model, d_model)
        self.linear_k = nn.Linear(d_model, d_model)
        self.linear_v = nn.Linear(d_model, d_model)
        self.linear_out = nn.Linear(d_model, d_model)
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        query = self.linear_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        key = self.linear_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        value = self.linear_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e4)
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, value)
        
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        return self.linear_out(output)


class MoE(nn.Module):
    def __init__(self, d_model, num_experts, d_ff, top_k=2, dropout=0.1):
        super(MoE, self).__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU() if i % 2 == 0 else nn.SiLU(),
                nn.Linear(d_ff, d_model)
            )
            for i in range(num_experts)
        ])
        self.gate = nn.Linear(d_model, num_experts)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        batch_size, seq_len, d_model = x.size()
        # Compute gating scores
        gate_scores = self.gate(x)  # (batch_size, seq_len, num_experts)
        top_k_scores, top_k_indices = torch.topk(gate_scores, self.top_k, dim=-1)  # (batch_size, seq_len, top_k)
        top_k_scores = F.softmax(top_k_scores, dim=-1)  # (batch_size, seq_len, top_k)

        # Initialize output
        output = torch.zeros_like(x)

        # Flatten batch and sequence dimensions
        x_flat = x.view(-1, d_model)  # (batch_size * seq_len, d_model)
        output_flat = output.view(-1, d_model)
        top_k_indices_flat = top_k_indices.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
        top_k_scores_flat = top_k_scores.view(-1, self.top_k)  # (batch_size * seq_len, top_k)

        for k in range(self.top_k):
            expert_idx_flat = top_k_indices_flat[:, k]  # (batch_size * seq_len)
            expert_scores_flat = top_k_scores_flat[:, k]  # (batch_size * seq_len)
            for e in range(self.num_experts):
                mask = (expert_idx_flat == e)  # Boolean mask
                if mask.any():
                    x_masked = x_flat[mask]  # Select tokens for expert e
                    expert_output = self.experts[e](x_masked)  # Apply expert e
                    output_flat[mask] += expert_scores_flat[mask].unsqueeze(-1) * expert_output

        output = output_flat.view(batch_size, seq_len, d_model)
        return self.dropout(output)


class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_experts, dropout=0.1, top_k=2):
        super(TransformerBlock, self).__init__()
        self.self_attention = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.cross_attention = MultiHeadAttention(d_model, num_heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.moe = MoE(d_model, num_experts, d_ff, top_k, dropout)
        self.norm3 = nn.LayerNorm(d_model)

    def forward(self, x, mask=None, enc_output=None, enc_mask=None):
        # Self-attention
        attn_output = self.self_attention(x, x, x, mask)
        x = self.norm1(x + attn_output)
        # Cross-attention (only in decoder)
        if enc_output is not None:
            cross_attn_output = self.cross_attention(x, enc_output, enc_output, enc_mask)
            x = self.norm2(x + cross_attn_output)
        # Feedforward/MoE
        moe_output = self.moe(x)
        return self.norm3(x + moe_output)


class Transformer(nn.Module):
    def __init__(self, input_dim, d_model, num_heads, num_layers, d_ff, num_experts, output_dim, dropout=0.1, top_k=2):
        super(Transformer, self).__init__()
        self.embedding = nn.Embedding(input_dim, d_model, padding_idx=input_dim - 1)
        self.rotary_positional_encoding = RotaryPositionalEncoding(d_model)
        self.encoder_layers = nn.ModuleList(
            [TransformerBlock(d_model, num_heads, d_ff, num_experts, dropout, top_k) for _ in range(num_layers)]
        )
        self.decoder_layers = nn.ModuleList(
            [TransformerBlock(d_model, num_heads, d_ff, num_experts, dropout, top_k) for _ in range(num_layers)]
        )
        self.output_layer = nn.Linear(d_model, output_dim)
        self.d_model = d_model

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        # Encoder
        src = self.embedding(src) * math.sqrt(self.d_model)
        src = src.transpose(0, 1)  # (batch_size, seq_len, d_model) -> (seq_len, batch_size, d_model)
        src = self.rotary_positional_encoding(src)
        src = src.transpose(0, 1)  # (seq_len, batch_size, d_model) -> (batch_size, seq_len, d_model)
        for layer in self.encoder_layers:
            src = layer(src, src_mask)

        # Decoder
        tgt = self.embedding(tgt) * math.sqrt(self.d_model)
        tgt = tgt.transpose(0, 1)
        tgt = self.rotary_positional_encoding(tgt)
        tgt = tgt.transpose(0, 1)
        for layer in self.decoder_layers:
            tgt = layer(tgt, tgt_mask, src, src_mask)
        output = self.output_layer(tgt)
        return output

    def generate(self, src, tokenizer, max_length=20, temperature=1.0):
        """
        Generate sequences using differentiable sampling (Gumbel-Softmax).

        Args:
            src (torch.Tensor): Source input tensor of shape (batch_size, seq_len)
            tokenizer (transformers.PreTrainedTokenizer): Tokenizer to access special tokens
            max_length (int): Maximum length of the generated sequence
            temperature (float): Temperature parameter for Gumbel-Softmax

        Returns:
            torch.Tensor: Generated sequences of shape (batch_size, max_length)
            torch.Tensor: Entropy values for each time step
            torch.Tensor: Variance values for each time step
        """
        batch_size = src.size(0)

        # Encode the source
        src_enc = self.embedding(src) * math.sqrt(self.d_model)
        src_enc = src_enc.transpose(0, 1)
        src_enc = self.rotary_positional_encoding(src_enc)
        src_enc = src_enc.transpose(0, 1)
        for layer in self.encoder_layers:
            src_enc = layer(src_enc)

        # Initialize decoder input with <sos> tokens
        tgt_seq = torch.full((batch_size, 1), tokenizer.bos_token_id, dtype=torch.long, device=src.device)
        entropies = []
        variances = []

        for _ in range(max_length):
            tgt_emb = self.embedding(tgt_seq) * math.sqrt(self.d_model)
            tgt_emb = tgt_emb.transpose(0, 1)
            tgt_emb = self.rotary_positional_encoding(tgt_emb)
            tgt_emb = tgt_emb.transpose(0, 1)
            tgt_dec = tgt_emb
            for layer in self.decoder_layers:
                tgt_dec = layer(tgt_dec, None, src_enc, None)
            output = self.output_layer(tgt_dec)  # (batch_size, seq_len, vocab_size)
            logits = output[:, -1, :]  # Get logits for the last time step

            # Compute token probabilities
            probs = F.softmax(logits / temperature, dim=-1)  # (batch_size, vocab_size)

            # Compute entropy
            entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)  # (batch_size)
            entropies.append(entropy)

            # Sample token using Gumbel-Softmax
            gumbel_noise = -torch.log(-torch.log(torch.rand_like(probs) + 1e-9) + 1e-9)
            y = (logits + gumbel_noise) / temperature
            y = F.softmax(y, dim=-1)  # (batch_size, vocab_size)

            # Compute variance
            variance = torch.var(y, dim=-1)  # (batch_size)
            variances.append(variance)

            # Get token indices (argmax for hard selection)
            next_tokens = torch.argmax(y, dim=-1, keepdim=True)  # (batch_size, 1)
            tgt_seq = torch.cat([tgt_seq, next_tokens], dim=1)

        # Stack entropies and variances
        entropies = torch.stack(entropies, dim=1)  # (batch_size, max_length)
        variances = torch.stack(variances, dim=1)  # (batch_size, max_length)

        return tgt_seq[:, 1:], entropies, variances  # Exclude the initial <sos> token


# Objective Functions

class InfoNCE_Loss(nn.Module):
    def __init__(self, temperature=0.07):
        super(InfoNCE_Loss, self).__init__()
        self.temperature = temperature
        self.cross_entropy = nn.CrossEntropyLoss()
    
    def forward(self, z_i, z_j):
        """
        Args:
            z_i (torch.Tensor): Flattened representations from view i, shape (2n, embed_dim)
            z_j (torch.Tensor): Flattened representations from view j, shape (2n, embed_dim)
        
        Returns:
            torch.Tensor: InfoNCE loss
        """
        n = z_i.size(0)
        z = torch.cat([z_i, z_j], dim=0)  # Shape: (2n, embed_dim)

        z = F.normalize(z, dim=1)
        similarity_matrix = torch.matmul(z, z.T)  # Shape: (2n, 2n)

        # Create a mask to exclude self-similarity
        mask = torch.eye(2 * n, device=z.device, dtype=torch.bool)
        similarity_matrix = similarity_matrix.masked_fill(mask, -1e4)  # Use a manageable negative value

        # Create labels for contrastive learning
        labels = torch.arange(n, device=z.device)
        labels = torch.cat([labels + n, labels], dim=0)  # Shape: (2n,)

        # Apply temperature scaling
        similarity_matrix /= self.temperature

        # Compute cross-entropy loss
        loss = self.cross_entropy(similarity_matrix, labels)
        return loss



class CovarianceRegularization(nn.Module):
    def __init__(self, lambda_reg=1e-3):
        super(CovarianceRegularization, self).__init__()
        self.lambda_reg = lambda_reg
    
    def forward(self, embeddings):
        """
        Args:
            embeddings (torch.Tensor): Embedding tensor, shape (batch_size, embed_dim)
        
        Returns:
            torch.Tensor: Covariance regularization loss
        """
        batch_size, embed_dim = embeddings.size()
        mean = embeddings.mean(dim=0)
        embeddings_centered = embeddings - mean
        cov = (embeddings_centered.T @ embeddings_centered) / (batch_size - 1)
        cov_loss = torch.sum(cov ** 2) - torch.sum(torch.diag(cov) ** 2)
        return self.lambda_reg * cov_loss


class DynamicsPerformanceLoss(nn.Module):
    def __init__(self, lambda_var=1e-3):
        super(DynamicsPerformanceLoss, self).__init__()
        self.lambda_var = lambda_var
    
    def forward(self, true_next_state, predicted_next_state):
        """
        Args:
            true_next_state (torch.Tensor): Ground truth next state, shape (batch_size, state_dim)
            predicted_next_state (torch.Tensor): Predicted next state, shape (batch_size, state_dim)
        
        Returns:
            torch.Tensor: Dynamics performance loss
        """
        mse_loss = F.mse_loss(predicted_next_state, true_next_state)
        variance_loss = torch.var(predicted_next_state, dim=0).mean()
        return mse_loss + self.lambda_var * variance_loss


class ThoughtConsistencyLoss(nn.Module):
    def __init__(self):
        super(ThoughtConsistencyLoss, self).__init__()
    
    def forward(self, true_next_state, perturbed_next_state):
        """
        Args:
            true_next_state (torch.Tensor): Ground truth next state, shape (batch_size, state_dim)
            perturbed_next_state (torch.Tensor): Perturbed next state, shape (batch_size, state_dim)
        
        Returns:
            torch.Tensor: Thought-consistency loss
        """
        return F.mse_loss(true_next_state, perturbed_next_state)


class PolicyValueJointLoss(nn.Module):
    def __init__(self, lambda_value=0.5):
        super(PolicyValueJointLoss, self).__init__()
        self.lambda_value = lambda_value
        self.cross_entropy = nn.CrossEntropyLoss()
        self.mse_loss = nn.MSELoss()
    
    def forward(self, policy_logits, true_policy, value_pred, true_value):
        """
        Args:
            policy_logits (torch.Tensor): Logits from the policy network, shape (batch_size * seq_len, num_actions)
            true_policy (torch.Tensor): Ground truth policy, shape (batch_size * seq_len, num_actions)
            value_pred (torch.Tensor): Predicted values, shape (batch_size * seq_len)
            true_value (torch.Tensor): Ground truth values, shape (batch_size * seq_len)
        
        Returns:
            torch.Tensor: Combined policy and value loss
        """
        policy_logits = policy_logits.view(-1, policy_logits.size(-1))
        true_policy = true_policy.view(-1, true_policy.size(-1))
        value_pred = value_pred.view(-1)
        true_value = true_value.view(-1)

        policy_loss = self.cross_entropy(policy_logits, true_policy.argmax(dim=1))
        value_loss = self.mse_loss(value_pred, true_value)
        return policy_loss + self.lambda_value * value_loss



class ActionDiversityReward(nn.Module):
    def __init__(self, lambda_div=1e-3):
        super(ActionDiversityReward, self).__init__()
        self.lambda_div = lambda_div
    
    def forward(self, action_embeddings):
        """
        Args:
            action_embeddings (torch.Tensor): Embeddings of actions, shape (batch_size, embed_dim)
        
        Returns:
            torch.Tensor: Action diversity loss
        """
        similarity_matrix = F.cosine_similarity(action_embeddings.unsqueeze(1), action_embeddings.unsqueeze(0), dim=2)
        # Zero out self-similarity
        similarity_matrix = similarity_matrix - torch.eye(similarity_matrix.size(0)).to(action_embeddings.device)
        diversity_loss = torch.sum(similarity_matrix ** 2)
        return self.lambda_div * diversity_loss


class ExpectedThoughtValueLoss(nn.Module):
    def __init__(self):
        super(ExpectedThoughtValueLoss, self).__init__()
    
    def forward(self, mcts_best_values):
        """
        Args:
            mcts_best_values (torch.Tensor): Best values from MCTS, shape (batch_size)
        
        Returns:
            torch.Tensor: ETV loss
        """
        return -mcts_best_values.mean()


class ExplorationRegularization(nn.Module):
    def __init__(self, lambda_expl=1e-3):
        super(ExplorationRegularization, self).__init__()
        self.lambda_expl = lambda_expl
    
    def forward(self, visit_counts):
        """
        Args:
            visit_counts (torch.Tensor): Visit counts for actions, shape (batch_size, num_actions)
        
        Returns:
            torch.Tensor: Exploration regularization loss
        """
        reward = torch.sum(1.0 / (visit_counts + 1), dim=-1)
        return self.lambda_expl * reward.mean()


class KL_DivergenceLoss(nn.Module):
    def __init__(self):
        super(KL_DivergenceLoss, self).__init__()
    
    def forward(self, old_policy, new_policy):
        """
        Args:
            old_policy (torch.Tensor): Old policy probabilities, shape (batch_size, num_actions)
            new_policy (torch.Tensor): New policy probabilities, shape (batch_size, num_actions)
        
        Returns:
            torch.Tensor: KL divergence loss
        """
        kl_div = F.kl_div(new_policy.log(), old_policy, reduction='batchmean')
        return kl_div

# MuZero

class ActionEncoder(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super(ActionEncoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
    
    def forward(self, action_sequences):
        """
        Args:
            action_sequences (torch.Tensor): Tensor of shape (batch_size, seq_len)
        
        Returns:
            torch.Tensor: Encoded actions of shape (batch_size, seq_len, embed_dim)
        """
        return self.embedding(action_sequences) #.half()  # Convert to half-precision

class RepresentationNetwork(nn.Module):
    def __init__(self, vocab_dim, d_model, state_dim):
        super(RepresentationNetwork, self).__init__()
        self.proj = nn.Linear(vocab_dim, d_model)  # Project from vocab_dim to d_model
        self.linear = nn.Linear(d_model, state_dim)  # Project from d_model to state_dim
        self.norm = nn.LayerNorm(state_dim)
    
    def forward(self, transformer_output):
        """
        Args:
            transformer_output (torch.Tensor): Shape (batch_size, seq_len, vocab_dim)
        
        Returns:
            torch.Tensor: Encoded state of shape (batch_size, seq_len, state_dim)
        """
        # First project down from vocab_dim to d_model
        projected_output = self.proj(transformer_output)
        # Then project down from d_model to state_dim
        state = self.linear(projected_output)
        state = self.norm(state)
        return state


class DynamicsNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim):
        super(DynamicsNetwork, self).__init__()
        self.rms_norm = nn.LayerNorm(state_dim)
        self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim)
        self.activation = nn.GELU()
        self.fc2 = nn.Linear(hidden_dim, state_dim)
    
    def forward(self, state, action):
        """
        Args:
            state (torch.Tensor): Current state, shape (batch_size, seq_len, state_dim)
            action (torch.Tensor): Action embedding, shape (batch_size, seq_len, action_dim)
        
        Returns:
            torch.Tensor: Predicted next state, shape (batch_size, seq_len, state_dim)
        """
        norm_state = self.rms_norm(state)
        combined = torch.cat([norm_state, action], dim=-1)
        hidden = self.activation(self.fc1(combined))
        next_state = self.fc2(hidden)
        return next_state

class PredictionNetwork(nn.Module):
    def __init__(self, state_dim, policy_dim, value_dim):
        super(PredictionNetwork, self).__init__()
        self.state_dim = state_dim
        self.rms_norm = nn.LayerNorm(state_dim)
        self.policy_head = nn.Linear(state_dim, policy_dim)
        self.value_head = nn.Linear(state_dim, value_dim)
    
    def forward(self, state):
        """
        Args:
            state (torch.Tensor): Predicted state, shape (batch_size, seq_len, state_dim)
        
        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Policy logits and value estimates
        """
        norm_state = self.rms_norm(state)
        policy_logits = self.policy_head(norm_state)
        value_estimates = self.value_head(norm_state)
        return policy_logits, value_estimates


class MCTSNode:
    def __init__(self, state, parent=None, action=None):
        """
        Initialize an MCTS node.

        Args:
            state (State): The current state representation.
            parent (MCTSNode, optional): The parent node. Defaults to None.
            action (int, optional): The action taken to reach this node. Defaults to None.
        """
        self.state = state  # Instance of State class
        self.parent = parent  # Parent MCTSNode
        self.action = action  # Action taken to reach this node
        self.children = {}  # Dict mapping actions to MCTSNode
        self.visit_count = 0
        self.value_sum = 0.0
        self.prior = 0.0  # Prior probability from policy network

    def expand(self, actions, priors):
        """
        Expand the node with possible actions and their priors.

        Args:
            actions (list): List of possible actions (action indices).
            priors (list): List of prior probabilities corresponding to actions.
        """
        for action, prior in zip(actions, priors):
            if action not in self.children:
                child_state = self.state.apply_action(action)  # Apply action to get new state
                child_node = MCTSNode(state=child_state, parent=self, action=action)
                child_node.prior = float(prior)  # Ensure that prior is a float value
                self.children[action] = child_node

    def is_leaf(self):
        """
        Check if the node is a leaf node (i.e., has no children).

        Returns:
            bool: True if leaf, False otherwise.
        """
        return len(self.children) == 0

    def ucb_score(self, total_visits, exploration_constant=math.sqrt(2)):
        """
        Calculate the UCB (Upper Confidence Bound) score for the node.

        Args:
            total_visits (int): Total number of visits to the parent node.
            exploration_constant (float, optional): Exploration parameter. Defaults to math.sqrt(2).

        Returns:
            float: The UCB score.
        """
        if self.visit_count == 0:
            return float('inf')
        average_value = self.value_sum / self.visit_count
        exploration_term = exploration_constant * self.prior * math.sqrt(total_visits) / (1 + self.visit_count)
        return average_value + exploration_term

class MCTS:
    def __init__(self, prediction_network, dynamics_network, action_encoder, num_iterations=10, exploration_constant=math.sqrt(2)):
        """
        Initialize the MCTS.

        Args:
            prediction_network (nn.Module): The Prediction Network.
            dynamics_network (nn.Module): The Dynamics Network.
            num_iterations (int): Number of MCTS iterations per search.
            exploration_constant (float): Exploration parameter for UCB.
        """
        self.action_encoder = action_encoder
        self.prediction_network = prediction_network
        self.dynamics_network = dynamics_network
        self.num_iterations = num_iterations
        self.exploration_constant = exploration_constant

    def search(self, root_state):
        """
        Perform MCTS starting from the root_state.

        Args:
            root_state: The initial state from which to start MCTS.

        Returns:
            The best action determined by MCTS.
        """
        self.root = MCTSNode(state=root_state)

        for _ in range(self.num_iterations):
            node = self.select(self.root)
            value = self.evaluate(node)
            self.backpropagate(node, value)

        return self.best_action()

    def select(self, node):
        """
        Traverse the tree to select a node for evaluation.

        Args:
            node: The starting node for selection.

        Returns:
            The node selected for evaluation.
        """
        while not node.is_leaf():
            best_action, best_node = max(node.children.items(),
                                         key=lambda item: item[1].ucb_score(node.visit_count, self.exploration_constant))
            node = best_node
        return node

    def evaluate(self, node):
        """
        Evaluate the node by expanding it and predicting its value.

        Args:
            node: The node to evaluate.

        Returns:
            The value estimate of the node.
        """
        # Use the prediction network to get policy logits and value estimate
        policy_logits, value_estimate = self.prediction_network(node.state.representation)

        # Convert logits to probabilities
        policy = F.softmax(policy_logits, dim=-1).detach().cpu().numpy()

        # Expand the node with possible actions and their priors
        actions = list(range(policy.shape[-1]))  # Assuming actions are indexed from 0 to num_actions-1
        priors = policy[0].flatten().tolist()  # Convert to a 1D list of floats

        node.expand(actions, priors)

        return value_estimate.mean().item()


    def backpropagate(self, node, value):
        """
        Backpropagate the value up the tree.

        Args:
            node: The node to start backpropagation from.
            value (float): The value to backpropagate.
        """
        while node is not None:
            node.visit_count += 1
            node.value_sum += value
            node = node.parent

    def best_action(self):
        """
        Choose the action with the highest visit count.

        Returns:
            The best action.
        """
        best_child = max(self.root.children.values(), key=lambda n: n.visit_count)
        return best_child.action

class State:
    def __init__(self, representation, dynamics_network, action_encoder):
        """
        Initialize the State.

        Args:
            representation (torch.Tensor): Encoded state representation, shape (batch_size, seq_len, state_dim)
            dynamics_network (nn.Module): The Dynamics Network to predict next states
            action_encoder (nn.Module): The Action Encoder to encode actions
        """
        self.representation = representation  # Shape: (batch_size, seq_len, state_dim)
        self.dynamics_network = dynamics_network  # Reference to Dynamics Network
        self.action_encoder = action_encoder

    def apply_action(self, action):
        """
        Apply an action to the current state to get a new state.

        Args:
            action (int): The action to apply (e.g., token index)

        Returns:
            State: The new state after applying the action
        """
        # Create action sequence filled with action index
        batch_size, seq_len, _ = self.representation.size()
        action_sequence = torch.full((batch_size, seq_len), action, dtype=torch.long, device=self.representation.device)
        # Encode action
        action_embedding = self.action_encoder(action_sequence)
        # Predict the next state using the Dynamics Network
        with torch.no_grad():
            next_state_representation = self.dynamics_network(self.representation, action_embedding)
        return State(next_state_representation, self.dynamics_network, self.action_encoder)




class PPOAgent:
    def __init__(self, policy_network, optimizer, clip_epsilon=0.2, entropy_coef=0.01, value_coef=0.5):
        self.policy_network = policy_network
        self.optimizer = optimizer
        self.clip_epsilon = clip_epsilon
        self.entropy_coef = entropy_coef
        self.value_coef = value_coef
    
    def compute_loss(self, states, old_log_probs, actions, returns, advantages):
        # Get policy logits and value estimates
        policy_logits, value_estimates = self.policy_network(states)
        batch_size, seq_len, num_actions = policy_logits.size()
        
        # Flatten tensors
        policy_logits = policy_logits.view(-1, num_actions)  # Shape: (batch_size * seq_len, num_actions)
        value_estimates = value_estimates.view(-1)           # Shape: (batch_size * seq_len)
        actions = actions.view(-1)                           # Shape: (batch_size * seq_len)
        old_log_probs = old_log_probs.view(-1)               # Shape: (batch_size * seq_len)
        returns = returns.view(-1)                           # Shape: (batch_size * seq_len)
        advantages = advantages.view(-1)                     # Shape: (batch_size * seq_len)
        
        # Compute new log probabilities
        new_log_probs_all = F.log_softmax(policy_logits, dim=-1)  # Shape: (batch_size * seq_len, num_actions)
        new_log_probs = new_log_probs_all.gather(1, actions.unsqueeze(-1)).squeeze(-1)  # Shape: (batch_size * seq_len)
        
        # Compute ratios
        ratios = torch.exp(new_log_probs - old_log_probs)
        
        # PPO surrogate loss
        surr1 = ratios * advantages
        surr2 = torch.clamp(ratios, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages
        policy_loss = -torch.min(surr1, surr2).mean()
        
        # Value loss
        value_loss = F.mse_loss(value_estimates, returns)
        
        # Entropy loss
        entropy = -(new_log_probs * torch.exp(new_log_probs)).mean()
        
        # Total loss
        total_loss = policy_loss + self.value_coef * value_loss - self.entropy_coef * entropy
        return total_loss




def compute_loss_world_model(predicted_next_state, true_next_state, policy_logits, true_policy, value_estimates, true_value, 
                             alpha, beta, temperature, lambda_reg, lambda_var, lambda_div, lambda_expl):
    """
    Compute the combined loss for the World Model.

    Args:
        predicted_next_state (torch.Tensor): Predicted next state, shape (batch_size, state_dim)
        true_next_state (torch.Tensor): Ground truth next state, shape (batch_size, state_dim)
        policy_logits (torch.Tensor): Policy logits, shape (batch_size, num_actions)
        true_policy (torch.Tensor): Ground truth policy, shape (batch_size, num_actions)
        value_estimates (torch.Tensor): Value estimates, shape (batch_size)
        true_value (torch.Tensor): Ground truth value, shape (batch_size)
        alpha (float): Entropy regularization weight
        beta (float): Variance regularization weight
        temperature (float): Temperature parameter
        lambda_reg (float): Covariance regularization weight
        lambda_var (float): Dynamics variance loss weight
        lambda_div (float): Action diversity reward weight
        lambda_expl (float): Exploration regularization weight

    Returns:
        torch.Tensor: Combined loss
    """
    # Cross-entropy loss
    ce_loss = F.cross_entropy(policy_logits, true_policy.argmax(dim=1))
    
    # Entropy loss
    probs = F.softmax(policy_logits / temperature, dim=-1)  # (batch_size, num_actions)
    entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)  # (batch_size)
    entropy_loss = -alpha * torch.mean(entropy)
    
    # Variance loss
    variance = torch.var(probs, dim=-1)  # (batch_size)
    variance_loss = -beta * torch.mean(variance)
    
    # Covariance Regularization
    cov_reg = CovarianceRegularization(lambda_reg)(predicted_next_state)
    
    # Dynamics Performance Loss
    dynamics_loss = DynamicsPerformanceLoss(lambda_var)(true_next_state, predicted_next_state)
    
    # Thought-Consistency Loss
    perturbed_next_state = predicted_next_state + torch.randn_like(predicted_next_state) * 0.01
    thought_loss = ThoughtConsistencyLoss()(true_next_state, perturbed_next_state)
    
    # Policy-Value Joint Loss
    pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates, true_value)
    
    # Action Diversity Reward
    action_embeddings = predicted_next_state  # Assuming actions are derived from state
    action_diversity = ActionDiversityReward(lambda_div)(action_embeddings)
    
    # Expected Thought Value (ETV) Loss
    # Placeholder: Replace with actual MCTS best values
    mcts_best_values = torch.zeros(value_estimates.size(0)).to(device)
    etv = ExpectedThoughtValueLoss()(mcts_best_values)
    
    # Exploration Regularization
    # Placeholder: Replace with actual visit counts
    visit_counts = torch.ones(predicted_next_state.size(0), input_dim).to(device)
    exploration = ExplorationRegularization(lambda_expl)(visit_counts)
    
    # KL Divergence Regularization
    # Placeholder: Replace with actual old and new policies
    old_policy = F.softmax(policy_logits.detach(), dim=-1)
    new_policy = F.softmax(policy_logits, dim=-1)
    kl_loss = KL_DivergenceLoss()(old_policy, new_policy)
    
    # Total Loss
    total_loss = (
        ce_loss +
        entropy_loss +
        variance_loss +
        cov_reg +
        dynamics_loss +
        thought_loss +
        pv_loss +
        action_diversity +
        etv +
        exploration +
        kl_loss
    )
    
    return total_loss


def train_epoch_world_model(world_model_components, train_loader, optimizer, scheduler, scaler, args, model_transformer, state_dim, embed_dim, input_dim):
    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent = world_model_components
    representation_network.train()
    dynamics_network.train()
    prediction_network.train()
    action_encoder.train()
    ppo_agent.policy_network.train()

    mcts = MCTS(prediction_network, dynamics_network, action_encoder, num_iterations=args.mcts_iterations, exploration_constant=args.mcts_exploration_constant)

    total_loss = 0.0
    optimizer.zero_grad()
    print(f"Starting World Model training epoch with {len(train_loader)} batches...")

    for i, batch in enumerate(train_loader):
        print(f"Processing batch {i+1}/{len(train_loader)}...")
        
        # Ensure batches are on the appropriate device for the Transformer
        src_batch = batch['input_ids'].to('cpu')  # Move to CPU for Transformer model
        tgt_batch = batch['labels'].to('cpu')     # Move to CPU for Transformer model

        with autocast(device_type='cuda'):
            print("Forward pass through Transformer (frozen)...")
            with torch.no_grad():
                transformer_output = model_transformer(src_batch, tgt_batch[:, :-1])

            # Move transformer output to the GPU for further processing
            transformer_output = transformer_output.to(device)

            # Encode actions directly on the GPU
            encoded_actions = action_encoder(tgt_batch[:, :-1].to(device))  # Move labels to GPU for encoding

            # World Model - Representation
            state_representation = representation_network(transformer_output)  # On GPU

            batch_size, seq_len, _ = state_representation.size()

            # Initialize list to collect predicted next states for the batch
            predicted_next_states = []

            # Iterate over each sample in the batch for MCTS
            for b in range(batch_size):
                # Create a State instance for the current sample
                current_state = State(state_representation[b].unsqueeze(0), dynamics_network, action_encoder)

                # Perform MCTS to find the best action
                best_action = mcts.search(current_state)

                # Create action sequence filled with best_action
                action_sequence = torch.full((1, seq_len), best_action, dtype=torch.long, device=device)

                # Get action embedding
                action_embedding = action_encoder(action_sequence)

                # Apply dynamics network
                predicted_next_state = dynamics_network(current_state.representation, action_embedding)

                predicted_next_states.append(predicted_next_state)

            # Concatenate predicted next states to form a batch
            predicted_next_state_batch = torch.cat(predicted_next_states, dim=0)

            # Prediction Network - Policy logits and value
            policy_logits, value_estimates = prediction_network(predicted_next_state_batch)


            # Define true_policy and true_value as placeholders on the GPU
            true_policy = torch.zeros_like(policy_logits).to(device)
            true_value = torch.zeros_like(value_estimates).to(device)

            # Compute PPO loss
            actions = torch.argmax(policy_logits, dim=-1)
            old_log_probs = torch.zeros_like(actions, dtype=torch.float32).to(device)
            returns = torch.zeros_like(actions, dtype=torch.float32).to(device)
            advantages = torch.zeros_like(actions, dtype=torch.float32).to(device)

            # Compute PPO loss using states
            ppo_loss = ppo_agent.compute_loss(state_representation, old_log_probs, actions, returns, advantages)

            # Compute InfoNCE Loss
            z_i = state_representation.view(batch_size * seq_len, state_dim)  # Shape: (batch_size * seq_len, state_dim)
            z_j = F.dropout(z_i, p=0.1, training=True)
            info_nce = InfoNCE_Loss()(z_i, z_j)

            # Compute other losses
            covariance = CovarianceRegularization()(predicted_next_state_batch.view(-1, predicted_next_state_batch.size(-1)))
            dynamics_loss = DynamicsPerformanceLoss()(torch.zeros_like(predicted_next_state_batch).to(device), predicted_next_state_batch)
            perturbed_next_state = predicted_next_state_batch + torch.randn_like(predicted_next_state_batch) * 0.01
            thought_loss = ThoughtConsistencyLoss()(torch.zeros_like(predicted_next_state_batch).to(device), perturbed_next_state)
            pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates.squeeze(-1), true_value.squeeze(-1))
            action_diversity = ActionDiversityReward()(encoded_actions.view(-1, embed_dim))
            mcts_best_values = torch.zeros(actions.size(0)).to(device)
            etv = ExpectedThoughtValueLoss()(mcts_best_values)
            visit_counts = torch.ones(actions.size(0), policy_logits.size(-1)).to(device)
            exploration = ExplorationRegularization()(visit_counts)
            old_policy = F.softmax(policy_logits.detach(), dim=-1)
            new_policy = F.softmax(policy_logits, dim=-1)
            kl_loss = KL_DivergenceLoss()(old_policy, new_policy)

            # Total Loss
            loss = (
                ppo_loss +
                info_nce +
                covariance +
                dynamics_loss +
                thought_loss +
                pv_loss +
                action_diversity +
                etv +
                exploration +
                kl_loss
            )
            loss = loss / args.accumulation_steps

        print("Backward pass...")
        scaler.scale(loss).backward()

        if (i + 1) % args.accumulation_steps == 0:
            print("Gradient clipping...")
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(
                [param for group in optimizer.param_groups for param in group['params']],
                args.max_grad_norm
            )

            print("Optimizer step...")
            scaler.step(optimizer)
            scaler.update()

            print("Zeroing gradients...")
            optimizer.zero_grad()

            print("Updating learning rate...")
            scheduler.step()

        total_loss += loss.item() * args.accumulation_steps
        print(f"Batch {i+1} completed. Current loss: {loss.item():.4f}")

    avg_loss = total_loss / len(train_loader)
    print(f"World Model training epoch completed. Average loss: {avg_loss:.4f}")
    return avg_loss



def evaluate_world_model(world_model_components, model_transformer, eval_loader, args):
    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent = world_model_components
    representation_network.eval()
    dynamics_network.eval()
    prediction_network.eval()
    action_encoder.eval()
    ppo_agent.policy_network.eval()
    
    total_loss = 0.0
    with torch.no_grad():
        for batch in eval_loader:
            src_batch = batch['input_ids'].to(device)
            tgt_batch = batch['labels'].to(device)
    
            # Forward pass through Transformer (on CPU)
            transformer_output = model_transformer(src_batch, tgt_batch[:, :-1])
    
            # Encode actions
            encoded_actions = action_encoder(tgt_batch[:, :-1].to(device))  # Move to GPU if necessary
    
            # World Model - Representation
            state = representation_network(transformer_output.to(device))
    
            # Dynamics Network - Predict next state
            predicted_next_state = dynamics_network(state, encoded_actions)
    
            # Prediction Network - Policy logits and value
            policy_logits, value_estimates = prediction_network(predicted_next_state)
    
            # Placeholder: Define true_policy and true_value
            # Replace these with actual targets from your environment or dataset
            true_policy = torch.zeros_like(policy_logits).to(device)
            true_value = torch.zeros_like(value_estimates).to(device)
    
            # Compute PPO loss
            # Placeholder: Replace with actual old_log_probs, actions, returns, and advantages
            old_log_probs = torch.zeros_like(policy_logits).to(device)
            actions = torch.argmax(policy_logits, dim=-1)
            returns = torch.zeros(actions.size(0)).to(device)
            advantages = torch.zeros(actions.size(0)).to(device)
    
            ppo_loss = ppo_agent.compute_loss(old_log_probs, actions, returns, advantages)
    
            # Compute other losses
            info_nce = InfoNCE_Loss()(state, state)  # Placeholder: replace with actual positive pairs
            covariance = CovarianceRegularization()(predicted_next_state.view(-1, predicted_next_state.size(-1)))
            dynamics_loss = DynamicsPerformanceLoss()(torch.zeros_like(predicted_next_state).to(device), predicted_next_state)
            perturbed_next_state = predicted_next_state + torch.randn_like(predicted_next_state) * 0.01
            thought_loss = ThoughtConsistencyLoss()(torch.zeros_like(predicted_next_state).to(device), perturbed_next_state)
            pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates.squeeze(-1), true_value.squeeze(-1))
            action_diversity = ActionDiversityReward()(encoded_actions.view(-1, encoded_actions.size(-1)))
            mcts_best_values = torch.zeros(actions.size(0)).to(device)  # Placeholder: replace with actual MCTS values
            etv = ExpectedThoughtValueLoss()(mcts_best_values)
            visit_counts = torch.ones(actions.size(0), policy_logits.size(-1)).to(device)  # Placeholder: replace with actual visit counts
            exploration = ExplorationRegularization()(visit_counts)
            old_policy = F.softmax(policy_logits.detach(), dim=-1)
            new_policy = F.softmax(policy_logits, dim=-1)
            kl_loss = KL_DivergenceLoss()(old_policy, new_policy)
    
            # Total Loss
            loss = (
                ppo_loss +
                info_nce +
                covariance +
                dynamics_loss +
                thought_loss +
                pv_loss +
                action_diversity +
                etv +
                exploration +
                kl_loss
            )
    
            total_loss += loss.item()
    
    avg_loss = total_loss / len(eval_loader)
    print(f"World Model evaluation completed. Average loss: {avg_loss:.4f}")
    return avg_loss


def main():
    args = parse_args()
    print("Arguments parsed successfully.")

    # Create save directory
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    print(f"Save directory created: {args.save_dir}")

    # Load tokenizer
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    print("Tokenizer loaded successfully.")

    # Define padding_idx and input dimension based on tokenizer
    padding_idx = tokenizer.pad_token_id
    input_dim = len(tokenizer)

    # Load data
    print("Loading and preprocessing data...")
    train_loader, eval_loader = load_data(args, tokenizer)
    print("Data loaded and preprocessed successfully.")

    # Define model parameters
    d_model = 512 # half to save space
    num_heads = 8
    num_layers = 6
    d_ff = 2048
    num_experts = 4
    output_dim = input_dim
    dropout = 0.1
    top_k = 2
    state_dim = 128
    action_dim = d_model
    hidden_dim = 512
    vocab_dim = len(tokenizer)
    # Initialize and load the Transformer model (on CPU)
    print("Initializing and loading Transformer model...")
    model_transformer = Transformer(input_dim, d_model, num_heads, num_layers, d_ff, num_experts, output_dim, dropout, top_k)
    model_transformer.load_state_dict(torch.load(args.transformer_model_path, map_location='cpu'))
    model_transformer.eval()
    model_transformer.to('cpu')
    print("Transformer model loaded and moved to CPU.")

    # Define World Model components
    representation_network = RepresentationNetwork(vocab_dim, d_model, state_dim).to(device)
    dynamics_network = DynamicsNetwork(state_dim, action_dim, hidden_dim).to(device)
    prediction_network = PredictionNetwork(state_dim, input_dim, 1).to(device)
    action_encoder = ActionEncoder(input_dim, action_dim).to(device)

    # Define Optimizers and Schedulers
    optimizer = optim.AdamW(
        list(representation_network.parameters()) +
        list(dynamics_network.parameters()) +
        list(prediction_network.parameters()) +
        list(action_encoder.parameters()), 
        lr=args.learning_rate, weight_decay=args.weight_decay
    )
    scheduler = CosineAnnealingLR(optimizer, T_max=args.num_epochs)
    scaler = GradScaler()

    # Initialize PPO Agent
    ppo_agent = PPOAgent(
        policy_network=prediction_network,
        optimizer=optim.AdamW(prediction_network.parameters(), lr=args.learning_rate),
        clip_epsilon=0.2,
        entropy_coef=0.01,
        value_coef=0.5
    )

    # Bundle World Model components
    world_model_components = (representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent)

    print("Setup complete. Starting training...")

    for epoch in range(args.num_epochs):
        print(f"Epoch {epoch + 1}/{args.num_epochs} started.")
        
        # Train World Model
        avg_train_loss = train_epoch_world_model(
            world_model_components, 
            train_loader, 
            optimizer, 
            scheduler, 
            scaler, 
            args, 
            model_transformer, 
            state_dim, 
            d_model, # this is the embedding dimension
            input_dim
        )

        print(f"World Model training epoch {epoch + 1} completed. Average loss: {avg_train_loss:.4f}")

        # Evaluate World Model
        avg_eval_loss = evaluate_world_model(
            world_model_components, 
            model_transformer, 
            eval_loader, 
            args
        )
        print(f"Evaluation for epoch {epoch + 1} completed. Average loss: {avg_eval_loss:.4f}")

        print(f"Epoch {epoch + 1}/{args.num_epochs}, Train Loss: {avg_train_loss:.4f}, Eval Loss: {avg_eval_loss:.4f}")

        # Save Models
        save_all_models(model_transformer, representation_network, dynamics_network, prediction_network, action_encoder, args.save_dir, epoch + 1)
        print(f"Models saved for epoch {epoch + 1}")

    print("Training completed.")
   

if __name__ == '__main__':
    main()