src/pino/pimt_model.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch.utils.data import Dataset

from .embeddings import OlfactoryEmbeddingEngine, get_openpom_embedding
from .semantics import get_objective_targets, sample_noisy_psychometric_targets

logger = logging.getLogger("pino.pimt_model")

DEFAULT_DATASET = Path(__file__).parent.parent.parent / "data" / "synthetic_dataset_v2.jsonl"

# The default token dimension is 151-D: 138-D structural + 7-D 3D shape + 2-D physics +
# 4-D functional-group flags.
DEFAULT_EMBEDDING_DIM = 151


def objective_targets_to_pyramid(targets: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
    """
    Normalize objective labels to the current static pyramid shape (3, 138).

    Older records store a time-resolved dry-down trajectory (T, 138).  The
    pyramid head predicts top/middle/base tiers, so legacy trajectories are
    folded into early/middle/late windows using the strongest descriptor
    evidence in each window.
    """
    if targets.ndim != 2:
        raise ValueError(f"objective targets must be 2-D, got shape {tuple(targets.shape)}")
    if targets.shape[0] == 3:
        return targets
    width = targets.shape[1]
    if width not in (138, 575):
        raise ValueError(f"objective targets must have 138 (pyramid) or 575 (all-tags) width, got {width}")

    if isinstance(targets, torch.Tensor):
        chunks = torch.tensor_split(targets, 3, dim=0)
        return torch.stack([chunk.max(dim=0).values for chunk in chunks], dim=0)

    chunks = np.array_split(targets, 3, axis=0)
    return np.stack([chunk.max(axis=0) for chunk in chunks], axis=0)


def _load_poucher_substantivity_targets() -> dict[str, float]:
    """Measured Poucher log10 substantivity coefficients keyed by CAS."""
    path = Path(__file__).parent.parent.parent / "data" / "poucher_substantivity.jsonl"
    if not path.exists():
        return {}
    targets: dict[str, float] = {}
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            row = json.loads(line)
            coeff = float(row["poucher_coefficient"])
            targets[str(row["cas"])] = float(np.log10(coeff))
    return targets


class FragranceTrajectoryDataset(Dataset):
    """
    Streaming PyTorch dataset over the PINO synthetic trajectory corpus.

    Each record is parsed line-by-line from the JSON-L file. For every
    ingredient we look up a static 151-D molecular embedding (138-D structural
    + 7-D 3D shape + 2-D physics + 4-D functional-group flags) and combine it with the time-resolved physical state
    vectors to build a token tensor.
    """

    def __init__(
        self,
        data_path: str | Path | None = DEFAULT_DATASET,
        records: list[dict[str, Any]] | None = None,
        max_ingredients: int = 32,
        embedding_dim: int = DEFAULT_EMBEDDING_DIM,
        state_dim: int = 4,
        use_embedding_fallback: bool = True,
        structural_source: str = "morgan",
        label_noise: bool = True,
        genre_map: dict[str, int] | None = None,
        objective_dim: int = 138,
        use_gamma: bool = False,
    ) -> None:
        self.data_path = Path(data_path) if data_path else None
        self.records = records
        self.max_ingredients = max_ingredients
        self.embedding_dim = embedding_dim
        self.state_dim = state_dim
        self.objective_dim = objective_dim
        self.use_gamma = use_gamma
        self.embedding_engine = OlfactoryEmbeddingEngine(use_fallback=use_embedding_fallback, structural_source=structural_source)
        self.embedding_dim = self.embedding_engine.embedding_dim
        self.label_noise = label_noise
        self.genre_map = genre_map or {}
        self.poucher_substantivity = _load_poucher_substantivity_targets()
        self._index_offsets: list[int] = []
        if self.records is None:
            if self.data_path is None:
                raise ValueError("Either data_path or records must be provided")
            self._build_index()
        else:
            logger.info("Using %d in-memory trajectory records", len(self.records))

    def _build_index(self) -> None:
        """Build a byte-offset index so the dataset can be randomly accessed."""
        offset = 0
        with self.data_path.open("r") as f:  # type: ignore[union-att]
            for line in f:
                self._index_offsets.append(offset)
                offset += len(line.encode("utf-8"))
        logger.info("Indexed %d trajectory records from %s", len(self._index_offsets), self.data_path)

    def __len__(self) -> int:
        return len(self._index_offsets) if self.records is None else len(self.records)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        if self.records is not None:
            record = self.records[idx]
        else:
            with self.data_path.open("r") as f:  # type: ignore[union-att]
                f.seek(self._index_offsets[idx])
                record = json.loads(f.readline())

        if "formula" not in record or "trajectory" not in record:
            raise IndexError(f"Record {idx} is not a formulation sample")
        formula = record["formula"]
        trajectory = record["trajectory"]

        # Static embeddings: one per ingredient (solvent included as a token).
        tokens: list[torch.Tensor] = []
        for comp in formula:
            smiles = comp.get("smiles", "")
            cas = comp.get("cas", "")
            if not smiles:
                smiles = self._smiles_for_cas(cas)
            z = self.embedding_engine.get_embedding(smiles, cas=cas)
            tokens.append(torch.from_numpy(z))

        # Dynamic states: time x ingredient x state, sized to the actual K.
        # Channels: [x_liquid, log10 OAV] plus an optional 3rd channel carrying
        # the UNIFAC blend-interaction term (log10 proportional activity
        # coefficient gamma), recovered by augment_gamma_channel.py. The gamma
        # channel is the NON-redundant physics: it depends on the whole mixture
        # and is not recoverable from single-molecule embeddings (unlike
        # volatility, which is ~76% MW-recoverable).
        t_steps = len(trajectory)
        n_mol = len(tokens)
        gamma_traj = record.get("physics_gamma")  # (T, S) or None
        state_dim = 3 if (self.use_gamma and gamma_traj is not None) else 2
        states = np.zeros((t_steps, n_mol, state_dim), dtype=np.float32)
        for t, step in enumerate(trajectory):
            for j, comp in enumerate(formula):
                if j >= n_mol:
                    break
                name = comp.get("cas", f"ing_{j}")
                x_liq = step["x_liquid"].get(name, 0.0)
                oav = step["OAV"].get(name, 0.0)
                states[t, j, 0] = x_liq
                states[t, j, 1] = np.log10(max(oav, 1e-10))
                if state_dim == 3 and t < len(gamma_traj) and j < len(gamma_traj[t]):
                    states[t, j, 2] = float(gamma_traj[t][j])

        strategy = record.get("metadata", {}).get("generation_strategy", "wildcard")

        # Prefer the new static pyramid targets (3, 138) when present; otherwise
        # fall back to the legacy time-resolved objective trajectory. For the
        # pom_alltags arm (objective_dim=575) use the 575-dim molequles tag targets.
        if self.objective_dim == 575 and "alltags_targets" in record:
            objective_targets = np.array(record["alltags_targets"], dtype=np.float32)
        elif "pyramid_targets" in record:
            objective_targets = np.array(record["pyramid_targets"], dtype=np.float32)
        elif "objective_targets" in record:
            objective_targets = np.array(record["objective_targets"], dtype=np.float32)
        else:
            objective_targets = get_objective_targets(strategy, t_steps)
        objective_targets = objective_targets_to_pyramid(objective_targets).astype(np.float32, copy=False)

        if "psychometric_targets" in record:
            psychometric_targets = np.array(record["psychometric_targets"], dtype=np.float32)
        elif self.label_noise:
            psychometric_targets = sample_noisy_psychometric_targets(strategy)
        else:
            from .semantics import get_psychometric_targets
            psychometric_targets = get_psychometric_targets(strategy)

        measured_substantivity: list[tuple[float, float]] = []
        for comp in formula:
            cas = str(comp.get("cas", "")).replace("NATURAL:", "")
            if cas in self.poucher_substantivity:
                measured_substantivity.append((
                    float(comp.get("weight_fraction", 1.0)),
                    self.poucher_substantivity[cas],
                ))
        if measured_substantivity:
            weights = np.array([w for w, _v in measured_substantivity], dtype=np.float32)
            values = np.array([v for _w, v in measured_substantivity], dtype=np.float32)
            substantivity_target = float(np.average(values, weights=weights))
            substantivity_mask = 1.0
        else:
            substantivity_target = 0.0
            substantivity_mask = 0.0

        genre = record.get("metadata", {}).get("generation_strategy", "wildcard")
        return {
            "tokens": torch.stack(tokens, dim=0),              # (num_molecules, feature_dim)
            "physics": torch.from_numpy(states),                # (num_timesteps, num_molecules, physics_dim=2)
            "target_obj": torch.from_numpy(objective_targets),  # (3, 138)
            "target_sub": torch.from_numpy(psychometric_targets),  # (7,)
            "target_substantivity": torch.tensor(substantivity_target, dtype=torch.float32),
            "target_substantivity_mask": torch.tensor(substantivity_mask, dtype=torch.float32),
            "genre_label": self.genre_map.get(genre, 0),
            "genre": genre,
        }

    def _smiles_for_cas(self, cas: str) -> str:
        """Best-effort registry lookup for SMILES."""
        try:
            from .registry import AromaRegistry

            reg = AromaRegistry()
            rec = reg.get(cas)
            reg.close()
            return rec.get("smiles", "") if rec else ""
        except Exception:
            return ""


class PhysicsInformedFiLMBlock(torch.nn.Module):
    """
    Feature-wise Linear Modulation (FiLM) of static token embeddings by a
    time-varying physical state vector.

    The physics engine predicts a per-token, per-timestep scale (gamma) and
    shift (beta) that are applied to the static ingredient embeddings. The
    ``(1 + gamma) * tokens + beta`` form lets the block default to the identity
    when weights are near zero, while still giving the network full dynamic
    range as it learns. No positional encoding is added, so the model remains
    permutation-invariant with respect to ingredient order.
    """

    def __init__(self, embedding_dim: int, state_dim: int, hidden_dim: int = 64) -> None:
        super().__init__()
        self.embedding_dim = embedding_dim
        self.state_dim = state_dim

        # Stabilized per-channel affine transform for the raw physics states.
        # This standardizes the physics channels without compressing them into a
        # vector norm, preserving the genuine slope/curvature of the VLE trajectory.
        self.physics_scaler = torch.nn.Parameter(torch.ones(state_dim))  # type: ignore[arg-type]
        self.physics_bias = torch.nn.Parameter(torch.zeros(state_dim))  # type: ignore[arg-type]

        self.state_proj = torch.nn.Sequential(
            torch.nn.Linear(state_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim, embedding_dim * 2),
        )

        # Token normalization is applied BEFORE modulation (AdaLN style) so the
        # time-varying beta/shift term is not dampened by static token statistics.
        self.token_norm = torch.nn.LayerNorm(embedding_dim)

    def forward(self, tokens: torch.Tensor, states: torch.Tensor) -> torch.Tensor:
        """
        tokens: (B, S, E)
        states: (B, T, S, state_dim)
        Returns: (B, T, S, E)
        """
        b, seq_len, e = tokens.shape
        t_steps = states.size(1)

        # 1. Standardize physics channels independently.
        normed_states = states * self.physics_scaler + self.physics_bias

        # 2. Generate FiLM parameters and split into gamma (scale) and beta (shift).
        film_params = self.state_proj(normed_states)  # (B, T, S, 2*E)
        gamma, beta = torch.chunk(film_params, 2, dim=-1)

        # 3. Normalize tokens once, then broadcast across the time dimension.
        normed_tokens = self.token_norm(tokens)
        tokens_t = normed_tokens.unsqueeze(1).expand(b, t_steps, seq_len, e)

        # 4. Affine modulation outside the token norm loop.
        # Temporal variance from gamma and beta passes completely uninhibited.
        modulated = (1.0 + gamma) * tokens_t + beta
        return modulated


class PhysicsInformedMixtureTransformer(torch.nn.Module):
    """
    Permutation-invariant transformer encoder for fragrance dry-down trajectories.

    No positional encodings are added. The model attends over ingredient tokens
    gated by time-resolved physical states.
    """

    def __init__(
        self,
        embedding_dim: int = DEFAULT_EMBEDDING_DIM,
        state_dim: int = 4,
        hidden_dim: int = 256,
        num_heads: int = 4,
        num_layers: int = 4,
        dropout: float = 0.1,
        max_ingredients: int = 32,
    ) -> None:
        super().__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.max_ingredients = max_ingredients

        self.gating = PhysicsInformedFiLMBlock(embedding_dim, state_dim, hidden_dim=hidden_dim)
        self.input_proj = torch.nn.Linear(embedding_dim, hidden_dim)

        encoder_layer = torch.nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=dropout,
            batch_first=True,
        )
        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

    def forward(
        self,
        tokens: torch.Tensor,
        states: torch.Tensor,
        src_key_padding_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor] | torch.Tensor:
        """
        Training/inference forward compatible with the dual-head trainer.

        tokens: (B, S, E) static molecular embeddings
        states: (B, T, S, state_dim) time-resolved physical state vectors
        src_key_padding_mask: (B, S) bool mask, True for padding positions
        labels: optional objective targets for HF-style loss (not used by trainer)

        Returns the full latent tensor (B, T, S, H) for the PIMT heads, or a
        Hugging Face compatible dict when labels are provided.
        """
        gated = self.gating(tokens, states)  # (B, T, S, E)
        b, t, s, _ = gated.shape
        x = self.input_proj(gated)  # (B, T, S, H)
        x = x.reshape(b * t, s, self.hidden_dim)
        if src_key_padding_mask is not None:
            # TransformerEncoderLayer expects (B, S) bool mask where True = padding.
            mask = src_key_padding_mask.unsqueeze(1).expand(-1, t, -1).reshape(b * t, s)
            x = self.encoder(x, src_key_padding_mask=mask)
        else:
            x = self.encoder(x)
        x = x.reshape(b, t, s, self.hidden_dim)  # (B, T, S, H)

        if labels is not None:
            # HF-compatible fallback: pooled representation and placeholder logits.
            pooled = x.mean(dim=2)  # (B, T, H)
            logits = torch.nn.functional.linear(pooled, torch.randn(138, self.hidden_dim, device=x.device))
            logits = logits[:, :1, :].expand(-1, t, -1)  # (B, T, 138)
            loss = torch.nn.functional.mse_loss(logits, labels)
            return {"loss": loss, "logits": logits}

        return x