# Math imports
from math import ceil

# Numpy imports
import numpy as np


class convInterpolate:
    """
    This class is a simple tool for simulating the output of an array of active pixels of an event-based camera.
    It is encapsulates all the parameters relevant for reproducing the events reported by each pixel as it is subject
    to temporarily varying illumination. It is intended to be a versatile block for creating event data sets from frames
    generated in ray-tracing software such as PANGU or Blender.

    The data fed to the converter should represent the electric current exiting the photodiode as pixels are being
    illuminated. The exact units depend on the type of event-based vision sensor that is being simulated, so care should
    be taken to make sure that they are consistent with the camera parameters input in the init method (in particular
    the time units!)

    Output is generated by linearly interpolating between succesive photo-current values of the pixels to check if and
    when they cross the next threshold. The convInterpolate class tracks the state of the pixels to assure that latency
    times are not violated, but does NOT model any non-linearities in pixel response. There is no absolute threshold for
    pixel saturation nor does this class handle simulating dark current. All device-level simulation must be handled by
    the user externally.

    Event convention is the following:
        [x, y, t, p]
         |  |  |  |_ event polatory
         |  |  |____ event time
         |  |_______ location along the array height
         |__________ location along the array width


    Developed in the Advanced Concepts Team at the European Space Agency by Olaf Sikorski.

    """

    def __init__(self, theta, T, latency, initial_image, log_input=False):
        """
        Initialises camera parameters and pixel states.

        Args:
            theta: float
                Threshold for
            T: float
                Period between subsequent frames, unit depends on user interpretation.
            latency: float
                Lacency of the pixel reactions, same unit at T.
            initial_image: 2D numpy array
                First frane, used to initialize the states of all pixels.
            log_input: bool (optional, default=False)
                Specify if the incoming data is logarithmically transformed (log_input=True)
                or requires transformation (log_input=False)
        """

        self.x_size = initial_image.shape[0]
        self.y_size = initial_image.shape[1]
        self.theta = theta
        self.T = T
        self.latency = latency
        self.log_input = log_input

        initial_image = initial_image.astype('float32')

        # Has to be a numpy array and not contain zeroes
        if self.log_input:
            initial_image = np.log10(initial_image)  # Create matrix of logged values

        self.current_states = initial_image  # Stores the values that pixel differentiators are zeroed at
        self.refract_period = np.zeros((self.x_size, self.y_size),
                                       dtype='float32')  # Stores ENDS of refractive periods for the pixels
        self.time = 0  # Time at which frame was recorded
        self.quick_burst = np.zeros((self.x_size, self.y_size))

        # quick_burst = 0   ->  no burst
        # quick_burst = 1   ->  positive
        # quick_burst = -1  ->  negative

        self.threshold_p = self.current_states + self.theta  # Constant threshold
        self.threshold_n = self.current_states - self.theta  # Constant threshold

        self.previous_image = initial_image

    def __del__(self):
        pass

    def update(self, new_image, look_up=True):
        """
        Takes new frame to update pixel states and generates events.

            Note:
                The input image has to be of dimensions (x_size, y_size).

            Args:
                new_image: 2D numpy array
                    Ndarray containing the next frame
                look_up: bool (optional, default=True)

            Returns:
                PARTIAL_LIST: list of lists
                    Chronologically sorted list of events [[x1, y1, t1, p1], [x2, y2, t2, p2], ...]
                image_out: 2D numpy array
                    ndarray marking which pixels fired events over this interval and with what polarity.
                    0.0 -> negative polarity event
                    0.5 -> no event
                    1.0 -> positive polarity event

            """

        # returned variables
        PARTIAL_LIST = []
        image_out = np.ones([self.x_size, self.y_size]) * 0.5
        new_image = new_image.astype('float32')

        if self.log_input:
            new_image = np.log10(new_image)  # Create matrix of logged values

        delta = new_image - self.previous_image
        # threshold = current_states * theta        # Dynamic threshold

        for x in range(self.x_size):  # Iterate over x dimension
            for y in range(self.y_size):  # Iterate over y dimension

                # Slope used for piecewise linear interpolation of pixel intensity
                slope = delta[x, y] / self.T

                # Do quick burst if flag is up
                if self.quick_burst[x, y] > 0:
                    event_time = self.refract_period[x, y]
                    # Update pixel state and set new end for refractory period
                    self.current_states[x, y] = self.previous_image[x, y] + slope * (event_time - self.time)
                    self.refract_period[x, y] = event_time + self.latency
                    self.threshold_p[x, y] = self.current_states[x, y] + self.theta
                    self.threshold_n[x, y] = self.current_states[x, y] - self.theta

                    if self.quick_burst[x, y] == 1:
                        PARTIAL_LIST.append([x, y, event_time, -1])
                    if self.quick_burst[x, y] == 2:
                        PARTIAL_LIST.append([x, y, event_time, 1])
                    self.quick_burst[x, y] = 0

                # Keeps looking for events until pixel refract period extends into next frame
                while self.refract_period[x, y] < (self.time + self.T):

                    # Case for increasing brightness
                    if self.threshold_p[x, y] < new_image[x, y] and slope > 0:

                        # Dummy value event visualisation, might delete later
                        image_out[x, y] = 1.0

                        # Linear estimate of threshold crossing instance
                        dt = abs((self.threshold_p[x, y] - self.previous_image[x, y]) / slope)

                        # This section calculates the registration of the event depending on the pixels refractory period
                        if self.refract_period[x, y] > self.time + dt:
                            event_time = self.refract_period[x, y]
                            self.current_states[x, y] = self.previous_image[x, y] + slope * (event_time - self.time)
                            self.refract_period[x, y] = event_time + self.latency
                        else:
                            event_time = ceil(self.time + dt)
                            self.current_states[x, y] = self.current_states[x, y] + self.theta
                            self.refract_period[x, y] = event_time + self.latency

                        # Append event to event list and update thresholds
                        PARTIAL_LIST.append([x, y, event_time, 1])
                        self.threshold_p[x, y] = self.current_states[x, y] + self.theta
                        self.threshold_n[x, y] = self.current_states[x, y] - self.theta

                    # Case for decreasing brightness
                    elif self.threshold_n[x, y] > new_image[x, y] and slope < 0:

                        image_out[x, y] = 0.0
                        # Linear estimate of threshold crossing instance
                        dt = abs((self.threshold_n[x, y] - self.previous_image[x, y]) / slope)

                        # This section calculates the registration of the event depending on the pixels refractory period
                        if self.refract_period[x, y] > self.time + dt:
                            event_time = self.refract_period[x, y]
                            # Update pixel state and set new end for refractory period
                            self.current_states[x, y] = self.previous_image[x, y] + slope * (event_time - self.time)
                            self.refract_period[x, y] = event_time + self.latency
                        else:
                            event_time = ceil(self.time + dt)
                            # Update pixel state and set new end for refractory period
                            self.current_states[x, y] = self.current_states[x, y] - self.theta
                            self.refract_period[x, y] = event_time + self.latency

                        PARTIAL_LIST.append([x, y, event_time, -1])
                        self.threshold_p[x, y] = self.current_states[x, y] + self.theta
                        self.threshold_n[x, y] = self.current_states[x, y] - self.theta

                    else:
                        break

                # track if pixel should fire immediately at the end of refractory period
                if (self.refract_period[x, y] >= (self.time + self.T)) and (self.threshold_p[x, y] < new_image[x, y]):
                    self.quick_burst[x, y] = 1
                elif (self.refract_period[x, y] >= (self.time + self.T)) and (self.threshold_n[x, y] > new_image[x, y]):
                    self.quick_burst[x, y] = 2

        # Update time and previous image
        self.time = self.time + self.T
        self.previous_image = new_image

        # Sort the generated events and append to returned list
        if len(PARTIAL_LIST) > 0:
            PARTIAL_LIST.sort(key=lambda z: z[2])

        return PARTIAL_LIST, image_out
