import numpy as np
import itertools
from copy import copy
import datetime
from timeit import default_timer as timer
from sympy.ntheory import primefactors
from sympy.ntheory.residue_ntheory import primitive_root

""" forbidden CYCLE directed Ramsey algs """

def psieve():
    for n in [2, 3, 5, 7]:
        yield n
    D = {}
    ps = psieve()
    next(ps)
    p = next(ps)
    assert p == 3
    psq = p*p
    for i in itertools.count(9, 2):
        if i in D:      # composite
            step = D.pop(i)
        elif i < psq:   # prime
            yield i
            continue
        else:           # composite, = p*p
            assert i == psq
            step = 2*p
            p = next(ps)
            psq = p*p
        i += step
        while i in D:
            i += step
        D[i] = step
        


        
def check_p_m_v6(p,m,g):
    ''' checks a prime p with primitive root g for m colors '''
    ''' fastest version so far '''
    n = 2 * m
    X0 = np.array( [pow(g, i, p) for i in xrange(0, p-n, n) ] )
    certificates_master = np.array([ pow(g, i, p) for i in xrange(n) ])
    C_minus_X0 = ( ( certificates_master[:, np.newaxis] - X0 ) % p )
    C_minus_X0_sets = [ set(L) for L in C_minus_X0 ]
    for i in xrange(n):
        Xi = {pow(g, x+i, p) for x in xrange(0, p-n, n)}
        for j in xrange(i+m,n+m):
            if bool(Xi.intersection(C_minus_X0_sets[(j%n)])) == bool((j%n) == m):
                return False
    return True

       
def main(mikelist):
    

    lget = primitive_root     ### GIVE FUNCTIONS LOCAL NAMES ###
    lcheck = check_p_m_v6     ##################################
    
   
    with open("forbcycleRamsey.txt", 'a') as file:
        for mike in mikelist:
            n = 2 * mike
            primes = psieve()
            prime = next(primes)
               
            for prime in primes:
                if (prime-1) % n == 0 and (prime-1)/n % 2 == 1:
                    # print prime
                    gen = lget(prime)
                    # print prime, gen
                    p_out = lcheck(prime, mike, gen)
                    if p_out == True:
                        print mike, prime, gen
                        file.write(str(mike) + ', ' + str(prime) + ', ' + str(gen) + '\n')
                        break
                if prime > n**4+5:
                    print mike, 'impossible, up though', prime, '\n'
                    break
      

main(range(1,100))