"""Read timeseries data from gravitational wave frame files."""

# Copyright (C) 2024 Becca Ewing, Yun-Jing Huang

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar, Optional

import gwframe
import igwn_segments as segments
from sgnts.base import Audioadapter, Offset, SeriesBuffer, TSFrame, TSSource

if TYPE_CHECKING:
    from sgn.base import SourcePad

logger = logging.getLogger("sgn")


@dataclass
class CacheEntry:
    """Simple cache entry parser for frame cache files.

    Cache file format: observatory description gps_start duration path
    Example: L L1_GWOSC_16KHZ_R1 1240215487 32 ./path/to/file.gwf
    """

    observatory: str
    description: str
    gps_start: float
    duration: float
    path: str

    @classmethod
    def from_line(cls, line: str) -> CacheEntry:
        """Parse a cache file line into a CacheEntry."""
        parts = line.strip().split()
        if len(parts) != 5:
            msg = f"Invalid cache line format: {line}"
            raise ValueError(msg)
        return cls(
            observatory=parts[0],
            description=parts[1],
            gps_start=float(parts[2]),
            duration=float(parts[3]),
            path=parts[4],
        )

    @property
    def segment(self) -> segments.segment:
        """Return the time segment for this cache entry."""
        return segments.segment(self.gps_start, self.gps_start + self.duration)


@dataclass(kw_only=True)
class GWFrameSource(TSSource):
    """Read timeseries data from gravitational wave frame files

    Args:
        channels:
            list[str], a list of channel names of the data, e.g.,
            ["L1:GWOSC-16KHZ_R1_STRAIN", "L1:GWOSC-16KHZ_R1_DQMASK"]. Source pads will
            be automatically generated for each channel, with channel name as pad name.
        frames:
            str | list[str], one of:
            - Path to a .cache file containing frame file locations
            - Path to a single .gwf frame file
            - List of paths to .gwf frame files
        instrument:
            str, optional, only read gwf files from this instrument (e.g., "H1", "L1").
            Default: None
    """

    allow_dynamic_source_pads: ClassVar[bool] = False

    channels: list[str]
    frames: str | list[str]
    instrument: Optional[str] = None

    @property
    def static_source_pads(self) -> list[str]:  # type: ignore[override]
        """Define source pads from channels."""
        return self.channels

    def __post_init__(self) -> None:
        super().__post_init__()
        self.cnt = dict.fromkeys(self.source_pads, 0)

        if self.start is None:
            msg = "GWFrameSource requires a start time to be specified"
            raise ValueError(msg)
        self.last_epoch: float = self.start

        if self.instrument is not None:
            for channel in self.channels:
                if self.instrument not in channel:
                    msg = (
                        f"Instrument '{self.instrument}' does not match "
                        f"channel name '{channel}'"
                    )
                    raise ValueError(msg)

        # init analysis segment
        self.analysis_seg = segments.segment(self.start, self.end)

        # Parse frames input (cache file, single gwf, or list of gwf files)
        cache = self._load_frames(self.frames)

        if self.instrument is not None:
            # only keep files with the correct instrument
            cache = [
                entry for entry in cache if entry.observatory == self.instrument[0]
            ]

        # only keep files that intersect the analysis segment
        self.cache: list[CacheEntry] = []
        for entry in cache:
            try:
                self.analysis_seg & entry.segment
            except ValueError:
                continue
            else:
                self.cache.append(entry)

        # make sure it is sorted by gps time
        self.cache.sort(key=lambda x: x.segment[0])

        # Check if there are missing segments
        segment_remaining = self.analysis_seg
        missing_segments = []
        for c in self.cache:
            if segment_remaining in c.segment:
                # the cache contains all the rest of the proposed segment
                segment_remaining = segments.segment(0, 0)
            elif segment_remaining[0] < c.segment[0]:
                # there is a discontinuity
                missing_segments.append(
                    segments.segment(segment_remaining[0], c.segment[0])
                )
                if c.segment[1] <= segment_remaining[1]:
                    segment_remaining = segments.segment(
                        c.segment[1], segment_remaining[1]
                    )
                else:
                    segment_remaining = segments.segment(0, 0)
            else:
                segment_remaining -= c.segment

        if segment_remaining:
            missing_segments.append(segment_remaining)

        if missing_segments:
            self.logger.warning(
                "Cache has missing segments %s, padding with gaps",
                missing_segments,
            )

        self.A = {c: Audioadapter() for c in self.channels}

        # load first segment of data to read sample rate
        self.rates: dict[str, int] = {}
        self.load_gwf_data(self.cache[0])
        self.logger.info("Initialized sample rates per channel: %s", self.rates)

        # Set buffer parameters for each pad (supports different sample rates per pad)
        for pad in self.source_pads:
            self.set_pad_buffer_params(
                sample_shape=(), rate=self.rates[self.rsrcs[pad]], pad=pad
            )

        # now that we have loaded data from this frame,
        # remove it from the cache
        self.cache.pop(0)

    def _load_frames(self, frames: str | list[str]) -> list[CacheEntry]:
        """Load frame file information from various input formats.

        Args:
            frames: Either a cache file path, single gwf file path, or list of gwf paths

        Returns:
            List of CacheEntry objects
        """
        if isinstance(frames, str):
            frames_path = Path(frames)

            # Check if it's a cache file
            if frames_path.suffix == ".cache":
                self.logger.info("Loading frame cache from %s", frames)
                with open(frames) as f:
                    return [CacheEntry.from_line(line) for line in f if line.strip()]

            # Single .gwf file
            elif frames_path.suffix == ".gwf":
                self.logger.info("Loading single frame file %s", frames)
                return [self._create_cache_entry_from_gwf(frames)]

            else:
                msg = (
                    f"Unsupported file type: {frames_path.suffix}. "
                    "Expected .cache or .gwf"
                )
                raise ValueError(msg)

        elif isinstance(frames, list):
            # List of .gwf files
            self.logger.info("Loading %d frame files", len(frames))
            return [self._create_cache_entry_from_gwf(f) for f in frames]

        else:
            msg = f"frames must be str or list[str], got {type(frames)}"
            raise TypeError(msg)

    def _create_cache_entry_from_gwf(self, gwf_path: str) -> CacheEntry:
        """Create a CacheEntry from a gwf file by reading its metadata.

        Args:
            gwf_path: Path to gwf file

        Returns:
            CacheEntry object
        """
        # Read the file to get metadata
        frame_data = gwframe.read(gwf_path, channel=self.channels)

        # Get metadata from first channel
        first_channel_data = next(iter(frame_data.values()))
        gps_start = first_channel_data.t0
        duration = first_channel_data.duration

        # Extract observatory from channel name (e.g., "H1" from "H1:GDS-CALIB_STRAIN")
        first_channel_name = next(iter(frame_data.keys()))
        observatory = first_channel_name.split(":")[0]

        return CacheEntry(
            observatory=observatory,
            description="FRAME",
            gps_start=gps_start,
            duration=duration,
            path=gwf_path,
        )

    def load_gwf_data(self, frame_file: CacheEntry) -> None:
        """Load timeseries data from a gwf frame file.

        Args:
            frame_file:
                CacheEntry, the gwf frame file to read timeseries data from

        Returns:
            dict[str, np.ndarray], a dictionary with channel names as keys and
            numpy arrays of timeseries data as values
        """

        # get first cache entry
        segment = frame_file.segment

        intersection = self.analysis_seg & segment
        start = intersection[0]
        end = intersection[1]

        data_dict = gwframe.read(
            frame_file.path, channel=self.channels, start=start, end=end
        )

        if len(self.rates) == 0:
            for channel, data in data_dict.items():
                self.rates[channel] = int(data.sample_rate)

        for channel, data in data_dict.items():
            if self.last_epoch < start:
                self.logger.warning(
                    "Unexpected epoch: %f, expected: %f, sending gap buffer",
                    start,
                    self.last_epoch,
                )
                self.A[channel].push(
                    SeriesBuffer(
                        offset=Offset.fromsec(self.last_epoch),
                        sample_rate=self.rates[channel],
                        data=None,
                        shape=(int((start - self.last_epoch) * self.rates[channel]),),
                    )
                )
            elif self.last_epoch > start:
                msg = (
                    f"Unepected epoch: {start}, expected: {self.last_epoch}, sending "
                    "gap buffer"
                )
                raise ValueError(msg)
            self.A[channel].push(
                SeriesBuffer(
                    offset=Offset.fromsec(float(start)),
                    sample_rate=self.rates[channel],
                    data=data.array,
                )
            )

        self.last_epoch = end

    def internal(self) -> None:
        """Check if we need to read the next gw frame file in the cache. All channels
        are read at once.
        """

        # load next frame of data from disk when we have less than
        # one buffer length of data left
        read_new = False
        for channel, adapter in self.A.items():
            if adapter.size < self.num_samples(self.rates[channel]):
                read_new = True
                break

        if read_new and self.cache:
            # Read multiple channels at once
            self.load_gwf_data(self.cache[0])

            # now that we have loaded data from this frame,
            # remove it from the cache
            self.cache.pop(0)

    def new(self, pad: SourcePad) -> TSFrame:
        """New frames are created on "pad" with an instance specific count and a name
        derived from the channel name. "EOS" is set once we have procssed all data in
        the cache within the analysis segment.

        Args:
            pad:
                SourcePad, the pad for which to produce a new TSFrame

        Returns:
            TSFrame, the TSFrame that carries a list of SeriesBuffers
        """

        self.cnt[pad] += 1

        channel = self.rsrcs[pad]

        metadata = {"cnt": self.cnt[pad], "name": "'%s'" % pad.name}

        frame = self.prepare_frame(pad, metadata=metadata)

        if self.A[channel].end_offset >= frame.end_offset:
            bufs = self.A[channel].get_sliced_buffers((frame.offset, frame.end_offset))

            frame.set_buffers(list(bufs))

            self.A[channel].flush_samples_by_end_offset(frame.end_offset)

        return frame
