import torch
import time
from model.dynamic_expand import expand_model_if_needed, _last_expansion_time, get_optimizer, expand_lock
from model.brain_state import model, tokenizer, DEVICE, loss_fn
from model.brainmap import update_brainmap
from context.context import add_to_context, get_recent_context

LOSS_FILE = "data/logs/loss.log"
VOCAB_GROWTH_FILE = "data/logs/vocab_growth.log"
scheduler = torch.optim.lr_scheduler.StepLR(get_optimizer(), step_size=500, gamma=0.95)


def log_vocab_growth():
    with open(VOCAB_GROWTH_FILE, "a", encoding="utf-8") as f:
        f.write(f"{time.time()},{len(tokenizer.vocab)}\n")


def log_loss(value: float):
    with open(LOSS_FILE, "a", encoding="utf-8") as f:
        f.write(f"{time.time()},{round(value, 4)}\n")


def train_on_message(text: str, source: str = "user"):
    expand_model_if_needed()

    now = time.time()
    if now - _last_expansion_time < 5:
        print("[Train] Skipping to stabilize after expansion.")
        return

    if not expand_lock.acquire(timeout=0.5):
        print("[Train] Skipped training due to active expansion.")
        return

    try:
        model.train()

        context_texts = get_recent_context(10)
        augmented_text = "<start> " + " ".join(context_texts + [text]) + " <end>"

        tokens = tokenizer.tokenize(augmented_text)

        if len(tokens) < 2:
            return

        max_token_id = model.head.out_features - 1
        tokens = [t if t <= max_token_id else max_token_id for t in tokens]
        tokens = tokens[:128]

        if len(tokens) < 2:
            return

        input_tensor = torch.tensor(tokens[:-1], dtype=torch.long, device=DEVICE).unsqueeze(0)
        target_tensor = torch.tensor(tokens[1:], dtype=torch.long, device=DEVICE).unsqueeze(0)

        opt = get_optimizer()

        output = model(input_tensor)
        loss = loss_fn(output.view(-1, output.size(-1)), target_tensor.view(-1))
        if torch.isnan(loss):
            print("[Trainer] Detected NaN loss, skipping update.")
            return

        opt.zero_grad()
        loss.backward()
        opt.step()
        scheduler.step()

        log_loss(loss.item())
        log_vocab_growth()
        add_to_context(text, source=source)
        update_brainmap(augmented_text.split())

    finally:
        expand_lock.release()