import os
import re
import json
import math
import time
import glob
import random

from dataclasses import dataclass

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

torch.manual_seed(42)
random.seed(42)

NUM_THREADS = os.cpu_count() or 4
torch.set_num_threads(NUM_THREADS)
try:
    torch.set_num_interop_threads(max(1, NUM_THREADS // 2))
except RuntimeError:
    # The number of interop threads can only be set once at the start of the program
    pass

try:
    torch.backends.mkldnn.enabled = True  # Intel MKL-DNN acceleration (if available)
except Exception:
    pass

DEVICE = torch.device("cpu")

# Autocasting to bfloat16 on the CPU can speed up most matmul operations (if supported)
USE_BF16_AUTOCAST = True
try:
    _ = torch.zeros(1, dtype=torch.bfloat16) + torch.zeros(1, dtype=torch.bfloat16)
except Exception:
    USE_BF16_AUTOCAST = False

print(f"🧵 Number of CPU threads           : {NUM_THREADS}")
print(f"⚙️  bfloat16 autocast status    : {'active' if USE_BF16_AUTOCAST else 'inactive'}")

@dataclass
class OSW1Config:
    data_dir: str = "data"

    block_size: int = 128      
    d_model: int = 256         
    n_layer: int = 4           
    n_head: int = 4            
    d_ff: int = 1024           
    dropout: float = 0.1

    batch_size: int = 8
    grad_accum_steps: int = 2      
    epochs: int = 20
    max_lr: float = 3e-4
    min_lr: float = 3e-5
    warmup_ratio: float = 0.05
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    label_smoothing: float = 0.05

    checkpoint_prefix: str = "opensoftware_world_osw1"

TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)

def tokenize(text: str):
    return TOKEN_RE.findall(text.lower())

class Vocab:
    PAD, UNK, BOS, EOS = "<pad>", "<unk>", "<bos>", "<eos>"

    def __init__(self):
        self.stoi = {}
        self.itos = []

    def build(self, token_stream):
        specials = [Vocab.PAD, Vocab.UNK, Vocab.BOS, Vocab.EOS]
        counts = {}
        for tok in token_stream:
            counts[tok] = counts.get(tok, 0) + 1
        sorted_toks = sorted(counts.items(), key=lambda x: (-x[1], x[0]))
        self.itos = specials + [t for t, _ in sorted_toks]
        self.stoi = {t: i for i, t in enumerate(self.itos)}

    def encode(self, text, add_bos=False, add_eos=False):
        ids = [self.stoi.get(t, self.stoi[Vocab.UNK]) for t in tokenize(text)]
        if add_bos:
            ids = [self.stoi[Vocab.BOS]] + ids
        if add_eos:
            ids = ids + [self.stoi[Vocab.EOS]]
        return ids

    def decode(self, ids):
        toks = [self.itos[i] for i in ids if 0 <= i < len(self.itos)]
        toks = [t for t in toks if t != Vocab.PAD and t != Vocab.BOS]
        out = []
        for t in toks:
            if t == Vocab.EOS:
                break
            out.append(t)
        text = " ".join(out)
        text = re.sub(r"\s+([.,!?;:])", r"\1", text)
        return text

    def __len__(self):
        return len(self.itos)

def load_json_pairs(json_dir):
    pairs = []
    if not os.path.isdir(json_dir):
        return pairs
    for path in glob.glob(os.path.join(json_dir, "*.json")):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
        except Exception as e:
            print(f"⚠️  {path} could not be read: {e}")
            continue
        intents = data.get("intents", data if isinstance(data, list) else [])
        for intent in intents:
            patterns = intent.get("patterns", []) or []
            responses = intent.get("responses", []) or []
            if not patterns or not responses:
                continue
            for p in patterns:
                for r in responses:
                    pairs.append((p, r))
    return pairs


def load_txt_qa_pairs(qa_dir):
    pairs = []
    if not os.path.isdir(qa_dir):
        return pairs
    for path in glob.glob(os.path.join(qa_dir, "*.txt")):
        with open(path, "r", encoding="utf-8") as f:
            lines = [l.rstrip("\n") for l in f.readlines()]
        q, a = None, None
        for raw in lines:
            line = raw.strip()
            if line.startswith("Q:"):
                q = line[2:].strip()
            elif line.startswith("A:"):
                a = line[2:].strip()
                if q is not None and a:
                    pairs.append((q, a))
                q, a = None, None
    return pairs


def load_plain_texts(txt_dir):
    texts = []
    if not os.path.isdir(txt_dir):
        return texts
    for path in glob.glob(os.path.join(txt_dir, "*.txt")):
        with open(path, "r", encoding="utf-8") as f:
            content = f.read().strip()
            if content:
                texts.append(content)
    return texts


def build_corpus(cfg: OSW1Config, vocab: Vocab):
    json_dir = os.path.join(cfg.data_dir, "json")
    qa_dir = os.path.join(cfg.data_dir, "txt_qa")
    txt_dir = os.path.join(cfg.data_dir, "txt")

    qa_pairs = load_json_pairs(json_dir) + load_txt_qa_pairs(qa_dir)
    plain_texts = load_plain_texts(txt_dir)

    print(f"📚 JSON + txt_qa pair count   : {len(qa_pairs)}")
    print(f"📄 Plain text file count     : {len(plain_texts)}")

    if not qa_pairs and not plain_texts:
        raise RuntimeError(
            "No data found! Please populate the 'data/json', 'data/txt', 'data/txt_qa' "
            "folders with data for the model to learn from."
        )

    all_tokens = []
    for q, a in qa_pairs:
        all_tokens.extend(tokenize(q))
        all_tokens.extend(tokenize(a))
    for t in plain_texts:
        all_tokens.extend(tokenize(t))
    vocab.build(all_tokens)

    sequences = []

    for q, a in qa_pairs:
        ids = [vocab.stoi[Vocab.BOS]]
        ids += vocab.encode(q)
        ids += vocab.encode(a)
        ids += [vocab.stoi[Vocab.EOS]]
        if len(ids) >= 4:
            sequences.append(ids)

    for t in plain_texts:
        ids = [vocab.stoi[Vocab.BOS]] + vocab.encode(t) + [vocab.stoi[Vocab.EOS]]
        stride = max(1, cfg.block_size // 2)
        for i in range(0, max(1, len(ids) - 1), stride):
            chunk = ids[i:i + cfg.block_size + 1]
            if len(chunk) >= 8:
                sequences.append(chunk)

    random.shuffle(sequences)
    print(f"🧩 Total training sequences (sequence): {len(sequences)}")
    print(f"🔤 Vocab size                    : {len(vocab)}")
    return sequences

class SeqDataset(torch.utils.data.Dataset):
    def __init__(self, sequences, block_size):
        self.sequences = sequences
        self.block_size = block_size

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        ids = self.sequences[idx][: self.block_size + 1]
        return torch.tensor(ids, dtype=torch.long)


def make_collate(pad_id):
    def collate(batch):
        max_len = max(len(x) for x in batch)
        padded = torch.full((len(batch), max_len), pad_id, dtype=torch.long)
        for i, seq in enumerate(batch):
            padded[i, : len(seq)] = seq
        x = padded[:, :-1].contiguous()
        y = padded[:, 1:].contiguous()
        return x, y
    return collate

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head, dropout):
        super().__init__()
        assert d_model % n_head == 0, "d_model must be evenly divisible by n_head"
        self.n_head = n_head
        self.head_dim = d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)

    def forward(self, x, attn_mask):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        att = att.masked_fill(attn_mask, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        out = att @ v
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.proj(out))


class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_head, d_ff, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x, attn_mask):
        x = x + self.attn(self.ln1(x), attn_mask)
        x = x + self.mlp(self.ln2(x))
        return x


class OSW1Model(nn.Module):
    def __init__(self, vocab_size, cfg: OSW1Config, pad_id: int):
        super().__init__()
        self.cfg = cfg
        self.pad_id = pad_id

        self.tok_emb = nn.Embedding(vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(cfg.d_model, cfg.n_head, cfg.d_ff, cfg.dropout)
            for _ in range(cfg.n_layer)
        ])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight

        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device).unsqueeze(0)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))

        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=idx.device), diagonal=1)
        for block in self.blocks:
            x = block(x, mask)
        x = self.ln_f(x)
        logits = self.head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
                ignore_index=self.pad_id,
                label_smoothing=self.cfg.label_smoothing,
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.9, top_k=40, eos_id=None):
        was_training = self.training
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.cfg.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
            if eos_id is not None and next_id.item() == eos_id:
                break
        if was_training:
            self.train()
        return idx

def count_parameters(model: OSW1Model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    breakdown = {
        "Token + Position Embedding": model.tok_emb.weight.numel() + model.pos_emb.weight.numel(),
        f"Transformer Blocks ({len(model.blocks)} pieces)": sum(p.numel() for p in model.blocks.parameters()),
        "Final LayerNorm": sum(p.numel() for p in model.ln_f.parameters()),
        "Output Layer (shared with embedding, no extra parameters)": 0,
    }
    return total, trainable, breakdown

def human_readable_param_count(n: int):
    if n >= 1_000_000_000:
        return f"{n/1_000_000_000:.2f}B", f"{max(1, round(n/1_000_000_000))}b"
    elif n >= 1_000_000:
        return f"{n/1_000_000:.2f}M", f"{max(1, round(n/1_000_000))}m"
    elif n >= 1_000:
        return f"{n/1_000:.2f}K", f"{max(1, round(n/1_000))}k"
    else:
        return str(n), str(n)

def print_model_report(model: OSW1Model, cfg: OSW1Config, vocab_size: int):
    total, trainable, breakdown = count_parameters(model)
    pretty, short = human_readable_param_count(total)
    size_mb = total * 4 / (1024 ** 2)

    print("\n" + "=" * 64)
    print("🧠  OpenSoftware-World OSW1 — MODEL REPORT")
    print("=" * 64)
    print(f"  Vocab size             : {vocab_size:,}")
    print(f"  Context window (block)  : {cfg.block_size}")
    print(f"  Embedding size (d_model)    : {cfg.d_model}")
    print(f"  Number of layers (n_layer)   : {cfg.n_layer}")
    print(f"  Head count (n_head)      : {cfg.n_head}")
    print(f"  Feed-forward size (d_ff)       : {cfg.d_ff}")
    print("-" * 64)
    for name, count in breakdown.items():
        print(f"  {name:<50}: {count:,}")
    print("-" * 64)
    print(f"  TOTAL PARAMETER COUNT   : {total:,}   (~{pretty})")
    print(f"  TRAINABLE PARAMETERS    : {trainable:,}")
    print(f"  Estimated model size      : {size_mb:.2f} MB (float32)")
    print(f"  Checkpoint file label       : {short}  ->  {cfg.checkpoint_prefix}_{short}.pth")
    print("=" * 64 + "\n")
    return short

def lr_at_step(step, total_steps, warmup_steps, max_lr, min_lr):
    if step < warmup_steps:
        return max_lr * (step + 1) / max(1, warmup_steps)
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    progress = min(max(progress, 0.0), 1.0)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

def train(cfg: OSW1Config):
    vocab = Vocab()
    sequences = build_corpus(cfg, vocab)
    pad_id = vocab.stoi[Vocab.PAD]

    dataset = SeqDataset(sequences, cfg.block_size)
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=cfg.batch_size,
        shuffle=True,
        collate_fn=make_collate(pad_id),
        num_workers=0,      
        drop_last=True,
    )

    if len(loader) == 0:
        raise RuntimeError(
            "The dataset is too small to even create a batch.  "
            "Try reducing 'batch_size' or adding more data."
        )

    model = OSW1Model(len(vocab), cfg, pad_id=pad_id).to(DEVICE)
    compiled_model = model
    try:
        compiled_model = torch.compile(model, backend="inductor")
        print("🚀 torch.compile has been enabled (provides an extra speed boost if available).")
    except Exception as e:
        print(f"ℹ️  torch.compile could not be used, continuing in normal mode: {e}")

    size_tag = print_model_report(model, cfg, len(vocab))

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=cfg.max_lr,
        betas=(0.9, 0.95),
        weight_decay=cfg.weight_decay,
    )

    steps_per_epoch = max(1, len(loader) // cfg.grad_accum_steps)
    total_steps = steps_per_epoch * cfg.epochs
    warmup_steps = max(1, int(total_steps * cfg.warmup_ratio))

    print(f"⏱️  Total optimization steps : {total_steps} | Warmup steps: {warmup_steps}")
    print(f"🏋️  Training starting... ({cfg.epochs} epoch, batch={cfg.batch_size}, "
          f"grad_accum={cfg.grad_accum_steps})\n")

    global_step = 0
    train_start = time.time()

    for epoch in range(1, cfg.epochs + 1):
        epoch_start = time.time()
        epoch_loss, n_batches = 0.0, 0
        optimizer.zero_grad(set_to_none=True)

        for i, (x, y) in enumerate(loader):
            x, y = x.to(DEVICE), y.to(DEVICE)

            if USE_BF16_AUTOCAST:
                with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
                    _, loss = compiled_model(x, y)
            else:
                _, loss = compiled_model(x, y)

            loss_scaled = loss / cfg.grad_accum_steps
            loss_scaled.backward()

            if (i + 1) % cfg.grad_accum_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
                lr = lr_at_step(global_step, total_steps, warmup_steps, cfg.max_lr, cfg.min_lr)
                for g in optimizer.param_groups:
                    g["lr"] = lr
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1

            epoch_loss += loss.item()
            n_batches += 1

        avg_loss = epoch_loss / max(1, n_batches)
        ppl = math.exp(min(avg_loss, 20))
        epoch_time = time.time() - epoch_start
        elapsed_total = time.time() - train_start
        current_lr = optimizer.param_groups[0]["lr"]
        print(
            f"📈 Epoch {epoch:>3}/{cfg.epochs} | "
            f"loss={avg_loss:.4f} | ppl={ppl:.2f} | "
            f"lr={current_lr:.2e} | "
            f"time={epoch_time:.1f}s | total={elapsed_total/60:.1f}m"
        )

    total_time = time.time() - train_start
    print(f"\n✅ Training completed! Total time: "
          f"{total_time/60:.2f} minutes ({total_time:.1f} seconds)\n")

    ckpt_path = f"{cfg.checkpoint_prefix}_{size_tag}.pth"
    torch.save({
        "model_state_dict": model.state_dict(),
        "config": cfg.__dict__,
        "vocab_stoi": vocab.stoi,
        "vocab_itos": vocab.itos,
        "pad_id": pad_id,
        "param_count": sum(p.numel() for p in model.parameters()),
        "training_time_sec": total_time,
        "final_loss": avg_loss,
    }, ckpt_path)
    print(f"💾 Model saved: {ckpt_path}\n")

    return model, vocab, cfg, ckpt_path

def chat_loop(model: OSW1Model, vocab: Vocab, cfg: OSW1Config):
    print("=" * 64)
    print("💬 OSW1 with chat mode! Type 'exit' to quit.")
    print("=" * 64)
    model.eval()
    eos_id = vocab.stoi[Vocab.EOS]
    bos_id = vocab.stoi[Vocab.BOS]

    while True:
        try:
            user_in = input("\You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye!")
            break

        if user_in.lower() in ("exit", "quit"):
            print("👋 Goodbye!")
            break
        if not user_in:
            continue

        ids = [bos_id] + vocab.encode(user_in)
        x = torch.tensor([ids], dtype=torch.long)
        out = model.generate(x, max_new_tokens=60, temperature=0.85, top_k=40, eos_id=eos_id)
        answer_ids = out[0, len(ids):].tolist()
        answer = vocab.decode(answer_ids)
        print(f"OSW1: {answer if answer else '(...silence...)'}")

def load_checkpoint(path: str):
    ckpt = torch.load(path, map_location="cpu")
    cfg = OSW1Config(**ckpt["config"])
    vocab = Vocab()
    vocab.stoi = ckpt["vocab_stoi"]
    vocab.itos = ckpt["vocab_itos"]
    model = OSW1Model(len(vocab), cfg, pad_id=ckpt["pad_id"])
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()
    return model, vocab, cfg

def main():
    cfg = OSW1Config()
    model, vocab, cfg, ckpt_path = train(cfg)
    chat_loop(model, vocab, cfg)


if __name__ == "__main__":
    main()