app.py

import re
import io
import math
import os
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from PIL import Image
from dataclasses import dataclass
import gradio as gr

# ---------- Optional: tiny open model for nicer text (CPU-friendly) ----------
USE_LLM = os.getenv("USE_LLM", "false").lower() == "true"
MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

llm = None
if USE_LLM:
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
        tok = AutoTokenizer.from_pretrained(MODEL_ID)
        mdl = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
        llm = pipeline("text-generation", model=mdl, tokenizer=tok, max_new_tokens=220)
    except Exception:
        llm = None  # fall back to rules

# ---------- Simple parser & causal checker for conditional probability ----------
@dataclass
class Parsed:
    pA: float | None
    pB: float | None
    pAintB: float | None
    text: str

prob_pat = r"P\(\s*A\s*\)\s*=\s*([0-9]*\.?[0-9]+)"
prob_patB = r"P\(\s*B\s*\)\s*=\s*([0-9]*\.?[0-9]+)"
# Use non-capturing groups to keep numbering stable
prob_pAnB = r"(?:P\(\s*A\s*∩\s*B\s*\)|P\(\s*A\\cap B\s*\)|P\(\s*A\s*&\s*B\s*\)|P\(\s*A\s*and\s*B\s*\)|P\(\s*A\s*\,\s*B\s*\))"

def parse_input(txt: str) -> Parsed:
    t = txt
    pA = re.search(prob_pat, t)
    pB = re.search(prob_patB, t)
    pAnB_val = None
    m = re.search(rf"{prob_pAnB}\s*=\s*([0-9]*\.?[0-9]+)", t, re.IGNORECASE)
    if m:
        pAnB_val = float(m.group(1))
    return Parsed(
        pA=float(pA.group(1)) if pA else None,
        pB=float(pB.group(1)) if pB else None,
        pAintB=pAnB_val,
        text=txt.strip()
    )

def compute_reasoning(parsed: Parsed, user_reasoning: str | None = ""):
    res = {
        "can_compute": False,
        "p_cond": None,
        "independence_claimed": False,
        "independence_holds": None,
        "explanation": "",
    }
    if parsed.pA is None or parsed.pB is None or parsed.pAintB is None:
        res["explanation"] = "Missing P(A), P(B) or P(A∩B)."
        return res

    pA, pB, pAnB = parsed.pA, parsed.pB, parsed.pAintB
    res["p_cond"] = pAnB / pB if pB != 0 else None
    res["can_compute"] = pB != 0

    if user_reasoning:
        if re.search(r"independ", user_reasoning, re.IGNORECASE) or re.search(r"P\(A\|B\)\s*=\s*P\(A\)", user_reasoning):
            res["independence_claimed"] = True

    res["independence_holds"] = math.isclose(pAnB, pA * pB, rel_tol=1e-6, abs_tol=1e-6)

    return res

def make_graph_image(parsed: Parsed, info: dict):
    G = nx.DiGraph()
    G.add_node("Given", desc=f"P(A)={parsed.pA}, P(B)={parsed.pB}, P(A∩B)={parsed.pAintB}")
    G.add_node("Formula", desc="P(A|B)=P(A∩B)/P(B)")
    G.add_node("Compute", desc=f"P(A|B) = {parsed.pAintB}/{parsed.pB}")
    result = f"{(parsed.pAintB/parsed.pB):.4f}" if info["can_compute"] else "undefined"
    G.add_node("Result", desc=f"P(A|B)={result}")
    G.add_node("Independence", desc="Assume A ⟂ B (P(A|B)=P(A))")
    G.add_node("Check", desc=f"P(A)P(B)={parsed.pA*parsed.pB:.4f} vs P(A∩B)={parsed.pAintB:.4f}")

    G.add_edges_from([
        ("Given","Formula"),
        ("Formula","Compute"),
        ("Compute","Result"),
        ("Given","Check"),
        ("Independence","Result")
    ])

    pos = {
        "Given": (-0.9,  0.3),
        "Formula": (-0.2,  0.5),
        "Compute": ( 0.5,  0.3),
        "Result":  ( 0.6, -0.5),
        "Independence": (-0.7, -0.4),
        "Check":  (-0.1, -0.2),
    }

    fig, ax = plt.subplots(figsize=(5.6, 4.2), dpi=180)
    ax.axis('off')

    node_colors = []
    for n in G.nodes():
        if n == "Independence" and info["independence_claimed"] and not info["independence_holds"]:
            node_colors.append("#f8d7da")
        elif n == "Result":
            node_colors.append("#d1e7dd")
        else:
            node_colors.append("#f0f4ff")

    nx.draw_networkx_nodes(G, pos, node_size=2200, node_color=node_colors, linewidths=1.5, edgecolors="#213555")
    nx.draw_networkx_edges(G, pos, arrows=True, arrowstyle="->", width=1.2, edge_color="#213555")
    nx.draw_networkx_labels(G, pos, font_size=9, font_color="#1a2b3c")

    for n, (x,y) in pos.items():
        ax.text(x, y-0.12, G.nodes[n]["desc"], fontsize=8, ha="center", color="#334e68")

    buf = io.BytesIO()
    plt.tight_layout()
    plt.savefig(buf, format="png", bbox_inches="tight", dpi=180)
    plt.close(fig)
    buf.seek(0)

    img = Image.open(buf).convert("RGB")
    return np.array(img)

def rule_based_explanation(parsed: Parsed, info: dict):
    if not info["can_compute"]:
        return "Cannot compute: P(B)=0 or missing values."
    lines = [
        "• Using the definition, P(A|B)=P(A∩B)/P(B).",
        f"• Substituting: {parsed.pAintB} / {parsed.pB} → {parsed.pAintB/parsed.pB:.4f}."
    ]
    if info["independence_claimed"]:
        if info["independence_holds"]:
            lines += [
                "• Your assumption of independence holds (P(A∩B)=P(A)P(B)).",
                "• In this case P(A|B)=P(A) is consistent with the data."
            ]
        else:
            lines += [
                f"• Your assumption of independence is violated: P(A)P(B)={parsed.pA*parsed.pB:.4f} ≠ P(A∩B)={parsed.pAintB:.4f}.",
                "• Minimal fix: treat A and B as dependent; use P(A|B)=P(A∩B)/P(B).",
                f"• Counterfactual: if independence held, P(A|B) would equal P(A)={parsed.pA:.4f}."
            ]
    else:
        lines.append("• No independence assumption detected; standard conditional rule applied.")
    return "\n".join(lines)

def llm_explain(prompt):
    if llm is None:
        return None
    try:
        out = llm(prompt)[0]["generated_text"]
        return out
    except Exception:
        return None

EXAMPLE_PROBLEM = "Problem: Given P(A)=0.4, P(B)=0.5, P(A∩B)=0.18, find P(A|B)."
EXAMPLE_REASONING = "You assumed independence, so P(A|B)=P(A)=0.4."

def run(problem_text, user_reasoning):
    parsed = parse_input(problem_text or "")
    info = compute_reasoning(parsed, user_reasoning or "")
    graph_img = make_graph_image(parsed, info)
    rb_text = rule_based_explanation(parsed, info)
    llm_out = llm_explain(
        f"Explain the mistake and minimal fix clearly:\nProblem: {problem_text}\nUser reasoning: {user_reasoning}\n"
        f"Key facts: P(A)={parsed.pA}, P(B)={parsed.pB}, P(A∩B)={parsed.pAintB}.\\n"
        f"Computed P(A|B)={info['p_cond']}. Independence claimed: {info['independence_claimed']}. "
        f"Independence holds: {info['independence_holds']}.\\n"
        f"Give a concise, student-friendly diagnosis and a contrastive / counterfactual fix."
    )
    final_text = llm_out if llm_out else rb_text
    status = {
        "Parsed": f"P(A)={parsed.pA}, P(B)={parsed.pB}, P(A∩B)={parsed.pAintB}",
        "Independence claimed": info["independence_claimed"],
        "Independence holds": info["independence_holds"],
        "P(A|B)": f"{info['p_cond']:.4f}" if info["p_cond"] is not None else "undefined",
    }
    stat_str = "\n".join([f"{k}: {v}" for k,v in status.items()])
    return graph_img, final_text, stat_str

with gr.Blocks(theme=gr.themes.Soft(primary_hue='indigo')) as demo:
    gr.Markdown("## Aicher (MVP) — Visual Causal Tutor\n"
                "Enter a conditional-probability problem (or use the example). "
                "Optionally type your own reasoning to test independence errors.")

    with gr.Row():
        with gr.Column(scale=1):
            problem = gr.Textbox(label="Problem", value=EXAMPLE_PROBLEM, lines=4)
            reasoning = gr.Textbox(label="Your reasoning (optional)", value=EXAMPLE_REASONING, lines=3)
            btn = gr.Button("Explain reasoning")
            gr.Examples(
                examples=[[EXAMPLE_PROBLEM, EXAMPLE_REASONING],
                          ["P(A)=0.3, P(B)=0.25, P(A∩B)=0.05. Find P(A|B).", "No independence claimed."]],
                inputs=[problem, reasoning],
            )
        with gr.Column(scale=2):
            graph = gr.Image(label="Reasoning Graph", type="numpy")
        with gr.Column(scale=1):
            diagnosis = gr.Textbox(label="Diagnosis & Minimal Fix", lines=12)
            status = gr.Textbox(label="Parsed / Status", lines=6)

    btn.click(run, inputs=[problem, reasoning], outputs=[graph, diagnosis, status])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", share=True)