'''
Authors: Ing. Kateřina Hromasová, Jakub Svoboda
katerina.hromasova@fjfi.cvut.cz, svoboj68@fjfi.cvut.cz
Date: 17th March 2021

Note that you might have to change the URL in the load_basic function
depending on how the oscilloscope folder is called.
'''

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd # for reading csv files

URL_BASE = 'http://golem.fjfi.cvut.cz'  #base URL for GOLEM data access

def load_basic(name, shot):
    '''
    Loads signals from basic diagnostics osciloscope.

    Uses predefined folder for Basic diagnostic that can have different name for older shots.
    
    Parameters
    ----------
    name : str
        {'U_Loop, U_BtCoil, U_RogCoil'}
    shot : int
        golem shot number
    
    Returns
    -------
    t, np.ndarray
        time axis for given signal
    data, np.array
        oscilloscope data of given diagnostic name in given shot
    '''
    url = '{}/shots/{}/Diagnostics/BasicDiagnostics/{}.csv'.format(URL_BASE, shot, name)
    print(url)
    df = pd.read_csv(url, names=['time', name], index_col='time')
    t = np.array(df.index)
    data = np.transpose(np.array(df))[0]
    return t, data 

def load_rigol(name, shot, group='b'):
    '''
    Loads data from Rigol osciloscope controled by students.

    Uses the file structure for the summer semester 2021.
    
    Parameters
    ----------
    name, str
        {'U_Loop, U_BtCoil, U_RogCoil', U_photod}
    shot, int
        golem shot number
    group, str, optional
        number or letter of rigol osciloscope, should math that of group, by default 'b'
    
    Returns
    -------
    t, np.ndarray
        time axis for given signal
    data, np.array
        oscilloscope data of given diagnostic name in given shot
    '''
    url = f'{URL_BASE}/shots/{shot}/Diagnostics/BasicDiagnostics4TrainCourses/DAS_raw_data_dir-{group}/{name}-{group}.csv'
    print(url)
    df = pd.read_csv(url, names=['time', name], index_col='time')
    t = np.array(df.index)
    data = np.transpose(np.array(df))[0]
    return t, data  
    
if __name__ == '__main__':
    plt.ion()
    # Download oscilloscope data of loop voltage
    t, data = load_basic('U_Loop', 35091)

    # Plot their time trace
    fig = plt.figure('Loop voltage')
    plt.gca().set_title('Loop voltage')
    plt.plot(t*1000, data)
    plt.xlabel('$t$ [ms]')
    plt.ylabel('$U_l$ [V]')
    plt.grid(True)
    fig.tight_layout()