"""Longitudinal time-series with SequenceSpec — padded and torch-ready.

Demonstrates ragged time-series handling: different numbers of time-points
per subject, padded to a common length, with explicit padding masks.
"""

import numpy as np

from toxarrow.schemas.canonical import (
    CanonicalStudyGroup, Study, Unit, Endpoint, Observation,
)
from toxarrow.compile.arrow_store import ArrowStudyStore
from toxarrow.schemas.tensor import SequenceSpec
from toxarrow.materialize.materializer import TensorMaterializer

# Build a synthetic repeated-measures study: body weight over time
study = Study(study_id="GLP-001", source_type="synthetic", organism_or_system="rat")
units = [Unit(unit_id=f"rat_{i}", study_id="GLP-001") for i in range(4)]
eps = [Endpoint(endpoint_id="bw", endpoint_name="Body Weight")]
obs = []

# Rat 0: 5 time-points, Rat 1: 3, Rat 2: 1, Rat 3: 4
timepoints = {
    "rat_0": [0.0, 7.0, 14.0, 21.0, 28.0],
    "rat_1": [0.0, 7.0, 14.0],
    "rat_2": [0.0],
    "rat_3": [0.0, 7.0, 14.0, 28.0],
}
values = {
    "rat_0": [250.0, 260.0, 275.0, 285.0, 300.0],
    "rat_1": [200.0, 205.0, 198.0],
    "rat_2": [300.0],
    "rat_3": [230.0, 240.0, 255.0, 270.0],
}

for uid in units:
    for t, v in zip(timepoints[uid.unit_id], values[uid.unit_id]):
        obs.append(Observation(
            unit_id=uid.unit_id, endpoint_id="bw",
            time_from_start=t, value=v, value_unit="g",
        ))

group = CanonicalStudyGroup(study=study, units=units, endpoints=eps, observations=obs)
store = ArrowStudyStore.from_records(group)

# Padding mode: pad all sequences to max length (5 time-points)
spec = SequenceSpec(
    sample_by="unit_id",
    time_field="time_from_start",
    pad_output=True,
    padding_value=0.0,
)
batch = TensorMaterializer(store).materialize(spec)

print(f"Padded shape:  {batch['values'].shape}  (S={len(units)}, T_max=5, F=1)")
print(f"Padding mask shape: {batch['padding_mask'].shape}")
print(f"Sample IDs: {batch['sample_ids']}")
print()
for si, sid in enumerate(batch["sample_ids"]):
    lengths = int((~batch["padding_mask"][si]).sum())
    raw_times = batch["time_points"][si]
    raw_values = batch["values"][si, ~batch["padding_mask"][si], 0]
    print(f"  {sid}: {lengths} time-points")
    print(f"    times:  {raw_times}")
    print(f"    values: {raw_values}")
    print()

# Verify the padded region is correct
assert batch["padding_mask"][0, -1] == False  # rat_0 has all 5
assert batch["padding_mask"][1, -1] == True   # rat_1 padded at positions 3,4
assert batch["padding_mask"][2, -1] == True   # rat_2 padded after position 0

# Ready for PyTorch
try:
    import torch
    t = torch.from_numpy(batch["values"])
    pm = torch.from_numpy(batch["padding_mask"])
    print(f"Torch padded values: {tuple(t.shape)}, mask: {tuple(pm.shape)}")
except ImportError:
    print("(install torch for backend handoff)")