#!/usr/bin/python3

import time
import serial
import sys
import numpy as np
from contextlib import contextmanager


class HendiCtrl:
	def __init__(self, port, speed):
		self.ser = serial.Serial(port, speed)
		try:
			self.ser.open()
		except:
			self.ser.close()
			self.ser.open()

		self.ser.timeout = 1.000
		self.prompt = b':'

		powers = [500,  600,  700,  800,  900, 1000, 1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100, 3200, 3300,3400, 3500]
		digits = [4092, 3840, 3700, 3440, 3300, 3060, 3000, 2930, 2860, 2780, 2640, 2510, 2430, 2360, 2190, 2130, 2030, 1980, 1880, 1730, 1630, 1560, 1480, 1400, 1240, 1150, 1150, 1080, 920, 740, 670]
		self.caps = {
			"pwr_watts_min"		: powers[0],
			"pwr_watts_max"		: powers[-1],
			"pwr_list"			: powers,
			"pwr_watts_step"	: 100,
			"pwr_digits_min"	: digits[0],
			"pwr_digits_max"	: digits[-1],
			"digit_list"		: digits
		}
		self.poly_w2d = np.polyfit(powers, digits, 5)
		self.poly_d2w = np.polyfit(digits, powers, 5)
		self.sw_id = self.getSoftwareIdentifier()
		self.sw_ver = self.getSoftwareVersion()

	@contextmanager
	def remote_open(self):
		try:
			self.remoteEnable(True)
			yield None
		finally:
			self.remoteEnable(False)

	def toDigits(self, power_watts):
		return np.polyval(self.poly_w2d, float(power_watts))

	def toWatts(self, power_digits):
		return np.polyval(self.poly_d2w, float(power_digits))

	def reset(self):
		self.ser.dtr = True
		self.ser.rts = False
		time.sleep(0.01)
		self.ser.dtr = True
		self.ser.rts = True

	def enter_bootloader(self):
		self.ser.dtr = True
		self.ser.rts = False
		time.sleep(0.01)
		self.ser.dtr = False
		self.ser.rts = True

	def firmware_update(self, filename):
		print("Start firmware update")
		self.enter_bootloader()
		self.ser.readline()

		self.ser.flushInput()
		with open(filename, 'r') as fp:
			line_count = 0
			while (True):
				s = fp.readline()
				if s == '':
					break
				data = s.encode()
				self.ser.write(data)
				ack = self.ser.read()
				sys.stdout.write(ack.decode('utf-8'))
				sys.stdout.flush()
				if line_count == 64:
					line_count = 0
					sys.stdout.write("\n")

				line_count += 1

		self.ser.readline()
		self.ser.flushInput()
		self.ser.write(b'q')
		self.ser.flushOutput()

	def __write(self, s):
		self.ser.flushInput()
		data = s.encode()
		self.ser.write(data + b'\r')

	def __read(self):
		# Read line echo
		echo = self.ser.readline()

		# Read answer line
		answer = self.ser.readline().decode('utf-8').replace('\r', '').replace('\n', '')
		if ':' in answer:
			result = answer.split(':')
		else:
			result = answer

		return result

	def cmd(self, req):
		self.__write(req)
		rsp = self.__read()
		if "OK" not in rsp[0]:
			raise Exception(rsp[1])

		return rsp[1]
		
	def getCapabilties(self):
		return self.caps

	def getInfo(self):
		result = {
			"sw_id"			: self.getSoftwareIdentifier(),
			"sw_ver"		: self.getSoftwareVersion(),
			'state'			: self.getState(),
			'pwr_digits'	: self.getPowerDigits(),
			'Switch_state'	: self.getSwitch()
		}
		return result

	def _debugEnable(self, enable):
		return self.cmd("D" + str(int(enable)))

	def _isDebugEnable(self):
		return self.cmd("D?")

	def getSoftwareIdentifier(self):
		return self.cmd("I?")

	def getSoftwareVersion(self):
		return self.cmd("V?")

	def getState(self):
		return self.cmd("T?")

	def remoteEnable(self, enable):
		return self.cmd("R" + str(int(enable)))

	def isRemoteEnable(self):
		return '1' in self.cmd("R?")

	def setPowerDigits(self, power_digits):
		return self.cmd("P" + str(power_digits))

	def getPowerDigits(self):
		return self.cmd("P?")

	def setSwitch(self, enable):
		return self.cmd("S" + str(int(enable)))

	def getSwitch(self):
		return self.cmd("S?")

	def setPowerWatts(self, power_watts):
		pwr_digits_min = self.caps['pwr_digits_min']
		pwr_digits_max = self.caps['pwr_digits_max']
		max_digits = max(pwr_digits_max, pwr_digits_min)
		min_digits = min(pwr_digits_max, pwr_digits_min)
		pwr_digits = int(self.toDigits(power_watts))
		pwr_digits = int(min(max_digits, max(min_digits, pwr_digits)))
		return self.setPowerDigits(pwr_digits)

	def getPowerWatts(self):
		return int(self.toWatts(self.getPowerDigits()) + 0.5)


if __name__ == '__main__':

	hendi = HendiCtrl('/dev/ttyUSB0', 115200)

#	hendi.firmware_update("HendiCtrl.srec")
#	hendi.reset()
	time.sleep(1)

	hendi.getState()
	hendi.getSoftwareIdentifier()
	hendi.getSoftwareVersion()
	caps = hendi.getCapabilties()
	print("Caps: ", caps)
	info = hendi.getInfo()
	print("Info: ", info)

	print("Test converters")
	for pwr in range(500, 3600, 100):
		print("pwr", pwr)
		digits = hendi.toDigits(pwr)
		print("digits", digits)
		power = hendi.toWatts(digits)
		print("power", power)

	print("Test Hendi")
	with hendi.remote_open():
		hendi.setSwitch(1)
		time.sleep(1.0)
		for pwr in range(500, 3600, 100):
			hendi.setPowerWatts(pwr)
			pw = hendi.getPowerWatts()
			print("power_w", pw)
			time.sleep(1.0)

		hendi.setSwitch(0)

	if 0:
		hendi.remoteEnable(True)
		if hendi.isRemoteEnable():
			hendi.setSwitch(True)
			for p in range(0, 4096, 100):
				hendi.setPowerDigits(p)
				time.sleep(1)

			hendi.setSwitch(False)
			hendi.remoteEnable(False)

	print("End of program.")
