"""
gnaninet.py — Inference-only model loader for HuggingFace.

This is a self-contained, minimal inference class. No training code,
no gradient computation, no Clifford algebra — just forward pass and softmax.
"""

import json
import math
import numpy as np
from pathlib import Path


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


class GnaniNet:
    """
    GnaniNet inference-only model.

    A lightweight fully-connected classifier with ReLU activations
    and softmax output. Pure NumPy — no framework dependencies.
    """

    def __init__(self, input_dim, hidden_dims, output_dim, activation='relu'):
        self.input_dim = input_dim
        self.hidden_dims = list(hidden_dims)
        self.output_dim = output_dim
        self.activation = activation

        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 = []
        self.bs = []
        for fan_in, fan_out in self.layer_dims:
            self.Ws.append(np.zeros((fan_out, fan_in)))
            self.bs.append(np.zeros(fan_out))

    @classmethod
    def from_pretrained(cls, path):
        """Load model from a directory containing weights.txt and meta.json."""
        path = Path(path)
        with open(path / 'meta.json') as f:
            meta = json.load(f)

        model = cls(
            input_dim=meta['input_dim'],
            hidden_dims=meta['hidden_dims'],
            output_dim=meta['output_dim'],
            activation=meta.get('activation', 'relu'),
        )

        with open(path / 'weights.txt') as f:
            params = np.array([float(x) for x in f.read().split()])

        idx = 0
        for i, (fan_in, fan_out) in enumerate(model.layer_dims):
            n_W = fan_out * fan_in
            model.Ws[i] = params[idx:idx + n_W].reshape(fan_out, fan_in)
            idx += n_W
            model.bs[i] = params[idx:idx + fan_out].copy()
            idx += fan_out

        return model

    def forward(self, x):
        """Forward pass. Returns pre-softmax logits."""
        h = np.asarray(x, dtype=np.float64)
        for i, (W, b) in enumerate(zip(self.Ws, self.bs)):
            h = W @ h + b
            if i < self.n_layers - 1:
                if self.activation == 'relu':
                    h = np.maximum(0.0, h)
                else:
                    h = np.tanh(h)
        return h

    def predict_proba(self, x):
        """Return class probabilities."""
        return _softmax(self.forward(x))

    def predict(self, features, class_names=None):
        """
        High-level predict from human-readable transaction features.

        Parameters
        ----------
        features : list or dict
            If list: [amount_ratio, hour, day_of_week, location_delta,
                      velocity_1h, velocity_24h, merchant_risk,
                      international, card_present, device_match,
                      account_age_days, prev_fraud_score]
            If dict: keys matching the feature names above
        class_names : list, optional
            Names for output classes (default: class_0, class_1, ...)

        Returns
        -------
        dict with class name → probability
        """
        if isinstance(features, dict):
            raw = features
        else:
            keys = [
                'amount_ratio', 'hour', 'day_of_week', 'location_delta',
                'velocity_1h', 'velocity_24h', 'merchant_risk',
                'international', 'card_present', 'device_match',
                'account_age_days', 'prev_fraud_score',
            ]
            raw = dict(zip(keys, features))

        x = self._normalize(raw)
        proba = self.predict_proba(x)

        if class_names is None:
            class_names = [f'class_{i}' for i in range(self.output_dim)]

        return {name: round(float(proba[i]), 4) for i, name in enumerate(class_names)}

    @staticmethod
    def _normalize(raw):
        """Convert human-readable features to 14-dim normalized vector."""
        h_rad = 2 * math.pi * raw['hour'] / 24
        d_rad = 2 * math.pi * raw['day_of_week'] / 7
        return np.array([
            float(raw['amount_ratio']),
            math.sin(h_rad), math.cos(h_rad),
            math.sin(d_rad), math.cos(d_rad),
            float(raw['location_delta']),
            min(raw['velocity_1h'] / 10.0, 1.0),
            min(raw['velocity_24h'] / 30.0, 1.0),
            float(raw['merchant_risk']),
            1.0 if raw['international'] else 0.0,
            1.0 if raw['card_present'] else 0.0,
            1.0 if raw['device_match'] else 0.0,
            min(raw['account_age_days'] / 3650.0, 1.0),
            float(raw['prev_fraud_score']),
        ], dtype=np.float64)

    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'GnaniNet({arch}, {self.param_count():,} params)'