import os
import sys
import gradio as gr
import numpy as np
import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from modeling import load_model, validate_window

WEIGHTS = os.environ.get("FELA_PDM_WEIGHTS", ".")
MODELS = {}


def _try_load(variant):
    try:
        return load_model(WEIGHTS, variant=variant)
    except Exception:
        return None


for v in ("cmapss_FD001", "cwru"):
    MODELS[v] = _try_load(v)
EXAMPLE_RUL = np.tile(np.linspace(0.3, 0.7, 30)[:, None], (1, 14))


def score_rul(text):
    m = MODELS.get("cmapss_FD001")
    if m is None:
        return "RUL weights not found. Set FELA_PDM_WEIGHTS to a dir with cmapss_FD001.safetensors + config.json."
    try:
        rows = [r for r in text.strip().splitlines() if r.strip()]
        arr = np.array(
            [[float(x) for x in r.replace(",", " ").split()] for r in rows],
            dtype=np.float32,
        )
    except Exception:
        return "Could not parse. Provide 30 rows of 14 numbers (whitespace or comma separated)."
    if arr.shape != (30, 14):
        return f"Expected a (30, 14) window, got {arr.shape}."
    x = torch.from_numpy(arr).reshape(1, 30, 14)
    validate_window(x, m.cfg)
    rul = m.predict(x, task="rul")
    return f"estimated remaining useful life: {rul:.1f} cycles (capped at 125)"


def load_rul_example():
    return "\n".join((" ".join((f"{v:.3f}" for v in row)) for row in EXAMPLE_RUL))


def score_cwru(text, use_synth):
    m = MODELS.get("cwru")
    if m is None:
        return "Bearing weights not found. Set FELA_PDM_WEIGHTS to a dir with cwru.safetensors + config.json."
    if use_synth or not text.strip():
        t = np.linspace(0, 1, 2048)
        sig = np.sin(2 * np.pi * 120 * t) + 0.2 * np.random.randn(2048)
    else:
        try:
            sig = np.array(
                [float(x) for x in text.replace(",", " ").split()], dtype=np.float32
            )
        except Exception:
            return (
                "Could not parse. Provide 2048 whitespace- or comma-separated samples."
            )
        if sig.size != 2048:
            return f"Expected 2048 samples, got {sig.size}."
    sig = (sig - sig.mean()) / (sig.std() + 1e-06)
    x = torch.from_numpy(sig.astype(np.float32)).reshape(1, 2048, 1)
    validate_window(x, m.cfg)
    idx, prob = m.predict(x, task="cls")
    return f"predicted fault class index: {idx}  (probability {prob:.4f})"


with gr.Blocks(title="FELA-PdM playground") as demo:
    gr.Markdown(
        "# FELA-PdM playground\nOn-device predictive maintenance. Feed a sensor window and see the model's call. For research and illustration only, not a safety-critical controller; do not act on the remaining-useful-life number without independent validation."
    )
    with gr.Tab("Remaining useful life (C-MAPSS)"):
        gr.Markdown(
            "Paste 30 cycles of 14 normalized sensor values (30 rows, 14 numbers each). The FD001 head was not trained on FD002/FD003/FD004, so a window from those public NASA subsets is a real out-of-distribution test. The example below is synthetic and illustrative."
        )
        rul_in = gr.Textbox(label="sensor window (30 x 14)", lines=8)
        rul_out = gr.Textbox(label="result")
        with gr.Row():
            gr.Button("Load synthetic example").click(load_rul_example, outputs=rul_in)
            gr.Button("Score").click(score_rul, inputs=rul_in, outputs=rul_out)
    with gr.Tab("Bearing fault (CWRU)"):
        gr.Markdown(
            "Paste 2048 raw vibration samples (12 kHz), or tick the box to score a synthetic illustrative window. The output is a fault-class index (10 classes)."
        )
        cwru_in = gr.Textbox(label="vibration window (2048 samples)", lines=4)
        cwru_synth = gr.Checkbox(
            label="use a synthetic illustrative window", value=True
        )
        cwru_out = gr.Textbox(label="result")
        gr.Button("Score").click(
            score_cwru, inputs=[cwru_in, cwru_synth], outputs=cwru_out
        )
if __name__ == "__main__":
    demo.launch()