import numpy as np

from .utilities import convert_tmin_tmax_framerate_to_trange
from bmtk.simulator.filternet.io_tools import io


class Movie(object):
    def __init__(self, data, row_range=None, col_range=None, labels=('time', 'y', 'x'),
                 units=('second', 'pixel', 'pixel'), frame_rate=None, t_range=None, padding=False):
        self.data = data
        self.labels = labels
        self.units = units
        self.padding=padding
        assert(units[0] == 'second')

        if t_range is None:
            self.frame_rate = float(frame_rate)
            self.t_range = np.arange(data.shape[0])*(1./self.frame_rate)
        else:
            self.t_range = np.array(t_range)
            self.frame_rate = 1./np.mean(np.diff(t_range))
            
        if row_range is None:
            self.row_range = np.arange(data.shape[1])
        else:
            self.row_range = np.array(row_range)
        if col_range is None:
            self.col_range = np.arange(data.shape[2])
        else:
            self.col_range = np.array(col_range)

    @property
    def ranges(self):
        return self.t_range, self.row_range, self.col_range

    def imshow_summary(self, ax=None, show=True, xlabel=None):
        import matplotlib.pyplot as plt

        if ax is None:
            _, ax = plt.subplots(1,1)
        
        t_vals = self.t_range.copy()
        y_vals = self.data.mean(axis=2).mean(axis=1)
        ax.plot(t_vals, y_vals)
        ax.set_ylim(y_vals.min()-np.abs(y_vals.min())*.05, y_vals.max()+np.abs(y_vals.max())*.05)
        
        if not xlabel is None:
            ax.set_xlabel(xlabel)
            
        ax.set_ylabel('Average frame intensity')
        
        if show:
            plt.show()
            
        return ax, (t_vals, y_vals)
    
    def imshow(self, t, show=True, vmin=-1, vmax=1, cmap=None):
        import matplotlib.pyplot as plt

        cmap = cmap or plt.cm.gray
        ti = int(t*self.frame_rate)
        data = self.data[ti, :, :]
        plt.imshow(data, vmin=vmin, vmax=vmax, cmap=cmap)
        plt.colorbar()
        if show:
            plt.show()

    def play(self, interval=None, skip=None, playback_speed=1.0, cmap=None, vmin=-1.0, vmax=1.0, save_as=None,
             show=True):
        """

        :param interval:
        :param skip: skip every Nth frame, useful for very large data sets
        :param playback_speed:
        :param cmap: color map, default to gray
        :param save_as: If set to a valid (mp4 file) path animation will be saved. Otherwise it won't
        :param show: display the animation, default True
        """
        import matplotlib.pyplot as plt
        import matplotlib.animation as animation

        cmap = cmap or plt.cm.gray

        data = self.data
        times = self.t_range

        if skip is not None:
            data = data[::skip, :, :]
            times = self.t_range[::skip]

        if interval is None:
            interval = np.mean(np.diff(times))*1000

        frames = []
        fig = plt.figure()
        ax = fig.add_subplot(111)
        for i in range(len(times)):
            # TODO: Find a way to display a running clock as a x-label or title
            frame_img = ax.imshow(data[i, :, :], cmap=cmap, animated=True, vmin=vmin, vmax=vmax,
                                  extent=[self.col_range[0], self.col_range[-1], self.row_range[0], self.row_range[-1]])
            title = ax.text(0.99, 0.000, "{:05.2f} s".format(times[i]),
                            size=plt.rcParams["axes.titlesize"],
                            horizontalalignment='right',
                            verticalalignment='bottom',
                            transform=ax.transAxes, bbox={'boxstyle': 'round', 'facecolor': 'white'})

            frames.append([frame_img, title])

        ani = animation.ArtistAnimation(fig, frames, interval=interval*playback_speed, blit=True, repeat_delay=1000)

        if save_as is not None:
            ani.save(save_as)

        if show:
            plt.show()

    def __add__(self, other):
        assert self.labels == other.labels
        assert self.units == other.units
        assert self.frame_rate == other.frame_rate
        np.testing.assert_almost_equal(self.col_range, other.col_range)
        np.testing.assert_almost_equal(self.row_range, other.row_range)
        
        new_data = np.empty((len(self.t_range)+len(other.t_range)-1, len(self.row_range), len(self.col_range)))
        new_data[:len(self.t_range), :, :] = self.data[:, :, :]
        new_data[len(self.t_range):, :, :] = other.data[1:, :, :]
        
        return Movie(new_data, row_range=self.row_range.copy(), col_range=self.col_range.copy(), labels=self.labels,
                     units=self.units, frame_rate=self.frame_rate)

    def __getitem__(self, *args):
        return self.data.__getitem__(*args)

    @staticmethod
    def normalize_matrix(m_data, domain=None):
        """Attempts to take a numpy matrix "movie" and normalize the contrast values to range [-1.0, 1.0].

        :param m_data: A numpy matrix TxRxC containing an original "movie" data
        :param domain: If A tuple of values [min_contrast, max_contrast], will use those values to normalize. Otherwise
            will attempt to determine original contrast range from the m_data.
        :return: m_data that has been normalized
        """
        if isinstance(domain, (list, tuple, np.ndarray)):
            # Allow users to specify the contrast range of the original movie [c_min, c_max] that will be
            # converted to range [-1.0, +1.0]
            if len(domain) != 2:
                raise ValueError(
                    'Unable to normalize movie; please specify to minimum and maximum contrast for original movie with '
                    'list [min_contrast, max_contrast]'
                )
            c_min, c_max = domain[0], domain[1]
            if c_min >= c_max:
                raise ValueError(
                    'Unable to normalize movie; original movie range must be defined with min_contrast < max_contrast'
                )
            elif m_data.min() < c_min or m_data.max() > c_max:
                raise ValueError(
                    'Unable to normalize movie; trying to convert movie from contrast values ' +
                    '[{}, {}] --> [-1.0, 1.0], '.format(c_min, c_max) +
                    'but movie contains values outside of original domain.'
                )
        else:
            # Use the existing movie to try to best guess the domain the original movie is already in; in most
            # cases it will be [0.0, 1.0] or [0, 256]
            m_min, m_max = m_data.min(), m_data.max()
            if m_min >= 0.0 and m_max <= 1.0:
                c_min, c_max = 0.0, 1.0
            elif m_min >= -1.0 and m_max <= 1.0:
                # All values fall between [-1.0, 1.0], likely movie is likely already normalized.
                io.log_debug('Input movie is already normalized between contrast values [-1.0, 1.0].')
                return m_data
            elif m_min >= 0 and m_max < 256:
                c_min, c_max = 0.0, 255.0
            else:
                c_min, c_max = m_min, m_max

        if c_min == c_max:
            raise ValueError(
                'Unable to normalize movie; please specify a range for the contrast min and max of the input movie.'
            )

        io.log_debug('Normalizing input movie from contrast range [{}, {}] --> [-1.0, 1.0].')
        return (m_data - c_min)*2.0/(c_max - c_min) - 1.0


# TODO: Instead of using subclasses, convert the following movie type to a Factory class/function that returns a Movie
#   object
class FullFieldMovie(Movie):
    def __init__(self, f, rows, cols, frame_rate=24.0):
        if np.isscalar(rows):
            rows = np.arange(0, rows)

        if np.isscalar(cols):
            cols = np.arange(0, cols)

        self.row_range = rows
        self.col_range = cols
        self.frame_size = (len(self.row_range), len(self.col_range))
        self._frame_rate = frame_rate
        self.f = f
        
    @property
    def frame_rate(self):
        return self._frame_rate
    
    @property
    def data(self):
        return self
    
    def __getitem__(self, *args):
        t_inds, x_inds, y_inds = args[0]
        assert (len(x_inds) == len(y_inds)) and (len(y_inds) == len(t_inds))
        
        # Convert frame indices to times:
        t_vals = (1./self.frame_rate)*t_inds
        
        # Evaluate and return:
        return self.f(t_vals)
    
    def full(self, t_min=0.0, t_max=None):
        # Compute t_range
        t_range = convert_tmin_tmax_framerate_to_trange(t_min, t_max, self.frame_rate)
         
        nt = len(t_range)
        nr = len(self.row_range)
        nc = len(self.col_range)
        a, b, c = np.meshgrid(range(nt), range(nr), range(nc))
        af, bf, cf = map(lambda x: x.flatten(), [a, b, c])
        data = np.empty((nt, nr, nc))
        data[af, bf, cf] = self.f(t_range[af])
        
        return Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                     units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)

    def create_movie(self, t_max, t_min=0.0):
        return self.full(t_min=t_min, t_max=t_max)


class FullFieldFlashMovie(FullFieldMovie):
    def __init__(self, row_range, col_range, t_on, t_off, max_intensity=1, frame_rate=24):
        assert(t_on < t_off)

        def f(t):
            return np.piecewise(t, *zip(*[(t < t_on, 0), (np.logical_and(t_on <= t, t < t_off), max_intensity),
                                          (t_off <= t, 0)]))

        super(FullFieldFlashMovie, self).__init__(f, row_range, col_range, frame_rate=frame_rate)


class GratingMovie(Movie):
    def __init__(self, row_size, col_size, frame_rate=1000.0):
        self.row_size = row_size  # in degrees
        self.col_size = col_size  # in degrees
        self.frame_rate = float(frame_rate)  # in Hz

    def create_movie(self, t_min=0, t_max=1, gray_screen_dur=0, cpd=0.05, temporal_f=4, theta=45,
                     phase=0., contrast=1.0, degrees_per_pixel=None, row_size_new=None, col_size_new=None):
        """Create the grating movie with the desired parameters

        :param t_min: start time in seconds
        :param t_max: end time in seconds
        :param gray_screen_dur: Duration of gray screen before grating stimulus starts
        :param cpd: cycles per degree
        :param temporal_f: in Hz
        :param theta: orientation angle, in degrees
        :param phase: temporal phase, in degrees
        :param contrast: the maximum constrast, must be between 0 and 1.0
        :param degrees_per_pixel: pixel pitch of the movie in degrees (default (if None) is 1.0/(cpd*10))
        :param row_size_new: Use to truncate screen, by default leaves original row dimension
        :param col_size_new: Use to truncate screen, by default leaves original col dimension
        :return: Movie object of grating with desired parameters
        """
        assert contrast <= 1, "Contrast must be <= 1"
        assert contrast > 0, "Contrast must be > 0"

        if degrees_per_pixel is None:  # default behavior when not specified
            degrees_per_pixel = 1.0 / (float(cpd) * 10)  # To make sure no aliasing occurs
        self.row_range = np.linspace(0, self.row_size, int(self.row_size/degrees_per_pixel), endpoint=True)
        self.col_range = np.linspace(0, self.col_size, int(self.col_size/degrees_per_pixel), endpoint=True)
        numberFramesNeeded = int(round(self.frame_rate * (t_max - gray_screen_dur))) + 1
        time_range = np.linspace(0, t_max - gray_screen_dur, numberFramesNeeded, endpoint=True)

        # Creates a drifting grating panel
        tt, yy, xx = np.meshgrid(time_range, self.row_range, self.col_range, indexing='ij')
        theta_rad = np.pi*(180 - theta) / 180.
        phase_rad = np.pi*(180 - phase) / 180.
        xy = xx * np.cos(theta_rad) + yy * np.sin(theta_rad)
        data = contrast*np.sin(2*np.pi*(cpd * xy + temporal_f *tt) + phase_rad)

        # truncate
        if row_size_new != None:
            self.row_range = np.linspace(0, row_size_new, data.shape[1], endpoint=True)

        if col_size_new != None:
            self.col_range = np.linspace(0, col_size_new, data.shape[2], endpoint=True)

        if gray_screen_dur > 0:
            # just adding one or two seconds to gray screen so flash never "happens"
            m_gray = FullFieldFlashMovie(self.row_range, self.col_range, gray_screen_dur + 1, gray_screen_dur + 2,
                                         frame_rate=self.frame_rate).full(t_max=gray_screen_dur)
            mov = m_gray + Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                                 units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)
        else:
            mov = Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                        units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)

        return mov


class LoomingMovie(Movie):
    def __init__(self, row_size, col_size, frame_rate=1000.):
        self.row_size = row_size  # in degrees
        self.col_size = col_size  # in degrees
        self.frame_rate = float(frame_rate)  # in Hz

    def create_movie(self, t_looming=1, gray_screen_dur=0.5):
        """Create the looming movie with the desired parameters

        :param t_looming: duration of time looming
        :param gray_screen_dur:
        """
        degrees_per_pixel = 1.0  # To make sure no aliasing occurs
        self.row_range = np.linspace(0, self.row_size, int(self.row_size/degrees_per_pixel), endpoint=True)
        self.col_range = np.linspace(0, self.col_size, int(self.col_size/degrees_per_pixel), endpoint=True)
        loomingFramesNeeded = int(round(self.frame_rate * t_looming))
        grayScreenFrames = int(round(self.frame_rate * gray_screen_dur))
        time_range = np.linspace(0, t_looming, loomingFramesNeeded, endpoint=True)

        data = np.zeros((grayScreenFrames + loomingFramesNeeded, self.row_size, self.col_size))
        # mgrid is a mesh creation helper
        xx, yy = np.mgrid[:self.row_size, :self.col_size]
        # yy, xx = np.meshgrid(time_range, self.row_range, self.col_range, indexing='ij')
        # circles contains the squared distance to the (middle, middle) point
        # we are just using the circle equation learnt at school
        circle = (xx - self.row_size / 2) ** 2 + (yy - self.col_size / 2) ** 2
        radius = 25  # Cirlces grow to this size

        for i in range(loomingFramesNeeded):
            # Create filled circle
            data[grayScreenFrames + i, :, :] = (circle <= (i/(4./3)) % (radius**2))
        data *= -1   # Adjusting so have dark looming and not bright looming disk

        mov = Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                    units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)

        return mov
