src/reason_interventions.py

"""
Sentence edits for reason faithfulness tests (UI + 07_faithfulness.py).

Each reason type has one targeted edit; re-predict after edit to see if hardness drops.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

from syntax_complexity import find_hardest_span
from utils import REASON_ORDER


@dataclass
class InterventionResult:
    reason: str
    original_sentence: str
    edited_sentence: str
    edit_method: str  # "openai", "gemini", or "rule"


def _call_llm(prompt: str) -> tuple[str | None, str]:
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI

            client = OpenAI()
            resp = client.chat.completions.create(
                model=os.environ.get("FAITHFULNESS_LLM", "gpt-4o-mini"),
                messages=[{"role": "user", "content": prompt}],
                temperature=0,
                max_tokens=256,
            )
            return resp.choices[0].message.content.strip(), "openai"
        except Exception:
            pass
    if os.environ.get("GOOGLE_API_KEY"):
        try:
            import google.generativeai as genai

            genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
            model = genai.GenerativeModel(os.environ.get("FAITHFULNESS_GEMINI", "gemini-1.5-flash"))
            return model.generate_content(prompt).text.strip(), "gemini"
        except Exception:
            pass
    return None, "rule"


def _rule_synonym_swap(sentence: str, target_word: str) -> str:
    simple = {
        "frankincense": "incense",
        "astonishment": "surprise",
        "partiality": "bias",
        "dominion": "rule",
        "assemblies": "meetings",
        "ewe": "sheep",
        "scribe": "writer",
        "scrivener": "writer",
        "inflammation": "swelling",
        "treaty": "deal",
    }
    rep = simple.get(target_word.lower(), target_word)
    if target_word in sentence:
        return sentence.replace(target_word, rep, 1)
    return sentence


def edit_lexical_rarity(sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    prompt = (
        f"Replace ONLY the word '{target_word}' in this sentence with a common, simple synonym "
        f"of the same meaning and part of speech. Change nothing else.\n"
        f"Sentence: {sentence}\n"
        f"Output only the edited sentence."
    )
    method = "rule"
    edited = _rule_synonym_swap(sentence, target_word)
    if use_llm:
        llm_out, method = _call_llm(prompt)
        if llm_out:
            edited = llm_out
    return InterventionResult("Lexical Rarity", sentence, edited, method)


def edit_contextual_ambiguity(sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    prompt = (
        f"Add at most 3 words to this sentence to make the meaning of '{target_word}' clear. "
        f"Do NOT change or remove '{target_word}'.\n"
        f"Sentence: {sentence}\n"
        f"Output only the edited sentence."
    )
    method = "rule"
    edited = f"{sentence} (meaning: {target_word})"
    if use_llm:
        llm_out, method = _call_llm(prompt)
        if llm_out:
            edited = llm_out
    return InterventionResult("Contextual Ambiguity", sentence, edited, method)


def edit_syntactic_complexity(sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    span, start, end = find_hardest_span(sentence)
    method = "rule"
    simplified = span
    if use_llm and span.strip() != sentence.strip():
        prompt = (
            f"Simplify ONLY this phrase to plain English. Keep all key words including '{target_word}'.\n"
            f"Phrase: {span}\n"
            f"Output only the simplified phrase."
        )
        llm_out, method = _call_llm(prompt)
        if llm_out:
            simplified = llm_out
    edited = sentence[:start] + simplified + sentence[end:]
    return InterventionResult("Syntactic Complexity", sentence, edited, method)


EDIT_FNS = {
    "Lexical Rarity": edit_lexical_rarity,
    "Contextual Ambiguity": edit_contextual_ambiguity,
    "Syntactic Complexity": edit_syntactic_complexity,
}


def apply_intervention(reason: str, sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    if reason not in EDIT_FNS:
        raise ValueError(f"Unknown reason: {reason}. Choose from {REASON_ORDER}")
    return EDIT_FNS[reason](sentence, target_word, use_llm=use_llm)


def apply_all_interventions(sentence: str, target_word: str, use_llm: bool = True) -> list[InterventionResult]:
    return [fn(sentence, target_word, use_llm=use_llm) for fn in EDIT_FNS.values()]


@dataclass
class AlterComplexityResult:
    reason: str
    direction: str  # "increase" | "decrease"
    original_sentence: str
    original_target: str
    new_sentence: str
    new_target: str
    edit_method: str
    explanation: str


def _require_llm(prompt: str) -> tuple[str, str]:
    text, method = _call_llm(prompt)
    if not text or method == "rule":
        raise RuntimeError(
            "OpenAI edit failed. Set OPENAI_API_KEY on the backend (HF Space secrets)."
        )
    return text, method


def _parse_json_block(raw: str) -> dict:
    import json
    import re

    raw = raw.strip()
    if raw.startswith("```"):
        raw = re.sub(r"^```(?:json)?\s*", "", raw)
        raw = re.sub(r"\s*```$", "", raw)
    return json.loads(raw)


def alter_complexity(
    reason: str,
    sentence: str,
    target_word: str,
    direction: str,
) -> AlterComplexityResult:
    """User-guided edit: increase or decrease one complexity dimension via OpenAI."""
    if reason not in EDIT_FNS:
        raise ValueError(f"Unknown reason: {reason}. Choose from {REASON_ORDER}")
    if direction not in ("increase", "decrease"):
        raise ValueError("direction must be 'increase' or 'decrease'")

    if reason == "Lexical Rarity":
        if direction == "decrease":
            task = (
                f"Replace ONLY '{target_word}' with a MORE COMMON, simpler synonym "
                f"(same meaning and part of speech). Keep the rest of the sentence unchanged."
            )
        else:
            task = (
                f"Replace ONLY '{target_word}' with a RARER, more formal or technical synonym "
                f"(same meaning and part of speech). Keep the rest of the sentence unchanged."
            )
        prompt = (
            f"{task}\n"
            f"Sentence: {sentence}\n"
            f"Target word: {target_word}\n"
            f'Return ONLY valid JSON: {{"sentence": "full sentence with new word", '
            f'"target_word": "the new target word", "explanation": "one short sentence"}}'
        )
        raw, method = _require_llm(prompt)
        data = _parse_json_block(raw)
        new_sentence = str(data["sentence"]).strip()
        new_target = str(data["target_word"]).strip()
        explanation = str(data.get("explanation", "Lexical rarity adjusted."))
    elif reason == "Contextual Ambiguity":
        if direction == "decrease":
            task = (
                f"Rewrite the sentence so the meaning of '{target_word}' is CLEARER and less ambiguous. "
                f"You MUST keep the exact word '{target_word}' unchanged (same spelling)."
            )
        else:
            task = (
                f"Rewrite the sentence so '{target_word}' is MORE AMBIGUOUS in context "
                f"(multiple plausible meanings). You MUST keep the exact word '{target_word}' unchanged."
            )
        prompt = (
            f"{task}\n"
            f"Sentence: {sentence}\n"
            f'Return ONLY valid JSON: {{"sentence": "new sentence", '
            f'"explanation": "one short sentence"}}'
        )
        raw, method = _require_llm(prompt)
        data = _parse_json_block(raw)
        new_sentence = str(data["sentence"]).strip()
        new_target = target_word
        explanation = str(data.get("explanation", "Contextual ambiguity adjusted."))
    else:  # Syntactic Complexity
        if direction == "decrease":
            task = (
                f"Simplify the sentence structure to plain, short English. "
                f"You MUST keep the exact word '{target_word}' unchanged."
            )
        else:
            task = (
                f"Make the sentence syntactically MORE COMPLEX (subordinate clauses, passive, etc.). "
                f"You MUST keep the exact word '{target_word}' unchanged."
            )
        prompt = (
            f"{task}\n"
            f"Sentence: {sentence}\n"
            f'Return ONLY valid JSON: {{"sentence": "new sentence", '
            f'"explanation": "one short sentence"}}'
        )
        raw, method = _require_llm(prompt)
        data = _parse_json_block(raw)
        new_sentence = str(data["sentence"]).strip()
        new_target = target_word
        explanation = str(data.get("explanation", "Syntactic complexity adjusted."))

    return AlterComplexityResult(
        reason=reason,
        direction=direction,
        original_sentence=sentence,
        original_target=target_word,
        new_sentence=new_sentence,
        new_target=new_target,
        edit_method=method,
        explanation=explanation,
    )