"""Time and amplitude alignment transforms.

XAlign: Shifts the time axis so t=0 corresponds to a meaningful event.
YAlign: Removes DC offset by subtracting the mean over a baseline window.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np

from impakt.channel.model import Channel


@dataclass(frozen=True)
class XAlign:
    """Time-zero alignment — shifts t=0 to a reference point.

    Methods:
        - 'manual': Shift by an explicit time offset.
        - 'threshold': Find the first time the signal exceeds a threshold.
        - 'trigger': Use pre-trigger information (set t=0 at trigger).
    """

    method: str = "manual"
    reference_time: float = 0.0
    threshold_value: float = 0.0
    threshold_direction: str = "abs"  # 'abs', 'positive', 'negative'

    @property
    def name(self) -> str:
        return "x_align"

    @property
    def params(self) -> dict[str, Any]:
        return {
            "method": self.method,
            "reference_time": self.reference_time,
            "threshold_value": self.threshold_value,
            "threshold_direction": self.threshold_direction,
        }

    def apply(self, channel: Channel) -> Channel:
        if self.method == "manual":
            return self._align_manual(channel)
        elif self.method == "threshold":
            return self._align_threshold(channel)
        elif self.method == "trigger":
            return self._align_trigger(channel)
        else:
            raise ValueError(f"Unknown X-align method: {self.method}")

    def _align_manual(self, channel: Channel) -> Channel:
        """Shift time by a fixed offset."""
        new_time = channel.time - self.reference_time
        return channel.with_data(
            data=channel.data,
            time=new_time,
            transform_note=f"X-align manual ({self.reference_time:+.6f}s)",
        )

    def _align_threshold(self, channel: Channel) -> Channel:
        """Find first threshold crossing and set it as t=0."""
        if self.threshold_direction == "abs":
            signal = np.abs(channel.data)
            mask = signal >= abs(self.threshold_value)
        elif self.threshold_direction == "positive":
            mask = channel.data >= self.threshold_value
        else:
            mask = channel.data <= -abs(self.threshold_value)

        indices = np.where(mask)[0]
        if len(indices) == 0:
            # Threshold never crossed — no shift
            return channel.with_data(
                data=channel.data,
                transform_note=f"X-align threshold (no crossing found)",
            )

        crossing_time = float(channel.time[indices[0]])
        new_time = channel.time - crossing_time
        return channel.with_data(
            data=channel.data,
            time=new_time,
            transform_note=f"X-align threshold ({crossing_time:.6f}s)",
        )

    def _align_trigger(self, channel: Channel) -> Channel:
        """Set t=0 at the trigger point (already done if pre-trigger is set)."""
        # If the channel was loaded with pre-trigger, t=0 should already
        # be at the trigger. This is a no-op but documents the intent.
        return channel.with_data(
            data=channel.data,
            transform_note="X-align trigger",
        )


@dataclass(frozen=True)
class YAlign:
    """Y-axis zero-offset correction.

    Computes the mean of the signal over a baseline window and subtracts it,
    correcting for sensor DC offset or drift.

    Args:
        window: (t_start, t_end) time range for computing the baseline mean.
                If not provided, uses the first 10% of the signal or pre-trigger.
    """

    window_start: float | None = None
    window_end: float | None = None

    @property
    def name(self) -> str:
        return "y_align"

    @property
    def params(self) -> dict[str, Any]:
        return {
            "window_start": self.window_start,
            "window_end": self.window_end,
        }

    def apply(self, channel: Channel) -> Channel:
        t = channel.time
        d = channel.data

        if self.window_start is not None and self.window_end is not None:
            mask = (t >= self.window_start) & (t <= self.window_end)
        else:
            # Default: use pre-trigger region (t < 0) or first 10%
            if t[0] < 0:
                mask = t < 0
            else:
                n = max(1, len(t) // 10)
                mask = np.zeros(len(t), dtype=bool)
                mask[:n] = True

        if not np.any(mask):
            return channel.with_data(
                data=d,
                transform_note="Y-align (no baseline samples)",
            )

        baseline_mean = float(np.mean(d[mask]))
        corrected = d - baseline_mean

        window_desc = ""
        if self.window_start is not None:
            window_desc = f" [{self.window_start:.4f}, {self.window_end:.4f}]"

        return channel.with_data(
            data=corrected,
            transform_note=f"Y-align offset={baseline_mean:+.4f}{window_desc}",
        )


def x_align(
    channel: Channel,
    method: str = "manual",
    reference_time: float = 0.0,
    **kwargs: Any,
) -> Channel:
    """Convenience function for time alignment."""
    return XAlign(method=method, reference_time=reference_time, **kwargs).apply(channel)


def y_align(
    channel: Channel,
    window: tuple[float, float] | None = None,
) -> Channel:
    """Convenience function for Y-axis zero correction."""
    start, end = window if window else (None, None)
    return YAlign(window_start=start, window_end=end).apply(channel)
