import numpy as np
import pandas as pd

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

def load_custom(name, shot, oscope):
    '''
    Loads data from oscilloscope controlled by students.

    Uses the file structure for the summer semester 2022.
    
    Parameters
    ----------
    name, str
        {'U_Loop, U_BtCoil, U_RogCoil', U_photod}
    shot, int
        golem shot number
    oscope, str
        {a, b, c, d}
        number or letter of Rigol oscilloscope, 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
    '''
    dct = dict(
        a='RigolMSO5204-d',
        b='RigolMSO5204-b',
        c='RigolMSO5104-a',
        d='RigolDS1104Z-a',
    )
    url = f'{URL_BASE}/{shot}/Devices/Oscilloscopes/{dct[oscope]}/{name}.csv'
    print(f'Loading data from {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