agent.py

import json
import os
import re
from typing import Any, Dict, List, Optional

import anthropic
import numpy as np
import pandas as pd
from model import FEATURES, FEATURE_DISPLAY_NAMES, QUALITY_LABELS, predict
from tabpfn_client import TabPFNClassifier

MODEL = "claude-sonnet-4-6"


def _anthropic_client(api_key: Optional[str] = None) -> anthropic.Anthropic:
    resolved = (api_key or "").strip() or os.environ.get("ANTHROPIC_API_KEY", "").strip()
    if not resolved:
        raise RuntimeError("Anthropic API key is required.")
    return anthropic.Anthropic(api_key=resolved)

EXTRACTION_SYSTEM = """You are a wine chemistry assistant. Extract physicochemical measurements from
the user's natural-language wine description and return ONLY a JSON object with these keys:

fixed_acidity, volatile_acidity, citric_acid, residual_sugar, chlorides,
free_sulfur_dioxide, total_sulfur_dioxide, density, pH, sulphates, alcohol, wine_type

wine_type: 0 for red, 1 for white. Use null for values the user did not mention.

Output rules: respond with raw JSON only — no markdown fences, no commentary."""

DEFAULTS = {
    "fixed_acidity": 7.0,
    "volatile_acidity": 0.33,
    "citric_acid": 0.32,
    "residual_sugar": 5.4,
    "chlorides": 0.056,
    "free_sulfur_dioxide": 30.0,
    "total_sulfur_dioxide": 115.0,
    "density": 0.994,
    "pH": 3.21,
    "sulphates": 0.53,
    "alcohol": 10.5,
    "wine_type": 1,
}


def _response_text(response) -> str:
    """Concatenate all text blocks from an Anthropic message response."""
    parts = []
    for block in response.content:
        if getattr(block, "type", None) == "text":
            parts.append(block.text)
    return "\n".join(parts).strip()


def _parse_json_object(text: str) -> dict:
    """Parse a JSON object from model output (handles fences and preamble)."""
    if not text:
        raise ValueError("Claude returned an empty response during feature extraction.")

    candidate = text.strip()
    fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", candidate, re.DOTALL | re.IGNORECASE)
    if fenced:
        candidate = fenced.group(1).strip()
    else:
        start = candidate.find("{")
        end = candidate.rfind("}")
        if start != -1 and end != -1 and end > start:
            candidate = candidate[start : end + 1]

    try:
        parsed = json.loads(candidate)
    except json.JSONDecodeError as exc:
        preview = text[:200].replace("\n", " ")
        raise ValueError(
            "Could not parse feature JSON from Claude. "
            f"Preview: {preview!r}"
        ) from exc

    if not isinstance(parsed, dict):
        raise ValueError("Claude feature extraction did not return a JSON object.")
    return parsed


def extract_features(
    user_text: str,
    anthropic_api_key: Optional[str] = None,
) -> tuple[dict, list[str]]:
    """Parse user description into a feature dict. Returns (features, assumed_fields)."""
    text = ""
    for attempt in range(2):
        response = _anthropic_client(anthropic_api_key).messages.create(
            model=MODEL,
            max_tokens=512,
            system=EXTRACTION_SYSTEM,
            messages=[{"role": "user", "content": user_text}],
        )
        text = _response_text(response)
        if text:
            break
    raw = _parse_json_object(text)
    assumed = []
    features = {}
    for f in FEATURES:
        val = raw.get(f)
        if val is None:
            features[f] = DEFAULTS[f]
            assumed.append(FEATURE_DISPLAY_NAMES[f])
        else:
            features[f] = float(val)
    return features, assumed


ADVICE_SYSTEM = """You are a helpful wine quality consultant. Given physicochemical measurements
of a wine, a TabPFN model's predicted quality class (Low/Medium/High), its confidence scores,
and a list of counterfactual probes, write a concise plain-English response (3–5 sentences) that:
1. States the predicted quality and confidence.
2. Highlights the 1–2 most influential factors.
3. Gives one concrete actionable suggestion for improvement.
Keep the tone friendly and accessible to non-chemists."""


def generate_advice(
    features: dict,
    prediction: str,
    probabilities: dict,
    counterfactuals: list[dict],
    anthropic_api_key: Optional[str] = None,
) -> str:
    context = {
        "features": features,
        "prediction": prediction,
        "probabilities": probabilities,
        "counterfactuals": counterfactuals,
    }
    response = _anthropic_client(anthropic_api_key).messages.create(
        model=MODEL,
        max_tokens=512,
        system=ADVICE_SYSTEM,
        messages=[{"role": "user", "content": json.dumps(context)}],
    )
    return response.content[0].text


def run_counterfactuals(clf: TabPFNClassifier, features: dict) -> list[dict]:
    """Probe a few key dimensions ±10% — batched into one predict_proba call."""
    from model import FEATURES, QUALITY_LABELS

    probe_keys = ["alcohol", "volatile_acidity", "sulphates", "residual_sugar"]
    rows = []
    meta = []
    for key in probe_keys:
        original_val = features[key]
        for delta_pct, label in [(+0.10, "+10%"), (-0.10, "−10%")]:
            modified = dict(features)
            modified[key] = original_val * (1 + delta_pct)
            rows.append([modified[f] for f in FEATURES])
            meta.append((key, label, round(modified[key], 4)))

    probs_all = clf.predict_proba(np.array(rows, dtype=float))
    results = []
    for (key, label, new_value), probs in zip(meta, probs_all):
        pred_idx = int(np.argmax(probs))
        prob_map = {QUALITY_LABELS[i]: float(probs[i]) for i in range(len(probs))}
        results.append({
            "feature": FEATURE_DISPLAY_NAMES[key],
            "change": label,
            "new_value": new_value,
            "prediction": QUALITY_LABELS[pred_idx],
            "probabilities": prob_map,
        })
    return results


INTERPRETATION_SYSTEM = """You are a wine educator explaining machine-learning results to a curious non-expert.

You will receive JSON with TabPFN prediction outputs and shapiq Shapley feature attributions.
Write a clear plain-English summary (about 4–6 short paragraphs) covering:

1. **The prediction** — the quality class, confidence (probabilities), and what the estimated 0–10 quality index means in everyday wine terms.
2. **The waterfall chart** — explain that the model starts from an average baseline probability for this class, then each measured feature pushes that probability up (positive Shapley) or down (negative). Name the top 2–3 features helping quality and the top 2–3 holding it back, using accessible wine language.
3. **Interactions** — if only first-order Shapley values are provided (no pairwise terms), say this view treats each factor independently; the model may still combine them internally.
4. **Ground truth** (only if present) — how the prediction compares to the actual score and what that gap might suggest.
5. **What-if probes** (only if present) — briefly note which small changes would most likely shift the prediction.

Rules:
- Use only numbers from the input; do not invent measurements or probabilities.
- Avoid jargon like "Shapley" or "coalition" in the main text — say "contribution" or "push" instead.
- Friendly, concise tone. No bullet lists longer than 3 items."""


def _features_for_prompt(features: dict) -> dict:
    return {FEATURE_DISPLAY_NAMES[k]: round(float(features[k]), 4) for k in FEATURES}


def generate_layman_explanation(
    features: dict,
    prediction: str,
    probabilities: dict,
    quality_index: float,
    attributions: pd.DataFrame,
    shapley_meta: dict,
    *,
    user_description: Optional[str] = None,
    ground_truth_class: Optional[str] = None,
    ground_truth_index: Optional[int] = None,
    counterfactuals: Optional[List[dict]] = None,
    assumed_fields: Optional[List[str]] = None,
    anthropic_api_key: Optional[str] = None,
) -> str:
    """Explain prediction + Shapley attributions in accessible language."""
    shapley_rows = []
    if attributions is not None and not attributions.empty and "Shapley" in attributions.columns:
        shapley_rows = attributions[["Feature", "Shapley", "Effect"]].to_dict(orient="records")

    context: Dict[str, Any] = {
        "user_description": user_description,
        "assumed_fields": assumed_fields or [],
        "measurements": _features_for_prompt(features),
        "prediction": {
            "quality_class": prediction,
            "probabilities": probabilities,
            "estimated_quality_index_0_to_10": quality_index,
        },
        "shapley_attribution": {
            **shapley_meta,
            "contributions": shapley_rows,
            "how_to_read_waterfall": (
                "Baseline is the average model score for this quality class; "
                "each bar shows how a feature moves that score up or down."
            ),
        },
    }
    if ground_truth_class is not None:
        context["ground_truth"] = {
            "quality_class": ground_truth_class,
            "quality_index": ground_truth_index,
        }
    if counterfactuals:
        context["what_if_probes"] = counterfactuals

    response = _anthropic_client(anthropic_api_key).messages.create(
        model=MODEL,
        max_tokens=900,
        system=INTERPRETATION_SYSTEM,
        messages=[{"role": "user", "content": json.dumps(context, indent=2)}],
    )
    return response.content[0].text


def analyse_wine(clf: TabPFNClassifier, user_text: str) -> tuple[str, dict, float, int, list[str], str, dict]:
    """Full agentic pipeline. Returns (prediction, probabilities, quality_index, pred_idx, assumed, advice, features)."""
    features, assumed = extract_features(user_text)
    prediction, probabilities, quality_index, pred_idx = predict(clf, features)
    counterfactuals = run_counterfactuals(clf, features)
    advice = generate_advice(features, prediction, probabilities, counterfactuals)
    return prediction, probabilities, quality_index, pred_idx, assumed, advice, features