import contextlib
from functools import cached_property
from pathlib import Path
import yaml

import batoid
import batoid_rubin
import danish
import galsim
import ipywidgets
import matplotlib.pyplot as plt
import numpy as np
from scipy.linalg import lstsq
from scipy.optimize import least_squares
import astropy.units as u
from astropy.units import Quantity


class Sensor:
    def __init__(self, name, thx, thy, nphot, img, fiducial, focusz=0.0, wavelength=620e-9, wf_img=None):
        self.name = name
        self.thx = thx
        self.thy = thy
        self.nphot = nphot
        self.img = img
        self.fiducial = fiducial
        self.focusz = focusz
        self.wavelength = wavelength
        self.wf_img = wf_img
        self.bins = img.get_array().shape[0]
        bo2 = self.bins//2
        self.range = [[-bo2*10e-6, bo2*10e-6], [-bo2*10e-6, bo2*10e-6]]

        telescope = self.fiducial.withGloballyShiftedOptic(
            "Detector", [0, 0, focusz]
        )

        self.z_ref = batoid.zernikeTA(
            telescope, np.deg2rad(self.thx), np.deg2rad(self.thy),
            self.wavelength,
            nrad=20, naz=120,
            reference='chief',
            jmax=66, eps=0.61, focal_length=10.31
        )*self.wavelength  # meters

    def draw(self, telescope, seeing, rng=None):
        if rng is None:
            rng = np.random.default_rng()
        elif isinstance(rng, int):
            rng = np.random.default_rng(rng)

        telescope = telescope.withGloballyShiftedOptic(
            "Detector", [0, 0, self.focusz]
        )

        rv = batoid.RayVector.asPolar(
            optic=telescope, wavelength=self.wavelength,
            nrandom=self.nphot, rng=rng,
            theta_x=np.deg2rad(self.thx), theta_y=np.deg2rad(self.thy)
        )
        telescope.trace(rv)

        rv.x[:] -= np.mean(rv.x[~rv.vignetted])
        rv.y[:] -= np.mean(rv.y[~rv.vignetted])

        # Convolve in a Gaussian
        scale = 10e-6 * seeing/2.35/0.2
        rv.x[:] += rng.normal(scale=scale, size=len(rv))
        rv.y[:] += rng.normal(scale=scale, size=len(rv))

        # Bin rays
        psf, _, _ = np.histogram2d(
            rv.y[~rv.vignetted], rv.x[~rv.vignetted], bins=self.bins,
            range=self.range
        )
        self.img.set_array(psf/np.max(psf))

    def draw_wf(self, telescope):
        if self.wf_img is None:
            return
        telescope = telescope.withGloballyShiftedOptic(
            "Detector", [0, 0, self.focusz]
        )
        wf = batoid.wavefront(
            telescope, np.deg2rad(self.thx), np.deg2rad(self.thy),
            self.wavelength, reference='mean',
            nx=127
        )
        wfarr = wf.array
        w = ~wfarr.mask

        # Subtract piston/tip/tilt
        coords = wf.coords
        x = coords[..., 0]
        y = coords[..., 1]

        basis = galsim.zernike.zernikeBasis(
            6, x[w], y[w],
            R_outer=telescope.pupilSize/2,
        )
        coefs, _, _, _ = np.linalg.lstsq(basis.T, wfarr[w], rcond=None)
        coefs[4:] = 0.0
        ptt = np.dot(coefs, basis)
        wfarr[w] = wfarr[w] - ptt
        wfarr *= self.wavelength * 1e6 / 2.0  # +/-2 microns range
        self.wf_img.set_array(wfarr)


class ComCamAOS:
    def __init__(
        self,
        debug=None,
        control_log=None,
        rng=None,
        nthread=4,
        randomized_dof=list(range(10)),
        controlled_dof=list(range(10)),
        controlled_truncated_dof=[0,1,2,6,7,8,9],
        dz_terms = ((1, 4), (2, 4), (3, 4), (2, 5), (3, 5), (2, 6), (3, 6), (1, 7), (1, 8))
    ):
        batoid._batoid.set_nthreads(nthread)

        if debug is None:
            debug = contextlib.redirect_stdout(None)
        self.debug = debug
        if control_log is None:
            control_log = contextlib.redirect_stdout(None)
        self.control_log = control_log
        if isinstance(self.debug, ipywidgets.Output):
            self.debug.clear_output()
        if isinstance(self.control_log, ipywidgets.Output):
            self.control_log.clear_output()

        if rng is None:
            rng = np.random.default_rng()
        elif isinstance(rng, int):
            rng = np.random.default_rng(rng)
        self.rng = rng
        self.offset_rng = np.random.default_rng(rng.integers(2**63))

        self.randomized_dof = randomized_dof
        self.controlled_dof = controlled_dof
        self.controlled_truncated_dof = controlled_truncated_dof

        self.dz_terms = dz_terms

        self.fiducial = batoid.Optic.fromYaml("ComCam_r.yaml")
        self.builder = batoid_rubin.LSSTBuilder(self.fiducial)
        self.wavelength = 620e-9

        # Relative magnitude of perturbations in arcsec/micron
        self.amp = [25.0, 1000.0, 1000.0, 25.0, 25.0, 25.0, 4000.0, 4000.0, 25.0, 25.0]
        self.amp += [0.1]*20
        self.amp += [0.1]*20

        # widget variables
        self.m2_dz = 0.0
        self.m2_dx = 0.0
        self.m2_dy = 0.0
        self.m2_Rx = 0.0
        self.m2_Ry = 0.0

        self.cam_dz = 0.0
        self.cam_dx = 0.0
        self.cam_dy = 0.0
        self.cam_Rx = 0.0
        self.cam_Ry = 0.0

        # Controls
        kwargs = {'layout':{'width':'180px'}, 'style':{'description_width':'initial'}, 'format':'.3f'}
        self.m2_dz_control = ipywidgets.FloatText(value=self.m2_dz, description="M2 dz (µm)", step=10, **kwargs)
        self.m2_dx_control = ipywidgets.FloatText(value=self.m2_dx, description="M2 dx (µm)", step=500, **kwargs)
        self.m2_dy_control = ipywidgets.FloatText(value=self.m2_dy, description="M2 dy (µm)", step=500, **kwargs)
        self.m2_Rx_control = ipywidgets.FloatText(value=self.m2_Rx, description="M2 Rx (arcsec)", step=10, **kwargs)
        self.m2_Ry_control = ipywidgets.FloatText(value=self.m2_Ry, description="M2 Ry (arcsec)", step=10, **kwargs)
        self.cam_dz_control = ipywidgets.FloatText(value=self.cam_dz, description="Cam dz (µm)", step=10, **kwargs)
        self.cam_dx_control = ipywidgets.FloatText(value=self.cam_dx, description="Cam dx (µm)", step=2000, **kwargs)
        self.cam_dy_control = ipywidgets.FloatText(value=self.cam_dy, description="Cam dy (µm)", step=2000, **kwargs)
        self.cam_Rx_control = ipywidgets.FloatText(value=self.cam_Rx, description="Cam Rx (arcsec)", step=10, **kwargs)
        self.cam_Ry_control = ipywidgets.FloatText(value=self.cam_Ry, description="Cam Ry (arcsec)", step=10, **kwargs)
        self.zero_control = ipywidgets.Button(description="Zero")
        self.randomize_control = ipywidgets.Button(description="Randomize")
        self.control_truncated_control = ipywidgets.Button(description="Control w/ Trunc")
        self.control_penalty_control = ipywidgets.Button(description="Control w/ Penalty")
        self.control_ofc_control = ipywidgets.Button(description="Control w/ OFC")

        self.hex_controls = ipywidgets.VBox([
            self.m2_dz_control, self.m2_dx_control, self.m2_dy_control,
            self.m2_Rx_control, self.m2_Ry_control,
            self.cam_dz_control, self.cam_dx_control, self.cam_dy_control,
            self.cam_Rx_control, self.cam_Ry_control,
        ])

        self.m1m3_vals = [0.0]*20
        self.m1m3_controls = [
            ipywidgets.FloatText(value=self.m1m3_vals[i], description=f"Mode {i} (µm)", step=0.25, **kwargs)
            for i in range(20)
        ]
        self.m1m3_control_box1 = ipywidgets.VBox(self.m1m3_controls[:10])
        self.m1m3_control_box2 = ipywidgets.VBox(self.m1m3_controls[10:])

        self.m2_vals = [0.0]*20
        self.m2_controls = [
            ipywidgets.FloatText(value=self.m2_vals[i], description=f"Mode {i} (µm)", step=0.25, **kwargs)
            for i in range(20)
        ]
        self.m2_control_box1 = ipywidgets.VBox(self.m2_controls[:10])
        self.m2_control_box2 = ipywidgets.VBox(self.m2_controls[10:])

        self.component_controls = ipywidgets.Tab(layout=ipywidgets.Layout(min_width="240pt", height="290pt"))
        self.component_controls.children = [
            self.hex_controls,
            self.m1m3_control_box1, self.m1m3_control_box2,
            self.m2_control_box1, self.m2_control_box2,
        ]
        self.component_controls.set_title(0, "Hex")
        self.component_controls.set_title(1, "M1M3a")
        self.component_controls.set_title(2, "M1M3b")
        self.component_controls.set_title(3, "M2a")
        self.component_controls.set_title(4, "M2b")

        self.controls = ipywidgets.VBox([
            self.component_controls,
            ipywidgets.VBox([
                self.zero_control, self.randomize_control,
                self.control_truncated_control, self.control_penalty_control, self.control_ofc_control
            ])
        ])

        # Observers
        self.m2_dz_control.observe(lambda change: self.handle_event(change, 'm2_dz'), 'value')
        self.m2_dx_control.observe(lambda change: self.handle_event(change, 'm2_dx'), 'value')
        self.m2_dy_control.observe(lambda change: self.handle_event(change, 'm2_dy'), 'value')
        self.m2_Rx_control.observe(lambda change: self.handle_event(change, 'm2_Rx'), 'value')
        self.m2_Ry_control.observe(lambda change: self.handle_event(change, 'm2_Ry'), 'value')
        self.cam_dz_control.observe(lambda change: self.handle_event(change, 'cam_dz'), 'value')
        self.cam_dx_control.observe(lambda change: self.handle_event(change, 'cam_dx'), 'value')
        self.cam_dy_control.observe(lambda change: self.handle_event(change, 'cam_dy'), 'value')
        self.cam_Rx_control.observe(lambda change: self.handle_event(change, 'cam_Rx'), 'value')
        self.cam_Ry_control.observe(lambda change: self.handle_event(change, 'cam_Ry'), 'value')
        for i in range(20):
            self.m1m3_controls[i].observe(lambda change, i=i: self.handle_event(change, ('m1m3_vals', i)), 'value')
            self.m2_controls[i].observe(lambda change, i=i: self.handle_event(change, ('m2_vals', i)), 'value')
        self.zero_control.on_click(self.zero)
        self.randomize_control.on_click(self.randomize)
        self.control_truncated_control.on_click(self.control_truncated)
        self.control_penalty_control.on_click(self.control_penalty)
        self.control_ofc_control.on_click(self.control_ofc)

        self.view = self._view()
        self._pause_handler = False
        self._control_history = []

    def zero(self, b=None):
        self.m2_dz = 0.0
        self.m2_dx = 0.0
        self.m2_dy = 0.0
        self.m2_Rx = 0.0
        self.m2_Ry = 0.0
        self.cam_dz = 0.0
        self.cam_dx = 0.0
        self.cam_dy = 0.0
        self.cam_Rx = 0.0
        self.cam_Ry = 0.0
        self.m1m3_vals = [0.0]*20
        self.m2_vals = [0.0]*20
        self.update()
        self._control_history = []
        self._control_history.append(
            (
                self.wfe,
                self.m2_dz,
                self.m2_dx,
                self.m2_dy,
                self.m2_Rx,
                self.m2_Ry,
                self.cam_dz,
                self.cam_dx,
                self.cam_dy,
                self.cam_Rx,
                self.cam_Ry,
                self.m1m3_vals,
                self.m2_vals,
            )
        )

    def randomize(self, b=None):
        # amplitudes for all 50 dof
        offsets = np.zeros(50)
        offsets[self.randomized_dof] = self.offset_rng.normal(scale=self.amp)[self.randomized_dof]
        offsets = np.round(offsets, 3)

        self.m2_dz = offsets[0]
        self.m2_dx = offsets[1]
        self.m2_dy = offsets[2]
        self.m2_Rx = offsets[3]
        self.m2_Ry = offsets[4]
        self.cam_dz = offsets[5]
        self.cam_dx = offsets[6]
        self.cam_dy = offsets[7]
        self.cam_Rx = offsets[8]
        self.cam_Ry = offsets[9]
        self.m1m3_vals = list(offsets[10:30])
        self.m2_vals = list(offsets[30:50])
        self.update()
        self._control_history = []
        self._control_history.append(
            (
                self.wfe,
                self.m2_dz,
                self.m2_dx,
                self.m2_dy,
                self.m2_Rx,
                self.m2_Ry,
                self.cam_dz,
                self.cam_dx,
                self.cam_dy,
                self.cam_Rx,
                self.cam_Ry,
                self.m1m3_vals,
                self.m2_vals,
            )
        )

    def control_truncated(self, b=None):
        dz_fit = self.fit_dz()
        sens = np.array(self.sens)

        # Don't use M2 tilt or camera piston.
        sens = sens[:, self.controlled_truncated_dof]
        dof_fit, _, _, _ = lstsq(sens, dz_fit)
        dof_fit = np.round(dof_fit, 3)
        full_dof = np.zeros(50)
        full_dof[self.controlled_truncated_dof] = dof_fit
        self.apply_dof(-full_dof)
        self._plot_control_history()

    def control_penalty(self, b=None):
        # Add rows to sens matrix to penalize large dof
        dz_fit = self.fit_dz()
        ndz = len(dz_fit)
        sens = np.zeros((ndz+len(self.controlled_dof), len(self.controlled_dof)))
        sens[:ndz, :] = self.sens[:, self.controlled_dof]
        penalty = 1./np.array(self.amp)
        # Let's increase the penalty for the mirrors
        penalty[10:] *= 10
        alpha = 1e-10 # overall strength of penalty
        for i in range(ndz, ndz+len(self.controlled_dof)):
            sens[i, i-ndz] = alpha*penalty[self.controlled_dof[i-ndz]]
        dof_fit, _, _, _ = lstsq(sens, np.concatenate([dz_fit, [0]*len(self.controlled_dof)]))
        dof_fit = np.round(dof_fit, 3)
        full_dof = np.zeros(50)
        full_dof[self.controlled_dof] = dof_fit
        # with self.debug:
        #     print(f"WFE resids = {sens[:ndz] @ dof_fit - dz_fit}")
        #     print(f"penalty = {sens[ndz:] @ dof_fit}")

        self.apply_dof(-full_dof)
        self._plot_control_history()

    def control_ofc(self, b=None):
        from lsst.ts.ofc import OFCData, OFC
        ofc_data = OFCData("comcam")

        js = np.unique([j for k, j in self.dz_terms])
        maxk = max([k for k, j in self.dz_terms])
        maxj = max(js)
        ofc_data.zn_selected = js
        ofc = OFC(ofc_data)

        dz_fit = self.fit_dz()
        coefs = np.zeros((maxk+1, maxj+1))
        for (k, j), val in zip(self.dz_terms, dz_fit):
            coefs[k, j] = val * 1e6
        dz = galsim.zernike.DoubleZernike(
            coefs, uv_outer=0.4, xy_outer=4.18, xy_inner=2.5498
        )

        wfe = []
        sensor_names = []
        for k, sensor in self._sensors.items():
            if sensor.focusz == 0.0:
                wfe.append(dz(sensor.thx, sensor.thy).coef[4:])
                sensor_names.append(k.replace("foc", "R22_"))
        corrections = ofc.calculate_corrections(
            np.array(wfe),
            sensor_names,
            "R",
            1.0,
            0.0
        )
        self.apply_dof(ofc.ofc_controller.dof_state)
        self._plot_control_history()

    def _plot_control_history(self):
        self._control_history.append(
            (
                self.wfe,
                self.m2_dz,
                self.m2_dx,
                self.m2_dy,
                self.m2_Rx,
                self.m2_Ry,
                self.cam_dz,
                self.cam_dx,
                self.cam_dy,
                self.cam_Rx,
                self.cam_Ry,
                self.m1m3_vals,
                self.m2_vals,
            )
        )
        with self.control_log:
            with contextlib.suppress(AttributeError):
                fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(8, 2))
                axes[0].plot([x[0] for x in self._control_history], c='k')
                axes[1].plot([x[1] for x in self._control_history], c='b')
                axes[1].plot([x[6] for x in self._control_history], c='r')
                axes[2].plot([x[2] for x in self._control_history], c='b', ls='--')
                axes[2].plot([x[3] for x in self._control_history], c='b', ls=':')
                axes[2].plot([x[7] for x in self._control_history], c='r', ls='--')
                axes[2].plot([x[8] for x in self._control_history], c='r', ls=':')
                axes[3].plot([x[4] for x in self._control_history], c='b', ls='--')
                axes[3].plot([x[5] for x in self._control_history], c='b', ls=':')
                axes[3].plot([x[9] for x in self._control_history], c='r', ls='--')
                axes[3].plot([x[10] for x in self._control_history], c='r', ls=':')

                axes[0].set_ylabel("WFE (µm)")
                axes[1].set_ylabel("dz (µm)")
                axes[2].set_ylabel("dx, dy (µm)")
                axes[3].set_ylabel("Rx, Ry (arcsec)")
                for ax in axes:
                    ax.set_xlabel("Iteration")

                fig.tight_layout()
                plt.show(fig)

    def fit_dz(self):
        # Wavefront estimation part of the control loop.
        Rubin_mask_params = yaml.safe_load(open(Path(danish.datadir)/'ComCamObsc.yaml'))
        factory = danish.DonutFactory(
            R_outer=4.18, R_inner=2.5498,
            mask_params=Rubin_mask_params,
            focal_length=10.31, pixel_scale=10e-6
        )
        sky_level = 1.0

        dz_terms = self.dz_terms

        thxs = []
        thys = []
        z_refs = []
        imgs = []
        names = []
        for sensor in self._sensors.values():
            # Fit all intra/extra focal; skip in-focus
            if sensor.focusz == 0.0:
                continue

            thxs.append(np.deg2rad(sensor.thx))
            thys.append(np.deg2rad(sensor.thy))
            z_refs.append(sensor.z_ref)
            imgs.append(sensor.img.get_array().data)
            names.append(sensor.name)

        fitter = danish.MultiDonutModel(
            factory, z_refs=z_refs, dz_terms=dz_terms,
            field_radius=np.deg2rad(0.4), thxs=thxs, thys=thys
        )
        nstar = len(thxs)
        guess = [0.0]*nstar + [0.0]*nstar + [0.5] + [0.0]*len(dz_terms)
        sky_levels = [sky_level]*nstar

        with self.control_log:
            print()
            result = least_squares(
                fitter.chi, guess, jac=fitter.jac,
                ftol=1e-3, xtol=1e-3, gtol=1e-3,
                max_nfev=20, verbose=2,
                args=(imgs, sky_levels)
            )

        dxs_fit, dys_fit, fwhm_fit, dz_fit = fitter.unpack_params(result.x)

        with self.control_log:
            with contextlib.suppress(AttributeError):
                mods = fitter.model(
                    dxs_fit, dys_fit, fwhm_fit, dz_fit
                )
                fig, axes = plt.subplots(nrows=4, ncols=9, figsize=(9, 4))
                for i in range(18):
                    j = 2 * (i//9)
                    axes[j,i%9].imshow(imgs[i]/np.sum(imgs[i]), origin='upper')
                    axes[j+1,i%9].imshow(mods[i]/np.sum(mods[i]), origin='upper')
                    axes[j,i%9].text(
                        0.02, 0.92, names[i],
                        transform=axes[j,i%9].transAxes, fontsize=6, color='white'
                    )
                for ax in axes.ravel():
                    ax.set_xticks([])
                    ax.set_yticks([])
                    ax.set_aspect('equal')
                axes[0, 0].set_ylabel('Data')
                axes[1, 0].set_ylabel('Model')
                axes[2, 0].set_ylabel('Data')
                axes[3, 0].set_ylabel('Model')
                fig.tight_layout(pad=0.2)
                plt.show(fig)

            print()
            print("Found double Zernikes:")
            for z_term, val in zip(dz_terms, dz_fit):
                print(f"DZ({z_term[0]:2d}, {z_term[1]:2d}): {val*1e9:10.2f} nm")

        return dz_fit

    @cached_property
    def sens(self):
        dz_terms = self.dz_terms
        ndz = len(dz_terms)
        sens = np.zeros((ndz, 50))
        jmax = max([j for k, j in dz_terms])
        kmax = max([k for k, j in dz_terms])
        dz_ref = batoid.doubleZernike(
            self.fiducial, np.deg2rad(0.4), self.wavelength,
            jmax=jmax, kmax=kmax, eps=0.61
        )*self.wavelength
        for idof in range(50):
            dof = np.zeros(50)
            dof[idof] = self.amp[idof]  # microns or arcsec
            builder = self.builder.with_aos_dof(dof.tolist())
            telescope = builder.build()
            dz_p = batoid.doubleZernike(
                telescope, np.deg2rad(0.4), self.wavelength,
                jmax=jmax, kmax=kmax, eps=0.61
            )*self.wavelength
            for i, (k, j) in enumerate(dz_terms):
                sens[i, idof] = (dz_p - dz_ref)[k, j]/self.amp[idof]
        return sens

    def apply_dof(self, dof):
        dof = np.round(dof, 3)
        with self.control_log:
            print()
            print("Applying DOF:")
            print(f"M2 dz:        {dof[0]:10.2f} µm")
            print(f"M2 dx:        {dof[1]:10.2f} µm")
            print(f"M2 dy:        {dof[2]:10.2f} µm")
            print(f"M2 Rx:        {dof[3]:10.2f} arcsec")
            print(f"M2 Ry:        {dof[4]:10.2f} arcsec")
            print(f"cam dz:       {dof[5]:10.2f} µm")
            print(f"cam dx:       {dof[6]:10.2f} µm")
            print(f"cam dy:       {dof[7]:10.2f} µm")
            print(f"cam Rx:       {dof[8]:10.2f} arcsec")
            print(f"cam Ry:       {dof[9]:10.2f} arcsec")
            for i in range(20):
                print(f"M1M3 mode {i:2d}: {dof[10+i]:10.2f} µm")
            for i in range(20):
                print(f"M2 mode   {i:2d}: {dof[30+i]:10.2f} µm")
        self.m2_dz += dof[0]
        self.m2_dx += dof[1]
        self.m2_dy += dof[2]
        self.m2_Rx += dof[3]
        self.m2_Ry += dof[4]
        self.cam_dz += dof[5]
        self.cam_dx += dof[6]
        self.cam_dy += dof[7]
        self.cam_Rx += dof[8]
        self.cam_Ry += dof[9]
        self.m1m3_vals = [x+y for x, y in zip(self.m1m3_vals, dof[10:30])]
        self.m2_vals = [x+y for x, y in zip(self.m2_vals, dof[30:50])]
        self.update()

    def apply_M2(self, dz=None, dx=None, dy=None, rx=None, ry=None, dof=None):
        if dof is not None:
            if any(kw is not None for kw in [dz, dx, dy, rx, ry]):
                raise ValueError("Cannot set kw and dof at same time")
        else:
            dof = [0] * 5
        for i, (kw, unit) in enumerate(zip((dz, dx, dy, rx, ry), [u.micron]*3+[u.arcsec*2])):
            if kw is not None:
                if isinstance(kw, Quantity):
                    kw = kw.to_value(unit)
                dof[i] = kw
        dof = np.concatenate([dof, [0]*45])
        self.apply_dof(dof)

    def apply_camera(self, dz=None, dx=None, dy=None, rx=None, ry=None, dof=None):
        if dof is not None:
            if any(kw is not None for kw in [dz, dx, dy, rx, ry]):
                raise ValueError("Cannot set kw and dof at same time")
        else:
            dof = [0] * 5
        for i, (kw, unit) in enumerate(zip((dz, dx, dy, rx, ry), [u.micron]*3+[u.arcsec*2])):
            if kw is not None:
                if isinstance(kw, Quantity):
                    kw = kw.to_value(unit)
                dof[i] = kw
        dof = np.concatenate([[0]*5, dof, [0]*40])
        self.apply_dof(dof)

    def apply_m1m3(self, dof=None, **kwargs):
        if dof is not None:
            if len(kwargs) > 0:
                raise ValueError("Cannot set kw and dof at same time")
        else:
            dof = [0] * 20
        for k, v in kwargs.items():
            # Format of k should be bXY
            idx = int(k[1:])
            if isinstance(v, Quantity):
                v = v.to_value(u.micron)
            dof[idx] = v
        dof = np.concatenate([[0]*10, dof, [0]*20])
        self.apply_dof(dof)

    def apply_m2(self, dof=None, **kwargs):
        if dof is not None:
            if len(kwargs) > 0:
                raise ValueError("Cannot set kw and dof at same time")
        else:
            dof = [0] * 20
        for k, v in kwargs.items():
            # Format of k should be bXY
            idx = int(k[1:])
            if isinstance(v, Quantity):
                v = v.to_value(u.micron)
            dof[idx] = v
        dof = np.concatenate([[0]*30, dof])
        self.apply_dof(dof)

    def handle_event(self, change, attr):
        if self._pause_handler:
            return
        if isinstance(attr, tuple):
            current = getattr(self, attr[0])
            current[attr[1]] = change['new']
            setattr(self, attr[0], current)
        else:
            setattr(self, attr, change['new'])
        self.update()

    def _view(self):
        sensspec = [["S00", "S10", "S20"],
                    ["S01", "S11", "S21"],
                    ["S02", "S12", "S22"]]

        with plt.ioff():
            intra_out = ipywidgets.Output()
            focal_out = ipywidgets.Output()
            extra_out = ipywidgets.Output()
            wf_out = ipywidgets.Output()
            self.zk_out = ipywidgets.Output()
            zk_out = self.zk_out

            with intra_out:
                self._intra_fig = plt.figure(constrained_layout=True, figsize=(3.5, 3.5))
                self._intra_axes = self._intra_fig.subplot_mosaic(sensspec)
                # Determine spacing
                center = (self._intra_axes["S11"].transAxes + self._intra_fig.transFigure.inverted()).transform([0.5, 0.5])
                s01 = (self._intra_axes["S01"].transAxes + self._intra_fig.transFigure.inverted()).transform([0.5, 0.5])
                dx = s01[0] - center[0]  # make this 0.3 degrees  (far enough to see vignetting)
                factor = 0.3/dx
            with focal_out:
                self._focal_fig = plt.figure(constrained_layout=True, figsize=(3.5, 3.5))
                self._focal_axes = self._focal_fig.subplot_mosaic(sensspec)
            with extra_out:
                self._extra_fig = plt.figure(constrained_layout=True, figsize=(3.5, 3.5))
                self._extra_axes = self._extra_fig.subplot_mosaic(sensspec)
            with wf_out:
                self._wf_fig = plt.figure(constrained_layout=True, figsize=(3.5, 3.5))
                self._wf_axes = self._wf_fig.subplot_mosaic(sensspec)
            with zk_out:
                zk_out.clear_output()
                print("Double Zernikes")

        self._sensors = {}
        for fig, axes, wf_axes, focusz, prefix, ctx in zip(
            [self._intra_fig, self._focal_fig, self._extra_fig],
            [self._intra_axes, self._focal_axes, self._extra_axes],
            [None, self._wf_axes, None],
            [-1.5e-3, 0.0, 1.5e-3],
            ["in", "foc", "ex"],
            [intra_out, focal_out, extra_out]
        ):
            with contextlib.ExitStack() as stack:
                stack.enter_context(ctx)
                for k, ax in axes.items():
                    if wf_axes is not None:
                        wf_ax = wf_axes[k]
                    ax.set_xticks([])
                    ax.set_yticks([])

                    mytrans = ax.transAxes + fig.transFigure.inverted()
                    x, y = mytrans.transform([0.5, 0.5])
                    if focusz == 0.0:
                        nphot = 1000
                        nx = 41
                        wf_img = wf_ax.imshow(np.zeros((127, 127)), vmin=-1, vmax=1, cmap='bwr')
                        wf_ax.set_xticks([])
                        wf_ax.set_yticks([])
                    else:
                        nphot = 50000
                        nx = 181
                        wf_img = None
                    thx=-(x-center[0])*factor
                    thy=(y-center[1])*factor

                    self._sensors[prefix+k] = Sensor(
                        prefix+k, thx, thy, nphot=nphot, fiducial=self.fiducial,
                        focusz=focusz,
                        img=ax.imshow(np.zeros((nx, nx)), vmin=0, vmax=1),
                        wf_img=wf_img
                    )

                    ax.text(0.02, 0.92, k, transform=ax.transAxes, fontsize=6, color='white')

        with intra_out:
            self._intra_canvas = self._intra_fig.canvas
            self._intra_canvas.header_visible = False
            self._intra_fig.show()
        with focal_out:
            self._focal_canvas = self._focal_fig.canvas
            self._focal_canvas.header_visible = False
            self._focal_fig.show()
        with extra_out:
            self._extra_canvas = self._extra_fig.canvas
            self._extra_canvas.header_visible = False
            self._extra_fig.show()
        with wf_out:
            self._wf_canvas = self._wf_fig.canvas
            self._wf_canvas.header_visible = False
            self._wf_fig.show()

        plot_out = ipywidgets.Tab(
            layout=ipywidgets.Layout(
                width="400px",
                height="420px"
            )
        )

        plot_out.children = [intra_out, focal_out, extra_out, wf_out, zk_out]
        plot_out.set_title(0, "Intra")
        plot_out.set_title(1, "Focal")
        plot_out.set_title(2, "Extra")
        plot_out.set_title(3, "WF")
        plot_out.set_title(4, "Zk")

        self.text = ""
        self.textout = ipywidgets.Textarea(
            value=self.text,
            layout={'height': '65px', 'width': '100%'}
        )
        out = ipywidgets.VBox([plot_out, self.textout])

        return out

    def update(self):
        dof = [self.m2_dz, self.m2_dx, self.m2_dy, self.m2_Rx, self.m2_Ry]
        dof += [self.cam_dz, self.cam_dx, self.cam_dy, self.cam_Rx, self.cam_Ry]
        dof += self.m1m3_vals
        dof += self.m2_vals

        builder = self.builder.with_aos_dof(dof)
        telescope = builder.build()

        self.dz = batoid.doubleZernike(
            telescope,
            np.deg2rad(0.4),
            self.wavelength,
            kmax=6,
            jmax=37,
            eps=0.61
        ) * self.wavelength * 1e6 # microns
        self.dz[0] = 0.0
        self.dz[:, :4] = 0.0
        self.wfe = np.sqrt(np.sum(np.square(self.dz)))

        fwhms = []
        for sensor in self._sensors.values():
            sensor.draw(telescope, seeing=0.5, rng=self.rng)
            sensor.draw_wf(telescope)
            if sensor.focusz == 0.0:
                img = galsim.Image(sensor.img.get_array().data)
                try:
                    mom = galsim.hsm.FindAdaptiveMom(img)
                    fwhms.append(mom.moments_sigma*2.355*0.2)
                except:
                    fwhms.append(np.nan)
        self._intra_canvas.draw_idle()
        self._focal_canvas.draw_idle()
        self._extra_canvas.draw_idle()

        self.text = f"WFE: {self.wfe:.2f} µm\n"
        self.text += f"FWHM: {np.median(fwhms):.2f} arcsec"
        self.textout.value = self.text

        self._pause_handler = True
        self.m2_dz_control.value = self.m2_dz
        self.m2_dx_control.value = self.m2_dx
        self.m2_dy_control.value = self.m2_dy
        self.m2_Rx_control.value = self.m2_Rx
        self.m2_Ry_control.value = self.m2_Ry
        self.cam_dz_control.value = self.cam_dz
        self.cam_dx_control.value = self.cam_dx
        self.cam_dy_control.value = self.cam_dy
        self.cam_Rx_control.value = self.cam_Rx
        self.cam_Ry_control.value = self.cam_Ry
        for i, val in enumerate(self.m1m3_vals):
            self.m1m3_controls[i].value = val
        for i, val in enumerate(self.m2_vals):
            self.m2_controls[i].value = val
        self._pause_handler = False

        with self.zk_out:
            self.zk_out.clear_output()
            print("Double Zernikes")
            print("  k   j    dz[µm]   tot[µm]")
            print("===========================")
            asort = np.argsort(np.abs(self.dz).ravel())[::-1]
            ks, js = np.unravel_index(asort[:15], self.dz.shape)
            cumsum = 0.0
            for k, j in zip(ks, js):
                val = self.dz[k, j]
                cumsum += val**2
                print("{:3d} {:3d}  {:8.4f}  {:8.4f}".format(k, j, val, np.sqrt(cumsum)))
            print()
            print("RSS [µm]: {:8.4f}".format(np.sqrt(np.sum(np.square(self.dz)))))


    def display(self):
        from IPython.display import display
        self.app = ipywidgets.HBox([
            self.view,
            self.controls,
        ])

        self.update()
        return self.app
