# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: test/test_model_chat.py
# ─────────────────────────────────────────────────────────

"""
AXIOM-1 Interactive Chat CLI

Downloads the Phase 1 checkpoint + tokenizer from the
Hugging Face Hub (ErebusTN/AXIOM-1) and runs a simple
interactive chat loop for testing the trained model.

Usage:
    python test/test_model_chat.py
    python test/test_model_chat.py --device cpu
    python test/test_model_chat.py --temperature 0.6 --top_k 40
"""

import os
import sys
import argparse
import math
import time

import torch
from huggingface_hub import hf_hub_download, snapshot_download

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.inference.generate import generate
from axiom1.tokenizer.tokenizer import AXIOM1Tokenizer


# ── Constants ────────────────────────────────────────────────────────────────
REPO_ID = "ErebusTN/AXIOM-1"
CKPT_FILENAME = "AXIOM-1_phase1_final.pt"
TOKENIZER_SUBDIR = "tokenizer"
CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "axiom1")


def parse_args() -> argparse.Namespace:
    """Parse CLI arguments."""
    p = argparse.ArgumentParser(
        description="AXIOM-1 Interactive Chat CLI — EGen Labs | ErebusTN"
    )
    p.add_argument(
        "--device", type=str, default="auto",
        help="Device: 'cuda', 'cpu', or 'auto' (default: auto)"
    )
    p.add_argument(
        "--temperature", type=float, default=0.8,
        help="Sampling temperature (default: 0.8)"
    )
    p.add_argument(
        "--top_k", type=int, default=50,
        help="Top-k filtering (default: 50)"
    )
    p.add_argument(
        "--top_p", type=float, default=0.9,
        help="Nucleus sampling threshold (default: 0.9)"
    )
    p.add_argument(
        "--max_tokens", type=int, default=256,
        help="Maximum tokens to generate per reply (default: 256)"
    )
    p.add_argument(
        "--num_loops", type=int, default=4,
        help="Number of loop iterations (4 for Phase 1, 12 for Phase 2)"
    )
    p.add_argument(
        "--checkpoint", type=str, default=None,
        help="Path to a local .pt checkpoint (skips HF download)"
    )
    p.add_argument(
        "--tokenizer_dir", type=str, default=None,
        help="Path to a local tokenizer directory (skips HF download)"
    )
    return p.parse_args()


def resolve_device(requested: str) -> str:
    """Pick the best available device."""
    if requested != "auto":
        return requested
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def download_checkpoint(cache_dir: str) -> str:
    """Download phase1_final.pt from HuggingFace Hub."""
    print(f"📦 Downloading checkpoint from {REPO_ID}...")
    path = hf_hub_download(
        repo_id=REPO_ID,
        filename=CKPT_FILENAME,
        cache_dir=cache_dir,
    )
    print(f"  ✅ Checkpoint cached: {path}")
    return path


def download_tokenizer(cache_dir: str) -> str:
    """Download tokenizer directory from HuggingFace Hub."""
    print(f"📦 Downloading tokenizer from {REPO_ID}...")
    tok_dir = snapshot_download(
        repo_id=REPO_ID,
        allow_patterns=f"{TOKENIZER_SUBDIR}/*",
        cache_dir=cache_dir,
    )
    tok_path = os.path.join(tok_dir, TOKENIZER_SUBDIR)
    if not os.path.isdir(tok_path):
        tok_path = tok_dir
    print(f"  ✅ Tokenizer cached: {tok_path}")
    return tok_path


def load_model(
    ckpt_path: str,
    num_loops: int,
    device: str,
) -> tuple[AXIOM1, AXIOM1Config]:
    """Load AXIOM-1 model from a checkpoint."""
    print("🏗️  Loading AXIOM-1 model...")

    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)

    if "cfg" in checkpoint:
        cfg = checkpoint["cfg"]
    else:
        cfg = AXIOM1Config()

    cfg.num_loops = num_loops

    model = AXIOM1(cfg)
    model.load_state_dict(checkpoint["model_state_dict"])
    model = model.to(device)
    model.eval()

    total_params = sum(p.numel() for p in model.parameters())
    step = checkpoint.get("step", "?")
    loss = checkpoint.get("final_loss", checkpoint.get("loss", "?"))

    print(f"  ✅ Model loaded")
    print(f"     Params     : {total_params / 1e6:.2f}M")
    print(f"     Loops      : {cfg.num_loops}")
    print(f"     Checkpoint : step={step}, loss={loss}")

    return model, cfg


def load_tokenizer(tok_path: str) -> AXIOM1Tokenizer:
    """Load the AXIOM-1 tokenizer."""
    print("🔤 Loading tokenizer...")
    tokenizer = AXIOM1Tokenizer.load(tok_path)
    print(f"  ✅ Tokenizer loaded (vocab={tokenizer.vocab_size:,})")
    return tokenizer


def stream_generate(
    model: AXIOM1,
    prompt_ids: torch.Tensor,
    cfg: AXIOM1Config,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    use_sim: bool,
):
    """Yields generated tokens one by one, along with internal loop statistics."""
    model.eval()
    device = prompt_ids.device
    ids = prompt_ids.clone()
    
    loop_data = []
    
    def gate_hook(module, input, output):
        if not use_sim: return
        # SwiGLU gate output: (batch, seq, d_inner)
        acts = output[0, -1, :] 
        active_pct = (acts > 0).float().mean().item() * 100
        max_act = acts.max().item()
        loop_data.append({"active": active_pct, "max": max_act})
        
    def out_hook(module, input, output):
        if not use_sim: return
        h = output[0, -1, :]
        l2_norm = h.norm().item()
        if len(loop_data) > 0:
            loop_data[-1]["norm"] = l2_norm
            
    # Attach hooks to capture metrics per loop iteration
    h1 = model.block.ffn.gate.register_forward_hook(gate_hook)
    h2 = model.block.register_forward_hook(out_hook)
    
    try:
        dev_type = getattr(device, 'type', str(device))
        autocast_device = 'cuda' if dev_type == 'cuda' else 'cpu'
        
        for _ in range(max_new_tokens):
            input_ids = ids[:, -cfg.max_seq_len:]
            loop_data.clear()
            
            with torch.autocast(device_type=autocast_device, dtype=torch.bfloat16):
                logits = model(input_ids)
                
            next_logits = logits[:, -1, :].float() / max(temperature, 1e-5)
            
            if top_k > 0:
                vals, _ = torch.topk(next_logits, top_k)
                next_logits[next_logits < vals[:, -1:]] = float('-inf')
                
            probs = torch.softmax(next_logits, dim=-1)
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cumprobs = sorted_probs.cumsum(dim=-1)
            remove = cumprobs - sorted_probs > top_p
            sorted_probs[remove] = 0.0
            sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)
            probs.scatter_(-1, sorted_idx, sorted_probs)
            
            next_token = torch.multinomial(probs, num_samples=1)
            ids = torch.cat([ids, next_token], dim=1)
            
            yield next_token, list(loop_data)
            
            if next_token.item() == cfg.eos_token_id:
                break
    finally:
        h1.remove()
        h2.remove()


def chat_loop(
    model: AXIOM1,
    cfg: AXIOM1Config,
    tokenizer: AXIOM1Tokenizer,
    device: str,
    temperature: float,
    top_k: int,
    top_p: float,
    max_tokens: int,
) -> None:
    """Run the interactive chat CLI."""

    CYAN = "\033[96m"
    YELLOW = "\033[93m"
    DIM = "\033[2m"
    RESET = "\033[0m"
    BOLD = "\033[1m"
    GREEN = "\033[92m"

    print()
    print("=" * 60)
    print(f"  {BOLD}AXIOM-1  ·  Interactive Chat  ·  EGen Labs | ErebusTN{RESET}")
    print("=" * 60)
    print("  Type your message and press Enter.")
    print(f"  Commands: {DIM}/quit  /clear  /config  /raw <text>  /sim on|off{RESET}")
    print("=" * 60)
    print()

    history_text = ""
    use_sim = False

    while True:
        try:
            user_input = input(f"{CYAN}[YOU]{RESET} ")
        except (KeyboardInterrupt, EOFError):
            print(f"\n{DIM}👋 Goodbye!{RESET}")
            break

        user_input = user_input.strip()
        if not user_input:
            continue

        # ── Commands ──────────────────────────────────────────────────────
        if user_input.lower() in ("/quit", "/exit", "/q"):
            print(f"{DIM}👋 Goodbye!{RESET}")
            break

        if user_input.lower() == "/clear":
            history_text = ""
            print(f"{DIM}🧹 Conversation cleared.{RESET}\n")
            continue

        if user_input.lower() == "/sim on":
            use_sim = True
            print(f"{GREEN}🔬 Neuron Route Simulator ENABLED.{RESET}\n")
            continue

        if user_input.lower() == "/sim off":
            use_sim = False
            print(f"{DIM}🔬 Neuron Route Simulator DISABLED.{RESET}\n")
            continue

        if user_input.lower() == "/config":
            print(f"  temperature : {temperature}")
            print(f"  top_k       : {top_k}")
            print(f"  top_p       : {top_p}")
            print(f"  max_tokens  : {max_tokens}")
            print(f"  num_loops   : {cfg.num_loops}")
            print(f"  device      : {device}")
            print(f"  simulator   : {'ON' if use_sim else 'OFF'}")
            print()
            continue

        # ── /raw mode: send text directly without chat formatting ─────────
        if user_input.lower().startswith("/raw "):
            raw_text = user_input[5:]
            prompt_ids = tokenizer.encode(raw_text, add_special_tokens=True)
            prompt_tensor = torch.tensor([prompt_ids], dtype=torch.long).to(device)

            print(f"{YELLOW}[RAW]{RESET} ", end="", flush=True)
            
            reply_ids = []
            if use_sim:
                print("\n" * (cfg.num_loops + 2))
                sys.stdout.write(f"\033[{cfg.num_loops + 2}A")
                
            t0 = time.time()
            for token_tensor, sim_data in stream_generate(
                model, prompt_tensor, cfg, max_tokens, temperature, top_p, top_k, use_sim
            ):
                tok_id = token_tensor.item()
                reply_ids.append(tok_id)
                chunk = tokenizer.decode([tok_id])
                sys.stdout.write(chunk)
                sys.stdout.flush()
                
                if use_sim and sim_data:
                    sys.stdout.write("\033[s") # save cursor
                    sys.stdout.write("\n\n")
                    for i, d in enumerate(sim_data):
                        bar_len = min(int(d['max']), 20)
                        bar = "█" * bar_len + "░" * (20 - bar_len)
                        sys.stdout.write(f"\033[K  {DIM}Loop {i:02d} | Firing: {d['active']:>4.1f}% | L2 Norm: {d.get('norm', 0):>5.1f} | Peak:{RESET} {GREEN}{bar}{RESET} {DIM}({d['max']:.1f}){RESET}\n")
                    sys.stdout.write("\033[u") # restore cursor
                    sys.stdout.flush()
                    time.sleep(0.05) # visual ease
                    
            elapsed = time.time() - t0
            if use_sim:
                sys.stdout.write(f"\033[{cfg.num_loops + 2}B\n")
            else:
                print()
                
            n_toks = len(reply_ids)
            tps = n_toks / max(elapsed, 1e-6)
            print(f"{DIM}  ({n_toks} tokens in {elapsed:.2f}s — {tps:.1f} tok/s){RESET}\n")
            continue

        # ── Build chat prompt ─────────────────────────────────────────────
        history_text += f"user: {user_input}\n\nassistant:"
        prompt_ids = tokenizer.encode(history_text, add_special_tokens=True)

        # Truncate to max_seq_len if too long
        if len(prompt_ids) > cfg.max_seq_len:
            prompt_ids = prompt_ids[-cfg.max_seq_len:]

        prompt_tensor = torch.tensor([prompt_ids], dtype=torch.long).to(device)

        # ── Generate ──────────────────────────────────────────────────────
        print(f"{YELLOW}[AXIOM-1]{RESET} ", end="", flush=True)
        
        reply_ids = []
        if use_sim:
            print("\n" * (cfg.num_loops + 2))
            sys.stdout.write(f"\033[{cfg.num_loops + 2}A")
            
        t0 = time.time()
        for token_tensor, sim_data in stream_generate(
            model, prompt_tensor, cfg, max_tokens, temperature, top_p, top_k, use_sim
        ):
            tok_id = token_tensor.item()
            if tok_id == cfg.eos_token_id:
                break
            reply_ids.append(tok_id)
            chunk = tokenizer.decode([tok_id])
            sys.stdout.write(chunk)
            sys.stdout.flush()
            
            if use_sim and sim_data:
                sys.stdout.write("\033[s") # save cursor
                sys.stdout.write("\n\n")
                for i, d in enumerate(sim_data):
                    bar_len = min(int(d['max']), 20)
                    bar = "█" * bar_len + "░" * (20 - bar_len)
                    # \033[K clears the line to the right to prevent artifacting from previous longer tokens
                    sys.stdout.write(f"\033[K  {DIM}Loop {i:02d} | Firing: {d['active']:>4.1f}% | L2 Norm: {d.get('norm', 0):>5.1f} | Peak:{RESET} {GREEN}{bar}{RESET} {DIM}({d['max']:.1f}){RESET}\n")
                sys.stdout.write("\033[u") # restore cursor
                sys.stdout.flush()
                time.sleep(0.05) # visual ease

        elapsed = time.time() - t0
        
        if use_sim:
            # Move down past the dashboard to prevent next prompt from overwriting it
            sys.stdout.write(f"\033[{cfg.num_loops + 2}B\n")
        else:
            print()
            
        n_toks = len(reply_ids)
        tps = n_toks / max(elapsed, 1e-6)
        
        reply = tokenizer.decode(reply_ids).strip()
        history_text += f" {reply}\n\n"
        
        # ── Display Summary ───────────────────────────────────────────────
        print(f"{DIM}  ({n_toks} tokens in {elapsed:.2f}s — {tps:.1f} tok/s){RESET}\n")


def main() -> None:
    """Entry point."""
    args = parse_args()
    device = resolve_device(args.device)

    print()
    print("╔════════════════════════════════════════════════════════════╗")
    print("║         AXIOM-1  ·  Model Test CLI  ·  EGen Labs          ║")
    print("║     EGen-Core/AXIOM-1  ·  LoopLLM  ·  bfloat16 only      ║")
    print("╚════════════════════════════════════════════════════════════╝")
    print()
    print(f"  Device : {device.upper()}")
    if device == "cuda":
        gpu = torch.cuda.get_device_properties(0)
        print(f"  GPU    : {gpu.name} ({gpu.total_memory / 1024**3:.1f} GB)")
    print()

    # ── Resolve checkpoint ────────────────────────────────────────────────
    os.makedirs(CACHE_DIR, exist_ok=True)

    if args.checkpoint:
        ckpt_path = args.checkpoint
        print(f"📂 Using local checkpoint: {ckpt_path}")
    else:
        ckpt_path = download_checkpoint(CACHE_DIR)

    if args.tokenizer_dir:
        tok_path = args.tokenizer_dir
        print(f"📂 Using local tokenizer: {tok_path}")
    else:
        tok_path = download_tokenizer(CACHE_DIR)

    # ── Load model + tokenizer ────────────────────────────────────────────
    model, cfg = load_model(ckpt_path, args.num_loops, device)
    tokenizer = load_tokenizer(tok_path)

    # ── Start interactive chat ────────────────────────────────────────────
    chat_loop(
        model=model,
        cfg=cfg,
        tokenizer=tokenizer,
        device=device,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
    )


if __name__ == "__main__":
    main()
