"""AngstromE1-Nano: one-click Kaggle setup → train → download"""
import os, sys, subprocess, time, re, json, math
from pathlib import Path
import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast

# ═══════════════════════════════════════════════════════════════════
# CONFIG
# ═══════════════════════════════════════════════════════════════════
GIT_REPO       = "https://github.com/er-del/angstrom.git"
BRANCH         = "main"
MODEL_CONFIG   = "medium"
TRAIN_STEPS    = 50000
BATCH_SIZE     = 16
SEQ_LEN        = 256
SAVE_EVERY     = 5000
SSH_PASSWORD   = "ChangeMe123!"
KAGGLE_OUT     = Path("/kaggle/working/angstrom_output")
LOCAL_OUT      = Path("checkpoints")

def log(msg): print(f"\033[0;32m[+]\033[0m {msg}")
def info(msg): print(f"    {msg}")
def err(msg): print(f"\033[0;31m[-]\033[0m {msg}")

# ═══════════════════════════════════════════════════════════════════
# 1. SYSTEM SETUP
# ═══════════════════════════════════════════════════════════════════
log("=== 1/7: System setup ===")
log("Checking CUDA + GPUs...")
n_gpus = torch.cuda.device_count()
for i in range(n_gpus):
    p = torch.cuda.get_device_properties(i)
    info(f"GPU {i}: {p.name}  {p.total_memory/1e9:.1f}GB  CUDA {torch.version.cuda}")
log(f"{n_gpus} GPU(s) available")

log("Installing system packages...")
os.system("apt-get update -qq > /dev/null 2>&1")
os.system("apt-get install -y -qq openssh-server curl wget git-lfs > /dev/null 2>&1")

# ═══════════════════════════════════════════════════════════════════
# 2. CLONE REPO
# ═══════════════════════════════════════════════════════════════════
log("=== 2/7: Clone repo ===")
if os.path.exists("angstrom"):
    log("angstrom/ already exists, pulling latest...")
    os.system("cd angstrom && git pull 2>/dev/null")
else:
    os.system(f"git clone {GIT_REPO} angstrom 2>&1")
    os.system(f"cd angstrom && git checkout {BRANCH} 2>/dev/null")

os.chdir("angstrom")
REPO = Path.cwd()
info(f"Working dir: {REPO}")

# ═══════════════════════════════════════════════════════════════════
# 3. PYTHON DEPENDENCIES
# ═══════════════════════════════════════════════════════════════════
log("=== 3/7: Python dependencies ===")
os.system("pip install -q --upgrade pip 2>/dev/null")
os.system("pip install -q datasets>=2.16.0 safetensors>=0.4.0 tqdm>=4.65.0 accelerate>=0.25.0 2>/dev/null")
os.system("pip install -q -e . 2>/dev/null || true")

try:
    from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM
    from angstrom_nano.tokenizer import AngstromNanoTokenizer
    log("All imports OK")
except Exception as e:
    err(f"Import failed: {e}")
    sys.exit(1)

# ═══════════════════════════════════════════════════════════════════
# 4. DOWNLOAD DATA
# ═══════════════════════════════════════════════════════════════════
log("=== 4/7: Download training data ===")
data_script = REPO / "download_data.py"
if data_script.exists():
    info("Running download_data.py...")
    os.system(f"python {data_script} 2>&1")

data_path = REPO / "data" / "corpus.txt"
if not data_path.exists():
    err("Data download failed")
    info("Writing fallback mini-corpus...")
    (REPO / "data").mkdir(exist_ok=True)
    data_path.write_text("Hello world.\n" * 1000000)
text = data_path.read_text()
info(f"Data: {len(text):,} chars ({len(text)/1e6:.1f}MB)")

# ═══════════════════════════════════════════════════════════════════
# 5. TRAIN MODEL
# ═══════════════════════════════════════════════════════════════════
log("=== 5/7: Train model ===")

CONFIGS = {
    "small": {
        "vocab_size": 8192, "hidden_size": 256, "intermediate_size": 1024,
        "num_hidden_layers": 8, "num_attention_heads": 8, "num_key_value_heads": 2,
        "head_dim": 32, "num_local_experts": 4, "num_experts_per_tok": 2,
        "max_position_embeddings": 2048, "sliding_window": 512,
        "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True,
        "tie_word_embeddings": True,
    },
    "medium": {
        "vocab_size": 4096, "hidden_size": 192, "intermediate_size": 768,
        "num_hidden_layers": 6, "num_attention_heads": 6, "num_key_value_heads": 3,
        "head_dim": 32, "num_local_experts": 4, "num_experts_per_tok": 2,
        "max_position_embeddings": 2048, "sliding_window": 512,
        "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True,
        "tie_word_embeddings": True,
    },
    "large": {
        "vocab_size": 16384, "hidden_size": 512, "intermediate_size": 2048,
        "num_hidden_layers": 12, "num_attention_heads": 16, "num_key_value_heads": 4,
        "head_dim": 32, "num_local_experts": 8, "num_experts_per_tok": 2,
        "max_position_embeddings": 4096, "sliding_window": 1024,
        "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True,
        "tie_word_embeddings": True,
    },
}

config_dict = CONFIGS[MODEL_CONFIG]
info(f"Using config: {MODEL_CONFIG}")

# Tokenizer
tok_path = REPO / "checkpoints/tokenizer.json"
if tok_path.exists():
    tok = AngstromNanoTokenizer.from_bpe_file(str(tok_path))
    config_dict["vocab_size"] = len(tok)
    info(f"Loaded tokenizer: {len(tok)} vocab")
else:
    info("Training new tokenizer...")
    tok_path.parent.mkdir(parents=True, exist_ok=True)
    tok = AngstromNanoTokenizer.train_bpe(
        [str(data_path)], vocab_size=config_dict["vocab_size"],
        out_path=str(tok_path))
    info(f"Trained tokenizer: {len(tok)} vocab")

# Dataset
class TextDataset(torch.utils.data.Dataset):
    def __init__(self, token_ids, seq_len):
        self.token_ids = token_ids
        self.seq_len = seq_len
        self.n_samples = len(token_ids) - seq_len - 1
    def __len__(self): return self.n_samples
    def __getitem__(self, idx):
        return (self.token_ids[idx:idx+self.seq_len],
                self.token_ids[idx+1:idx+self.seq_len+1])

info("Tokenizing data...")
ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long)
info(f"Tokens: {len(ids):,} ({len(ids)/1e6:.1f}M)")

dataset = TextDataset(ids, SEQ_LEN)
dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=BATCH_SIZE, shuffle=True,
    num_workers=2, pin_memory=True, drop_last=True)
info(f"Samples: {len(dataset):,}")

# Model
device = torch.device("cuda")
cfg = AngstromNanoConfig(**config_dict)
model: nn.Module = AngstromNanoForCausalLM(cfg)
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
info(f"Model: {n_params:,} params ({n_params*4/1e6:.1f}MB FP32)")
model = model.to(device)

if n_gpus > 1:
    model = nn.DataParallel(model, device_ids=list(range(n_gpus)))
    info(f"DataParallel across {n_gpus} GPUs")

# Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.1, betas=(0.9, 0.95))
scaler = GradScaler(enabled=True)

WARMUP = 500
MIN_LR = 3e-4
GRAD_CLIP = 1.0

def get_lr(step):
    if step < WARMUP:
        return 3e-3 * step / max(1, WARMUP)
    progress = (step - WARMUP) / max(1, TRAIN_STEPS - WARMUP)
    return MIN_LR + 0.5 * (3e-3 - MIN_LR) * (1.0 + math.cos(math.pi * progress))

log("Starting training loop...")
model.train()
t0 = time.time()
running_loss = 0.0
running_steps = 0

LOCAL_OUT.mkdir(parents=True, exist_ok=True)
KAGGLE_OUT.mkdir(parents=True, exist_ok=True)

data_iter = iter(dataloader)
for step in range(1, TRAIN_STEPS + 1):
    try:
        x, y = next(data_iter)
    except StopIteration:
        data_iter = iter(dataloader)
        x, y = next(data_iter)

    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)

    with autocast(enabled=True, dtype=torch.float16):
        out = model(x, labels=y, output_router_logits=True)
        loss = out["loss"]
        aux_loss = out.get("aux_loss", torch.tensor(0.0))

    optimizer.zero_grad()
    scaler.scale(loss).backward()

    if GRAD_CLIP > 0:
        scaler.unscale_(optimizer)
        nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)

    scaler.step(optimizer)
    scaler.update()

    lr = get_lr(step)
    for pg in optimizer.param_groups:
        pg["lr"] = lr

    running_loss += loss.item()
    running_steps += 1

    if step % 100 == 0 or step == 1:
        avg = running_loss / running_steps
        ppl = math.exp(min(avg, 20))
        elapsed = time.time() - t0
        tok_s = (BATCH_SIZE * SEQ_LEN * running_steps) / max(1, elapsed)
        gpu_mem = torch.cuda.memory_allocated(0) / 1e9
        aux_val = aux_loss.item() if isinstance(aux_loss, torch.Tensor) else 0.0
        print(f"  step {step:>6d}/{TRAIN_STEPS}  loss={avg:.4f}  ppl={ppl:.2f}  "
              f"aux={aux_val:.6f}  lr={lr:.1e}  tok/s={tok_s:.0f}  "
              f"gpu={gpu_mem:.1f}GB  {elapsed:.0f}s")
        running_loss = 0.0
        running_steps = 0

    if step % SAVE_EVERY == 0:
        log(f"Saving checkpoint step {step}...")
        sd = model.state_dict()
        from safetensors.torch import save_file
        fname = f"checkpoint-{step}.safetensors"
        save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(LOCAL_OUT / fname))
        save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(KAGGLE_OUT / fname))
        (LOCAL_OUT / "config.json").write_text(json.dumps(config_dict, indent=2))
        (KAGGLE_OUT / "config.json").write_text(json.dumps(config_dict, indent=2))

# ── Final save ────────────────────────────────────────────────────
log("Training complete! Saving final model...")
sd = model.state_dict()
from safetensors.torch import save_file
for out_dir in [LOCAL_OUT, KAGGLE_OUT]:
    save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(out_dir / "model_final.safetensors"))
    (out_dir / "config.json").write_text(json.dumps(config_dict, indent=2))
    if tok_path.exists():
        import shutil
        shutil.copy(tok_path, out_dir / "tokenizer.json")

total_h = (time.time() - t0) / 3600
info(f"Total time: {total_h:.1f} hours")
info(f"Model saved to: {KAGGLE_OUT}/  ← download from Kaggle 'Output' tab!")

# ═══════════════════════════════════════════════════════════════════
# 6. SSH TUNNEL
# ═══════════════════════════════════════════════════════════════════
log("=== 6/7: SSH tunnel ===")
os.system("sed -i '/^PermitRootLogin\\|^PasswordAuthentication\\|^UseDNS/d' /etc/ssh/sshd_config")
with open("/etc/ssh/sshd_config", "a") as f:
    f.write("PermitRootLogin yes\nPasswordAuthentication yes\nUseDNS no\n")
os.system(f"echo 'root:{SSH_PASSWORD}' | chpasswd")
os.system("mkdir -p /var/run/sshd && /usr/sbin/sshd 2>/dev/null")

log("Starting bore tunnel...")
os.system("curl -sL https://github.com/ekzhang/bore/releases/download/v0.5.2/bore-v0.5.2-x86_64-unknown-linux-musl.tar.gz | tar xz -C /usr/local/bin/ 2>/dev/null")

port = ""
if os.path.exists("/usr/local/bin/bore"):
    bore_proc = subprocess.Popen(
        ["bore", "local", "22", "--to", "bore.pub"],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
    )
    start = time.time()
    while time.time() - start < 30:
        assert bore_proc.stdout is not None
        line = bore_proc.stdout.readline()
        if line:
            print(f"    {line.rstrip()}")
            m = re.search(r"bore\.pub:(\d+)", line)
            if m:
                port = m.group(1)
                break

if port:
    print()
    print("=" * 60)
    print("  SSH TUNNEL READY")
    print("=" * 60)
    print(f"\n  ssh root@bore.pub -p {port} -o StrictHostKeyChecking=no")
    print(f"  Password: {SSH_PASSWORD}")
    print(f"\n  Model files at: {KAGGLE_OUT}/")
    print("  KEEP THIS CELL RUNNING\n")
    print("=" * 60)
    try:
        while True: time.sleep(10)
    except KeyboardInterrupt:
        pass

# ═══════════════════════════════════════════════════════════════════
# 7. SUMMARY
# ═══════════════════════════════════════════════════════════════════
log("=== 7/7: Complete ===")
info("Download from Kaggle: 'Output' tab → angstrom_output/")
for f in sorted(KAGGLE_OUT.glob("*")):
    info(f"  {f.name} ({f.stat().st_size / 1e6:.1f}MB)")