Fitting¤

class ModelFit(zdx.Base):
    """Base class for fitting model parameters to a single exposure."""

    @abstractmethod
    def __call__(self, model, exposure):
        """Compute the forward model for one exposure."""
        pass

    def get_key(self, exposure, param):
        """Map a model parameter name to the exposure-specific storage key."""
        match param:
            case "positions":
                return exposure.key
            case "aberrations":
                return exposure.key
            case "defocus":
                return exposure.key
            case "pupil_delta":
                return exposure.key
            case "spectrum":
                return f"{exposure.filter}"
            case "fluxes":
                return exposure.key
            case _:
                raise ValueError(f"Parameter {param} has no key")

    def map_param(self, exposure, param):
        """Return the fully-qualified parameter name for a fit item."""
        if param in [
            "fluxes",
            "positions",
            "spectrum",
            "aberrations",
            "defocus",
        ]:
            return f"{param}.{exposure.get_key(param)}"
        return param

    def update_optics_zernikes(self, model, exposure):
        """Apply Zernike-aberration and defocus updates to the optics model."""
        optics = model.optics
        if "aberrations" in model.params.keys():
            coefficients = model.aberrations[self.get_key(exposure, "aberrations")]
            _ = lax.stop_gradient(coefficients[0, 0])
            optics = optics.set("pupil.coefficients", coefficients)

        if "defocus" in model.params.keys():
            disp = model.defocus[self.get_key(exposure, "defocus")]
            optics = optics.set("defocus", disp)

        return optics

class SinglePointFilterFit(ModelFit):
    """Pixel-basis fitter for point-source PSFs."""

    source: dl.Telescope = eqx.field(static=True)
    nwavels: int = eqx.field(static=True)

    def __init__(self, nwavels: int = 1):
        """Initialise the fitter with a single-point source and a wavelength grid."""
        self.source = dl.PointSource(wavelengths=[1.0])
        self.nwavels = int(nwavels)

    def update_optics(self, model, exposure):
        """Update the optics with pupil amplitude, OPD, and optional defocus terms."""
        optics = model.optics

        # ----------------------------
        # 1) Start from the current pupil transmission
        # ----------------------------
        base_amp = optics.layers["pupil"].transmission  # P0(x)
        amp = base_amp

        # ----------------------------
        # 2) Optional: learn pupil amplitude via pupil_delta
        # ----------------------------
        # --- pupil amplitude correction ---
        if "pupil_delta" in model.params.keys():
            base_amp = optics.layers["pupil"].transmission  # A0(x)
            pupil_mask = base_amp > 0

            delta = model.pupil_delta

            # scale: choose eps so raw delta stays O(1)
            eps_amp = 0.05  # 5% per unit in log-space (tune)
            amp = base_amp * jnp.exp(eps_amp * delta)

            # hard zero outside pupil
            amp = jnp.where(pupil_mask, amp, 0.0)

            # IMPORTANT: renormalize mean amplitude inside pupil to 1
            mean_amp = jnp.sum(amp) / (jnp.sum(pupil_mask) + 1e-12)
            amp = amp / (mean_amp + 1e-12)

            optics = optics.set("pupil.transmission", amp)

        # ----------------------------
        # 3) OPD (aberrations) + pupil-plane shear
        # ----------------------------
        if "aberrations" in model.params.keys():
            opd_map = model.aberrations[self.get_key(exposure, "aberrations")]
        else:
            # if you ever call update_optics without aberrations present
            opd_map = optics.layers["pupil"].opd

        # --- Apply pupil-plane shear as a coordinate warp of pupil-plane arrays ---
        if "pupil_shear" in model.params.keys():
            s_raw = model.pupil_shear

            # scale raw -> physical dimensionless shear coefficient
            # Start small; tune later (1e-4 to 1e-2 are typical exploration ranges)
            eps_shear = 1e-3
            shx = eps_shear * s_raw
            shy = 0.0

            # Shear pupil transmission (smooth is better for gradients)
            # If you want *strict* binary display later, threshold outside optimisation.
            amp = apply_pupil_shear(amp, shx=shx, shy=shy, order=1, cval=0.0)

            # Shear OPD map
            opd_map = apply_pupil_shear(opd_map, shx=shx, shy=shy, order=1, cval=0.0)

        # ----------------------------
        # 4) Piston removal using the (possibly sheared) pupil mask
        # ----------------------------
        pupil_mask = amp > 0
        mean_val = jnp.sum(opd_map * pupil_mask) / (jnp.sum(pupil_mask) + 1e-12)
        opd_map = opd_map - mean_val

        # ----------------------------
        # 5) Write back into optics
        # ----------------------------
        optics = optics.set("pupil.transmission", amp)
        optics = optics.set("pupil.opd", opd_map)

        # ----------------------------
        # 6) Defocus with optional global scale
        # ----------------------------
        if "defocus" in model.params.keys():
            disp = model.defocus[
                self.get_key(exposure, "defocus")
            ]  # nominal defocus param

            scale = 1.0
            if "defocus_scale" in model.params.keys():
                k_raw = model.defocus_scale
                eps_k = 5e-3
                scale = 1.0 + eps_k * k_raw

            optics = optics.set("defocus", disp * scale)

        return optics

    # Forward model
    def __call__(self, model, exposure):
        source = self.source
        nw = self.nwavels

        # 1) Flux
        log_flux = model.get(exposure.fit.map_param(exposure, "fluxes"))
        flux = jnp.exp(log_flux * jnp.log(10.0))
        source = source.set("flux", flux)

        # 2) Position
        pos = model.get(exposure.fit.map_param(exposure, "positions"))
        source = source.set("position", pos * dlu.arcsec2rad(0.031))

        # 3) Polynomial spectrum in log10 space
        wv, filt = calc_throughput(exposure.filter, nwavels=nw)
        wv = jnp.asarray(wv)  # shape (nw,)
        filt = jnp.asarray(filt)
        filt = filt / (jnp.sum(filt) + 1e-12)  # base: normalised filter throughput

        # --- Get polynomial coefficients for this exposure ---
        if "spectrum" in model.params.keys():
            spec_param = exposure.fit.map_param(exposure, "spectrum")
            coeffs = model.get(spec_param)
            coeffs = jnp.atleast_1d(coeffs)  # ensure 1D, handles (1,) or (2,) etc
        else:
            coeffs = jnp.zeros((1,), dtype=jnp.float64)  # default = flat in log10
            # (log10_I = 0 → I = 1)

        # --- Build dimensionless wavelength coordinate ---
        lambda0 = jnp.mean(wv)
        x = (wv - lambda0) / (lambda0 + 1e-12)  # shape (nw,)

        # --- Evaluate log10 intensity p(x) and convert to linear ---
        log10_I = eval_poly_log10(x, coeffs)  # shape (nw,)

        # (Optional, but nice): remove the intercept so polynomial only changes shape,
        # and the overall normalisation is left to the flux parameter.
        log10_I = log10_I - jnp.mean(log10_I)

        I = jnp.power(10.0, log10_I)
        I = jnp.where(jnp.isfinite(I), I, 0.0)  # paranoia against NaN/inf
        I = jnp.clip(I, 0.0, jnp.inf)

        # --- Combine source SED with filter throughput ---
        weights = filt * I
        weights = weights / (jnp.sum(weights) + 1e-12)  # PSF weights sum to 1

        source = source.set("spectrum", dl.Spectrum(wv, weights))

        # 4) Optics, PSF and shear
        optics = self.update_optics(model, exposure)
        psfs = optics.model(source, return_psf=True)
        data = psfs.data

        if data.ndim == 3:
            # shape: (nwavels, ny, nx) -> integrate over wavelength
            psf = data.sum(axis=0)
        elif data.ndim == 2:
            # already a 2D PSF (ny, nx)
            psf = data
        else:
            raise ValueError(
                f"Unexpected PSF data ndim={data.ndim}, shape={data.shape}"
            )

        pixel_scale = psfs.pixel_scale.mean()

        if "jitter_raw" in model.params.keys():
            # map raw -> positive jitter in arcsec (choose a scale that makes raw~O(1))
            # e.g. 1 mas = 1e-3 arcsec
            jitter_arcsec = 1e-3 * jnn.softplus(model.jitter_raw)  # >= 0
            jitter_rad = dlu.arcsec2rad(jitter_arcsec)

            sigma_pix = jitter_rad / (pixel_scale + 1e-30)
            psf = gaussian_blur_fft(psf, sigma_pix)

        # if "primary_shear" in model.params.keys():
        #     shear_raw = model.primary_shear  # scalar
        #     eps_shear = 7.1e-3                 # physical shear = eps_shear * raw
        #     shx = eps_shear * shear_raw
        #     psf = apply_shear(psf, shx=shx, shy=0.0)

        if "radial_scale" in model.params.keys():
            raw = model.radial_scale  # scalar
            eps_rs = 5e-3  # 0.5% per unit
            scale = 1.0 + eps_rs * raw
            psf = radial_zoom(psf, scale)

        psf_obj = dl.PSF(psf, pixel_scale)
        return dlu.downsample(psf_obj.data, 4, mean=False)

class BaseModeller(zdx.Base):
    """Mixin that exposes a nested parameter dictionary through attribute access."""

    params: dict

    def __init__(self, params):
        """Store a parameter dictionary on the model object."""
        self.params = params

    def __getattr__(self, key):
        """Resolve parameter values without forcing a custom __getattribute__ path."""
        if key in self.params:
            return self.params[key]
        for k, val in self.params.items():
            if hasattr(val, key):
                return getattr(val, key)
        raise AttributeError(
            f"Attribute {key} not found in params of {self.__class__.__name__} object"
        )

    def __getitem__(self, key):
        """Fetch a nested parameter value by key, returning a dict of matches."""
        values = {}
        for param, item in self.params.items():
            if isinstance(item, dict) and key in item.keys():
                values[param] = item[key]

        return values

class ModelParams(BaseModeller):

    def __getitem__(self, key):
        return self.params[key]

    def __getattr__(self, key):

        # Make the object act like a real dictionary
        if hasattr(self.params, key):
            return getattr(self.params, key)

        if key in self.params.keys():
            return self.params[key]

        for sub_key, val in self.params.items():
            if hasattr(val, key):
                return getattr(val, key)

        raise AttributeError(
            f"Attribute {key} not found in params of {self.__class__.__name__} object"
        )

    def replace(self, values):
        # Takes in a super-set class and updates this class with input values
        return self.set(
            "params", dict([(param, getattr(values, param)) for param in self.keys()])
        )

    def from_model(self, values):
        return self.set(
            "params", dict([(param, values.get(param)) for param in self.keys()])
        )

    def __add__(self, values):
        matched = self.replace(values)
        return jax.tree.map(lambda x, y: x + y, self, matched)

    def __iadd__(self, values):
        return self.__add__(values)

    def __mul__(self, values):
        matched = self.replace(values)
        return jax.tree.map(lambda x, y: x * y, self, matched)

    def __imul__(self, values):
        return self.__mul__(values)

    def map(self, fn):
        return jax.tree.map(lambda x: fn(x), self)

    def inject(self, other):
        # Injects the values of this class into another class
        return other.set(list(self.keys()), list(self.values()))

    def partition(self, params):
        """params can be a model params object or a list of keys"""
        if isinstance(params, ModelParams):
            params = list(params.params.keys())
        return (
            ModelParams({param: self[param] for param in params}),
            ModelParams(
                {param: self[param] for param in self.keys() if param not in params}
            ),
        )

    def combine(self, params2):
        return ModelParams({**self.params, **params2.params})

    def jacfwd(self, fn, n_batch=1):
        X, unravel_fn = ravel_pytree(self)
        Xs = jnp.array_split(X, n_batch)
        rebuild = lambda X_batch, index: X.at[index : index + len(X_batch)].set(X_batch)
        lens = jnp.cumsum(jnp.asarray([len(x) for x in Xs], dtype=jnp.int32))[:-1]
        starts = jnp.concatenate([jnp.asarray([0], dtype=jnp.int32), lens])

        @eqx.filter_jacfwd
        def batched_jac_fn(x, index):
            model_params = unravel_fn(rebuild(x, index))
            return eqx.filter_jit(fn)(model_params)

        return jnp.concatenate(
            [batched_jac_fn(x, index) for x, index in zip(Xs, starts)], axis=-1
        )