import numpy as np
import pandas as pd
from datetime import datetime, timezone
from astropy.time import Time

from pybrotlib import Telemetry
from .pointing import Pointing, PointingModel


class AltAzPointing(Pointing):
    def __init__(self):
        super().__init__(
            telescope_params=["ObsTime", "Alt", "Az", "AltOff", "AzOff"],
            model_params=["AOFF", "BNP", "AN_A", "AE_A", "NPAE", "EOFF", "AN_E", "AE_E", "TF"],
        )

    def get_telescope_data(self, telemetry: Telemetry) -> dict[str, float]:
        return {
            "ObsTime": Time(val=datetime.now(timezone.utc), format="datetime", scale="utc").isot,
            "Az": telemetry.OBJECT.HORIZONTAL.AZ,
            "Alt": telemetry.OBJECT.HORIZONTAL.ALT,
            "AzOff": telemetry.POSITION.INSTRUMENTAL.AZ.OFFSET / np.cos(np.radians(telemetry.POSITION.HORIZONTAL.ALT))
            + telemetry.POINTING.OFFSETS.AZ,
            "AltOff": telemetry.POSITION.INSTRUMENTAL.ALT.OFFSET + telemetry.POINTING.OFFSETS.ALT,
        }

    def fit_model(self, data: pd.DataFrame, latitude: float) -> pd.DataFrame:
        pm_az = PointingModelAzimuth(1, 1, 1, 1, 1)
        pm_az.fit(data["Az"], data["Alt"], data["AzOff"], verbose=True)
        print("RMS Az: ", pm_az.rms(data["Az"], data["Alt"], data["AzOff"]))

        pm_alt = PointingModelElevation(1, 1, 1, 1)
        pm_alt.fit(data["Az"], data["Alt"], data["AltOff"], verbose=True)
        print("RMS Alt: ", pm_alt.rms(data["Az"], data["Alt"], data["AltOff"]))

        return pd.DataFrame(
            [
                [
                    pm_az.aoff,
                    pm_az.bnp,
                    pm_az.an_a,
                    pm_az.ae_a,
                    pm_az.npae,
                    pm_alt.eoff,
                    pm_alt.an_e,
                    pm_alt.ae_e,
                    pm_alt.tf,
                ]
            ],
            columns=self.model_params,
        )


class PointingModelAzimuth(PointingModel):
    def __init__(self, aoff: float, bnp: float, an_a: float, ae_a: float, npae: float):
        """Creates a pointing model for the azimuth.

        Args:
            aoff: azimuth encoder zero offset
            bnp: collimation error of the electromagnetic axis
            an_a: azimuth axis offset/misalignment north-south
            ae_a: azimuth axis offset/misalignment east-west,
            npae: Non-perpendicularity between the mount azimuth and elevation axes
        """
        PointingModel.__init__(self, aoff=aoff, bnp=bnp, an_a=an_a, ae_a=ae_a, npae=npae)

    def __call__(self, az: float, el: float) -> float:
        # if el < 89.9:
        tan_elevation = np.tan(np.radians(el))
        cos_elevation = np.cos(np.radians(el))
        # else:
        # clamp values so that we do not run into DIV by zero or overflow
        #    tan_elevation = 572.9572133543033
        #    cos_elevation = 0.0017453283658982613

        return (
            self.aoff
            - self.bnp / cos_elevation
            + self.an_a * np.sin(np.radians(az)) * tan_elevation
            - self.ae_a * np.cos(np.radians(az)) * tan_elevation
            + self.npae * tan_elevation
        )


class PointingModelElevation(PointingModel):
    def __init__(self, eoff: float, an_e: float, ae_e: float, tf: float):
        """Creates a pointing model for the elevation.

        Args:
            eoff: elevation encoder zero offset
            an_e: azimuth axis offset/misalignment north-south
            ae_e: azimuth axis offset/misalignment east-west
            tf: gravitational flexure correction at the horizon
        """
        PointingModel.__init__(self, eoff=eoff, an_e=an_e, ae_e=ae_e, tf=tf)

    def __call__(self, az: float, el: float) -> float:
        return (
            self.eoff
            + self.an_e * np.cos(np.radians(az))
            + self.ae_e * np.sin(np.radians(az))
            + self.tf * np.cos(np.radians(el))
        )
