from abc import ABC, abstractmethod
from collections import OrderedDict
from datetime import timedelta
from uuid import UUID, uuid4

import astropy.units as u  # type: ignore[import-untyped]
import numpy as np
import plotly.graph_objects as go
from astropy.coordinates import SkyCoord  # type: ignore[import-untyped]
from astropy.time import Time, TimeDelta  # type: ignore[import-untyped]
from pydantic import Field, model_validator

from across.tools.core.schemas import AstropyDateTime, AstropyTimeDelta
from across.tools.core.schemas.visibility import VisibilityComputedValues

from ..core.enums.constraint_type import ConstraintType
from ..core.plotting import plot_visibility_windows
from ..core.schemas import (
    ConstrainedDate,
    ConstraintReason,
    VisibilityWindow,
    Window,
)
from ..core.schemas.base import BaseSchema


class Visibility(ABC, BaseSchema):
    """
    Abstract base class for visibility calculations.

    Parameters
    ----------
    ra
        Right ascension of the target in degrees.
    dec
        Declination of the target in degrees.
    coordinate
        SkyCoord object representing the position of the target in sky
        coordinates. Optional if ra and dec are provided.
    begin
        Start time of visibility period.
    end
        End time of visibility period.
    step_size
        Time step size for visibility calculations. Default is 60 seconds.
    min_vis
        Minimum visibility duration in seconds for a window to be included.
        Default is 0.
    observatory_id
        Unique Observatory ID for which this visibility is calculated.
    observatory_name
        Name of the observatory for which this visibility is calculated.

    Attributes
    ----------
    timestamp
        Array of time values used in the visibility calculation.
    inconstraint
        Boolean array indicating whether the target is not constrained at each time step.
    calculated_constraints
        Ordered dictionary mapping constraint types to boolean arrays.
    visibility_windows
        List of visibility windows computed from the constraints.
    computed_values
        Container (VisibilityComputedValues) providing access to geometric quantities
        calculated during visibility analysis. After running a visibility calculation,
        you can inspect these values for further analysis or visualization. Fields include:

        - sun_angle: Angular separation between target and Sun (Quantity array)
        - moon_angle: Angular separation between target and Moon (Quantity array)
        - earth_angle: Angular separation between target and Earth limb (Quantity array)
        - alt_az: Altitude-azimuth coordinates of target from observatory (SkyCoord)

        Not all fields are populated for every calculation; values are only computed
        and stored by constraints that need them.

    Methods
    -------
    visible(t)
        Check if the target is visible at a given time.
    compute()
        Perform visibility calculation.

    """

    # Parameters
    ra: float | None = Field(default=None, ge=0, lt=360)
    dec: float | None = Field(default=None, ge=-90, le=90)
    step_size: AstropyTimeDelta = TimeDelta(60 * u.s)
    coordinate: SkyCoord | None = Field(default=None, exclude=True)
    begin: AstropyDateTime
    end: AstropyDateTime
    min_vis: int = 0
    observatory_id: UUID = Field(default_factory=uuid4, exclude=True)
    observatory_name: str

    # Computed values
    timestamp: AstropyDateTime | None = Field(default=None, exclude=True)
    inconstraint: np.typing.NDArray[np.bool_] = Field(default=np.array([]), exclude=True)
    calculated_constraints: OrderedDict[ConstraintType, np.typing.NDArray[np.bool_]] = Field(
        default_factory=OrderedDict, exclude=True
    )
    visibility_windows: list[VisibilityWindow] = []
    computed_values: VisibilityComputedValues = Field(default_factory=VisibilityComputedValues)

    @model_validator(mode="before")
    @classmethod
    def validate_parameters(cls, values: dict[str, float | SkyCoord]) -> dict[str, float | SkyCoord]:
        """
        Validate and synchronize SkyCoord, RA and DEC values.
        This method ensures that both coordinate representations (SkyCoord object and separate RA/DEC values)
        are present and consistent. If one representation is missing, it is
        computed from the other. In addition, it validates the step size and
        ensures it is a positive value. The begin and end times are rounded to
        the nearest step size, to ensure consistency.

        Parameters
        ----------
        values
            Dictionary containing the input values to validate, including coordinate,
            ra, dec, step_size, begin, and end.

        Returns
        -------
            Validated and synchronized dictionary of values.

        """
        if not isinstance(values, dict):
            values = values.__dict__

        # Check that either coordinate or ra/dec values are present
        if not values.get("coordinate") and values.get("ra") and values.get("dec"):
            values["coordinate"] = SkyCoord(ra=values["ra"] * u.deg, dec=values["dec"] * u.deg)
        elif not values.get("dec") and not values.get("ra") and values.get("coordinate"):
            values["ra"] = values["coordinate"].icrs.ra.deg
            values["dec"] = values["coordinate"].icrs.dec.deg
        else:
            raise ValueError("Must supply either coordinate or ra/dec values")

        # Extract TimeDelta from values
        if values.get("step_size") is None:
            values["step_size"] = cls.model_fields["step_size"].default

        # convert to step_size TimeDelta
        if isinstance(values["step_size"], timedelta):
            values["step_size"] = TimeDelta(values["step_size"])
        elif isinstance(values["step_size"], int):
            values["step_size"] = TimeDelta(values["step_size"] * u.s)

        # Check that step_size is a positive value
        if values["step_size"] < TimeDelta(0 * u.s):
            raise ValueError("Step size must be a positive value")

        # Round begin/end to step_size
        if values.get("begin") is not None and values.get("end") is not None:
            begin = Time(values["begin"]).unix
            end = Time(values["end"]).unix
            step_size = values["step_size"].to_value(u.s)

            values["begin"] = Time(begin // step_size * step_size, format="unix")
            values["end"] = Time((end // step_size) * step_size, format="unix")

        return values

    def visible(self, t: Time) -> bool | np.typing.NDArray[np.bool_]:
        """
        For a given time or array of times, is the target visible?

        Parameters
        ----------
        t
            Time to check

        Returns
        -------
            True if visible, False if not
        """
        if not t.isscalar:
            # Return an array of visibility booleans for multiple times
            result = np.zeros(len(t), dtype=bool)

            for win in self.visibility_windows:
                mask = (t >= win.window.begin.datetime) & (t <= win.window.end.datetime)
                result |= mask

            return result
        else:
            # Return if visible for a single time
            for win in self.visibility_windows:
                if win.window.begin.datetime <= t <= win.window.end.datetime:
                    return True
            return False

    def _compute_timestamp(self) -> None:
        """
        Compute timestamp array for visibility calculation. The base
        implementation is to create a timestamp array from the begin and end
        times with a step size defined by the step_size attribute.
        """
        step_seconds = self.step_size.to_value(u.s)
        offsets = np.arange(0.0, (self.end - self.begin).to_value(u.s), step_seconds)
        self.timestamp = Time(self.begin.unix + offsets, format="unix")

    def _quantize_to_step(self, t: Time) -> Time:
        """
        Quantize a time to the nearest configured step size relative to begin.
        """
        step_seconds = self.step_size.to_value(u.s)
        steps = np.round((t.unix - self.begin.unix) / step_seconds)
        return Time(self.begin.unix + steps * step_seconds, format="unix")

    def index(self, t: Time) -> int:
        """
        For a given time, return the index in the ephemeris.

        Parameters
        ----------
        t
            Time to find the index for

        Returns
        -------
            Index of the nearest time in the ephemeris
        """
        if self.timestamp is None:
            raise ValueError("Timestamp not computed")
        index = int(np.round((t.unix - self.timestamp[0].unix) // (self.step_size.to_value(u.s))))
        assert index >= 0 and index < len(self.timestamp), "Time outside of allowed range"
        return index

    @abstractmethod
    def _constraint(self, i: int) -> ConstraintType:
        """
        For a given index, return the constraint at that time.

        Parameters
        ----------
        i
            Index in the timestamp array.

        Returns
        -------
            The constraint type at the given index.

        """
        raise NotImplementedError("Subclasses must implement this method.")  # pragma: no cover

    @abstractmethod
    def _merge_computed_values(self) -> None:
        """
        Abstract method to merge computed values from constraints.
        """
        raise NotImplementedError("Subclasses must implement this method.")  # pragma: no cover

    @abstractmethod
    def prepare_data(self) -> None:
        """
        Abstract method to perform visibility calculation.
        """
        raise NotImplementedError("Subclasses must implement this method.")  # pragma: no cover

    def _get_id(self, i: int) -> UUID:
        """
        For a given index, get the ID of the observatory or instrument.

        Parameters
        ----------
        i
            Index in the timestamp array.

        Returns
        -------
            The observatory or instrument ID.

        """
        return self.observatory_id  # pragma: no cover

    def _get_name(self, i: int) -> str:
        """
        For a given index, get the name of the observatory or instrument.

        Parameters
        ----------
        i
            Index in the timestamp array.

        Returns
        -------
            The observatory or instrument name.

        """
        return self.observatory_name  # pragma: no cover

    def _make_windows(self) -> list[VisibilityWindow]:
        """
        Create visibility windows from the inconstraint array.

        Returns
        -------
            List of VisibilityWindow objects representing periods when the target is visible.

        """
        # Check that timestamp is set and constraints
        if self.timestamp is None:
            raise ValueError("Timestamp not set")

        # Find the start and end of the visibility windows
        buff: np.typing.NDArray[np.bool_] = np.concatenate(
            ([False], np.logical_not(self.inconstraint), [False])
        )
        begin = np.flatnonzero(~buff[:-1] & buff[1:])
        end = np.flatnonzero(buff[:-1] & ~buff[1:])
        indices = np.column_stack((begin, end - 1))

        visibility_windows = []
        for i in indices:
            constrained_date_begin = ConstrainedDate(
                datetime=self._quantize_to_step(self.timestamp[i[0]]).datetime,
                constraint=self._constraint(i[0] - 1),
                observatory_id=self._get_id(i[0] - 1),
            )
            constrained_date_end = ConstrainedDate(
                datetime=self._quantize_to_step(self.timestamp[i[1]]).datetime,
                constraint=self._constraint(i[1] + 1),
                observatory_id=self._get_id(i[1] + 1),
            )
            window = Window(begin=constrained_date_begin, end=constrained_date_end)
            visibility = int((self.timestamp[i[1]] - self.timestamp[i[0]]).to_value(u.s))
            constraint_reason = ConstraintReason(
                start_reason=f"{self._get_name(i[0] - 1)} {self._constraint(i[0] - 1).value}",
                end_reason=f"{self._get_name(i[1] + 1)} {self._constraint(i[1] + 1).value}",
            )

            visibility_window = VisibilityWindow(
                window=window,
                max_visibility_duration=visibility,
                constraint_reason=constraint_reason,
            )
            # Only add windows if they are > min_vis seconds long
            if self.timestamp[i[1]].datetime - self.timestamp[i[0]].datetime > timedelta(
                seconds=self.min_vis
            ):
                visibility_windows.append(visibility_window)

        return visibility_windows

    def plot(
        self, fig: go.Figure | None = None, offset: int | float = 0, width: int = 700, height: int = 1000
    ) -> go.Figure:
        """
        Method to visualize visibility windows using plotly.
        Calls the across-tools plotting core functionality and configures the plot
        layout to user specifications.

        Parameters
        ----------
        fig : go.Figure, optional
            An existing plotly figure to add to, by default None
        offset : int | float, optional
            The x-axis offset to plot new visibility windows, by default 0
        width: int, optional
            The width of the plot, in pixels. Defaults to 700
        height: int, optional
            The height of the plot, in pixels. Defaults to 1000

        Returns
        -------
        go.Figure
            The plotly figure containing the footprint plot
        """
        if fig is None:
            fig = go.Figure()

        fig = plot_visibility_windows(
            visibility_windows=[window.model_dump() for window in self.visibility_windows],
            observatory_name=self.observatory_name,
            fig=fig,
            offset=offset,
        )

        fig.update_layout(
            title="Visibility Windows",
            yaxis=dict(
                title="Time (UTC)",
                range=[self.end.to_datetime(), self.begin.to_datetime()],  # descending time
                type="date",
                autorange=False,  # don't resize
            ),
            xaxis=dict(
                title="Visibility Windows",
                tickvals=[offset],
                ticktext=[self.observatory_name],
            ),
            width=width,
            height=height,
        )
        return fig

    def compute(self) -> None:
        """
        Perform visibility calculation.
        """
        self._compute_timestamp()
        self.prepare_data()
        self.visibility_windows = self._make_windows()
        self._merge_computed_values()
