#!/usr/bin/env python

import numpy as np
import time
from rebin import rebin
from plico.utils.zernike_generator import ZernikeGenerator
from plico.types.zernike_coefficients import ZernikeCoefficients
from pysilico_server.devices.base_simulated_camera import BaseSimulatedCamera


__version__ = "$Id: simulated_camera.py 293 2017-06-21 17:10:57Z lbusoni $"





class SimulatedPyramidWfsCamera(BaseSimulatedCamera):
    SENSOR_H= 1024
    SENSOR_W= 1360
    DTYPE= np.uint16
    MAX_VALUE= 4096


    def __init__(self, name='Simulated Pyramid Wfs Camera'):
        BaseSimulatedCamera.__init__(self, name)
        self._pupilRadius= None
        pupilsSeparation= 520
        opticalAxis= (self.SENSOR_H / 2, self.SENSOR_W / 2)
        self._pupilsCenter= self._computeCenters(pupilsSeparation,
                                                 opticalAxis)
        self._tipTiltCoefficients= np.zeros(2)
        self._wavefront= None
        self._zernikeGenerator= None

        self.setPupilsRadiusInUnbinnedPixels(80)
        self.setScaleInMeterPerPixel(
            10e-3/ (2* self.getPupilsRadiusInUnbinnedPixels()))
        self.setSlopeSaturationInRadians(40e-3)

        self._logger.notice('Pupil centers %s' % (self._pupilsCenter))
        self._logger.notice('Pupil radius %s' % (self._pupilRadius))



    def setPupilsRadiusInUnbinnedPixels(self, radius):
        self._pupilRadius= radius
        self._wavefront= np.zeros((2* self._pupilRadius, 2* self._pupilRadius))
        self._zernikeGenerator= ZernikeGenerator(2* self._pupilRadius)


    def getPupilsRadiusInUnbinnedPixels(self):
        return self._pupilRadius


    def setTilt(self, tipTiltCoefficients):
        self._tipTiltCoefficients= tipTiltCoefficients


    def setWavefront(self, wavefront):
        self._wavefront= wavefront


    def _computeWavefrontFromZernikeCoefficients(self, coeff, zg):
        wf= 0.* zg.getZernike(1)
        for y in coeff.zernikeIndexes():
            wf+= coeff.getZ([y])* zg.getZernike(y)
        return wf


    def setWavefrontFromZernikeVector(self, zernikeVector):
        wf= self._computeWavefrontFromZernikeCoefficients(
            ZernikeCoefficients.fromNumpyArray(np.array(zernikeVector)),
            self._zernikeGenerator)
        self.setWavefront(wf)


    def _computeCenters(self, pupilsSeparation, opticalAxis):
        pupDist= int(0.5 * pupilsSeparation)
        center= opticalAxis
        res= np.zeros((4, 2))
        res[0, :]= [center[0] + pupDist, center[1] - pupDist]
        res[1, :]= [center[0] + pupDist, center[1] + pupDist]
        res[2, :]= [center[0] - pupDist, center[1] - pupDist]
        res[3, :]= [center[0] - pupDist, center[1] + pupDist]
        return res


    def setPupilsCenterInUnbinnedPixels(self, centers):
        assert centers.shape == (4, 2)
        self._pupilsCenter= centers


    def getPupilsCenterInUnbinnedPixels(self):
        return self._pupilsCenter


    def _computeFluxesForTilt(self):
        z2= self._tipTiltCoefficients[0]
        z3= self._tipTiltCoefficients[1]
        tot= 4.0
        tl= 0.25 * (z3 - z2 + 1) * tot
        tr= 0.25 * (z3 + z2 + 1) * tot
        bl= 0.25 * (-z3 - z2 + 1) * tot
        br= 0.25 * (-z3 + z2 + 1) * tot
        self._logger.debug(
            "tl,tr,bl,br %g,%g,%g,%g - z2,z3: %g,%g" %
            (tl, tr, bl, br, (tr+ br- tl- bl) / tot,
             (tl+ tr- bl- br) / tot))
        return tl, tr, bl, br


    def _pupilImagesFromWavefront(self, wavefront):
        Amat= np.array([[1., 0, 1, 0],
                        [0., 1, 0, 1],
                        [0., 0, 1, 1],
                        [1., 1, 0, 0]])


        (dyMap, dxMap)= np.gradient(wavefront / self._scaleInMeterPerPixel)
        dxMapClipped= dxMap.clip(min=-self._slopeSaturationInRad,
                                 max=self._slopeSaturationInRad)
        dyMapClipped= dyMap.clip(min=-self._slopeSaturationInRad,
                                 max=self._slopeSaturationInRad)
        dx=2. / np.pi * np.arcsin(dxMapClipped/ self._slopeSaturationInRad)
        dy=2. / np.pi * np.arcsin(dyMapClipped/ self._slopeSaturationInRad)
        (sz0, sz1)= dxMap.shape
        totalIntensity= 1.* sz0* sz1
        B= totalIntensity / (sz0 * sz1) / 2 * np.array(
            [1- dx, 1+ dx, 1- dy, 1+ dy])

        res=np.linalg.lstsq(Amat, B.reshape((4, sz0 * sz1)), rcond=-1)
        pupils= res[0].reshape((4, sz0, sz1))

        if isinstance(wavefront, np.ma.MaskedArray):
            for i in range(4):
                pupils[i][wavefront.mask]=0

        return pupils


    def _computeFrameFromWavefront(self):
        radius= self._pupilRadius
        cc= self._pupilsCenter
        maxNoiseInCount= self._noiseInCount

        pupils= self._pupilImagesFromWavefront(self._wavefront)
        pixels = np.zeros((self.SENSOR_H, self.SENSOR_W),
                          dtype=np.float)

        pixels[int(cc[0, 0] - radius): int(cc[0, 0] + radius),
               int(cc[0, 1] - radius): int(cc[0, 1] + radius)]= pupils[0]
        pixels[int(cc[1, 0] - radius): int(cc[1, 0] + radius),
               int(cc[1, 1] - radius): int(cc[1, 1] + radius)]= pupils[1]
        pixels[int(cc[2, 0] - radius): int(cc[2, 0] + radius),
               int(cc[2, 1] - radius): int(cc[2, 1] + radius)]= pupils[2]
        pixels[int(cc[3, 0] - radius): int(cc[3, 0] + radius),
               int(cc[3, 1] - radius): int(cc[3, 1] + radius)]= pupils[3]

        normalize= self._totalFluxPerMilliSecond * \
            self.exposureTime() / pixels.sum()
        pixels*= normalize

        if maxNoiseInCount >0:
            pixels+= np.random.normal(
                3* maxNoiseInCount,
                maxNoiseInCount,
                size=(self.SENSOR_H, self.SENSOR_W))

        if self._binning != 1:
            pixels= rebin(pixels, self._binning) * self._binning**2
        pixels= pixels.clip(min=0, max=self.MAX_VALUE)

        self._counter += 1
        time.sleep(1./ self.getFrameRate())
        return pixels.astype(self.DTYPE)


    def setSlopeSaturationInRadians(self, slopeSaturationInRadians):
        self._slopeSaturationInRad= slopeSaturationInRadians


    def getSlopeSaturationInRadians(self):
        return self._slopeSaturationInRad


    def setScaleInMeterPerPixel(self, scaleInMeterPerPixel):
        self._scaleInMeterPerPixel= scaleInMeterPerPixel


    def getScaleInMeterPerPixel(self):
        return self._scaleInMeterPerPixel
