#!/usr/bin/env python

import _init_paths
import lxml.etree as ET
import datasets.imdb as imdb
import numpy as np
import os
import ast
import scipy
try:
	import dill as pickle
except ImportError:
	try:
		import cPickle as pickle
	except ImportError:
		import pickle

CACHE_DIR = os.path.join('.', 'cache')

cls = ('__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor')

class xmldb(imdb):
	def __init__(self, xml, name=None, clazz='VOC', load_cache=True):
		if not isinstance(xml, list):
			xml = [xml]
		root, self.pklname, tmpname = self.multiplexml(xml)
		if os.path.exists(os.path.join(self.cache_dir(), self.pklname + '.pkl')) and load_cache:
			self.load_cache(self.pklname)
		else:
			if name is None:
				name = tmpname
			imdb.__init__(self, name)
			if clazz == 'VOC':
				self._classes = cls
			else:
				self._classes = clazz
			self._num_classes = len(self._classes)
			self._class_to_ind = dict(zip(self.classes, xrange(self.num_classes)))
			
			self.load_data(root)
			self.save_cache(self.pklname)
		
	
	def multiplexml(self, xmllist):
		xml = ET.Element('data')
		tmpname = []
		names = []
		for x in xmllist:
			a = ET.parse(x).getroot()
			tmpname.append(a.get('dataset'))
			names.append(os.path.splitext(os.path.basename(os.path.abspath(x)))[0])
			xml.extend(a[0])
		return xml, '-'.join(map(lambda x: x[0]+'_'+x[1], zip(names,tmpname))), '_'.join(tmpname)
		

	def load_data(self, root):
		index = []
		roidb = []
		for f in root:
			if 'framecount' not in f.keys():
				index.append(f.get('path'))
				objs = list(f)
				num_objs = len(objs)

				boxes = np.zeros((num_objs, 4), dtype=np.uint16)
				gt_cls = np.zeros((num_objs), dtype=np.int32)
				overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32)
				
				for i, obj in enumerate(objs):
					bbox = ast.literal_eval(obj.get('bbox'))
					boxes[i,:] = [int(round(b,0)) for b in bbox]
					cls = self._class_to_ind[obj.get('class')]
					gt_cls[i] = cls
					overlaps[i,cls] = 1.0
				overlaps = scipy.sparse.csr_matrix(overlaps)

				roidb.append({'boxes' : boxes, 'gt_classes': gt_cls, 'gt_overlaps' : overlaps, 'flipped' : False})
		self._roidb = roidb
		self._image_index = index
	
	def image_path_at(self, i):
		return self._image_index[i]
	
	def cache_dir(self):
		cdir = os.path.join('.', 'cache')
		if not os.path.exists(cdir):
			os.makedirs(cdir)
		return cdir
	
	def load_cache(self, name):
		with open(os.path.join(self.cache_dir(), name + '.pkl'), mode='rb') as f:
			self.__dict__.update(pickle.load(f))
	
	def save_cache(self, name):
		with open(os.path.join(self.cache_dir(), name + '.pkl'), mode='wb') as f:
			pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL)
