'''
Created on Jul 23, 2020

@author: vladislavkargin
'''

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import cm

import rgs.mndrpy.Pairing as pr
#import rgs.mndrpy.Utility as ut
import rgs.mndrpy.DyckPaths as dp

class Meander(object):
    '''
    classdocs
    '''


    def __init__(self, uPairing, dPairing, uprngArray = None,
                    dprngArray = None):
        '''
        Constructor
        Makes the meander from upper and lower pairings
        arguments:
        uPairing, dPairing -- upper and lower parings (as objects or lists)
        
        The other pair of arguments is obsolete (poor design),
        kept only for compatibility with some other functions. 
        uprgnArray, dprngArray -- upper and lower pairings (as lists)
        '''
        if type(uPairing) is list:
            self.uPairing = pr.Pairing(prngArray = uPairing)
        elif uPairing == None:
            self.uPairing = pr.Pairing(prngArray = uprngArray)
        else: 
            self.uPairing = uPairing
            
        if type(dPairing) is list:
            self.dPairing = pr.Pairing(prngArray = dPairing)
        elif uPairing == None:
            self.dPairing = pr.Pairing(prngArray = dprngArray)
        else: 
            self.dPairing = dPairing
        
        
    def __str__(self):
        return str(self.uPairing) + ", " + str(self.dPairing)
    
    def __eq__(self, other):
        '''Overrides the default implementation'''
        if isinstance(other, Meander):
            if self.uPairing == other.uPairing and self.dPairing == other.dPairing: 
                return True
            else: 
                return False
        return False
    
    def __hash__(self):
        return hash(self.uPairing) + 17 * hash(self.dPairing)
    
    def size(self):
        ''' returns n for the meander system on 2n bridges '''
        return self.uPairing.size()
     
    def reflected(self):
        '''returns a meander with lower and upper pairings exchanged
        '''     
        new_mndr = Meander(uPairing = self.dPairing, dPairing = self.uPairing)
        return new_mndr
    
    def draw(self, ax = None, drawCycles = False, palette = "jet", 
             asPolygon = False):
        if ax == None:
            fig = plt.figure()
            ax = fig.add_subplot(1, 1, 1)
        fig = ax.figure
        self.uPairing.draw(ax = ax, asPath = asPolygon)
        self.dPairing.draw(ax = ax, up = False, asPath = asPolygon)
        if asPolygon:
            ax.grid(which='both')
        if drawCycles and not asPolygon:
            colormap = cm.get_cmap(palette, 128)
            cycles, _ = self.findCycles()
            for i in range(len(cycles)):
                col = colormap(i/len(cycles))
                #print(i, col)
                self.drawCycle(cycles[i], ax, color = col) 
        return fig, ax
    
    def drawCycle(self, cycle, ax = None, color = "red"):
        if ax == None:
            fig = plt.figure()
            ax = fig.add_subplot(1, 1, 1)
            width = len(self.uPairing.prng)
            ax.set_xlim(0, width)
            ax.set_ylim(-width/2, width/2)
        for c, i in enumerate(cycle):
            j = cycle[(c + 1) % len(cycle)]
            xy = (i + (j - i)/2, 0)
            w = np.abs((j - i)/2)
            if c % 2 == 0:
                theta1 = 0.0
                theta2 = 180.0
            else:
                theta1 = 180.0
                theta2 = 360.0
            arc = mpatches.Arc(xy, 2*w, 2*w, 0, theta1, theta2, color = color)
            ax.add_patch(arc)
    
    def isProper(self):
        '''checks if this system has only one cycle'''  
        N = len(self.uPairing.prng)
        #print(N)
        next_v = self.uPairing.prng[0]
        up = False 
        counter = 0    
        while next_v != 0 and counter <= 2*N:
            counter = counter + 1
            if up == True:
                next_v = self.uPairing.prng[next_v]
                up = False
            else:
                next_v = self.dPairing.prng[next_v]
                up = True
        if counter == N - 1:
            return True
        else:
            return False
        
    def  isIrreducible(self):
        '''checks if the meander is irreducible. (NB. The definition is different
        from irreducible meanders in the sense of Zvonkin-Lando, so it is
        more of exploratory nature.)'''
        if self.isProper():
            return True   
        cycles, _ = self.findCycles()
        for cycle in cycles:
            if len(cycle) == max(cycle) - min(cycle) + 1:
                return False
        return True
    
        
    def findCycles(self):
        ''' finds all cycles in the meander and stores them in two structures
        cycles - which is an array of arrays that represent cycles
        cyclesInd - which is an 2n-array with each element marked by 
        the number of the cycle that it belongs to.
        '''
        N = len(self.uPairing.prng)
        c = 0
        up = True #shows which pairing we are going to use in tracking 
        cycles = []
        cyclesInd = [-1] * N
        
        for i in range(N):
            if cyclesInd[i] == -1:
                c = c + 1
                cyclesInd[i] = c
                cycle = [i]
                if up:
                    next_v =  self.uPairing.prng[i]
                    up = False
                else:
                    next_v =  self.dPairing.prng[i]
                    up = True
                #the following cycle can go in infinite loop if one of the
                #pairing is invalid for this reason we break it 
                # if it go for more than N iterations.
                # The result is invalid and likely to break the following programs
                counter = 0
                while (next_v != i and counter <= 2 * N):
                    cyclesInd[next_v] = c
                    cycle.append(next_v)
                    if up:
                        next_v =  self.uPairing.prng[next_v]
                        up = False
                    else:
                        next_v =  self.dPairing.prng[next_v]
                        up = True
                    counter = counter + 1
                cycles.append(cycle)
                if counter > 2 * N:
                    print( "findCycles Warning:")
                    print(" An infinite loop in find components was broken")
                    print(" A likely problem with validity of one of the pairings")
        return cycles, cyclesInd
     
    def calcCycleLengths(self):
        ''' caclulate the lengths of all cycles in the meander'''
        (cycles, _) = self.findCycles()
        lengths = []
        for cycle in cycles:
            lengths.append(len(cycle))
        return lengths
        
def areaMeander(mndr):
    '''calculates the area of a meander system defined as the area 
    of the polygone that has the Dyck path for the upper pairing
    as its upper boundary and the reflected Dyck path for the lower
    pairing as its lower boundary. The size of the meander is subtracted
    for normalization, so the area of the standard meander system is 0'''
    area = dp.area(pr.prng2path(mndr.uPairing))
    area += dp.area(pr.prng2path(mndr.dPairing))
    #area += mndr.size()
    return area
         
def randomMeander(n, w = None, seed = None):
    ''' generates a uniformly random meander on half-length n.
    If 2-tuple w != None, generates a meander with pairings distributed 
    according to weights w1 and w2'''
    if seed != None:
        np.random.seed(seed)
    if w == None:
        uP = pr.randomPairing(n)
        dP = pr.randomPairing(n)
    else:
        uP = pr.randomPairing(n, w[0])
        dP = pr.randomPairing(n, w[1])
    mndr = Meander(uP, dP)
    return mndr

def randomCrossSemiMeander(n, seed = None):
    '''generate a random crossing semi-meander'''
    if seed != None:
        np.random.seed(seed)
    uP = pr.randomCrossPairing(n)
    dP = pr.rainbowPairing(n)
    mndr = Meander(uP, dP)
    return mndr
    

def randomComb(n, seed = None):
    ''' generates a random comb meander of half-length n'''
    if seed != None:
        np.random.seed(seed)
    uP = pr.randomPairing(n)
    dP = pr.standardPairing(n)
    mndr = Meander(uP, dP)
    return mndr

def randomSemiMeander(n, seed = None):
    ''' generates a random meander of half-length n that have a random pairing
    on top and the rainbow pairing in the bottom'''
    
    if seed != None:
        np.random.seed(seed)
    uP = pr.randomPairing(n)
    dP = pr.rainbowPairing(n)
    mndr = Meander(uP, dP)
    return mndr

def rainbowMeander(uWidths, dWidths):
    '''
    generates a rainbow meander.
    Parameters:
    uWidths, dWidths -- sizes of the meanders on the top and on the bottom.
    returns a meander
    '''  
    N = sum(uWidths)
    M = sum(dWidths)
    if M != N:
        print("Rainbow meander says: the total sizes of upper and lower meanders")
        print("should be the same. N = ", N, " and M = ", M)
        return 
    uP = pr.rainbows(uWidths)
    dP = pr.rainbows(dWidths)
    return Meander(uP, dP)

def smallRainbowsMeander(amounts, sizes):
    '''
    creates a random meander whose pairings are collections of small rainbows
    '''
    uP = pr.smallRainbows(amounts, sizes)
    dP = pr.smallRainbows(amounts, sizes)
    return Meander(uP, dP)

def drawAsPolygon(mndr):
    ''' draws meander as polygon. The upper pairing is represented
    by a Dyck path above the diagonal y = x, and the lower pairing
    is represented by a Dyck path below the main diagonal. 
    '''
    upath = pr.prng2path(mndr.uPairing)
    ax = dp.plotDyckPath(upath)
    dpath = pr.prng2path(mndr.dPairing)
    ax = dp.plotDyckPath(dpath, ax = ax, method = "lowerLatticePath")
    return ax    
    
'''
For testing methods
'''
def main():
    
    #n = 100
    #seed = None    
    #mndr = randomMeander(n, seed = seed)
    #mndr = randomComb(n, seed = seed)
    #mndr = randomSemiMeander(n, seed = seed)
    #mndr = randomMeander(n, w = ([1,1,1],[1,1,1]), seed = seed)
    #mndr = randomCrossSemiMeander(n)
    amounts = [20, 20]
    sizes = [1, 5]
    mndr = smallRainbowsMeander(amounts, sizes)
    
    print(str(mndr))
    _, ax = mndr.draw(drawCycles = True)
    
    (cycles, cyclesInd) = mndr.findCycles()
    print("Cycles = " + str(cycles))
    print("CyclesInd = " + str(cyclesInd))
    
    lengths = mndr.calcCycleLengths()
    print("Cycle lengths = " + str(lengths))
    
    maxL = max(lengths)
    i = lengths.index(maxL)
    
    print("The length of the largest cycle is " + str(maxL))
    
    mndr.drawCycle(cycles[i], ax = ax)
    
    print("Hash = ", hash(mndr))
    print("Meander? ", mndr.isProper())
    print("Irreducible? ", mndr.isIrreducible())
    
    
    
    ''' Picture for the paper about meanders'''
    
    uWidths = [15, 5, 11]
    dWidths = [18, 13]
    mndr = rainbowMeander(uWidths, dWidths)
    _, ax = mndr.draw(drawCycles = True)  
    (cycles, cyclesInd) = mndr.findCycles()
    print("Cycles = " + str(cycles))
    print("CyclesInd = " + str(cyclesInd))
    lengths = mndr.calcCycleLengths()
    print("Cycle lengths = " + str(lengths))
    maxL = max(lengths)
    i = lengths.index(maxL)
    print("The length of the largest cycle is " + str(maxL))
    mndr.drawCycle(cycles[i], ax = ax)
    
    drawAsPolygon(mndr)
    drawAsPolygon(mndr.reflected())
    
    plt.show()


if __name__ == '__main__':
    main()          
        
        