"""ITZ-property sub-model."""

from __future__ import annotations

from typing import Optional

from torch import Tensor, nn

from ..missing_features import MissingFeatureEncoder
from ..schema import DEFAULT_SCHEMA, SchemaSpec
from .outputs import transform_itz_outputs


class ITZSubModel(nn.Module):
    """Predict ITZ thickness, porosity and elastic modulus."""

    def __init__(
        self,
        schema: SchemaSpec = DEFAULT_SCHEMA,
        hidden_dim: int = 64,
        input_dim: int = 24,
    ):
        super().__init__()
        self.schema = schema
        self.input_dim = input_dim
        self.encoder = MissingFeatureEncoder(input_dim, hidden_dim)
        self.body = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
        )
        self.head = nn.Linear(hidden_dim, 3)

    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        h = self.encoder(x, mask)
        h = self.body(h)
        return transform_itz_outputs(self.head(h))
