'''This code accompanies the paper "An approach to root multiplicities for symmetric
Kac-Moody algebras through quiver varieties" by Peter Tingley. The code was written my
Colin Williams and Peter Tingley. It calculates estimates for imaginary root space
multiplicities for symmetric rank two hyperbolic Kac-Moody algebras using the methods
described in that paper. The function "Root_Multiplicity_Estimates"
calculates the estimates exactly by searching through all Dyck paths. The function
"Monte_Carlo_Estimate_of_multiplicity" estimates the estimates using Monte-Carlo sampling.
It will run on much larger roots than "Root_Multiplicity_Estimates" can handle.
All the other functions here are to support these. This was written using python 2.7'''

import time
import random
import numpy
import math



def DyckPaths(zeros,ones,slope):
    '''
    DyckPaths(2,3,1.5) -> [[1, 1, 0, 1, 0], [1, 1, 1, 0, 0]]

    Takes a number of zeros and a number of ones, and a slope, and returns the list of all words
    such that, for any prefix, the number of ones divided by the number of zeros in the prefix
    is at most the slope. To get rational Dyck paths the spole should be equal to ones/zeros, but
    other values of slope are needed in the recursive use of the function.

    The function works by recursively building up the Dyck path by adding chuncks of the form
    100...0. If there are no zeros, or the number of zeros is such that 100...0 itself is a
    valid prefix for a Dyck path of slope "slope", then adding that to any valid prefix still
    gives a valid prefix. In the case where the chunk being added is not a valid prefix,
    it can result in a non-valid path, so we must apply the isDyckPath function to check.
    '''
    answer=[]
    if zeros >=1:
        if ones >=1:
            for i in range(zeros+1):
                partialperms=DyckPaths(zeros-i,ones-1,slope)
                if i==0:
                    for partial in partialperms:
                        partial.append(1)
                        for k in range(i):
                            partial.append(0)
                        answer.append(partial)
                elif slope<=float(1)/i:
                    for partial in partialperms:
                        partial.append(1)
                        for k in range(i):
                            partial.append(0)
                        answer.append(partial)
                else:
                    for partial in partialperms:
                        partial.append(1)
                        for k in range(i):
                            partial.append(0)
                        if isDyckPath(partial,slope):
                            answer.append(partial)
            return(answer)
        else:
            return([])
    else:
        someones=[]
        for k in range(ones):
            someones.append(1)
        return([someones])


def isDyckPath(path,slope):
    '''
    isDyckPath('0010101', 7, .75) -> True
    isDyckPath('1010010', 7, .75) -> False

    Function used by test_paths() function to determine if  a path is a rational Dyck path.
    path is a list of 0's and 1's, and slope is the number of 1's divided by the number of 0's
    slope can of course be calculated from the path, but it is convenient to pass it in.
    '''

    cur_zeros = 0
    cur_ones = 0

    '''Checks each edge to see whether the step violates the Dyck Path condition'''
    for steps in range(0,len(path)):

        '''updates numer of zeros and ones to the next step in the path'''
        if path[steps] == 0:
            cur_zeros += 1
        elif path[steps] ==1:
            cur_ones += 1
        '''checks the rational Dyck path condition and returns False if it fails'''
        if slope *cur_zeros - cur_ones > 0:
            return False

    '''returns True if all checks passed, so it is a Dyck path'''
    return True

def KashiwaraData(path):
    '''
    KashiwaraData([1,1,0,1,0,0,0,1,1,0,0,0]) -> [2,1,1,3,2,3]

    Returns path as list of numbers of consecutive similar edges. This is useful
    because the conditions describing when a Dyck path fails to correspond to valid
    Kashiwara data are best described in terms of the Kashiwara data itself.
    '''


    pathKD = [0]
    current_type=1
    index=0

    for steps in range(0,len(path)):
        if path[steps] == current_type:
            pathKD[index]+=1
        else:
            current_type= path[steps]
            pathKD.append(1)
            index+=1

    return pathKD

def cond1violated(KD,n):
    '''
    cond1violated([2,1,1,1,1,1],3) -> False
    cond1violated([1,1,3,2],3) -> True

    KD is the Kashiwara Data for a path and n is the parameter in the Cartan matrix

    Returns True if the KD violates condition 1, as described in the paper (that is, if
    an entry in the Kashiwara data is too much larger than the previous entry)
    and False Otherwise.

    Used by test_paths()
    '''

    for steps in range(len(KD)-1):
        if KD[steps+1] > KD[steps]:
            if 2*(KD[steps]**2) + 2*(KD[steps+1]**2) - (2*n*KD[steps]*KD[steps+1]) >=0:
                return True

    return False

def cond2violated(KD,slope,n):
    '''
    cond2violated([1,1,2,1,1,1], .75, 3) -> True
    cond2violated([4,3], .75, 3) -> False

    Takes Kashiwara Data for a path, the slope, and the parameter n of the Cartan matrix.
    Returns True if the path violates condition 2, as described in the paper.
    Used by test_paths() to test a path that has passed condition 1 for
    violation of condition 2. slope is redundant data, but passing it in is useful.

    Comparing this code to the statement of Theorem 4.14, the interval (x,y) corresponds
    to left_end_of_interval and right_end_of_interval here by
    left_end_of_interval=2x-2, right_end_of_interval=2y
    '''

    '''The two loops below run over every possible interval [d_k,...d_j] in the Kashiwara
    where both j and k are even, so correspond to strings of 1s in the path.
    zeros_so_far and ones_so_far count zeros and ones before left_end_of_interval,
    which is initially 0'''
    zeros_so_far=0
    ones_so_far=0
    for left_end_of_interval in range(0,len(KD)-2,2):
        for right_end_of_interval in range(left_end_of_interval+2,len(KD),2):
            '''interval zeros and interval ones are zeros and ones corresponding to the
            Kashiwara data d_(left_end _of_interval+1) + d_(right_end _of_interval)'''
            interval_zeros = 0
            interval_ones = 0
            for steps in range(left_end_of_interval,right_end_of_interval,2):
                interval_ones += KD[steps+2]
                interval_zeros += KD[steps+1]

            '''Image ones is the zones coming from Kashiwara data in the interval which
            are in the sub-module generted by the basis vectors corresponding to all the
            0's in the interval. You will need to read the paper to understand why...'''
            image_ones = n*interval_zeros - interval_ones

            '''This is checking is the submodule generated by Kashiwara data before
            left_end_of_interval and the 1's in the interval is unstable. If it is
            unstable, this is a violation of condition 2, so the function returns True.'''
            if slope *(zeros_so_far+interval_zeros) - ones_so_far-image_ones > 0:
                return True
                break
        '''place is about to change, so update the number of zeros and ones before
        place'''
        ones_so_far+=KD[left_end_of_interval]
        zeros_so_far+=KD[left_end_of_interval+1]
    '''returns False if all passed, so the Kashiwara data doesn't violate condition 2'''
    return False

def Sample_Dyck_Paths(zeros,ones):
    '''Samples rational Dyck paths by sampling all paths and then performing
    the unique cyclic shift that gives a rational Dyck path. It only works if "zeros"
    and "ones" are relatively prime. Otherwise, it will under-sample paths that return
    to the diagonal. This could be fixed by returning all possibly cyclic shifts that
    give Dyck paths, not just the first.'''


    slope =float(ones)/float(zeros)

    '''randomly creates a list with a 0s and b 1s'''
    Ones_Places=random.sample(range(0,zeros+ones), ones)
    path=[]
    for k in range (0,zeros+ones):
        path.append(0)
    for k in range (0,ones):
        path[Ones_Places[k]]=1

    '''finds the corner of the path that is the farthest above the diagonal'''
    max_dist = 0
    cur_dist = 0
    shift_amount = 0
    cur_zeros = 0
    cur_ones = 0
    for steps in range(zeros+ones):
        if path[steps] ==1:
            cur_ones += 1
        elif path[steps] == 0:
            cur_zeros += 1
            if ((slope *cur_zeros) - cur_ones) > 0:
                cur_dist = slope * cur_zeros - cur_ones
                if cur_dist > max_dist:
                    max_dist = cur_dist
                    shift_amount = steps+1

    '''Performs unique cyclic shift to transform the path into a Dyck path'''
    if max_dist != 0:
        path = path[shift_amount:] + path[:shift_amount]
    return path

def Root_Multiplicity_Estimates(zeros,ones,n):
    '''This calculates our upper bounds on the multiplicity of
    zeros \alpha_0+ones \alpha_1 in the Kac-Moody algebra with Cartan Matrix
    ((2,-n),(-n,2)). It returns both the estimate using only Theorem 4.5 and the estimate using
    Theorem 4.14 as well.


    Root_Multiplicity_Estimates(10,9,3) ->  ('Number of rational Dyck paths', 4862)
                                            ('Estimate Using Theorem 4.4', 1278)
                                            ('Estimate Using Both Theorems', 1267)
                                            ('Difference in Estimates', 11)
    '''

    slope = float(ones)/zeros

    '''numpaths will count the number of rational Dyck paths. cond1_violators will count
    the number of rational Dyck paths which violate condition 1. cond2_violators will
    count the number of rational Dyck path that violate condition 2 but not condition 1.
    Passed will count rational Dyck paths that satisfy both conditions. All are
    initialized to 0'''
    num_paths = 0
    cond1_violators = 0
    cond2_violators = 0
    passed = 0

    '''loop is over all paths with the right number of zeros and ones'''
    for path in DyckPaths(zeros, ones,slope):
        num_paths += 1
        '''created tha Kashiwara data of the path'''
        pathKD = KashiwaraData(path)
        if cond1violated(pathKD,n):
            cond1_violators += 1
        else:
            if cond2violated(pathKD,slope,n):
               cond2_violators += 1
            else:
               passed += 1

    print('Number of rational Dyck paths', num_paths)
    print('Estimate Using Theorem 4.4', num_paths - cond1_violators)
    print('Estimate Using Both Theorems',passed)
    print('Difference in Estimates', cond2_violators)




def Monte_Carlo_Estimate_of_multiplicity(zeros,ones,n,runs):
    '''
    This function estimates our upper bounds on the multiplicity of
    zeros \alpha_0+ones \alpha_1 in the Kac-Moody algebra with Cartan Matrix
    ((2,-n),(-n,2)). It uses Monte-Carlo sampling of `runs' Dyck paths to estimate the
    proportion that satisfy our conditions 1 and 2, and then uses this to output estimates
    for the two upper bounds on the root multiplicity. It also returns the number of paths
    that passes one or both conditions.

    note: this algorithm is randomized, so output will vary.
    sample_paths(8,9,3,10000) ->    ('Paths Passing Theorem 4.4', 4509)
                                    ('Estimate Using Only Theorem 4.4', 644.787)
                                    ('Paths Passing Both Theorems', 3208)
                                    ('Estimate Using Both Theorems', 458.744)
                                    ('Run time', 0.20244000000000006)
    '''

    start = time.clock()

    slope = float(ones)/float(zeros)
    cond1_violators = 0
    cond2_violators = 0
    passed = 0
    number_of_dyck_paths = math.factorial(zeros+ones)/(math.factorial(zeros)*math.factorial(ones)*(zeros+ones))

    for r in range(runs):



        '''samples from Dyck paths'''
        path=Sample_Dyck_Paths(zeros,ones)

        '''Creates Kashiwara Data for the Path'''
        pathKD = KashiwaraData(path)

        '''checks the conditions'''
        if cond1violated(pathKD,n):
            cond1_violators +=1
        else:
            if cond2violated(pathKD,slope,n):
                cond2_violators += 1
            else:
                passed += 1
                print (r,  passed, float(number_of_dyck_paths*(r-cond1_violators))/r,float(number_of_dyck_paths*passed)/r)

    '''Calculates estimates of passing paths, or equivalents of the root multiplicity)'''
    cond1_est = float(number_of_dyck_paths*(runs-cond1_violators))/runs
    cond2_est = float(number_of_dyck_paths*passed)/runs

    print('Paths Passing Theorem 4.4', runs - cond1_violators)
    print('Estimate Using Only Theorem 4.4', cond1_est)
    print('Paths Passing Both Theorems', passed)
    print('Estimate Using Both Theorems',cond2_est)



    print('Run time', time.clock() - start)


