# -*- coding: utf-8 -*-
from objsimpy.int_mod_n import IntModuloN
from random import choice
import pylab
from stim_sequence import StimulusSequence
import numpy as np    



class SequenceGenerator_Gauss:
    def __init__(self, Nx, Ny, NRepetitions, NoJumpDuration=10, XContinuous=1, XGap=None, YGap=None, n_high_num=2):
        self.Nx = Nx
        self.Ny = Ny
        self.n_high_num = n_high_num
        self.n_low_num = 1
        self.NRepetitions = NRepetitions
        self.NoJumpDuration = NoJumpDuration
        self.XContinuous = XContinuous        
        self.X = IntModuloN(0, self.Nx)
        self.Y = IntModuloN(0, self.Ny)
        self.StimHistogram = StimHistogram(self.Nx*self.Ny)
        self.NoJumpHold = self.NoJumpDuration
        self.NStimuli = self.Nx*self.Ny
        self.XStimuli = range(self.Nx)
        self.YStimuli = range(self.Ny)
        self.StimulusSequence=[]
        self.Directions = [-1,1]
        self.Direction = choice(self.Directions)        
        self.Stimuli = range(self.NStimuli)
        self.StimReservoirPattern = range(self.NStimuli)*self.n_high_num
        if not XGap is None:
            for i in range(self.n_high_num)[self.n_low_num:]:
                for y in self.YStimuli:
                    for x in range(XGap[0],XGap[1]): self.StimReservoirPattern.remove(y*self.Nx + x)
        print "YGap is", YGap
        if not YGap is None: 
            for i in range(self.n_high_num)[self.n_low_num:]:
                for stim in range(YGap[0]*self.Nx,YGap[1]*self.Nx): 
                    self.StimReservoirPattern.remove(stim)
        self.StimReservoir = self.StimReservoirPattern[:]
        self.setRandomStimulus()
        
        
    def setRandomStimulus(self):
        CurStim = choice(self.StimReservoir)
        self.CurX = IntModuloN(CurStim % self.Nx, self.Nx)
        self.CurY = IntModuloN(CurStim / self.Nx, self.Ny)
        if (self.XContinuous):
            self.X = choice(self.XStimuli)
        else:
            self.Y = choice(self.YStimuli)

    def generate(self):
        RandomJumpDurRange=range(int(0.5*self.NoJumpDuration))
        #while self.StimHistogram.getMinRepetitions()<self.NRepetitions :
        for i in range(self.NRepetitions):
            #print i
            self.NoJumpHold-=1
            if self.NoJumpHold<1:
                self.NoJumpHold=self.NoJumpDuration+choice(RandomJumpDurRange)
                self.Direction = choice(self.Directions)
                self.setRandomStimulus()
            else:
                if self.XContinuous == 1:
                    self.CurX += self.Direction
                else:
                    self.CurY += self.Direction
                self.appendStimXY(self.CurX, self.CurY)
        print("Sequence generated. Length=", len(self.StimulusSequence))
    
    def appendStimXY(self, X, Y):
        CurStim = Y.int()*self.Nx + X.int()
        self.StimHistogram.addStimulus(CurStim)
        self.StimulusSequence.append(CurStim)     
        try:
            self.StimReservoir.remove(CurStim)   
        except:
            self.StimReservoir += self.StimReservoirPattern
            self.StimReservoir.remove(CurStim)   
    
    
    def __repr__(self):
        return str(self.StimulusSequence)
    
    def plot(self, x0=0, x1=300):
        XSeq = map(lambda x: x%self.Nx, self.StimulusSequence)
        YSeq = map(lambda x: x/self.Nx, self.StimulusSequence)
        x = range(x0,x1)
        pylab.subplot(211) #first subplot
        pylab.ylabel('y parameter')
        pylab.plot(x, YSeq[x0:x1], '.')
        pylab.ylabel("Y Parameter")
        
        pylab.subplot(212) #second subplot 
        pylab.ylabel('x parameter')
        pylab.plot(x, XSeq[x0:x1], '+')
        pylab.grid(True)
        pylab.ylabel("X Parameter")
        
        pylab.figure(3)
        n, bins, patches = pylab.hist(XSeq, self.Nx, normed=1, facecolor='g', alpha=0.75)
        print("n=", n, " bins=", bins, " patches=", patches)
        pylab.xlabel('x parameter')
        pylab.ylabel('Probability')
        pylab.title('Histogram of X stimuli')
        pylab.axis([0, 19, 0, 0.15])
        pylab.grid(True)

        pylab.figure(4)
        n, bins, patches = pylab.hist(YSeq, self.Ny, normed=1, facecolor='r', alpha=0.75)
        print("n=", n, " bins=", bins, " patches=", patches)
        pylab.xlabel('y parameter')
        pylab.ylabel('Probability')
        pylab.title('Histogram of Y stimuli')
        pylab.axis([0, 19, 0, 0.15])
        pylab.grid(True)
        
        pylab.figure(5)
        n, bins, patches = pylab.hist(self.StimReservoirPattern, 100, normed=1, facecolor='r', alpha=0.75)
        print("n=", n, " bins=", bins, " patches=", patches)
        
        
        pylab.show()
        


class StimHistogram:
    def __init__(self, NStimuli):
        self.NStimuli = NStimuli
        self.Histogram = [0]*self.NStimuli
        self.MinRepetitions = 0
        self.MinRepIndex=0
    def addStimulus(self, StimIndex):
        self.Histogram[StimIndex] += 1
        if StimIndex==self.MinRepIndex:
            self.findMinRepIndex()
    def findMinRepIndex(self):
        for Index in range(self.MinRepIndex+1, self.NStimuli):
            if self.Histogram[Index]<=self.MinRepetitions:
                assert self.Histogram[Index]==self.MinRepetitions, "kein Histogram-Wert darf kleiner sein als self.MinRepetitions"                    
                self.MinRepIndex = Index
                break
        else:
            self.MinRepetitions += 1
            for Index in range(0, self.NStimuli):
                if self.Histogram[Index]<=self.MinRepetitions:
                    self.MinRepIndex = Index
                    break
            else:
                raise "inconsistent histogram data"
    def getMinRepetitions(self):
        return self.MinRepetitions

def generateMovieSequence_Gauss(Nx, Ny, NRepetitions, NoJumpDuration=10, XContinuous=1, XGap=None, YGap=None, n_high_num=2):
    gen = SequenceGenerator_Gauss(Nx, Ny, NRepetitions, NoJumpDuration, XContinuous, XGap=XGap, YGap=YGap, n_high_num=n_high_num)
    gen.generate()
    Seq = StimulusSequence(gen.StimulusSequence)
    SeqFileName = "test_gauss20x20"
    SeqFileName += "_nojump"+str(NoJumpDuration)
    if (XContinuous==1):
        SeqFileName += "_xcont"
    else:
        SeqFileName += "_ycont"
    if (XGap):
        print "XGap=", XGap
        SeqFileName += "_xgap"
    if (YGap):
        print "YGap=", YGap
        SeqFileName += "_ygap"
    SeqFileName += "_high" + str(n_high_num)
    SeqFileName += ".sequence"
    DirName="/home/frank/prog/objsim/data/movies/"
    SeqFileName = DirName + SeqFileName
    print("SeqFileName=", SeqFileName)
    Seq.tofile(SeqFileName)
    gen.plot()

            
                

def generateMovieSequence_Diagonal(Nx, Ny, NRepetitions):
    max_diagonal_index = min(Nx, Ny)
    sequence = []
    for val in range(max_diagonal_index):
        sequence.append(val*Nx + val)
    DirName = "/home/frank/prog/objsim/data/movies/"
    FileName = "diagonal.sequence"
    path = DirName + FileName
    seq = StimulusSequence(sequence)
    seq.tofile(path)
    seq.show()


class RectStateMachine:
    def __init__(self, n_x_elements, n_y_elements, n_x_groups, n_y_groups):
        if n_x_elements % n_x_groups != 0:
            raise ValueError
        if n_y_elements % n_y_groups != 0:
            raise ValueError
        elements = range(n_x_elements * n_y_elements)
        self.groups = []        
        self.n_x_elements = n_x_elements
        self.n_y_elements = n_y_elements            
        for x_group_nr in range(n_x_groups):
            for y_group_nr in range(n_y_groups):
                group_start_x = x_group_nr * n_x_elements / n_x_groups
                group_stop_x = (x_group_nr + 1) * n_x_elements / n_x_groups
                group_start_y = y_group_nr * n_y_elements / n_y_groups
                group_stop_y = (y_group_nr + 1) * n_y_elements / n_y_groups
                #print "(",group_start_x, ",",group_stop_x, ",",group_start_y, ",",group_stop_y,")"
                group = []
                for x_element in range(group_start_x, group_stop_x):
                    #print "x=",x_element
                    for y_element in range(group_start_y, group_stop_y):
                        #print "y=",y_element
                        element_nr = x_element + y_element * n_x_elements
                        #print element_nr
                        group.append(element_nr)
                        elements.remove(element_nr)
                self.groups.append(group)
        if elements:
            raise RuntimeError("elements should be empty now")  
        self.n_groups = len(self.groups)
        self.group_ids = range(self.n_groups)
        self.cur_group = choice(self.group_ids)
        self.group_ids.remove(self.cur_group)
        self.cur_elements = list(self.groups[self.cur_group])
        self.stay_in_group_period = 2 * len(self.cur_elements)
        self.stay_in_group = self.stay_in_group_period-1
       
    def set_stay_in_group_factor(self, f):
        self.stay_in_group_period = f * len(self.cur_elements)
        
    def next_group(self):
        if not self.group_ids:
            self.group_ids = range(self.n_groups)            
        self.cur_group = choice(self.group_ids)        
        self.group_ids.remove(self.cur_group)    
        return self.cur_group    
    
    def next_element(self):
        if not self.cur_elements:
            self.cur_elements = list(self.groups[self.cur_group])
        self.cur_element = choice(self.cur_elements)
        self.cur_elements.remove(self.cur_element)
        return self.cur_element
    
    def next(self):
        if not self.stay_in_group:
            self.next_group()
            self.stay_in_group = self.stay_in_group_period
        self.stay_in_group -= 1        
        return self.next_element()

    def show(self):
        fig = pylab.figure()
        for group in self.groups:
            g = np.array(group)
            y = g / self.n_x_elements
            x = g % self.n_x_elements
            pylab.plot(x,y, '-o')
        pylab.show() 

    
def generate_movie_sequence_rectangles(n_x_elements=20, n_y_elements=20, n_x_groups=4, n_y_groups=4, n_repetitions=5):
    """ generates a stimulus sequence with groups of rectangles
        
        Stimulus space is a rectangle of numbers.
        This rectangle is devided into subrectangles.     
    """
    rect_state = RectStateMachine(n_x_elements, n_y_elements, n_x_groups, n_y_groups)
    f = 3
    rect_state.set_stay_in_group_factor(f)
    sequence = []
    n = 20 * n_x_elements * n_y_elements * n_repetitions * f
    for i in range(n):
        sequence.append(rect_state.next())
        
        
    DirName = "/home/frank/prog/objsim/data/movies/"
    FileName = "rectangle.sequence"
    path = DirName + FileName
    seq = StimulusSequence(sequence)
    seq.tofile(path)
    seq.show()
    return sequence
   
        
        
        
if __name__=='__main__':
    generateMovieSequence_Gauss(20,20,50000, 20, 1, n_high_num=2)
    #generateMovieSequence_Gauss(20,20,50000, 20, 1, YGap=(4,16), n_high_num=2)
    #generateMovieSequence_Diagonal(20,20,1)
    #sequence = generate_movie_sequence_rectangles(20,20,4,4,1)