import json
import hashlib
import os
import sys
import torch
from tokenizers import Tokenizer

sys.path.append("/data/foundational_model")
from model.model import ModelConfig, Transformer

phrases = [
    "namaste, aap kaise hain?",
    "ramayana kisne likhi hai?",
    "mahabharata me arjun kaun the?",
    "kya kar rahe ho?",
    "shri ram ke pita ka naam kya tha?",
    "gita me bhagwan krishna ne kya kaha?",
    "aaj ka mausam kaisa hai?",
    "mujhe ek achhi kahani sunao.",
    "ravan ka rajya kahan tha?",
    "pandav kitne bhai the?",
    "dharma ka kya matlab hai?",
    "kaise ho bhai?",
    "life me khush kaise rahe?",
    "hanuman ji ke bare me batao.",
    "karna kaun tha?",
    "aaj ka din kaisa raha?",
    "sita ji ka janma kahan hua tha?",
    "bheeshma pitamah kaun the?",
    "ek accha vichar bataiye.",
    "alvida, phir milenge!"
]

def sample_next_token(logits, temperature=0.7, top_p=0.9, generated_ids=[], no_repeat_ngram_size=2):
    logits = logits.clone()

    # N-gram blocking to eliminate token repetition
    if len(generated_ids) >= no_repeat_ngram_size:
        prev_ngram = tuple(generated_ids[-(no_repeat_ngram_size - 1):])
        for i in range(len(generated_ids) - no_repeat_ngram_size + 1):
            if tuple(generated_ids[i:i + no_repeat_ngram_size - 1]) == prev_ngram:
                forbidden_token = generated_ids[i + no_repeat_ngram_size - 1]
                logits[forbidden_token] = -float("Inf")

    # Apply temperature
    logits = logits / temperature
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)

    sorted_indices_to_remove = cumulative_probs > top_p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    indices_to_remove = sorted_indices[sorted_indices_to_remove]
    logits[indices_to_remove] = -float("Inf")

    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1).item()

def generate_direct():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer_path = "/data/foundational_model/tokenizer/hinglish_tokenizer.json"
    checkpoint_path = "/data/foundational_model/checkpoints/model_125m_final.pt"

    print(f"Loading PyTorch scratch model on {device}...")
    tokenizer = Tokenizer.from_file(tokenizer_path)

    config = ModelConfig(
        vocab_size=16384,
        dim=256,
        n_layers=6,
        n_heads=8,
        n_kv_heads=8,
        max_seq_len=128
    )

    model = Transformer(config).to(device)
    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
    model.eval()

    results = []
    response_texts = []

    print(f"Evaluating {len(phrases)} phrases on Scratch Model with 2-Gram Blocking...\n")

    torch.manual_seed(42)

    for idx, phrase in enumerate(phrases, 1):
        prompt_tokens = [4] + tokenizer.encode(f" {phrase}").ids + [5]
        generated = prompt_tokens[:]

        for _ in range(30):
            context = generated[-128:]
            logits = model(torch.tensor([context], dtype=torch.long, device=device))[0, -1, :]
            gen_ids = generated[len(prompt_tokens):]
            next_tok = sample_next_token(logits, temperature=0.7, top_p=0.9, generated_ids=gen_ids, no_repeat_ngram_size=2)
            if next_tok in (2, 3):  # </s> or <pad>
                break
            generated.append(next_tok)

        resp_text = tokenizer.decode(generated[len(prompt_tokens):]).strip()

        print(f"[{idx:02d}/{len(phrases)}] Prompt: '{phrase}'")
        print(f"     Response: '{resp_text}'\n")

        results.append({
            "index": idx,
            "prompt": phrase,
            "response": resp_text
        })
        response_texts.append(resp_text)

    combined = "\n---RESPONSE_SEP---\n".join(response_texts)
    md5_hash = hashlib.md5(combined.encode("utf-8")).hexdigest()

    print("=" * 60)
    print(f"MD5 Checksum of Scratch Model 20 Responses: {md5_hash}")
    print("=" * 60)

    out_file = "/data/foundational_model/eval/evaluation_results.json"
    with open(out_file, "w", encoding="utf-8") as f:
        json.dump({
            "md5": md5_hash,
            "results": results
        }, f, indent=2, ensure_ascii=False)
    print(f"Saved results to {out_file}")

if __name__ == "__main__":
    generate_direct()