"""Model X-Ray -- health, heat & complexity report for Keras models.

Answers, in one report:
  * which neurons are HOT (dominating), COLD (barely firing) or DEAD (silent)
  * which input features actually drive the first layer
  * which biases are overshooting (bias dominates the weighted sum)
  * which units are saturated (sigmoid/tanh stuck at the rails)
  * which layers hold more capacity than they use (over-parameterization
    / overfitting risk), via effective-rank analysis of the kernels

Usage:
    python model_xray.py MODEL.keras [--data X.npy | X.npz] [--out report.html]
                                     [--sample 512]

  --data   sample of *inputs* the model was trained/evaluated on.
           .npy for single-input models; .npz for multi-input models
           (arrays matched to inputs by name, else by order).
           Without data you still get the full static (weights/bias/rank)
           analysis; activation heat needs data.
  --sample max number of rows drawn from the data (default 512).
  --out    HTML report path (default: xray_report.html next to the model).

Or from Python / a notebook:
    from model_xray import xray
    xray(model, x_sample)            # prints summary, writes xray_report.html

Requires only: keras (3.x, any backend) or tf.keras, and numpy.
"""

from __future__ import annotations

import argparse
import html as _html
import math
import os
import sys

import numpy as np

try:
    import keras
except ImportError:  # fall back to the Keras bundled with TensorFlow
    from tensorflow import keras  # type: ignore

# ---------------------------------------------------------------- thresholds
NEAR_ZERO_W = 1e-3      # |w| below this counts as a near-zero weight
DEAD_MAX_ABS = 1e-7     # unit never exceeding this (on data) is DEAD
COLD_PCTL = 10          # bottom X% of units by mean |activation| -> COLD
HOT_PCTL = 90           # top X% -> HOT
BIAS_RATIO_UNIT = 2.0   # |b| / (w_std * sqrt(fan_in)) above this -> overshoot
BIAS_RATIO_HARD = 5.0
SAT_LEVEL = 0.99        # tanh |a|>0.99 / sigmoid a<0.01 or >0.99 -> saturated
SAT_UNIT_FRAC = 0.95    # unit saturated on >95% of samples -> stuck
RANK_UTIL_LOW = 0.30    # effective rank / full rank below this -> oversized
RANK_MIN_DIM = 16       # only judge rank utilization on layers this wide

RELU_FAMILY = {"relu", "relu6", "leaky_relu", "elu", "selu", "gelu",
               "silu", "swish", "hard_silu", "hard_swish", "mish"}
SAT_SIGMOID = {"sigmoid", "hard_sigmoid"}
SAT_TANH = {"tanh", "hard_tanh"}


def _np(x):
    try:
        return np.asarray(keras.ops.convert_to_numpy(x))
    except Exception:
        return np.asarray(x)


def _walk_layers(container, prefix=""):
    """Yield (qualified_name, layer) for every leaf layer, entering submodels."""
    for layer in container.layers:
        name = f"{prefix}{layer.name}"
        if hasattr(layer, "layers") and layer.layers:
            yield from _walk_layers(layer, prefix=name + "/")
        else:
            yield name, layer


def _activation_name(layer):
    cls = layer.__class__.__name__.lower()
    if cls in ("relu", "leakyrelu", "elu", "prelu", "softmax"):
        return cls.replace("leakyrelu", "leaky_relu")
    cfg_act = getattr(layer, "activation", None)
    if cfg_act is not None:
        return getattr(cfg_act, "__name__", str(cfg_act)).lower()
    if cls == "activation":
        try:
            return layer.get_config()["activation"]
        except Exception:
            pass
    return None


# ------------------------------------------------------------ static analysis

def analyze_weights(model):
    """Per-layer kernel/bias statistics + rank utilization + bias overshoot."""
    rows = []
    total_params = max(1, model.count_params())
    for qname, layer in _walk_layers(model):
        if not layer.weights:
            continue
        kernel, bias = None, None
        for w in layer.weights:
            arr = _np(w)
            wname = w.name if hasattr(w, "name") else w.path
            if arr.ndim >= 2 and ("kernel" in wname or kernel is None):
                if kernel is None or "kernel" in wname:
                    kernel = arr
            elif arr.ndim == 1 and "bias" in wname:
                bias = arr

        n_params = int(sum(int(np.prod(_np(w).shape)) for w in layer.weights))
        row = {
            "layer": qname,
            "type": layer.__class__.__name__,
            "activation": _activation_name(layer),
            "params": n_params,
            "param_share": n_params / total_params,
            "flags": [],
        }

        if kernel is not None:
            k2d = kernel.reshape(-1, kernel.shape[-1])  # (fan_in*, units)
            w_std = float(k2d.std())
            row.update(
                w_mean=float(k2d.mean()), w_std=w_std,
                w_max_abs=float(np.abs(k2d).max()),
                w_near_zero=float(np.mean(np.abs(k2d) < NEAR_ZERO_W)),
                units=int(k2d.shape[1]), fan_in=int(k2d.shape[0]),
            )
            # rank utilization: how much of the layer's capacity is used
            if min(k2d.shape) >= 2:
                s = np.linalg.svd(k2d, compute_uv=False)
                s_sum = s.sum()
                if s_sum > 0:
                    p = s / s_sum
                    p = p[p > 0]
                    eff_rank = float(np.exp(-(p * np.log(p)).sum()))
                    full = min(k2d.shape)
                    row["eff_rank"] = eff_rank
                    row["rank_util"] = eff_rank / full
                    if full >= RANK_MIN_DIM and row["rank_util"] < RANK_UTIL_LOW:
                        row["flags"].append(
                            f"OVERSIZED: uses ~{eff_rank:.0f} of {full} "
                            f"dimensions ({row['rank_util']:.0%}) -- layer "
                            f"could likely be much smaller")
            if bias is not None and w_std > 0:
                expected = w_std * math.sqrt(max(1, k2d.shape[0]))
                ratios = np.abs(bias) / expected
                row.update(
                    b_mean=float(bias.mean()), b_std=float(bias.std()),
                    b_max_abs=float(np.abs(bias).max()),
                    b_overshoot_frac=float(np.mean(ratios > BIAS_RATIO_UNIT)),
                    b_ratio_max=float(ratios.max()),
                )
                if row["b_ratio_max"] > BIAS_RATIO_HARD or \
                        row["b_overshoot_frac"] > 0.10:
                    row["flags"].append(
                        f"BIAS OVERSHOOT: {row['b_overshoot_frac']:.0%} of "
                        f"units have |bias| > {BIAS_RATIO_UNIT}x the typical "
                        f"weighted-sum scale (worst {row['b_ratio_max']:.1f}x)")
        rows.append(row)
    return rows


# ----------------------------------------------------------- dynamic analysis

def _load_data(path, model):
    if path.endswith(".npz"):
        z = np.load(path)
        names = [getattr(t, "name", f"input_{i}") for i, t in
                 enumerate(model.inputs)]
        if all(n in z.files for n in names):
            return [z[n] for n in names]
        return [z[k] for k in z.files]
    return np.load(path)


def analyze_activations(model, x, sample=512):
    """Run x through every layer; per-unit heat, dead/saturated detection."""
    xs = x if isinstance(x, (list, tuple)) else [x]
    n = min(sample, len(xs[0]))
    idx = np.random.default_rng(0).choice(len(xs[0]), size=n, replace=False)
    xs = [np.asarray(a)[idx] for a in xs]

    taps, names, layer_map = [], [], []
    for layer in model.layers:
        if isinstance(layer, keras.layers.InputLayer):
            continue
        try:
            out = layer.output
        except Exception:
            continue
        outs = out if isinstance(out, (list, tuple)) else [out]
        for j, o in enumerate(outs):
            taps.append(o)
            names.append(layer.name if len(outs) == 1
                         else f"{layer.name}[{j}]")
            layer_map.append(layer)

    probe = keras.Model(model.inputs, taps)
    acts = probe.predict(xs if len(xs) > 1 else xs[0],
                         batch_size=min(64, n), verbose=0)
    if not isinstance(acts, list):
        acts = [acts]

    results = []
    for name, layer, a in zip(names, layer_map, acts):
        a = np.asarray(a)
        if a.ndim == 1:
            a = a[:, None]
        flat = a.reshape(-1, a.shape[-1])          # (samples*spatial, units)
        mean_abs = np.abs(flat).mean(axis=0)       # per-unit heat
        max_abs = np.abs(flat).max(axis=0)
        zero_frac = (flat == 0).mean(axis=0)

        act_fn = _activation_name(layer)
        dead = max_abs < DEAD_MAX_ABS
        if act_fn in RELU_FAMILY or (flat >= 0).all():
            dead = dead | (zero_frac > 0.999)

        sat_frac = None
        if act_fn in SAT_TANH:
            sat_frac = (np.abs(flat) > SAT_LEVEL).mean(axis=0)
        elif act_fn in SAT_SIGMOID:
            sat_frac = ((flat > SAT_LEVEL) | (flat < 1 - SAT_LEVEL)).mean(axis=0)
        stuck = int((sat_frac > SAT_UNIT_FRAC).sum()) if sat_frac is not None else 0

        lo, hi = np.percentile(mean_abs, [COLD_PCTL, HOT_PCTL])
        res = {
            "layer": name, "type": layer.__class__.__name__,
            "activation": act_fn, "shape": list(a.shape[1:]),
            "units": int(flat.shape[1]),
            "mean_abs": mean_abs, "act_mean": float(flat.mean()),
            "act_std": float(flat.std()),
            "dead_n": int(dead.sum()), "dead_mask": dead,
            "cold_n": int(((mean_abs <= lo) & ~dead).sum()),
            "hot_n": int((mean_abs >= hi).sum()),
            "stuck_n": stuck,
            "sat_overall": float(sat_frac.mean()) if sat_frac is not None else None,
            "flags": [],
        }
        if res["dead_n"]:
            pct = res["dead_n"] / res["units"]
            res["flags"].append(
                f"DEAD UNITS: {res['dead_n']}/{res['units']} ({pct:.0%}) "
                f"never fire on this data" +
                (" -- wasted capacity / dying-ReLU" if pct > 0.2 else ""))
        if stuck:
            res["flags"].append(
                f"SATURATED: {stuck}/{res['units']} units pinned at the "
                f"{act_fn} rails on >{SAT_UNIT_FRAC:.0%} of samples "
                f"(gradient starvation)")
        results.append(res)
    return results


def input_feature_drive(model, x):
    """For a Dense first layer on 2-D input: how hard each input feature
    drives the network = std(feature) * L1 norm of its outgoing weights."""
    if isinstance(x, (list, tuple)):
        return None
    x = np.asarray(x)
    if x.ndim != 2:
        return None
    for layer in model.layers:
        if isinstance(layer, keras.layers.InputLayer):
            continue
        if layer.__class__.__name__ == "Dense" and layer.weights:
            W = _np(layer.weights[0])
            if W.shape[0] != x.shape[1]:
                return None
            drive = x.std(axis=0) * np.abs(W).sum(axis=1)
            tot = drive.sum()
            return drive / tot if tot > 0 else drive
        return None  # first real layer isn't Dense
    return None


# ------------------------------------------------------------------ reporting

def _heat_color(v):
    """v in [0,1] -> blue(cold) .. red(hot) css color."""
    hue = 215 - 205 * min(1.0, max(0.0, v))
    return f"hsl({hue:.0f},80%,{55 - 10 * v:.0f}%)"


def _strip(values, dead_mask=None, cap=512):
    """Row of colored cells, one per unit."""
    vals = np.asarray(values, dtype=float)
    shown = vals[:cap]
    vmax = vals.max() if vals.max() > 0 else 1.0
    cells = []
    for i, v in enumerate(shown):
        if dead_mask is not None and dead_mask[i]:
            c, title = "#999", f"unit {i}: DEAD"
        else:
            c = _heat_color(v / vmax)
            title = f"unit {i}: mean|a|={v:.4g}"
        cells.append(f'<i style="background:{c}" title="{title}"></i>')
    note = (f'<span class="note">first {cap} of {len(vals)} units</span>'
            if len(vals) > cap else "")
    return f'<div class="strip">{"".join(cells)}</div>{note}'


def _bar(frac, color="#4a7"):
    return (f'<div class="bar"><b style="width:{min(100, frac * 100):.1f}%;'
            f'background:{color}"></b></div>')


def render_html(model, w_rows, a_rows, drive, out_path):
    e = _html.escape
    a_by_name = {r["layer"]: r for r in (a_rows or [])}
    css = """
    body{font:14px/1.5 system-ui,sans-serif;margin:24px auto;max-width:1100px;
         color:#1a1a1a;background:#fafafa;padding:0 16px}
    h1{font-size:22px} h2{font-size:17px;margin-top:32px}
    table{border-collapse:collapse;width:100%;background:#fff;font-size:13px}
    th,td{border:1px solid #ddd;padding:5px 8px;text-align:right;
          white-space:nowrap}
    th{background:#f0f0f0} td:first-child,th:first-child{text-align:left}
    .flag{display:inline-block;background:#fde8e8;color:#b42318;border:1px solid
          #f5b5b1;border-radius:4px;padding:1px 7px;margin:2px 4px 2px 0;
          font-size:12px}
    .ok{color:#1a7f37;font-size:12px}
    .card{background:#fff;border:1px solid #e2e2e2;border-radius:8px;
          padding:12px 16px;margin:10px 0}
    .strip{display:flex;flex-wrap:wrap;gap:1px;margin:6px 0}
    .strip i{width:7px;height:16px;display:block}
    .bar{background:#eee;border-radius:3px;height:9px;width:140px;
         display:inline-block;vertical-align:middle}
    .bar b{display:block;height:100%;border-radius:3px}
    .note,.sub{color:#777;font-size:12px}
    .legend i{display:inline-block;width:12px;height:12px;border-radius:2px;
              vertical-align:-2px;margin:0 4px 0 12px}
    code{background:#eee;padding:0 4px;border-radius:3px}
    """
    total_flags = sum(len(r["flags"]) for r in w_rows) + \
        sum(len(r["flags"]) for r in (a_rows or []))
    parts = [f"<title>Model X-Ray: {e(model.name)}</title><style>{css}</style>",
             f"<h1>Model X-Ray &mdash; <code>{e(model.name)}</code></h1>",
             f"<p>{model.count_params():,} parameters &middot; "
             f"{len(list(_walk_layers(model)))} layers &middot; "
             f"{total_flags} issue(s) flagged</p>",
             '<p class="legend">Heat scale: <i style="background:'
             f'{_heat_color(0)}"></i>cold <i style="background:'
             f'{_heat_color(0.5)}"></i>warm <i style="background:'
             f'{_heat_color(1)}"></i>hot <i style="background:#999"></i>dead'
             '</p>']

    # ---- input feature drive
    if drive is not None:
        order = np.argsort(drive)[::-1]
        parts.append("<h2>Input feature drive</h2>"
                     '<p class="sub">std(feature) &times; &Sigma;|outgoing '
                     "weights| into the first Dense layer &mdash; which raw "
                     "features the network actually listens to.</p>"
                     '<div class="card">' + _strip(drive) +
                     "<p class='sub'>strongest: " +
                     ", ".join(f"f{int(i)} ({drive[i]:.1%})"
                               for i in order[:8]) +
                     " &middot; weakest: " +
                     ", ".join(f"f{int(i)} ({drive[i]:.2%})"
                               for i in order[-5:]) + "</p></div>")

    # ---- weights & complexity table
    parts.append("<h2>Weights, biases &amp; capacity</h2>"
                 "<table><tr><th>layer</th><th>type</th><th>params</th>"
                 "<th>share</th><th>w std</th><th>near-zero w</th>"
                 "<th>max|bias|</th><th>bias ratio</th><th>rank use</th>"
                 "<th>issues</th></tr>")
    for r in w_rows:
        flags = " ".join(f'<span class="flag">{e(f)}</span>'
                         for f in r["flags"]) or '<span class="ok">ok</span>'
        parts.append(
            "<tr><td>{}</td><td>{}</td><td>{:,}</td>"
            "<td>{}{:.1%}</td><td>{}</td><td>{}</td><td>{}</td><td>{}</td>"
            "<td>{}</td><td style='text-align:left'>{}</td></tr>".format(
                e(r["layer"]), e(r["type"]), r["params"],
                _bar(r["param_share"], "#68a"), r["param_share"],
                f"{r['w_std']:.3f}" if "w_std" in r else "&ndash;",
                f"{r['w_near_zero']:.0%}" if "w_near_zero" in r else "&ndash;",
                f"{r['b_max_abs']:.3f}" if "b_max_abs" in r else "&ndash;",
                f"{r['b_ratio_max']:.1f}x" if "b_ratio_max" in r else "&ndash;",
                (f"{r['rank_util']:.0%}" if "rank_util" in r else "&ndash;"),
                flags))
    parts.append("</table>"
                 '<p class="sub"><b>share</b>: fraction of all model '
                 "parameters. <b>bias ratio</b>: worst |bias| vs the typical "
                 "weighted-sum magnitude (&gt;2 means bias dominates the "
                 "signal). <b>rank use</b>: effective rank / full rank of the "
                 "kernel &mdash; low values mean the layer learned far fewer "
                 "directions than it has, i.e. it is oversized for the "
                 "problem (shrink it or regularize).</p>")

    # ---- activations
    if a_rows:
        parts.append("<h2>Activation heat (per unit, on your data)</h2>")
        for r in a_rows:
            flags = " ".join(f'<span class="flag">{e(f)}</span>'
                             for f in r["flags"])
            sat = (f" &middot; saturation {r['sat_overall']:.0%}"
                   if r["sat_overall"] is not None else "")
            parts.append(
                f'<div class="card"><b>{e(r["layer"])}</b> '
                f'<span class="sub">{e(r["type"])}'
                f'{" / " + e(r["activation"]) if r["activation"] else ""} '
                f'&middot; out {r["shape"]} &middot; mean {r["act_mean"]:.3f} '
                f'&middot; std {r["act_std"]:.3f}{sat} &middot; '
                f'hot {r["hot_n"]} / cold {r["cold_n"]} / dead {r["dead_n"]}'
                f'</span><br>{flags}'
                + _strip(r["mean_abs"], r["dead_mask"]) + "</div>")
    else:
        parts.append("<h2>Activation heat</h2><p>No sample data was given "
                     "&mdash; rerun with <code>--data X.npy</code> (a few "
                     "hundred training rows) to light up per-neuron heat, "
                     "dead-unit and saturation detection.</p>")

    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n".join(parts))
    return out_path


def print_summary(model, w_rows, a_rows, out_path):
    print(f"\nModel X-Ray: {model.name} "
          f"({model.count_params():,} params)")
    print("-" * 64)
    any_flag = False
    for r in w_rows:
        for f in r["flags"]:
            any_flag = True
            print(f"  [!] {r['layer']:<28} {f}")
    for r in (a_rows or []):
        for f in r["flags"]:
            any_flag = True
            print(f"  [!] {r['layer']:<28} {f}")
    if not any_flag:
        print("  no structural issues flagged")
    if a_rows:
        hottest = max(a_rows, key=lambda r: np.max(r["mean_abs"]))
        print(f"  [i] hottest layer: {hottest['layer']} "
              f"(peak unit mean|a| = {np.max(hottest['mean_abs']):.3g})")
    big = max(w_rows, key=lambda r: r["param_share"], default=None)
    if big:
        print(f"  [i] heaviest layer: {big['layer']} "
              f"({big['param_share']:.0%} of all parameters)")
    print(f"\n  full report -> {out_path}\n")


# ----------------------------------------------------------------- entrypoint

def xray(model, x=None, sample=512, out="xray_report.html"):
    """Programmatic entry point. Returns path of the HTML report."""
    w_rows = analyze_weights(model)
    a_rows = analyze_activations(model, x, sample) if x is not None else None
    drive = input_feature_drive(model, x) if x is not None else None
    render_html(model, w_rows, a_rows, drive, out)
    print_summary(model, w_rows, a_rows, out)
    return out


def main(argv=None):
    ap = argparse.ArgumentParser(
        description="X-ray a saved Keras model: neuron heat, dead units, "
                    "bias overshoot, saturation, layer capacity.")
    ap.add_argument("model", help="path to .keras / .h5 / .pkl model file")
    ap.add_argument("--data", help=".npy (single input) or .npz (multi input)"
                                   " sample of model inputs")
    ap.add_argument("--sample", type=int, default=512,
                    help="max rows sampled from --data (default 512)")
    ap.add_argument("--out", help="HTML report path")
    args = ap.parse_args(argv)

    if args.model.endswith(".pkl"):
        import pickle
        with open(args.model, "rb") as f:
            model = pickle.load(f)
    else:
        model = keras.models.load_model(args.model, compile=False)
    x = _load_data(args.data, model) if args.data else None
    out = args.out or os.path.join(
        os.path.dirname(os.path.abspath(args.model)), "xray_report.html")
    xray(model, x, sample=args.sample, out=out)


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