"""Missing-feature handling.

For every input feature we store a binary mask (1 = observed, 0 = missing).
Observed values are first standardized per channel (z-score using frozen stats
fit on the training set), so a small-magnitude but informative feature (e.g.
water/binder ratio ~0.4, admixture dosage ~0.01) is not swamped by a large one
(cement content ~500) in the input projection.  Missing entries are then
replaced by a learnable per-channel embedding, observed entries pick up a
learnable per-channel bias, and the value-plus-mask pair is projected through a
small Linear + LayerNorm + SiLU stack.  The mask is fed in both as a
multiplicative gate and as a side channel so the network can still learn whether
a value was measured, imputed, or structurally absent.
"""

from __future__ import annotations

from typing import Optional, Tuple

import torch
from torch import nn


def masked_feature_stats(
    x: torch.Tensor, mask: torch.Tensor, eps: float = 1e-5
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Per-feature (mean, std) over observed entries only (``mask > 0``).

    Columns with no observed entries (or with near-constant values) return
    ``mean=0, std=1`` so standardization is a no-op there.
    """

    mask = (mask > 0).to(x.dtype)
    x = torch.nan_to_num(x, nan=0.0)
    count = mask.sum(dim=0)
    safe = count.clamp(min=1.0)
    mean = (x * mask).sum(dim=0) / safe
    var = (((x - mean) ** 2) * mask).sum(dim=0) / safe
    std = var.clamp(min=0.0).sqrt()
    has = count > 0
    mean = torch.where(has, mean, torch.zeros_like(mean))
    std = torch.where(has & (std > eps), std, torch.ones_like(std))
    return mean, std


def masked_feature_stats_by_type(
    x: torch.Tensor,
    mask: torch.Tensor,
    type_index: torch.Tensor,
    num_types: int,
    eps: float = 1e-5,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Per-(type, feature) masked stats, shape ``(num_types, D)`` each.

    Used where one encoder multiplexes several feature schemas into the same
    columns (e.g. aggregate vs mortar nodes), so each row type is standardized
    against its own statistics. Types with no rows fall back to ``mean=0, std=1``.
    """

    in_dim = x.size(1)
    means = torch.zeros(num_types, in_dim, dtype=x.dtype)
    stds = torch.ones(num_types, in_dim, dtype=x.dtype)
    for t in range(num_types):
        sel = type_index == t
        if bool(sel.any()):
            means[t], stds[t] = masked_feature_stats(x[sel], mask[sel], eps)
    return means, stds


class MissingFeatureEncoder(nn.Module):
    """Encode a possibly-missing feature vector into a dense hidden representation.

    Parameters
    ----------
    in_dim: int
        Number of raw input channels.
    out_dim: int
        Output hidden dimension.
    """

    def __init__(self, in_dim: int, out_dim: int, num_stat_groups: int = 1):
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_stat_groups = num_stat_groups

        self.missing_embedding = nn.Parameter(torch.zeros(in_dim))
        nn.init.normal_(self.missing_embedding, std=0.02)
        self.observed_bias = nn.Parameter(torch.zeros(in_dim))

        # Frozen per-feature input standardization (z-score). Default identity;
        # call ``set_feature_stats`` with training-set stats to activate. Stored
        # as buffers (shape ``(num_stat_groups, in_dim)``) so they travel with
        # .to(device) and persist in checkpoints. ``num_stat_groups`` > 1 lets one
        # encoder hold per-node-type / per-edge-type stats, selected per row via a
        # ``type_index`` in ``forward`` (the columns mean different features for
        # different types, so a single shared z-score would blend them).
        self.register_buffer("feat_mean", torch.zeros(num_stat_groups, in_dim))
        self.register_buffer("feat_std", torch.ones(num_stat_groups, in_dim))

        self.proj = nn.Sequential(
            nn.Linear(in_dim * 2, out_dim),
            nn.LayerNorm(out_dim),
            nn.SiLU(),
        )

    @torch.no_grad()
    def set_feature_stats(
        self, mean: torch.Tensor, std: torch.Tensor, eps: float = 1e-5
    ) -> None:
        """Freeze standardization stats (from the train split).

        ``mean`` / ``std`` are ``(in_dim,)`` for a single group or
        ``(num_stat_groups, in_dim)`` for per-type stats.
        """

        mean = torch.as_tensor(mean, dtype=self.feat_mean.dtype, device=self.feat_mean.device)
        std = torch.as_tensor(std, dtype=self.feat_std.dtype, device=self.feat_std.device)
        if mean.dim() == 1:
            mean = mean.unsqueeze(0)
        if std.dim() == 1:
            std = std.unsqueeze(0)
        std = torch.where(std > eps, std, torch.ones_like(std))
        self.feat_mean.copy_(mean.expand_as(self.feat_mean))
        self.feat_std.copy_(std.expand_as(self.feat_std))

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        type_index: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if mask is None:
            mask = torch.ones_like(x)
        x = torch.nan_to_num(x, nan=0.0)
        if type_index is None:
            mean, std = self.feat_mean[0], self.feat_std[0]
        else:
            mean = self.feat_mean.index_select(0, type_index)
            std = self.feat_std.index_select(0, type_index)
        x = (x - mean) / std
        imputed = x * mask + self.missing_embedding * (1.0 - mask)
        imputed = imputed + self.observed_bias * mask
        joined = torch.cat([imputed, mask], dim=-1)
        return self.proj(joined)


def apply_random_missingness(
    x: torch.Tensor,
    rate: float = 0.15,
    generator: Optional[torch.Generator] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Randomly drop entries of ``x`` and return ``(x_masked, mask)``."""

    if generator is None:
        mask = (torch.rand_like(x) > rate).float()
    else:
        rnd = torch.rand(x.shape, generator=generator, device=x.device)
        mask = (rnd > rate).float()
    return x * mask, mask