from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from allytools.units import Length
from typing import Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    from lensguild.sensor.Sensor import Sensor


def _axis_1d(
        start_end: Tuple[float, float],
    step_mm: Length,
    *,
    reverse: bool,
    shift_to_zero: bool,
) -> np.ndarray:
    if not isinstance(step_mm, Length):
        raise TypeError("step_mm must be a Length instance.")
    if step_mm.value_mm <= 0:
        raise ValueError("step_mm must be positive.")

    start, end = start_end
    vals = np.arange(start, end + 1e-9, step_mm.value_mm, dtype=np.float64)
    if reverse:
        vals = vals[::-1]
    if shift_to_zero:
        vals = vals - start  # 0 at quadrant origin (sensor center)
    return vals


@dataclass(frozen=True, slots=True)
class Grid:
    _sensor: Sensor

    # ---------- pixel pitch helpers ----------
    @staticmethod
    def _px_pitch_x(sensor) -> float:
        if hasattr(sensor.pixel, "width") and isinstance(sensor.pixel.width, Length):
            return sensor.pixel.width.value_mm
        if hasattr(sensor.pixel, "length") and isinstance(sensor.pixel.length, Length):
            return sensor.pixel.length.value_mm
        raise AttributeError("Pixel must have .width or .length as Length for X pitch.")

    @staticmethod
    def _px_pitch_y(sensor) -> float:
        if hasattr(sensor.pixel, "height") and isinstance(sensor.pixel.height, Length):
            return sensor.pixel.height.value_mm
        if hasattr(sensor.pixel, "length") and isinstance(sensor.pixel.length, Length):
            return sensor.pixel.length.value_mm
        raise AttributeError("Pixel must have .height or .length as Length for Y pitch.")

    @staticmethod
    def _quadrant_start_end(n_pix: int, pitch_mm: float) -> Tuple[float, float]:
        mid = n_pix / 2.0
        start = (mid - 0.5) * pitch_mm
        end   = (n_pix - 0.5) * pitch_mm
        return start, end

    @property
    def x_start_end(self) -> Tuple[float, float]:
        """Start and end X values of the top-right quadrant (mm)."""
        s = self._sensor
        px = self._px_pitch_x(s)
        return self._quadrant_start_end(s.width_pix, px)

    @property
    def y_start_end(self) -> Tuple[float, float]:
        """Start and end Y values of the top-right quadrant (mm)."""
        s = self._sensor
        py = self._px_pitch_y(s)
        return self._quadrant_start_end(s.height_pix, py)

    # ---------- internal 1D axis builder ----------

    # ---------- public 1D APIs ----------
    def x1d(self, step_mm: Length) -> np.ndarray:
        """
        1D X axis (0 → +X), same geometry as quadrant_grid, custom step.
        """
        return _axis_1d(self.x_start_end, step_mm, reverse=False, shift_to_zero=True)

    def y1d(self, step_mm: Length) -> np.ndarray:
        """
        1D Y axis (+Y → 0), top-to-bottom, same geometry as quadrant_grid, custom step.
        """
        return _axis_1d(self.y_start_end, step_mm, reverse=True, shift_to_zero=True)

    # ---------- 2D grids on demand ----------
    def x(self, step_mm: Length) -> np.ndarray:
        x1 = self.x1d(step_mm)
        y1 = self.y1d(step_mm)
        return np.tile(x1, (y1.size, 1))

    def y(self, step_mm: Length) -> np.ndarray:
        x1 = self.x1d(step_mm)
        y1 = self.y1d(step_mm)
        return np.tile(y1[:, None], (1, x1.size))

    # ---------- radial sequence ----------
    def get_radial(self, step_mm: Length) -> np.ndarray:
        if not isinstance(step_mm, Length):
            raise TypeError("step_mm must be a Length instance.")
        if step_mm.value_mm <= 0:
            raise ValueError("step_mm must be positive.")

        r_max_mm = 0.5 * self._sensor.diagonal.value_mm
        seq = np.arange(0.0, r_max_mm + 1e-9, step_mm.value_mm, dtype=np.float64)
        return seq[seq <= r_max_mm + 1e-9]
