import numpy as np
import pandas as pd


from datetime import datetime, timezone
from astropy.time import Time
from .pointing import Pointing, PointingModel
from pybrotlib import Telemetry

latitude = 0.0


class HaDecPointing(Pointing):
    def __init__(self):
        super().__init__(
            telescope_params=["ObsTime", "Ha", "Dec", "HaOff", "DecOff"],
            model_params=[
                "C_mAl_tau",
                "C_mEl_tau",
                "C_Coll_tau",
                "C_Refr_tau",
                "C_Enc_tau",
                "C_Flex_tau",
                "C_NP_tau",
                "C_mAl_dec",
                "C_mEl_dec",
                "C_Coll_dec",
                "C_Refr_dec",
                "C_Flex_dec",
            ],
        )

    def get_telescope_data(self, telemetry: Telemetry) -> dict[str, float]:
        return {
            "ObsTime": Time(val=datetime.now(timezone.utc), format="datetime", scale="utc").isot,
            "Ha": telemetry.OBJECT.EQUATORIAL.HA,
            "Dec": telemetry.OBJECT.EQUATORIAL.DEC,
            "HaOff": telemetry.POSITION.INSTRUMENTAL.HA.OFFSET / np.cos(np.radians(telemetry.OBJECT.EQUATORIAL.DEC))
            + telemetry.POINTING.OFFSETS.HA,
            "DecOff": telemetry.POSITION.INSTRUMENTAL.DEC.OFFSET + telemetry.POINTING.OFFSETS.DEC,
        }

    def fit_model(self, data: pd.DataFrame, latitude_: float) -> pd.DataFrame:
        global latitude
        latitude = latitude_
        pm_ha = PointingModelHa(0, 0, 0, 0, 0, 0, 0)
        pm_ha.fit(data["Ha"], data["Dec"], data["HaOff"], verbose=True)
        print("RMS Ha: ", pm_ha.rms(data["Ha"], data["Dec"], data["HaOff"]))

        pm_dec = PointingModelDec(0, 0, 0, 0, 0)
        pm_dec.fit(data["Ha"], data["Dec"], data["DecOff"], verbose=True)
        print("RMS Dec: ", pm_dec.rms(data["Ha"], data["Dec"], data["DecOff"]))

        return pd.DataFrame(
            [
                [
                    pm_ha.C_mAl_tau,
                    pm_ha.C_mEl_tau,
                    pm_ha.C_Coll_tau,
                    pm_ha.C_Refr_tau,
                    pm_ha.C_Enc_tau,
                    pm_ha.C_Flex_tau,
                    pm_ha.C_NP_tau,
                    pm_dec.C_mAl_dec,
                    pm_dec.C_mEl_dec,
                    pm_dec.C_Coll_dec,
                    pm_dec.C_Refr_dec,
                    pm_dec.C_Flex_dec,
                ]
            ],
            columns=self.model_params,
        )


class PointingModelHa(PointingModel):
    def __init__(
        self,
        C_mAl_tau: float,
        C_mEl_tau: float,
        C_Coll_tau: float,
        C_Refr_tau: float,
        C_Enc_tau: float,
        C_Flex_tau: float,
        C_NP_tau: float,
    ):
        """Creates a pointing model for the ha dec.

        Args:
            C_Enc_tau: Encoder error
            C_mAl_tau: Misalignment of Polar Axis in Azimuth
            C_mEl_tau: Misalignment of Polar Axis in Elevation
            C_Refr_tau: Refraction correction
            C_Coll_tau: Collimation correction
            C_Flex_tau: Tube Flexure correction
            C_NP_tau: Skewness of the Axes
        """
        PointingModel.__init__(
            self,
            C_mAl_tau=C_mAl_tau,
            C_mEl_tau=C_mEl_tau,
            C_Coll_tau=C_Coll_tau,
            C_Refr_tau=C_Refr_tau,
            C_Enc_tau=C_Enc_tau,
            C_Flex_tau=C_Flex_tau,
            C_NP_tau=C_NP_tau,
        )

    def __call__(self, ha: float, dec: float) -> float:
        global latitude
        omega = np.radians(360.0 / 86164.099)  # rotationrate of equatorial axis
        # to rad
        TAU_obj = np.radians(ha)
        DEC_obj = np.radians(dec)
        # pre-calculate repeating sine and cosine functions
        sin_TAU = np.sin(TAU_obj)
        cos_TAU = np.cos(TAU_obj)
        sin_DEC = np.sin(DEC_obj)
        cos_DEC = np.cos(DEC_obj)

        latitude_ = np.radians(latitude)
        sin_lat = np.sin(latitude_)
        cos_lat = np.cos(latitude_)

        x = np.sin(DEC_obj) * sin_lat + np.cos(DEC_obj) * cos_lat * cos_TAU
        altitude = np.arcsin(x)
        alt_diff = -omega * (1 / np.sqrt(1 - x**2)) * np.cos(DEC_obj) * cos_lat * np.sin(TAU_obj)
        sin_alt = np.sin(altitude)
        cos_alt = np.cos(altitude)

        return (
            self.C_mEl_tau * ((sin_TAU * sin_DEC) / cos_DEC)
            - self.C_mAl_tau * ((cos_TAU * sin_DEC) / cos_DEC)
            + self.C_Coll_tau
            + self.C_Refr_tau * ((sin_TAU * cos_lat) / (sin_alt * cos_DEC))
            + self.C_Flex_tau * ((sin_TAU * cos_lat) / cos_DEC)
            + self.C_Enc_tau / cos_DEC
            - self.C_NP_tau * (sin_DEC / cos_DEC)
        )


class PointingModelDec(PointingModel):
    def __init__(self, C_mAl_dec: float, C_mEl_dec: float, C_Coll_dec: float, C_Refr_dec: float, C_Flex_dec: float):
        """Creates a pointing model for the dec.

        Args:
            C_mAl_dec: Misalignment of Polar Axis in Azimuth
            C_mEl_dec: Misalignment of Polar Axis in Elevation
            C_Refr_dec: Refraction correction
            C_Coll_dec: Collimation correction
            C_Flex_dec: Tube Flexure correction
        """
        PointingModel.__init__(
            self,
            C_mAl_dec=C_mAl_dec,
            C_mEl_dec=C_mEl_dec,
            C_Coll_dec=C_Coll_dec,
            C_Refr_dec=C_Refr_dec,
            C_Flex_dec=C_Flex_dec,
        )

    def __call__(self, ha: float, dec: float) -> float:
        global latitude
        omega = np.radians(360.0 / 86164.099)  # rotationrate of equatorial axis
        # to rad
        TAU_obj = np.radians(ha)
        DEC_obj = np.radians(dec)
        # pre-calculate repeating sine and cosine functions
        sin_TAU = np.sin(TAU_obj)
        cos_TAU = np.cos(TAU_obj)
        sin_DEC = np.sin(DEC_obj)
        cos_DEC = np.cos(DEC_obj)

        latitude_ = np.radians(latitude)
        sin_lat = np.sin(latitude_)
        cos_lat = np.cos(latitude_)

        x = np.sin(DEC_obj) * sin_lat + np.cos(DEC_obj) * cos_lat * cos_TAU
        altitude = np.arcsin(x)
        alt_diff = -omega * (1 / np.sqrt(1 - x**2)) * np.cos(DEC_obj) * cos_lat * np.sin(TAU_obj)
        sin_alt = np.sin(altitude)
        cos_alt = np.cos(altitude)

        return (
            self.C_mAl_dec * sin_TAU
            + self.C_mEl_dec * cos_TAU
            + self.C_Coll_dec
            + self.C_Refr_dec * ((sin_lat - sin_alt * sin_DEC) / (cos_DEC * sin_alt))
            - self.C_Flex_dec * ((sin_lat - sin_alt * sin_DEC) / cos_DEC)
        )
