"""FlawX -- turn any pickled/saved model into an interactive visual explorer.

One self-contained HTML file, three toggleable views:

  1. ARCHITECTURE  -- the standard structural diagram (Netron / plot_model style)
  2. X-RAY         -- health diagnostics: neuron heat strips, dead units,
                      bias overshoot, saturation, layer capacity (rank use)
  3. GALAXY        -- the whole network drawn neuron-by-neuron: heat-mapped
                      nodes, weight "strings" colored by sign, pan/zoom,
                      click any neuron/tree-node for every parameter in it

Supported inputs:
  * Keras models        (.keras / .h5 / pickled)  -> full three-view support
  * sklearn tree models (DecisionTree*, RandomForest*, GradientBoosting*)
                        -> interactive heat-mapped tree with per-node info
  * sklearn linear models (coef_ / intercept_) -> drawn as a one-layer network
  * anything else pickled -> parameter card (graceful fallback)

Usage:
    python flawx.py MODEL.pkl [--data X.npy|.npz] [--features names.txt]
                                 [--sample 512] [--out flawx.html]

    from flawx import flawx
    flawx(model, x=X_train, feature_names=x_cols)   # notebook use

Only needs numpy (+ keras and/or sklearn for those model types).
"""

from __future__ import annotations

import argparse
import json
import os
import pickle
import sys

import numpy as np

try:
    from model_xray import (analyze_weights, analyze_activations,
                            input_feature_drive, _np, _activation_name)
    HAVE_XRAY = True
except ImportError:
    HAVE_XRAY = False

TOP_EDGES = 8          # strongest incoming weights kept per neuron
MAX_TREES = 40         # forest estimators exposed in the picker


# ----------------------------------------------------------------- extractors

def _model_kind(model):
    if hasattr(model, "layers") and hasattr(model, "count_params"):
        return "keras"
    if hasattr(model, "tree_"):
        return "sk_tree"
    if hasattr(model, "estimators_"):
        return "sk_forest"
    if hasattr(model, "coef_"):
        return "sk_linear"
    return "generic"


def extract_keras(model, x=None, sample=512, feature_names=None):
    if not HAVE_XRAY:
        raise RuntimeError("model_xray.py must sit next to flawx.py")
    w_rows = {r["layer"]: r for r in analyze_weights(model)}
    a_rows = {}
    if x is not None:
        a_rows = {r["layer"]: r for r in analyze_activations(model, x, sample)}
    drive = input_feature_drive(model, x) if x is not None else None

    # adjacency from tensor ownership (functional / sequential)
    owner = {}
    for L in model.layers:
        try:
            outs = L.output if isinstance(L.output, (list, tuple)) else [L.output]
            for o in outs:
                owner[id(o)] = L.name
        except Exception:
            pass
    inbound = {}
    for L in model.layers:
        srcs = []
        try:
            ins = L.input if isinstance(L.input, (list, tuple)) else [L.input]
            for t in ins:
                s = owner.get(id(t))
                if s and s != L.name:
                    srcs.append(s)
        except Exception:
            pass
        inbound[L.name] = srcs

    layers = []
    for L in model.layers:
        cls = L.__class__.__name__
        w = w_rows.get(L.name, {})
        a = a_rows.get(L.name)
        try:
            oshape = list(L.output.shape[1:])
        except Exception:
            oshape = None
        entry = {
            "name": L.name, "type": cls,
            "activation": _activation_name(L),
            "out_shape": oshape,
            "params": w.get("params", 0),
            "param_share": w.get("param_share", 0),
            "inbound": inbound.get(L.name, []),
            "flags": list(w.get("flags", [])) + (list(a["flags"]) if a else []),
            "w_stats": {k: w[k] for k in
                        ("w_mean", "w_std", "w_max_abs", "w_near_zero",
                         "eff_rank", "rank_util", "b_mean", "b_std",
                         "b_max_abs", "b_ratio_max", "b_overshoot_frac")
                        if k in w},
        }
        if a is not None:
            entry["act"] = {
                "mean": a["act_mean"], "std": a["act_std"],
                "sat": a["sat_overall"],
                "heat": [round(float(v), 5) for v in a["mean_abs"]],
                "dead": [bool(b) for b in a["dead_mask"]],
            }
        # neuron-level detail for Dense layers
        if cls == "Dense" and L.weights:
            K = _np(L.weights[0])
            b = _np(L.weights[1]) if len(L.weights) > 1 else np.zeros(K.shape[1])
            order = np.argsort(-np.abs(K), axis=0)[:TOP_EDGES]  # per column
            entry["neurons"] = {
                "bias": [round(float(v), 5) for v in b],
                "w_l1": [round(float(v), 4) for v in np.abs(K).sum(axis=0)],
                "top_in": [[[int(order[k, u]),
                             round(float(K[order[k, u], u]), 5)]
                            for k in range(min(TOP_EDGES, K.shape[0]))]
                           for u in range(K.shape[1])],
                "fan_in": int(K.shape[0]),
            }
        layers.append(entry)

    payload = {
        "kind": "network", "name": model.name,
        "params": int(model.count_params()),
        "layers": layers,
        "feature_names": list(feature_names) if feature_names else None,
        "drive": ([round(float(v), 5) for v in drive]
                  if drive is not None else None),
        "has_data": x is not None,
    }
    return payload


def extract_linear(model, feature_names=None):
    coef = np.atleast_2d(np.asarray(model.coef_))          # (targets, features)
    inter = np.atleast_1d(np.asarray(getattr(model, "intercept_", np.zeros(coef.shape[0]))))
    n_f, n_t = coef.shape[1], coef.shape[0]
    K = coef.T                                             # (features, targets)
    order = np.argsort(-np.abs(K), axis=0)[:TOP_EDGES]
    classes = [str(c) for c in getattr(model, "classes_", range(n_t))]
    layers = [
        {"name": "input", "type": "InputLayer", "activation": None,
         "out_shape": [n_f], "params": 0, "param_share": 0,
         "inbound": [], "flags": [], "w_stats": {}},
        {"name": "coefficients", "type": "Dense",
         "activation": "sigmoid/softmax",
         "out_shape": [n_t], "params": int(K.size + n_t),
         "param_share": 1.0, "inbound": ["input"],
         "flags": [], "w_stats": {"w_mean": float(K.mean()),
                                  "w_std": float(K.std()),
                                  "w_max_abs": float(np.abs(K).max())},
         "neurons": {
             "bias": [round(float(v), 5) for v in inter],
             "w_l1": [round(float(v), 4) for v in np.abs(K).sum(axis=0)],
             "top_in": [[[int(order[k, u]), round(float(K[order[k, u], u]), 5)]
                         for k in range(min(TOP_EDGES, n_f))]
                        for u in range(n_t)],
             "fan_in": n_f}},
    ]
    return {"kind": "network", "name": type(model).__name__,
            "params": int(K.size + n_t), "layers": layers,
            "feature_names": list(feature_names) if feature_names else None,
            "unit_names": {"coefficients": classes},
            "drive": [round(float(v), 5) for v in
                      (np.abs(K).sum(axis=1) / max(1e-12, np.abs(K).sum()))],
            "has_data": False}


def _extract_one_tree(tree, feature_names, classes):
    t = tree.tree_
    nodes = []
    for i in range(t.node_count):
        val = np.asarray(t.value[i]).ravel()
        tot = val.sum()
        nodes.append({
            "id": i,
            "left": int(t.children_left[i]), "right": int(t.children_right[i]),
            "feature": (feature_names[t.feature[i]] if feature_names and
                        t.feature[i] >= 0 else
                        (f"f{t.feature[i]}" if t.feature[i] >= 0 else None)),
            "threshold": (round(float(t.threshold[i]), 4)
                          if t.feature[i] >= 0 else None),
            "impurity": round(float(t.impurity[i]), 4),
            "samples": int(t.n_node_samples[i]),
            "value": [round(float(v), 3) for v in val],
            "pred": (classes[int(val.argmax())] if classes is not None and tot
                     else round(float(val[0]), 3) if len(val) else None),
        })
    return {"nodes": nodes, "max_depth": int(t.max_depth),
            "n_leaves": int(t.n_leaves)}


def extract_tree(model, feature_names=None):
    classes = ([str(c) for c in model.classes_]
               if hasattr(model, "classes_") else None)
    return {"kind": "tree", "name": type(model).__name__,
            "criterion": getattr(model, "criterion", None),
            "classes": classes,
            "trees": [_extract_one_tree(model, feature_names, classes)],
            "feature_names": list(feature_names) if feature_names else None}


def extract_forest(model, feature_names=None):
    classes = ([str(c) for c in model.classes_]
               if hasattr(model, "classes_") else None)
    ests = list(np.asarray(model.estimators_).ravel())[:MAX_TREES]
    return {"kind": "tree", "name": type(model).__name__,
            "criterion": getattr(model, "criterion", None),
            "classes": classes,
            "trees": [_extract_one_tree(e, feature_names, classes)
                      for e in ests],
            "n_total_trees": len(np.asarray(model.estimators_).ravel()),
            "feature_names": list(feature_names) if feature_names else None}


def extract_generic(model):
    params = None
    if hasattr(model, "get_params"):
        try:
            params = {k: repr(v) for k, v in model.get_params().items()}
        except Exception:
            pass
    return {"kind": "generic", "name": type(model).__name__,
            "repr": repr(model)[:4000], "params": params}


# ------------------------------------------------------------------ html app

def build_html(payload, out_path, data_note=""):
    here = os.path.dirname(os.path.abspath(__file__))
    tpl_path = os.path.join(here, "flawx_app.html")
    with open(tpl_path, encoding="utf-8") as f:
        tpl = f.read()
    blob = json.dumps(payload, separators=(",", ":"))
    html = tpl.replace("__DATA_NOTE__", data_note).replace(
        "__PAYLOAD__", blob.replace("</", "<\\/"))
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(html)
    return out_path


def flawx(model, x=None, feature_names=None, sample=512,
             out="flawx.html", json_out=None):
    kind = _model_kind(model)
    note = ""
    if kind == "keras":
        payload = extract_keras(model, x, sample, feature_names)
        if x is None:
            note = ("No sample data given -- neuron heat shows weight "
                    "magnitude, not activations. Pass --data for live heat.")
    elif kind == "sk_tree":
        payload = extract_tree(model, feature_names)
    elif kind == "sk_forest":
        payload = extract_forest(model, feature_names)
    elif kind == "sk_linear":
        payload = extract_linear(model, feature_names)
    else:
        payload = extract_generic(model)
    build_html(payload, out, note)
    if json_out:                       # droppable on the FlawX web viewer
        with open(json_out, "w", encoding="utf-8") as f:
            json.dump(payload, f, separators=(",", ":"))
        print(f"FlawX: JSON export -> {json_out}")
    print(f"FlawX: {payload['name']} ({kind}) -> {out}")
    return out


def _load_model(path):
    if path.endswith((".keras", ".h5", ".hdf5")):
        import keras
        return keras.models.load_model(path, compile=False)
    with open(path, "rb") as f:
        return pickle.load(f)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("model", help=".pkl / .keras / .h5 model file")
    ap.add_argument("--data", help=".npy or .npz sample of model inputs")
    ap.add_argument("--features", help="text file, one feature name per line")
    ap.add_argument("--sample", type=int, default=512)
    ap.add_argument("--out", help="output HTML path")
    ap.add_argument("--json", dest="json_out",
                    help="also write the payload as JSON (droppable on the "
                         "FlawX web viewer)")
    args = ap.parse_args(argv)

    model = _load_model(args.model)
    x = None
    if args.data:
        x = (np.load(args.data) if args.data.endswith(".npy")
             else [np.load(args.data)[k] for k in np.load(args.data).files])
    names = None
    if args.features:
        with open(args.features, encoding="utf-8") as f:
            names = [ln.strip() for ln in f if ln.strip()]
    out = args.out or os.path.splitext(os.path.abspath(args.model))[0] + \
        "_flawx.html"
    flawx(model, x, names, args.sample, out, json_out=args.json_out)


if __name__ == "__main__":
    sys.exit(main())

netscope = flawx  # old name, kept as an alias