"""Pepaman byte-level causal transformer. Random weights, no pretrained assets."""
import hashlib
import json
import math
import random
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F


@dataclass
class Config:
    width: int = 192
    layers: int = 4
    heads: int = 6
    context: int = 128
    dropout: float = 0.1


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.heads = c.heads
        self.norm1 = nn.LayerNorm(c.width)
        self.norm2 = nn.LayerNorm(c.width)
        self.qkv = nn.Linear(c.width, 3 * c.width)
        self.proj = nn.Linear(c.width, c.width)
        self.ff = nn.Sequential(nn.Linear(c.width, 4*c.width), nn.GELU(),
                                nn.Linear(4*c.width, c.width), nn.Dropout(c.dropout))
        self.dropout = c.dropout

    def forward(self, x):
        b, t, w = x.shape
        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)
        q, k, v = [z.view(b, t, self.heads, w//self.heads).transpose(1, 2)
                   for z in (q, k, v)]
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                  dropout_p=self.dropout if self.training else 0.0)
        x = x + self.proj(a.transpose(1, 2).contiguous().view(b, t, w))
        return x + self.ff(self.norm2(x))


class ScratchGPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.width % config.heads:
            raise ValueError('Width must be divisible by number of heads.')
        self.config = config
        self.tokens = nn.Embedding(256, config.width)
        self.positions = nn.Embedding(config.context, config.width)
        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.layers)])
        self.norm = nn.LayerNorm(config.width)
        self.output = nn.Linear(config.width, 256, bias=False)
        self.output.weight = self.tokens.weight
        self.apply(self._init)

    @staticmethod
    def _init(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(self, tokens, targets=None):
        if tokens.size(1) > self.config.context:
            raise ValueError('Sequence exceeds context.')
        positions = torch.arange(tokens.size(1), device=tokens.device)
        x = self.tokens(tokens) + self.positions(positions)
        logits = self.output(self.norm(self.blocks(x)))
        loss = None if targets is None else F.cross_entropy(
            logits.reshape(-1, 256), targets.reshape(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, prompt, length=300, temperature=0.8, seed=123):
        if temperature <= 0:
            raise ValueError('Temperature must be positive.')
        self.eval()
        device = next(self.parameters()).device
        ids = list(prompt.encode('utf-8')) or [10]
        tokens = torch.tensor([ids], device=device, dtype=torch.long)
        rng = torch.Generator(device=device).manual_seed(seed)
        for _ in range(length):
            logits, _ = self(tokens[:, -self.config.context:])
            probabilities = F.softmax(logits[:, -1] / temperature, dim=-1)
            next_token = torch.multinomial(probabilities, 1, generator=rng)
            tokens = torch.cat([tokens, next_token], dim=1)
        return bytes(tokens[0].tolist()).decode('utf-8', errors='replace')


def load_splits(folder, context):
    result = {}
    fingerprints = {}
    for name in ('train', 'validation', 'test'):
        raw = (Path(folder) / (name + '.txt')).read_bytes()
        raw.decode('utf-8')  # Reject non-text inputs.
        if len(raw) <= context:
            raise ValueError(f'{name}.txt needs more than {context} UTF-8 bytes.')
        fingerprints[name] = hashlib.sha256(raw).hexdigest()
        result[name] = torch.tensor(list(raw), dtype=torch.long)
    if len(set(fingerprints.values())) != 3:
        raise ValueError('Training, validation and test files must differ.')
    # Catch exact copied paragraphs; human review must also catch paraphrases.
    paragraphs = {}
    for name in result:
        text = (Path(folder) / (name + '.txt')).read_text(encoding='utf-8')
        paragraphs[name] = {' '.join(p.lower().split()) for p in text.split('\n\n')
                            if len(p.strip()) >= 80}
    for a, b in [('train', 'validation'), ('train', 'test'), ('validation', 'test')]:
        if paragraphs[a] & paragraphs[b]:
            raise ValueError(f'Copied paragraphs overlap between {a} and {b}.')
    return result, fingerprints


def batch(data, context, size, device, generator=None):
    starts = torch.randint(len(data)-context, (size,), generator=generator)
    x = torch.stack([data[i:i+context] for i in starts])
    y = torch.stack([data[i+1:i+context+1] for i in starts])
    return x.to(device), y.to(device)


@torch.no_grad()
def evaluate(model, data, batches=8, size=4):
    was_training = model.training
    model.eval()
    rng = torch.Generator().manual_seed(987)
    losses = []
    for _ in range(batches):
        x, y = batch(data, model.config.context, size, next(model.parameters()).device, rng)
        losses.append(float(model(x, y)[1]))
    model.train(was_training)
    return sum(losses) / len(losses)


def train(data_dir, output_dir, config=None, steps=500, batch_size=8,
          accumulation=2, learning_rate=3e-4, save_every=50, resume=False,
          allow_cpu=False, seed=42):
    """steps is total target step, also when resuming. No network or paid APIs."""
    if min(steps, batch_size, accumulation, save_every) < 1:
        raise ValueError('Steps, batch size, accumulation and save interval must be positive.')
    config = config or Config()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if device == 'cpu' and not allow_cpu:
        raise RuntimeError('Select a GPU runtime, or explicitly set allow_cpu=True for a small test.')
    random.seed(seed)
    torch.manual_seed(seed)
    if device == 'cpu':
        torch.set_num_threads(min(4, torch.get_num_threads()))
    data, fingerprints = load_splits(data_dir, config.context)
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    checkpoint = output / 'latest.pt'
    if checkpoint.exists() and not resume:
        raise ValueError('An existing run is present. Resume it or choose a different output folder.')
    model = ScratchGPT(config).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.1)
    start = 0
    history = []
    settings = dict(batch_size=batch_size, accumulation=accumulation, learning_rate=learning_rate)
    if resume:
        if not checkpoint.exists():
            raise FileNotFoundError('Restore latest.pt into the output folder before resuming.')
        saved = torch.load(checkpoint, map_location=device, weights_only=True)
        if saved['config'] != asdict(config) or saved['data_sha256'] != fingerprints:
            raise ValueError('Resume requires the same architecture and exact dataset files.')
        if saved['training_settings'] != settings:
            raise ValueError('Resume requires the same batch, accumulation and learning rate.')
        model.load_state_dict(saved['model'])
        optimizer.load_state_dict(saved['optimizer'])
        torch.set_rng_state(saved['rng'].cpu())
        if device == 'cuda' and saved.get('cuda_rng'):
            if len(saved['cuda_rng']) == torch.cuda.device_count():
                torch.cuda.set_rng_state_all([r.cpu() for r in saved['cuda_rng']])
        start, history = saved['step'], saved['history']
    params = sum(p.numel() for p in model.parameters())
    print(f'Device: {device}; parameters: {params:,}; completed steps: {start}')
    if len(data['train']) < 1_000_000:
        print('DEMO CORPUS: too small for a useful language or science model. Expect memorization.')
    model.train()
    started = time.time()

    def save(step):
        state = dict(model=model.state_dict(), optimizer=optimizer.state_dict(),
                     config=asdict(config), data_sha256=fingerprints,
                     step=step, history=history, training_settings=settings,
                     rng=torch.get_rng_state(),
                     cuda_rng=torch.cuda.get_rng_state_all() if device == 'cuda' else [])
        temporary = output / 'latest.tmp'
        torch.save(state, temporary)
        temporary.replace(checkpoint)

    for step in range(start+1, steps+1):
        optimizer.zero_grad(set_to_none=True)
        total = 0.0
        for _ in range(accumulation):
            x, y = batch(data['train'], config.context, batch_size, device)
            _, loss = model(x, y)
            if not torch.isfinite(loss):
                raise RuntimeError('Nonfinite loss. Restore the previous checkpoint and inspect data/settings.')
            (loss / accumulation).backward()
            total += float(loss.detach()) / accumulation
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        if step % save_every == 0 or step == steps:
            val_loss = evaluate(model, data['validation'])
            entry = dict(step=step, train_loss=total, validation_loss=val_loss)
            history.append(entry)
            save(step)
            print(json.dumps(entry))

    # Test is deliberately separate, run once after selecting the final model.
    report = dict(parameters=params, device=device, completed_steps=max(start, steps),
                  elapsed_seconds=round(time.time()-started, 2), history=history,
                  data_sha256=fingerprints, config=asdict(config),
                  warning='Loss is a text-prediction metric, not a factual-accuracy or truth score.')
    (output / 'metrics.json').write_text(json.dumps(report, indent=2), encoding='utf-8')
    torch.save(dict(model=model.state_dict(), config=asdict(config)), output / 'model.pt')
    return model, report


def load_model(path, device='cpu'):
    saved = torch.load(path, map_location=device, weights_only=True)
    model = ScratchGPT(Config(**saved['config'])).to(device)
    model.load_state_dict(saved['model'])
    model.eval()
    return model