import math
import time
import torch
import torch.nn.functional as F
from transformers import GPT2LMHeadModel, GPT2TokenizerFast
from config import BATCH_SIZE, LEARNING_RATE, STEPS

REPLY_PREFIX = "RAM-ALTMAN:"
HOOK_TOKENS = 8
HOOK_WEIGHT = 4.0
LOG_EVERY = 250


def read_segments(data_file: str) -> tuple[list[str], list[bool]]:
    texts, is_reply = [], []

    with open(data_file, encoding="utf-8") as f:
        for line in f:
            if line.startswith(REPLY_PREFIX):
                texts += [REPLY_PREFIX, line.removeprefix(REPLY_PREFIX)]
                is_reply += [False, True]
            else:
                texts.append(line)
                is_reply.append(False)

    return texts, is_reply


def reply_weights(length: int) -> list[float]:
    # The hook weighs more than the rant
    if length <= HOOK_TOKENS:
        return [1.0] * length
    return [HOOK_WEIGHT] * HOOK_TOKENS + [1.0] * (length - HOOK_TOKENS)


def load_token_ids(hf_tokenizer: GPT2TokenizerFast, data_file: str, device: str) -> tuple[torch.Tensor, torch.Tensor]:
    texts, is_reply = read_segments(data_file)
    ids, weights = [], []

    for segment_ids, reply in zip(hf_tokenizer(texts)["input_ids"], is_reply):
        ids += segment_ids
        weights += reply_weights(len(segment_ids)) if reply else [0.0] * len(segment_ids)

    return torch.tensor(ids, device=device), torch.tensor(weights, device=device)


def cosine_lr(step: int) -> float:
    return LEARNING_RATE * 0.5 * (1 + math.cos(math.pi * step / STEPS))


def random_batch(token_ids: torch.Tensor, weights: torch.Tensor, context_length: int) -> tuple[torch.Tensor, torch.Tensor]:
    starts = torch.randint(0, len(token_ids) - context_length - 1, (BATCH_SIZE,))
    inputs = torch.stack([token_ids[i:i + context_length] for i in starts])
    input_weights = torch.stack([weights[i:i + context_length] for i in starts])
    return inputs, input_weights


def weighted_loss(logits: torch.Tensor, inputs: torch.Tensor, input_weights: torch.Tensor) -> torch.Tensor:
    predictions = logits[:, :-1].reshape(-1, logits.size(-1))
    targets = inputs[:, 1:].reshape(-1)
    target_weights = input_weights[:, 1:].reshape(-1)

    token_losses = F.cross_entropy(predictions, targets, reduction="none")
    return (token_losses * target_weights).sum() / target_weights.sum().clamp_min(1e-8)


def train_loop(model: GPT2LMHeadModel, token_ids: torch.Tensor, weights: torch.Tensor, context_length: int) -> None:
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01)
    start = time.time()

    for step in range(1, STEPS + 1):
        for group in optimizer.param_groups:
            group["lr"] = cosine_lr(step)

        inputs, input_weights = random_batch(token_ids, weights, context_length)
        loss = weighted_loss(model(input_ids=inputs).logits, inputs, input_weights)

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        if step == 1 or step % LOG_EVERY == 0:
            print(f"step {step}/{STEPS} - loss {loss.item():.3f} <> {time.time() - start:.0f}s", flush=True)