"""Inference script for Diffusion-LM Riddle Solver (Hugging Face model).

Usage:
    python3 inference.py --riddle "i speak without a mouth what am i"
"""

import json
import sys
import argparse
import warnings


def load_model(model_dir: str = "."):
    """Load model weights and config from HF model directory."""
    import torch
    
    with open(f"{model_dir}/config.json") as f:
        config = json.load(f)
    
    with open(f"{model_dir}/vocab.json") as f:
        vocab = json.load(f)
    
    inv_vocab = {int(v): k for k, v in vocab.items()}
    
    # Add riddle_diffusion.py to Python path if running from HF directory
    sys.path.insert(0, model_dir)
    from riddle_diffusion import DiffusionRiddleModel
    from riddle_diffusion import get_schedule
    
    model = DiffusionRiddleModel(
        vocab_size=config["vocab_size"],
        d_model=config["d_model"],
        n_layers=config["n_layers"],
        d_ff=config["d_ff"],
        n_heads=config["n_heads"],
        a_len=config["a_len"],
        q_len=config["q_len"],
        T=config["T"],
    )
    
    state = torch.load(f"{model_dir}/model.safetensors", map_location="cpu",
                       weights_only=True)
    model.load_state_dict(state)
    model.eval()
    
    return model, config, vocab, inv_vocab


def predict(model, config, vocab, inv_vocab, riddle: str, k_samples: int = 10,
            device: str = "cpu"):
    """Run prediction on a single riddle."""
    import torch
    import torch.nn.functional as F
    
    model.to(device)
    
    # Tokenize
    tokens = [vocab.get(w, vocab.get("<UNK>", 1)) for w in riddle.lower().split()]
    if len(tokens) > config["q_len"]:
        tokens = tokens[:config["q_len"]]
    
    q_tokens = torch.tensor([tokens], device=device)
    
    # Diffusion schedule (sqrt power law)
    betas = torch.sqrt(torch.linspace(1e-4, 0.02, config["T"])).to(device)
    alphas = 1.0 - betas
    alpha_bars = torch.cumprod(alphas, dim=0)
    
    # Reverse diffusion
    x_t = torch.randn(k_samples, config["a_len"], config["d_model"], device=device)
    
    for t in reversed(range(config["T"])):
        t_tensor = torch.full((k_samples,), t, device=device, dtype=torch.long)
        pred_x0 = model(x_t, t_tensor, q_tokens)
        
        # Euclidean clamping
        logits = 2.0 * F.linear(pred_x0, model.emb.weight)
        logits = logits - model.emb.weight.square().sum(dim=-1).unsqueeze(0).unsqueeze(0)
        x0_tokens = logits.argmax(dim=-1)
        x0_emb = model.emb(x0_tokens)
        
        if t > 0:
            alpha_bar = alpha_bars[t]
            alpha_bar_prev = alpha_bars[t - 1]
            beta_tilde = betas[t] * (1 - alpha_bar_prev) / (1 - alpha_bar)
            noise = torch.randn_like(x_t)
            coef1 = torch.sqrt(alpha_bar_prev) * betas[t] / (1 - alpha_bar)
            coef2 = torch.sqrt(alpha_bar) * (1 - alpha_bar_prev) / (1 - alpha_bar)
            mu = coef1 * x0_emb + coef2 * x_t
            x_t = mu + torch.sqrt(beta_tilde) * noise
        else:
            x_t = x0_emb
    
    # Decode
    pred_tokens = []
    for b in range(k_samples):
        pred = ""
        for pos in range(config["a_len"]):
            tok_id = x0_tokens[b, pos].item()
            if tok_id == 0:
                break
            pred += inv_vocab.get(tok_id, "?") + " "
        pred_tokens.append(pred.strip())
    
    # Majority vote
    from collections import Counter
    counts = Counter(pred_tokens)
    winner = counts.most_common(1)[0][0]
    
    return winner, pred_tokens


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--riddle", required=True, help="Riddle text")
    parser.add_argument("--model-dir", default=".", help="Model directory")
    parser.add_argument("--device", default="cpu", help="Device (cpu, mps, cuda)")
    parser.add_argument("--k-samples", type=int, default=10)
    args = parser.parse_args()
    
    model, config, vocab, inv_vocab = load_model(args.model_dir)
    answer, candidates = predict(model, config, vocab, inv_vocab,
                                 args.riddle, args.k_samples, args.device)
    
    print(f"Riddle:  {args.riddle}")
    print(f"Answer:  {answer}")
    if len(set(candidates)) > 1:
        from collections import Counter
        counts = Counter(candidates)
        print(f"Candidates ({args.k_samples} samples):")
        for text, count in counts.most_common():
            print(f"  {text:<20} ({count} votes)")


if __name__ == "__main__":
    main()