#!/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import io, os
import argparse

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)

psu_tty_dev = find_tty_dev('KORAD')

if psu_tty_dev != '':
    print('found PSU TTY device: ' + psu_tty_dev)
else:
    print('couldn\'t find PSU TTY device')
    exit()

psu_tty = io.TextIOWrapper(io.FileIO(os.open(psu_tty_dev, os.O_NOCTTY | os.O_RDWR), "r+"),
                           line_buffering=True, write_through=True)

def psu_read(cmd):
    psu_tty.write(cmd + '\n')
    psu_tmp_line = psu_tty.readline()
    return(psu_tmp_line.strip())

def psu_write(cmd):
    psu_tty.write(cmd + '\n')

def psu_get_id(val):
    print('PSU ID: ' + 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(val):
    return(float(psu_read('VOUT1?')))

def psu_get_i1(val):
    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(val):
    return(float(psu_read('VOUT2?')))

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


def psu_turn_off(val):
    psu_set_v1(0)
    psu_set_v2(0)

def psu_on_default(val):
    psu_set_v1(12.0)
    psu_set_v2(5.0)


parser = argparse.ArgumentParser(description='Do something with the PSU.')
parser.add_argument('action', choices=['getID', 'setV1', 'setV2', 'setI1', 'setI2', 'getV1', 'getV2', 'getI1', 'getI2', 'off', 'on_default'],
                    help='what to do.')
parser.add_argument('value', type=float,
                    help='value to set something to (set to 0 if not setting something).')


fn_switch = {
    'getID': psu_get_id,
    'setV1': psu_set_v1,
    'setV2': psu_set_v2,
    'setI1': psu_set_i1,
    'setI2': psu_set_i2,
    'getV1': psu_get_v1,
    'getV2': psu_get_v2,
    'getI1': psu_get_i1,
    'getI2': psu_get_i2,
    'off':   psu_turn_off,
    'on_default':    psu_on_default}


args = parser.parse_args()

fn_switch.get(args.action)(args.value)

