"""Resampling and trimming transforms."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np
from scipy.signal import resample as scipy_resample

from impakt.channel.model import Channel


@dataclass(frozen=True)
class Trim:
    """Extract a time range from a channel.

    Args:
        t_start: Start time (inclusive). None = beginning.
        t_end: End time (inclusive). None = end.
    """

    t_start: float | None = None
    t_end: float | None = None

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

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

    def apply(self, channel: Channel) -> Channel:
        t = channel.time
        mask = np.ones(len(t), dtype=bool)

        if self.t_start is not None:
            mask &= t >= self.t_start
        if self.t_end is not None:
            mask &= t <= self.t_end

        if not np.any(mask):
            raise ValueError(
                f"Trim range [{self.t_start}, {self.t_end}] contains no samples. "
                f"Channel time range: [{t[0]:.6f}, {t[-1]:.6f}]"
            )

        return channel.with_data(
            data=channel.data[mask],
            time=channel.time[mask],
            transform_note=f"trim [{self.t_start}, {self.t_end}]",
        )


@dataclass(frozen=True)
class Resample:
    """Change the sample rate of a channel.

    Uses scipy.signal.resample for Fourier-based resampling.

    Args:
        target_rate: Desired sample rate in Hz.
    """

    target_rate: float

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

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

    def apply(self, channel: Channel) -> Channel:
        if self.target_rate <= 0:
            raise ValueError(f"Target rate must be positive, got {self.target_rate}")

        duration = channel.duration
        if duration <= 0:
            return channel

        new_num_samples = int(round(duration * self.target_rate))
        if new_num_samples < 2:
            new_num_samples = 2

        resampled_data = scipy_resample(channel.data, new_num_samples)
        new_time = np.linspace(channel.time[0], channel.time[-1], new_num_samples)

        return Channel(
            name=channel.name,
            code=channel.code,
            data=resampled_data.astype(np.float64),
            time=new_time,
            unit=channel.unit,
            sample_rate=self.target_rate,
            cfc_class=channel.cfc_class,
            metadata=channel.metadata,
            source_test_id=channel.source_test_id,
            transform_history=channel.transform_history
            + (f"resample {channel.sample_rate:.0f} -> {self.target_rate:.0f} Hz",),
        )


def trim(channel: Channel, t_start: float | None = None, t_end: float | None = None) -> Channel:
    """Convenience function to trim a channel to a time range."""
    return Trim(t_start=t_start, t_end=t_end).apply(channel)


def resample(channel: Channel, target_rate: float) -> Channel:
    """Convenience function to resample a channel."""
    return Resample(target_rate=target_rate).apply(channel)
