"""
convnext_sensor.py — Inference wrapper for the trained ConvNeXt-1D β-TCVAE.

Usage
-----
    sensor = ConvNextSensor(model_path)
    z_mu, z_post_std = sensor.encode_window(df_1m, end_row)
    # z_mu: (32,) float64 — latent mean for the 32-bar window ending at end_row
    # z_post_std: float   — mean posterior std (OOD indicator, >1 = wide/uncertain)

Key z-dim assignments (from convnext_query.py, ep=17 checkpoint):
    z[10]  r=+0.973  proxy_B  (instability_50 - v750_velocity)
    z[30]  r=-0.968  proxy_B  (anti-correlated)
    z[24]  r=+0.942  proxy_B
    ...10+ dims encoding proxy_B trajectory at >0.86

Architecture: ConvNeXtVAE  C_in=11  T_in=32  z_dim=32  base_ch=32  n_blocks=3
Input channels:
    ch0-3  v50/v150/v300/v750 lambda_max_velocity
    ch4    vel_div
    ch5    instability_50
    ch6    instability_150
    ch7    proxy_B  (= instability_50 - v750_lambda_max_velocity)
    ch8    dvol_btc      (ExF, broadcast constant)
    ch9    fng           (ExF, broadcast constant)
    ch10   funding_btc   (ExF, broadcast constant)
"""

import os
import sys
import json
import numpy as np

_DVAE_DIR = os.path.dirname(os.path.abspath(__file__))
if _DVAE_DIR not in sys.path:
    sys.path.insert(0, _DVAE_DIR)

from convnext_dvae import ConvNeXtVAE

FEATURE_COLS = [
    'v50_lambda_max_velocity',
    'v150_lambda_max_velocity',
    'v300_lambda_max_velocity',
    'v750_lambda_max_velocity',
    'vel_div',
    'instability_50',
    'instability_150',
]
EXF_COLS = ['dvol_btc', 'fng', 'funding_btc']
T_WIN    = 32
N_CH     = 11   # 7 FEATURE + proxy_B + 3 ExF

# z-dim index of the primary proxy_B encoding (r=+0.973)
PROXY_B_DIM = 10


class ConvNextSensor:
    """
    Stateless inference wrapper.  Thread-safe (model weights are read-only numpy).
    """

    def __init__(self, model_path: str):
        with open(model_path) as f:
            meta = json.load(f)

        arch = meta.get('architecture', {})
        self.model = ConvNeXtVAE(
            C_in     = arch.get('C_in',     N_CH),
            T_in     = arch.get('T_in',     T_WIN),
            z_dim    = arch.get('z_dim',    32),
            base_ch  = arch.get('base_ch',  32),
            n_blocks = arch.get('n_blocks', 3),
            seed     = 42,
        )
        self.model.load(model_path)

        self.norm_mean = np.array(meta['norm_mean'], dtype=np.float64) if 'norm_mean' in meta else None
        self.norm_std  = np.array(meta['norm_std'],  dtype=np.float64) if 'norm_std'  in meta else None
        self.epoch     = meta.get('epoch', '?')
        self.val_loss  = meta.get('val_loss', float('nan'))
        self.z_dim     = arch.get('z_dim', 32)

    # ── low-level: encode a (1, N_CH, T_WIN) array ──────────────────────────
    def encode_raw(self, arr: np.ndarray):
        """
        arr: (N_CH, T_WIN) float64, already in raw (un-normalised) units.
        Returns z_mu (z_dim,), z_post_std float.
        """
        x = arr[np.newaxis].astype(np.float64)          # (1, C, T)
        if self.norm_mean is not None:
            x = (x - self.norm_mean[None, :, None]) / self.norm_std[None, :, None]
            np.clip(x, -6.0, 6.0, out=x)
        z_mu, z_logvar = self.model.encode(x)            # (1, D)
        z_post_std = float(np.exp(0.5 * z_logvar).mean())
        return z_mu[0], z_post_std

    # ── high-level: encode from a 1m DataFrame row ──────────────────────────
    def encode_window(self, df_1m, end_row: int,
                      exf_dvol: float = 0., exf_fng: float = 0.,
                      exf_funding: float = 0.):
        """
        Build a (N_CH, T_WIN) window ending at end_row (inclusive) from df_1m.
        Missing columns are treated as zero.

        Parameters
        ----------
        df_1m      : DataFrame with FEATURE_COLS as columns
        end_row    : integer row index (loc-style), window = [end_row-T_WIN+1 : end_row+1]
        exf_*      : ExF scalars broadcast across the window (set to 0 if unavailable)

        Returns
        -------
        z_mu       : (z_dim,) float64
        z_post_std : float  (>1 suggests OOD regime)
        """
        start = max(0, end_row - T_WIN + 1)
        rows  = df_1m.iloc[start : end_row + 1]

        T_actual = len(rows)
        arr = np.zeros((T_WIN, N_CH - 3), dtype=np.float64)   # (T_WIN, 8)

        for i, col in enumerate(FEATURE_COLS):
            if col in rows.columns:
                vals = rows[col].values.astype(np.float64)
                arr[T_WIN - T_actual:, i] = vals

        # proxy_B = instability_50 - v750_lambda_max_velocity  (ch7)
        arr[:, 7] = arr[:, 5] - arr[:, 3]

        # ExF channels broadcast as scalar across T_WIN
        exf = np.array([exf_dvol, exf_fng, exf_funding], dtype=np.float64)
        full = np.concatenate([arr, np.tile(exf, (T_WIN, 1))], axis=1)  # (T_WIN, 11)

        return self.encode_raw(full.T)   # (N_CH, T_WIN)

    # ── convenience scalar: primary proxy_B z-dim ────────────────────────────
    def z_proxy_b(self, df_1m, end_row: int, **exf_kwargs) -> float:
        """Return scalar z[PROXY_B_DIM] for the window ending at end_row."""
        z_mu, _ = self.encode_window(df_1m, end_row, **exf_kwargs)
        return float(z_mu[PROXY_B_DIM])