#!/usr/bin/python


import xml.dom.minidom
import sys
from optparse import OptionParser
import random

from hadoop_conf import *

chunk_size = []

def xml_children(node, children_name):
	"""return list of node's children nodes with name of children_name"""
	return node.getElementsByTagName(children_name)

def xml_text(node):
	return node.childNodes[0].nodeValue

def xml_child_text(node, child_name):
	"""probably encoded in utf-8, be careful."""
	return xml_text(xml_children(node, child_name)[0])


class empty_t:
	pass

class hnode_t:
	"""HDFS node for a HDFS tree.

5-level hierarchy: rack_group (multiple identical racks), rack, node_group
(multiple identical nodes), node, and disk.
disk should be initiated with a capacity. Other nodes' capacity are
calculated by summing up children's capacity."""
	def __init__(self, parent, capacity=None, num=1):
		self.parent = parent
		self._capacity = capacity
		self._num = num
		self._children = []
		self.used = 0
		self.end = None
		self.reserved = None
		if parent <> None:
			self.index_stack = parent.index_stack[:] + [len(parent.children())]
			parent.children().append(self)
			if parent._capacity <> None:
				parent._capacity = None
		else:
			self.index_stack = []

	def clone(self, parent=None):
		'''clone a node from self, and append it to parent's children'''
		if parent == None:
			parent = self.parent
		node = hnode_t(parent, self._capacity)
		node._children = []
		if self._children <> []:
			for child in self._children:
				#print self, self.parent, self._children
				child.clone(node)
				#node._children.append(child.clone(node))  ## wrong!!!
		node.used = 0
		node.reserved = self.reserved
		return node

	def capacity(self):
		if self._capacity <> None:
			return self._capacity
		else :
			assert self._children <> []
			self._capacity = 0
			for child in self._children:
				self._capacity += child.capacity()
			return self._capacity

	def children(self):
		return self._children;

	def add_chunk(self):
		if self.used >= self.capacity():
			print 'error: node full' + self.index_stack
		self.used += chunk_size
		parent = self.parent
		if parent != None:
			parent.add_chunk()

	def name(self):
		if len(self.index_stack) == 5:   #disk
			return 'd_rg%d_%d_ng%d_%d_disk%d' % tuple(self.index_stack)
		elif len(self.index_stack) == 4:   #node
			return 'n_rg%d_%d_ng%d_%d' % tuple(self.index_stack)
		elif len(self.index_stack) == 3:   #node group template
			return 'n_rg%d_%d_ng%d' % tuple(self.index_stack)
		elif len(self.index_stack) == 2:   #rack
			return 'r_rg%d_%d' % tuple(self.index_stack)
		elif len(self.index_stack) == 1:   #rack_group
			return 'rg_rg%d' % tuple(self.index_stack)
		else:
			print 'error: request name for unknown node type. (' \
					+ self.index_stack + ')'

	def dump(self, level=0):
		if options.verbose == False:
			return
		print self.index_stack, self.used, self._capacity, len(self.children())
		node = self
		if node.children() <> []:
			for child in node.children():
				child.dump()

	def prev_node(self):
		if self.index_stack == []:
			return None

		myindex = self.index_stack[-1]
		if myindex == 0:
			return self.parent.prev_node()

		siblings = self.parent.children()
		return siblings[myindex-1]

	def global_end(self):
		'''global index at the end of a node'''
		if self.end <> None:
			return self.end

		# end should be previous node's end + self.capacity()
		prev = self.prev_node()
		if prev <> None:
			self.end = prev.global_end() + self.capacity()
		else:
			# Otherwise, this is a first node
			self.end = self.capacity()
		return self.end

	def choose_disk(self):
		'''when a node is chosen for replication, it needs to choose a disk to put the data.'''
		if self.used >= self.capacity():
			return None

		disk_id = random.randrange(len(self.children()))
		disk = self.children()[disk_id]
		if disk.used < disk.capacity():
			return disk
		else:
			return self.choose_disk()


class machine_type_t:
	def __init__(self, mt):
		disk = xml_children(mt, u'disk')[0]
		self.disk = empty_t()
		self.disk.type = str(xml_child_text(disk, u'type'))
		self.disk.capacity = int(xml_child_text(disk, u'capa'))*1024 # in MB
		self.disk.num = int(xml_child_text(disk, u'num'))

		cpu = xml_children(mt, u'cpu')[0]
		self.cpu = empty_t()
		self.cpu.type = str(xml_child_text(cpu, u'type'))
		self.cpu.cores = int(xml_child_text(cpu, u'number_of_cores'))
		self.cpu.num = int(xml_child_text(cpu, u'num'))

		mem = xml_children(mt, u'mem')[0]
		self.mem = empty_t()
		self.mem.type = str(xml_child_text(mem, u'type'))
		self.mem.capacity = str(xml_child_text(mem, u'capa')) # in MB

		# TODO: other parts of machine_type

class topology_t:
	def __init__(self, topo_xml):
		root = xml.dom.minidom.parse(topo_xml)
		self.htree = hnode_t(None)
		self.dmt = {} # dict of machine type
		topo = root.getElementsByTagName(u"topo")[0]

		# populate dict of machine type
		list_machine_type = topo.getElementsByTagName(u'machine_type')
		for mt_node in list_machine_type:
			name = str(xml_child_text(mt_node, u'name'))
			self.dmt[name] = machine_type_t(mt_node)

		# topology
		for rack_group in xml_children(topo, u"rack_group"):
			rg_node = hnode_t(self.htree)
			# rgname not in use currently. maybe a name-node map is needed.
			rg_node.rgname = str(xml_child_text(rack_group, u'name'))
			num_rack = len(xml_children(rack_group, u"rack_index"))
			self.racks = num_rack
			rack_node = hnode_t(rg_node)

			# populate the first rack_node
			for node_group in xml_children(rack_group, u"compute_node_group"):
				ng_node = hnode_t(rack_node)
				# machine type and disk
				mt_name = str(xml_child_text(node_group, u'machine_type_name'))
				mt = self.dmt[mt_name]
				ng_node.reserved = mt

				num_node = len(xml_children(node_group, u'node_index'))
				self.nodes = num_node
				node_node = hnode_t(ng_node)

				# populate the first node_node
				for i in range(mt.disk.num):
					disk_node = hnode_t(node_node, mt.disk.capacity)

				#self.htree.dump()

				# clone other node_nodes
				for i in range(num_node-1):
					new_node_node = node_node.clone()

				#self.htree.dump()

			# clone other rack_nodes
			for i in range(num_rack-1):
				new_rack_node = rack_node.clone()

		#self.htree.dump()

		self.routers = []
		for router in xml_children(topo, u'router'):
			rt = empty_t()
			rt.connect_to_groups = []
			for connect_to_group in xml_children(router, u'connect_to_group'):
				rgname = str(xml_child_text(connect_to_group, u'rack_group_name'))
				switch = empty_t()
				switch.rg = self.find_hnode(tuple([int(rgname[5:])]))
				switch.index = int(xml_child_text(connect_to_group, u'switch_index'))
				rt.connect_to_groups.append(switch)

			rt.name = str(xml_child_text(router, u'name'))
			self.routers.append(rt)

		self.data_nodes = int(xml_child_text(topo, u'data_nodes'))
		self.job_tracker = str(xml_child_text(topo, u'job_tracker'))
		topology = xml_children(topo, u'topology')
		if len(topology) > 0 :
			self.topology = str(xml_text(topology[0]))
		else:
			self.topology = None

	def find_hnode(self, index_stack):
		if len(index_stack) > 5:
			print 'Wrong index stack' + index_stack
			return None

		node = self.htree
		for i in index_stack:
			children = node.children()
			node = children[i]
		return node

	def totcl(self, topo_tcl, netsize):
		f = open(topo_tcl, 'w')

		#f.write('set int_bw %s\n' % (int_bw))
		#f.write('set int_latency %s\n' % (int_latency))
		f.write('source router.topo.tcl\n')
		f.write('array set routersArr [create-router-topology]\n\n')
		f.write('set numRouters [array size routersArr]\n')
		f.write('set toList [array get routersArr]\n\n')
		f.write('set i 0\n')
		f.write('set counter 0\n')
		f.write('foreach router [array get routersArr] {\n')
		f.write('\tif { ($i % 2) > 0 } {\n')
		f.write('\t\tset routers($counter) $router\n')
		f.write('\t\tset counter [expr $counter + 1]\n')
		f.write('\t}\n\tset i [expr $i + 1]\n')
		f.write('}\n\n')
		
		num_of_nodes = 0

		if self.topology == 'dcell':
			# special case, assume everything is symmetric
			#  take the first ng to get mt (machine type)
			#  number of nodes in a rack matters,
			#  number of racks does not matter.
			rg = self.htree.children()[0]
			racks = len(rg.children())
			r = rg.children()[0]
			ng = r.children()[0]
			nodes = len(ng.children())
			mt = ng.reserved
			f.write('set cpu_freq %f\n' % (freq_table[mt.cpu.type]))
			f.write('set cpu_cores %d\n' % (mt.cpu.cores * mt.cpu.num))
			f.write('set rbw %f\n' % (read_bw_table[mt.disk.type]))
			f.write('set wbw %f\n' % (write_bw_table[mt.disk.type]))
			f.write("\nset num_of_nodes %d\n" % (self.data_nodes))
			f.write('setup_2level_dcell %d\n' % (nodes))
			f.write('\n')
			f.write('set jt $%s\n' % (self.job_tracker))
			f.write('set racks %d\n' % (racks))
			f.write('set nodes %d\n' % (nodes))
			f.write('set data_nodes %d\n' % (self.data_nodes))
			f.write('set_mapnodes %d %d %d\n' % (racks, nodes, self.data_nodes))
			f.write('\n')
			f.close()
			return

		for rg in self.htree.children():
			#self.racks = len(rg.children())
			#print "self.htree.children loop %s" % self.racks
			self.nodes = len(rg.children())
			for r in rg.children():
				self.racks = len(r.children())
				#print "rg.children loop %s" % self.racks
				# spalmer commented out
				#f.write('set %s [$ns node]\n' % (r.name()))
				nodeCtr = 0
				for ng in r.children():
					#print "r.children loop"
					#self.racks = len(ng.children())
					#self.nodes = len(ng.children())
					mt = ng.reserved 
					# cpu information for all nodes in a node group
					freq = freq_table[mt.cpu.type]
					cores = mt.cpu.cores * mt.cpu.num
					# disk read and write bandwidths
					rbw = read_bw_table[mt.disk.type]
					wbw = write_bw_table[mt.disk.type]
					
					f.write('# cpu.type = %s cpu.cores = %s hd.type = %s\n' % (mt.cpu.type, mt.cpu.cores, mt.disk.type))

					f.write('for {set i 0} {$i < %d} {incr i} {\n' \
								% (len(ng.children())))
					# Sam Palmer
					#hostSpeeds = {0 : '6Mb', 1 : '30Mb', 2 : '4Mb', 3 : '5Mb', 4 : '5Mb', 5 : '6Mb', 6 : '6Mb', 7 : '6Mb', 8 : '7Mb', 9 : '7Mb', 10 : '8Mb', 11 : '10Mb', 12 : '10Mb', 13 : '14Mb', 14 : '14Mb', 15 : '14Mb', 16 : '30Mb', 17 : '30Mb', 18 : '30Mb',}
					f.write('\tnewnode "n_rg0_%d_ng0_$i" $routers([expr %s %% $numRouters]) %sMb %sms\n' % (nodeCtr, random.randint(0,int(netsize)-1), abs(random.normalvariate(16.0, 8.0)), abs(random.normalvariate(0.25, 0.07))))
					nodeCtr = nodeCtr + 1
#					f.write('\tnewnode "%s_$i" $%s\n' % (ng.name(), r.name()))
					num_of_nodes += len(ng.children())
					#f.write('\t$n30 set freq %f\n' % (freq))
					f.write('\t$n30 set tasklist [new MRPerf/TaskList %f %d]\n' % (freq, cores))
					f.write('\tfor {set j 0} {$j < %d} {incr j} {\n' \
								% (mt.disk.num))
					f.write('\t\t$n30 newdisk %f %f\n' % (rbw, wbw))
					f.write('\t}\n')
					f.write('}\n')
				f.write('\n')
		if True:
			# Guanying 2009.3.10: add a dedicated jobtracker
			# it does not count into num_of_nodes
			rg = self.htree.children()[0]
			r = rg.children()[0]
			ng = r.children()[0]
			mt = ng.reserved
			# cpu information for all nodes in a node group
			freq = freq_table[mt.cpu.type]
			cores = mt.cpu.cores * mt.cpu.num
			# disk read and write bandwidths
			rbw = read_bw_table[mt.disk.type]
			wbw = write_bw_table[mt.disk.type]

			'''
			jt = ng.name()+'_jobtracker'
			f.write('\nnewnode "%s" $%s\n' % (jt, r.name()))
			f.write('set jt $%s\n' % (jt))
			f.write('$jt set tasklist [new MRPerf/TaskList %f %d]\n' % (freq, cores))
			f.write('for {set j 0} {$j < %d} {incr j} {\n' \
						% (mt.disk.num))
			f.write('\t$jt newdisk %f %f\n' % (rbw, wbw))
			f.write('}\n')
			'''
		#f.write("\nset num_of_nodes %d\n" % (num_of_nodes))
		f.write("\nset num_of_nodes %d\n" % (self.data_nodes))

		for rt in self.routers:
			f.write('set %s [$ns node]\n' % (rt.name))
			f.write('$%s shape hexagon\n' % (rt.name))
			f.write('\n')
			for switch in rt.connect_to_groups:
				for r in switch.rg.children():
					f.write('$ns duplex-link $%s $%s %s %s DropTail\n' \
								% (r.name(), rt.name, ext_bw, ext_latency))

		f.write('\n')
		f.write('set jt $%s\n' % (self.job_tracker))
		f.write('set racks %d\n' % (self.racks))
		f.write('set nodes %d\n' % (self.nodes))
		f.write('set data_nodes %d\n' % (self.data_nodes))
		f.write('set_mapnodes %d %d %d\n' % (self.racks, self.nodes, self.data_nodes))
		f.write('\n')
		f.close()

	def totcl2(self, mapnodes_tcl):
		f = open(mapnodes_tcl, 'w')

		for rg_id in range(len(self.htree.children())):
			rg = self.htree.children()[rg_id]
			racks = len(rg.children())
			#racks = len(rg.children()[0].children())
			f.write('for {set i 0} {$i < %d} {incr i} {\n' % (racks))
			r = rg.children()[0]
			for ng_id in range(len(r.children())):
				ng = r.children()[ng_id]
				nodes = len(ng.children())
				f.write('\tfor {set j 0} {$j < %d} {incr j} {\n' % (nodes))
				n = ng.children()[0]

				'''
		set mn [format "%s%s%s%s" "\$n_rg0_" $i "_ng0_" $j]
		set tcp0 [new Agent/TCP/FullTcp]
		set dummy [new MRPerf/NodeApp $tcp0]
		eval "$dummy set hnode $mn"
		set app11 [$dummy new-connection $jt]
		$ns at 0.05 "$app11 snd {heartbeat}"
'''

				f.write('\t\tset mn [format "%%s%%s%%s%%s" "\\$n_rg%d_" $i "_ng%d_" $j]\n' % (rg_id, ng_id))
				f.write('\t\tset tcp0 [new Agent/TCP/FullTcp]\n')
				f.write('\t\tset dummy [new MRPerf/NodeApp $tcp0]\n')
				f.write('\t\teval "$dummy set hnode $mn"\n')
				f.write('\t\tset app11 [$dummy new-connection $jt]\n')
				f.write('\t\t$ns at 0.05 "$app11 send_heartbeat"\n')
				f.write('\t}\n')
			f.write('}\n')

		f.write('\n')
		f.close()

class conf_t:
	def __init__(self, gen_xml):
		root = xml.dom.minidom.parse(gen_xml)