#!/usr/bin/env python3
import os
import sys
import json
import re
import csv
import subprocess
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ------------------------------------------------------------------
# 1. Install auto-gptq (PyPI is reachable at runtime)
# ------------------------------------------------------------------
try:
    import auto_gptq
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "auto-gptq"])
    import auto_gptq

# ------------------------------------------------------------------
# 2. Offline mode - no internet allowed during evaluation
# ------------------------------------------------------------------
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

# ------------------------------------------------------------------
# 3. Load model from local directory (the repo root)
# ------------------------------------------------------------------
MODEL_DIR = "."
print("Loading tokenizer and model...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_DIR,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.float16,
)
model.eval()
print("Model loaded.", flush=True)

# ------------------------------------------------------------------
# 4. Helper: count numbered items in query
# ------------------------------------------------------------------
def count_items(query: str) -> int:
    numbers = re.findall(r'(?:^|\n)\s*(?:\d+)\.', query)
    if not numbers:
        lines = query.strip().splitlines()
        count = 0
        for ln in lines:
            if re.match(r'^\s*\d+[\.\)]', ln):
                count += 1
        return count
    return len(numbers)

# ------------------------------------------------------------------
# 5. Parse model output to a JSON list
# ------------------------------------------------------------------
def extract_answers(text: str, expected_count: int) -> list:
    json_match = re.search(r'\[\s*".*?"\s*(?:,\s*".*?"\s*)*\]', text, re.DOTALL)
    if json_match:
        try:
            ans = json.loads(json_match.group())
            if isinstance(ans, list) and all(isinstance(x, str) for x in ans):
                return ans
        except:
            pass

    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    clean = []
    for ln in lines:
        m = re.match(r'^\s*\d+[\.\)]\s*(.*)', ln)
        if m:
            clean.append(m.group(1).strip())
        else:
            clean.append(ln)
    if len(clean) > expected_count:
        clean = clean[:expected_count]
    elif len(clean) < expected_count:
        clean += [""] * (expected_count - len(clean))
    return clean

# ------------------------------------------------------------------
# 6. Optional explanation generation (set to True to opt in)
# ------------------------------------------------------------------
EXPLAIN = False   # change to True if you want to include explanations

def generate_explanation(context: str, query: str, answer_list: list) -> str:
    prompt = (
        f"{context}\n\n{query}\n\n"
        "My answers are: " + ", ".join(answer_list) + "\n\n"
        "Now, provide a concise explanation (max 3 bullet points) of the reasoning behind these answers. "
        "Do not repeat the answers, only the reasoning."
    )
    messages = [
        {"role": "system", "content": "You are an expert linguist. Provide a short, human-readable explanation of your reasoning for the given linguistic problem."},
        {"role": "user", "content": prompt},
    ]
    ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(ids, max_new_tokens=300, do_sample=False)
    expl = tokenizer.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
    return expl

# ------------------------------------------------------------------
# 7. Main loop over test.csv
# ------------------------------------------------------------------
TEST_CSV = "/tmp/data/test.csv"
df = pd.read_csv(TEST_CSV, dtype=str).fillna("")

rows = []
for idx, r in df.iterrows():
    context = r["context"].strip()
    query = r["query"].strip()
    num_items = count_items(query)

    system_msg = (
        "You solve International Linguistics Olympiad problems. "
        "Answer every numbered item in the query. "
        "Output the answers as a JSON list of strings, e.g., [\"answer1\", "answer2\"]. "
        "Do not include anything else."
    )
    user_msg = f"{context}\n\n{query}"

    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_msg},
    ]
    ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(ids, max_new_tokens=512, do_sample=False)
    generated = tokenizer.decode(outputs[0][ids.shape[-1]:], skip_special_tokens=True)

    answers = extract_answers(generated, num_items)
    print(f"Problem {r['id']}: generated {len(answers)} answers for {num_items} items", flush=True)

    explanation = ""
    if EXPLAIN:
        explanation = generate_explanation(context, query, answers)

    rows.append({
        "id": r["id"],
        "pred": json.dumps(answers, ensure_ascii=False),
        "explanation": explanation,
    })

pd.DataFrame(rows).to_csv("submission.csv", index=False)
print("wrote submission.csv", flush=True)