agents/explainer_agent.py

"""
Trade-off Explainer Agent.

Responsibility: read the ALREADY-DECIDED AllocationPlan (produced entirely
by optimization/resource_optimizer.py) and write a plain-English explanation
of why locations received what they received. This agent receives numbers,
it does not produce them. If the LLM is unavailable, a deterministic
template-based explanation is used instead — the UI never goes silent.
"""

from __future__ import annotations
import logging
from core.schemas import AllocationPlan, NeedProfile
from models.model_manager import get_llm, ModelLoadError
from models.zero_gpu import gpu_decorator

logger = logging.getLogger("explainer_agent")

EXPLAIN_PROMPT = """You are an operations analyst. Below is a resource allocation decision that was
already computed by a mathematical optimizer. Explain it in plain, concise English (5-8 sentences).
Do not suggest any different allocation. Only explain the numbers given.

Priority ranking (highest need first):
{priority_lines}

Allocation results:
{allocation_lines}

Overall coverage: {overall_coverage}%
Resources available: {available}
Resources used: {used}

Explain: (1) why the highest-priority locations received what they did, (2) which locations have
unmet needs and why, (3) what the overall coverage percentage means in practice."""


@gpu_decorator(duration=30)
def _call_llm(prompt: str, max_new_tokens: int = 260) -> str:
    import torch
    bundle = get_llm()
    model, tokenizer = bundle["model"], bundle["tokenizer"]
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def _template_fallback(plan: AllocationPlan, profiles: list[NeedProfile]) -> str:
    """Deterministic explanation used if the LLM is unavailable — always accurate
    because it's built directly from the same numbers, just without LLM prose."""
    by_id = {p.location_id: p for p in profiles}
    lines = [f"Overall coverage reached {plan.overall_coverage_pct}% of total requested resources."]
    sorted_allocs = sorted(plan.allocations, key=lambda a: by_id[a.location_id].priority_score, reverse=True)
    top = sorted_allocs[0] if sorted_allocs else None
    if top:
        p = by_id[top.location_id]
        lines.append(
            f"{p.display_name} had the highest priority score ({p.priority_score}/100) and received "
            f"{top.assigned_medical_teams} medical, {top.assigned_rescue_teams} rescue, and "
            f"{top.assigned_supply_trucks} supply units ({top.coverage_pct}% of its request)."
        )
    unmet = [a for a in plan.allocations if (a.unmet_medical + a.unmet_rescue + a.unmet_supply) > 0]
    if unmet:
        names = ", ".join(by_id[a.location_id].display_name for a in unmet)
        lines.append(
            f"{names} have unmet needs because the available resource pool "
            f"(medical: {plan.resources_available.medical_teams}, rescue: {plan.resources_available.rescue_teams}, "
            f"supply: {plan.resources_available.supply_trucks}) was insufficient to fully cover every location "
            f"at once, so the optimizer prioritized higher-need locations first."
        )
    else:
        lines.append("All locations had their full resource request met.")
    return " ".join(lines)


def explain_allocation(plan: AllocationPlan, profiles: list[NeedProfile]) -> tuple[str, str]:
    """Returns (explanation_text, source) where source is 'llm' or 'template_fallback'."""
    by_id = {p.location_id: p for p in profiles}
    priority_lines = "\n".join(
        f"- {by_id[a.location_id].display_name}: priority {by_id[a.location_id].priority_score}/100"
        for a in sorted(plan.allocations, key=lambda a: by_id[a.location_id].priority_score, reverse=True)
    )
    allocation_lines = "\n".join(
        f"- {by_id[a.location_id].display_name}: assigned medical={a.assigned_medical_teams}/{a.required_medical_teams}, "
        f"rescue={a.assigned_rescue_teams}/{a.required_rescue_teams}, supply={a.assigned_supply_trucks}/{a.required_supply_trucks} "
        f"({a.coverage_pct}% covered)"
        for a in plan.allocations
    )
    prompt = EXPLAIN_PROMPT.format(
        priority_lines=priority_lines,
        allocation_lines=allocation_lines,
        overall_coverage=plan.overall_coverage_pct,
        available=plan.resources_available,
        used=plan.resources_used,
    )
    try:
        text = _call_llm(prompt)
        if not text or len(text) < 20:
            raise ValueError("LLM explanation too short/empty")
        return text, "llm"
    except (ModelLoadError, Exception) as e:
        logger.warning(f"Explainer LLM failed ({e}); using deterministic template fallback")
        return _template_fallback(plan, profiles), "template_fallback"