'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
'	TODO
'	Eric J Anderson
'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''

import os
import os.path
import sys
import re
import argparse

# Global
g_bucket_size = [4 << 10, 8 << 10, 16 << 10, 32 << 10, 64 << 10, 128 << 10, 256 << 10, 512 << 10, 1 << 20, 2 << 20, 4 << 20, 8 << 20, 16 << 20, 32 << 20]
g_bucket_name = []

g_all = []
g_free = []
g_used = []
g_used_normal = []
g_used_pooled = []
g_used_moved = []
g_used_external = []
g_used_no_defrag = []
g_used_no_delete = []

g_res_pool_size = (16 << 20)
g_res_pool_ptr = 0
g_res_pool_path = None
g_res_pool = []

# Flags
FLAG_NONE = 0				#  0
FLAG_NORMAL = (1 << 0)		#  1
FLAG_NO_DEFRAG = (1 << 1)	#  2
FLAG_EXTERNAL = (1 << 2)	#  4
FLAG_MOVED = (1 << 3)		#  8
FLAG_NO_DELETE = (1 << 4)	# 16
FLAG_POOLED = (1 << 5)		# 32

# Classes
class CAlloc:
	def __init__(self, status):		
		self.Clear()
		self.status = status

	def __str__(self):		
		aligned = self.size

		while True:
			if (self.address == 0 or self.address % (aligned * 2)) != 0:
				break
			aligned *= 2			

		value = "{0:08X}".format(self.address) + ", " + str(self.size) + ", " + str(aligned) + ", " + self.status.upper()
		if self.IsUsed():
			value +=  ", " + self.bucket
			flags = self.GetFlags()
			if len(flags) > 0:
				value += ", " + flags.upper()

		return value

	def __eq__(self, other):
		return (self.size == other.size) and (self.hash == other.hash)

	def  __hash__(self):
		#value = str(self.id) + str(self.address)
		return self.hash.__hash__()

	def Clear(self):
		self.id = 0
		self.flag = 0
		self.bucket = ""
		self.address = 0
		self.size = 0
		self.callstack = []
		self.hash = ""

	def Parse(self, data):
		self.Clear()
		items = mysplit(data, ',')

		self.id = int(items[0])
		self.flag = int(items[1])
		self.bucket = items[2]
		self.frame = int(items[3])
		self.address = int(items[4])
		self.size = int(items[5])
		
		i = 6
		hash = ""
		while i < len(items):
			hash += items[i]
			i += 1

		if hash.startswith("\"") and hash.endswith("\""):
			hash = hash[1:len(hash)-1]
		self.SetCallstack(hash)

	def SetCallstack(self, callstack):
		self.hash = callstack
		self.callstack = []
		items = callstack.split("|")
		for item in items:
			self.callstack.append(item)

	#def GetCallstack(self):
	#	value = ""
	#	for item in self.callstack:
	#		value += item + "|"
	#	return value

	def HasFlag(self, flags):
		return (self.flag & flags) == flags

	def GetFlags(self):
		flag = ""
		#if self.HasFlag(FLAG_NORMAL):
		#	flag += "NORMAL | "

		if self.HasFlag(FLAG_NO_DEFRAG):
			flag += "no_defrag | "

		if self.HasFlag(FLAG_EXTERNAL):
			flag += "external | "

		if self.HasFlag(FLAG_MOVED):
			flag += "moved | "

		if self.HasFlag(FLAG_NO_DELETE):
			flag += "no_delete | "

		if self.HasFlag(FLAG_POOLED):
			flag += "pooled | "

		return flag[0:len(flag)-3]

	def IsUsed(self):
		return self.status == "used"

	def IsFree(self):
		return self.status == "free"

# Utility
def FlipWord(w):
   return ((w << 8) & 0x0FF00) | (w >> 8)

def FlipDWord(d):
	return ((((d)<<24) & 0xFF000000) | (((d)<<8) & 0x00FF0000) | (((d)>>8) & 0x0000FF00) | (d)>>24)

def mysplit(s, delim=None):
    return [x for x in s.split(delim) if x]

def Hex(s):
    lst = []
    for ch in s:
        hv = hex(ord(ch)).replace('0x', '')
        if len(hv) == 1:
            hv = '0'+hv
        lst.append(hv)
    
    return reduce(lambda x,y:x+y, lst)

# Flags
def InListEq(obj, list):	
	for item in list:
		if obj == item:
			return True
	return False

def IsPooled(alloc):
	return (g_res_pool_ptr > 0) and (alloc.address >= g_res_pool_ptr) and (alloc.address < (g_res_pool_ptr + g_res_pool_size))

def GetUsed(path):
	file = open(path, 'r')
	file.readline()

	data = []
	global g_res_pool_ptr
	last = None

	for line in file.readlines():
		used = CAlloc("used")
		used.Parse(line.strip())

		if last != None:
			last.size = used.size
			last = None

		# HACK to ignore resource pools
		if used.size == (16 << 20):
			g_res_pool_ptr = used.address
			last = used
		
		if IsPooled(used):
			used.flag |= FLAG_POOLED
		
		data.append(used)

	file.close()
	return data

def GetUsedType(incFlags, excFlags = FLAG_NONE):
	data = []
	for used in g_used:
		if used.HasFlag(incFlags):
			if (excFlags != FLAG_NONE) and used.HasFlag(excFlags):
				continue
			data.append(used)
	
	return data

def GetPools(path):
	file = open(path, 'r')
	file.readline()

	data = []
	for line in file.readlines():
		used = CAlloc("used")
		used.Parse(line.strip())
		data.append(used)

	file.close()
	return data

# Free
def ComputeFree(data, address, total):
	# OLD CODE - DON'T DELETE
	#
	#i = 0
	#for size in g_bucket_size:
	#	if size >= total:
	#		if size > total:
	#			size = g_bucket_size[i - 1]

	#		free = CAlloc()
	#		free.SetAddress(address)
	#		free.SetSize(size)
	#		data.append(free)

	#		address += size
	#		total -= size
	#		break
	#	i += 1

	#value = "{0:08X}".format(address) + " [" + str(total) + "]"
	#print value

	i = len(g_bucket_size) - 1
	while i >= 0:
		size = g_bucket_size[i]
		if (size <= total) and (address % size == 0):
			free = CAlloc("free")
			free.address = address
			free.size = size
			data.append(free)

			address += size
			total -= size
			break
		i -= 1

	if total > 0:
		ComputeFree(data, address, total)

def GetFree(path):
	file = open(path, 'r')
	file.readline()

	next = 0
	data = []
	for line in file.readlines():
		used = CAlloc("used")
		used.Parse(line.strip())

		if next < used.address:
			#print "next = ", next, ", used = ", used.address, ", free size = ", used.address - next
			#ComputeFree(data, used.address, used.address - next)
			ComputeFree(data, next, used.address - next)

		next = used.address + used.size

	file.close()
	return data

# Misc
def ModifyExt(path, ext):
	pos = path.rfind('.')
	if pos == -1:
		pos = len(path)

	return path[0:pos] + "." + ext

def ModifyPath(output, suffix, ext):
	pos = output.rfind('.')
	if pos == -1:
		pos = len(output)

	return output[0:pos] + " (" + suffix + ")." + ext

def SortByAddress(x, y):
	return cmp(x.address, y.address)

def MergeAndSort(list1, list2):
	return sorted(list1 + list2, SortByAddress)

def GetTotal(allocs):
	total = 0
	for alloc in allocs:
		total += alloc.size
	return total

def GetBuckets(allocs):
	buckets = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

	for alloc in allocs:
		pos = 0
		for size in g_bucket_size:
			if alloc.size == size:
				buckets[pos] += 1
				break
			pos += 1

	return buckets

def Setup(args):
	physical = "physical" in args.filename

	# HACK: We store mipmaps in physical memory
	if physical:
		i = 0
		while i < len(g_bucket_size):
			g_bucket_size[i] = int(g_bucket_size[i] * 1.343750)
			i += 1

	i = 0
	while i < len(g_bucket_size):
		size = g_bucket_size[i]
		kb = size / 1024
		g_bucket_name.append(str(kb))
		i += 1

# Print
def PrintData(title, data, heapSize):
	margin = 24

	print
	print "{0:{width}}".format(title, width=margin), "{:>{}}".format("Count", 8), "{:>{}}".format("KB", 12), "{:>{}}".format("MB", 12), "{:>{}}".format("% Total", 12), "{:>{}}".format("% Section", 12)
	print "-------------------------------------------------------------------------------------"	
	
	buckets = GetBuckets(data)	
	sectionSize = 0
	totalPercent = 0.0
	totalBytes = 0

	pos = 0
	for bucket in buckets:
		sectionSize += bucket * g_bucket_size[pos]
		pos += 1

	pos = 0
	for bucket in buckets:
		bytes = bucket * g_bucket_size[pos]
		totalBytes += bytes

		percentOfSection = 0.0
		if sectionSize > 0:
			percentOfSection = (float(bytes) / float(sectionSize)) * 100.0

		percentOfTotal = (float(bytes) / float(heapSize)) * 100.0
		totalPercent += percentOfTotal

		kb = bytes >> 10
		mb = kb / 1024.0
		print "{0:{width}}".format(g_bucket_name[pos], width=margin), "{:>{}}".format(bucket, 8), "{:>{}}".format(kb, 12), "{0:{width}.2f}".format(mb, width=12), "{0:{width}.2f}".format(percentOfTotal, width=12), "{0:{width}.2f}".format(percentOfSection, width=12)
		pos += 1

	print "-------------------------------------------------------------------------------------"
	print "{0:{width}}".format("", width=margin), "{:>{}}".format(len(data), 8), "{:>{}}".format(totalBytes >> 10, 12), "{0:{width}.2f}".format(totalBytes / 1024.0 / 1024.0, width=12), "{0:{width}.2f}".format(totalPercent, width=12), "{0:{width}.2f}".format(100.0, width=12)
	print

def Print(total):
	print "Total KB = ", total >> 10

	PrintData("All", g_all, total)
	PrintData("Free", g_free, total)
	PrintData("Used", g_used, total)	
	PrintData("Used (Normal)", g_used_normal, total)
	PrintData("Used (Pooled)", g_used_pooled, total)	
	PrintData("Used (Moved)", g_used_moved, total)
	PrintData("Used (External)", g_used_external, total)
	PrintData("Used (No Defrag)", g_used_no_defrag, total)
	PrintData("Used (No Delete)", g_used_no_delete, total)

# Frag Analysis
def GetIndexByAddress(address, allocs):
	i = 0
	for alloc in allocs:
		if alloc.address == address:
			return i
		i += 1

	return -1

def FragAnal(args):
	if len(g_free) == 0:
		return

	path = ModifyPath(args.output, "Free", "frag")
	file = open(path, "w")
	
	border = 7
	for alloc in g_free:
		pos = GetIndexByAddress(alloc.address, g_all)
		
		i = border - 1
		while i > 0:
			if (pos - i) >= 0:
				suspect = g_all[pos - i]
				flags = suspect.IsUsed() and (suspect.HasFlag(FLAG_EXTERNAL) or suspect.HasFlag(FLAG_NO_DEFRAG))
				if flags:
					file.write("*")

				j = 0
				while j < i:
					if flags and j == 0:
						file.write("  ")
					else:
						file.write("   ")
					j += 1
				file.write(str(suspect) + "\n")
			i -= 1

		file.write(str(g_all[pos]) + "\n")
		
		i = 1
		while i < border:
			if (pos + i) < len(g_all):
				suspect = g_all[pos + i]
				flags = suspect.IsUsed() and (suspect.HasFlag(FLAG_EXTERNAL) or suspect.HasFlag(FLAG_NO_DEFRAG))
				if flags:
					file.write("*")

				j = 0
				while j < i:
					if flags and j == 0:
						file.write("  ")
					else:
						file.write("   ")
					j += 1
				file.write(str(suspect) + "\n")
			i += 1

		file.write("\n")

	file.close()

def DispersionAnal(args, freeTotal):
	if len(g_free) == 0:
		return

	path = ModifyPath(args.output, "Free", "disp")
	file = open(path, "w")

	desired = args.dispersion << 10	
	badTotal = 0
	allocs = {}

	# Compute
	for alloc in g_free:				
		pos = GetIndexByAddress(alloc.address, g_all)
		
		allocs[alloc] = []

		# Prev
		i = pos - 1	
		address = alloc.address

		while not (address % desired) == 0:
			if i < 0:
				break

			suspect = g_all[i]
			if suspect.HasFlag(FLAG_EXTERNAL) or suspect.HasFlag(FLAG_NO_DEFRAG):
				allocs[alloc].append(suspect)
			address = g_all[i].address
			i -= 1

		# Next
		i = pos + 1
		address = alloc.address

		while not (address % desired) == 0:
			if i >= len(g_all):
				break

			suspect = g_all[i]
			if suspect.HasFlag(FLAG_EXTERNAL) or suspect.HasFlag(FLAG_NO_DEFRAG):
				allocs[alloc].append(suspect)
			address = g_all[i].address
			i += 1

		if len(allocs[alloc]) > 0:
			badTotal += alloc.size

	# Print
	file.write(str(badTotal >> 10) + " KB is unable to be coalesced due to non-defragmentable blocks within " + str(args.dispersion) + " KB\n")
	file.write("This represents {0:.2f}% of the total free memory\n".format((float(badTotal) / float(freeTotal)) * 100.0))
	file.write("\n")

	for alloc in allocs:
		suspects = allocs[alloc]
		if len(suspects) > 0:
			for suspect in suspects:
				if suspect.address > alloc.address:
					break
				file.write("  -> " + str(suspect) + "\n")
			
			file.write(str(alloc) + "\n")

			for suspect in suspects:
				if suspect.address > alloc.address:
					file.write("  -> " + str(suspect) + "\n")
			
			file.write("\n")

	file.close()
		
# CSV
def SaveCSVData(file, title, data, heapSize):
	if len(data) == 0:
		return

	# Header
	file.write(title + ",")
	subtitles = ["COUNT", "KB", "MB", "%TOTAL", "%SECTION"]
	for subtitle in subtitles:
		file.write(subtitle + ",")
	file.write("\n")

	buckets = GetBuckets(data)	
	sectionSize = 0
	totalPercent = 0.0
	totalBytes = 0

	pos = 0
	for bucket in buckets:
		sectionSize += bucket * g_bucket_size[pos]
		pos += 1

	pos = 0
	for bucket in buckets:
		bytes = bucket * g_bucket_size[pos]
		totalBytes += bytes

		percentOfSection = 0.0
		if sectionSize > 0:
			percentOfSection = (float(bytes) / float(sectionSize)) * 100.0

		percentOfTotal = (float(bytes) / float(heapSize)) * 100.0
		totalPercent += percentOfTotal

		kb = bytes >> 10
		mb = kb / 1024.0
		file.write(g_bucket_name[pos] + "," + str(bucket) + "," + str(kb) + "," + str(mb) + "," + str(percentOfTotal) + "," + str(percentOfSection) + "\n")
		pos += 1

	file.write(" " + "," + str(len(data)) + "," + str(totalBytes >> 10) + "," + str(totalBytes / 1024.0 / 1024.0) + "," + str(totalPercent) + "," + str(100.0) + "\n")
	file.write("\n")


def SaveCSV(args, total):
	file = open(args.output + ".csv", "w")

	SaveCSVData(file, "All", g_all, total)
	SaveCSVData(file, "Free", g_free, total)
	SaveCSVData(file, "Used", g_used, total)	
	SaveCSVData(file, "Used Normal", g_used_normal, total)
	SaveCSVData(file, "Used Moved", g_used_moved, total)
	SaveCSVData(file, "Used External", g_used_external, total)
	SaveCSVData(file, "Used No Defrag", g_used_no_defrag, total)
	SaveCSVData(file, "Used No Delete", g_used_no_delete, total)
	
	file.close()

# Assets
def GetAsset(alloc):
	for func in alloc.callstack:
		if (not func.startswith(".") and "." in func) or "platform:" in func:
			return func.lower()
	return None

def GetAssets(data):
	assets = {}

	for alloc in data:
		asset = GetAsset(alloc)
		if asset != None:
			if assets.get(asset) == None:
				assets[asset] = []
			assets[asset].append(alloc)

	return assets

def SaveAssetData(args, suffix, data):
	assets = GetAssets(data)
	if len(assets) == 0:
		return

	path = ModifyPath(args.output, suffix, "ass")
	file = open(path, "w")

	sortedList = sorted(list(assets))

	for asset in sortedList:
		allocs = assets[asset]

		total = 0
		for alloc in allocs:
			total += alloc.size

		file.write("{0},{1:d}\n".format(asset, total))
		for alloc in allocs:
			file.write("{0:08X},{1},{2:d}\n".format(alloc.address, alloc.bucket, alloc.size))
		file.write("\n")

	file.close()

def SaveAssets(args):
	SaveAssetData(args, "All", g_all)
	SaveAssetData(args, "Free", g_free)
	SaveAssetData(args, "Used", g_used)	
	SaveAssetData(args, "Used Normal", g_used_normal)
	SaveAssetData(args, "Used Moved", g_used_moved)
	SaveAssetData(args, "Used External", g_used_external)
	SaveAssetData(args, "Used No Defrag", g_used_no_defrag)
	SaveAssetData(args, "Used No Delete", g_used_no_delete)

# Callstacks
def SaveCallstackData(args, suffix, data):
	if len(data) == 0:
		return

	path = ModifyPath(args.output, suffix, "call")
	file = open(path, "w")

	for alloc in data:
		file.write("{0:08X},{1},{2:d}\n".format(alloc.address, alloc.bucket, alloc.size))

		if len(alloc.callstack) > 0:
			for func in alloc.callstack:
				file.write(func + "\n")
			file.write("\n")

	file.close()

def SaveCallstacks(args):
	SaveCallstackData(args, "All", g_all)
	SaveCallstackData(args, "Free", g_free)
	SaveCallstackData(args, "Used", g_used)	
	SaveCallstackData(args, "Used Normal", g_used_normal)
	SaveCallstackData(args, "Used Moved", g_used_moved)
	SaveCallstackData(args, "Used External", g_used_external)
	SaveCallstackData(args, "Used No Defrag", g_used_no_defrag)
	SaveCallstackData(args, "Used No Delete", g_used_no_delete)

# Main
def ParseArgs(argv):
	parser = argparse.ArgumentParser( description = 'Analyzes MemVisualize CSV resource memory files' )
	parser.add_argument( '-a', '--assets', action = 'store_true', help = 'Output assets')
	parser.add_argument( '-c', '--callstacks', action = 'store_true', help = 'Output callstacks')
	parser.add_argument( '-d', '--dispersion', type=int, help = 'Perform dispersion KB test on free memory')
	parser.add_argument( '-f', '--frag', action = 'store_true', help = 'Perform fragmentation analysis')
	parser.add_argument( '-o', '--output', help = 'Output output')
	parser.add_argument("filename")
	return parser.parse_args( argv )

def Main():	
	# Global
	global g_all
	global g_free
	global g_used
	global g_used_normal
	global g_used_pooled
	global g_used_moved
	global g_used_external
	global g_used_no_defrag
	global g_used_no_delete
	global g_bucket_size
	global g_bucket_name

	global g_res_pool_size
	global g_res_pool_ptr
	global g_res_pool_path
	global g_res_pool

	# Input
	args = ParseArgs(sys.argv[1:])	
	Setup(args)
	
	if not args.output:
		args.output = "output"

	# Data
	path = args.filename.lower()
	g_res_pool_path = args.filename[0:args.filename.rfind('.')] + " Pool.csv"
	if os.path.exists(g_res_pool_path):
		g_res_pool = GetPools(g_res_pool_path)

	g_used = GetUsed(path)	
	g_free = GetFree(path)		
	g_all = MergeAndSort(g_free, g_used)

	g_used_pooled = GetUsedType(FLAG_POOLED)
	g_used_normal = GetUsedType(FLAG_NORMAL, FLAG_POOLED)
	g_used_moved = GetUsedType(FLAG_MOVED, FLAG_POOLED)
	g_used_external = GetUsedType(FLAG_EXTERNAL)
	g_used_no_defrag = GetUsedType(FLAG_NO_DEFRAG)
	g_used_no_delete = GetUsedType(FLAG_NO_DELETE)

	# Total
	usedTotal = GetTotal(g_used)
	freeTotal = GetTotal(g_free)
	total = usedTotal + freeTotal

	# Output
	Print(total);

	# CSV
	SaveCSV(args, total)

	# Frag Analysis
	if args.frag:
		FragAnal(args)

	# Dispersion Analysis
	if args.dispersion:
		DispersionAnal(args, freeTotal)

	# Assets
	if args.assets:
		SaveAssets(args)

	# Callstacks
	if args.callstacks:
		SaveCallstacks(args)

# Run
if __name__ == "__main__":
	Main()