"""
Appliance Health & Leakage Monitor — plain "government service" interface.

Goal: anyone can use it first time. A clear table of every appliance (estimated vs
actual watts + status), tick-box selection, plain-language fault labels, and a
high-contrast status banner. The simulation engine (engine.py) is unchanged.

Run locally:   python app.py
Data:          expects data/ampds2_hourly.parquet
"""
import os
import matplotlib
matplotlib.use("Agg")
from matplotlib.figure import Figure
import numpy as np
import pandas as pd
import gradio as gr

import engine

DATA_PATH = os.environ.get("AMPDS2_PARQUET", "data/ampds2_hourly.parquet")

DF = MAINS = None
SUBS = []
LOAD_ERR = ""
try:
    DF, MAINS, SUBS = engine.load_hourly(DATA_PATH)
except Exception as e:
    LOAD_ERR = str(e)


def nice(code):
    n = engine.AMPDS2_NAMES.get(code, code)
    return n if n != code else code.replace("_", " ").title()


CHOICES = [(nice(c), c) for c in SUBS]

# ---- plain-language maps ----
PROBLEMS = ["Working fine",
            "Drawing too much (worn)",
            "Drawing too little (weak)",
            "Not running (broken)"]
KIND = {"Drawing too much (worn)": "Over-draw",
        "Drawing too little (weak)": "Under-draw",
        "Not running (broken)": "Off"}
OVER_PCT = {"Low": 25, "Medium": 60, "High": 120}
UNDER_PCT = {"Low": 20, "Medium": 45, "High": 75}
LEAK_W = {"None": 0, "Small": 150, "Medium": 450, "Large": 900}
SPEED = {"Slow": 2, "Normal": 6, "Fast": 12}
STATUS_PLAIN = {"OK": ("Normal", "s-ok"), "OVER": ("Drawing too much", "s-bad"),
                "UNDER": ("Drawing too little", "s-warn"), "OFF": ("Not running", "s-bad"),
                "idle": ("Not in use", "s-idle")}

# Force the plain light theme (avoids the dark-mode low-contrast problem entirely).
FORCE_LIGHT = """() => {
  const u = new URL(window.location.href);
  if (u.searchParams.get('__theme') !== 'light') { u.searchParams.set('__theme','light'); window.location.replace(u.href); }
}"""

CSS = """
.gradio-container{max-width:1040px !important;margin:0 auto !important;
  font-family:Arial,"Helvetica Neue",Helvetica,sans-serif;background:#ffffff !important;}
.govhead{background:#0b0c0c;color:#fff;padding:16px 20px;border-bottom:6px solid #1d70b8;}
.govhead h1{margin:0;font-size:24px;font-weight:700;color:#fff;}
.govhead .sub{color:#cfd2d4;font-size:14px;margin-top:4px;}
.step{font-size:18px;font-weight:700;color:#0b0c0c;border-bottom:2px solid #0b0c0c;padding-bottom:4px;margin:4px 0 2px;}
.help{color:#505a5f;font-size:14px;margin:2px 0 8px;}
.stats{display:flex;border:1px solid #b1b4b6;margin:4px 0;}
.stat{flex:1;padding:12px 14px;border-right:1px solid #b1b4b6;background:#ffffff;}
.stat:last-child{border-right:none;}
.stat .l{font-size:12px;color:#505a5f;text-transform:uppercase;letter-spacing:.03em;}
.stat .v{font-size:22px;font-weight:700;color:#0b0c0c;margin-top:4px;}
.banner{padding:14px 16px;border-left:6px solid #00703c;background:#f3f9f4;color:#0b0c0c;font-size:16px;margin:4px 0;}
.banner.bad{border-left-color:#d4351c;background:#fdf2f1;}
.banner h3{margin:0 0 6px;font-size:18px;font-weight:700;color:#0b0c0c;}
.banner ul{margin:6px 0 0;padding-left:20px;} .banner li{margin:3px 0;}
table.t{width:100%;border-collapse:collapse;font-size:15px;background:#fff;}
table.t th{text-align:left;border-bottom:2px solid #0b0c0c;padding:8px 10px;font-weight:700;color:#0b0c0c;}
table.t td{border-bottom:1px solid #b1b4b6;padding:8px 10px;color:#0b0c0c;}
table.t td.num{text-align:right;font-variant-numeric:tabular-nums;}
.s-ok{color:#00703c;font-weight:700;} .s-bad{color:#d4351c;font-weight:700;}
.s-warn{color:#8a6100;font-weight:700;} .s-idle{color:#505a5f;}
button.primary{background:#00703c !important;color:#fff !important;border:0 !important;border-radius:0 !important;font-weight:700 !important;}
button.secondary{border-radius:0 !important;}
"""


# ----------------------------------------------------------------- HTML builders
def stats_html(k, leak, n_problems):
    leakv = f"{max(0, leak['now'] - leak['expected']):.0f} W" if leak["flag"] else "None"
    return (f'<div class="stats">'
            f'<div class="stat"><div class="l">Date &amp; time</div><div class="v">{k["time"]:%d %b %Y, %H:%M}</div></div>'
            f'<div class="stat"><div class="l">Whole house now</div><div class="v">{k["load_W"]:,.0f} W</div></div>'
            f'<div class="stat"><div class="l">Possible leak</div><div class="v">{leakv}</div></div>'
            f'<div class="stat"><div class="l">Problems found</div><div class="v">{n_problems}</div></div>'
            f'</div>')


def banner_html(rows, leak):
    issues = []
    if leak["flag"]:
        issues.append(f"Possible electricity leak of about {max(0, leak['now'] - leak['expected']):.0f} W "
                      f"somewhere in the house — power being used that no appliance accounts for.")
    for r in rows:
        st, nm = r["status"], nice(r["code"])
        if st == "OVER":
            issues.append(f"{nm} is drawing too much power — about {r['live_W']:,} W "
                          f"(normally about {r['rated_W']:,} W).")
        elif st == "UNDER":
            issues.append(f"{nm} is drawing less power than normal — about {r['live_W']:,} W "
                          f"(normally about {r['rated_W']:,} W).")
        elif st == "OFF":
            issues.append(f"{nm} is not running, although it normally would use about {r['rated_W']:,} W.")
    if not issues:
        return ('<div class="banner"><h3>&#10003; System normal</h3>'
                'All equipment is working as expected and no electricity leak was detected.</div>')
    lis = "".join(f"<li>{x}</li>" for x in issues)
    return f'<div class="banner bad"><h3>&#9888; {len(issues)} problem(s) detected</h3><ul>{lis}</ul></div>'


def table_html(rows):
    body = ""
    for r in sorted(rows, key=lambda r: -r["live_W"]):
        label, cls = STATUS_PLAIN.get(r["status"], (r["status"], "s-idle"))
        body += (f'<tr><td>{nice(r["code"])}</td>'
                 f'<td class="num">{r["rated_W"]:,}</td>'
                 f'<td class="num">{r["live_W"]:,}</td>'
                 f'<td class="{cls}">{label}</td></tr>')
    return (f'<table class="t"><thead><tr><th>Equipment</th>'
            f'<th>Estimated watts</th><th>Now using (watts)</th><th>Status</th></tr></thead>'
            f'<tbody>{body}</tbody></table>')


def chart(sim, cursor):
    code = sim.focus
    name = nice(code)
    t = sim.play_index
    y = sim.live()["sub"][code].to_numpy()
    fig = Figure(figsize=(7, 2.9), dpi=100)
    fig.patch.set_facecolor("white")
    ax = fig.add_subplot(111)
    ax.set_facecolor("white")
    ax.plot(t[:cursor + 1], y[:cursor + 1], color="#1d70b8", lw=1.5, label="power used now")
    if sim.rated.get(code, 0) > 0:
        ax.axhline(sim.rated[code], color="#0b0c0c", lw=1.0, ls="--",
                   label=f"estimated normal ({sim.rated[code]:.0f} W)")
    ax.axvline(t[cursor], color="#505a5f", lw=0.8)
    ax.set_title(f"{name}: power used over time", fontsize=11, color="#0b0c0c")
    ax.set_ylabel("Watts", fontsize=9)
    ax.tick_params(labelsize=8, colors="#0b0c0c")
    for s in ax.spines.values():
        s.set_color("#b1b4b6")
    ax.legend(fontsize=8, loc="upper left")
    fig.autofmt_xdate(rotation=0, ha="center")
    fig.tight_layout()
    return fig


def render(sim, cursor):
    cursor = int(min(max(cursor, 0), sim.play_len - 1))
    _, rows, leak = sim.detect(cursor)
    k = sim.kpis(cursor)
    n = sum(1 for r in rows if r["status"] in ("OVER", "UNDER", "OFF")) + (1 if leak["flag"] else 0)
    return stats_html(k, leak, n), banner_html(rows, leak), table_html(rows), chart(sim, cursor)


# --------------------------------------------------------------------- handlers
def initialize():
    if LOAD_ERR:
        msg = f'<div class="banner bad"><h3>&#9888; System data not found</h3>{LOAD_ERR}</div>'
        return [None, 0, gr.update(), gr.update(), "", msg, "", None,
                "Could not load the system."]
    sim = engine.Simulator(DF, MAINS, SUBS)
    sim.setup(45, 60)
    model = engine.get_timesfm()
    cyc = [c for c in SUBS if c in engine.CYCLIC_CODES and sim.rated.get(c, 0) > 0]
    sim.build_envelopes(list(dict.fromkeys(cyc + [sim.focus])), model)
    s, b, tab, fig = render(sim, 0)
    note = ("System ready. Every appliance is listed below with its estimated normal power and what "
            "it is using now. Set a problem on the left if you like, then press Play.")
    choices = [(nice(c), c) for c in sorted(SUBS, key=lambda c: -sim.rated.get(c, 0))]
    return [sim, 0,
            gr.update(choices=choices, value=[]),
            gr.update(choices=choices, value=sim.focus),
            s, b, tab, fig, note]


def on_tick(sim, cursor, speed):
    if sim is None:
        return [cursor, gr.update(), gr.update(), gr.update(), gr.update(), gr.update()]
    nc = int(min(cursor + SPEED.get(speed, 6), sim.play_len - 1))
    s, b, tab, fig = render(sim, nc)
    return [nc, s, b, tab, fig, gr.Timer(active=(nc < sim.play_len - 1))]


def on_reset_time(sim):
    if sim is None:
        return [0, gr.update(), gr.update(), gr.update(), gr.update()]
    return [0, *render(sim, 0)]


def on_chart(sim, cursor, code):
    if sim is None or code is None:
        return [sim, gr.update(), gr.update(), gr.update(), gr.update()]
    sim.focus = code
    if code not in sim.envelopes:
        sim.build_envelopes([code], engine.get_timesfm())
    return [sim, *render(sim, cursor)]


def on_apply_fault(sim, cursor, codes, problem, severity):
    if sim is None or not codes:
        return [sim, gr.update(), gr.update(), gr.update(), gr.update()]
    for code in codes:
        if problem == "Working fine":
            sim.injections.pop(code, None)
            sim._dirty = True
        elif problem == "Not running (broken)":
            sim.set_injection(code, "Off", 0, 0)
        elif problem == "Drawing too much (worn)":
            sim.set_injection(code, "Over-draw", OVER_PCT.get(severity, 60), 0)
        elif problem == "Drawing too little (weak)":
            sim.set_injection(code, "Under-draw", UNDER_PCT.get(severity, 45), 0)
    return [sim, *render(sim, cursor)]


def on_reset_faults(sim, cursor):
    if sim is None:
        return [sim, gr.update(), gr.update(), gr.update(), gr.update()]
    sim.clear_injections()
    return [sim, *render(sim, cursor)]


def on_set_leak(sim, cursor, size):
    if sim is None:
        return [sim, gr.update(), gr.update(), gr.update(), gr.update()]
    if size == "None":
        sim.clear_leak()
    else:
        sim.set_leak(LEAK_W.get(size, 0), 0)
    return [sim, *render(sim, cursor)]


# --------------------------------------------------------------------- layout
with gr.Blocks(css=CSS, theme=gr.themes.Default(primary_hue="green", neutral_hue="slate"),
               title="Appliance Power Monitor") as demo:
    sim_state = gr.State(None)
    cursor_state = gr.State(0)

    gr.HTML('<div class="govhead"><h1>Appliance Power &amp; Leak Monitor</h1>'
            '<div class="sub">Checks each appliance against its normal power use, and watches the '
            'whole house for unaccounted electricity (leaks).</div></div>')

    gr.HTML('<div class="step">Step 1 — Load the system</div>'
            '<div class="help">This reads the household data and works out the normal power for each appliance. '
            'It can take up to a minute the first time.</div>')
    with gr.Row():
        load_btn = gr.Button("Load the system", variant="primary", scale=0)
        init_status = gr.Markdown("Press **Load the system** to begin."
                                  + (f"  \n**{LOAD_ERR}**" if LOAD_ERR else ""))

    with gr.Row():
        # -------- left: simple controls --------
        with gr.Column(scale=1):
            gr.HTML('<div class="step">Step 2 — Simulate a problem (optional)</div>'
                    '<div class="help">Tick the equipment, choose what is wrong and how bad, then Apply.</div>')
            equip = gr.CheckboxGroup(choices=CHOICES, label="Equipment to affect")
            problem = gr.Radio(PROBLEMS, value="Working fine", label="What is wrong with it?")
            severity = gr.Radio(["Low", "Medium", "High"], value="Medium",
                                label="How bad? (for over/under power)")
            with gr.Row():
                apply_btn = gr.Button("Apply to ticked equipment", variant="primary")
                reset_faults_btn = gr.Button("Set all back to normal", variant="secondary")

            gr.HTML('<div class="step" style="margin-top:14px">Electricity leak in the house</div>'
                    '<div class="help">Simulate power being lost somewhere with no appliance (faulty wiring, '
                    'an unmetered load, theft).</div>')
            leak = gr.Radio(list(LEAK_W.keys()), value="None", label="Leak size")

            gr.HTML('<div class="step" style="margin-top:14px">Run</div>'
                    '<div class="help">Play steps through the records so you can watch the table and status change.</div>')
            with gr.Row():
                play_btn = gr.Button("Play", variant="primary")
                pause_btn = gr.Button("Pause", variant="secondary")
            with gr.Row():
                step_btn = gr.Button("Step", variant="secondary")
                reset_time_btn = gr.Button("Reset time", variant="secondary")
            speed = gr.Radio(list(SPEED.keys()), value="Normal", label="Speed")

        # -------- right: what you watch --------
        with gr.Column(scale=2):
            stats_box = gr.HTML()
            banner_box = gr.HTML()
            gr.HTML('<div class="step">Equipment status</div>')
            table_box = gr.HTML()
            chart_dd = gr.Dropdown(choices=CHOICES, label="Show power chart for",
                                   value=(CHOICES[0][1] if CHOICES else None))
            chart_box = gr.Plot(show_label=False)

    timer = gr.Timer(0.8, active=False)

    with gr.Accordion("How this works (plain English)", open=False):
        gr.Markdown(
            "**Estimated watts** is the appliance's normal power, learned from a healthy period of data.\n\n"
            "**Now using** is the live power as the records play.\n\n"
            "An appliance is flagged when its power stays clearly above its normal value (worn / drawing too much), "
            "below it (weak), or at zero when it should be running (broken).\n\n"
            "A **leak** is found by comparing the whole-house meter to the sum of all appliances. A faulty "
            "appliance raises the meter by the same amount, so it does not look like a leak — only truly "
            "unaccounted power does.\n\n"
            "There are no real fault labels in the data, so you test the system by simulating problems above "
            "and watching it respond.")

    # ---- wiring ----
    vis = [stats_box, banner_box, table_box, chart_box]
    load_btn.click(initialize, None,
                   [sim_state, cursor_state, equip, chart_dd, *vis, init_status])

    play_btn.click(lambda: gr.Timer(active=True), None, timer)
    pause_btn.click(lambda: gr.Timer(active=False), None, timer)
    timer.tick(on_tick, [sim_state, cursor_state, speed], [cursor_state, *vis, timer])
    step_btn.click(on_tick, [sim_state, cursor_state, speed], [cursor_state, *vis, timer])
    reset_time_btn.click(on_reset_time, [sim_state], [cursor_state, *vis])

    apply_btn.click(on_apply_fault, [sim_state, cursor_state, equip, problem, severity],
                    [sim_state, *vis])
    reset_faults_btn.click(on_reset_faults, [sim_state, cursor_state], [sim_state, *vis])
    leak.change(on_set_leak, [sim_state, cursor_state, leak], [sim_state, *vis])
    chart_dd.change(on_chart, [sim_state, cursor_state, chart_dd], [sim_state, *vis])

    demo.load(None, None, None, js=FORCE_LIGHT)

if __name__ == "__main__":
    demo.launch()