import numpy as np
from scipy.io.idl import readsav as readsav


# this is code to read in the mako multitone data sav files generated by Steve's idl code and store it in something useful for python


# Please note the date and details of any changes
# writen by Jordan Wheeler 12/13/2016
# add catch for instance when there are not

# To Do
# Would it be nicer to have these spit out a dictionary? or have another code that spits out everything as a python dictionary
# that would probably be better so that we don't have a backwards compatibility problem whenever we add new output.

# Change log
# 2017-06-02-jdw added catch for if there are no blind tone sin read_raw
# 2017-06-21-jdw changed file prefix so that it doesn't have to be in current directory
# 2017-06-21-jdw add keyword to allow loading of nonlinear fit extrapolated f0 calculated values

def read_raw(filename):  # filename is just the date/time code

    # Generate which resonators are blind and which are not so we don't look at blind tones
    # assumes it will be the same for each file
    raw = readsav(filename + "/rawdata.sav", verbose=False)
    bins = raw['multitone_data_raw']['bins'][0]
    blindbins = raw['multitone_data_raw']['blindbin'][0]
    if len(blindbins) == 0:  # there are no blind bins
        isblind = np.ones(bins.size)
    else:
        isblind = np.ones(bins.size)
        for i in range(0, bins.size):
            for j in range(0, blindbins.size):
                if bins[i] == blindbins[j]:
                    isblind[i] = 0

    blind_index = np.where(isblind < 0.5)[0]
    non_blind_index = np.where(isblind > 0)[0]

    f_res_all = raw['multitone_data_raw']['frlist'][0]
    # reads the data in 
    f_res = f_res_all[non_blind_index]  # only collect the f_res of the non blind bins
    f_res_blind = f_res_all[blind_index]  # only collect the f_res of the non blind bins

    raw_i = raw['multitone_data_raw'][0]['streamdata']['stream_data_concat'][0]['s21i'][0]
    raw_q = raw['multitone_data_raw'][0]['streamdata']['stream_data_concat'][0]['s21q'][0]
    print(raw_i.shape)
    non_blind_raw_i = raw_i[:, non_blind_index]
    non_blind_raw_q = raw_q[:, non_blind_index]

    blind_raw_i = raw_i[:, blind_index]
    blind_raw_q = raw_q[:, blind_index]

    # get the fine calibration data
    fine_i = np.zeros(
        (bins.shape[0], raw['multitone_data_raw']['fine_caldata'][0]['data_calibration'][0]['i'].shape[0]))
    fine_q = np.zeros(
        (bins.shape[0], raw['multitone_data_raw']['fine_caldata'][0]['data_calibration'][0]['q'].shape[0]))
    for b in range(0, raw['multitone_data_raw']['fine_caldata'][0]['data_calibration'][0]['i'].shape[0]):
        fine_i[:, b] = raw['multitone_data_raw']['fine_caldata'][0]['data_calibration'][0]['i'][b]
        fine_q[:, b] = raw['multitone_data_raw']['fine_caldata'][0]['data_calibration'][0]['q'][b]

    non_blind_fine_i = fine_i[non_blind_index, :]
    non_blind_fine_q = fine_q[non_blind_index, :]
    blind_fine_i = fine_i[blind_index, :]
    blind_fine_q = fine_q[blind_index, :]

    return non_blind_raw_i, non_blind_raw_q, blind_raw_i, blind_raw_q, f_res, f_res_blind, non_blind_fine_i, non_blind_fine_q, blind_fine_i, blind_fine_q


def read_cal(filename, **keywords):  # reading the calibrated file requires the raw file as well
    if ('model' in keywords):
        load_model = 1
    else:
        load_model = 0
    try:
        raw = readsav(filename + "/rawdata.sav", verbose=False)
    except:
        print("Reading the calibration file requires the raw file")
    cal = readsav(filename + "/calibrateddata.sav", verbose=False)

    # have to find out which bins are blind again
    bins = raw['multitone_data_raw']['bins'][0]
    blindbins = raw['multitone_data_raw']['blindbin'][0]
    isblind = np.ones(bins.size)

    for i in range(0, bins.size):
        for j in range(0, blindbins.size):
            if bins[i] == blindbins[j]:
                isblind[i] = 0

    blind_index = np.where(isblind < 0.5)[0]
    non_blind_index = np.where(isblind > 0)[0]

    S21_tan = np.imag(cal['multitone_data_calibrated'][0]['master_s21_stream_corr_rot_resampled']).astype('float64')[:,
              0, :]
    S21_norm = np.real(cal['multitone_data_calibrated'][0]['master_s21_stream_corr_rot_resampled']).astype('float64')[:,
               0, :]

    non_blind_S21_tan = S21_tan[:, non_blind_index]
    non_blind_S21_norm = S21_norm[:, non_blind_index]
    blind_S21_tan = S21_tan[:, blind_index]
    blind_S21_norm = S21_norm[:, blind_index]

    # f0_calc = cal['multitone_data_calibrated'][0]['master_f0_calculated_resampled'].astype('float64')[:,0,:]
    if load_model == 0:
        f0_calc = cal['multitone_data_calibrated'][0]['MASTER_F0_CALCULATED'].astype('float64')[:, 0, :]
    else:
        f0_calc = cal['multitone_data_calibrated'][0]['master_f0_calculated_fine_seriesx_nonlinear'].astype('float64')[
                  :, 0, :]
    try:
        nan_loc = np.where(np.isnan(f0_calc[:, -1]) == True)[0][0]
        f0_calc = f0_calc[0:nan_loc, :]
    except:
        f0_calc = cal['multitone_data_calibrated'][0]['master_f0_calculated_resampled'].astype('float64')[:, 0, :]

    f0_calc = f0_calc[:, non_blind_index]

    return non_blind_S21_tan, non_blind_S21_norm, blind_S21_tan, blind_S21_norm, f0_calc
