import gc, torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
import re, json
import pandas as pd
from collections import Counter

# Free a model already on the GPU, so re-running this cell doesn't stack a second copy and run out of
# memory. (If you still hit "out of memory", do Runtime -> Restart session and run this cell just once.)
model = tok = None
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

MODEL_ID = "."

MAX_NEW_TOKENS = 1500

generation_config = dict(
    do_sample=True,
    temperature=0.3,
    top_p=0.9,
    top_k=40,
    repetition_penalty=1.05,
)

tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto",
).eval()
print("loaded", MODEL_ID, "| VRAM", round(torch.cuda.max_memory_allocated() / 1e9, 1), "GB")

df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")

SOLVER_SYSTEM = """
You are an expert International Linguistics Olympiad (IOL) solver.
Your task is to infer the hidden linguistic rules ONLY from the provided examples.
REQUIRED WORKFLOW:
1. SCRATCHPAD: List all recurring units (morphemes, words, or sounds) and their meanings.
2. RULE VERIFICATION: Write down rules for combining these units. Test against examples.
3. FINAL DERIVATION: Step-by-step derivation for each query item.
IMPORTANT RULES:
- Never rely on outside linguistic knowledge.
- If the task is 'match_letters', the FINAL ANSWERS must be ONLY the letter (e.g., A, B, C) that corresponds to each query item, one per line. Do NOT output the word itself.
- If the target language is phonetic (uses brackets [] or special symbols), keep that notation exactly.
- ABSOLUTELY NO ENGLISH in the FINAL ANSWERS section unless the target language is English.
- Output exactly one answer per query item.
Output format:
FINAL ANSWERS:
[Answer 1]
[Answer 2]
... (one per line)
"""

VALIDATOR_SYSTEM = """
You are an expert IOL solution validator.
Check if the FINAL ANSWERS match the expected format of the query:
- For 'match_letters', are they ONLY single letters (A, B, C...)? If they are words, it is INVALID.
- For 'translation' or 'fill_blanks', are they in the target language (not English)? If there is English, it is INVALID.
- Is the number of answers correct?
Output ONLY 'VALID' or 'INVALID' followed by specific contradictions.
"""

CORRECTOR_SYSTEM = """
You are correcting an IOL solution.
- If the task is 'match_letters', replace words with the corresponding labels (A, B, C).
- Ensure the FINAL ANSWERS contain ONLY the target language forms.
- Remove all English translations, explanations, or labels like 'Item 1:'.
Output exactly:
FINAL ANSWERS:
followed by one answer per line.
"""

def parse_answers(text):
    m = list(re.finditer(r"(?im)^\s*FINAL ANSWERS\s*:?\s*$", text))
    if m:
        text = text[m[-1].end():]

    answers = []
    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue

        # Clean formatting
        line = re.sub(r"^\d+[.)]\s*", "", line)
        line = re.sub(r"^[-*•]\s*", "", line)
        line = line.strip("'\" ")

        # Heuristic to filter out leaked reasoning:
        # 1. Skip lines that contain markdown bolding or italics
        if '*' in line or '_' in line:
            continue
        # 2. Skip lines that look like full sentences (too many spaces)
        # unless it's a translation task where the target is a sentence.
        if line.count(' ') > 5 and len(line) > 50:
            continue

        answers.append(line)

    return answers

def constraint_check(problem, answers):

    errors = []

    if len(answers) == 0:
        errors.append("No answers generated.")

    # No empty answers

    for i,a in enumerate(answers):

        if len(a.strip()) == 0:
            errors.append(f"Answer {i+1} is empty.")

    # Remove duplicate consecutive answers

    for i in range(1,len(answers)):

        if answers[i] == answers[i-1]:
            errors.append("Duplicate consecutive answers.")

    # Very long outputs

    for a in answers:

        if len(a) > 120:
            errors.append("Answer too long.")

    return errors

def generate(system_prompt, user_prompt):

    messages = [
        {"role":"system","content":system_prompt},
        {"role":"user","content":user_prompt},
    ]

    text = tok.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tok(
        text,
        return_tensors="pt"
    ).to(model.device)

    outputs = model.generate(

        **inputs,

        max_new_tokens=MAX_NEW_TOKENS,

        do_sample=True,
        # Lowered temperature for more stable linguistic reasoning
        temperature=0.1,
        top_p=0.95,
        top_k=50,

        repetition_penalty=1.05,
    )

    output = tok.decode(
        outputs[0][inputs.input_ids.shape[1]:],
        skip_special_tokens=True
    )

    return output

def self_consistency(problem, n=3):
    candidates_with_raw = [] # Store (parsed_answers_tuple, raw_text_string)

    for _ in range(n):
        raw_out_text = generate(
            SOLVER_SYSTEM,
            problem
        )
        parsed_ans = tuple(parse_answers(raw_out_text))
        candidates_with_raw.append((parsed_ans, raw_out_text))

    # Count occurrences of parsed answers
    parsed_ans_counts = Counter(item[0] for item in candidates_with_raw)
    best_parsed_ans = parsed_ans_counts.most_common(1)[0][0]

    # Find the raw text that produced the best_parsed_ans (take the first one if multiple)
    best_raw_text = None
    for parsed_ans, raw_text_candidate in candidates_with_raw:
        if parsed_ans == best_parsed_ans:
            best_raw_text = raw_text_candidate
            break

    return list(best_parsed_ans), best_raw_text

def validate(problem, answers):

    prompt = f"""
Problem
{problem}
Candidate solution
FINAL ANSWERS:
{chr(10).join(answers)}
"""

    result = generate(
        VALIDATOR_SYSTEM,
        prompt
    )

    return result

def correct(problem, answers, validator_output):

    prompt = f"""
Problem
{problem}
Previous answer
FINAL ANSWERS:
{chr(10).join(answers)}
Validation
{validator_output}
"""

    result = generate(
        CORRECTOR_SYSTEM,
        prompt
    )

    return parse_answers(result)

def solve(problem):

    answers, raw_output_for_best_ans = self_consistency(
        problem,
        n=1
    )
    raw_text = raw_output_for_best_ans

    for _ in range(1):

        validator = validate(
            problem,
            answers
        )

        constraints = constraint_check(
            problem,
            answers
        )

        if validator.strip() == "VALID" and len(constraints) == 0:

            return answers, raw_text

        answers = correct(
            problem,
            answers,
            validator + "\n" + "\n".join(constraints)
        )

    return answers, raw_text
results = []
for i, r in df.iterrows():
    print(f"\n{'=' * 72}\nPROBLEM {i + 1}/{len(df)} -- {r['task_type']}\n{'=' * 72}", flush=True)
    problem_text = f"{r['context'].strip()}\n\n{r['query'].strip()}"
    answers, raw_text = solve(problem_text)
    results.append({
        "id": r["id"],
        "query": r["query"],
        "raw": raw_text,
        "pred": answers,
    })
    print(f"\n--> parsed {len(answers)} answers", flush=True)
SUMMARIZE = (
    "Summarize the following reasoning into a few short bullet points: the rule or pattern found "
    "in the data and the key evidence for the answer. Be concise and structured -- do not repeat "
    "the full reasoning."
)


submission = pd.DataFrame([
    {
        "id": res["id"],
        "pred": json.dumps(res["pred"], ensure_ascii=False)
    }
    for res in results
])

submission.to_csv("submission.csv", index=False)

print(submission.head())
print(f"Saved {len(submission)} predictions to submission.csv")