"""Minimal dual-clock field sampling for legacy CelerisLab rollouts."""
from __future__ import annotations

from collections.abc import Iterable
from pathlib import Path

import numpy as np

from CCD_analysis.utils.cfd_interface import get_velocity_field


def field_steps_from_interval(total_steps: int, field_interval: int) -> tuple[int, ...]:
    """Return absolute field-save steps within a finite rollout."""
    if type(total_steps) is not int or total_steps < 1:
        raise ValueError("total_steps must be a positive integer")
    if type(field_interval) is not int or field_interval < 1:
        raise ValueError("field_interval must be a positive integer")
    return tuple(range(field_interval, total_steps + 1, field_interval))


class DualClockCollector:
    """Keep policy decisions fixed while reading fields at independent steps."""

    def __init__(
        self,
        flow_field,
        *,
        control_interval: int,
        control_count: int,
        u0: float,
        field_steps: Iterable[int],
    ):
        if type(control_interval) is not int or control_interval < 1:
            raise ValueError("control_interval must be a positive integer")
        if type(control_count) is not int or control_count < 1:
            raise ValueError("control_count must be a positive integer")
        self.flow_field = flow_field
        self.control_interval = control_interval
        self.control_count = control_count
        self.u0 = u0
        self.field_steps = tuple(field_steps)
        horizon = control_interval * control_count
        if any(type(step) is not int for step in self.field_steps):
            raise ValueError("field steps must be integers")
        if tuple(sorted(set(self.field_steps))) != self.field_steps:
            raise ValueError("field steps must be strictly increasing and unique")
        if any(step < 1 or step > horizon for step in self.field_steps):
            raise ValueError("field steps must lie within the rollout")
        self._next_field = 0
        self._next_control = 0
        self.ux: list[np.ndarray] = []
        self.uy: list[np.ndarray] = []
        self.saved_steps: list[int] = []

    def run_interval(self, control_index: int, action_target: np.ndarray) -> np.ndarray:
        """Advance exactly one policy interval, pausing only at due field steps."""
        if control_index != self._next_control:
            raise ValueError("control intervals must run sequentially")
        start = control_index * self.control_interval
        stop = start + self.control_interval
        due = []
        while self._next_field < len(self.field_steps):
            step = self.field_steps[self._next_field]
            if step > stop:
                break
            if step > start:
                due.append(step)
            self._next_field += 1

        ff = self.flow_field
        ff.begin_control_interval(self.control_interval, action_target)
        current = start
        boundaries = due if due and due[-1] == stop else [*due, stop]
        for step in boundaries:
            if step > current:
                ff.run_control_segment(step - current)
                current = step
            if step in due:
                ux, uy = get_velocity_field(ff, u0=self.u0)
                self.ux.append(ux)
                self.uy.append(uy)
                self.saved_steps.append(step)
        obs = ff.end_control_interval().copy()
        self._next_control += 1
        return obs

    def save(self, path: str | Path) -> None:
        """Save fields and their absolute lattice-step timeline."""
        if self._next_control != self.control_count:
            raise RuntimeError("cannot save an incomplete rollout")
        destination = Path(path)
        destination.parent.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(
            destination,
            ux=np.asarray(self.ux, dtype=np.float32),
            uy=np.asarray(self.uy, dtype=np.float32),
            lattice_steps=np.asarray(self.saved_steps, dtype=np.int64),
            control_indices=(np.asarray(self.saved_steps, dtype=np.int64) - 1) // self.control_interval,
            control_offsets=(np.asarray(self.saved_steps, dtype=np.int64) - 1) % self.control_interval + 1,
            control_interval=np.asarray(self.control_interval, dtype=np.int64),
        )
