```
# Imports
import os
import json
import numpy as np
import torch
import GPyOpt
import time
import math
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F
from sklearn.cluster import KMeans

# Define hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_clusters = 4
total_experts = 8
aux_loss_weight = 0.0001

GLOBAL_ITER = 1  # Declare at the top of the script

# Set random seed for reproducibility
torch.manual_seed(1337)

# Load data and define data processing functions
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Create a character-level vocabulary
chars = sorted(list(set(text)))
vocab_size = len(chars)

# Create mappings from characters to indices and vice versa
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

# Encode and decode functions to convert between text and indices
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# Convert entire text into a tensor of indices
data = torch.tensor(encode(text), dtype=torch.long)

# Split the data into training and validation sets
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

# Function to generate a batch of data for training or validation
def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

# Function to estimate the model's loss on training and validation data
@torch.no_grad()
def estimate_loss(model):
    out = {}
    model.eval()  # Set model to evaluation mode
    for split in ['train', 'val']:
        total_loss = 0.0
        total_ce_loss = 0.0
        total_aux_loss = 0.0
        for k in range(eval_iters):
            X, Y = get_batch(split)
            # Unpack all outputs from the model
            logits, loss, ce_loss, aux_loss_value = model(X, Y)
            total_loss += loss.item()
            total_ce_loss += ce_loss.item()
            total_aux_loss += aux_loss_value.item()
        avg_loss = total_loss / eval_iters
        avg_ce_loss = total_ce_loss / eval_iters
        avg_aux_loss = total_aux_loss / eval_iters
        out[split] = {
            'avg_loss': avg_loss,
            'avg_ce_loss': avg_ce_loss,
            'avg_aux_loss': aux_loss_weight * avg_aux_loss
        }
    model.train()  # Set model back to training mode
    return out

def compute_metrics(model):
    model.eval()
    total_ce_loss = 0.0
    correct_predictions = 0
    total_tokens = 0

    with torch.no_grad():
        for _ in range(eval_iters):
            inputs, targets = get_batch('val')
            inputs = inputs.to(device)
            targets = targets.to(device)

            logits, _, ce_loss, _ = model(inputs, targets)
            total_ce_loss += ce_loss.item()

            predictions = torch.argmax(logits, dim=-1)
            correct_predictions += (predictions == targets).sum().item()
            total_tokens += targets.numel()

    avg_ce_loss = total_ce_loss / eval_iters
    try:
        perplexity = math.exp(avg_ce_loss)
    except OverflowError:
        perplexity = float('inf')  # Handle overflow by setting perplexity to infinity
    accuracy = correct_predictions / total_tokens
    return {'perplexity': perplexity, 'accuracy': accuracy}

class LiquidExpert(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, tau=0.1):
        super(LiquidExpert, self).__init__()
        self.hidden_dim = hidden_dim
        self.tau = tau

        # Define weights
        self.W = nn.Parameter(torch.randn(hidden_dim, hidden_dim) * 0.1)
        self.U = nn.Parameter(torch.randn(hidden_dim, input_dim) * 0.1)
        self.b = nn.Parameter(torch.zeros(hidden_dim))

        # Output layer
        self.output_layer = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        B, T, _ = x.shape
        h = torch.zeros(B, self.hidden_dim, device=x.device)  # Initial hidden state

        for t in range(T):  # Process each timestep
            u_t = x[:, t, :]  # Input at timestep t
            h_dot = -h + torch.tanh(self.W @ h.T + self.U @ u_t.T + self.b.unsqueeze(1))
            h = h + self.tau * h_dot.T

        output = self.output_layer(h)  # Final output
        return output

class ClusteredSwitchLiquid(nn.Module):
    """Clustered Mixture-of-Experts (MoE) with teams of specialized experts."""
    def __init__(self, n_clusters, total_experts, n_embd, dropout):
        super().__init__()
        self.n_clusters = n_clusters
        self.total_experts = total_experts
        self.experts = nn.ModuleList([LiquidExpert(n_embd, dropout) for _ in range(self.total_experts)])
        self.expert_distributions = nn.Parameter(torch.ones(n_clusters, total_experts) / total_experts)

    def forward(self, x):
        B, T, D = x.shape
        if not hasattr(self, 'cluster_assignments'):
            cluster_assignments = KMeans(n_clusters=self.n_clusters).fit_predict(x.view(-1, D).detach().cpu().numpy())
            self.cluster_assignments = torch.tensor(cluster_assignments, device=x.device).view(B, T)

        outputs_per_cluster = []
        for cluster_id in range(self.n_clusters):
            assigned_indices = (self.cluster_assignments == cluster_id).nonzero(as_tuple=True)[0]
            if len(assigned_indices) == 0:
                continue
            
            cluster_inputs = x[assigned_indices]
            num_experts = torch.round(self.expert_distributions[cluster_id] * self.total_experts).long().tolist()
            active_experts = [self.experts[i] for i in range(num_experts)]

            even_split_inputs = torch.chunk(cluster_inputs, len(active_experts), dim=0)
            cluster_outputs = []
            for idx, expert in enumerate(active_experts):
                expert_out = expert(even_split_inputs[idx])
                cluster_outputs.append(expert_out)

            outputs_per_cluster.append(sum(cluster_outputs))

        if outputs_per_cluster:
            final_output = torch.cat(outputs_per_cluster)
        else:
            final_output = torch.zeros(B, T, D, device=x.device)

        return final_output

class LiquidHead(nn.Module):
    """One head of self-attention with liquid dynamics."""
    def __init__(self, head_size, dropout):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

        # Liquid neural dynamics parameters
        self.alpha_k = nn.Parameter(torch.randn(head_size))
        self.beta_k = nn.Parameter(torch.randn(head_size))
        self.alpha_q = nn.Parameter(torch.randn(head_size))
        self.beta_q = nn.Parameter(torch.randn(head_size))

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)

        k = self.alpha_k * k + self.beta_k * torch.tanh(k)
        q = self.alpha_q * q + self.beta_q * torch.tanh(q)

        wei = q @ k.transpose(-2, -1) * C**-0.5 
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)

        v = self.value(x)
        out = wei @ v
        return out

class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""
    def __init__(self, num_heads, head_size, dropout):
        super().__init__()
        self.heads = nn.ModuleList([LiquidHead(head_size, dropout) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

class Block_ClusteredSwitch(nn.Module):
    """Transformer block with Clustered Mixture-of-Experts FeedForward layer."""
    def __init__(self, n_embd, n_head, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        self.sa = MultiHeadAttention(n_head, (n_embd // n_head), dropout)
        self.ffwd = ClusteredSwitchLiquid(n_clusters, total_experts, n_embd, dropout)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = self.ffwd(self.ln2(x))
        return x

class Clustered_ST(nn.Module):
    """Clustered Switch Transformer with Dynamic Mixture-of-Experts."""
    def __init__(self, n_embd, n_head, n_layer, dropout):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Block_ClusteredSwitch(n_embd, n_head, dropout) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))
        x = tok_emb + pos_emb

        for block in self.blocks:
            x = block(x)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
            ce_loss = None
            aux_loss_value = None
        else:
            B, T, C = logits.shape
            logits_flat = logits.view(B * T, C)
            targets_flat = targets.view(B * T)
            ce_loss = F.cross_entropy(logits_flat, targets_flat)
            aux_loss_value = 0  
            loss = ce_loss + aux_loss_weight * aux_loss_value  

        return logits, loss, ce_loss, aux_loss_value

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            outputs = self(idx_cond)
            logits = outputs[0]
            logits = logits[:, -1, :]  
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)  
        return idx

def objective_function(hyperparameters):
    global GLOBAL_ITER  
    results = []

    for idx, params in enumerate(hyperparameters):
        learning_rate, n_embd, n_head, n_layer, weight_decay, dropout = params[:6]
        expert_distribution = params[6:]

        # Enforcing that expert distributions sum to total_experts
        if any(x < 1 for x in expert_distribution) or sum(expert_distribution) != total_experts:
            print(f"Invalid expert distribution for params {params}: {expert_distribution}") 
            results.append(float('inf'))  
            continue

        # Initialize model
        model = Clustered_ST(
            n_embd=int(n_embd),
            n_head=int(n_head),
            n_layer=int(n_layer),
            dropout=dropout
        ).to(device)

        optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

        best_acc = 0
        best_step = 0

        # Train the model
        for step in range(max_iters + 1):
            if step % eval_interval == 0:
                metrics = compute_metrics(model)
                print(f"Step {step} metrics: {metrics}")  
                if metrics['accuracy'] > best_acc:
                    best_acc = metrics['accuracy']
                    best_step = step

        if best_acc > 0: 
            results.append(1 - best_acc)  
            print(f"BO Iteration {GLOBAL_ITER}: Best accuracy {best_acc:.4f} at step {best_step}/{max_iters}")
        else:
            results.append(float('inf'))  

        GLOBAL_ITER += 1

    return np.array(results).reshape(-1, 1)

def save_best_hyperparameters(best_params, filepath="best_hyperparams.json"):
    with open(filepath, "w") as f:
        json.dump(best_params, f)

def load_best_hyperparameters(filepath="best_hyperparams.json"):
    try:
        with open(filepath, "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return None

def episodic_bayesian_optimization(bo_iterations=10, bounds=None, prev_best_params=None, save_path="best_hyperparams.json"):
    if bounds is None:
        bounds = [
            {'name': 'learning_rate', 'type': 'continuous', 'domain': (1e-5, 1e-2)},
            {'name': 'n_embd', 'type': 'discrete', 'domain': [128, 256, 512, 1024]},
            {'name': 'n_head', 'type': 'discrete', 'domain': [2, 4, 8, 16]},
            {'name': 'n_layer', 'type': 'discrete', 'domain': (2, 4, 8, 12)},
            {'name': 'weight_decay', 'type': 'continuous', 'domain': (1e-6, 1e-2)},
            {'name': 'dropout', 'type': 'continuous', 'domain': (0.1, 0.5)},
        ]
        bounds.extend(
            {'name': f'expert_cluster_{i}', 'type': 'discrete', 'domain': (1, total_experts)}
            for i in range(n_clusters)
        )

    print(f"Starting Bayesian Optimization for {bo_iterations} iterations...\n")

    bo = GPyOpt.methods.BayesianOptimization(
        f=objective_function,
        domain=bounds,
        acquisition_type='EI'  
    )

    bo.run_optimization(max_iter=bo_iterations)

    best_hyperparams = {
        'learning_rate': float(bo.x_opt[0]),
        'n_embd': int(bo.x_opt[1]),
        'n_head': int(bo.x_opt[2]),
        'n_layer': int(bo.x_opt[3]),
        'weight_decay': float(bo.x_opt[4]),
        'dropout': float(bo.x_opt[5]),
    }

    for i in range(n_clusters):
        best_hyperparams[f'expert_cluster_{i+1}'] = int(bo.x_opt[6 + i])

    # Ensure the expert distribution sums correctly
    expert_distribution_check = [best_hyperparams[f'expert_cluster_{i+1}'] for i in range(n_clusters)]
    if sum(expert_distribution_check) != total_experts:
        print(f"Invalid expert distribution from BO: {expert_distribution_check}")

    save_best_hyperparameters(best_hyperparams, filepath=save_path)
    print(f"\nBest hyperparameters: {best_hyperparams}")

    best_accuracy = 1 - bo.fx_opt if bo.fx_opt is not None else None
    print(f"Best accuracy: {best_accuracy}")

    return best_hyperparams, best_accuracy

# Define a function to append results to a log file
def log_results(episode_number, best_params, best_accuracy, log_filepath="bo_log.json"):
    entry = {
        "episode_number": episode_number,
        "best_hyperparameters": best_params,
        "accuracy": best_accuracy
    }

    if not os.path.exists(log_filepath):
        with open(log_filepath, "w") as f:
            json.dump([], f)

    with open(log_filepath, "r") as f:
        data = json.load(f)

    data.append(entry)

    with open(log_filepath, "w") as f:
        json.dump(data, f, indent=2)

# Run the Bayesian Optimization (current episode: 1)
best_params, best_accuracy = episodic_bayesian_optimization(bo_iterations=3, save_path="best_params.json")
log_results(1, best_params, best_accuracy, log_filepath="results_log.json")
```
 Starting Bayesian Optimization for 3 iterations...

Invalid expert distribution for params [7.47613533e-03 5.12000000e+02 8.00000000e+00 4.00000000e+00
 5.99895604e-03 3.00868867e-01 1.00000000e+00 8.00000000e+00
 1.00000000e+00 8.00000000e+00]: [1. 8. 1. 8.]
Invalid expert distribution for params [8.17146282e-03 1.02400000e+03 1.60000000e+01 8.00000000e+00
 3.83824920e-03 1.83144628e-01 8.00000000e+00 8.00000000e+00
 8.00000000e+00 1.00000000e+00]: [8. 8. 8. 1.]
Invalid expert distribution for params [5.57097234e-03 2.56000000e+02 1.60000000e+01 1.20000000e+01
 7.55903516e-03 1.40731956e-01 1.00000000e+00 1.00000000e+00
 8.00000000e+00 1.00000000e+00]: [1. 1. 8. 1.]
Invalid expert distribution for params [6.42297162e-03 5.12000000e+02 4.00000000e+00 8.00000000e+00
 1.19529855e-03 3.58026453e-01 8.00000000e+00 1.00000000e+00
 1.00000000e+00 1.00000000e+00]: [8. 1. 1. 1.]
Invalid expert distribution for params [3.03098938e-03 2.56000000e+02 2.00000000e+00 8.00000000e+00
 5.05250921e-03 2.04518795e-01 1.00000000e+00 1.00000000e+00
 1.00000000e+00 1.00000000e+00]: [1. 1. 1. 1.]
Invalid expert distribution for params [5.39918266e-04 2.56000000e+02 1.60000000e+01 2.00000000e+00
 4.53087917e-03 1.60743776e-01 1.00000000e+00 1.00000000e+00
 8.00000000e+00 1.00000000e+00]: [1. 1. 8. 1.]
Invalid expert distribution for params [7.79884602e-03 1.28000000e+02 2.00000000e+00 8.00000000e+00
 4.73628294e-03 3.16611058e-01 1.00000000e+00 8.00000000e+00
 8.00000000e+00 8.00000000e+00]: [1. 8. 8. 8.]
Invalid expert distribution for params [8.87914794e-03 5.12000000e+02 8.00000000e+00 4.00000000e+00
 1.69657690e-03 1.15494026e-01 8.00000000e+00 1.00000000e+00
 1.00000000e+00 8.00000000e+00]: [8. 1. 1. 8.]
Invalid expert distribution from BO: [1, 8, 1, 8]

Best hyperparameters: {'learning_rate': 0.007476135329267061, 'n_embd': 512, 'n_head': 8, 'n_layer': 4, 'weight_decay': 0.0059989560365188615, 'dropout': 0.3008688671818145, 'expert_cluster_1': 1, 'expert_cluster_2': 8, 'expert_cluster_3': 1, 'expert_cluster_4': 8}
Best accuracy: -inf