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
from typing import List, Tuple

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='inference', 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')
    parser.add_argument('--beam_size', type=int, default=5, help='Beam size for beam search')
    parser.add_argument('--n_tokens_predict', type=int, default=3, help='Number of tokens to predict at each step')
    parser.add_argument('--load_model', type=str, default=None, 
                        help='Path to load saved model. If not provided, a new model will be initialized.')


    # 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_with_beam_search(self, src, tokenizer, beam_size=5, max_length=20, n_tokens_predict=3, temperature=1.0):
        """
        Generate sequences using beam search with multi-token prediction.
        
        Args:
            src (torch.Tensor): Source input tensor of shape (batch_size, seq_len)
            tokenizer: Tokenizer to access special tokens
            beam_size (int): Size of the beam for beam search
            max_length (int): Maximum length of the generated sequence
            n_tokens_predict (int): Number of tokens to predict at each step
            temperature (float): Temperature parameter for softmax
        
        Returns:
            List[Tuple[torch.Tensor, float]]: List of (sequence, score) tuples
        """
        batch_size = src.size(0)
        device = src.device
        vocab_size = self.output_layer.out_features

        # Encode the source
        src_enc = self.encode(src)

        # Initialize beam
        beam = [(torch.full((batch_size, 1), tokenizer.bos_token_id, dtype=torch.long, device=device),
                 0.0,  # log probability
                 torch.zeros(batch_size, device=device),  # cumulative entropy
                 torch.zeros(batch_size, device=device))]  # cumulative variance

        for _ in range(max_length // n_tokens_predict):
            all_candidates = []
            for seq, score, cum_entropy, cum_variance in beam:
                if seq[:, -1].item() == tokenizer.eos_token_id:
                    all_candidates.append((seq, score, cum_entropy, cum_variance))
                    continue

                # Predict next n tokens
                logits = self.predict_next_n_tokens(src_enc, seq, n_tokens_predict)
                
                # Calculate probabilities, entropy, and variance
                probs = F.softmax(logits / temperature, dim=-1)
                entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)
                variance = torch.var(probs, dim=-1)

                # Sample top-k tokens for each position
                topk_probs, topk_indices = torch.topk(probs, k=beam_size, dim=-1)
                
                # Generate all possible continuations
                for i in range(beam_size ** n_tokens_predict):
                    indices = [i // (beam_size ** j) % beam_size for j in range(n_tokens_predict)]
                    new_tokens = topk_indices[:, range(n_tokens_predict), indices]
                    new_seq = torch.cat([seq, new_tokens], dim=-1)
                    new_score = score + torch.sum(torch.log(topk_probs[:, range(n_tokens_predict), indices]))
                    new_entropy = cum_entropy + torch.sum(entropy[:, indices])
                    new_variance = cum_variance + torch.sum(variance[:, indices])
                    
                    all_candidates.append((new_seq, new_score, new_entropy, new_variance))

            # Select top beam_size candidates
            beam = sorted(all_candidates, key=lambda x: x[1] - 0.1 * x[2] + 0.05 * x[3], reverse=True)[:beam_size]

            # Stop if all beams have ended
            if all(seq[:, -1].item() == tokenizer.eos_token_id for seq, _, _, _ in beam):
                break

        return [(seq, score) for seq, score, _, _ in beam]

    def encode(self, src):
        src_emb = self.embedding(src) * math.sqrt(self.d_model)
        src_emb = src_emb.transpose(0, 1)
        src_emb = self.rotary_positional_encoding(src_emb)
        src_emb = src_emb.transpose(0, 1)
        src_enc = src_emb
        for layer in self.encoder_layers:
            src_enc = layer(src_enc)
        return src_enc

    def predict_next_n_tokens(self, src_enc, tgt_seq, n_tokens):
        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[:, -1:])
        return output.repeat(1, n_tokens, 1)

# 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)  # Shape: (batch_size, seq_len, d_model)
        # Then project down from d_model to state_dim
        state = self.linear(projected_output)  # Shape: (batch_size, seq_len, state_dim)
        state = self.norm(state)  # Shape: (batch_size, seq_len, state_dim)
        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, state_dim)
            action (torch.Tensor): Action embedding, shape (batch_size, action_dim)

        Returns:
            torch.Tensor: Predicted next state, shape (batch_size, 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, 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, action_vocab_size)
        value_estimates = self.value_head(norm_state).squeeze(-1)  # Shape: (batch_size)
        return policy_logits, value_estimates




class MCTSNode:
    __slots__ = [
        'state',
        'parent',
        'action',
        'children',
        'visit_count',
        'value_sum',
        'prior',
        'cached_policy',
        'cached_value',
        'thought_node',
        'entropy',
        'variance'
    ]

    def __init__(self, state, thought_node, parent=None, action=None):
        self.state = state
        self.thought_node = thought_node
        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
        self.entropy = 0.0
        self.variance = 0.0

    def expand(self, priors):
        for child_thought_node in self.thought_node.children:
            action = child_thought_node.name
            if action not in self.children:
                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))
                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')  # Ensure unvisited nodes are selected first
        avg_value = self.value_sum / self.visit_count
        exploration_term = exploration_constant * self.prior * math.sqrt(total_visits) / (1 + self.visit_count)
        entropy_term = -0.1 * self.entropy  # Slightly prefer lower entropy
        variance_term = 0.05 * self.variance  # Slightly prefer higher variance
        return avg_value + exploration_term + entropy_term + variance_term


class MCTS:
    def __init__(self, prediction_network, dynamics_network, action_encoder, num_iterations=10, exploration_constant=math.sqrt(2), beam_size=5, n_tokens_predict=3):
        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.beam_size = beam_size
        self.n_tokens_predict = n_tokens_predict
        self.cache = {}

    def search_with_beam(self, root_state):
        root_node = MCTSNode(state=root_state, thought_node=root_state.thought_node)
        
        # Evaluate the root node and backpropagate
        value_estimate = self.evaluate(root_node)  # Evaluate and expand root_node
        self.backpropagate(root_node, value_estimate)  # Backpropagate the value

        beam = [(root_node, 0.0, 0.0, 0.0, [])]  # (node, score, cum_entropy, cum_variance, action_sequence)

        for iteration in range(self.num_iterations):
            all_candidates = []
            for node, score, cum_entropy, cum_variance, action_sequence in beam:
                if node.is_leaf():
                    value_estimate = self.evaluate(node)
                    self.backpropagate(node, value_estimate)  # Backpropagate after evaluation
                if len(node.children) == 0:
                    continue  # No children to expand

                total_visits = sum(child.visit_count for child in node.children.values())
                # Select top actions based on UCB score
                sorted_children = sorted(
                    node.children.items(),
                    key=lambda item: item[1].ucb_score(total_visits, self.exploration_constant),
                    reverse=True
                )[:self.beam_size]

                for selected_action, selected_node in sorted_children:
                    current_node = selected_node
                    current_sequence = action_sequence + [selected_action]
                    current_score = score
                    current_entropy = cum_entropy + selected_node.entropy
                    current_variance = cum_variance + selected_node.variance

                    # Predict n_tokens_predict actions
                    for _ in range(self.n_tokens_predict):
                        if current_node.is_leaf():
                            value_estimate = self.evaluate(current_node)
                            self.backpropagate(current_node, value_estimate)  # Backpropagate after evaluation
                        if len(current_node.children) == 0:
                            break  # No more actions
                        total_visits = sum(child.visit_count for child in current_node.children.values())
                        next_action, next_node = max(
                            current_node.children.items(),
                            key=lambda item: item[1].ucb_score(total_visits, self.exploration_constant)
                        )
                        current_sequence.append(next_action)
                        
                        # Prevent division by zero by ensuring visit_count > 0
                        if next_node.visit_count > 0:
                            current_score += next_node.value_sum / next_node.visit_count
                        else:
                            # Assign a default value or handle the zero division case
                            current_score += 0.0  # Alternatively, use a small epsilon or skip

                        current_entropy += next_node.entropy
                        current_variance += next_node.variance
                        current_node = next_node

                    all_candidates.append((current_node, current_score, current_entropy, current_variance, current_sequence))

            if not all_candidates:
                break  # No more candidates to expand

            # Select top beam_size candidates
            beam = sorted(all_candidates, key=lambda x: x[1] - 0.1 * x[2] + 0.05 * x[3], reverse=True)[:self.beam_size]
            print(f"Iteration {iteration + 1}: Beam size after sorting: {len(beam)}")  # Debug

        if beam:
            best_sequence = beam[0][4]
            return best_sequence
        else:
            return []



    def search(self, 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)
        
        return self.best_action_sequence(root_node)

    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):
        # Extract the last time step
        state_representation = node.state.representation[:, -1, :]  # Shape: (batch_size=1, state_dim)
        print(f"Evaluating node with state_representation shape: {state_representation.shape}")  # Debug
        policy_logits, value_estimate = self.prediction_network(state_representation)
        print(f"Policy logits shape: {policy_logits.shape}, Value estimate shape: {value_estimate.shape}")  # Debug
        value_estimate = value_estimate.item()  # Now safe as batch_size=1

        policy_probs = F.softmax(policy_logits, dim=-1).squeeze(0)  # Shape: (action_vocab_size,)
        print(f"Policy probabilities shape: {policy_probs.shape}")  # Debug

        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)

        node.expand(priors)

        # Calculate entropy and variance
        entropy = -torch.sum(policy_probs * torch.log(policy_probs + 1e-9))
        variance = torch.var(policy_probs)
        node.entropy = entropy.item()
        node.variance = variance.item()

        print(f"Node entropy: {node.entropy}, variance: {node.variance}")  # Debug

        return value_estimate  # Return the value estimate for backpropagation


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

    def best_action_sequence(self, root_node):
        sequences = []
        self._generate_sequences(root_node, [], sequences)
        
        # Score sequences based on visit counts, entropy, and variance
        scored_sequences = []
        for seq in sequences:
            score = sum(node.visit_count for node in seq)
            entropy = sum(node.entropy for node in seq)
            variance = sum(node.variance for node in seq)
            adjusted_score = score - 0.1 * entropy + 0.05 * variance
            scored_sequences.append((seq, adjusted_score))
        
        # Sort sequences by adjusted score and select top beam_size
        best_sequences = sorted(scored_sequences, key=lambda x: x[1], reverse=True)[:self.beam_size]
        
        # Return the actions of the best sequence
        best_sequence = best_sequences[0][0]
        return [node.action for node in best_sequence[1:self.n_tokens_predict+1]]  # Exclude root node

    def _generate_sequences(self, node, current_sequence, sequences):
        current_sequence.append(node)
        if len(current_sequence) > self.n_tokens_predict or not node.children:
            sequences.append(current_sequence)
        else:
            for child in node.children.values():
                self._generate_sequences(child, current_sequence.copy(), sequences)

class State:
    def __init__(self, representation, dynamics_network, action_encoder, thought_node):
        self.representation = representation
        self.dynamics_network = dynamics_network
        self.action_encoder = action_encoder
        self.thought_node = thought_node

    def apply_action(self, action):
        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.")

        # Adjust action_index and action_embedding shapes
        action_index = torch.tensor([action_to_index[action]], device=self.representation.device)
        action_embedding = self.action_encoder(action_index)  # Shape: (batch_size=1, action_dim)

        # Extract the last time step of the state
        state = self.representation[:, -1, :]  # Shape: (batch_size, state_dim)

        # Ensure action_embedding matches the state dimension
        next_state_representation = self.dynamics_network(state, action_embedding)  # Shape: (batch_size, state_dim)

        # Append the new state to the representation history
        new_representation = torch.cat([self.representation, next_state_representation.unsqueeze(1)], dim=1)  # Shape: (batch_size, seq_len+1, state_dim)

        return State(
            representation=new_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


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



def infer(query, world_model_components, root_thought_node, tokenizer, max_length=20, inference_mode='world_model', beam_size=5, n_tokens_predict=3, mcts_iterations=10, exploration_constant=1.414):


    """
    Perform inference given a query, utilizing the Tree of Thought and MCTS with multi-token beam search.

    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')
        beam_size (int): Size of the beam for beam search
        n_tokens_predict (int): Number of tokens to predict at each step

    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 beam search
        with torch.no_grad():
            generated_sequences = model_transformer.generate_with_beam_search(
                src=input_ids,
                tokenizer=tokenizer,
                beam_size=beam_size,
                max_length=max_length,
                n_tokens_predict=n_tokens_predict,
                temperature=args.temperature
            )
        best_sequence, best_score = generated_sequences[0]
        generated_text = tokenizer.decode(best_sequence[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_representation = initial_representation[:, -1, :].unsqueeze(1)  # Shape: (batch_size=1, 1, 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 and multi-token beam search
                mcts = MCTS(prediction_network, dynamics_network, action_encoder, num_iterations=mcts_iterations, exploration_constant=exploration_constant)

                current_state = initial_state
                thought_sequence = []

                for _ in range(max_length // n_tokens_predict):
                    best_actions = mcts.search_with_beam(current_state)

                    thought_sequence.extend(best_actions)

                    # Apply the best actions to get the next state
                    for action in best_actions:
                        current_state = current_state.apply_action(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, but with multi-token beam search
                beam = [(initial_state, 0.0, torch.zeros(1, device=device), torch.zeros(1, device=device))]  # (state, score, cum_entropy, cum_variance)
                
                for _ in range(max_length // n_tokens_predict):
                    all_candidates = []
                    for state, score, cum_entropy, cum_variance in beam:
                        policy_logits, _ = prediction_network(state.representation)
                        probs = F.softmax(policy_logits / args.temperature, dim=-1)
                        entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)
                        variance = torch.var(probs, dim=-1)
                        
                        topk_probs, topk_indices = torch.topk(probs, k=beam_size, dim=-1)
                        
                        for i in range(beam_size ** n_tokens_predict):
                            indices = [i // (beam_size ** j) % beam_size for j in range(n_tokens_predict)]
                            new_actions = [index_to_action[topk_indices[0, j, indices[j]].item()] for j in range(n_tokens_predict)]
                            new_score = score + torch.sum(torch.log(topk_probs[0, range(n_tokens_predict), indices]))
                            new_entropy = cum_entropy + torch.sum(entropy[0, indices])
                            new_variance = cum_variance + torch.sum(variance[0, indices])
                            
                            new_state = state
                            for action in new_actions:
                                new_state = new_state.apply_action(action)
                            
                            all_candidates.append((new_state, new_score, new_entropy, new_variance, new_actions))
                    
                    # Select top beam_size candidates
                    beam = sorted(all_candidates, key=lambda x: x[1] - 0.1 * x[2] + 0.05 * x[3], reverse=True)[:beam_size]
                    
                    # Accumulate actions
                    if not thought_sequence:
                        thought_sequence = [b[4] for b in beam]
                    else:
                        for i, b in enumerate(beam):
                            thought_sequence[i].extend(b[4])
                
                # Return the top sequence
                return thought_sequence[0]


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)

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

            # Get action embeddings
            action_embeddings = action_encoder(action_sequences)

            # Apply dynamics network
            predicted_next_state_batch = dynamics_network(state_representation, action_embeddings)

            # 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 = F.one_hot(true_actions, num_classes=input_dim).float()
            true_value = torch.zeros_like(value_estimates).to(device)

            # Compute individual losses
            ppo_loss = ppo_agent.compute_loss(
                state_representation,
                torch.zeros_like(true_actions, dtype=torch.float32).to(device),
                true_actions,
                torch.zeros_like(value_estimates, dtype=torch.float32).to(device),
                torch.zeros_like(value_estimates, dtype=torch.float32).to(device)
            )

            info_nce = InfoNCE_Loss()(
                state_representation.view(-1, state_dim),
                F.dropout(state_representation.view(-1, state_dim), p=0.1, training=True)
            )

            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(true_actions.size(0)).to(device)
            etv = ExpectedThoughtValueLoss()(mcts_best_values)
            
            visit_counts = torch.ones(true_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 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 individual losses and total loss for this batch
        print(f"Batch {i+1} completed. Losses:")
        print(f"  PPO Loss: {ppo_loss.item():.4f}")
        print(f"  InfoNCE Loss: {info_nce.item():.4f}")
        print(f"  Covariance Loss: {covariance.item():.4f}")
        print(f"  Dynamics Loss: {dynamics_loss.item():.4f}")
        print(f"  Thought Consistency Loss: {thought_loss.item():.4f}")
        print(f"  Policy-Value Loss: {pv_loss.item():.4f}")
        print(f"  Action Diversity Loss: {action_diversity.item():.4f}")
        print(f"  Expected Thought Value Loss: {etv.item():.4f}")
        print(f"  Exploration Loss: {exploration.item():.4f}")
        print(f"  KL Divergence Loss: {kl_loss.item():.4f}")
        print(f"  Total 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)

    print(f"Current mode: {args.mode}")
    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':
        print("Entering inference mode...")
        # Build Tree of Thought if needed
        print("Building Tree of Thought...")
        tree_root = build_tree_of_thought()
        print("Tree of Thought built successfully.")

        # Generate action list
        print("Generating action list...")
        action_list = []
        traverse_tree(tree_root, action_list)
        print(f"Action list generated. Total actions: {len(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)
        print(f"Action mappings created. Vocabulary size: {action_vocab_size}")

        # Initialize or load models based on the load_model argument
        if args.load_model:
            print(f"Loading saved model from {args.load_model}")
            # Load the saved models
            model_transformer.load_state_dict(torch.load(os.path.join(args.load_model, 'transformer_model.pt')))
            representation_network.load_state_dict(torch.load(os.path.join(args.load_model, 'representation_network.pt')))
            dynamics_network.load_state_dict(torch.load(os.path.join(args.load_model, 'dynamics_network.pt')))
            
            # Load prediction network and adjust its size if necessary
            saved_state_dict = torch.load(os.path.join(args.load_model, 'prediction_network.pt'))
            saved_vocab_size = saved_state_dict['policy_head.weight'].size(0)
            if saved_vocab_size != action_vocab_size:
                print(f"Adjusting prediction network size from {saved_vocab_size} to {action_vocab_size}")
                prediction_network = PredictionNetwork(state_dim, saved_vocab_size, 1).to(device)
                prediction_network.load_state_dict(saved_state_dict)
                prediction_network.policy_head = nn.Linear(prediction_network.state_dim, action_vocab_size).to(device)
            else:
                prediction_network = PredictionNetwork(state_dim, action_vocab_size, 1).to(device)
                prediction_network.load_state_dict(saved_state_dict)

            action_encoder.load_state_dict(torch.load(os.path.join(args.load_model, 'action_encoder.pt')))
        else:
            print("Using newly initialized models")

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

        print("Starting inference loop...")
        while True:
            if args.query:
                query = args.query
                args.query = None  # Reset query for next iteration
            else:
                query = input("Please enter your query (or type 'exit' to quit): ")
                if query.lower() == 'exit':
                    break

            print(f"Processing query: {query}")
            result = infer(query, world_model_components, tree_root, tokenizer, 
                          max_length=args.max_length,
                          inference_mode=args.inference_mode, 
                          beam_size=args.beam_size, 
                          n_tokens_predict=args.n_tokens_predict,
                          mcts_iterations=args.mcts_iterations,
                          exploration_constant=args.mcts_exploration_constant)


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

            print("\n")  # Add a newline for better readability between queries

        print("Inference completed.")
    
    else:
        print(f"Invalid mode: {args.mode}. Please choose 'train' or 'inference'.")

if __name__ == '__main__':
    main()