from search_proto import load
import sys,random,time

def covered_by_col(colors, locks, D, pairs):
    L=len(locks); masks=[]
    for A in locks:
        m=0
        for k in A: m|=1<<colors[k]
        masks.append(m)
    cov=[]
    for idx,(l,k) in enumerate(pairs):
        if not ((masks[l]>>colors[k])&1): cov.append(idx)
    return cov

def solve(D,P,K,locks,limit=20):
    allpairs=[(l,k) for l,A in enumerate(locks) for k in range(K) if k not in A]
    rng=random.Random(3); t0=time.time(); bestleft=10**9; bestseq=None
    while time.time()-t0<limit:
        uncovered=allpairs[:]; seq=[]; prev=None
        for p in range(P):
            allowed=[list(range(D)) if prev is None else [d for d in range(D) if abs(d-prev[k])<=1] for k in range(K)]
            colors=[rng.choice(allowed[k]) for k in range(K)]
            bestcov=-1; stagn=0
            # local improve
            for it in range(2000):
                covset=set(covered_by_col(colors,locks,D,uncovered)); cov=len(covset)
                if cov>bestcov: bestcov=cov; stagn=0
                else: stagn+=1
                if stagn>200: break
                k=rng.randrange(K); old=colors[k]; bestd=old; b= cov
                for d in allowed[k]:
                    if d==old: continue
                    colors[k]=d
                    c=len(covered_by_col(colors,locks,D,uncovered))
                    if c>b: b=c; bestd=d
                colors[k]=bestd
            covidx=set(covered_by_col(colors,locks,D,uncovered))
            uncovered=[pair for i,pair in enumerate(uncovered) if i not in covidx]
            seq.append(colors[:]); prev=colors[:]
            if len(uncovered)<bestleft:
                bestleft=len(uncovered); bestseq=[s[:] for s in seq]
                print('bestleft',bestleft,'pos',p,'time',time.time()-t0, flush=True)
            if not uncovered:
                # transpose to key cuts
                return [tuple(seq[p][k] for p in range(P)) for k in range(K)]
    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 20)
    print('SOL',sol is not None)
    if sol: print(sol)
