import numpy as np
import h5py
import json

from qolab.hardware.slm import stretch2mat
from qolab.hardware.slm import hadamard, matDotVec


def spi_hadamard(slm, method="positiveHalfHadamard"):
    """
    Recover an image which is rows x cols logical pixel by illuminating
    with Hadamard matrix patterns.
    Important: the reconstruction of the very first pixel is super sensitive to
    the measurement noise!

    Returns: dictionary with illumination masks, brightness, areas of logical pixels,
    and method one of [fullHaddamard, positiveHalfHadamard]. I.e. dictionary with fields
    {"masks", "brightness", "logical_pixels_area", "method"}

    The desired image does not need to be square but its rows*cols should be power of 2.
    For example 16 x 16 is ok, but so are 8 x 16 or 16 x 8; 1 x 4 is fine too.
    We will treat them as logical pixels since they will be put on top
    of physical SLM and thus can be non squared in physical space.

    If method="Hadamard", we recover illuminate with Hadamard>0 and Hadamard<0 matrix masks combo.
    If method="positiveHalfHadamard" we illuminate only with Hadamard>0 half.

    Important: each logical pixel must be the same in physical space, otherwise
    Hadamard mask assumption where all weight are 1 is broken, and we have to do much more
    complex math to recover image.
    """

    spi_results = {}
    spi_results.update({"slm": slm.getConfig()})

    if method != "positiveHalfHadamard":
        method = "fullHaddamard"
    # to recover rows x cols image, we need the same number of linearly independent sample masks
    Npixels = np.prod(slm.pxlsLogical.shape)
    hMasks = hadamard(
        Npixels, dtype=np.int8
    )  # Hadamard matrix with -1, 0, 1 and size rows x cols
    # mask are unwrapped and stacked, one row per rows x cols image

    """ On a computer with plenty of memory we can do just this:
    hMasksPositive = (hMasks == 1).astype(np.int8)
    hMasksNegative = (hMasks == -1).astype(np.int8)

    bP = measureMasksBrightness(hMasksPositive, slm=slm)
    bN = measureMasksBrightness(hMasksNegative, slm=slm)
    # but for an embedded computer with just 512 MB hadamard(128*128) takes 256 MB
    # so above will not leave us any memory,
    # since hMasks, hMasksPositive, and hMasksPositive will take 768 MB
    """
    # we will do some trickery, to modify matrix in place to be memory efficient
    hMasks += 1
    hMasks //= 2  # this creates hMasksPositive

    print("Collecting positive masks")
    bP = slm.measureMasksBrightness(hMasks)
    bP = np.array(bP)

    a = slm.getLogicalPxlsArea()

    if method == "positiveHalfHadamard":
        print("Measuring background illumination")
        bDark = slm.measureMasksBrightness(
            np.zeros((1, hMasks.shape[1]))
        )  # dark background level
        bP -= bDark[0]  # recall that measureMasksBrightness gives back list

        spi_results.update(
            {
                "masks": hMasks,
                "brightness": bP,
                "logical_pixels_area": a,
                "method": method,
            }
        )
        return spi_results

    hMasks -= 1
    hMasks *= -1  # this creates hMasksNegative
    print("Collecting negative masks")
    bN = slm.measureMasksBrightness(hMasks)

    # we do not need to subtract background since we are using bP - bN which does it automatically
    # bDark = bN[0]  # 1st negative Hadamard is special: it is all zeros
    # bP -= bDark
    # bN -= bDark
    b = np.array(bP) - np.array(bN)

    hMasks *= -2
    hMasks += 1  # now we are back to Hadamard masks

    spi_results.update(
        {
            "masks": hMasks,
            "brightness": b,
            "logical_pixels_area": a,
            "method": method,
        }
    )
    return spi_results


def saveSPIresults(filename, spi_results):
    """Save SPI results
    spi_results must have the following fields:
       {"masks": masks, "brightness": b, "logical_pixels_area": a, "method": method}
    """
    masks = spi_results["masks"]
    areas = spi_results["logical_pixels_area"]
    b = spi_results["brightness"]
    method = spi_results["method"]

    with h5py.File(filename, "w") as hdf5_file:
        hdf5_file.attrs["method"] = method
        hdf5_file.attrs["slm"] = json.dumps(spi_results["slm"])
        hdf5_file.create_dataset(
            "masks", data=masks, compression="gzip", compression_opts=9
        )
        hdf5_file.create_dataset(
            "logical_pixels_area", data=areas, compression="gzip", compression_opts=9
        )
        hdf5_file.create_dataset(
            "brightness", data=b, compression="gzip", compression_opts=9
        )


def loadSPIresults(filename):
    """Load SPI results
    returns spi_results which  must have the following fields:
       {"masks": masks, "brightness": b, "logical_pixels_area": a, "method": method}
    """
    with h5py.File(filename, "r") as hdf5_file:
        method = hdf5_file.attrs["method"]
        slm = json.loads(hdf5_file.attrs["slm"])
        masks = np.array(hdf5_file.get("masks"))
        b = np.array(hdf5_file.get("brightness"))
        areas = np.array(hdf5_file.get("logical_pixels_area"))
    return {
        "masks": masks,
        "brightness": b,
        "logical_pixels_area": areas,
        "method": method,
        "slm": slm,
    }


def reconstructLogicalImageFromSPIresults(spi_results):
    """Produce resulting image from the spi_results dictionary
    spi_results must have the following fields  {"masks": masks, "brightness": b, "logical_pixels_area": a, "method": method}
    """

    masks = spi_results["masks"]
    areas = spi_results["logical_pixels_area"]
    b = spi_results["brightness"]
    method = spi_results["method"]
    num_of_diff_areas = np.unique(areas).size

    # To find and image, we are solving b = mask @ img => img = inv(masks) @ b
    if (method == "fullHaddamard") and (num_of_diff_areas == 1):
        # when all areas are the same we can use Hadamard property inv(H) = H, which speeds up calculation
        img = matDotVec(masks, b).reshape(areas.shape)
        img /= masks.shape[0] * areas[0, 0]  # normalization of Hadamard matrix
    if (method == "positiveHalfHadamard") and (num_of_diff_areas == 1):
        # we can recover Hadamard matrix and would be measured corresponding brightness
        bAllOn = b[0]  # the very 1st mask is everybody on
        hMasks = masks * 2 - 1
        bH = b * 2 - bAllOn
        # when all areas are the same we can use Hadamard property inv(H) = H, which speeds up calculation
        img = matDotVec(hMasks, bH).reshape(areas.shape)
        img /= masks.shape[0] * areas[0, 0]  # normalization of Hadamard matrix
    else:  # fallback to general method. It works but slow for large matrices
        # solving the equation b = mask @ img is faster then img = inv(masks) @ B
        if num_of_diff_areas == 1:
            img = np.linalg.solve(masks, b).reshape(areas.shape) / areas[0, 0]
        else:
            Nlp = np.prod(areas.shape)  # number of logical pixels
            # masks are weighted by the areas of logical pixels
            img = np.linalg.solve(masks * areas.reshape(Nlp), b).reshape(areas.shape)
    return img


def reconstructPhysicalImageFromSPIresults(spi_results):
    """Produce resulting physical space image from the spi_results dictionary also return logical pixel image.
    spi_results must have the following fields  {"masks": masks, "brightness": b, "logical_pixels_area": a, "method": method}
    Logical image calculated as an intermediate results, so we return it too.
    """
    img_logical = reconstructLogicalImageFromSPIresults(spi_results)

    slm = spi_results["slm"]
    img = np.zeros(slm["DeviceConfig"]["PhysicalShape"])
    roi_pos = slm["DeviceConfig"]["ROIspecs"]["pos"]
    roi_shape = slm["DeviceConfig"]["ROIspecs"]["shape"]
    roi = img[
        roi_pos[0] : roi_pos[0] + roi_shape[0], roi_pos[1] : roi_pos[1] + roi_shape[1]
    ]
    stretch2mat(img_logical, roi)
    return img, img_logical


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from slm import fakeSLM

    slm = fakeSLM((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_hadamard(slm, method="positiveHalfHadamard")
    # spi_results = spi_hadamard(slm, method="fullHaddamard")
    saveSPIresults("spi_data.h5", spi_results)
    spi_results = loadSPIresults("spi_data.h5")
    imgp, imgl = 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")
