"""
This module implements directories as items for partitioner.py and can 
split a file collection (e.g., from touchlogger) into collections of 
directories.
"""

from itertools import *
from sets import Set
import os.path

from partitioner import Item


class Directory(Item):
	def __init__(self, rootDir, dirName, files):
		self.rootDir = rootDir
		self.dirName = dirName
		self.files = files
		Item.__init__(self, self.dirName, self._computeSize())

	def __iter__(self):
		return iter(self.files)

	def __repr__(self):
		return "Directory('%s', '%s', %s)"%(self.rootDir, self.dirName,
			self.files)
	
	def _computeSize(self):
		return sum(imap(
			lambda fName: os.path.getsize(os.path.join(
				self.rootDir, self.dirName, fName)), self.files))
	
	def getDirName(self):
		return self.dirName
	
	def getRootDir(self):
		return self.rootDir
	

class DirCollection(Set):
	def __init__(self, fileSet, rootDir, topLevel=0):
		Set.__init__(self)
		for dirName, fNames in self._makeDirDict(fileSet, topLevel).iteritems():
			self.add(Directory(rootDir, dirName, fNames))

	def __repr__(self):
		return str(self)

	def _makeDirDict(self, fileSet, topLevel):
		dirDict = {}
		for fPath in fileSet:
			if topLevel:
				head, tail = fPath.split('/', 1)
			else:
				head, tail = os.path.split(fPath)
			dirDict.setdefault(head, []).append(tail)
		return dirDict
	
	def getTotalSize(self):
		return sum(imap(lambda a: a.getSize(), self))


def _test():
	import touchlogger
	fileSet = touchlogger.RelativeFileSet("/usr/media/mpg3")
	dc = DirCollection(fileSet, fileSet.getRootDir())
	dcTop = DirCollection(fileSet, fileSet.getRootDir(), topLevel=1)
	for dir in dc:
		print "%-50s  %d"%(dir.getDirName(), dir.getSize()/1000000)
	print "========"
	for dir in dcTop:
		print "%-50s  %d"%(dir.getDirName(), dir.getSize()/1000000)


if __name__=="__main__":
	_test()