from __future__ import annotations
from typing import Optional, Any, TYPE_CHECKING
from zempy.zosapi.analysis.fftpsf.psf_sampling import  PsfSampling
from zempy.zosapi.analysis.fftpsf.psf_rotation import  PsfRotation
from zempy.zosapi.analysis.fftpsf.fft_psf_type import  FftPsfType
from zempy.zosapi.analysis.ias.adapters.ias import IAS

if TYPE_CHECKING:

    from zempy.zosapi.analysis.ias.ias_field import IASField
    from zempy.zosapi.analysis.ias.ias_surface import IASSurface
    from zempy.zosapi.analysis.ias.ias_wavelength import IASWavelength

class PSFSettings(IAS):
    def __init__(self, ias_settings: Any, *,
                 field: Optional[IASField] = None,
                 wavelength: Optional[IASWavelength] = None,
                 image_surface: Optional[IASSurface] = None,
                 sample_size: PsfSampling = PsfSampling.PsfS_64x64,
                 output_size: Optional[PsfSampling] = None,
                 rotation: PsfRotation = PsfRotation.CW0,
                 image_delta: float = 0.0,
                 use_polarization: bool = False,
                 normalize: bool = True,
                 psf_type: FftPsfType = FftPsfType.Linear):
        super().__init__(ias_settings)

        self.field: Optional[int] = field
        self.wavelength: Optional[int] = wavelength
        self.image_surface: Optional[int] = image_surface
        self.sample_size: PsfSampling = sample_size
        self.output_size: PsfSampling = output_size or sample_size
        self.rotation: PsfRotation = rotation
        self.image_delta: float = image_delta if image_delta >= 0 else 1.0
        self.use_polarization: bool = bool(use_polarization)
        self.normalize: bool = bool(normalize)
        self.psf_type: FftPsfType = psf_type

    @staticmethod
    def _get_casted_settings(analysis_or_settings) -> Any:
        """Return the concrete IAS_FftPsf if available, else the base IAS_* implementation."""
        base = (analysis_or_settings.GetSettings()
                if hasattr(analysis_or_settings, "GetSettings")
                else analysis_or_settings)
        cast = getattr(base, "As_ZOSAPI_Analysis_Settings_Psf_IAS_FftPsf", None)
        if callable(cast):
            obj = cast()
            if obj is not None:
                return obj
        return getattr(base, "__implementation__", base)

    @staticmethod
    def _py_to_zos_enum(zosapi, enum_cls, value):
        """Use the base-enum's native conversion; no aliasing, no normalization."""
        return enum_cls.to_native(zosapi, value)

    # ---------- factory: snapshot from live analysis/settings ----------
    @classmethod
    def from_analysis(cls, analysis) -> "PSFSettings":
        s = cls._get_casted_settings(analysis)

        # No normalization: read names and index directly into the Python enums.
        smp = PsfSampling[s.SampleSize.ToString()]
        out = PsfSampling[s.OutputSize.ToString()]
        rot = PsfRotation[s.Rotation.ToString()]
        typ = FftPsfType[s.Type.ToString()]

        # Field / Wavelength / Surface via wrappers
        try:
            field = IASField.wrap(s.Field).get()
        except Exception:
            field = None
        try:
            wl = IASWavelength.wrap(s.Wavelength).get()
        except Exception:
            wl = None
        try:
            surf = IASSurface.wrap(s.Surface).get()
        except Exception:
            surf = None

        return cls(
            s,
            field=field,
            wavelength=wl,
            image_surface=surf,
            sample_size=smp,
            output_size=out,
            rotation=rot,
            image_delta=float(s.ImageDelta),
            use_polarization=bool(s.UsePolarization),
            normalize=bool(s.Normalize),
            psf_type=typ,
        )

    @classmethod
    def from_settings(cls, ias_settings) -> "PSFSettings":
        """Same as from_analysis but accepts an IAS_* settings object directly."""
        s = cls._get_casted_settings(ias_settings)
        return cls.from_analysis(s)

    # ---------- apply current Python fields back to live ZOSAPI settings ----------
    def apply(self, zosapi) -> None:
        target = self._get_casted_settings(self._ias)
        target.SampleSize = self._py_to_zos_enum(zosapi, PsfSampling, self.sample_size)
        target.OutputSize = self._py_to_zos_enum(zosapi, PsfSampling, self.output_size)
        target.Rotation   = self._py_to_zos_enum(zosapi, PsfRotation, self.rotation)
        target.Type       = self._py_to_zos_enum(zosapi, FftPsfType, self.psf_type)
        target.ImageDelta = float(self.image_delta)
        target.UsePolarization = bool(self.use_polarization)
        target.Normalize = bool(self.normalize)

        try:
            if self.field is None:
                IASField.wrap(target.Field).use_all()
            else:
                IASField.wrap(target.Field).set(self.field)
        except Exception:
            pass

        if self.wavelength is not None:
            try:
                IASWavelength.wrap(target.Wavelength).set(self.wavelength)
            except Exception:
                pass

        if self.image_surface is not None:
            try:
                IASSurface.wrap(target.Surface).set(self.image_surface)
            except Exception:
                pass

    # ---------- utility ----------
    def __str__(self) -> str:
        return (
            "PSF settings:(\n"
            f"  field={self.field},\n"
            f"  wavelength={self.wavelength},\n"
            f"  image_surface={self.image_surface},\n"
            f"  sample_size={self.sample_size.name},\n"
            f"  output_size={self.output_size.name},\n"
            f"  rotation={self.rotation.name},\n"
            f"  image_delta={self.image_delta},\n"
            f"  use_polarization={self.use_polarization},\n"
            f"  normalize={self.normalize},\n"
            f"  psf_type={self.psf_type.name}\n"
            ")"
        )
