# SPDX-FileCopyrightText: © 2024-2025 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
# This file is part of SPECTRE
# SPDX-License-Identifier: GPL-3.0-or-later

import datetime
import dataclasses
import typing

import numpy as np
import numpy.typing as npt
import astropy.io

import spectre_core.spectrograms
from ._base import Base, BatchFile


@dataclasses.dataclass(frozen=True)
class _Extension:
    FITS = "fits"


class _FitsFile(BatchFile[spectre_core.spectrograms.Spectrogram]):
    """Stores spectrogram data in the FITS format generated by the e-Callisto network."""

    def read(self) -> spectre_core.spectrograms.Spectrogram:
        """Parses a FITS file to generate a `Spectrogram`.

        Reverses the spectra along the frequency axis and converts units to linearised
        values if necessary. Infers the spectrum type from the `BUNIT` header.

        :raises NotImplementedError: If the `BUNIT` header value represents an unsupported spectrum type.
        :return: A `Spectrogram` instance containing the parsed FITS file data.
        """
        with astropy.io.fits.open(self.file_path, mode="readonly") as hdulist:
            primary_hdu = hdulist["PRIMARY"]

            date_obs = primary_hdu.header["DATE-OBS"]
            time_obs = primary_hdu.header["TIME-OBS"]
            spectrogram_start_datetime = datetime.datetime.strptime(
                f"{date_obs}T{time_obs}", "%Y/%m/%dT%H:%M:%S.%f"
            )

            bintable_hdu = hdulist[1]
            times = bintable_hdu.data["TIME"][0]
            frequencies = bintable_hdu.data["FREQUENCY"][0] * 1e6  # Convert to Hz

            bunit = primary_hdu.header["BUNIT"]

            # bunit is interpreted as a SpectrumUnit
            spectrum_unit = spectre_core.spectrograms.SpectrumUnit(bunit)
            if spectrum_unit == spectre_core.spectrograms.SpectrumUnit.DIGITS:
                dynamic_spectra_linearised = self._convert_units_to_linearised(
                    primary_hdu.data
                )

                return spectre_core.spectrograms.Spectrogram(
                    dynamic_spectra_linearised[
                        ::-1, :
                    ],  # reverse the spectra along the frequency axis
                    times,
                    frequencies[::-1],  # sort the frequencies in ascending order
                    spectrum_unit,
                    spectrogram_start_datetime,
                )
            else:
                raise NotImplementedError(
                    f"Spectrum type with BUNITS '{spectrum_unit}' is not currently supported."
                )

    def _convert_units_to_linearised(
        self, raw_digits: npt.NDArray[np.float32]
    ) -> npt.NDArray[np.float32]:
        """Converts spectrogram data from raw digit values to linearised units.

        Applies a transformation based on ADC specifications to convert raw values
        to dB and then to linearised units.

        :param dynamic_spectra: Raw dynamic spectra in digit values.
        :return: The dynamic spectra with linearised units.
        """
        # conversion as per ADC specs [see email from C. Monstein]
        dB = (raw_digits / 255) * (2500 / 25)
        return 10 ** (dB / 10)


class CallistoBatch(Base):
    def __init__(self, batches_dir_path: str, start_time: str, tag: str) -> None:
        """A batch of data generated by the e-Callisto network.

        Supports the following file extensions:
        - `.fits`

        :param start_time: The start time of the batch.
        :param tag: The batch name tag.
        """
        super().__init__(batches_dir_path, start_time, tag)

        self.add_file(_FitsFile, _Extension.FITS)

    @property
    def fits_file(self) -> _FitsFile:
        """The batch file corresponding to the `.fits` extension."""
        return typing.cast(_FitsFile, self.get_file(_Extension.FITS))

    @property
    def spectrogram_file(self) -> BatchFile[spectre_core.spectrograms.Spectrogram]:
        return self.fits_file
