# Copyright 2016-2018 Fawzi Mohamed, Carl Poelking, Daria Tomecka
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

from __future__ import print_function
from __future__ import division
from future import standard_library
standard_library.install_aliases()
from builtins import zip
from builtins import str
from past.utils import old_div
from builtins import object
# See git https://github.com/capoe/momo.git
import os
import sys
import subprocess
import argparse
import time
import numpy as np
try:
	from lxml import etree
except ImportError:
    pass

boolean_dict = \
    {'true' : True,   '1' : True,  'yes' : True,
     'false' : False, '0' : False, 'no' : False, 'none' : False }

# =============================================================================
# XML WRAPPERS
# =============================================================================

class ExtendableNamespace(argparse.Namespace):
	def AddNamespace(self,  **kwargs):
		for name in kwargs:
			att = getattr(self, name, None)
			if att is None:
				setattr(self, name, kwargs[name])
			else:
				setattr(self, name, kwargs[name].As(type(att)))
		return
	def Add(self, name, value):
		att = getattr(self, name, None)
		if att is None:
			setattr(self, name, value)
		else:
			att.Add(name, value)
		return value

def GenerateTreeDict(tree, element, path='', paths_rel_to=None):
	if type(element) == etree._Comment: return [], {}
	# Update path
	if path == '':
		if element.tag != paths_rel_to:
			path += element.tag
	else:
		path += '/' + element.tag
	# Containers for lower levels
	tag_node = {}
	nodes = []
	# Construct Node
	xmlnode = XmlNode(element, path) # tree.getpath(element))
	nodes.append(xmlnode)
	if len(element) == 0:
		tag_node[path] = xmlnode
	#else:
	#	print "len 0", xmlnode.path
	#	tag_node[path] = xmlnode
	# Iterate over children
	for child in element:
		child_elements, childtag_element = GenerateTreeDict(tree, child, path)
		nodes = nodes + child_elements
		for key in childtag_element.keys():
			if key in tag_node:
				if type(tag_node[key]) != list:
					tag_node[key] = [ tag_node[key], childtag_element[key] ]
				else:
					tag_node[key].append(childtag_element[key])
			else:
				tag_node[key] = childtag_element[key]
	return nodes, tag_node

def NamespaceFromDict(tree_dict):
	nspace = ExtendableNamespace()
	for key in tree_dict.keys():
		#print "NSPACE Path = %s" % key
		sections = key.split('/')
		values = [ None for s in sections ]
		values[-1] = tree_dict[key]
		add_to_nspace = nspace
		for s,v in zip(sections, values):
			if v == None:
				#print "NSPACE Check for existing"
				if getattr(add_to_nspace, s, None):
					#print "NSPACE '%s' already exists" % s
					add_to_nspace = getattr(add_to_nspace, s, None)
				else:
					#print "NSPACE '%s' created" % s
					sub_nspace = ExtendableNamespace()
					add_to_nspace = add_to_nspace.Add(s, sub_nspace)
			else:
				#print "NSPACE '%s' value added" % s
				add_to_nspace.Add(s, v)
	return nspace


class XmlTree(list):
	def __init__(self, xmlfile, paths_rel_to=None):
		self.xmlfile = xmlfile
		self.xtree = etree.parse(xmlfile)
		self.xroot = self.xtree.getroot()
		self.nodes, self.tag_node = GenerateTreeDict(self.xtree, self.xroot, '', paths_rel_to)
		self.xspace = NamespaceFromDict(self.tag_node)
	def SelectByTag(self, tag):
		selection = [ e for e in self.nodes if e.tag == tag ]
		return selection
	def __getitem__(self, key):
		return self.tag_node[key]
	def keys(self):
		return list(self.tag_node.keys())


class XmlNode(object):
	def __init__(self, element, path):
		self.path = path
		self.node = element
		self.tag = element.tag
		self.value = element.text
		self.attributes = element.attrib
	def As(self, typ):
		if typ == np.array:
			sps = self.value.split()
			return typ([ float(sp) for sp in sps ])
		elif typ == bool:
			return boolean_dict.get(self.value.lower())
		else:
			return typ(self.value)
	def AsArray(self, typ, sep=' ', rep='\t\n'):
		for r in rep:
			self.value = self.value.replace(r, sep)
		sp = self.value.split(sep)
		return [ typ(s) for s in sp if str(s) != '' ]
	def SetNodeValue(self, new_value):
		self.value = new_value
		if self.node != None:
			self.node.firstChild.nodeValue = new_value
		return
	def __getitem__(self, key):
		return self.node.get(key)


# =============================================================================
# COMMAND LINE & XML INPUT INTERFACE
# =============================================================================

class CLIO_HelpFormatter(argparse.HelpFormatter):
	def _format_usage(self, usage, action, group, prefix):
		return "%s : Command Line Interface\n" % sys.argv[0]


class OptionsInterface(object):
	def __init__(self):
		# COMMAND-LINE ARGUMENTS
		self.is_connected_to_cmd_ln = False
		self.cmd_ln_args = None
		self.cmd_ln_opts = None
		self.cmd_ln_nicknames = [ '-h' ]
		self.boolean_translator = boolean_dict
		self.subtype = str
		# XML OPTIONS FILE
		self.is_connected_to_xml = False
		self.xmlfile = None
		self.tree = None
		self.xdict = None
		self.xspace = None
		# JOINED OPTIONS
		self.opts = ExtendableNamespace()
	def Connect(self, xmlfile=None):
		self.ConnectToCmdLn()
		self.ConnectToOptionsFile(xmlfile)
	def Parse(self, xkey='options'):
		if self.is_connected_to_cmd_ln:
			self.ParseCmdLn()
		if self.is_connected_to_xml:
			self.ParseOptionsFileXml(xkey)
		if self.is_connected_to_cmd_ln and not self.is_connected_to_xml:
			return self.cmd_ln_opts
		elif self.is_connected_to_xml and not self.is_connected_to_cmd_ln:
			return self.xspace
		else:
			return self.cmd_ln_opts, self.xspace
	def ParseOptionsFile(self, xmlfile, xkey):
		self.xmlfile = xmlfile
		self.is_connected_to_xml = True
		self.ParseOptionsFileXml(xkey)
		return self.xspace
	# COMMAND-LINE CONVENIENCE ================================================
	def __call__(self):
		return self.cmd_ln_opts
	def ConnectToCmdLn(self, prog=sys.argv[0], descr=None):
		self.cmd_ln_args = argparse.ArgumentParser(prog=sys.argv[0],
			formatter_class=lambda prog: CLIO_HelpFormatter(prog,max_help_position=70))
		self.is_connected_to_cmd_ln = True
		return
	def ParseCmdLn(self):
		self.cmd_ln_opts = self.cmd_ln_args.parse_args()
	def InterpretAsBoolean(self, expr):
		try:
			return self.boolean_translator.get(expr.lower())
		except KeyError:
			raise ValueError('CLIO does not know how to convert %s into a boolean.' % expr)
	def InterpretAsNumpyArray(self, expr):
		# print("Interpret", expr)
		array = [ float(e) for e in expr ]
		array = np.array(array)
		return array
	def InterpretAsList(self, expr):
		array = [ self.subtype(e) for e in expr ]
		return array
	def AddArg(self, name, typ=str, nickname=None,
			default=None, destination=None, help=None):
		# Sort out <name> (e.g. --time) vs <destination> (e.g., time)
		if '--' != name[0:2]:
			dest = name
			name = '--' + name
		else:
			dest = name[2:]
		# Sort out <default> vs <required>
		if default == None: required = True
		else: required = False
		# Construct default <help> it not given
		if help == None: help = str(typ)
		# Construct <nickname> if not given
		if nickname == None:
			nickname = '-'
			for char in dest:
				nickname += char
				if not nickname in self.cmd_ln_nicknames:
					break
			if nickname in self.cmd_ln_nicknames:
				raise ValueError('CLIO could not construct nickname from %s option'\
					% name)
			self.cmd_ln_nicknames.append(nickname)
		# Process type
		if typ in [int, float, str]:
			nargs = None
		elif typ == bool:
			typ = self.InterpretAsBoolean
			nargs = None
		elif typ == np.array:
			raise NotImplementedError
			typ = float # self.InterpretAsNumpyArray
			nargs = 3
		elif typ == list:
			typ = str
			nargs = '*'
		elif len(typ) == 2 and typ[0] == list:
			typ = typ[1]
			nargs = '*'
		else:
			raise NotImplementedError("CLIO does not know how to generate type '%s'"\
				 % typ)
		self.cmd_ln_args.add_argument(nickname, name,
				dest=dest,
				action='store',
				nargs=nargs,
				required=required,
				type=typ,
				metavar=dest[0:1].upper(),
				default=default,
				help=help)
		return
	# COMMAND-LINE CONVENIENCE ================================================
	def ConnectToOptionsFile(self, xmlfile):
		if xmlfile == None or xmlfile == '': return
		self.xmlfile = xmlfile
		self.is_connected_to_xml = True
		return
	def ParseOptionsFileXml(self, xkey='options'):
		if self.xmlfile == None: return
		self.tree = XmlTree(self.xmlfile, paths_rel_to=xkey)
		self.xdict = self.tree.tag_node
		self.xspace = self.tree.xspace
		return
	def __getitem__(self, key):
		try:
			return self.xspace.__dict__[key]
		except KeyError:
			return self.cmd_ln_opts.__dict__[key]
		except KeyError:
			raise AttributeError('No such option registered: \'%s\'' % key)
		return None


# =============================================================================
# OS SHELL INTERFACE
# =============================================================================

class ShellInterface(object):
	def __init__(self):
		# =====================================================================
		# PRINTER ATTRIBUTES
		self.color_dict = { \
			'pp' : '\033[95m',
			'mb' : '\033[34m',
			'lb' : '\033[1;34m',
			'my' : '\033[1;33m',
			'mg' : '\033[92m',
			'mr' : '\033[91m',
			'ww' : '\033[0;1m',
			'ok' : '\033[92m',
			'xx' : '\033[91m',
			'warning' : '\033[93m',
			'error' : '\033[95m',
			'endcolor' : '\033[0;1m' }
		self.justify_dict = { \
			'o' : '  o ',
			'.' : '... ',
            'r' : '\r',
			'ro' : '\r  o '}
		self.pp = OS_COLOR('pp')
		self.lb = OS_COLOR('lb')
		self.mb = OS_COLOR('mb')
		self.mg = OS_COLOR('mg')
		self.my = OS_COLOR('my')
		self.mr = OS_COLOR('mr')
		self.ww = OS_COLOR('ww')
		self.ok = OS_COLOR('ok')
		self.xx = OS_COLOR('xx')
		self.colors = [ OS_COLOR(c) for c in sorted(self.color_dict.keys()) ]
		self.item = '  o '
		self.iitem = '      - '
		self.endl = OS_LINE_CHAR('\n')
		self.flush = OS_LINE_CHAR('')
		self.back = OS_LINE_CHAR('\r')
		self.trail = ' '
		# CURRENT STYLE SELECTION
		self.sel_color = None
		self.sel_justify = None
		self.sel_header = False
		self.sel_trim = "="
		# EXE ATTRIBUTES
		self.catch = OS_EXE_CATCH()
		self.assert_zero = OS_EXE_ASSERT()
		self.dev = OS_EXE_DEV('')
		self.nodev = OS_EXE_DEV('')
		self.devnull = OS_EXE_DEV()
		self.devfile = OS_EXE_DEV('')
		self.os_exe_get = False
		self.os_exe_assert_zero = False
		self.os_exe_dev = ''
		self.os_exe_verbose = False
		# LOGGING
		self.logfile = None
		# DIRECTORY HOPPING
		self.paths_visited = [ os.getcwd() ]
		self.exe_root_path = self.paths_visited[0]
		self.N_store_paths_visited = 1+5
	def __call__(self, mssg, c=None, j=None, h=False, t="="):
		# c=color, j=justify, h=header, t=trim, u=upper-case
		if j:
			mssg = self.justify_dict[j] + mssg
		if c != None:
			mssg = self.color_dict[c] + mssg + self.color_dict['endcolor']
		if h:
			mssg = self.os_generate_header(mssg, t)
		# print(mssg)
	# LOGFILE ADAPTOR =========================================================
	def ConnectToFile(self, logfile):
		self.logfile = logfile
		sys.stdout = open(logfile, 'w')
		self.devfile = OS_EXE_DEV(' >> {log} 2>> {log}'.format(log=logfile))
		return
	def DisconnectFromFile(self):
		if self.logfile != None:
			self.devfile = OS_EXE_DEV('')
			self.logfile = None
			sys.stdout = sys.__stdout__
		else: pass
		return
	# PRINTER METHODS =========================================================
	def __lshift__(self, mssg):
		if type(mssg) == OS_LINE_CHAR:
			# <LOG MESSAGE HERE>
			sys.stdout.write(str(mssg))
			sys.stdout.flush()
			self.sel_color = None
			return self
		elif type(mssg) == OS_COLOR:
			self.sel_color = str(mssg)
			return self
		mssg = str(mssg)
		if self.sel_justify != None:
			mssg = self.justify_dict[j] + mssg
		mssg += self.trail
		if self.sel_color != None:
			mssg = self.color_dict[self.sel_color] \
				+ mssg + self.color_dict['endcolor']
		if self.sel_header:
			mssg = self.os_generate_header(mssg, self.sel_trim)
		# <LOG MESSAGE HERE>
		sys.stdout.write(mssg)
		return self
	def os_print(self, mssg, c=None, j=None, h=False, t="="):
		# c=color, j=justify, h=header, t=trim, u=upper-case
		if j:
			mssg = OS_JUSTIFY_DICT[j] + mssg
		if c != None:
			mssg = OS_COLOR_DICT[c] + mssg + OS_COLOR_DICT['endcolor']
		if h:
			mssg = os_generate_header(mssg, t)
		# print(mssg)
		return
	def os_print_config(self, c=None, j=None, h=False, t="=", tl=' '):
		self.sel_color = c
		self.sel_justify = j
		self.sel_header = h
		self.sel_trim = t
		self.trail = tl
		return
	def os_print_reset(self):
		self.sel_color = None
		self.sel_justify = None
		self.sel_header = False
		self.sel_trim = "="
		self.trail = ' '
		return
	def os_generate_header(title, trim="="):
		try:
		    height, width = os.popen('stty size', 'r').read().split()
		    width = int(width)
		    leftright = ((width - len(title)-2)) // 2
		except ValueError:
		    leftright = 40
		return trim*leftright + " " + title + " " + trim*leftright
	# EXECUTE METHODS =========================================================
	def __rshift__(self, cmmd):
		if type(cmmd) == OS_EXE_CATCH:
			self.os_exe_get = True
			return self
		elif type(cmmd) == OS_EXE_DEV:
			self.dev = cmmd
			return self
		elif type(cmmd) == OS_EXE_ASSERT:
			self.os_exe_assert_zero = True
			return self
		# Redirect command as requested (not applicable if output is asked for)
		if not self.os_exe_get:
			if str(self.dev) != '':
				cmmd += str(self.dev)
				self.dev = self.nodev
			else:
				cmmd += str(self.devfile)
		# Execute
		if self.debug: self << self.my << "exe:" << cmmd << endl
		if self.os_exe_get:
			output = subprocess.getoutput(cmmd)
			self.os_exe_get = False
			return output
		else:
			sign = os.system(cmmd)
		if self.os_exe_assert_zero:
			if str(sign) != '0':
				raise RuntimeError("<OSIO> '%s' returned '%s'" % (cmmd, sign))
			self.os_exe_assert_zero = False
		return sign
	# EXIT PROTOCOLS ==========================================================
	def okquit(self, what=''):
		if what != '': self << self.ok << what << self.endl
		self.DisconnectFromFile()
		sys.exit(0)
		return
	def xxquit(self, what=''):
		if what != '': self << self.xx << "ERROR" << what << self.endl
		self.DisconnectFromFile()
		sys.exit(1)
		return
	# DIRECTORY HOPPING =======================================================
	def cd(self, d):
		# Current working directory, for archiving ... =>
		cwd = os.getcwd()
		if type(d) == int:
			# Change to previously visited path
			os.chdir(self.paths_visited[d])
		elif type(d) == str:
			# Change to path as specified explicitly
			os.chdir(d)
		else:
			raise NotImplementedError
		# <= ... previous path
		self.paths_visited.append(cwd)
		if len(self.paths_visited) > self.N_store_paths_visited:
			self.paths_visited.pop(1) # 0 stores root
		if self.debug: self << self.my << "cd: " << os.getcwd() << self.endl
		return
	def pwd(self):
		return self.cwd()
	def cwd(self):
		return os.getcwd()
	def root(self):
		self.cd(self.exe_root_path)
		return
	def abspath(self, file):
		if not os.path.exists(file):
			raise IOError("<osio::abspath> No such item in local directory: '%s'" % file)
		return os.path.join(self.cwd(), file)
	def mkcd(self, directory):
		self >> self.assert_zero >> 'mkdir -p %s' % directory
		self.cd(directory)
		return directory

class OS_EXE_DEV(object):
	def __init__(self, dev=' > /dev/null 2> /dev/null'):
		self.dev = dev
	def __str__(self):
		return self.dev

class OS_EXE_CATCH(object):
	def __init__(self):
		self.catch = True

class OS_EXE_ASSERT(object):
	def __init__(self):
		self.assert_0 = True

class OS_COLOR(object):
	def __init__(self, colstr):
		self.colstr = colstr
	def __str__(self):
		return self.colstr

class OS_LINE_CHAR(object):
	def __init__(self, char):
		self.char = char
	def __str__(self):
		return self.char


class OSIO(ShellInterface, OptionsInterface):
	def __init__(self):
		self.debug = False
		ShellInterface.__init__(self)
		OptionsInterface.__init__(self)
		return
	def sleep(self, dt):
		time.sleep(dt)
		return


osio = OSIO()
endl = OS_LINE_CHAR('\n')
flush = OS_LINE_CHAR('')
back = OS_LINE_CHAR('\r')
catch = OS_EXE_CATCH()
devnull = OS_EXE_DEV()

