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.cuda.amp import autocast, GradScaler
from datasets import load_dataset
from transformers import AutoTokenizer

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

def parse_args():
    parser = argparse.ArgumentParser(description='Train or Inference with World Model and Tree of Thought.')
    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=4, 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=3, help='Number of MCTS Iterations')
    parser.add_argument('--mcts_exploration_constant', type=float, default=1.414, help='Exploration constant for MCTS')
    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('--mode', type=str, choices=['train', 'inference'], default='train', help='Mode: train or inference')
    parser.add_argument('--inference_mode', type=str, choices=['world_model', 'without_world_model', 'world_model_tree_of_thought'], default='world_model_tree_of_thought', help='Inference mode')
    parser.add_argument('--query', type=str, default='', help='Input query for inference')
    parser.add_argument('--train_mode', type=str, choices=['world_model', 'language_model'], default='world_model', help='Train world model or language model only')

    # Use parse_known_args to ignore unknown arguments
    args, unknown = parser.parse_known_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']

    def data_collator(data):
        return {
            '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,
        pin_memory=True,  # Speeds up transfer to GPU
        num_workers=4
    )
    eval_loader = DataLoader(
        eval_dataset,
        shuffle=False,
        batch_size=args.batch_size,
        collate_fn=data_collator,
        pin_memory=True,
        num_workers=4
    )

    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 Components

class ActionEncoder(nn.Module):
    def __init__(self, action_vocab_size, embed_dim):
        super(ActionEncoder, self).__init__()
        self.embedding = nn.Embedding(action_vocab_size, embed_dim)

    def forward(self, action_indices):
        """
        Args:
            action_indices (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_indices)

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, action_vocab_size, 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, action_vocab_size)  # Output size is action_vocab_size
        self.value_head = nn.Linear(state_dim, value_dim)

    def forward(self, state):
        """
        Args:
            state (torch.Tensor): State representation, 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)  # Shape: (batch_size, seq_len, action_vocab_size)
        value_estimates = self.value_head(norm_state).squeeze(-1)  # Shape: (batch_size, seq_len)
        return policy_logits, value_estimates


# Tree of Thought Components

class ThoughtNode:
    def __init__(self, name):
        self.name = name
        self.children = []
        self.parent = None

    def add_child(self, child_node):
        child_node.parent = self
        self.children.append(child_node)

# Function to build the Tree of Thought from your detailed structure
def build_tree_of_thought():
    # Create the root node
    root = ThoughtNode('Problem-Solving Process')

    # Level 1 nodes
    problem_identification = ThoughtNode('Problem Identification')
    problem_analysis = ThoughtNode('Problem Analysis')
    solution_generation = ThoughtNode('Solution Generation')
    implementation = ThoughtNode('Implementation')
    evaluation_adjustment = ThoughtNode('Evaluation and Adjustment')

    root.add_child(problem_identification)
    root.add_child(problem_analysis)
    root.add_child(solution_generation)
    root.add_child(implementation)
    root.add_child(evaluation_adjustment)

    # Problem Identification children
    B1 = ThoughtNode('Define the Problem')
    B2 = ThoughtNode('Identify Stakeholders')
    B3 = ThoughtNode('Determine Constraints')
    B4 = ThoughtNode('Recognize Problem Type')
    B5 = ThoughtNode('Historical Context')
    problem_identification.add_child(B1)
    problem_identification.add_child(B2)
    problem_identification.add_child(B3)
    problem_identification.add_child(B4)
    problem_identification.add_child(B5)

    # Define the Problem children
    B1a = ThoughtNode('Problem Statement Formulation')
    B1b = ThoughtNode('Scope Definition')
    B1c = ThoughtNode('Objective Setting')
    B1.add_child(B1a)
    B1.add_child(B1b)
    B1.add_child(B1c)

    # Identify Stakeholders children
    B2a = ThoughtNode('Stakeholder Mapping')
    B2b = ThoughtNode('Interest and Influence Analysis')
    B2c = ThoughtNode('Engagement Strategy')
    B2.add_child(B2a)
    B2.add_child(B2b)
    B2.add_child(B2c)

    # Determine Constraints children
    B3a = ThoughtNode('Resource Limitations')
    B3b = ThoughtNode('Time Constraints')
    B3c = ThoughtNode('Legal and Regulatory Constraints')
    B3.add_child(B3a)
    B3.add_child(B3b)
    B3.add_child(B3c)

    # Recognize Problem Type children
    B4a = ThoughtNode('Simple vs Complex')
    B4b = ThoughtNode('Known vs Unknown')
    B4c = ThoughtNode('Tame vs Wicked Problems')
    B4.add_child(B4a)
    B4.add_child(B4b)
    B4.add_child(B4c)

    # Historical Context children
    B5a = ThoughtNode('Previous Attempts')
    B5b = ThoughtNode('Lessons Learned')
    B5c = ThoughtNode('Environmental Factors')
    B5.add_child(B5a)
    B5.add_child(B5b)
    B5.add_child(B5c)

    # Problem Analysis children
    C1 = ThoughtNode('Root Cause Analysis')
    C2 = ThoughtNode('System Mapping')
    C3 = ThoughtNode('Data Collection')
    C4 = ThoughtNode('Impact Assessment')
    C5 = ThoughtNode('Theoretical Framework')
    problem_analysis.add_child(C1)
    problem_analysis.add_child(C2)
    problem_analysis.add_child(C3)
    problem_analysis.add_child(C4)
    problem_analysis.add_child(C5)

    # Root Cause Analysis children
    C1a = ThoughtNode('5 Whys Technique')
    C1b = ThoughtNode('Fishbone Diagram')
    C1c = ThoughtNode('Pareto Analysis')
    C1.add_child(C1a)
    C1.add_child(C1b)
    C1.add_child(C1c)

    # System Mapping children
    C2a = ThoughtNode('Causal Loop Diagrams')
    C2b = ThoughtNode('Stock and Flow Models')
    C2c = ThoughtNode('Network Analysis')
    C2.add_child(C2a)
    C2.add_child(C2b)
    C2.add_child(C2c)

    # Data Collection children
    C3a = ThoughtNode('Quantitative Data')
    C3b = ThoughtNode('Qualitative Data')
    C3c = ThoughtNode('Data Validation')
    C3.add_child(C3a)
    C3.add_child(C3b)
    C3.add_child(C3c)

    # Quantitative Data children
    C3a1 = ThoughtNode('Surveys and Questionnaires')
    C3a2 = ThoughtNode('Experimental Data')
    C3a3 = ThoughtNode('Big Data Analytics')
    C3a.add_child(C3a1)
    C3a.add_child(C3a2)
    C3a.add_child(C3a3)

    # Qualitative Data children
    C3b1 = ThoughtNode('Interviews')
    C3b2 = ThoughtNode('Focus Groups')
    C3b3 = ThoughtNode('Observational Studies')
    C3b.add_child(C3b1)
    C3b.add_child(C3b2)
    C3b.add_child(C3b3)

    # Data Validation children
    C3c1 = ThoughtNode('Statistical Validation')
    C3c2 = ThoughtNode('Cross-Validation')
    C3c3 = ThoughtNode('Expert Review')
    C3c.add_child(C3c1)
    C3c.add_child(C3c2)
    C3c.add_child(C3c3)

    # Impact Assessment children
    C4a = ThoughtNode('Environmental Impact')
    C4b = ThoughtNode('Social Impact')
    C4c = ThoughtNode('Economic Impact')
    C4.add_child(C4a)
    C4.add_child(C4b)
    C4.add_child(C4c)

    # Theoretical Framework children
    C5a = ThoughtNode('Literature Review')
    C5b = ThoughtNode('Conceptual Modeling')
    C5c = ThoughtNode('Hypothesis Formation')
    C5.add_child(C5a)
    C5.add_child(C5b)
    C5.add_child(C5c)

    # Solution Generation children
    D1 = ThoughtNode('Creative Problem Solving')
    D2 = ThoughtNode('Analytical Approach')
    D3 = ThoughtNode('Mathematical Computation')
    D4 = ThoughtNode('Decision Making')
    solution_generation.add_child(D1)
    solution_generation.add_child(D2)
    solution_generation.add_child(D3)
    solution_generation.add_child(D4)

    # Action Planning, Resource Allocation, Change Management children (implementation phase)
    E1 = ThoughtNode('Action Planning')
    E2 = ThoughtNode('Resource Allocation')
    E3 = ThoughtNode('Change Management')
    implementation.add_child(E1)
    implementation.add_child(E2)
    implementation.add_child(E3)

    # Verification, Performance Metrics, Feedback Loops, Continuous Improvement children (evaluation phase)
    F1 = ThoughtNode('Verification')
    F2 = ThoughtNode('Performance Metrics')
    F3 = ThoughtNode('Feedback Loops')
    F4 = ThoughtNode('Continuous Improvement')
    evaluation_adjustment.add_child(F1)
    evaluation_adjustment.add_child(F2)
    evaluation_adjustment.add_child(F3)
    evaluation_adjustment.add_child(F4)

    # Cross-Cutting Considerations children
    G = ThoughtNode('Cross-Cutting Considerations')
    root.add_child(G)

    # Cross-Cutting Considerations children
    G1 = ThoughtNode('Ethical Framework')
    G2 = ThoughtNode('Stakeholder Management')
    G3 = ThoughtNode('Interdisciplinary Connections')
    G4 = ThoughtNode('Technological Integration')
    G5 = ThoughtNode('Emotional Intelligence')
    G6 = ThoughtNode('Collaborative Problem Solving')
    G7 = ThoughtNode('Computational Considerations')  # Assuming H was intended as G7
    G8 = ThoughtNode('Order of Operations')  # Assuming I was intended as G8
    G9 = ThoughtNode('Critical Thinking')  # Assuming J was intended as G9
    G10 = ThoughtNode('Future Perspective')  # Assuming K was intended as G10
    G11 = ThoughtNode('Learning and Adaptation')  # Assuming L was intended as G11
    G.add_child(G1)
    G.add_child(G2)
    G.add_child(G3)
    G.add_child(G4)
    G.add_child(G5)
    G.add_child(G6)
    G.add_child(G7)
    G.add_child(G8)
    G.add_child(G9)
    G.add_child(G10)
    G.add_child(G11)

    # Ethical Framework children
    G1a = ThoughtNode('Value-based Decision Making')
    G1b = ThoughtNode('Long-term Consequences')
    G1.add_child(G1a)
    G1.add_child(G1b)

    # Value-based Decision Making children
    G1a1 = ThoughtNode('Ethical Theories Application')
    G1a2 = ThoughtNode('Moral Dilemma Resolution')
    G1a.add_child(G1a1)
    G1a.add_child(G1a2)

    # Long-term Consequences children
    G1b1 = ThoughtNode('Sustainability Assessment')
    G1b2 = ThoughtNode('Intergenerational Impact')
    G1b.add_child(G1b1)
    G1b.add_child(G1b2)

    # Stakeholder Management children
    G2a = ThoughtNode('Direct Stakeholders')
    G2b = ThoughtNode('Indirect Stakeholders')
    G2c = ThoughtNode('Conflicting Interests')
    G2.add_child(G2a)
    G2.add_child(G2b)
    G2.add_child(G2c)

    # Conflicting Interests children
    G2c1 = ThoughtNode('Negotiation Strategies')
    G2c2 = ThoughtNode('Conflict Resolution Techniques')
    G2c.add_child(G2c1)
    G2c.add_child(G2c2)

    # Interdisciplinary Connections children
    G3a = ThoughtNode('Related Fields')
    G3b = ThoughtNode('Cross-disciplinary Impact')
    G3.add_child(G3a)
    G3.add_child(G3b)

    # Related Fields children
    G3a1 = ThoughtNode('Cross-domain Knowledge Transfer')
    G3a2 = ThoughtNode('Interdisciplinary Collaboration')
    G3a.add_child(G3a1)
    G3a.add_child(G3a2)

    # Cross-disciplinary Impact children
    G3b1 = ThoughtNode('Synergy Identification')
    G3b2 = ThoughtNode('Holistic Impact Assessment')
    G3b.add_child(G3b1)
    G3b.add_child(G3b2)

    # Technological Integration children
    G4a = ThoughtNode('AI-assisted Problem Solving')
    G4b = ThoughtNode('Data-driven Insights')
    G4c = ThoughtNode('Digital Collaboration Tools')
    G4.add_child(G4a)
    G4.add_child(G4b)
    G4.add_child(G4c)

    # AI-assisted Problem Solving children
    G4a1 = ThoughtNode('Machine Learning Models')
    G4a2 = ThoughtNode('Natural Language Processing')
    G4a.add_child(G4a1)
    G4a.add_child(G4a2)

    # Data-driven Insights children
    G4b1 = ThoughtNode('Big Data Analytics')
    G4b2 = ThoughtNode('Predictive Modeling')
    G4b.add_child(G4b1)
    G4b.add_child(G4b2)

    # Digital Collaboration Tools children
    G4c1 = ThoughtNode('Project Management Platforms')
    G4c2 = ThoughtNode('Virtual Reality Collaboration')
    G4c.add_child(G4c1)
    G4c.add_child(G4c2)

    # Emotional Intelligence children
    G5a = ThoughtNode('Self-Awareness')
    G5b = ThoughtNode('Empathy')
    G5c = ThoughtNode('Stress Management')
    G5.add_child(G5a)
    G5.add_child(G5b)
    G5.add_child(G5c)

    # Self-Awareness children
    G5a1 = ThoughtNode('Emotional Recognition')
    G5a2 = ThoughtNode('Personal Bias Identification')
    G5a.add_child(G5a1)
    G5a.add_child(G5a2)

    # Empathy children
    G5b1 = ThoughtNode('Perspective Taking')
    G5b2 = ThoughtNode('Active Listening')
    G5b.add_child(G5b1)
    G5b.add_child(G5b2)

    # Stress Management children
    G5c1 = ThoughtNode('Mindfulness Techniques')
    G5c2 = ThoughtNode('Resilience Building')
    G5c.add_child(G5c1)
    G5c.add_child(G5c2)

    # Collaborative Problem Solving children
    G6a = ThoughtNode('Team Dynamics')
    G6b = ThoughtNode('Communication Strategies')
    G6c = ThoughtNode('Conflict Resolution')
    G6.add_child(G6a)
    G6.add_child(G6b)
    G6.add_child(G6c)

    # Team Dynamics children
    G6a1 = ThoughtNode('Team Formation Strategies')
    G6a2 = ThoughtNode('Role Assignment')
    G6a.add_child(G6a1)
    G6a.add_child(G6a2)

    # Communication Strategies children
    G6b1 = ThoughtNode('Clear Messaging')
    G6b2 = ThoughtNode('Feedback Mechanisms')
    G6b.add_child(G6b1)
    G6b.add_child(G6b2)

    # Conflict Resolution children
    G6c1 = ThoughtNode('Mediation Techniques')
    G6c2 = ThoughtNode('Consensus Building')
    G6c.add_child(G6c1)
    G6c.add_child(G6c2)

    # Computational Considerations children
    G7a = ThoughtNode('CPU Operations')
    G7b = ThoughtNode('GPU Parallelization')
    G7c = ThoughtNode('Floating-Point Precision')
    G7.add_child(G7a)
    G7.add_child(G7b)
    G7.add_child(G7c)

    # CPU Operations children
    G7a1 = ThoughtNode('Instruction Set Architecture')
    G7a2 = ThoughtNode('Pipelining and Parallelism')
    G7a.add_child(G7a1)
    G7a.add_child(G7a2)

    # GPU Parallelization children
    G7b1 = ThoughtNode('CUDA Programming')
    G7b2 = ThoughtNode('OpenCL Framework')
    G7b.add_child(G7b1)
    G7b.add_child(G7b2)

    # Floating-Point Precision children
    G7c1 = ThoughtNode('IEEE 754 Standard')
    G7c2 = ThoughtNode('Error Propagation Analysis')
    G7c.add_child(G7c1)
    G7c.add_child(G7c2)

    # Order of Operations children
    G8a = ThoughtNode('Parentheses')
    G8b = ThoughtNode('Exponents')
    G8c = ThoughtNode('Multiplication and Division')
    G8d = ThoughtNode('Addition and Subtraction')
    G8.add_child(G8a)
    G8.add_child(G8b)
    G8.add_child(G8c)
    G8.add_child(G8d)

    # Critical Thinking children
    G9a = ThoughtNode('Assumptions Questioning')
    G9b = ThoughtNode('Bias Recognition')
    G9.add_child(G9a)
    G9.add_child(G9b)

    # Assumptions Questioning children
    G9a1 = ThoughtNode('Socratic Questioning')
    G9a2 = ThoughtNode('Devil\'s Advocate Approach')
    G9a.add_child(G9a1)
    G9a.add_child(G9a2)

    # Bias Recognition children
    G9b1 = ThoughtNode('Cognitive Bias Identification')
    G9b2 = ThoughtNode('Debiasing Techniques')
    G9b.add_child(G9b1)
    G9b.add_child(G9b2)

    # Future Perspective children
    G10a = ThoughtNode('Short-term Projections')
    G10b = ThoughtNode('Long-term Scenarios')
    G10c = ThoughtNode('Potential Impacts')
    G10.add_child(G10a)
    G10.add_child(G10b)
    G10.add_child(G10c)

    # Short-term Projections children
    G10a1 = ThoughtNode('Trend Analysis')
    G10a2 = ThoughtNode('Scenario Planning')
    G10a.add_child(G10a1)
    G10a.add_child(G10a2)

    # Long-term Scenarios children
    G10b1 = ThoughtNode('Futures Wheel')
    G10b2 = ThoughtNode('Backcasting')
    G10b.add_child(G10b1)
    G10b.add_child(G10b2)

    # Potential Impacts children
    G10c1 = ThoughtNode('Risk Assessment')
    G10c2 = ThoughtNode('Opportunity Identification')
    G10c.add_child(G10c1)
    G10c.add_child(G10c2)

    # Learning and Adaptation children
    G11a = ThoughtNode('Reflective Practice')
    G11b = ThoughtNode('Knowledge Transfer')
    G11c = ThoughtNode('Adaptive Problem Solving')
    G11.add_child(G11a)
    G11.add_child(G11b)
    G11.add_child(G11c)

    # Reflective Practice children
    G11a1 = ThoughtNode('After Action Review')
    G11a2 = ThoughtNode('Learning Journals')
    G11a.add_child(G11a1)
    G11a.add_child(G11a2)

    # Knowledge Transfer children
    G11b1 = ThoughtNode('Best Practice Documentation')
    G11b2 = ThoughtNode('Mentoring Programs')
    G11b.add_child(G11b1)
    G11b.add_child(G11b2)

    # Adaptive Problem Solving children
    G11c1 = ThoughtNode('Iterative Approaches')
    G11c2 = ThoughtNode('Flexibility in Methodology')
    G11c.add_child(G11c1)
    G11c.add_child(G11c2)

    return root

def traverse_tree(node, action_list):
    if node.name not in action_list:
        action_list.append(node.name)
    for child in node.children:
        traverse_tree(child, action_list)

class MCTSNode:
    __slots__ = [
        'state',
        'parent',
        'action',
        'children',
        'visit_count',
        'value_sum',
        'prior',
        'cached_policy',
        'cached_value',
        'thought_node'  # Added to keep track of the current thought node
    ]

    def __init__(self, state, thought_node, parent=None, action=None):
        self.state = state
        self.thought_node = thought_node  # Reference to the ThoughtNode
        self.parent = parent
        self.action = action
        self.children = {}
        self.visit_count = 0
        self.value_sum = 0.0
        self.prior = 0.0
        self.cached_policy = None
        self.cached_value = None

    def expand(self, priors):
        """
        Expand the node by adding all valid child nodes from the thought tree.

        Args:
            priors (dict): A dictionary mapping action names to prior probabilities.
        """
        for child_thought_node in self.thought_node.children:
            action = child_thought_node.name  # Action name
            if action not in self.children:
                # Assume batch size of 1 for individual nodes
                child_state = self.state.apply_action(action)
                child_node = MCTSNode(
                    state=child_state,
                    thought_node=child_thought_node,
                    parent=self,
                    action=action
                )
                child_node.prior = priors.get(action, 1.0 / len(self.thought_node.children))  # Default prior if not provided
                self.children[action] = child_node

    def is_leaf(self):
        return len(self.children) == 0

    def ucb_score(self, total_visits, exploration_constant=math.sqrt(2)):
        if self.visit_count == 0:
            return float('inf')
        avg_value = self.value_sum / self.visit_count
        exploration_term = exploration_constant * self.prior * math.sqrt(total_visits) / (1 + self.visit_count)
        return avg_value + exploration_term

class MCTS:
    def __init__(self, prediction_network, dynamics_network, action_encoder, num_iterations=10, exploration_constant=math.sqrt(2)):
        self.prediction_network = prediction_network
        self.dynamics_network = dynamics_network
        self.action_encoder = action_encoder
        self.num_iterations = num_iterations
        self.exploration_constant = exploration_constant
        self.cache = {}

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

        Args:
            root_state (State): The root state from which to start the search.

        Returns:
            str: The best action to take from the root state.
        """
        root_node = MCTSNode(state=root_state, thought_node=root_state.thought_node)

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

        best_action = self.best_action(root_node)
        return best_action

    def select(self, node):
        while not node.is_leaf():
            total_visits = sum(child.visit_count for child in node.children.values())
            _, node = max(
                node.children.items(),
                key=lambda item: item[1].ucb_score(total_visits, self.exploration_constant)
            )
        return node

    def evaluate(self, node):
        # Use the prediction network to get policy and value estimates
        state_representation = node.state.representation  # Shape: (batch_size=1, seq_len, state_dim)
        policy_logits, value_estimate = self.prediction_network(state_representation)
        value_estimate = value_estimate.item()  # Convert tensor to scalar

        # Convert policy logits to probabilities
        policy_probs = F.softmax(policy_logits, dim=-1).squeeze(0)  # Shape: (seq_len, action_vocab_size)
        # For simplicity, use the last time step's policy
        policy_probs = policy_probs[-1]  # Shape: (action_vocab_size,)

        # Map policy probabilities to the actions available from the current thought node
        priors = {}
        for child in node.thought_node.children:
            action_name = child.name
            action_idx = action_to_index.get(action_name, None)
            if action_idx is not None and action_idx < policy_probs.size(0):
                priors[action_name] = policy_probs[action_idx].item()
            else:
                priors[action_name] = 1.0 / len(node.thought_node.children)  # Uniform prior if not found

        # Expand the node
        node.expand(priors)

        return value_estimate

    def backpropagate(self, node, value):
        while node is not None:
            node.visit_count += 1
            node.value_sum += value
            node = node.parent

    def best_action(self, root_node):
        # Select the child with the highest visit count
        best_child = max(root_node.children.values(), key=lambda n: n.visit_count)
        return best_child.action

class State:
    def __init__(self, representation, dynamics_network, action_encoder, thought_node):
        """
        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
            thought_node (ThoughtNode): The current node in the Tree of Thought
        """
        self.representation = representation  # Shape: (batch_size, seq_len, state_dim)
        self.dynamics_network = dynamics_network
        self.action_encoder = action_encoder
        self.thought_node = thought_node  # Current position in the Tree of Thought

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

        Args:
            action (str): The action to apply (the name of the ThoughtNode)

        Returns:
            State: The new state after applying the action
        """
        # Find the corresponding child node in the thought tree
        next_thought_node = None
        for child in self.thought_node.children:
            if child.name == action:
                next_thought_node = child
                break
        if next_thought_node is None:
            raise ValueError(f"Action '{action}' is not valid from the current thought node.")

        # Encode action
        action_index = torch.tensor([[action_to_index[action]]], device=self.representation.device)
        action_embedding = self.action_encoder(action_index)

        # Predict the next state using the Dynamics Network
        next_state_representation = self.dynamics_network(self.representation, action_embedding)

        return State(
            representation=next_state_representation,
            dynamics_network=self.dynamics_network,
            action_encoder=self.action_encoder,
            thought_node=next_thought_node
        )

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 using reshape
        policy_logits = policy_logits.reshape(-1, num_actions)  # Shape: (batch_size * seq_len, num_actions)
        value_estimates = value_estimates.view(-1)
        actions = actions.reshape(-1)                           # Shape: (batch_size * seq_len)
        old_log_probs = old_log_probs.reshape(-1)               # Shape: (batch_size * seq_len)
        returns = returns.view(-1)
        advantages = advantages.reshape(-1)                     # Shape: (batch_size * seq_len)

        # Ensure value_estimates and returns are the same size
        if value_estimates.size() != returns.size():
            print(f"Shape mismatch: value_estimates shape: {value_estimates.size()}, returns shape: {returns.size()}")
            value_estimates = value_estimates[:returns.size(0)]

        # 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 infer(query, world_model_components, root_thought_node, tokenizer, max_length=20, inference_mode='world_model'):
    """
    Perform inference given a query, utilizing the Tree of Thought and MCTS.

    Args:
        query (str): The input query or prompt.
        world_model_components (tuple): Tuple containing the model components.
        root_thought_node (ThoughtNode): The root node of the Tree of Thought.
        tokenizer (transformers.PreTrainedTokenizer): The tokenizer used.
        max_length (int): Maximum length for the generated sequence.
        inference_mode (str): Inference mode ('world_model', 'without_world_model', 'world_model_tree_of_thought')

    Returns:
        List[str] or str: The sequence of actions (thoughts) selected or generated text.
    """
    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent, model_transformer = world_model_components

    # Tokenize and encode the query
    input_ids = tokenizer.encode(query, return_tensors='pt').to(device)
    attention_mask = (input_ids != tokenizer.pad_token_id).long()

    if inference_mode == 'without_world_model':
        # Directly use the transformer model to generate text
        with torch.no_grad():
            generated_ids, entropies, variances = model_transformer.generate(src=input_ids, tokenizer=tokenizer, max_length=max_length, temperature=args.temperature)
        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        return generated_text

    else:
        # Use the world model components
        with torch.no_grad():
            transformer_output = model_transformer(input_ids, input_ids)
            # Get the initial state representation
            initial_representation = representation_network(transformer_output)  # Shape: (batch_size=1, seq_len, state_dim)
            initial_state = State(
                representation=initial_representation,
                dynamics_network=dynamics_network,
                action_encoder=action_encoder,
                thought_node=root_thought_node
            )

            if inference_mode == 'world_model_tree_of_thought':
                # Use MCTS with Tree of Thought
                mcts = MCTS(prediction_network, dynamics_network, action_encoder, num_iterations=args.mcts_iterations, exploration_constant=args.mcts_exploration_constant)
                current_state = initial_state
                thought_sequence = []

                for _ in range(max_length):
                    best_action = mcts.search(current_state)
                    thought_sequence.append(best_action)

                    # Apply the best action to get the next state
                    current_state = current_state.apply_action(best_action)

                    # Check if we've reached a leaf node (no further actions)
                    if len(current_state.thought_node.children) == 0:
                        break

                return thought_sequence
            else:
                # Use the world model without Tree of Thought
                # For simplicity, we will generate actions based on the prediction network
                policy_logits, _ = prediction_network(initial_state.representation)
                policy_probs = F.softmax(policy_logits, dim=-1)
                # Select actions with highest probabilities
                top_actions = torch.argmax(policy_probs, dim=-1)
                generated_actions = [index_to_action[idx.item()] for idx in top_actions[0]]
                return generated_actions

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()

    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)}...")

        # Move batches to the device
        src_batch = batch['input_ids'].to(device)
        tgt_batch = batch['labels'].to(device)

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

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

            # For simplicity, let's assume true actions are provided (e.g., next tokens)
            true_actions = tgt_batch[:, :-1]  # Shape: (batch_size, seq_len)
            action_sequences = true_actions

            # Get action embeddings
            action_embeddings = action_encoder(action_sequences)  # Shape: (batch_size, seq_len, embed_dim)

            # Apply dynamics network
            predicted_next_state_batch = dynamics_network(state_representation, action_embeddings)  # Shape: (batch_size, seq_len, state_dim)

            # Prediction Network - Policy logits and value
            policy_logits, value_estimates = prediction_network(predicted_next_state_batch)
    # value_estimates now has shape (batch_size, seq_len)

            # Define true_policy and true_value as placeholders on the GPU
            true_policy = F.one_hot(true_actions, num_classes=input_dim).float()  # Shape: (batch_size, seq_len, input_dim)
            true_value = torch.zeros_like(value_estimates).to(device)

            # Compute PPO loss
            actions_selected = true_actions  # Shape: (batch_size, seq_len)
            old_log_probs = torch.zeros_like(actions_selected, dtype=torch.float32).to(device)
            returns = torch.zeros_like(actions_selected, dtype=torch.float32).to(device)
            advantages = torch.zeros_like(actions_selected, dtype=torch.float32).to(device)

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

            # Compute InfoNCE Loss
            z_i = state_representation.view(-1, 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()(state_representation, predicted_next_state_batch)
            perturbed_next_state = predicted_next_state_batch + torch.randn_like(predicted_next_state_batch) * 0.01
            thought_loss = ThoughtConsistencyLoss()(predicted_next_state_batch, perturbed_next_state)
            pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates.squeeze(-1), true_value.squeeze(-1))
            action_diversity = ActionDiversityReward()(action_embeddings.view(-1, embed_dim))
            mcts_best_values = torch.zeros(actions_selected.size(0)).to(device)
            etv = ExpectedThoughtValueLoss()(mcts_best_values)
            visit_counts = torch.ones(actions_selected.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 or (i + 1) == len(train_loader):
            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 train_epoch_language_model(model, train_loader, optimizer, scheduler, scaler, args):
    model.train()
    total_loss = 0.0
    optimizer.zero_grad()
    print(f"Starting Language Model training epoch with {len(train_loader)} batches...")

    for i, batch in enumerate(train_loader):
        input_ids = batch['input_ids'].to(device)
        labels = batch['labels'].to(device)

        with autocast():
            outputs = model(input_ids, input_ids)
            logits = outputs.view(-1, outputs.size(-1))
            labels = labels.view(-1)
            loss = F.cross_entropy(logits, labels, ignore_index=model.embedding.padding_idx)
            loss = loss / args.accumulation_steps

        scaler.scale(loss).backward()

        if (i + 1) % args.accumulation_steps == 0 or (i + 1) == len(train_loader):
            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
            )
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
            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"Language Model training epoch completed. Average loss: {avg_loss:.4f}")
    return avg_loss



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

    # Create save directory
    os.makedirs(args.save_dir, exist_ok=True)
    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)

    # Initialize the Transformer model on GPU
    print("Initializing Transformer model...")
    model_transformer = Transformer(
        input_dim=input_dim,
        d_model=128,
        num_heads=4,
        num_layers=4,
        d_ff=256,
        num_experts=2,
        output_dim=input_dim,
        dropout=0.1,
        top_k=2
    ).to(device)
    model_transformer.train()
    print("Transformer model initialized on device.")

    # Define model parameters (adjusted for speed)
    d_model = 128
    state_dim = 128
    action_dim = d_model
    hidden_dim = 256
    vocab_dim = input_dim
    embed_dim = d_model

    # 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)

    # 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, model_transformer)

    if args.mode == 'train':
        print("Loading and preprocessing data...")
        train_loader, eval_loader = load_data(args, tokenizer)
        print("Data loaded and preprocessed successfully.")

        # Optimizer and Scheduler
        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
        ) if args.train_mode == 'world_model' else optim.AdamW(model_transformer.parameters(), lr=args.learning_rate)
        scheduler = CosineAnnealingLR(optimizer, T_max=args.num_epochs)
        scaler = GradScaler()

        print(f"Starting {args.train_mode} training...")

        for epoch in range(args.num_epochs):
            if args.train_mode == 'world_model':
                avg_loss = train_epoch_world_model(
                    world_model_components,
                    train_loader,
                    optimizer,
                    scheduler,
                    scaler,
                    args,
                    model_transformer,
                    state_dim,
                    embed_dim,
                    input_dim
                )
            else:
                avg_loss = train_epoch_language_model(
                    model_transformer,
                    train_loader,
                    optimizer,
                    scheduler,
                    scaler,
                    args
                )

            print(f"{args.train_mode.capitalize()} training epoch {epoch + 1} completed. Average loss: {avg_loss:.4f}")

            if args.train_mode == 'world_model':
                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}")
            else:
                torch.save(model_transformer.state_dict(), os.path.join(args.save_dir, f'language_model_epoch_{epoch + 1}.pt'))
                print(f"Language model saved for epoch {epoch + 1}")

        print("Training completed.")

    elif args.mode == 'inference':
        # Build Tree of Thought if needed
        tree_root = build_tree_of_thought()
        # Generate action list
        action_list = []
        traverse_tree(tree_root, action_list)

        # Create mappings
        global action_to_index, index_to_action
        action_to_index = {action: idx for idx, action in enumerate(action_list)}
        index_to_action = {idx: action for action, idx in action_to_index.items()}
        action_vocab_size = len(action_list)

        # Update action encoder and prediction network with new vocab size
        action_encoder = ActionEncoder(action_vocab_size, action_dim).to(device)
        prediction_network = PredictionNetwork(state_dim, action_vocab_size, 1).to(device)

        # Load the saved models
        # Assuming the models are saved after training
        # You need to adjust the paths and epoch numbers as necessary
        model_transformer.load_state_dict(torch.load(os.path.join(args.save_dir, 'transformer_model_epoch_2.pt')))
        representation_network.load_state_dict(torch.load(os.path.join(args.save_dir, 'representation_network_epoch_2.pt')))
        dynamics_network.load_state_dict(torch.load(os.path.join(args.save_dir, 'dynamics_network_epoch_2.pt')))
        saved_state_dict = torch.load(os.path.join(args.save_dir, 'prediction_network_epoch_2.pt'))
        prediction_network.policy_head = nn.Linear(prediction_network.state_dim, 50257)  # Update to match saved model size
        prediction_network.load_state_dict(saved_state_dict, strict=False)

        # Resize `policy_head` back to 158 after loading
        prediction_network.policy_head = nn.Linear(prediction_network.state_dim, 158).to(device)



        action_encoder.load_state_dict(torch.load(os.path.join(args.save_dir, 'action_encoder_epoch_2.pt')))

        # Prepare the components
        world_model_components = (representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent, model_transformer)

        # Perform inference
        if not args.query:
            args.query = input("Please enter your query: ")

        result = infer(args.query, world_model_components, tree_root, tokenizer, inference_mode=args.inference_mode)

        if args.inference_mode == 'without_world_model':
            print("Generated Text:")
            print(result)
        else:
            print("Generated Thought Sequence:")
            for thought in result:
                print(thought)

if __name__ == '__main__':
    main()