#!/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import io, os
import serial
import argparse
import time
import math
from datetime import datetime
from datetime import timedelta
from tabulate import tabulate
from apscheduler.schedulers.background import BackgroundScheduler

def find_tty_dev(dev_name):
    tty_enum = os.popen('./get_tty_devs.sh').read().strip().split('\n')
    # only keep the last entry that way
    tty_dev = ''
    for a in tty_enum:
        if a.find(dev_name) > 0:
            tty_dev = '/dev/' + a.split(':')[0]
    return(tty_dev)


def find_psu_tty():
    global psu_tty_dev
    psu_tty_dev = find_tty_dev('KORAD')
    if psu_tty_dev != '':
        print('found the PSU TTY device: ' + psu_tty_dev)
    else:
        print('couldn\'t find the PSU TTY device')
        exit()

def find_mm_tty():
    global mm_tty_dev
    mm_tty_dev = find_tty_dev('DIOT monitoring module (monimod)')
    if mm_tty_dev != '':
        print('found the MoniMod debug bridge TTY device: ' + mm_tty_dev)
    else:
        print('couldn\'t find the MoniMod debug bridge TTY device')
        exit()

SER_DEFAULT_IMPL = 0

def init_mm_tty():
    global mm_tty
    if SER_DEFAULT_IMPL:
        mm_tty = io.TextIOWrapper(io.FileIO(os.open(mm_tty_dev, os.O_NOCTTY | os.O_RDWR), "r+"),
                                  line_buffering=True, write_through=True)
    else:
        mm_tty = serial.Serial(mm_tty_dev, timeout=0.1)
        mm_tty.close()
        mm_tty.open()
        if mm_tty.isOpen():
            print(mm_tty.portstr, "PSU TTY connection successful.")
            return(0)
        else:
            return(-1)

def init_psu_tty():
    global psu_tty
    if SER_DEFAULT_IMPL:
        psu_tty = io.TextIOWrapper(io.FileIO(os.open(psu_tty_dev, os.O_NOCTTY | os.O_RDWR), "r+"),
                                   line_buffering=True, write_through=True)
    else:
        psu_tty = serial.Serial(psu_tty_dev, 115200)
        psu_tty.close()
        psu_tty.open()
        if psu_tty.isOpen():
            print(psu_tty.portstr, "PSU TTY connection successful.")
            return(0)
        else:
            return(-1)

def reset_mm_controller():
    print('resetting the MM controller...')
    mm_write('S')
    mm_tty.close()
    for i in range(5):
        try:
            print('Trying to reopen the test controller tty...')
            time.sleep(2)
            if SER_DEFAULT_IMPL:
                init_mm_tty()
            else:
                mm_tty.open()
            print('Enabling test mode...')
            mm_write('L')
            time.sleep(0.1)
            print('Enabling PEC...')
            mm_write('D')
            time.sleep(0.1)
            mm_write('E')
            time.sleep(0.1)
            mm_write('d')
            time.sleep(0.1)
            mm_write('E')
            time.sleep(0.1)
            print('Controller successfully reset...')
            break
        except:
            print('Failed to reopen the test controller tty...')

def psu_read(cmd):
    if SER_DEFAULT_IMPL:
        psu_tty.write(cmd + '\n')
        psu_tty.flush()
        time.sleep(0.2)
        psu_tmp_line = psu_tty.readline()
        return(psu_tmp_line.strip())
    else:
        psu_tty.write((cmd + '\n').encode('utf-8'))
        psu_tmp_line = psu_tty.readline()
        return(psu_tmp_line.decode('utf-8').strip())

def psu_write(cmd):
    if SER_DEFAULT_IMPL:
        psu_tty.write(cmd + '\n')
    else:
        psu_tty.write((cmd + '\n').encode('utf-8'))


def mm_write(cmd):
    if SER_DEFAULT_IMPL:
        mm_tty.write(cmd)
    else:
        mm_tty.write(cmd.encode('utf-8'))

def mm_read(cmd, timeout=1.5):
    if SER_DEFAULT_IMPL:
        mm_tty.write(cmd)
        # while not mm_tty.in_waiting:
        #     time.sleep(0.01)
        time.sleep(timeout)
        mm_tmp_line = mm_tty.read(256).strip()
        mm_tmp_line = mm_tmp_line.split('\n')
        mm_tmp_line[0] = mm_tmp_line[0].strip()
        mm_tmp_line[1] = mm_tmp_line[1].strip()
        return(mm_tmp_line)
    else:
        mm_tty.write(cmd.encode('utf-8'))
        # while not mm_tty.in_waiting:
        #     time.sleep(0.01)
        time.sleep(timeout)
        mm_tmp_line = mm_tty.read(256).decode('utf-8').strip().split('\n')
        return(mm_tmp_line)


def psu_read_id():
    return(psu_read('*IDN?'))


def psu_set_v1(val):
    if ((val >= 0.0) and (val < 13.2)):
        psu_write('VSET1:{:05.2f}'.format(val))

def psu_set_i1(val):
    if ((val >= 0.0) and (val < 5)):
        psu_write('ISET1:{:05.2f}'.format(val))

def psu_get_v1():
    return(float(psu_read('VOUT1?')))

def psu_get_i1():
    return(float(psu_read('IOUT1?')))


def psu_set_v2(val):
    if ((val >= 0.0) and (val < 5.5)):
        psu_write('VSET2:{:05.2f}'.format(val))

def psu_set_i2(val):
    if ((val >= 0.0) and (val < 5)):
        psu_write('ISET2:{:05.2f}'.format(val))

def psu_get_v2():
    return(float(psu_read('VOUT2?')))

def psu_get_i2():
    return(float(psu_read('IOUT2?')))

def mm_select(val):
    if (val == 0):
        mm_write('d')
    else:
        mm_write('D')

def mm_set_fan(fan, val):
    mm_write('s')
    mm_write('{:01d}'.format(fan))
    for a in '{:04d}'.format(int(val)):
        mm_write(a)


psu_tmp_v1 = 12.0
psu_tmp_v2 = 5.0

def psu_turn_off():
    global psu_tmp_v1
    global psu_tmp_v2
    psu_tmp_v1 = psu_get_v1()
    psu_tmp_v2 = psu_get_v2()
    psu_set_v1(0)
    psu_set_v2(0)

def psu_turn_back_on():
    global psu_tmp_v1
    global psu_tmp_v2
    psu_set_v1(psu_tmp_v1)
    psu_set_v2(psu_tmp_v2)


def check_range(low, high, arg):
    try:
        value = int(arg)
    except ValueError as err:
        raise argparse.ArgumentTypeError(str(err))

    if value < low or value > high:
        message = "Expected {} <= value <= {}, got value = {}".format(low, high, value)
        raise argparse.ArgumentTypeError(message)

    return value

check_range_2s_30s = lambda x: check_range(2, 30, x)
check_range_2s_1h  = lambda x: check_range(2, 3600, x)
check_range_0_1h   = lambda x: check_range(0, 3600, x)
check_range_0p_10p = lambda x: check_range(0, 10, x)
check_range_0p_50p = lambda x: check_range(0, 50, x)
check_range_0_1000 = lambda x: check_range(0, 1000, x)


parser = argparse.ArgumentParser(description='Control the MoniMod test setup and log the results.')
parser.add_argument('--log-interval', type=check_range_2s_30s, default=5,
                    help='Logging interval (2s-30s, default=5).')

parser.add_argument('--v1-period', type=check_range_2s_1h, default=600,
                    help='12V fluctuation period (2s-3600s, default=600).')
parser.add_argument('--v1-range', type=check_range_0p_10p, default=10,
                    help='12V +- fluctuation percent (0%%-10%%, default=10%%).')

parser.add_argument('--v2-period', type=check_range_2s_1h, default=600,
                    help='5V fluctuation period (2s-3600s, default=600).')
parser.add_argument('--v2-range', type=check_range_0p_10p, default=10,
                    help='5V +- fluctuation percent (0%%-10%%, default=10%%).')

parser.add_argument('--f1-setpoint', type=check_range_0_1000, default=500,
                    help='Fan 1 setpoint (0-1000, default=500).')
parser.add_argument('--f1-period', type=check_range_2s_1h, default=300,
                    help='Fan 1 fluctuation period (2s-3600s, default=300).')
parser.add_argument('--f1-range', type=check_range_0p_50p, default=10,
                    help='Fan 1 +- fluctuation percent (0%%-50%%, default=10%%).')

parser.add_argument('--f2-setpoint', type=check_range_0_1000, default=400,
                    help='Fan 2 setpoint (0-1000, default=400).')
parser.add_argument('--f2-period', type=check_range_2s_1h, default=300,
                    help='Fan 2 fluctuation period (2s-3600s, default=300).')
parser.add_argument('--f2-range', type=check_range_0p_50p, default=10,
                    help='Fan 2 +- fluctuation percent (0%%-50%%, default=10%%).')

parser.add_argument('--max-cons-errors', type=int, default=2,
                    help='Maximum consecutive errors before power cycling the MoniMods (default is 2).')
parser.add_argument('--power-cycle-period', type=check_range_0_1h, default=600,
                    help='Forcibly power cycle every POWER_CYCLE_PERIOD seconds (0-3600s, 0 for off, default is 600).')

parser.add_argument('--psu-log-file', type=argparse.FileType('w'),
                    default=open('log_' + datetime.utcnow().strftime("%d.%m.%Y_%H:%M") + '_psu.txt', 'w'),
                    help='PSU log file.')

parser.add_argument('--mm-log-file', type=argparse.FileType('w'),
                    default=open('log_' + datetime.utcnow().strftime("%d.%m.%Y_%H:%M") + '_mm.txt', 'w'),
                    help='MoniMod log file.')

parser.add_argument('--dry-run', action='store_true',
                    help='Don\'t write anything to the PSU or MoniMods.')

type=argparse.FileType('w')
args = parser.parse_args()

find_psu_tty()
init_psu_tty()
find_mm_tty()
init_mm_tty()

print('Device code: ' + psu_read_id())
print()

psu_data_headers = ['Rail', 'Voltage [V]', 'Current [mA]']
psu_data = [['P12V', 0.0, 0.0],
            ['P5V',  0.0, 0.0]]

def log_psu(f = None):
    global psu_data

    psu_data[0][1] = psu_get_v1()
    psu_data[0][2] = 1000.0*psu_get_i1()
    psu_data[1][1] = psu_get_v2()
    psu_data[1][2] = 1000.0*psu_get_i2()

    if f:
        f.write(datetime.utcnow().strftime("%d.%m.%Y %H:%M:%S.%f") + ',')
        f.write('{:.2f},'.format(psu_data[0][1]))
        f.write('{:.2f},'.format(psu_data[0][2]))
        f.write('{:.2f},'.format(psu_data[1][1]))
        f.write('{:.2f}\n'.format(psu_data[1][2]))
        f.flush()


mm_data_headers = ['Value', 'MM 0', 'MM 1']
mm_data = [['N', 0, 0],
           ['TMR_ERR', 0, 0],
           ['I2C_RES', 0, 0],
           ['V1', 0.0, 0.0],
           ['V2', 0.0, 0.0],
           ['V3', 0.0, 0.0],
           ['I1', 0.0, 0.0],
           ['I2', 0.0, 0.0],
           ['I3', 0.0, 0.0],
           ['T1', 0.0, 0.0],
           ['T2', 0.0, 0.0],
           ['T3', 0.0, 0.0],
           ['F1', 0.0, 0.0],
           ['F2', 0.0, 0.0],
           ['F3', 0.0, 0.0]]

mm_error_msgs = []              # messages to print after the table
mm_i2c_err_cnt = 0              # error cnt
mm_i2c_cons_err_cnt = 0         # consecutive error cnt

def log_mm(f = None, pcyc = False):
    global mm_data
    global mm_error_msgs
    global mm_i2c_err_cnt
    global mm_i2c_cons_err_cnt

    if pcyc:
        reset_mm_controller()
        mm_ret = mm_read('t', 3.5)
        mm_error_msgs += ['Power cycled']
    else:
        mm_ret = mm_read('a')

    if len(mm_ret) > 1:
        mm_ret[0] = mm_ret[0].split('\t')
        mm_ret[1] = mm_ret[1].split('\t')

        if (mm_data[1][1] != int(mm_ret[0][2])) or (mm_data[1][2] != int(mm_ret[1][2])):
            mm_error_msgs += ['TMR error count changed ({} != {}, {} != {})'.format(mm_data[1][1], int(mm_ret[0][2]), mm_data[1][2], int(mm_ret[1][2]))]

        if (int(mm_ret[0][1]) != 0) or (int(mm_ret[1][1]) != 1):
            mm_error_msgs += ['MM misalignment in test controller response ({}, {})'.format(mm_ret[0][1], mm_ret[1][1])]

        for i in range(2):
            mm_data[0][i+1] = int(mm_ret[i][0])
            mm_data[1][i+1] = int(mm_ret[i][2])
            mm_data[2][i+1] = int(mm_ret[i][3])
            for j in range(12):
                mm_data[j+3][i+1] = float(mm_ret[i][j+4])
    else:                       # bad reply
        for i in range(2):
            mm_data[0][i+1] = -1
            mm_data[1][i+1] = -1
            mm_data[2][i+1] = -1
            for j in range(12):
                mm_data[j+3][i+1] = -1.0

    if f:
        for i in range(0, 2):
            f.write(datetime.utcnow().strftime("%d.%m.%Y %H:%M:%S.%f") + ',')
            f.write('{:d},'.format(i))
            for j in range(0, 3):
                f.write('{:d},'.format(mm_data[j][i+1]))
            for j in range(3, 14):
                f.write('{:.3f},'.format(mm_data[j][i+1]))
            f.write('{:.3f}\n'.format(mm_data[14][i+1]))
        f.flush()

    if (mm_data[2][1] != 0) or (mm_data[2][2] != 0):
        mm_i2c_err_cnt += 1
        mm_i2c_cons_err_cnt += 1
        mm_error_msgs += ['I2C error ({} and {})'.format(mm_data[2][1], mm_data[2][2])]
    else:
        mm_i2c_cons_err_cnt = 0

    if (mm_data[0][1] % 2 != 0) or (mm_data[0][2] % 2 != 1):
        mm_error_msgs += ['N not even for 1st MM or not odd for 2nd MM']


cnt = 0
secs = 0
pc_pos = 0
v1_pos = 0
v2_pos = 0
f1_pos = 0
f2_pos = 0

def main_loop():
    global args
    global cnt
    global secs
    global pc_pos
    global v1_pos
    global v2_pos
    global f1_pos
    global f2_pos
    global psu_data
    global mm_data
    global mm_error_msgs
    global mm_i2c_err_cnt
    global mm_i2c_cons_err_cnt

    mm_error_msgs = []

    # skip first iteration
    if cnt:
        log_psu(args.psu_log_file)
        log_mm(args.mm_log_file)

    if mm_i2c_cons_err_cnt > args.max_cons_errors:
        print('too many errors, power cycling...')
        log_mm(args.mm_log_file, pcyc=True)
    elif (args.power_cycle_period) and (pc_pos + args.log_interval >= args.power_cycle_period):
        print('power cycling as per the schedule...')
        log_mm(args.mm_log_file, pcyc=True)

    pc_pos = secs % args.power_cycle_period

    v1_pos = secs % args.v1_period
    v2_pos = secs % args.v2_period
    f1_pos = secs % args.f1_period
    f2_pos = secs % args.f2_period

    v1_val = 12.0*(1.0 + 0.01 * float(args.v1_range) * math.sin(2.0 * math.pi * float(v1_pos) / float(args.v1_period)))
    v2_val = 5.0*(1.0 + 0.01 * float(args.v2_range) * math.sin(2.0 * math.pi * float(v2_pos) / float(args.v2_period)))
    f1_val = float(args.f1_setpoint)*(1.0 + 0.01 * float(args.f1_range) * math.sin(2.0 * math.pi * float(f1_pos) / float(args.f1_period)))
    f2_val = float(args.f2_setpoint)*(1.0 + 0.01 * float(args.f2_range) * math.sin(2.0 * math.pi * float(f2_pos) / float(args.f2_period)))

    if (not args.dry_run):
        psu_set_v1(v1_val)
    if (not args.dry_run):
        psu_set_v2(v2_val)

    if (not args.dry_run):
        mm_select(0)
        time.sleep(0.1)
        mm_set_fan(0, f1_val)
        time.sleep(0.1)
        mm_set_fan(1, f2_val)
        time.sleep(0.1)

        mm_select(1)
        time.sleep(0.1)
        mm_set_fan(0, f1_val)
        time.sleep(0.1)
        mm_set_fan(1, f2_val)
        time.sleep(0.1)

    os.system('clear')

    if (args.dry_run):
        print("Dry run mode")
        print()

    print('Time elapsed: {}'.format(str(timedelta(seconds=secs))))
    print()

    print(tabulate(psu_data, psu_data_headers, floatfmt=".2f"), flush=True)

    print(tabulate(mm_data, mm_data_headers, floatfmt=".2f"), flush=True)

    print()
    print('---------------------------')
    print('V1 has been set to = {:.2f}'.format(v1_val))
    print('V2 has been set to = {:.2f}'.format(v2_val))
    print('F1 has been set to = {:d}'.format(int(f1_val)))
    print('F2 has been set to = {:d}'.format(int(f2_val)))
    print('---------------------------')
    print('total I2C errors so far: {}'.format(mm_i2c_err_cnt))
    print('consecutive I2C errors: {}'.format(mm_i2c_cons_err_cnt))
    print('power cycle at {} consecutive errors'.format(args.max_cons_errors + 1))
    if args.power_cycle_period:
        print('scheduled power cycle in {}'.format(str(timedelta(seconds=(args.power_cycle_period-pc_pos)))))
    else:
        print('scheduled power cycles disabled')
    print('---------------------------')
    print('other errors:')
    for a in mm_error_msgs:
        print(a)

    cnt = cnt + 1
    secs = secs + args.log_interval


def init():
    psu_turn_off()
    time.sleep(2)
    psu_turn_back_on()
    time.sleep(1.0)
    reset_mm_controller()
    time.sleep(1.0)
    mm_read('t', 3.5)
    time.sleep(1.0)

init()

sched = BackgroundScheduler()
sched.add_job(main_loop, 'interval', seconds=args.log_interval)
sched.start()

while True:
    None

sched.shutdown()
