#!/usr/bin/env python3
"""
Sixpert K2 - Example Generation Script
=======================================
Demonstrates how to load and run inference with Sixpert K2 (Q4_K_M GGUF, MoE).

Sixpert K2 uses Mixture-of-Experts architecture with 16 experts and activates
only 2 per token, enabling ~8.9B total parameters while maintaining fast
inference speeds comparable to ~1.2B dense models.

Usage:
    pip install llama-cpp-python
    python generate.py --prompt "Explain the theory of relativity"
"""

import argparse
import time
import sys

try:
    from llama_cpp import Llama
except ImportError:
    print("Installing llama-cpp-python...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
    from llama_cpp import Llama


def format_prompt(messages: list[dict]) -> str:
    """Format messages into Sixpert chat template."""
    formatted = ""
    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        if role == "system":
            formatted += f"<|im_start|>system\n{content}<|im_end|>\n"
        elif role == "user":
            formatted += f"<|im_start|>user\n{content}<|im_end|>\n"
        elif role == "assistant":
            formatted += f"<|im_start|>assistant\n{content}<|im_end|>\n"
    formatted += "<|im_start|>assistant\n"
    return formatted


def run_generation(
    model_path: str,
    prompt: str,
    max_tokens: int = 4096,
    temperature: float = 0.6,
    top_p: float = 0.85,
    top_k: int = 50,
    repeat_penalty: float = 1.08,
    gpu_layers: int = -1,
    threads: int = 8,
    verbose: bool = True,
):
    """Run text generation with Sixpert K2."""

    print(f"Loading Sixpert K2 from: {model_path}")
    print(f"Architecture: MoE (16 experts, 2 active per token)")
    print(f"Quantization: Q4_K_M | Layers: {gpu_layers if gpu_layers > 0 else 'All (offload)'}")
    print(f"Total params: ~8.9B | Active per token: ~1.2B")
    print("-" * 60)

    llm = Llama(
        model_path=model_path,
        n_ctx=131072,
        n_gpu_layers=gpu_layers,
        n_threads=threads,
        verbose=False,
    )

    messages = [
        {
            "role": "system",
            "content": "You are Sixpert K2, a deep reasoning engine developed by Sixpert AI. "
            "You are a Mixture-of-Experts model with exceptional capabilities in: "
            "deep reasoning and multi-step problem solving, "
            "long-context document analysis (up to 1M tokens), "
            "complex mathematical proofs and derivations, "
            "advanced code generation and system design, "
            "scientific research and analysis, "
            "agentic workflows with tool use. "
            "You always think deeply before responding, exploring multiple "
            "reasoning paths before arriving at your answer.",
        },
        {"role": "user", "content": prompt},
    ]

    formatted_prompt = format_prompt(messages)

    if verbose:
        print(f"\nPrompt:\n{prompt}\n")
        print("Generating response (deep reasoning mode)...")
        print("-" * 40)

    start_time = time.time()

    stream = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repeat_penalty=repeat_penalty,
        stream=True,
    )

    full_response = ""
    for chunk in stream:
        delta = chunk["choices"][0]["delta"].get("content", "")
        if delta:
            full_response += delta
            if verbose:
                print(delta, end="", flush=True)

    elapsed = time.time() - start_time

    if verbose:
        print("\n")
        print("-" * 60)
        print(f"Generation completed in {elapsed:.2f}s")
        print(f"Output: {len(full_response.split())} words | {len(full_response)} chars")
        print(f"Note: MoE architecture used ~1.2B active params per token")

    return full_response


def main():
    parser = argparse.ArgumentParser(description="Sixpert K2 Generation Script")
    parser.add_argument(
        "--model", type=str, default="SixpertK2.gguf", help="Path to GGUF model file"
    )
    parser.add_argument("--prompt", type=str, default="What is your name and what makes you special?", help="Input prompt")
    parser.add_argument("--max-tokens", type=int, default=4096, help="Maximum tokens to generate")
    parser.add_argument("--temperature", type=float, default=0.6, help="Sampling temperature")
    parser.add_argument("--top-p", type=float, default=0.85, help="Top-p sampling")
    parser.add_argument("--top-k", type=int, default=50, help="Top-k sampling")
    parser.add_argument("--gpu-layers", type=int, default=-1, help="GPU layers to offload (-1 for all)")
    parser.add_argument("--threads", type=int, default=8, help="CPU threads")
    parser.add_argument("--verbose", action="store_true", default=True, help="Verbose output")

    args = parser.parse_args()

    run_generation(
        model_path=args.model,
        prompt=args.prompt,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        gpu_layers=args.gpu_layers,
        threads=args.threads,
        verbose=args.verbose,
    )


if __name__ == "__main__":
    main()