inference.py

"""
inference.py — Self-contained model loader for KestrelNet/GoshawkNet benchmarks.

Pure NumPy. No framework dependencies. Supports both standard FC (KestrelNet)
and multivector product (GoshawkNet) architectures.

Usage:
    from inference import load_model
    model = load_model("ecg-heartbeat")
    proba = model.predict_proba(x)
"""

import numpy as np
from pathlib import Path

ROOT = Path(__file__).resolve().parent

# ── Model configs (architecture used for each benchmark) ────────────────────

CONFIGS = {
    "ecg-heartbeat": {
        "input_dim": 187,
        "hidden_dims": [16, 8],
        "output_dim": 5,
        "algebra": (0, 2),  # Cl(0,2) quaternion
        "class_names": ["Normal", "Supraventricular", "Ventricular", "Fusion", "Unknown"],
    },
    "eeg-emotions": {
        "input_dim": 2548,
        "hidden_dims": [16, 8],
        "output_dim": 3,
        "algebra": (0, 2),
        "class_names": ["Negative", "Neutral", "Positive"],
    },
    "eye-state": {
        "input_dim": 14,
        "hidden_dims": [16, 8],
        "output_dim": 2,
        "algebra": (0, 2),
        "class_names": ["Eyes Open", "Eyes Closed"],
    },
    "seizure-prediction": {
        "input_dim": 178,
        "hidden_dims": [16, 8],
        "output_dim": 2,
        "algebra": (0, 2),
        "class_names": ["Non-seizure", "Seizure"],
    },
    "har-smartphones": {
        "input_dim": 228,
        "hidden_dims": [16, 8],
        "output_dim": 6,
        "algebra": (0, 2),
        "class_names": ["Walking", "Walking Upstairs", "Walking Downstairs",
                        "Sitting", "Standing", "Laying"],
    },
}


# ── Clifford algebra (inference-only, minimal) ─────────────────────────────

class _CliffordAlgebra:
    """Minimal Cl(p,q) for inference. Precomputes Cayley tensor."""

    def __init__(self, p, q):
        self.p, self.q = p, q
        self.n = p + q
        self.dim = 1 << self.n

        self.cayley = np.zeros((self.dim, self.dim, self.dim), dtype=np.float64)
        for i in range(self.dim):
            for j in range(self.dim):
                sign, k = self._blade_product(i, j)
                self.cayley[k, i, j] = sign

        self.cayley_flat = self.cayley.reshape(self.dim * self.dim, self.dim)

    def _blade_product(self, a, b):
        n_swaps = 0
        temp = a >> 1
        while temp:
            n_swaps += bin(temp & b).count('1')
            temp >>= 1
        sign = -1 if n_swaps % 2 else 1
        common = a & b
        for i in range(self.n):
            if (common >> i) & 1 and i >= self.p:
                sign = -sign
        return sign, a ^ b


# ── Softmax ────────────────────────────────────────────────────────────────

def _softmax(logits):
    m = np.max(logits)
    e = np.exp(logits - m)
    return e / e.sum()


# ── GoshawkNet (inference-only) ────────────────────────────────────────────

class GoshawkNet:
    """Multivector product neural network — inference only."""

    def __init__(self, input_dim, hidden_dims, output_dim, p=0, q=2):
        self.input_dim = input_dim
        self.hidden_dims = list(hidden_dims)
        self.output_dim = output_dim

        self.algebra = _CliffordAlgebra(p, q)
        self.D = self.algebra.dim

        dims = [input_dim] + list(hidden_dims) + [output_dim]
        self.layer_dims = list(zip(dims[:-1], dims[1:]))
        self.n_layers = len(self.layer_dims)

        self.Ws = [np.zeros((fo, fi, self.D)) for fi, fo in self.layer_dims]
        self.bs = [np.zeros((fo, self.D)) for _, fo in self.layer_dims]

    def set_params(self, v):
        idx = 0
        for l, (fi, fo) in enumerate(self.layer_dims):
            n_W = fo * fi * self.D
            self.Ws[l] = v[idx:idx + n_W].reshape(fo, fi, self.D)
            idx += n_W
            n_b = fo * self.D
            self.bs[l] = v[idx:idx + n_b].reshape(fo, self.D)
            idx += n_b

    def predict_proba(self, x):
        x = np.asarray(x, dtype=np.float64)
        D = self.D
        cf = self.algebra.cayley_flat

        # Lift input to scalar multivectors
        h = np.zeros((self.input_dim, D))
        h[:, 0] = x

        for l in range(self.n_layers):
            W, b = self.Ws[l], self.bs[l]
            fo, fi = W.shape[0], W.shape[1]

            Rh = (h @ cf.T).reshape(fi, D, D)
            Rh_mat = Rh.transpose(0, 2, 1).reshape(fi * D, D)
            W_mat = W.reshape(fo, fi * D)
            z = W_mat @ Rh_mat + b

            if l < self.n_layers - 1:
                h = np.maximum(0.0, z)
            else:
                h = z

        return _softmax(h[:, 0])

    def predict(self, x):
        return int(np.argmax(self.predict_proba(x)))

    def param_count(self):
        return sum(W.size + b.size for W, b in zip(self.Ws, self.bs))

    def __repr__(self):
        dims = [self.input_dim] + self.hidden_dims + [self.output_dim]
        arch = ' > '.join(str(d) for d in dims)
        return f'GoshawkNet({arch}, Cl({self.algebra.p},{self.algebra.q}), {self.param_count():,} params)'


# ── Loader ─────────────────────────────────────────────────────────────────

def load_model(name):
    """
    Load a benchmark model by name.

    Parameters
    ----------
    name : str
        One of: 'ecg-heartbeat', 'eeg-emotions', 'eye-state',
                'seizure-prediction', 'har-smartphones'

    Returns
    -------
    model : GoshawkNet with loaded weights
    """
    if name not in CONFIGS:
        available = ', '.join(sorted(CONFIGS.keys()))
        raise ValueError(f"Unknown model '{name}'. Available: {available}")

    cfg = CONFIGS[name]
    p, q = cfg["algebra"]

    model = GoshawkNet(
        input_dim=cfg["input_dim"],
        hidden_dims=cfg["hidden_dims"],
        output_dim=cfg["output_dim"],
        p=p, q=q,
    )

    weights_path = ROOT / name / "weights.txt"
    with open(weights_path) as f:
        params = np.array([float(x) for x in f.read().split()])
    model.set_params(params)

    model.class_names = cfg["class_names"]
    return model


def list_models():
    """List available benchmark models with their configs."""
    for name, cfg in CONFIGS.items():
        p, q = cfg["algebra"]
        model = GoshawkNet(cfg["input_dim"], cfg["hidden_dims"],
                           cfg["output_dim"], p=p, q=q)
        print(f"  {name:<25} {model}  classes={cfg['output_dim']}")


if __name__ == "__main__":
    print("Available models:\n")
    list_models()

    print("\n\nQuick test — loading all models:\n")
    for name in CONFIGS:
        model = load_model(name)
        x = np.random.randn(model.input_dim)
        proba = model.predict_proba(x)
        top = model.class_names[np.argmax(proba)]
        print(f"  {name:<25} {top:<20} (prob={proba.max():.3f}, params={model.param_count():,})")