"""
Kaggle GPU Training Script - Run 1: Chinchilla Optimal (100 Million Tokens)
5M Parameter Model trained on Real Streamed Datasets + Terminal Commands.
Auto-uploads to HuggingFace repository: 'kipasyangin5/5m-terminal-lm-chinchilla'
"""
import os
import math
import time
import random
import json
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 transformers import PreTrainedTokenizerFast
from datasets import load_dataset
from huggingface_hub import HfApi, login

# ==========================================
# 1. Environment & Credentials Setup
# ==========================================
HF_TOKEN = os.environ.get("HF_TOKEN", "YOUR_HF_TOKEN")
HF_REPO_ID = os.environ.get("HF_REPO_ID", "kipasyangin5/5m-terminal-lm-chinchilla")

if HF_TOKEN:
    try:
        login(token=HF_TOKEN)
        print(f"[HF Login] Authenticated as '{HF_REPO_ID.split('/')[0]}'.")
    except Exception as e:
        print(f"[HF Login Warning] {e}")

# ==========================================
# 2. Terminal Commands Data Generator
# ==========================================
COMMAND_TEMPLATES = [
    ("How do I navigate up one directory?", "cd .."),
    ("How do I go to the home directory?", "cd ~"),
    ("How do I check my current directory path?", "pwd"),
    ("How do I list all files including hidden files?", "ls -la"),
    ("How do I list files with human readable file sizes?", "ls -lh"),
    ("How do I create a nested directory structure?", "mkdir -p path/to/nested/directory"),
    ("How do I print directory tree structure?", "tree -L 2"),
    ("How do I copy a directory recursively?", "cp -r source_dir/ target_dir/"),
    ("How do I move or rename a file?", "mv old_filename.txt new_filename.txt"),
    ("How do I force remove a folder and all contents?", "rm -rf target_folder/"),
    ("How do I create an empty file?", "touch index.js"),
    ("How do I inspect the first 20 lines of a file?", "head -n 20 logfile.log"),
    ("How do I monitor a log file in real-time?", "tail -f /var/log/syslog"),
    ("How do I count lines in a text file?", "wc -l dataset.txt"),
    ("How do I recursively search for text in files?", "grep -rn \"search_term\" ."),
    ("How do I find all python files in the current folder?", "find . -type f -name \"*.py\""),
    ("How do I sort lines and remove duplicates?", "sort input.txt | uniq -c"),
    ("How do I make a shell script executable?", "chmod +x script.sh"),
    ("How do I check system RAM usage?", "free -h"),
    ("How do I check disk space usage in human readable format?", "df -h"),
    ("How do I check disk usage of current directories?", "du -sh * | sort -hr"),
    ("How do I download a file silently with curl?", "curl -sSL https://example.com/file.tar.gz -o file.tar.gz"),
    ("How do I check repository status in git?", "git status"),
    ("How do I stage all changed files in git?", "git add ."),
    ("How do I commit staged changes with a message?", "git commit -m \"feat: implement terminal parser\""),
    ("How do I push commits to remote main branch?", "git push origin main")
]

SHELL_INTERACTIONS = [
    "$ cd ..\n$ pwd\n/home/user\n$ ls -la\ntotal 32\ndrwxr-xr-x 4 user user 4096 Aug 2 00:00 .\ndrwxr-xr-x 8 user user 4096 Aug 2 00:00 ..\n-rw-r--r-- 1 user user  220 Aug 2 00:00 .bashrc",
    "$ mkdir project && cd project\n$ git init\nInitialized empty Git repository in /home/user/project/.git/\n$ touch main.py README.md\n$ git status\nOn branch main\nUntracked files:\n  (use \"git add <file>...\" to include in what will be committed)\n\tREADME.md\n\tmain.py",
    "$ grep -rn \"import torch\" src/\nsrc/model.py:1:import torch\nsrc/train.py:2:import torch\nsrc/utils.py:1:import torch",
    "$ chmod +x build.sh\n$ ./build.sh\n[INFO] Building release binary...\n[SUCCESS] Build completed in 2.4s."
]

def generate_terminal_samples(num_samples=20000):
    samples = []
    for _ in range(num_samples):
        qa = random.choice(COMMAND_TEMPLATES)
        fmt = random.choice([
            f"User: {qa[0]}\nAssistant: Run `{qa[1]}`\n",
            f"Question: {qa[0]}\nAnswer:\n```bash\n{qa[1]}\n```\n",
            f"$ {qa[1]}\n# Executed successfully\n"
        ])
        samples.append(fmt)
    for _ in range(num_samples // 2):
        s = random.choice(SHELL_INTERACTIONS)
        samples.append(f"```session\n{s}\n```\n")
    return samples

# ==========================================
# 3. Model Architecture (5.0M Parameters)
# ==========================================
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

class SwiGLUMLP(nn.Module):
    def __init__(self, dim, inter_dim):
        super().__init__()
        self.gate_proj = nn.Linear(dim, inter_dim, bias=False)
        self.up_proj = nn.Linear(dim, inter_dim, bias=False)
        self.down_proj = nn.Linear(inter_dim, dim, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

class CausalSelfAttention(nn.Module):
    def __init__(self, dim, n_head):
        super().__init__()
        self.dim = dim
        self.n_head = n_head
        self.head_dim = dim // n_head
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(y)

class TransformerBlock(nn.Module):
    def __init__(self, dim, n_head, inter_dim):
        super().__init__()
        self.attn = CausalSelfAttention(dim, n_head)
        self.mlp = SwiGLUMLP(dim, inter_dim)
        self.norm1 = RMSNorm(dim)
        self.norm2 = RMSNorm(dim)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

class TerminalLM5M(nn.Module):
    def __init__(self, vocab_size=4096, dim=256, n_layer=6, n_head=8, inter_dim=512, max_seq_len=512):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.tok_embeddings = nn.Embedding(vocab_size, dim)
        self.pos_embeddings = nn.Embedding(max_seq_len, dim)
        self.layers = nn.ModuleList([
            TransformerBlock(dim, n_head, inter_dim) for _ in range(n_layer)
        ])
        self.norm = RMSNorm(dim)
        self.lm_head = nn.Linear(dim, vocab_size, bias=False)
        self.lm_head.weight = self.tok_embeddings.weight

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, targets=None):
        B, T = input_ids.shape
        device = input_ids.device
        pos = torch.arange(0, T, dtype=torch.long, device=device)
        
        h = self.tok_embeddings(input_ids) + self.pos_embeddings(pos)
        for layer in self.layers:
            h = layer(h)
        h = self.norm(h)
        logits = self.lm_head(h)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1))

        return logits, loss

class TextDataset(Dataset):
    def __init__(self, token_ids, seq_len=256):
        self.seq_len = seq_len
        self.num_samples = (len(token_ids) - 1) // seq_len
        self.inputs = []
        self.targets = []
        for i in range(self.num_samples):
            start = i * seq_len
            end = start + seq_len
            self.inputs.append(token_ids[start:end])
            self.targets.append(token_ids[start+1:end+1])

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

    def __getitem__(self, idx):
        return torch.tensor(self.inputs[idx], dtype=torch.long), torch.tensor(self.targets[idx], dtype=torch.long)

# ==========================================
# 4. Tokenizer Construction
# ==========================================
def build_tokenizer(vocab_size=4096):
    print(f"[Tokenizer] Training 4096 BPE Tokenizer...")
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)

    special_tokens = ["<pad>", "<s>", "</s>", "<unk>", "<cmd>", "</cmd>", "<user>", "<assistant>"]

    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=special_tokens,
        min_frequency=2,
        show_progress=False
    )

    corpus = generate_terminal_samples(15000)
    tokenizer.train_from_iterator(corpus, trainer=trainer)

    fast_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<pad>",
        unk_token="<unk>",
        mask_token="<mask >",
        additional_special_tokens=["<cmd>", "</cmd>", "<user>", "<assistant>"]
    )
    return fast_tokenizer

# ==========================================
# 5. Main Training Routine
# ==========================================
def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("==================================================")
    print(" RUN 1: CHINCHILLA OPTIMAL (100M TOKENS)")
    print(f" Target HF Repo: {HF_REPO_ID}")
    print(f" Device: {device.upper()}")
    if device == "cuda":
        print(f" GPU Device: {torch.cuda.get_device_name(0)}")
    print("==================================================")

    # 1. Build Tokenizer
    tokenizer = build_tokenizer(vocab_size=4096)
    vocab_size = len(tokenizer)

    # 2. Build Dataset (Streaming Real Datasets + Terminal Commands)
    print("[Dataset] Building dataset from real sources + Terminal engine...")
    terminal_samples = generate_terminal_samples(num_samples=30000)
    
    # Try streaming real wikitext from HF
    wikitext_text = ""
    try:
        print("[Dataset] Streaming real 'wikitext-2-raw-v1' from Hugging Face...")
        ds_wiki = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train', streaming=True)
        wiki_lines = []
        for idx, item in enumerate(ds_wiki):
            if idx >= 5000: break
            if item['text'].strip():
                wiki_lines.append(item['text'])
        wikitext_text = "\n".join(wiki_lines)
        print(f"[Dataset] Streamed {len(wiki_lines)} lines of real Wiki text.")
    except Exception as e:
        print(f"[Dataset Warning] {e}")

    full_text = "\n".join(terminal_samples) + "\n" + wikitext_text
    tokens = tokenizer.encode(full_text)

    seq_len = 256
    batch_size = 64 if device == "cuda" else 8
    dataset = TextDataset(tokens, seq_len=seq_len)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)

    # 3. Model Setup
    model = TerminalLM5M(
        vocab_size=vocab_size,
        dim=256,
        n_layer=6,
        n_head=8,
        inter_dim=512,
        max_seq_len=seq_len
    ).to(device)

    params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[Model] Trainable Parameters: {params:,} (~{params/1e6:.2f}M)")

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
    scaler = torch.amp.GradScaler('cuda') if device == "cuda" else None

    # Target: 6,104 steps @ 16,384 tokens/step = ~100 Million Tokens
    total_steps = 6104 if device == "cuda" else 30
    model.train()
    step = 0
    t0 = time.time()
    data_iter = iter(dataloader)

    while step < total_steps:
        try:
            x, y = next(data_iter)
        except StopIteration:
            data_iter = iter(dataloader)
            x, y = next(data_iter)

        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()

        if device == "cuda":
            with torch.amp.autocast('cuda', dtype=torch.float16):
                logits, loss = model(x, y)
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            logits, loss = model(x, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

        step += 1

        if step % 200 == 0 or step == 1:
            t1 = time.time()
            dt = t1 - t0
            t0 = t1
            tok_s = (200 * batch_size * seq_len) / (dt if dt > 0 else 1.0) if step > 1 else 0
            ppl = math.exp(min(loss.item(), 20.0))
            tokens_so_far = step * batch_size * seq_len
            print(f"Step {step:5d}/{total_steps} | Tokens: {tokens_so_far:,}/100,000,000 | Loss: {loss.item():.4f} | PPL: {ppl:.2f} | Speed: {tok_s:.0f} tok/s")

    print("\n[SUCCESS] Run 1 (Chinchilla Optimal 100M Tokens) Completed!")

    # Save and Upload
    output_dir = "saved_5m_model_chinchilla"
    os.makedirs(output_dir, exist_ok=True)
    torch.save(model.state_dict(), os.path.join(output_dir, "model.pt"))
    tokenizer.save_pretrained(output_dir)

    config_dict = {
        "model_type": "terminal_lm_5m_chinchilla",
        "tokens_trained": step * batch_size * seq_len,
        "vocab_size": vocab_size,
        "dim": 256,
        "n_layer": 6,
        "n_head": 8,
        "total_parameters": params
    }
    with open(os.path.join(output_dir, "config.json"), "w") as f:
        json.dump(config_dict, f, indent=2)

    if HF_TOKEN:
        try:
            print(f"[HuggingFace] Pushing model to '{HF_REPO_ID}'...")
            api = HfApi()
            api.create_repo(repo_id=HF_REPO_ID, exist_ok=True)
            api.upload_folder(
                folder_path=output_dir,
                repo_id=HF_REPO_ID,
                commit_message=f"Upload 5M Chinchilla Optimal model (100M tokens, Loss: {loss.item():.4f})"
            )
            print(f"🚀 [HF Upload Complete] Model live at: https://huggingface.co/{HF_REPO_ID}")
        except Exception as e:
            print(f"[HF Upload Error] {e}")

if __name__ == "__main__":
    main()