#!/usr/bin/env python3
"""
Generate reasoning traces for multihop temporal reasoning QA samples.

For each sample in the multihop task CSVs:
  1. Load question, task, question_type, answer, and source_categories
  2. Compute a deterministic symbolic trace from trace_templates.json
  3. Pass ONLY the symbolic trace + question + answer to Llama-3.1-8B-Instruct
  4. Llama verbalizes (does NOT solve) the trace
  5. Validate output
  6. Append symbolic_trace and verbal_trace columns to all 3 CSVs

Usage:
    python generate_reasoning_traces.py \
        --dataset_dir /home/debarpanb1/TREA_2.0/pipeline/dataset_v5 \
        --trace_templates /home/debarpanb1/TREA_2.0/pipeline/trace_templates.json \
        [--tasks conditional_count conditional_duration ...] \
        [--batch_size 8] [--dry_run]
"""

import argparse
import ast
import json
import os
import re
import sys
from pathlib import Path

import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


# ── Tasks that require reasoning traces ──────────────────────────────────────
MULTIHOP_TASKS = [
    "conditional_count",
    "conditional_duration",
    "between_events",
    "event_density",
    "duration_gap",
    "temporal_arithmetic",
    "temporal_loudness",
    "multi_hop",
]

# ── System prompt for Llama verbalization ────────────────────────────────────
SYSTEM_PROMPT = (
    "You are a reasoning-trace verbalizer for an audio temporal reasoning dataset.\n\n"
    "Your job is only to convert the provided symbolic trace into a short natural-language explanation.\n\n"
    "Rules:\n"
    "- Do not solve the question yourself.\n"
    "- Do not add any event, sound, time, duration, count, or comparison not present in the input.\n"
    "- Do not change the answer.\n"
    "- Use only the provided symbolic trace and answer.\n"
    "- Keep the explanation 2 to 4 sentences.\n"
    '- End with: "Therefore, the answer is <answer>."\n'
    "- Output only the trace text. No JSON. No extra commentary."
)

USER_TEMPLATE = (
    "Question: {question}\n\n"
    "Answer: {answer}\n\n"
    "Symbolic trace:\n{symbolic_trace}\n\n"
    "Allowed sound labels:\n{allowed_labels}\n\n"
    "Write a short grounded reasoning trace."
)


# ═══════════════════════════════════════════════════════════════════════════════
#  Helpers: parse categories / questions to extract placeholder values
# ═══════════════════════════════════════════════════════════════════════════════

def safe_parse_list(val):
    """Parse a stringified Python list from CSV."""
    if isinstance(val, list):
        return val
    if pd.isna(val):
        return []
    try:
        return ast.literal_eval(val)
    except Exception:
        return [x.strip().strip("'\"") for x in val.strip("[]").split(",") if x.strip()]


def extract_placeholder(question: str, template_pattern: str):
    placeholder_names_all = re.findall(r"\{(\w+)\}", template_pattern)
    if not placeholder_names_all:
        return {}

    patterns_to_try = [template_pattern]
    
    if "{target_sound}" in template_pattern:
        # If target sound is completely omitted, e.g., "How many sounds..." instead of "How many {target_sound} sounds..."
        patterns_to_try.append(template_pattern.replace("{target_sound} sounds", "sounds"))
        patterns_to_try.append(template_pattern.replace(" {target_sound} ", " "))
        patterns_to_try.append(template_pattern.replace("{target_sound} ", ""))
        patterns_to_try.append(template_pattern.replace(" {target_sound}", ""))

    for pat in patterns_to_try:
        placeholder_names = re.findall(r"\{(\w+)\}", pat)
        regex = re.escape(pat)
        seen = set()
        for name in placeholder_names:
            token = re.escape("{" + name + "}")
            if name not in seen:
                regex = regex.replace(token, f"(?P<{name}>.+?)", 1)
                seen.add(name)
            else:
                regex = regex.replace(token, f"(?P={name})", 1)
        regex = "^" + regex + "$"
        m = re.match(regex, question)
        if m:
            return m.groupdict()
            
    return {}

def try_extract_placeholders(question: str, templates: dict, question_type: str):
    """Try to extract placeholders from question using config templates."""
    # Try the specific question_type templates first
    candidates = []
    if question_type in templates:
        t = templates[question_type]
        if isinstance(t, list):
            candidates.extend(t)
        else:
            candidates.append(t)

    for tmpl in candidates:
        result = extract_placeholder(question, tmpl)
        if result:
            return result
    return {}


# ═══════════════════════════════════════════════════════════════════════════════
#  Symbolic trace computation
# ═══════════════════════════════════════════════════════════════════════════════

def compute_symbolic_trace(
    task: str,
    question_type: str,
    question: str,
    answer,
    categories: list,
    trace_templates: dict,
    config_templates: dict,
) -> list:
    """Compute a filled symbolic trace for one QA sample.

    Uses the trace_templates.json skeleton and fills placeholders from
    the question text + categories list.
    """
    # Get template steps
    task_traces = trace_templates.get(task, {})
    template_steps = task_traces.get(question_type)
    if not template_steps:
        return [f"Answer the question. The answer is {answer}."]

    # ── Extract placeholders from question text ──
    mcq_templates = config_templates.get(task, {}).get("mcq_questions", {})
    open_templates = config_templates.get(task, {}).get("open_text_questions", {})

    placeholders = try_extract_placeholders(question, open_templates, question_type)
    if not placeholders:
        placeholders = try_extract_placeholders(question, mcq_templates, question_type)

    # ── Clean underscores from placeholders and question ──
    for k, v in placeholders.items():
        if isinstance(v, str):
            placeholders[k] = v.replace("_", " ")

    categories = [str(c).replace("_", " ") for c in categories]

    # ── Derive additional placeholders deterministically from categories ──
    answer_str = str(answer).replace("_", " ")
    placeholders["answer"] = answer_str

    # selected_events: events in the relevant region
    if "selected_events" not in placeholders:
        placeholders["selected_events"] = ", ".join(categories) if categories else "none"

    # For multi_hop: derive pivot-based placeholders
    if task == "multi_hop":
        _fill_multi_hop_placeholders(placeholders, question_type, categories, answer_str)

    # ── Fill template ──
    filled = []
    for step in template_steps:
        try:
            filled_step = step.format(**placeholders)
        except KeyError:
            # Fill what we can, leave unknowns as-is
            filled_step = step.format_map(_SafeDict(placeholders))
            
        # Clean up missing target_sound literal if it wasn't extracted
        if "target_sound" not in placeholders:
            filled_step = filled_step.replace("{target_sound} sounds ", "sounds ")
            filled_step = filled_step.replace("{target_sound} events ", "events ")
            filled_step = filled_step.replace(" {target_sound} ", " ")
            filled_step = filled_step.replace("{target_sound} ", "")
            filled_step = filled_step.replace("{target_sound}", "")
            # Ensure we don't leave double spaces except after period if any
            filled_step = re.sub(r'\s+', ' ', filled_step).strip()
            
        filled.append(filled_step)
    return filled


class _SafeDict(dict):
    """Dict that returns '{key}' for missing keys in str.format_map."""
    def __missing__(self, key):
        return "{" + key + "}"


def _fill_multi_hop_placeholders(placeholders, question_type, categories, answer_str):
    """Fill derived placeholders specific to multi_hop task."""
    if question_type in ("after_longest", "before_longest", "count_after_longest"):
        # We don't know which is longest from CSV alone; use a generic label
        if "longest_sound" not in placeholders:
            placeholders["longest_sound"] = "the longest event"
    if question_type == "after_shortest":
        if "shortest_sound" not in placeholders:
            placeholders["shortest_sound"] = "the shortest event"
    if question_type == "before_loudest" or question_type == "count_before_loudest":
        if "loudest_sound" not in placeholders:
            placeholders["loudest_sound"] = "the loudest event"
    if question_type == "after_longest_gap":
        if "longest_gap_before_sound" not in placeholders:
            placeholders["longest_gap_before_sound"] = "the event before the longest silence"
        if "longest_gap_after_sound" not in placeholders:
            placeholders["longest_gap_after_sound"] = answer_str
    if question_type == "overlap_after_anchor":
        if "event_after_anchor" not in placeholders:
            placeholders["event_after_anchor"] = "the event after the anchor"


# ═══════════════════════════════════════════════════════════════════════════════
#  LLM verbalization
# ═══════════════════════════════════════════════════════════════════════════════

def verbalize_trace(
    tokenizer, model, device,
    question: str, answer: str,
    symbolic_trace: list, allowed_labels: list,
) -> str:
    """Use Llama-3.1-8B-Instruct to verbalize a symbolic trace."""
    trace_text = "\n".join(f"- {s}" for s in symbolic_trace)
    labels_text = ", ".join(sorted(set(allowed_labels)))

    user_msg = USER_TEMPLATE.format(
        question=question,
        answer=answer,
        symbolic_trace=trace_text,
        allowed_labels=labels_text,
    )

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

    input_len = inputs.shape[1]
    with torch.no_grad():
        output = model.generate(
            inputs,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    generated = output[0, input_len:]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


def validate_verbal_trace(verbal_trace: str, answer: str) -> bool:
    """Basic validation: trace should mention the answer and end properly."""
    answer_norm = str(answer).replace("_", " ").lower()
    trace_lower = verbal_trace.lower()
    # Check answer appears somewhere
    if answer_norm not in trace_lower:
        return False
    # Check it ends with the canonical closing (loosely)
    if "therefore" not in trace_lower and "the answer is" not in trace_lower:
        return False
    # Not too short / too long
    if len(verbal_trace.split()) < 8 or len(verbal_trace.split()) > 120:
        return False
    return True


# ═══════════════════════════════════════════════════════════════════════════════
#  Main pipeline
# ═══════════════════════════════════════════════════════════════════════════════

def load_config_templates(config_path: str) -> dict:
    """Load question templates from config.yaml keyed by task name."""
    import yaml
    with open(config_path) as f:
        config = yaml.safe_load(f)
    templates = {}
    for task_name, task_cfg in config.get("tasks", {}).items():
        templates[task_name] = {
            "mcq_questions": task_cfg.get("mcq_questions", {}),
            "open_text_questions": task_cfg.get("open_text_questions", {}),
        }
    return templates


def process_task(
    task: str,
    dataset_dir: Path,
    trace_templates: dict,
    config_templates: dict,
    tokenizer, model, device,
    dry_run: bool = False,
    max_retries: int = 2,
):
    """Process a single task: compute traces and add columns to CSVs."""
    task_dir = dataset_dir / task
    if not task_dir.exists():
        print(f"  [SKIP] {task}: directory not found")
        return

    # Identify CSV files
    mcq_csv = task_dir / f"{task}_mcq.csv"
    open_csv = task_dir / f"{task}_open_text.csv"
    meta_csv = task_dir / f"{task}_metadata.csv"

    # We compute traces from the open_text CSV (has question, answer, question_type, source_categories)
    if not open_csv.exists():
        print(f"  [SKIP] {task}: open_text CSV not found")
        return

    df_open = pd.read_csv(open_csv)
    df_mcq = pd.read_csv(mcq_csv) if mcq_csv.exists() else None
    df_meta = pd.read_csv(meta_csv) if meta_csv.exists() else None

    print(f"  Processing {task}: {len(df_open)} samples")

    symbolic_traces = []
    verbal_traces = []

    for idx, row in df_open.iterrows():
        question = str(row["question"])
        answer = str(row["answer"]).replace("_", " ")
        question_type = str(row["question_type"])
        categories = safe_parse_list(row.get("source_categories", "[]"))

        # 1. Compute symbolic trace
        sym_trace = compute_symbolic_trace(
            task, question_type, question, answer,
            categories, trace_templates, config_templates,
        )
        symbolic_traces.append(json.dumps(sym_trace))

        # 2. Verbalize with LLM
        if dry_run:
            verbal = "[DRY RUN] " + " ".join(sym_trace)
            verbal_traces.append(verbal)
        else:
            verbal = ""
            clean_question = question.replace("_", " ")
            clean_answer = answer.replace("_", " ")
            clean_categories = [str(c).replace("_", " ") for c in categories]
            for attempt in range(max_retries + 1):
                verbal = verbalize_trace(
                    tokenizer, model, device,
                    clean_question, clean_answer, sym_trace, clean_categories,
                )
                if validate_verbal_trace(verbal, clean_answer):
                    break
                if attempt < max_retries:
                    print(f"    [RETRY] sample {row['id']} attempt {attempt+1}")
            verbal_traces.append(verbal)

        if (idx + 1) % 10 == 0:
            print(f"    {idx+1}/{len(df_open)} done")

    # 3. Add columns to open_text CSV
    df_open["symbolic_trace"] = symbolic_traces
    df_open["verbal_trace"] = verbal_traces
    df_open.to_csv(open_csv, index=False)
    print(f"    Saved {open_csv}")

    # 4. Add columns to MCQ CSV (join on id)
    if df_mcq is not None:
        trace_map = df_open.set_index("id")[["symbolic_trace", "verbal_trace"]]
        df_mcq = df_mcq.merge(trace_map, left_on="id", right_index=True, how="left")
        df_mcq.to_csv(mcq_csv, index=False)
        print(f"    Saved {mcq_csv}")

    # 5. Add columns to metadata CSV (join on id)
    if df_meta is not None:
        trace_map = df_open.set_index("id")[["symbolic_trace", "verbal_trace"]]
        df_meta = df_meta.merge(trace_map, left_on="id", right_index=True, how="left")
        df_meta.to_csv(meta_csv, index=False)
        print(f"    Saved {meta_csv}")


def main():
    parser = argparse.ArgumentParser(
        description="Generate reasoning traces for multihop temporal reasoning QA samples"
    )
    parser.add_argument(
        "--dataset_dir", type=str,
        default="/home/debarpanb1/TREA_2.0/pipeline/dataset_v5",
        help="Path to dataset directory",
    )
    parser.add_argument(
        "--trace_templates", type=str,
        default="/home/debarpanb1/TREA_2.0/pipeline/trace_templates.json",
        help="Path to trace_templates.json",
    )
    parser.add_argument(
        "--config", type=str,
        default="/home/debarpanb1/TREA_2.0/pipeline/config.yaml",
        help="Path to config.yaml (for question templates)",
    )
    parser.add_argument(
        "--tasks", nargs="+", default=None,
        help=f"Tasks to process (default: all multihop). Options: {MULTIHOP_TASKS}",
    )
    parser.add_argument("--dry_run", action="store_true", help="Skip LLM, use raw symbolic trace")
    parser.add_argument("--max_retries", type=int, default=2, help="Max retries on validation failure")
    args = parser.parse_args()

    dataset_dir = Path(args.dataset_dir)
    tasks = args.tasks or MULTIHOP_TASKS

    # Load trace templates
    with open(args.trace_templates) as f:
        trace_templates = json.load(f)

    # Load config question templates
    config_templates = load_config_templates(args.config)

    # Load model (unless dry run)
    tokenizer, model, device = None, None, None
    if not args.dry_run:
        print("Loading meta-llama/Llama-3.1-8B-Instruct...")
        tokenizer = AutoTokenizer.from_pretrained(
            "meta-llama/Llama-3.1-8B-Instruct", use_fast=False
        )
        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.1-8B-Instruct",
            torch_dtype="auto",
            device_map="auto",
        )
        model.eval()
        device = next(model.parameters()).device
        print(f"Model loaded on {device}")

    # Process each task
    for task in tasks:
        print(f"\n{'='*60}")
        print(f"Task: {task}")
        print(f"{'='*60}")
        process_task(
            task, dataset_dir, trace_templates, config_templates,
            tokenizer, model, device,
            dry_run=args.dry_run,
            max_retries=args.max_retries,
        )

    print("\n✓ All done!")


if __name__ == "__main__":
    main()