import urllib2
import time
import os
#from ds1054z import DS1054Z
from rigol import*
import matplotlib.pyplot as plt
import numpy as np

home = os.getcwd()

noscp = False

IP = "http://192.168.254.2/" # ip adress of quido eth
osc_ip = '192.168.254.10' # ip adress of osciloscope
sc = 1 # relay for shortcutting the condenser, defualt is shortcutted
ch = 2 # relay for charging the condenser
trig = 6 # relay for simulating trigger
cur_stab = 7 # relay for switching curent stability, defualt is no stabilization
dch = 8 # # relay for conecting condenser to discharging circuits

u_coil = 1 # osciloscope channel with coil voltage

if not noscp:
    try:
        scope = DS1054Z(osc_ip)
    except:
        print('Unable to connect osciloscope. Check osciloscope IP adress, should be set to {}(static)'.format(osc_ip))

gdata = 'gte_data'
folder = os.path.join(home, gdata)
if not os.path.exists(folder):
    os.mkdir(folder)

def get_num( name = 'counter.txt' ):
    '''
    Loads shor number from file.
    Input:
    name = counter.txt -- string, name of file containign counter
    '''
    cnt = os.path.join(folder, name)
    if not os.path.exists(cnt):
        counter = open(cnt,'w')
        counter.write('1')
        counter.close()
        num = 1
    else:
        counter = open(cnt,'r')
        num = int(counter.read())
        counter.close()
    return num

def add_num(num, name = 'counter.txt' ):
    '''
    Adds 1 to num ands saves as shot number into counter file.
    Input:
    num -- int, number of previous shot
    name = counter.txt -- string, name of file containign counter
    '''
    cnt = os.path.join(folder, name)
    counter = open(cnt,'w')
    counter.write(str(num + 1))
    counter.close()
    return

def close_all(url = IP):
    """
    Switches all relays to defaul position, called 'off'.
    Input:
        url = IP -- ip adress of I/O module
    """
    urllib2.urlopen(url + "set.xml?type=R").read()
    return

def switch_on(num, url = IP):
    """
    Switches on relay numbered num. DO NOT use separately unless
    you are absolutely sure what will be the outcome.
    Input:
        num -- number of relay to be switched
        url = IP -- ip adress of I/O module
    """
    urllib2.urlopen( url + "set.xml?type=s&id={}".format(num) )
    return

def switch_off(num, url = IP):
    """
    Switches off relay numbered num. DO NOT use separately unless
    you are absolutely sure what will be the outcome.
    Input:
        num -- number of relay to be switched
        url = IP -- ip adress of I/O module
    """
    urllib2.urlopen(url + "set.xml?type=r&id={}".format(num))
    return

def input_read(num, url = IP):
    """
    Gets input messured on I/O module's input numbered num
    Input:
        num -- number of input
        url = IP -- ip adress of I/O module
    Output:
        val -- int, 1 or 0 representing logical values on or off
    """
    status = urllib2.urlopen(url + 'fresh.xml').read()
    inputs = [i for i in status.split("<") if i.startswith("din")]
    val = int(inputs[num-1].split("val=")[1][1])
    return val

def shot(wait = 0.1, chwait = 2, dchwait = 0.5, url = IP, stab_i = False, cstab = cur_stab, graphs = [], save = False, noscope = noscp):
    """
    Executes Golem tokamak like shot: charges and discharges condenser.
    Input:
        wait = 0.1 -- delay between charge and discharge in seconds
        chwait = 2 -- duration of charging in seconds
        dchwait = 0.5 -- duration of discharge in seconds
        url = IP -- ip adress of I/O module
        stab_i = False -- toggles current stabilisation
        cstab = cur_stab -- number of relay switching current stabilisation on
        graphs = [] -- numbers of channels on osciloscope to be ploted
        save = False -- enables saving graphs to file
        noscope = False -- disables osciloscope
    Output:
        none or graphs
    """
    fig = plt.figure()
    i = get_num()
    print('Shot number: {}'.format(i))
    try:
        charge(wait = chwait)
        time.sleep(wait)
        if not noscope:
            prepare_osc(scope)
        discharge(wait = dchwait, stab = stab_i)
        if graphs and not noscope:
            fig = make_graphs(graphs, save)
    except:
        close_all()
    add_num(i)
    plt.show()
    return

def charge(wait = 2, url = IP, shortcut = sc, charging = ch):
    """
    Charges condenser.
    Input:
        wait = 2 -- duration of charging in seconds
        url = IP -- ip adress of I/O module
        shortcut = sc -- number of relay switching shortcutting
        charging = ch -- number of relay switching charging
    """
    print('charging')
    try:
        close_all()
        switch_on(shortcut)
        time.sleep(0.01)
        switch_on(charging)
        time.sleep(wait)
        switch_off(charging)
    except:
        close_all()
    return

def discharge(wait = 0.5, stab = False, url = IP, disch = dch, cstab = cur_stab, delay = 0.1):
    """
    Discharges condenser.
    Input:
        wait = 0.5 -- delay between opening discharging relay and shortcutting condenser
        stab_i = False -- toggles current stabilisation
        url = IP -- ip adress of I/O module
        disch = dch -- number of relay switching discharging
        thyr = th --number of relay switching thyristor
    """
    try:
        if stab:
            switch_on(cstab)
            time.sleep(delay)
        switch_on(disch)
        time.sleep(2)
        switch_on(trig)
        switch_off(trig)
        print('discharging')
        time.sleep(wait)
        close_all()
    except:
        close_all()
    return

def make_graphs(graphs, save = False):
    """
    Draws graphs from osciloskope data on selected channels specified in input graphs
    Input:
        graphs -- numbers of channels to plot
        save -- enables saving of ploted figures
    """
    print("plotting")
    if isinstance(graphs, int):
        graphs = [graphs]
    for channel in graphs:
        time, data = get_data(scope, channel)
        #print('data loaded')
        #fig = plt.figure()
        plt.plot(time,data)
        if save:
            try:
                i = get_num()
                picname = os.path.join(gdata, 'shot_{}_channel{}.pdf'.format(i,channel))
                dataname = os.path.join(gdata, 'shot_{}_channel{}.npz'.format(i,channel))
                if not os.path.exists(picname):
                    plt.savefig()
                if not os.path.exists(dataname):
                    np.savez(dataname, data = data, time = time)
            except:
                print('saving error')
    return