#!/usr/bin/env python3
from search_proto import load, seqs
import sys,random,time

def solve(D,P,K,locks,limit=60,seed=1):
    codes=seqs(D,P); allkeys=set(range(K)); unauth=[list(allkeys-A) for A in locks]
    contains=[[] for _ in range(K)]
    for l,A in enumerate(locks):
        for k in A: contains[k].append(l)
    rng=random.Random(seed); t0=time.time(); best=10**9; bestassign=None
    # precompute code masks? value at p
    while time.time()-t0<limit:
        assign=[rng.randrange(len(codes)) for _ in range(K)]
        # maybe initialize all same? random from spread
        total=None
        for step in range(1000000):
            masks=[]; viol_by_key=[0]*K; viol=[]; total=0
            for l,A in enumerate(locks):
                m=[0]*P
                for kk in A:
                    c=codes[assign[kk]]
                    for p,d in enumerate(c): m[p]|=1<<d
                masks.append(m)
                for kk in unauth[l]:
                    c=codes[assign[kk]]
                    ok=True
                    for p,d in enumerate(c):
                        if not (m[p]>>d)&1: ok=False; break
                    if ok:
                        viol.append((l,kk)); viol_by_key[kk]+=1; total+=1
            if total<best:
                best=total; bestassign=assign[:]; print('best',best,'step',step,'time',time.time()-t0, flush=True)
                if total==0: return [codes[i] for i in assign]
            if time.time()-t0>=limit: break
            # choose among keys involved, sometimes authorized keys of bad locks
            if not viol: return [codes[i] for i in assign]
            if rng.random()<0.4:
                l,k=rng.choice(viol)
                pool=list(locks[l])+[k]
                k=rng.choice(pool)
            else:
                mx=max(viol_by_key); cand=[i for i,v in enumerate(viol_by_key) if v==mx]
                k=rng.choice(cand) if mx else rng.randrange(K)
            affected=set(contains[k])
            for l in range(len(locks)):
                if k not in locks[l]: affected.add(l)  # own violation maybe all noncontains
            # since k unauth for almost all, all locks affected for own status; but only O(L)
            unaffected=0
            for l in range(len(locks)):
                if l in affected: continue
                m=masks[l]
                for kk in unauth[l]:
                    c=codes[assign[kk]]
                    if all((m[p]>>d)&1 for p,d in enumerate(c)): unaffected+=1
            bestci=assign[k]; bestscore=10**9
            # sample all for small K, random 500 for large
            cand_iter=range(len(codes)) if K<60 else rng.sample(range(len(codes)), min(len(codes),600))
            for ci in cand_iter:
                old=assign[k]; assign[k]=ci; sc=unaffected
                for l in affected:
                    A=locks[l]; m=[0]*P
                    for kk in A:
                        c=codes[assign[kk]]
                        for p,d in enumerate(c): m[p]|=1<<d
                    for kk in unauth[l]:
                        c=codes[assign[kk]]
                        if all((m[p]>>d)&1 for p,d in enumerate(c)):
                            sc+=1
                            if sc>bestscore: break
                    if sc>bestscore: break
                assign[k]=old
                # anneal/noise
                if sc<bestscore or (sc==bestscore and rng.random()<0.005): bestscore=sc; bestci=ci
            if bestscore>=total and rng.random()<0.03: bestci=rng.randrange(len(codes))
            assign[k]=bestci
    return None
if __name__=='__main__':
    D,P,K,locks=load(sys.argv[1]); print('params',D,P,K,len(locks),'domain',len(seqs(D,P)))
    sol=solve(D,P,K,locks,float(sys.argv[2]) if len(sys.argv)>2 else 60, int(sys.argv[3]) if len(sys.argv)>3 else 1)
    print('SOL',sol is not None)
    if sol: print(sol)
