import torch
from model.brain_architecture import TinyTransformer
from model.brain_state import model, tokenizer, DEVICE

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)


def get_optimizer():
    global optimizer
    return optimizer


def expand_model_if_needed():
    global model, optimizer

    current_vocab_size = len(tokenizer.vocab) + 10  # Buffer
    old_vocab_size = model.head.out_features

    if current_vocab_size <= old_vocab_size:
        return

    print(f"Expanding model from {old_vocab_size} -> {current_vocab_size}")

    old_state = model.state_dict()
    new_model = TinyTransformer(vocab_size=current_vocab_size).to(DEVICE)

    # Transfer matching parameters
    with torch.no_grad():
        for name, param in new_model.named_parameters():
            if name in old_state and old_state[name].shape == param.shape:
                param.copy_(old_state[name])

    model = new_model
    opt =  get_optimizer()

    print("Model expanded and optimizer rebuilt.")