from typing import Annotated

import numpy as np
from pydantic import BaseModel, Field, AfterValidator

from aaredaqlib.coordinate import Coordinate, SmargonCoordinate, positive_coords


class SampleGeometryModel(BaseModel):
    # Beam location in the camera coordinates
    beam_location_pxl: Coordinate

    # 1 pixel in mm
    pixel_in_mm: Annotated[float, Field(gt=0.0)]

    aerotech: Coordinate

    aerotech_meas: Coordinate
    # Current smargon coordinates/target
    smargon: SmargonCoordinate

    omega_deg: float

    # Beam size in mm
    beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)]

    @property
    def beam_size_pxl(self) -> Coordinate:
        return self.beam_size_mm / self.pixel_in_mm


    def smargon_z(self, pic_y: float) -> Coordinate:
        y_mm = (pic_y - self.beam_location_pxl.y) * self.pixel_in_mm
        return self.smargon.sh_mm + self.smargon_nudge(Coordinate(z=y_mm))

    # This goes from image coordinates to relative beam coordinates in image plane
    def picture_to_sample(self, pic_pxl: Coordinate) -> Coordinate:
        # Relative to beam location
        beam_mm = (pic_pxl - self.beam_location_pxl) * self.pixel_in_mm + self.aerotech
        return beam_mm

    # This goes from relative beam coordinates in image plane to image
    def sample_to_picture(self, beam_mm: Coordinate) -> Coordinate:
        beam_pxl = (beam_mm - self.aerotech) / self.pixel_in_mm
        return beam_pxl + self.beam_location_pxl

    def smargon_nudge(self, coord: Coordinate) -> Coordinate:
        phi = np.radians(np.around(self.smargon.phi_deg, decimals=1))
        chi = np.radians(np.around(self.smargon.chi_deg, decimals=1))
        omega = np.radians(np.around(self.omega_deg, decimals=1))

        offset = Coordinate()
        co = np.cos(omega)
        so = np.sin(omega)
        cp = np.cos(phi)
        sp = np.sin(phi)
        cc = np.cos(chi)
        sc = np.sin(chi)

        offset.x = (
            -coord.z * co * sp
            - coord.y * so * sp
            + coord.x * cp * sc
            - coord.y * co * cc * cp
            + coord.z * so * cc * cp
        )
        offset.y = (
            -coord.z * co * cp
            - coord.y * so * cp
            - coord.x * sc * sp
            + coord.y * co * cc * sp
            - coord.z * so * cc * sp
        )
        offset.z = -coord.x * cc - coord.y * co * sc + coord.z * so * sc
        return offset

    def beamline_to_smargon(self, coord: Coordinate) -> Coordinate:
        return self.smargon.sh_mm + self.smargon_nudge(coord)

    def smargon_to_beamline(self, coord: Coordinate) -> Coordinate:
        rel = coord - self.smargon.sh_mm

        vec_x = self.smargon_nudge(Coordinate(x=1)).normalize()
        vec_y = self.smargon_nudge(Coordinate(y=1)).normalize()

        return Coordinate(x=rel * vec_x, y=rel * vec_y)

    def picture_to_smargon(self, coord: Coordinate) -> Coordinate:
        return self.beamline_to_smargon(self.picture_to_sample(coord))

    def smargon_to_picture(self, coord: Coordinate) -> Coordinate:
        return self.sample_to_picture(self.smargon_to_beamline(coord))