page2_risk_register.py

"""
page2_risk_register.py — Risk Register, 3×3 Matrix, Rules, Weights
"""

import streamlit as st
import plotly.graph_objects as go
import pandas as pd
from styles import (inject_global_css, page_header, metric_row,
                    WCO_GOLD, WCO_BLUE, WCO_GREEN, WCO_RED, WCO_CARD_BG,
                    WCO_BORDER, WCO_MUTED, WCO_GREY_BG)
from simulation_engine import RISK_AREAS, get_default_weights

# ── Static risk matrix data ──────────────────────────────────────────────────
RISK_MATRIX_DATA = {
    "Drugs & Narcotics":           {"likelihood": 3, "impact": 3, "priority": 1, "color": "#C8102E"},
    "Revenue Leakage":             {"likelihood": 3, "impact": 2, "priority": 2, "color": "#F5A800"},
    "IPR Enforcement":             {"likelihood": 2, "impact": 2, "priority": 3, "color": "#9B59B6"},
    "Environmental/Plastic Waste": {"likelihood": 2, "impact": 3, "priority": 4, "color": "#00843D"},
    "Wildlife Smuggling":          {"likelihood": 1, "impact": 3, "priority": 5, "color": "#E67E22"},
}

RISK_PRIORITY_DATA = [
    {
        "Priority": 1, "Risk Area": "Drugs & Narcotics",
        "Likelihood": "High (3)", "Impact": "High (3)",
        "Risk Level": "CRITICAL", "Strategy": "Exploit + Explore",
        "Annual Target": "Zero tolerance – 100% risk-score flagging",
        "WCO Reference": "Compendium Vol.1 Ch.3",
    },
    {
        "Priority": 2, "Risk Area": "Revenue Leakage",
        "Likelihood": "High (3)", "Impact": "Medium (2)",
        "Risk Level": "HIGH", "Strategy": "DATE Exploitation",
        "Annual Target": "≥85% recovery rate on flagged shipments",
        "WCO Reference": "Compendium Vol.2 Ch.6",
    },
    {
        "Priority": 3, "Risk Area": "IPR Enforcement",
        "Likelihood": "Medium (2)", "Impact": "Medium (2)",
        "Risk Level": "MEDIUM", "Strategy": "Hybrid (DATE 80/gATE 20)",
        "Annual Target": "500 seizures/year, 30% uplift YoY",
        "WCO Reference": "TRIPS/WCO IPR Guidelines",
    },
    {
        "Priority": 4, "Risk Area": "Environmental/Plastic Waste",
        "Likelihood": "Medium (2)", "Impact": "High (3)",
        "Risk Level": "HIGH", "Strategy": "gATE Exploration",
        "Annual Target": "Basel Convention compliance 95%",
        "WCO Reference": "Compendium Vol.3 Ch.9",
    },
    {
        "Priority": 5, "Risk Area": "Wildlife Smuggling",
        "Likelihood": "Low (1)", "Impact": "High (3)",
        "Risk Level": "MEDIUM-HIGH", "Strategy": "Hybrid (DATE 70/gATE 30)",
        "Annual Target": "CITES compliance 100%; 200 cases/year",
        "WCO Reference": "CITES/WCO Wildlife Guidelines",
    },
]

RISK_COLORS_LEVEL = {
    "CRITICAL": "#C8102E", "HIGH": "#F5A800",
    "MEDIUM-HIGH": "#E67E22", "MEDIUM": "#9B59B6", "LOW": "#00843D",
}


def risk_matrix_chart():
    """3×3 WCO Risk Matrix with plotted risks."""
    fig = go.Figure()

    # Background zones
    zone_colors = [
        # (x0,y0,x1,y1, fill)
        (0.5, 2.5, 3.5, 3.5, "rgba(200,16,46,0.18)"),   # High-High
        (0.5, 1.5, 3.5, 2.5, "rgba(245,168,0,0.13)"),   # Medium row
        (0.5, 0.5, 3.5, 1.5, "rgba(0,132,61,0.12)"),    # Low row
    ]
    for x0, y0, x1, y1, fill in zone_colors:
        fig.add_shape(type="rect", x0=x0, y0=y0, x1=x1, y1=y1,
                      fillcolor=fill, line_width=0, layer="below")

    # Grid lines
    for v in [1.5, 2.5]:
        fig.add_shape(type="line", x0=0.5, y0=v, x1=3.5, y1=v,
                      line=dict(color="#1E3A6E", width=1))
        fig.add_shape(type="line", x0=v, y0=0.5, x1=v, y1=3.5,
                      line=dict(color="#1E3A6E", width=1))

    # Plot risks
    for name, d in RISK_MATRIX_DATA.items():
        short = name.replace("Environmental/Plastic Waste","Env/Plastic").replace("Revenue Leakage","Revenue")
        fig.add_trace(go.Scatter(
            x=[d["likelihood"]], y=[d["impact"]],
            mode="markers+text",
            marker=dict(size=38, color=d["color"], opacity=0.88,
                        line=dict(color="#C8A951", width=1.5)),
            text=[f"<b>P{d['priority']}</b>"],
            textfont=dict(color="white", size=13, family="Georgia"),
            textposition="middle center",
            name=short,
            hovertemplate=(f"<b>{name}</b><br>Likelihood: {d['likelihood']}"
                           f"<br>Impact: {d['impact']}<br>Priority: {d['priority']}<extra></extra>"),
        ))

    fig.update_layout(
        paper_bgcolor="#070E1C", plot_bgcolor="#0B1220",
        font=dict(family="IBM Plex Sans", color="#D0DCF0"),
        height=420,
        margin=dict(l=60, r=20, t=50, b=60),
        xaxis=dict(
            title="<b>Likelihood →</b>", range=[0.5, 3.5],
            tickvals=[1,2,3], ticktext=["Low","Medium","High"],
            gridcolor="#1E3A6E", zerolinecolor="#1E3A6E",
            title_font_color=WCO_GOLD,
        ),
        yaxis=dict(
            title="<b>← Impact</b>", range=[0.5, 3.5],
            tickvals=[1,2,3], ticktext=["Low","Medium","High"],
            gridcolor="#1E3A6E", zerolinecolor="#1E3A6E",
            title_font_color=WCO_GOLD,
        ),
        title=dict(text="<b>WCO 3×3 Risk Matrix — Customs Risk Areas</b>",
                   font=dict(color=WCO_GOLD, size=15, family="Playfair Display, Georgia"),
                   x=0.5),
        showlegend=True,
        legend=dict(bgcolor="#0F1C35", bordercolor=WCO_BORDER, borderwidth=1,
                    font=dict(size=11), x=1.01, y=0.98),
    )
    return fig


def rules_weight_editor(default_weights: dict) -> dict:
    """Render editable weight sliders per rule and return current weights."""
    st.markdown('<div class="section-title">⚙️ Risk Rules Configuration & Annual Weights</div>',
                unsafe_allow_html=True)
    st.markdown("""<div class="alert-gold">
        Adjust annual weights for each risk rule. Weights are automatically recalibrated 
        after each simulation based on detection efficiency (self-learning feedback loop).
    </div>""", unsafe_allow_html=True)

    updated = {}
    for area_name, area_cfg in RISK_AREAS.items():
        icon = area_cfg["icon"]
        color = area_cfg["color"]
        with st.expander(f"{icon}  {area_name}", expanded=False):
            sub_label = ""
            if "sub_areas" in area_cfg:
                sub_label = f"  *(Sub-areas: {', '.join(area_cfg['sub_areas'])})*"
            st.markdown(f"<span style='color:{color};font-size:13px;'>{sub_label}</span>",
                        unsafe_allow_html=True)

            rule_data = []
            for rule in area_cfg["rules"]:
                rid   = rule["id"]
                rname = rule["name"]
                sub   = rule.get("sub", "")
                col1, col2, col3 = st.columns([1, 3, 2])
                with col1:
                    st.markdown(f"<span style='color:{color};font-weight:600;font-size:12px;'>{rid}</span>",
                                unsafe_allow_html=True)
                    if sub:
                        st.caption(sub)
                with col2:
                    st.markdown(f"<span style='color:#D0DCF0;font-size:13px;'>{rname}</span>",
                                unsafe_allow_html=True)
                with col3:
                    w = st.slider(
                        f"Weight [{rid}]", 0.05, 0.60,
                        float(default_weights.get(rid, rule["weight"])),
                        step=0.01, key=f"w_{rid}",
                        label_visibility="collapsed",
                    )
                    updated[rid] = w
                    rule_data.append((rid, rname, sub, w))
                st.markdown("<hr style='margin:4px 0;border-color:#1E3A6E;'>", unsafe_allow_html=True)

            # Mini weight chart
            if rule_data:
                ids   = [r[0] for r in rule_data]
                names = [r[1][:30] for r in rule_data]
                vals  = [r[3] for r in rule_data]
                mini = go.Figure(go.Bar(
                    x=ids, y=vals,
                    marker_color=color, opacity=0.85,
                    text=[f"{v:.2f}" for v in vals],
                    textposition="outside",
                    textfont=dict(color="#D0DCF0", size=10),
                    hovertext=names, hoverinfo="text+y",
                ))
                mini.update_layout(
                    paper_bgcolor="transparent", plot_bgcolor="#0B1220",
                    height=180, margin=dict(l=10, r=10, t=10, b=30),
                    font=dict(color="#D0DCF0", size=10),
                    xaxis=dict(gridcolor="#1E3A6E"),
                    yaxis=dict(gridcolor="#1E3A6E", range=[0, 0.65]),
                    showlegend=False,
                )
                st.plotly_chart(mini, use_container_width=True)

    return updated


def show():
    inject_global_css()
    page_header("📋", "Risk Register & Risk Matrix",
                "WCO RISK MANAGEMENT COMPENDIUM · 5 RISK AREAS · ANNUAL WEIGHT CONFIGURATION")

    metric_row([
        ("5",     "Risk Areas",       WCO_GOLD),
        ("26",    "Risk Rules Total", WCO_BLUE),
        ("3×3",   "Risk Matrix",      WCO_RED),
        ("Annual","Weight Cycle",     WCO_GREEN),
        ("WCO",   "Compendium Ref.",  WCO_GOLD),
    ])

    # ── Risk Matrix ───────────────────────────────────────────────
    st.markdown('<div class="section-title">📊 3×3 Risk Priority Matrix</div>', unsafe_allow_html=True)
    c_mat, c_legend = st.columns([2, 1])
    with c_mat:
        st.plotly_chart(risk_matrix_chart(), use_container_width=True)
    with c_legend:
        st.markdown("""
        <div class="wco-card">
        <h3>Matrix Legend</h3>
        <div style="font-size:12px;color:#D0DCF0;line-height:2.2;">
          <div style="background:rgba(200,16,46,0.2);padding:4px 10px;border-radius:5px;border-left:3px solid #C8102E;margin:4px 0;">
            🟥 <b>CRITICAL</b> — Immediate action, maximum resources
          </div>
          <div style="background:rgba(245,168,0,0.15);padding:4px 10px;border-radius:5px;border-left:3px solid #F5A800;margin:4px 0;">
            🟧 <b>HIGH</b> — Priority 2, enhanced monitoring
          </div>
          <div style="background:rgba(155,89,182,0.15);padding:4px 10px;border-radius:5px;border-left:3px solid #9B59B6;margin:4px 0;">
            🟪 <b>MEDIUM-HIGH</b> — Balanced exploit/explore
          </div>
          <div style="background:rgba(0,132,61,0.15);padding:4px 10px;border-radius:5px;border-left:3px solid #00843D;margin:4px 0;">
            🟩 <b>MEDIUM</b> — Exploration-led discovery
          </div>
        </div>
        <div style="margin-top:14px;color:#6B85AA;font-size:11px;">
          P1 = Drugs &amp; Narcotics (High-High)<br/>
          P2 = Revenue Leakage (High-Medium)<br/>
          P3 = IPR Enforcement (Med-Med)<br/>
          P4 = Environmental (Med-High)<br/>
          P5 = Wildlife (Low-High)
        </div>
        </div>""", unsafe_allow_html=True)

    # ── Priority Table ────────────────────────────────────────────
    st.markdown('<div class="section-title">🗂️ Risk Prioritisation Table</div>', unsafe_allow_html=True)
    df_prio = pd.DataFrame(RISK_PRIORITY_DATA)

    rows_html = ""
    for _, row in df_prio.iterrows():
        level_col = RISK_COLORS_LEVEL.get(row["Risk Level"], WCO_MUTED)
        area_col  = RISK_MATRIX_DATA.get(row["Risk Area"], {}).get("color", WCO_MUTED)
        rows_html += f"""
        <tr>
          <td style="text-align:center;font-weight:700;color:{WCO_GOLD};">P{row['Priority']}</td>
          <td><span style="color:{area_col};font-weight:600;">{row['Risk Area']}</span></td>
          <td style="text-align:center;">{row['Likelihood']}</td>
          <td style="text-align:center;">{row['Impact']}</td>
          <td><span style="background:{level_col}22;color:{level_col};border:1px solid {level_col};
                           border-radius:5px;padding:2px 8px;font-size:11px;font-weight:600;">
              {row['Risk Level']}</span></td>
          <td style="color:#CC77FF;font-size:12px;">{row['Strategy']}</td>
          <td style="color:#8BAAD4;font-size:12px;">{row['Annual Target']}</td>
          <td style="color:#557;font-size:11px;">{row['WCO Reference']}</td>
        </tr>"""

    st.markdown(f"""
    <table class="wco-table">
      <thead><tr>
        <th>Priority</th><th>Risk Area</th><th>Likelihood</th><th>Impact</th>
        <th>Risk Level</th><th>Selection Strategy</th><th>Annual Target</th><th>WCO Reference</th>
      </tr></thead>
      <tbody>{rows_html}</tbody>
    </table>""", unsafe_allow_html=True)

    st.markdown("<br/>", unsafe_allow_html=True)

    # ── Rules & Weights Editor ────────────────────────────────────
    if "rule_weights" not in st.session_state:
        st.session_state.rule_weights = get_default_weights()

    updated_weights = rules_weight_editor(st.session_state.rule_weights)
    if st.button("💾  Save Weights to Session", type="primary"):
        st.session_state.rule_weights = updated_weights
        st.success("✅ Weights saved — navigate to Page 3 to run simulation with updated weights.")

    # ── Risk area summary cards ───────────────────────────────────
    st.markdown('<div class="section-title">📌 Risk Area Summary Cards</div>', unsafe_allow_html=True)
    cols = st.columns(5)
    for i, (area_name, area_cfg) in enumerate(RISK_AREAS.items()):
        color = area_cfg["color"]
        icon  = area_cfg["icon"]
        n_rules = len(area_cfg["rules"])
        illicit = area_cfg["base_illicit_rate"]
        sub_txt = ", ".join(area_cfg.get("sub_areas", []))
        with cols[i]:
            st.markdown(f"""
            <div style="background:#0F1C35;border:1px solid {color};border-radius:11px;
                        padding:18px 14px;text-align:center;">
              <div style="font-size:28px;">{icon}</div>
              <div style="color:{color};font-weight:700;font-size:12px;margin:8px 0 4px;
                           font-family:'Playfair Display',serif;">{area_name}</div>
              {'<div style="color:#6B85AA;font-size:10px;margin-bottom:6px;">'+sub_txt+'</div>' if sub_txt else ''}
              <div style="color:#D0DCF0;font-size:12px;margin:4px 0;">{n_rules} Risk Rules</div>
              <div style="color:#6B85AA;font-size:11px;">Base Rate: {illicit*100:.0f}%</div>
              <div style="background:{color};height:3px;border-radius:2px;margin-top:10px;
                          opacity:0.6;width:{int(illicit*400)}%;margin-left:auto;margin-right:auto;"></div>
            </div>""", unsafe_allow_html=True)