#
# mega unduplicator
#
# (C) C. Tismer 03/2004

"""
Finding duplicates in a huge data set.

The idea: computing a sha key is relatively cheap, and reduces
the amount of data while keeping the equality property
almost valid.
Another nice feature of this approach is that we don't need
any preprocessing of the records. Just compute the sha key
of any representation that is easy to obtain.

For really large data sets, we need a way to split the data
apart, in order to avoid swapping. This is easily done by a two-
phase algorithm:
In phase 1, we compute the sha key for every record and distribute
the result into bins, computed from a few bits of the key.
As a result, we know that there are for sure no equalities
between the bins.
>From time to time, we flush the bins out into an according number of
files.
In phase 2, we read every bin back in and insert all the keys into
a dictionary, to find clashes.
We can then either trust the uniqueness of sha keys, or we do a further
check against the data.
"""

import sha, cPickle as pickle

class ShaBins:
    def __init__(self):
        # simply make 256 files
        self.bins = [[] for i in range(256)]
        self.shafile = file("shafile.tmp", "wb")
        self.offsets = [[] for i in range(256)]
        self.flushes  = 0
        self.duplicates = 0
    def writedata(self, reader):
        recno = 1
        # work in chunks of 100000 records
        step = 100000
        newhash = sha.new
        bins = self.bins
        for line in iter(reader):
            digest = newhash(line).digest()
            num = ord(digest[0])
            bins[num].append( (digest, recno) )
            if recno % step == 1:
                print "processing record %d" % (recno)
                self.flush_bins()
            recno += 1
        self.flush_bins()
        print "processed %d records." % recno
        self.reopen_file()
    def flush_bins(self):
        f = self.shafile
        for i in range(len(self.bins)):
            self.offsets[i].append(f.tell())
            pickle.dump(self.bins[i], f, 2)
            self.bins[i] = []
        self.flushes += 1
    def reopen_file(self):
        f = self.shafile
        f.close()
        self.shafile = file(f.name, "rb")
    def process_duplicates(self):
        f = self.shafile
        for i in range(len(self.bins)):
            seen = {}
            has = seen.has_key
            print "processing bin %d of %d" % (i+1, len(self.bins))
            for k in range(self.flushes):
                f.seek(self.offsets[i][k])
                bin = pickle.load(f)
                for key, value in bin:
                    if has(key):
                        print "record %d is a duplicate of %d" % (value, seen[key])
                        self.duplicates += 1
                    else:
                        seen[key] = value
    def run(self, reader):
        try:
            self.writedata(reader)
            self.process_duplicates()
        finally:
            self.shafile.close()

if __name__ == "__main__":
    # simple test
    import os, time
    os.chdir("d:/tmp")
    fname = "bigfile.txt"
    nrecords = 10000000
    fluff = 300
    fluffstr = fluff * "x"
    if not os.path.exists(fname):
        f = file(fname, "w")
        print "creating file..."
        for i in xrange(nrecords):
            if i == 1234567:
                i = 7654321
            print >> f, fluffstr, i
    f = file(fname)
    u = ShaBins()
    start = time.time()
    u.run(f)
    delta = time.time() - start
    print "found %d duplicates" % u.duplicates
    print "%0.3f seconds for %d records of about %d bytes each" % (
        delta, nrecords, fluff)
    