from pathlib import Path

import torch
from loguru import logger
from torch import Tensor
from torch.nn import Linear
from torch.nn import Module
from torch.nn import ReLU
from torch.nn import Sequential


class SleepEpochClassifier(Module):
    """Per-epoch sleep stage classifier.

    Input: a fixed-size feature vector summarizing one epoch's context window
           (e.g. motion mean/max, HR mean/variance, time-of-night, previous
           predicted label — whatever features you settle on).
    Output: logits over sleep stage classes for that single epoch.

    No recurrence/attention — sequence smoothing (majority vote / HMM)
    happens as a separate post-processing step, not inside this network.
    """

    def __init__(self, n_features: int, n_classes: int, hidden_size: int = 16):
        super().__init__()
        self.net = Sequential(
            Linear(n_features, hidden_size),
            ReLU(),
            Linear(hidden_size, n_classes),
        )

    def forward(self, x: Tensor) -> Tensor:
        # x shape: (batch, n_features) -> returns (batch, n_classes) logits
        return self.net(x)

    def load(self, input_file: Path) -> list | None:
        payload = torch.load(input_file, weights_only=True)
        state_dict = payload.get("model_state_dict", payload)
        self.load_state_dict(state_dict)
        return payload.get("label_mapping")