import sys
import os
import re
import numpy as np


class HydTableData():

	def __init__(self):
		self.displayName = ''
		self.tcf = ''
		self.xsNos = []
		self.xsNames = []
		self.xsSources = []
		self.xsHydTa_headers = []
		self.xsHydTa = []
		self.xsSections = []
		self.xsSection_headers = []
		self.xsBeds= []
		self.channelNames = []
		self.channelXs = []
		self.channelIntp = []
		self.channelHydTa = []
		self.channelHydTa_headers = []
		self.loaded = False
		
	def load(self, fname):
		with open(fname, 'r') as fopen:
			xsHydTa_header = False  # Header for XS Hydraulic Tables has not been populated yet
			channelHydTa_header = False  # Header for Channel Hydraulic Tables has not been populated yet
			for line in fopen:
				if 'generated by' in line.lower():
					self.tcf = line.split('"')[1].strip()
					# self.displayName = os.path.basename(os.path.splitext(self.tcf)[0])
					self.displayName = re.sub(r"(_1d_ta_tables_check.csv)$", "", os.path.basename(fname),
					            flags=re.IGNORECASE)  # displayname
					self.displayName  = '{0}_1d_ta'.format(self.displayName)
				elif 'section' in line.lower():
					a = re.split('[\[\] ]', line)  # split by [ ] and space
					self.xsNos.append(a[1])
					self.xsSources.append(a[-2])
					self.xsNames.append(os.path.basename(a[-2]))
					header = True
					data = []
					xsData = []
					bed = []
					xsSection = False
					for blockLine in fopen:
						if blockLine == '\n':
							pass
						# elif header and not xsHydTa_header:
						elif header:
							if 'section' in blockLine.lower():
								continue
							else:
								a = blockLine.split(",")
								s1 = a.index('"Elevation"')  # get index of first occurrence
								s2 = len(a) - a[::-1].index('"Elevation"') - 1  # get index of last occurrence
								if s1 != s2:
									xsSection = True
									e1 = a.index('""')  # get index of blank column
									#for header in a[:e1]:
									#	self.xsSection_headers.append(header.strip().strip('"'))
									self.xsSection_headers.append([x.strip().strip('"') for x in a[:e1]])
								headerLen = len(a)
								#for i, header in enumerate(a[s2:]):
								#	self.xsHydTa_headers.append(header.strip().strip('"'))
								self.xsHydTa_headers.append([x.strip().strip('"') for x in a[s2:]])
								xsHydTa_header = True
								header = False
								continue
						# elif header:
						# 	if 'section' in blockLine.lower():
						# 		continue
						# 	else:
						# 		a = blockLine.split(",")
						# 		s1 = a.index('"Elevation"')  # get index of first occurrence
						# 		s2 = len(a) - a[::-1].index('"Elevation"') - 1  # get index of last occurrence
						# 		if s1 != s2:
						# 			xsSection = True
						# 			e1 = a.index('""')  # get index of blank column
						# 		headerLen = len(a)
						# 		header = False
						# 		continue
						else:
							a = blockLine.split(",")
							b = []
							c = []
							if xsSection:
								for i, dat in enumerate(a[:e1]):
									if i == 0:
										if dat.strip().strip('"').lower() == 'bed':
											bed.append([a[1].strip(), a[2].strip()])
									elif i > 0:
										if dat.strip().strip('""') != "" and dat != '\n':
											c.append(float(dat.strip()))
										elif dat == "":
											c.append(np.nan)
								xsData.append(c)
							if len(a) >= headerLen:
								for i, dat in enumerate(a[s2:]):
									if i < headerLen:
										if dat.strip().strip('"') != "" and dat != '\n':
											b.append(float(dat.strip()))
										elif dat == "":
											b.append(np.nan)
								data.append(b)
							continue
						array = np.array(data)
						xsArray = np.array(xsData)
						self.xsHydTa.append(array)
						self.xsSections.append(xsArray)
						self.xsBeds.append(bed)
						break  # break data block loop
				elif 'channel' in line.lower():
					a = re.split('[\(\) ]', line)  # split by ( ) and space
					self.channelNames.append(a[1])
					xs = []
					intp = []
					for i, h in enumerate(a):
						if h == "":
							continue
						if h[0:2] == 'XS':
							if len(h) < 3:
								h = h + a[i+1].strip()
							xs.append(h)
						elif h[-1] == '%':
							intp.append(h)
					self.channelXs.append(xs)
					self.channelIntp.append(intp)
					header = True
					data = []
					for blockLine in fopen:
						if blockLine == '\n':
							pass
						# elif header and not channelHydTa_header:
						elif header:
							if 'channel' in blockLine.lower():
								continue
							else:
								a = blockLine.split(",")
								try:
									s = len(a) - a[::-1].index('"Elevation"') - 1  # get index of last occurrence
								except:
									try:
										s = len(a) - a[::-1].index('Elevation') - 1
									except:
										s = 0
								headerLen = len(a[s:])
								#for i, header in enumerate(a[s:]):
								#	self.channelHydTa_headers.append(header.strip().strip('"'))
								self.channelHydTa_headers.append([x.strip().strip('"') for x in a[s:]])
								channelHydTa_header = True
								header = False
								continue
						# elif header:
						# 	if 'channel' in blockLine.lower():
						# 		continue
						# 	else:
						# 		a = blockLine.split(",")
						# 		try:
						# 			s = len(a) - a[::-1].index('"Elevation"') - 1  # get index of last occurrence
						# 		except:
						# 			try:
						# 				s = len(a) - a[::-1].index('Elevation') - 1
						# 			except:
						# 				s = 0
						# 		header = False
						# 		continue
						else:
							a = blockLine.split(",")
							b = []
							if len(a[s:]) < headerLen:
								continue
							for i, dat in enumerate(a[s:]):
								if i < headerLen:
									if dat.strip().strip('"') != "" and dat != '\n':
										b.append(float(dat.strip()))
									elif dat == "":
										b.append(np.nan)
							data.append(b)
							continue
						array = np.array(data)
						self.channelHydTa.append(array)
						break  # break data block loop
			if len(self.channelNames) > len(self.channelHydTa):
				array = np.array(data)
				self.channelHydTa.append(array)

		if self.xsSection_headers or self.xsHydTa_headers or self.channelHydTa_headers:
			self.loaded = True

	def getAllTypes(self):
		"""

		"""

		xsTaTypes = set(sum(self.xsHydTa_headers, []))
		if 'Elevation' in xsTaTypes:
			xsTaTypes.remove('Elevation')
		chTaTypes = set(sum(self.channelHydTa_headers, []))
		if 'Elevation' in chTaTypes:
			chTaTypes.remove('Elevation')

		return set(list(xsTaTypes) + list(chTaTypes))

	def getAllIdsForType(self, t):
		"""

		"""

		xsNames = [x for i, x in enumerate(self.xsNames) if t in self.xsHydTa_headers[i]]
		chNames = [x for i, x in enumerate(self.channelNames) if t in self.channelHydTa_headers[i]]

		return xsNames + chNames

	def plotProperty(self, id, prop, dom=''):
		"""

		"""

		x, y = [], []
		if dom.lower() == 'xs':
			ids = [x.lower() for x in self.xsNames]
		elif dom.lower() == 'ch':
			ids = [x.lower() for x in self.channelNames]
		else:
			ids = [x.lower() for x in self.xsNames + self.channelNames]

		if id.lower() in ids:
			i = ids.index(id)
			if dom.lower() == 'xs':
				props = self.xsHydTa_headers[:]
				data = self.xsHydTa[:]
			elif dom.lower() == 'ch':
				props = self.channelHydTa_headers[:]
				data = self.channelHydTa[:]
			elif i <= len(self.xsNames) - 1:
				props = self.xsHydTa_headers[:]
				data = self.xsHydTa[:]
			else:
				i -= len(self.xsNames) - 1
				props = self.channelHydTa_headers[:]
				data = self.channelHydTa[:]
			propH = [x.lower() for x in props[i]]
			if prop.lower() in propH:
				j = propH.index(prop.lower())
				x = data[i][:,j]

				if 'elevation' in propH:
					j = propH.index('elevation')
					y = data[i][:,j]

		return x, y

class HydTables():
	
	def __init__(self):
		self.loadedData = []
		
	def loadData(self, inFile):
		hydTab = HydTableData()
		hydTab.load(inFile)
		if hydTab.loaded:
			self.loadedData.append(hydTab)
			return hydTab, ''
		else:
			return None, 'Error loading 1D Hydraulic Tables - Check file is a _1d_ta_tables_check.csv and not empty'
		
	def closeData(self, displayName):
		for i, data in enumerate(self.loadedData):
			if data.displayName == displayName:
				del(self.loadedData[i])

	def getData(self, displayName):
		loadedDisplayNames = [x.displayName.lower() for x in self.loadedData]
		if displayName.lower() in loadedDisplayNames:
			i = loadedDisplayNames.index(displayName.lower())
			return self.loadedData[i]