import json, os, pickle, math, time, sys
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors
from datasets import load_dataset
import numpy as np
import warnings
warnings.filterwarnings("ignore")

os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "data")
TOKENIZER_DIR = os.path.join(BASE_DIR, "tokenizer")
MODEL_DIR = os.path.join(BASE_DIR, "model")
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(TOKENIZER_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)

torch.set_num_threads(4)

VOCAB_SIZE = 50000
HIDDEN_SIZE = 1536
NUM_LAYERS = 30
NUM_HEADS = 12
HEAD_DIM = HIDDEN_SIZE // NUM_HEADS
INTERMEDIATE_SIZE = 6144
MAX_SEQ_LEN = 128
NUM_SAMPLES = 10000
TRAIN_BATCH_SIZE = 2
GRAD_ACCUM_STEPS = 4
LEARNING_RATE = 4e-4
NUM_EPOCHS = 3
WARMUP_STEPS = 50

total_p = (VOCAB_SIZE * HIDDEN_SIZE +
           NUM_LAYERS * (4 * HIDDEN_SIZE * HIDDEN_SIZE + 3 * HIDDEN_SIZE * INTERMEDIATE_SIZE + 2 * HIDDEN_SIZE) +
           HIDDEN_SIZE * VOCAB_SIZE)
print(f"=== Sage 1B ({total_p/1e9:.3f}B params) ===")

# ====== STEP 1: Load English Dataset ======
print("\n--- Step 1: Loading English text dataset ---")
dataset = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
samples = []
start = time.time()
for i, example in enumerate(dataset):
    if i >= NUM_SAMPLES:
        break
    text = example.get("text", "").strip()
    if len(text) >= 100:
        samples.append(text)
    if (i+1) % 10000 == 0:
        print(f"  {i+1}/{NUM_SAMPLES} scanned, {len(samples)} valid ({time.time()-start:.0f}s)")

# Supplement with more if needed
if len(samples) < 10000:
    print(f"  Only {len(samples)} valid samples. Trying additional sources...")
    try:
        ds2 = load_dataset("wikipedia", "20220301.en", split="train", streaming=True)
        for i, ex in enumerate(ds2):
            if len(samples) >= NUM_SAMPLES:
                break
            text = ex.get("text", "").strip()
            if len(text) >= 200:
                samples.append(text[:2000])
            if (i+1) % 5000 == 0:
                print(f"  wiki: {i+1} scanned, {len(samples)} total")
    except Exception as e:
        print(f"  Wikipedia supplement failed: {e}")

print(f"Collected {len(samples)} samples in {time.time()-start:.0f}s")
with open(os.path.join(DATA_DIR, "raw_texts.pkl"), "wb") as f:
    pickle.dump(samples, f)

# ====== STEP 2: Train BPE Tokenizer ======
print("\n--- Step 2: Training BPE tokenizer ---")
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder = decoders.ByteLevel()
tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
trainer = trainers.BpeTrainer(
    vocab_size=VOCAB_SIZE,
    special_tokens=["<PAD>", "<UNK>", "<BOS>", "<EOS>"],
    min_frequency=2,
)
tokenizer.train_from_iterator(samples, trainer=trainer)
tokenizer.save(os.path.join(TOKENIZER_DIR, "tokenizer.json"))
print(f"Vocabulary size: {tokenizer.get_vocab_size()}")

# ====== STEP 3: Tokenize ======
print("\n--- Step 3: Tokenizing ---")
pad_id = tokenizer.token_to_id("<PAD>")
bos_id = tokenizer.token_to_id("<BOS>")
eos_id = tokenizer.token_to_id("<EOS>")
tokenized = []
for text in samples:
    ids = tokenizer.encode(text).ids
    if len(ids) > MAX_SEQ_LEN - 2:
        ids = ids[:MAX_SEQ_LEN - 2]
    ids = [bos_id] + ids + [eos_id]
    if len(ids) < MAX_SEQ_LEN:
        ids += [pad_id] * (MAX_SEQ_LEN - len(ids))
    tokenized.append(ids)
tensor_data = torch.tensor(tokenized, dtype=torch.long)
torch.save(tensor_data, os.path.join(DATA_DIR, "tokenized.pt"))
print(f"Tokenized {len(tokenized)} sequences, shape: {tensor_data.shape}")

# ====== STEP 4: Build Model ======
print("\n--- Step 4: Building Sage 1B model ---")

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=MAX_SEQ_LEN):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len = max_seq_len
        self._cos = None
        self._sin = None
    def get_cos_sin(self, x, seq_len=None):
        seq_len = seq_len or x.size(1)
        if self._cos is None or self._cos.size(-2) < seq_len:
            t = torch.arange(self.max_seq_len, device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)[None, None]
            self._cos = emb.cos()
            self._sin = emb.sin()
        return self._cos[..., :seq_len, :], self._sin[..., :seq_len, :]

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary(x, cos, sin):
    return (x * cos) + (rotate_half(x) * sin)

class Attention(nn.Module):
    def __init__(self):
        super().__init__()
        self.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
        self.k_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
        self.v_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
        self.o_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
    def forward(self, x, cos, sin, mask):
        B, T, _ = x.shape
        q = self.q_proj(x).reshape(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
        k = self.k_proj(x).reshape(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
        v = self.v_proj(x).reshape(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
        q, k = apply_rotary(q, cos, sin), apply_rotary(k, cos, sin)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(HEAD_DIM)
        attn = attn + mask[:, :, :T, :T]
        attn = F.softmax(attn, dim=-1)
        return self.o_proj(attn.matmul(v).transpose(1, 2).reshape(B, T, HIDDEN_SIZE))

class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.gate = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
        self.up = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
        self.down = nn.Linear(INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn_norm = nn.RMSNorm(HIDDEN_SIZE, eps=1e-6)
        self.ffn_norm = nn.RMSNorm(HIDDEN_SIZE, eps=1e-6)
        self.attn = Attention()
        self.ffn = FeedForward()
    def forward(self, x, cos, sin, mask):
        x = x + self.attn(self.attn_norm(x), cos, sin, mask)
        x = x + self.ffn(self.ffn_norm(x))
        return x

mask_cache = {}
def get_causal_mask(T, device):
    if T not in mask_cache:
        m = torch.triu(torch.full((T, T), float('-inf'), device=device), diagonal=1)
        mask_cache[T] = m
    return mask_cache[T][None, None]

class Sage1B(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed_tokens = nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE)
        self.layers = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
        self.norm = nn.RMSNorm(HIDDEN_SIZE, eps=1e-6)
        self.lm_head = nn.Linear(HIDDEN_SIZE, VOCAB_SIZE, bias=False)
        self.rotary = RotaryEmbedding(HEAD_DIM)
        self.max_seq_len = MAX_SEQ_LEN
        self.vocab_size = VOCAB_SIZE
        self.hidden_size = HIDDEN_SIZE

    def forward(self, input_ids, labels=None):
        B, T = input_ids.shape
        x = self.embed_tokens(input_ids) * math.sqrt(HIDDEN_SIZE)
        cos, sin = self.rotary.get_cos_sin(x, T)
        mask = get_causal_mask(T, x.device)
        for layer in self.layers:
            x = layer(x, cos, sin, mask)
        x = self.norm(x)
        logits = self.lm_head(x)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, VOCAB_SIZE), labels.view(-1), ignore_index=0)
        return loss, logits

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=50, temperature=0.8, top_k=40):
        self.eval()
        for _ in range(max_new_tokens):
            if input_ids.size(1) > MAX_SEQ_LEN:
                input_ids = input_ids[:, -MAX_SEQ_LEN:]
            _, logits = self.forward(input_ids)
            logits = logits[:, -1, :] / temperature
            if top_k > 0:
                vals = torch.topk(logits, top_k).values[:, -1:]
                logits[logits < vals] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, nxt], dim=1)
            if nxt.item() == 3:
                break
        return input_ids

model = Sage1B()
total_params = sum(p.numel() for p in model.parameters())
print(f"Parameters: {total_params:,} ({total_params/1e9:.3f}B)")

config = {
    "vocab_size": VOCAB_SIZE, "hidden_size": HIDDEN_SIZE,
    "num_hidden_layers": NUM_LAYERS, "num_attention_heads": NUM_HEADS,
    "head_dim": HEAD_DIM, "intermediate_size": INTERMEDIATE_SIZE,
    "max_position_embeddings": MAX_SEQ_LEN, "model_type": "sage_1b",
    "total_params": total_params, "torch_dtype": "float32",
}
with open(os.path.join(MODEL_DIR, "config.json"), "w") as f:
    json.dump(config, f, indent=2)

# Copy this file as modeling_sage_1b.py for HF distribution
with open(os.path.join(MODEL_DIR, "modeling_sage_1b.py"), "w") as f:
    f.write(open(os.path.abspath(__file__)).read())

# ====== STEP 5: Train ======
print("\n--- Step 5: Training ---")
data = torch.load(os.path.join(DATA_DIR, "tokenized.pt"))
print(f"Training samples: {len(data)}")

class TextDataset(Dataset):
    def __init__(self, data):
        self.data = data
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        t = self.data[idx]
        return t[:-1], t[1:]

tds = TextDataset(data)
loader = DataLoader(tds, batch_size=TRAIN_BATCH_SIZE, shuffle=True, drop_last=True)

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

def get_lr(step):
    if step < WARMUP_STEPS:
        return LEARNING_RATE * (step + 1) / WARMUP_STEPS
    return LEARNING_RATE * (1 - min(step, 10000) / 10000 * 0.9)

best_loss = float('inf')
global_step = 0

for epoch in range(NUM_EPOCHS):
    model.train()
    total_loss = 0
    n_batches = 0
    optimizer.zero_grad()
    epoch_start = time.time()

    for bidx, (inp, tgt) in enumerate(loader):
        loss, _ = model(inp, labels=tgt)
        loss = loss / GRAD_ACCUM_STEPS
        loss.backward()

        if (bidx + 1) % GRAD_ACCUM_STEPS == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            for pg in optimizer.param_groups:
                pg['lr'] = get_lr(global_step)
            optimizer.step()
            optimizer.zero_grad()
            global_step += 1

        total_loss += loss.item() * GRAD_ACCUM_STEPS
        n_batches += 1

        if (bidx + 1) % 200 == 0:
            elapsed = time.time() - epoch_start
            avg = total_loss / max(n_batches, 1)
            lr = optimizer.param_groups[0]['lr']
            print(f"  E{epoch+1} B{bidx+1}/{len(loader)} | Loss: {avg:.4f} | LR: {lr:.2e} | {elapsed:.0f}s")

    avg = total_loss / max(n_batches, 1)
    et = time.time() - epoch_start
    print(f"Epoch {epoch+1} | Avg loss: {avg:.4f} | Time: {et:.0f}s | Steps: {global_step}")

    if avg < best_loss:
        best_loss = avg
        sd = model.state_dict()
        torch.save(sd, os.path.join(MODEL_DIR, "pytorch_model.bin"))
        torch.save({k: v.half() if v.dtype == torch.float32 else v for k, v in sd.items()},
                   os.path.join(MODEL_DIR, "pytorch_model_state.bin"))
        print(f"  Best model saved (loss: {avg:.4f})")

# Final save
sd = model.state_dict()
torch.save(sd, os.path.join(MODEL_DIR, "pytorch_model.bin"))
torch.save({k: v.half() if v.dtype == torch.float32 else v for k, v in sd.items()},
           os.path.join(MODEL_DIR, "pytorch_model_state.bin"))

# Save tokenizer pickle
with open(os.path.join(TOKENIZER_DIR, "tokenizer.pkl"), "wb") as f:
    pickle.dump(tokenizer, f)

# Test generation
print("\n--- Test generation ---")
model.eval()
from tokenizers import Tokenizer as Tk
test_tokenizer = Tk.from_file(os.path.join(TOKENIZER_DIR, "tokenizer.json"))
prompt = "Once upon a time"
tokens = test_tokenizer.encode(prompt).ids
inp = torch.tensor([[1] + tokens[:20]], dtype=torch.long)
out = model.generate(inp, max_new_tokens=30, temperature=0.7)
gen_text = test_tokenizer.decode(out[0].tolist(), skip_special_tokens=True)
print(f"Prompt: {prompt}")
print(f"Generated: {gen_text}")

print(f"\n=== DONE ===")
print(f"Params: {total_params:,} ({total_params/1e9:.3f}B)")
print(f"Best loss: {best_loss:.4f}")
print(f"Model: {MODEL_DIR}")