import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import csv
import json
import re
import time
import warnings
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

warnings.filterwarnings("ignore")

MODEL_PATH = "."
TIME_LIMIT_SECONDS = 25 * 60  # Hard stop with 5 min buffer

# NF4 quantization for T4 16 GB
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

print("Loading tokenizer...", flush=True)
tok = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

print("Loading model...", flush=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True,
).eval()
print("Model ready.", flush=True)

SYSTEM_PROMPT = (
    "You are an expert computational linguist competing in the International Linguistics Olympiad (IOL). "
    "You solve linguistic puzzles by analyzing patterns in unfamiliar languages.\n\n"
    "REASONING FORMAT:\n"
    "First, provide your step-by-step linguistic analysis inside <think>...</think> tags. "
    "Structure this analysis clearly with bullet points or short paragraphs covering:\n"
    "- Key patterns observed (morphology, syntax, phonology, number system)\n"
    "- Rules inferred from the context examples\n"
    "- How you apply those rules to each query item\n"
    "Keep the analysis concise and readable by a human jury in 2-3 minutes.\n\n"
    "ANSWER FORMAT:\n"
    "After </think>, provide EXACTLY one answer per line, with NO numbering, NO bullet points, "
    "NO labels, and NO extra commentary. The number of answer lines must exactly match the number "
    "of numbered items in the query, in the same order.\n\n"
    "TASK RULES:\n"
    "- translation: output the full translated text\n"
    "- match_letters: output ONLY the matching capital letter (e.g., A, B, C)\n"
    "- fill_blanks: output ONLY the missing linguistic form\n"
    "- text_to_num: output digits only (e.g., 42, 111)\n"
    "- num_to_text: output the number written in the task language\n\n"
    "CRITICAL: Never skip an item. Always provide your best guess."
)

def count_query_items(query_text):
    """Count how many numbered items the query contains."""
    # Matches: 1. ..., 1) ..., (1) ..., etc., at line start
    pattern = r'(?:^|\n)\s*[(]?\s*\d+\s*[.)]\s+'
    matches = re.findall(pattern, '\n' + query_text)
    if matches:
        return len(matches)
    # Fallback: non-empty lines minus a likely header line
    lines = [l for l in query_text.split('\n') if l.strip()]
    return max(1, len(lines) - 1) if len(lines) > 1 else 1

def extract_answers(text, expected_count):
    """Parse final answers from generated text (think blocks removed)."""
    cleaned_text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
    if not cleaned_text:
        cleaned_text = text.strip()

    lines = [ln.strip() for ln in cleaned_text.split('\n') if ln.strip()]
    answers = []
    for line in lines:
        # Strip leading numbering and common prefixes
        line = re.sub(r'^\s*\d+[\.)]\s+', '', line)
        line = re.sub(r'^(?:Answer|ANSWER)\s*[:.\-]?\s*', '', line, flags=re.IGNORECASE)
        line = line.strip().strip('"').strip("'").strip('`')
        if line:
            answers.append(line)

    # Fallback: comma / semicolon separated values
    if len(answers) != expected_count and expected_count > 1:
        combined = ' '.join(lines)
        for sep in [';', ',']:
            if sep in combined:
                parts = [p.strip().strip('"').strip("'").strip('`') for p in combined.split(sep) if p.strip()]
                if len(parts) == expected_count:
                    answers = parts
                    break

    # Guarantee exact expected count
    if len(answers) < expected_count:
        while len(answers) < expected_count:
            answers.append(answers[-1] if answers else "")
    elif len(answers) > expected_count:
        answers = answers[:expected_count]

    return answers

def extract_explanation(text):
    """Pull reasoning from <think> block for the human-eval column."""
    m = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
    if m:
        exp = m.group(1).strip()
        exp = re.sub(r'\n{3,}', '\n\n', exp)
        if len(exp) > 2000:
            exp = exp[:2000] + "..."
        return exp
    return ""

# ------------------------------------------------------------------
# Main inference loop
# ------------------------------------------------------------------
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
expected_counts = [count_query_items(row["query"]) for _, row in df.iterrows()]

results = []
start_time = time.time()

for idx, row in df.iterrows():
    elapsed = time.time() - start_time
    if elapsed > TIME_LIMIT_SECONDS:
        print(f"Time buffer reached. Flushing remaining {len(df) - idx} rows as blanks.", flush=True)
        for j in range(idx, len(df)):
            results.append({
                "id": df.iloc[j]["id"],
                "pred": json.dumps([""] * expected_counts[j], ensure_ascii=False),
                "explanation": ""
            })
        break

    context = row["context"].strip()
    query = row["query"].strip()
    task_type = row.get("task_type", "")
    expected = expected_counts[idx]

    user_msg = (
        f"TASK TYPE: {task_type}\n\n"
        f"CONTEXT:\n{context}\n\n"
        f"QUERY:\n{query}\n\n"
        f"Provide your analysis in <think>...</think> tags, then list exactly {expected} answer(s) "
        f"below it, one per line, with no numbering and no extra text."
    )

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_msg},
    ]

    try:
        inputs = tok.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to("cuda")

        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=2048,
                do_sample=False,
                pad_token_id=tok.pad_token_id,
            )

        gen_text = tok.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
        answers = extract_answers(gen_text, expected)
        explanation = extract_explanation(gen_text)

    except Exception as e:
        print(f"ERROR on {row['id']}: {e}", flush=True)
        answers = [""] * expected
        explanation = ""

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

    print(f"[{idx+1}/{len(df)}] {row['id']} | items={expected} | answers={answers}", flush=True)

# ------------------------------------------------------------------
# Write submission
# ------------------------------------------------------------------
with open("submission.csv", "w", newline="", encoding="utf-8") as file_handle:
    writer = csv.DictWriter(file_handle, fieldnames=["id", "pred", "explanation"])
    writer.writeheader()
    for r in results:
        writer.writerow(r)

print("Wrote submission.csv", flush=True)