import time
import numpy as np
from tqdm import tqdm
from numba import jit
from qolab.hardware.basic import BasicInstrument


@jit(nopython=True, cache=True, fastmath=True)
def stretch2mat(sm, lrg):
    """Stretches/expands small matrix (sm) to the larger matrix (lrg).
    Note that lrg matrix should exist and will be overwritten."""
    sm_rows = sm.shape[0]
    sm_cols = sm.shape[1]
    lrg_rows = lrg.shape[0]
    lrg_cols = lrg.shape[1]
    for i in range(lrg_rows):
        first_index = i * sm_rows // lrg_rows
        for j in range(lrg_cols):
            lrg[i, j] = sm[first_index, j * sm_cols // lrg_cols]
    return lrg


class SLM(BasicInstrument):
    def __init__(self, *args, shape=(360, 640), **kwds):
        BasicInstrument.__init__(self, *args, **kwds)
        self.config["Device type"] = "SLM"
        self.config["Device model"] = "Generic SLM Without Hardware interface"
        self.config["SLM type"] = None  # could be amplitude, phase, DMD
        self.config["FnamePrefix"] = "slm"
        self.deviceProperties.update(
            {
                "PhysicalShape",
                "LogicalShape",
                "ROIspecs",
            }
        )
        self.dfltValue = 0  # pixel is off
        self.dfltValueFrameBuffer = 0
        # physical pixels
        self.pxlsPhysical = np.zeros(shape, dtype=np.uint8)
        # SLM often look like displays so we assign Frame Buffer
        self.frameBuffer = np.zeros(self.pxlsPhysical.shape, dtype=np.uint32)
        # logical pixels (super pixels) to which ROI is divided,
        # one logical pixel consists of  physical pixels array
        self.pxlsLogical = np.zeros((2, 2), dtype=np.uint8)  # size is just an example
        # ROI definition, by default the whole SLM
        self.roi = (
            self.pxlsPhysical
        )  # mapped part of physical space, the whole SLM by default

        # time for sending new pattern to SLM plus SLM settling time
        self.settleTime = 0.000  # in seconds
        self.fillWithDefault()

    def fillWithDefault(self):
        self.pxlsPhysical.fill(self.dfltValue)
        self.pxlsLogical.fill(self.dfltValue)
        self.physical2fb()

    def logical2ROI(self):
        stretch2mat(self.pxlsLogical, self.roi)

    def getPhysicalShape(self):
        return self.pxlsPhysical.shape

    def getLogicalShape(self):
        return self.pxlsLogical.shape

    def getROIspecs(self):
        offset = 0
        if self.roi.base is not None:
            # roi which should be slice of pxlsPhysical points into depth of that array
            offset = self.roi.ctypes.data - self.roi.base.ctypes.data
        r = offset // self.pxlsPhysical.shape[1]
        c = offset - self.pxlsPhysical.shape[1] * r
        pos = (r, c)
        return {"pos": pos, "shape": self.roi.shape}

    def getLogicalPxlsSize(self):
        """Return tuple containing arrays of logical row and column sizes in physical pixels.

        Note: make sure that it is mimics calculation in stretch2mat.
        """
        physNrows = self.roi.shape[0]
        physNcols = self.roi.shape[1]
        logcNrows = self.pxlsLogical.shape[0]
        logcNcols = self.pxlsLogical.shape[1]
        # luckily row index will be sorted by construction so we do not need to collect unique elements
        _, row_sizes = np.unique(
            np.array(range(physNrows)) * logcNrows // physNrows, return_counts=True
        )
        _, col_sizes = np.unique(
            np.array(range(physNcols)) * logcNcols // physNcols, return_counts=True
        )
        return row_sizes, col_sizes

    def getLogicalPxlsArea(self):
        """Return 2d array with area of each logical pixel measured in physical pixels."""
        logcNrows = self.pxlsLogical.shape[0]
        logcNcols = self.pxlsLogical.shape[1]
        rs, cs = self.getLogicalPxlsSize()
        return np.matmul(rs.reshape(logcNrows, 1), cs.reshape(1, logcNcols))

    # Hardware dependent
    # SLM often looks like displays so we need to convert SLM pixel values
    # to physical display (Frame Buffer) representation
    # Making framebuffer could be calculation heavy so we separate it from display
    def physical2fb(self):
        """convert physical pixels values to their frame buffer representation"""
        raise NotImplementedError("physical2fb function is not implemented")

    # Hardware dependent
    def display(self):
        """Sends data prepared in Frame Buffer to the SLM hardware"""
        raise NotImplementedError("display function is not implemented")

    # Hardware dependent
    def _measure(self):
        # grab data with already displayed on SLM mask
        raise NotImplementedError("_measure function is not implemented")

    def measureMasksBrightness(self, masks):
        """
        Measure brightness of the masks. They assumed to be row stacked and unwrapped along row.
        Returns list of measurements. I.e. it could be some sort of structure or class per measurement.

        Note: this method tries to be smart and use SLM settleTime to calculate next mask representation.
        """
        slm = self
        Nmasks = masks.shape[0]
        brightness = np.zeros((Nmasks, 1), float)
        brightness = []

        frame_buffer_ready = False
        for i in tqdm(range(Nmasks)):
            if not frame_buffer_ready:
                slm.pxlsLogical = masks[i, :].reshape(slm.pxlsLogical.shape)
                slm.logical2ROI()  # long operation about 10 mS on low end software
                slm.physical2fb()
                frame_buffer_ready = True
            slm.display()
            slm_ready_deadline = (
                time.perf_counter() + slm.settleTime
            )  # settle time could be as long as half a second on liquid crystal SLMs
            frame_buffer_ready = False

            # now we have time to do some calculations while SLM settles
            if i + 1 < Nmasks:
                slm.pxlsLogical = masks[i + 1, :].reshape(
                    slm.pxlsLogical.shape
                )  # long operation about 10 mS
                slm.logical2ROI()  # long operation about 10 mS on low end software
                slm.physical2fb()
                frame_buffer_ready = True

            # killing time until SLM is ready
            time.sleep(max(0, slm_ready_deadline - time.perf_counter()))

            brightness.append(slm._measure())
        return brightness


class fakeSLM(SLM):
    """SLM which pretends to be real. Handy for debugging purposes"""

    def __init__(self, *args, **kwds):
        super().__init__(*args, **kwds)
        self.config["Device model"] = "Fake SLM"
        self.config["SLM type"] = "amplitude"
        # hardware view is used to store "transmitted" to hardware SLM data
        self.hardwareView = np.zeros(self.pxlsPhysical.shape).astype(float) * 0
        # fake scene which SLM "illuminates"
        self.scene = np.zeros(self.pxlsPhysical.shape).astype(float) * 0
        self.doGraphic = True  # shall fake hardware draw on computer screen
        self.settleTime = 0.000  # in seconds

    def physical2fb(self):
        # convert physical pixels values to their frame buffer representation
        np.copyto(self.frameBuffer, self.pxlsPhysical)
        # now scale up
        self.frameBuffer *= 256**4 - 1  # convert on/off to display representation

    def display(self):
        import matplotlib.pyplot as plt

        if self.doGraphic:
            plt.figure(4)
            plt.clf()
            plt.imshow(self.pxlsPhysical)
            plt.colorbar()
            plt.title("SLM physical pixels")
            plt.figure(5)
            plt.clf()
            plt.imshow(self.frameBuffer)
            plt.colorbar()
            plt.title("Frame buffer view")
        # here we undo pxlsPhysical -> frameBuffer transformation
        np.copyto(self.hardwareView, self.frameBuffer / (256**4 - 1) * 1.0)

    def _measure(self):
        import matplotlib.pyplot as plt

        b = self.hardwareView * self.scene
        b += 10  # adding background
        b += 0.001 * np.random.normal()  # adding noise
        if self.doGraphic:
            plt.figure(7)
            plt.clf()
            plt.imshow(b)
            plt.colorbar()
            plt.title("Masked scene")
        meas = b.sum()
        return meas


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import spi

    slm = fakeSLM(shape=(360, 640))
    slm.doGraphic = False
    slm.roi = slm.pxlsPhysical[
        10:330, 0:640
    ]  # for speed of recovery,  make it divisible by 2 in some power
    slm.pxlsLogical = np.zeros((16, 16), dtype=np.uint8)
    np.fill_diagonal(slm.pxlsLogical, 1)

    test_image = np.copy(slm.pxlsPhysical).astype(float) * 0
    test_image[np.tril_indices(test_image.shape[0])] = 1.0
    test_image[:100, :100] = 1.0
    test_image[:100, -100:] = 2.0
    test_image[-100:, -100:] = 2.0
    slm.scene = test_image

    plt.figure(6)
    plt.clf()
    plt.imshow(test_image)
    plt.colorbar()
    plt.title("Scene")

    slm.logical2ROI()
    slm.physical2fb()
    slm.display()
    b = slm._measure()

    spi_results = spi.spi_hadamard(slm, method="positiveHalfHadamard")
    # spi_results = spi.spi_hadamard(slm, method="fullHaddamard")
    spi.saveSPIresults("spi_data.h5", spi_results)
    spi_results = spi.loadSPIresults("spi_data.h5")
    imgP, imgL = spi.reconstructPhysicalImageFromSPIresults(spi_results)

    plt.figure(8)
    plt.clf()
    plt.imshow(imgL)
    plt.colorbar()
    plt.title("Reconstruction Logical pixels")

    plt.figure(9)
    plt.clf()
    plt.imshow(imgP)
    plt.colorbar()
    plt.title("Reconstruction Physical pixels")
