"""Pydantic models for Babamul alerts."""

from datetime import timezone

import matplotlib.pyplot as plt
from astropy.time import Time

from .cutout_utils import plot_cutouts
from .raw_models import (
    EnrichedLsstAlert,
    EnrichedZtfAlert,
    LsstCandidate,
    LsstPhotometry,
    ZtfCandidate,
    ZtfPhotometry,
)


# here we just want to re-export the raw model we autogenerated from
# avro using pydantic-avro, as the main model
# to which we add extra functions, like a `get_photometry` accessor
# that takes care of combining the different photometry sources
# (prv_candidates, prv_nondetections, fp_hists)
class BabamulZtfAlert(EnrichedZtfAlert):
    """Pydantic model for a Babamul ZTF alert."""

    def get_photometry(self, deduplicated: bool = True) -> list[ZtfPhotometry]:
        """Combine and return all photometry data from the alert."""
        photometry = []
        # Add prv_candidates photometry
        photometry.extend(self.prv_candidates)
        # Add fp_hists photometry if available
        photometry.extend(self.fp_hists)
        # Add prv_nondetections photometry
        photometry.extend(self.prv_nondetections)

        # Sort photometry by Julian Date (jd)
        photometry.sort(key=lambda x: x.jd)

        # if deduplicated, remove duplicates based on (jd, band)
        if deduplicated:
            seen = set()
            deduped_photometry = []
            for p in photometry:
                key = (p.jd, p.band)
                if key not in seen:
                    seen.add(key)
                    deduped_photometry.append(p)
            photometry = deduped_photometry

        return photometry

    # let's add a `survey` property for convenience
    @property
    def survey(self) -> str:
        return "ZTF"

    @property
    def drb(self) -> float | None:
        """Return the reliability (DRB) score of the candidate, if available."""
        return self.candidate.drb

    def plot_cutouts(
        self,
        orientation: str = "horizontal",
        axes: list[plt.Axes] | None = None,
        show: bool = True,
        figsize: tuple | None = None,
        title: str | None = None,
    ) -> list[plt.Axes]:
        """Display the science, template, and difference cutouts for this alert.

        Parameters
        ----------
        orientation : str, default='horizontal'
            Layout orientation: 'horizontal' or 'vertical'. (overwritten if axes is not None)
        axes : list of matplotlib.axes.Axes, optional
            List of 3 axes to plot on. If None, creates new figure.
        show : bool, default=True
            Whether to call plt.show() after plotting.
        figsize : tuple, optional
            Figure size. If None, uses defaults based on orientation.
        title : str, optional
            Overall figure title. If None, uses objectId.

        Returns
        -------
        list of matplotlib.axes.Axes
            List of the three axes objects (science, template, difference).
        """
        return plot_cutouts(
            self,
            self.survey,
            False,
            axes,
            show,
            orientation,
            figsize,
            title or self.objectId,
        )

    def show_cutouts(
        self,
        orientation: str = "horizontal",
    ) -> None:
        """Display the cutouts in a new matplotlib figure.

        Parameters
        ----------
        orientation : str, default='horizontal'
            Layout orientation: 'horizontal' or 'vertical'.
        """
        self.plot_cutouts(orientation=orientation, show=True)


ZtfCandidate.datetime = property(
    lambda self: Time(self.jd, format="jd").to_datetime(timezone=timezone.utc)
)
ZtfPhotometry.datetime = property(
    lambda self: Time(self.jd, format="jd").to_datetime(timezone=timezone.utc)
)


class BabamulLsstAlert(EnrichedLsstAlert):
    """Pydantic model for a Babamul LSST alert."""

    def get_photometry(
        self, deduplicated: bool = True
    ) -> list[LsstPhotometry]:
        """Combine and return all photometry data from the alert."""
        photometry = []
        # Add prv_candidates photometry
        photometry.extend(self.prv_candidates)
        # Add fp_hists photometry if available
        photometry.extend(self.fp_hists)

        # Sort photometry by Julian Date (jd)
        photometry.sort(key=lambda x: x.jd)

        # if deduplicated, remove duplicates based on (jd, band)
        if deduplicated:
            seen = set()
            deduped_photometry = []
            for p in photometry:
                key = (p.jd, p.band)
                if key not in seen:
                    seen.add(key)
                    deduped_photometry.append(p)
            photometry = deduped_photometry

        return photometry

    # let's add a `survey` property for convenience
    @property
    def survey(self) -> str:
        return "LSST"

    @property
    def drb(self) -> float | None:
        """Return the reliability (DRB) score of the candidate, if available."""
        return self.candidate.reliability

    def plot_cutouts(
        self,
        orientation: str = "horizontal",
        use_rotation: bool = True,
        axes: list[plt.Axes] | None = None,
        show: bool = True,
        figsize: tuple | None = None,
        title: str | None = None,
    ) -> list[plt.Axes]:
        """Display the science, template, and difference cutouts for this alert.

        Parameters
        ----------
        orientation : str, default='horizontal'
            Layout orientation: 'horizontal' or 'vertical'. (overwritten if axes is not None)
        use_rotation : bool, default=True
            Whether to apply rotation based on FITS header (if available).
        axes : list of matplotlib.axes.Axes, optional
            List of 3 axes to plot on. If None, creates new figure.
        show : bool, default=True
            Whether to call plt.show() after plotting.
        figsize : tuple, optional
            Figure size. If None, uses defaults based on orientation.
        title : str, optional
            Overall figure title. If None, uses objectId.

        Returns
        -------
        list of matplotlib.axes.Axes
            List of the three axes objects (science, template, difference).
        """
        return plot_cutouts(
            self,
            self.survey,
            use_rotation,
            axes,
            show,
            orientation,
            figsize,
            title or self.objectId,
        )

    def show_cutouts(
        self,
        orientation: str = "horizontal",
        use_rotation: bool = True,
    ) -> None:
        """Display the cutouts in a new matplotlib figure.

        Parameters
        ----------
        orientation : str, default='horizontal'
            Layout orientation: 'horizontal' or 'vertical'.
        use_rotation : bool, default=True
            Whether to apply rotation based on FITS header (if available).
        """
        self.plot_cutouts(
            use_rotation=use_rotation, orientation=orientation, show=True
        )


LsstCandidate.datetime = property(
    lambda self: Time(self.jd, format="jd").to_datetime(timezone=timezone.utc)
)
LsstPhotometry.datetime = property(
    lambda self: Time(self.jd, format="jd").to_datetime(timezone=timezone.utc)
)
