import math
import operator
import datetime
from collections import defaultdict

# The goal of this code is to take in a threshold N and a set of avoidance M, and to produce,
# for each 1 <= n <= N, the largest symmetric two-dimensional generalized arithmetic
# progession (GAP) in [-n,n] containing no nonzero elements of M.

# Example: If you wanted M to be the squares and to do N=100000, you would simply do:

#   M=Squares(100000)
#   AllMaxGAP(100000,M,1)

# and just watch it go! See the description of AllMaxGAP for more info on the output

# sumset(A,B) is quite self-explantory, and is used to combine two symmetric arithmetic progressions
# into a single symmetric two-dimensional GAP

def sumset(A, B):
    sumset = [(a+b) for a in A for b in B]
    return set(sumset)

# Squares(N) builds the set of squares up to N. Squares were the first case we explored, but
# as mentioned, the code can accomodate any desired set of avoidance.

def Squares(N):
        M = set([a**2 for a in range(1,math.floor(math.sqrt(N))+1)])
        return M
    
# PM1(N) builds the set of numbers of the form p-1, up to N, where p is prime.

def PM1(N):
    numbers = set(range(N, 1, -1))
    M = []
    while numbers:
        p = numbers.pop()
        M.append(p-1)
        numbers.difference_update(set(range(p*2, N+1, p)))
    return M

# G keeps track of the "record" for the largest M-free GAP in [-N,N] for each N

G=defaultdict(int)

# OneDimMax takes as input the interval threshold N, a step size d, and a set of
# avoidance M and finds the length of the longest symmetric AP in [-N,N] with
# step size d which contains no element of M.

def OneDimMax(N,d,M):
    m=0
    for x in range(1,d+1):
        if x*d in M: break
        if -x*d in M: break
        if x*d > N: break
        m=m+1
    return m

# symAP simply takes in a step size d and a length L and builds the symmetric AP
# {xd : |d|\leq L}

def symAP(d,L):
    A=set([d*x for x in range (-L,L+1)])
    return A

# Opt2Dim takes in the interval threshold N, the set of avoidance M,
# the two step sizes d1 and d2, and the
# length in the d1 direction L1, and it determines how long the lenth
# in the d2 direction can possibly be while remaining contained in
# [-N,N] and remaining M free

def Opt2Dim(N,d1,L1,d2,M):
    A=symAP(d1,L1)
    j=0
    for y in range(1,math.floor(((N-d1*L1)/d2))+1):
        if sumset(A,set([-y*d2,y*d2])).intersection(M): break
        j=j+1
       
    return j    
        
H=defaultdict(set)
d1F=defaultdict(int)
d2F=defaultdict(int)
L1F=defaultdict(int)
L2F=defaultdict(int)

# MaxGAP is the main algorithm. It takes in the threshold N and the
# set of avoidance M and determines the largest M-free GAP in [-N,N].
# It runs through all choices of larger step size d1,
# all possible lengths L1 in the d1 direction (bounded by OneDimMax), then
# once d1 and L1 are fixed, it runs through all admissible smaller step
# sizes d2, and uses Opt2Dim to determine the maximal length L2 in the
# d2 direction. Along the way it keeps track of the record holding GAP,
# not just its size but its step sizes and lengths in each direction.
# Finally, it prints all the details of the record holder, and stores
# the maximal GAP in the dictionary H and the size in the dictionary G.

def MaxGAP(N,M):
    
    start_time = datetime.datetime.now()
    
    Gmax=0
    for d1 in range(2,N):
        #print(d1)
        for L1 in range(1,OneDimMax(N,d1,M)+1):
            s=min(d1,math.floor(N-d1*L1))
            for d2 in range(1, s+1):
                j=Opt2Dim(N,d1,L1,d2,M)
                A=sumset(symAP(d1,L1),symAP(d2,j))
                t=len(A)
                if t>Gmax:

                    Gmax=t
                    H[N]=A
                    d1F[N]=d1
                    d2F[N]=d2
                    L1F[N]=L1
                    L2F[N]=j
    G[N]=math.floor((Gmax-1)/2)
    print(N)
    print("Size")
    print(G[N])
    print("d1=")
    print(d1F[N])
    print("L1=")
    print(L1F[N])
    print("d2=")
    print(d2F[N])
    print("L2=")
    print(L2F[N])

    
    
    end_time = datetime.datetime.now()
    print(end_time-start_time)
    return

# AllMaxGAP takes in the interval threshold N, the set of avoidance M, and a lower bound K,
# and it finds the maximal M-free symmetric two-dimensional in [-n,n] for ALL K <= n <= N,
# printing the sizes, step sizes d1 and d2, and lengths in each direction L1, L2,
# of all the record holders along the way.

# It actually starts at the "top" and works downward, exploiting the fact that, for example,
# if you found that the largest M-free GAP in [-20000,20000] is actually contained in
# [-17000,17000], then its actually the record holder for all 17000 \leq N \leq 20000, so you
# might as well then turn your attention to N=16999. That's exactly what this algorithm does.

def AllMaxGAP(N,M,K):
    start_time = datetime.datetime.now()
    n=N
    while n > K:
        MaxGAP(n,M)
        j=max(H[n])
        for m in range(j,n+1):
            H[m]=H[n]
            d1F[m]=d1F[n]
            d2F[m]=d2F[n]
            L1F[m]=L1F[n]
            L2F[m]=L2F[n]
            n=j-1
    end_time = datetime.datetime.now()
    print(end_time-start_time)