ui/components.py

"""
Chart and map builders. Pure functions: take domain objects (NeedProfile,
AllocationPlan), return Plotly figures. No Gradio import here, so these
can be unit-tested without spinning up the UI.
"""

from __future__ import annotations
import plotly.graph_objects as go
from core.schemas import NeedProfile, AllocationPlan

DARK_BG = "#0b0f14"
PANEL_BG = "#11161d"
GRID_COLOR = "#232b36"
TEXT_COLOR = "#e6edf3"
ACCENT = "#3fb6ff"

SEVERITY_COLORS = {
    "critical": "#ff4d4f",  # >=75
    "high": "#ff9f43",      # 50-74
    "medium": "#ffd166",    # 30-49
    "low": "#3ddc97",       # <30
}


def _priority_band(score: float) -> str:
    if score >= 75:
        return "critical"
    if score >= 50:
        return "high"
    if score >= 30:
        return "medium"
    return "low"


def _base_layout(title: str, height: int = 340) -> dict:
    return dict(
        title=dict(text=title, font=dict(color=TEXT_COLOR, size=14)),
        paper_bgcolor=PANEL_BG,
        plot_bgcolor=PANEL_BG,
        font=dict(color=TEXT_COLOR, size=11),
        margin=dict(l=40, r=20, t=40, b=40),
        height=height,
        xaxis=dict(gridcolor=GRID_COLOR, zerolinecolor=GRID_COLOR),
        yaxis=dict(gridcolor=GRID_COLOR, zerolinecolor=GRID_COLOR),
    )


def build_situation_map(profiles: list[NeedProfile]) -> go.Figure:
    """
    Self-contained simulated map using local (lat, lon) coordinates on a
    Plotly scatter plot — no external map tiles, no Google Maps/Mapbox
    dependency, exactly per the 'no external paid map services' requirement.
    """
    fig = go.Figure()
    for p in profiles:
        band = _priority_band(p.priority_score)
        fig.add_trace(go.Scatter(
            x=[p.lon], y=[p.lat],
            mode="markers+text",
            marker=dict(size=18 + p.priority_score / 4, color=SEVERITY_COLORS[band],
                        line=dict(width=2, color="#ffffff")),
            text=[p.display_name.split(" - ")[0]],
            textposition="top center",
            textfont=dict(color=TEXT_COLOR, size=10),
            name=p.display_name,
            hovertext=(f"{p.display_name}<br>Priority: {p.priority_score}/100<br>"
                       f"Severity: {p.severity}/10<br>People affected: {p.people_affected}<br>"
                       f"Damage: {p.damage_type}"),
            hoverinfo="text",
        ))
    layout = _base_layout("Simulated Incident Map (local coordinates)", height=420)
    layout["showlegend"] = False
    layout["xaxis"]["title"] = "Relative East-West position"
    layout["yaxis"]["title"] = "Relative North-South position"
    fig.update_layout(**layout)
    return fig


def build_allocation_chart(plan: AllocationPlan, profiles: list[NeedProfile]) -> go.Figure:
    names = {p.location_id: p.display_name.split(" - ")[0] for p in profiles}
    locs = [a.location_id for a in plan.allocations]
    labels = [names.get(l, l) for l in locs]

    fig = go.Figure()
    for resource, assigned_attr, required_attr, color in [
        ("Medical", "assigned_medical_teams", "required_medical_teams", "#ff6b6b"),
        ("Rescue", "assigned_rescue_teams", "required_rescue_teams", "#ffa94d"),
        ("Supply", "assigned_supply_trucks", "required_supply_trucks", "#4dabf7"),
    ]:
        assigned = [getattr(a, assigned_attr) for a in plan.allocations]
        fig.add_trace(go.Bar(name=f"{resource} assigned", x=labels, y=assigned, marker_color=color))

    required_total = [sum([getattr(a, "required_medical_teams"), getattr(a, "required_rescue_teams"),
                            getattr(a, "required_supply_trucks")]) for a in plan.allocations]
    fig.add_trace(go.Scatter(name="Total required", x=labels, y=required_total, mode="markers",
                              marker=dict(color="white", size=10, symbol="line-ew", line=dict(width=2))))
    fig.update_layout(barmode="stack", **_base_layout("Resource Allocation vs Required, by Location"))
    return fig


def build_coverage_chart(plan: AllocationPlan, profiles: list[NeedProfile]) -> go.Figure:
    names = {p.location_id: p.display_name.split(" - ")[0] for p in profiles}
    sorted_allocs = sorted(plan.allocations, key=lambda a: a.coverage_pct)
    labels = [names.get(a.location_id, a.location_id) for a in sorted_allocs]
    values = [a.coverage_pct for a in sorted_allocs]
    colors = ["#ff4d4f" if v < 50 else "#ffd166" if v < 90 else "#3ddc97" for v in values]
    fig = go.Figure(go.Bar(x=values, y=labels, orientation="h", marker_color=colors,
                            text=[f"{v}%" for v in values], textposition="outside"))
    fig.update_layout(**_base_layout("Coverage % by Location"))
    fig.update_xaxes(range=[0, 110])
    return fig


def build_severity_distribution(profiles: list[NeedProfile]) -> go.Figure:
    bands = {"critical": 0, "high": 0, "medium": 0, "low": 0}
    for p in profiles:
        bands[_priority_band(p.priority_score)] += 1
    fig = go.Figure(go.Bar(
        x=list(bands.keys()), y=list(bands.values()),
        marker_color=[SEVERITY_COLORS[k] for k in bands.keys()],
        text=list(bands.values()), textposition="outside",
    ))
    fig.update_layout(**_base_layout("Locations by Priority Band", height=300))
    return fig


def build_resource_utilization(plan: AllocationPlan) -> go.Figure:
    resources = ["Medical Teams", "Rescue Teams", "Supply Trucks"]
    used = [plan.resources_used.medical_teams, plan.resources_used.rescue_teams, plan.resources_used.supply_trucks]
    available = [plan.resources_available.medical_teams, plan.resources_available.rescue_teams,
                 plan.resources_available.supply_trucks]
    fig = go.Figure()
    fig.add_trace(go.Bar(name="Used", x=resources, y=used, marker_color=ACCENT))
    fig.add_trace(go.Bar(name="Remaining", x=resources,
                          y=[a - u for a, u in zip(available, used)], marker_color="#334155"))
    fig.update_layout(barmode="stack", **_base_layout("Resource Pool Utilization", height=300))
    return fig


def build_unmet_need_chart(plan: AllocationPlan, profiles: list[NeedProfile]) -> go.Figure:
    names = {p.location_id: p.display_name.split(" - ")[0] for p in profiles}
    labels, unmet_totals = [], []
    for a in plan.allocations:
        total_unmet = a.unmet_medical + a.unmet_rescue + a.unmet_supply
        if total_unmet > 0:
            labels.append(names.get(a.location_id, a.location_id))
            unmet_totals.append(total_unmet)
    if not unmet_totals:
        fig = go.Figure()
        fig.add_annotation(text="No unmet needs — full coverage achieved", showarrow=False,
                            font=dict(color=TEXT_COLOR, size=13))
        fig.update_layout(**_base_layout("Unmet Need by Location", height=300))
        return fig
    fig = go.Figure(go.Bar(x=labels, y=unmet_totals, marker_color="#ff4d4f",
                            text=unmet_totals, textposition="outside"))
    fig.update_layout(**_base_layout("Unmet Need (total units short) by Location", height=300))
    return fig


def build_priority_ranking_table(profiles: list[NeedProfile]) -> list[list]:
    """Returns row data for a Gradio Dataframe: rank, name, priority, severity, people, damage type."""
    rows = []
    for i, p in enumerate(profiles, start=1):
        rows.append([i, p.display_name, p.priority_score, p.severity, p.people_affected,
                     p.damage_type.replace("_", " ").title()])
    return rows