"""Typed tensor contract shared by training, inference and validation banks."""

from dataclasses import dataclass
from typing import Any

import torch


@dataclass
class FlowTask:
    study_id: str
    context_ids: list[str]
    target_id: str
    context_time: torch.Tensor
    context_value: torch.Tensor
    context_mask: torch.Tensor
    target_time: torch.Tensor
    target_value: torch.Tensor
    target_mask: torch.Tensor
    protocol_events: torch.Tensor
    protocol_event_mask: torch.Tensor
    context_covariates: list[dict[str, Any]]
    target_covariates: dict[str, Any]
    concentration_scale: torch.Tensor
    log_center: torch.Tensor
    log_scale: torch.Tensor
    context_protocol_events: torch.Tensor | None = None
    context_protocol_event_mask: torch.Tensor | None = None
    target_protocol_events: torch.Tensor | None = None
    target_protocol_event_mask: torch.Tensor | None = None
    reporting_metadata: dict[str, float] | None = None
    context_covariate_values: torch.Tensor | None = None
    context_covariate_mask: torch.Tensor | None = None
    target_covariate_values: torch.Tensor | None = None
    target_covariate_mask: torch.Tensor | None = None
    context_cens: torch.Tensor | None = None
    assay_limit: torch.Tensor | None = None
    assay_known: torch.Tensor | None = None
    covariate_metadata: dict[str, Any] | None = None


@dataclass
class FlowBatch:
    study_id: list[str]
    context_ids: list[list[str]]
    target_id: list[str]
    context_time: torch.Tensor
    context_value: torch.Tensor
    context_mask: torch.Tensor
    context_individual_mask: torch.Tensor
    target_time: torch.Tensor
    target_value: torch.Tensor
    target_mask: torch.Tensor
    protocol_events: torch.Tensor
    protocol_event_mask: torch.Tensor
    context_covariates: list[list[dict[str, Any]]]
    target_covariates: list[dict[str, Any]]
    concentration_scale: torch.Tensor
    log_center: torch.Tensor
    log_scale: torch.Tensor
    context_protocol_events: torch.Tensor | None = None
    context_protocol_event_mask: torch.Tensor | None = None
    target_protocol_events: torch.Tensor | None = None
    target_protocol_event_mask: torch.Tensor | None = None
    reporting_metadata: list[dict[str, float] | None] | None = None
    context_covariate_values: torch.Tensor | None = None
    context_covariate_mask: torch.Tensor | None = None
    target_covariate_values: torch.Tensor | None = None
    target_covariate_mask: torch.Tensor | None = None
    context_cens: torch.Tensor | None = None
    assay_limit: torch.Tensor | None = None
    assay_known: torch.Tensor | None = None
    covariate_metadata: list[dict[str, Any] | None] | None = None