from search_proto import load,seqs
import sys,random,time

def evalx(x,D,P,K,locks,weights):
    wsum=0; cnt=0; bad=[]
    for l,A in enumerate(locks):
        Aset=set(A); masks=[0]*P
        for k in A:
            for p in range(P): masks[p]|=1<<x[k][p]
        for k in range(K):
            if k in Aset: continue
            if all((masks[p]>>x[k][p])&1 for p in range(P)):
                cnt+=1; wsum+=weights[(l,k)]; bad.append((l,k))
    return cnt,wsum,bad

def solve(D,P,K,locks,limit=120):
 rng=random.Random(5); codes=seqs(D,P); weights={(l,k):1 for l,A in enumerate(locks) for k in range(K) if k not in A}; t0=time.time(); best=999999
 while time.time()-t0<limit:
  x=[list(rng.choice(codes)) for _ in range(K)]
  cnt,ws,bad=evalx(x,D,P,K,locks,weights); stagnant=0
  for step in range(300000):
   if cnt<best:
    best=cnt; print('best',best,'w',ws,'step',step,'time',time.time()-t0, flush=True)
    if cnt==0: return [tuple(r) for r in x]
   if time.time()-t0>limit: break
   oldws=ws
   if bad and rng.random()<0.75:
    l,k=rng.choice(bad)
    if rng.random()<0.6: k=rng.choice(list(locks[l]))
   else: k=rng.randrange(K)
   p=rng.randrange(P)
   allowed=[d for d in range(D) if (p==0 or abs(x[k][p-1]-d)<=1) and (p==P-1 or abs(x[k][p+1]-d)<=1)]
   old=x[k][p]; bd=old; bcnt=cnt; bws=ws; bbad=bad
   for d in allowed:
    if d==old: continue
    x[k][p]=d; nc,nw,nb=evalx(x,D,P,K,locks,weights)
    if nw<bws or (nw==bws and nc<bcnt) or (nw==bws and nc==bcnt and rng.random()<0.01):
     bcnt=nc; bws=nw; bd=d; bbad=nb
   if bws>=ws:
    stagnant+=1
    if stagnant>200:
     for pair in bad: weights[pair]+=1
     cnt,ws,bad=evalx(x,D,P,K,locks,weights); stagnant=0
     continue
    if rng.random()<0.03: bd=rng.choice(allowed)
   else: stagnant=0
   x[k][p]=bd; cnt,ws,bad=evalx(x,D,P,K,locks,weights)
 return None
if __name__=='__main__':
 D,P,K,locks=load(sys.argv[1]); print(D,P,K,len(locks)); sol=solve(D,P,K,locks,float(sys.argv[2]) if len(sys.argv)>2 else 120); print('SOL',sol is not None); print(sol if sol else '')
