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

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

WEIGHTS = os.environ.get(
    "FELATAB_WEIGHTS", os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
TIER = os.environ.get("FELATAB_TIER", "small")
try:
    MODEL = load_model(WEIGHTS, tier=TIER)
except Exception as e:
    MODEL = None
    _ERR = str(e)
EXAMPLE = "weight_g, sugar_pct, width_mm, label\n151, 7.2, 12, Apple\n149, 7.0, 13, Apple\n158, 7.4, 11, Apple\n102, 6.1, 3, Lemon\n98, 5.9, 2, Lemon\n110, 6.3, 4, Lemon\n6, 2.1, 16, Grape\n5, 1.9, 17, Grape\n7, 2.2, 15, Grape\n153, 7.1, 12, ?"


def _parse(text):
    rows = [r.strip() for r in text.strip().splitlines() if r.strip()]
    header = [c.strip() for c in rows[0].replace(",", " ").split()]
    data = [[c.strip() for c in r.replace(",", " ").split()] for r in rows[1:]]
    return (header, data)


def run(text):
    if MODEL is None:
        return f"Weights not found ({_ERR}). Set FELATAB_WEIGHTS to the repo directory."
    try:
        header, data = _parse(text)
        label_col = len(header) - 1
        feats, labels, queries = ([], [], [])
        for r in data:
            xs = [float(v) for v in r[:label_col]]
            lab = r[label_col]
            if lab in ("?", "", "-"):
                queries.append(xs)
            else:
                feats.append(xs)
                labels.append(lab)
        if not feats or not queries:
            return (
                "Provide labelled support rows and at least one query row (label '?')."
            )
    except Exception as e:
        return f"Could not parse the table: {e}"
    Xtr = np.array(feats, dtype=np.float32)
    Xte = np.array(queries, dtype=np.float32)
    classes = sorted(set(labels))
    is_cls = not all((_is_num(l) for l in labels))
    if is_cls:
        cidx = {c: i for i, c in enumerate(classes)}
        ytr = np.array([cidx[l] for l in labels], dtype=np.int64)
        probs = predict(MODEL, Xtr, ytr, Xte, "classification", n_classes=len(classes))
        out = []
        for qi, p in enumerate(probs):
            top = int(p.argmax())
            dist = "  ".join(
                (f"{classes[c]}={p[c] * 100:.0f}%" for c in range(len(classes)))
            )
            out.append(
                f"query {qi + 1}: {classes[top]}  ({p[top] * 100:.0f}% confidence)\n    {dist}"
            )
        return "\n".join(out)
    else:
        ytr = np.array([float(l) for l in labels], dtype=np.float32)
        mean, std = predict(MODEL, Xtr, ytr, Xte, "regression")
        return "\n".join(
            (
                f"query {i + 1}: {mean[i]:.2f}  (likely within +/- {1.645 * std[i]:.2f})"
                for i in range(len(mean))
            )
        )


def _is_num(s):
    try:
        float(s)
        return True
    except Exception:
        return False


with gr.Blocks(title="FelaTab playground") as demo:
    gr.Markdown(
        "# FelaTab playground\nPaste a small labelled table and a row to predict. The model learns the pattern from your examples in a single pass, with no training or setup, and fills the answer with a confidence range. Put `?` in the label column for the rows you want predicted. If the label column is numbers, it predicts a value with an error bar (regression); if it is categories, it predicts a class with probabilities. Runs on CPU."
    )
    inp = gr.Textbox(
        label="your table (comma or space separated, last column = label, '?' = predict)",
        lines=12,
        value=EXAMPLE,
    )
    out = gr.Textbox(label="prediction", lines=8)
    with gr.Row():
        gr.Button("Load example").click(lambda: EXAMPLE, outputs=inp)
        gr.Button("Predict").click(run, inputs=inp, outputs=out)
if __name__ == "__main__":
    demo.launch()