import os
import matplotlib
matplotlib.use("Agg") # agg-backend, so we can create figures without x-server (no PDF, just PNG etc.)
from elasticite import EdgeGrid
e = EdgeGrid()
fps = 24.
loop = 1
autoplay = 0
duration = 4.
figpath = '../files/elasticite/'

import matplotlib as mpl
import matplotlib.pyplot as plt
from moviepy.video.io.bindings import mplfig_to_npimage
import moviepy.editor as mpy

fig_mpl, ax = plt.subplots(1, figsize=(1, 1), facecolor='white')

def draw_elementary_pattern(ax, center): 
    #ax.add_artist(mpl.patches.Wedge(center-1., 1., 0, 180, width=.1)) #center, r, theta1, theta2, width=None
    #ax.add_artist(mpl.patches.Wedge(-center+1., 1., 0, 180, width=.1))
    # class matplotlib.patches.RegularPolygon(xy, numVertices, radius=5, orientation=0, **kwargs)¶
    ax.add_artist(mpl.patches.RegularPolygon((.5,.5), 5, center, facecolor='r'))
    ax.add_artist(mpl.patches.RegularPolygon((.5,.5), 5, center/2, facecolor='w'))
    ax.add_artist(mpl.patches.RegularPolygon((.5,.5), 5, center/4, facecolor='g'))

def make_frame_mpl(t):
    ax = fig_mpl.add_axes([0., 0., 1., 1.], axisbg='w')
    ax.cla()
    plt.setp(ax, xticks=[])
    plt.setp(ax, yticks=[])
    #ax.axis(c='b', lw=0, frame_on=False)
    ax.grid(b=False, which="both")
    draw_elementary_pattern(ax, t/duration)
    return mplfig_to_npimage(fig_mpl) # RGB image of the figure

animation = mpy.VideoClip(make_frame_mpl, duration=duration)

name, vext = 'elasticite_test', '.mp4'
if not os.path.isfile(os.path.join(figpath, name + '.mp4')):
    animation.write_videofile(figpath + name + vext, fps=fps)
    animation.write_videofile(figpath + name + e.vext, fps=fps)
e.ipython_display(name)

"""
Explanation: On va maintenant utiliser une image naturelle comme entrée. 
At each time, the pipeline is the  following:

take an image, 
turn into blocks corresponding to the edges' centers,
into each block determine the most likely orientation

<!-- TEASER_END -->

Let's first create a dummy movie:
End of explanation
"""


import imageio

reader = imageio.get_reader(figpath + name + vext)
for i, im in enumerate(reader):
    print('Mean of frame %i is %1.1f' % (i, im.mean()))
    if i > 15: break

"""
Explanation: Now read this clip using imageio:
End of explanation
"""


def mat2ipn(mat):   
    # create a temporary file
    import tempfile
    filename = tempfile.mktemp(suffix='.png')
    # Use write_png to export your wonderful plot as png ! 
    import vispy.io as io
    imageio.imwrite(filename, mat)
    from IPython.core.display import display, Image
    return display(Image(filename))    
mat2ipn(im)
#from holoviews import Image
#from holoviews import HoloMap, Dimension
#%load_ext holoviews.ipython

"""
Explanation: Let's consider one frame, and a ROI defined by its center and width:
End of explanation
"""


from NeuroTools.parameters import ParameterSet
from SLIP import Image
slip = Image(ParameterSet({'N_X':im.shape[1], 'N_Y':im.shape[0]}))

import numpy as np
im_ = im.sum(axis=-1)
print(im_.shape)

X_=.3
Y_=.5
w=.2
im_ = im_ * np.exp(-.5*((slip.x-X_)**2+(slip.y-Y_)**2)/w**2)
mat2ipn(im_)

"""
Explanation: trying to guess orientations using LogGabors
End of explanation
"""


from LogGabor import LogGabor
lg = LogGabor(slip)
N_theta = 24
thetas, E = np.linspace(0, np.pi, N_theta), np.zeros((N_theta,))
for i_theta, theta in enumerate(thetas):
    params= {'sf_0':.3, 'B_sf': .3, 'theta':theta, 'B_theta': .1}
    FT_lg = lg.loggabor(0, 0, **params)
    E[i_theta] = np.sum(np.absolute(slip.FTfilter(im_.T, FT_lg, full=True))**2)
    print(theta*180/np.pi, E[i_theta])       

"""
Explanation: Now, we will test the energy of different orientations :
End of explanation
"""


print(np.argmax(E), thetas[np.argmax(E)]*180/np.pi)

"""
Explanation: Now select the most likely:
End of explanation
"""


from NeuroTools.parameters import ParameterSet
from SLIP import Image
from LogGabor import LogGabor
import numpy as np
slip = Image(ParameterSet({'N_X':im.shape[1], 'N_Y':im.shape[0]}))
lg = LogGabor(slip)
N_theta = 24
def theta_max(im, X_=.0, Y_=.0, w=.3):
    im_ = im.sum(axis=-1)
    im_ = im_ * np.exp(-.5*((slip.x-X_)**2+(slip.y-Y_)**2)/w**2)
    thetas, E = np.linspace(0, np.pi, N_theta), np.zeros((N_theta,))
    for i_theta, theta in enumerate(thetas):
        params= {'sf_0':.3, 'B_sf': .3, 'theta':theta, 'B_theta': .1}
        FT_lg = lg.loggabor(0, 0, **params)
        E[i_theta] = np.sum(np.absolute(slip.FTfilter(im_.T, FT_lg, full=True))**2)
    return np.pi/2 - thetas[np.argmax(E)]     

e.reader = imageio.get_reader(figpath + 'elasticite_test.mp4', loop=True)
for i, im in enumerate(reader):
    print(i, theta_max(im, X_=.3, Y_=.3, w=.3)*180./np.pi)
    if i > 5: break

"""
Explanation: wrapping things in one function:
End of explanation
"""


name = 'trame_loggabor'
import numpy as np
from EdgeGrid import EdgeGrid
e = EdgeGrid()

import imageio
e.reader = imageio.get_reader(figpath + 'elasticite_test.mp4', loop=True)
def make_lames(e):
    im = e.reader.get_next_data()
    for i in range(e.N_lame):
        e.lames[2, i] = e.theta_max(im, X_=e.lames[0, i], Y_=e.lames[1, i], w=.05)
    return e.lames[2, :]

e.make_anim(name, make_lames, duration=duration)
e.ipython_display(name)

"""
Explanation: We retrieve the centers and span of all edges from the EdgeGrid class:
End of explanation
"""


name = 'trame_sobel_orientations'
if not os.path.isfile(os.path.join(figpath, name + '.mp4')):
    from EdgeGrid import EdgeGrid
    e = EdgeGrid()

    import imageio
    e.reader = imageio.get_reader(figpath + 'elasticite_test.mp4', loop=True)

    import matplotlib.pyplot as plt
    import numpy as np
    from moviepy.video.io.bindings import mplfig_to_npimage
    import moviepy.editor as mpy

    # DRAW A FIGURE WITH MATPLOTLIB
    fps = 24.
    duration = 4.
    fig_mpl, ax = plt.subplots(1, 2, figsize=(10,5), facecolor='white')
    def make_frame_mpl(t):
        import numpy as np
        sobel = np.array([[1,   2,  1,],
                          [0,   0,  0,],
                          [-1, -2, -1,]])
        im = e.reader.get_next_data()

        im_ = im.sum(axis=-1)
        from scipy.signal import convolve2d
        #im_ = im_ * np.exp(-.5*((slip.x-X_)**2+(slip.y-Y_)**2)/w**2)
        ax[0].imshow(convolve2d(im_, sobel, 'same'))
        ax[1].imshow(convolve2d(im_, sobel.T, 'same'))
        return mplfig_to_npimage(fig_mpl) # RGB image of the figure

    animation = mpy.VideoClip(make_frame_mpl, duration=duration)
    animation.write_videofile(os.path.join(figpath, name + '.mp4'), fps=fps)
e.ipython_display(name)

"""
Explanation: trying to guess orientations using Sobel filters
dans ce cas, on voit que les filtres orientés sont corrects, mais c'est un peu overkill (et lent) donc on peut préférer utiliser des filtres orientés plus simples, les filtres de Sobel, soit pour les horizontales la matrice:
[1   2  1]
[0   0  0]
[-1 -2 -1]

et son transposé (pour les verticales).
End of explanation
"""


name = 'trame_sobel_orientation'
import os
if True or not os.path.isfile(os.path.join(figpath, name + '.mp4')):
    from EdgeGrid import EdgeGrid
    e = EdgeGrid()

    import imageio
    e.reader = imageio.get_reader(figpath + 'elasticite_test.mp4', loop=True)

    import matplotlib.pyplot as plt
    import numpy as np
    from moviepy.video.io.bindings import mplfig_to_npimage
    import moviepy.editor as mpy

    # DRAW A FIGURE WITH MATPLOTLIB
    fps = 24.
    duration = 4.
    
    
    def make_frame_mpl(t):
        fig_mpl, ax = plt.subplots(figsize=(5,5), facecolor='white')

        import numpy as np
        sobel = np.array([[1,   2,  1],
                          [0,   0,  0],
                          [-1, -2, -1]])
        im = e.reader.get_next_data()

        im_ = im.sum(axis=-1)
        N_X, N_Y = im_.shape
        x, y = np.mgrid[0:1:1j*N_X, 0:1:1j*N_Y]
        # mask = np.exp(-.5*((x-.5)**2+(y-.5)**2)/.1**2)       
        blur = np.array([[1, 2, 1],
                         [1, 8, 2],
                         [1, 2, 1]])

        from scipy.signal import convolve2d
        im_X = convolve2d(im_, sobel, 'same')
        im_Y = convolve2d(im_, sobel.T, 'same')
        for i in range(10):
            im_X = convolve2d(im_X, blur, 'same')
            im_Y = convolve2d(im_Y, blur, 'same')
        mappable = ax.imshow(np.arctan2(im_Y, im_X)*180/np.pi, origin='lower')
        fig_mpl.colorbar(mappable)
        return mplfig_to_npimage(fig_mpl) # RGB image of the figure

    animation = mpy.VideoClip(make_frame_mpl, duration=duration)
    animation.write_videofile(os.path.join(figpath, name + '.mp4'), fps=fps)
e.ipython_display(name)

"""
Explanation: The angle is derived as the arctan of the 2 components
End of explanation
"""


from EdgeGrid import EdgeGrid
e = EdgeGrid()
import numpy as np
np.set_printoptions(precision=2, suppress=True)
import imageio
e.reader = imageio.get_reader(figpath + 'elasticite_test.mp4', loop=True)
for i, im in enumerate(e.reader):
    print(i, e.theta_sobel(im, N_blur=10)*180/np.pi)
    if i>5: break

name = 'trame_sobel'
from EdgeGrid import EdgeGrid
e = EdgeGrid()

import imageio
e.reader = imageio.get_reader(figpath + 'elasticite_test.mp4', loop=True)
def make_lames(e):
    e.im = e.reader.get_next_data()
    return e.theta_sobel(e.im, N_blur=10)

duration = 4.
e.make_anim(name, make_lames, duration=duration)
e.ipython_display(name)

import imagen as ig
line=ig.Line(xdensity=5, ydensity=5, smoothing=0)
import numpy as np
np.set_printoptions(1)
import holoviews
%reload_ext holoviews.ipython

import numbergen as ng
from holoviews import NdLayout
import param
param.Dynamic.time_dependent=True
stim = ig.SineGrating(orientation=np.pi*ng.UniformRandom())
NdLayout(stim.anim(3))

name = 'trame_sobel_grating'
from EdgeGrid import EdgeGrid
e = EdgeGrid()
stim = ig.SineGrating(xdensity=64, ydensity=64)
def make_lames(e):
    stim.orientation=np.pi*e.t/4.
    e.im = stim()
    return e.theta_sobel(e.im, N_blur=5)

duration = 4.
e.make_anim(name, make_lames, duration=duration)
e.ipython_display(name)

%%opts Image.Pattern (cmap='Blues_r')
l1 = ig.Line(orientation=-np.pi/4)
l2 = ig.Line(orientation=+np.pi/4)
cross = l1 | l2
cross.orientation=ng.ScaledTime()*(np.pi/-20)
l1.anim(20) + l2.anim(20) + cross.anim(20)

name = 'trame_sobel_cross'
from EdgeGrid import EdgeGrid
e = EdgeGrid()
l1 = ig.Line(orientation=-np.pi/4)
l2 = ig.Line(orientation=+np.pi/4)
cross = l1 | l2

def make_lames(e):
    cross.orientation = np.pi*e.t/4.
    e.im = cross()
    return e.theta_sobel(e.im, N_blur=1)

duration = 4.
e.make_anim(name, make_lames, duration=duration)
e.ipython_display(name)

line.set_param(xdensity=72,ydensity=72,orientation=np.pi/4, thickness=0.02, smoothing=0.02)
line.x = .25

noise = ig.Composite(xdensity=72, ydensity=72,
                     operator=np.add,
                     generators=[ig.Gaussian(size=0.1,
                                          x=ng.UniformRandom(seed=i+1)-0.5,
                                          y=ng.UniformRandom(seed=i+2)-0.5,
                                          orientation=np.pi*ng.UniformRandom(seed=i+3))
                                for i in range(10)])

stim = line + 0.3*noise
NdLayout(stim.anim(4)).cols(5)

name = 'trame_sobel_line_tmp_4'
from EdgeGrid import EdgeGrid
e = EdgeGrid()

def make_lames(e):
    line.x = -.5 + e.t / 4.
    stim = line + noise
    e.im = stim()
    return e.theta_sobel(e.im, N_blur=1)

duration = 4.
e.make_anim(name, make_lames, duration=duration)
e.ipython_display(name)

"""
Explanation: This function is included in the EdgeGrid class:
End of explanation
"""


