functions_solo.py

import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from spacepy import pycdf
import cdflib
import spiceypy
# import os
import urllib.request
import os.path
import pickle
import glob

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.colors as colors
from mpl_toolkits.axes_grid1 import make_axes_locatable

from functions_general import load_path


"""
SOLAR ORBITER SERVER DATA PATH
"""

solo_path=load_path(path_name='solo_path')
print(f"Solo path loaded: {solo_path}")

# Load path once globally
kernels_path = load_path(path_name='kernels_path')
print(f"Kernels path loaded: {kernels_path}")


"""
SOLO BAD DATA FILTER
"""


def filter_bad_data(df, col, bad_val): #filter across whole df
    if bad_val < 0:
        mask = df[col] < bad_val  # boolean mask for all bad values
    else:
        mask = df[col] > bad_val  # boolean mask for all bad values
    cols = [x for x in df.columns if x != 'timestamp']
    df.loc[mask, cols] = np.nan
    return df


def filter_bad_col(df, col, bad_val): #filter by individual columns
    if bad_val < 0:
        mask_vals = df[col] < bad_val  # boolean mask for all bad values
    else:
        mask_vals = df[col] > bad_val  # boolean mask for all bad values
    df[col].mask(mask_vals)
    return df


"""
SOLO MAG DATA
# Potentially different SolO MAG file names: internal low latency, formagonly, and formagonly 1 min.
e.g.
- Internal files: solo_L2_mag-rtn-ll-internal_20230225_V00.cdf
- For MAG only files: solo_L2_mag-rtn-normal-formagonly_20200415_V01.cdf
- For MAG only 1 minute res files: solo_L2_mag-rtn-normal-1-minute-formagonly_20200419_V01.cdf
# All in RTN coords
# Should all follow same variable names within cdf
"""


#DOWNLOAD FUNCTIONS for 1min or 1sec data


def download_solomag_1min(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/l2/1min'): 
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        data_item_id = f'solo_L2_mag-rtn-normal-1-minute_{date_str}'
        if os.path.isfile(f"{path}/{data_item_id}.cdf") == True:
            print(f'{data_item_id}.cdf has already been downloaded.')
            start += timedelta(days=1)
        else:
            try:
                data_url = f'http://soar.esac.esa.int/soar-sl-tap/data?retrieval_type=PRODUCT&data_item_id={data_item_id}&product_type=SCIENCE'
                urllib.request.urlretrieve(data_url, f"{path}/{data_item_id}.cdf")
                print(f'Successfully downloaded {data_item_id}.cdf')
                start += timedelta(days=1)
            except Exception as e:
                print('ERROR', e, data_item_id)
                start += timedelta(days=1)


def download_solomag_1sec(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/l2/1sec'):
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        data_item_id = f'solo_L2_mag-rtn-normal_{date_str}'
        if os.path.isfile(f"{path}/{data_item_id}.cdf") == True:
            print(f'{data_item_id}.cdf has already been downloaded.')
            start += timedelta(days=1)
        else:
            try:
                data_url = f'http://soar.esac.esa.int/soar-sl-tap/data?retrieval_type=PRODUCT&data_item_id={data_item_id}&product_type=SCIENCE'
                urllib.request.urlretrieve(data_url, f"{path}/{data_item_id}.cdf")
                print(f'Successfully downloaded {data_item_id}.cdf')
                start += timedelta(days=1)
            except Exception as e:
                print('ERROR', e, data_item_id)
                start += timedelta(days=1)


#combined solomag download functions to specify resolution of data 
def download_solomag(start_timestamp, end_timestamp, res="1min"):
    if res == "1min":
        download_solomag_1min(start_timestamp, end_timestamp)
    elif res == "1sec":
        download_solomag_1sec(start_timestamp, end_timestamp)
            

#LOAD FUNCTIONS for MAG data 


#Load single file from specific path using pycdf from spacepy
def get_solomag(fp):
    """raw = rtn"""
    try:
        cdf = pycdf.CDF(fp)
        data = {df_col: cdf[cdf_col][:] for cdf_col, df_col in zip(['EPOCH'], ['time'])}
        df = pd.DataFrame.from_dict(data)
        bx, by, bz = cdf['B_RTN'][:].T
        df['bx'] = bx
        df['by'] = by
        df['bz'] = bz
        df['bt'] = np.linalg.norm(df[['bx', 'by', 'bz']], axis=1)
    except Exception as e:
        print('ERROR:', e, fp)
        df = None
    return df

# #Load single file from specific path using cdflib
# def get_solomag(fp):
#     """raw = rtn"""
#     try:
#         cdf = cdflib.CDF(fp)
#         t1 = cdflib.cdfepoch.to_datetime(cdf.varget('EPOCH'))
#         df = pd.DataFrame(t1, columns=['time'])
#         bx, by, bz = cdf['B_RTN'][:].T
#         df['bx'] = bx
#         df['by'] = by
#         df['bz'] = bz
#         df['bt'] = np.linalg.norm(df[['bx', 'by', 'bz']], axis=1)
#         cols = ['bx', 'by', 'bz', 'bt']
#         for col in cols:
#             df[col].mask(df[col] < -9.999E29 , pd.NA, inplace=True)
#     except Exception as e:
#         print('ERROR:', e, fp)
#         df = None
#     return df


#Load range of files using specified start and end dates/ timestamps
def get_solomag_range_formagonly_internal(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/ll'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        year = start.year
        date_str = f'{year}{start.month:02}{start.day:02}'
        fn = f'{path}/solo_L2_mag-rtn-ll-internal_{date_str}_V00.cdf'
        _df = get_solomag(fn)
        if _df is not None:
            if df is None:
                df = _df.copy(deep=True)
            else:
                df = pd.concat([df, _df])
        start += timedelta(days=1)
    return df


def get_solomag_range_formagonly(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/formagonly/full'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        year = start.year
        date_str = f'{year}{start.month:02}{start.day:02}'
        fn = glob.glob(f'{path}/solo_L2_mag-rtn-normal-formagonly_{date_str}_*.cdf')
        _df = get_solomag(fn[0])
        if _df is not None:
            if df is None:
                df = _df.copy(deep=True)
            else:
                df = pd.concat([df, _df])
        start += timedelta(days=1)
    return df


def get_solomag_range_formagonly_1min(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/formagonly/1min'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        year = start.year
        date_str = f'{year}{start.month:02}{start.day:02}'
        fn = glob.glob(f'{path}/solo_L2_mag-rtn-normal-1-minute-formagonly_{date_str}_*.cdf')
        _df = get_solomag(fn[0])
        if _df is not None:
            if df is None:
                df = _df.copy(deep=True)
            else:
                df = pd.concat([df, _df])
        start += timedelta(days=1)
    return df


def get_solomag_range_1sec(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/l2/1sec'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        fn = f'{path}/solo_L2_mag-rtn-normal_{date_str}.cdf'
        _df = get_solomag(fn)
        if _df is not None:
            if df is None:
                df = _df.copy(deep=True)
            else:
                df = pd.concat([df, _df])
        start += timedelta(days=1)
    return df


def get_solomag_range_1min(start_timestamp, end_timestamp, path=f'{solo_path}'+'mag/l2/1min'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        fn = f'{path}/solo_L2_mag-rtn-normal-1-minute_{date_str}.cdf'
        _df = get_solomag(fn)
        if _df is not None:
            if df is None:
                df = _df.copy(deep=True)
            else:
                df = pd.concat([df, _df])
        start += timedelta(days=1)
    return df


#combined solomag range function to specify level and resolution of data 
def get_solomag_range(start_timestamp, end_timestamp, level="l2", res="1min"):
    if level == "l2":
        if res == "1min":
            df = get_solomag_range_1min(start_timestamp, end_timestamp)
        elif res == "1sec":
            df = get_solomag_range_1sec(start_timestamp, end_timestamp)
    elif level == "ll":
        df = get_solomag_range_formagonly_internal(start_timestamp, end_timestamp)
    elif level == "formagonly":
        if res == "full":
            df = get_solomag_range_formagonly(start_timestamp, end_timestamp)
        elif res == "1min":
            df = get_solomag_range_formagonly_1min(start_timestamp, end_timestamp)
    return df 


"""
SOLO PLASMA DATA
# Level 2 science SWA PLAS grnd moment data
"""


#DOWNLOAD FUNCTION for swa/plas data
def download_soloplas(start_timestamp, end_timestamp, path=f'{solo_path}'+'swa/plas/l2'):
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        data_item_id = f'solo_L2_swa-pas-grnd-mom_{date_str}'
        if os.path.isfile(f"{path}/{data_item_id}.cdf") == True:
            print(f'{data_item_id}.cdf has already been downloaded.')
            start += timedelta(days=1)
        else:
            try:
                data_url = f'http://soar.esac.esa.int/soar-sl-tap/data?retrieval_type=PRODUCT&data_item_id={data_item_id}&product_type=SCIENCE'
                urllib.request.urlretrieve(data_url, f"{path}/{data_item_id}.cdf")
                print(f'Successfully downloaded {data_item_id}.cdf')
                start += timedelta(days=1)
            except Exception as e:
                print('ERROR', e, data_item_id)
                start += timedelta(days=1)


#Load single file from specific path using pycdf from spacepy
def get_soloplas_l2(fp):
    """raw = rtn"""
    try:
        cdf = pycdf.CDF(fp)
        data = {df_col: cdf[cdf_col][:] for cdf_col, df_col in zip(['Epoch', 'N', 'T'], ['time', 'np', 'tp'])}
        df = pd.DataFrame.from_dict(data)
        df['time'] = pd.to_datetime(df['time'])
        vx, vy, vz = cdf['V_RTN'][:].T
        df['vx'] = vx
        df['vy'] = vy
        df['vz'] = vz
        df['vt'] = np.linalg.norm(df[['vx', 'vy', 'vz']], axis=1)
    except Exception as e:
        print('ERROR:', e, fp)
        df = None
    return df


#Load single low latency file from specific path (no temp data)
def get_soloplas_ll(fp):
    """raw = rtn"""
    try:
        cdf = pycdf.CDF(fp)
        data = {df_col: cdf[cdf_col][:] for cdf_col, df_col in zip(['EPOCH', 'SWA_PAS_DENSITY'], ['time', 'np'])}
        df = pd.DataFrame.from_dict(data)
        df['time'] = pd.to_datetime(df['time'])
        vx, vy, vz = cdf['SWA_PAS_VELOCITY_RTN'][:].T
        df['vx'] = vx
        df['vy'] = vy
        df['vz'] = vz
    except Exception as e:
        print('ERROR:', e, fp)
        df = None
    df = filter_bad_col(df, 'np', -1E30)
    df = filter_bad_col(df, 'vx', -1E30)
    df = filter_bad_col(df, 'vy', -1E30)
    df = filter_bad_col(df, 'vz', -1E30)
    df['vt'] = np.linalg.norm(df[['vx', 'vy', 'vz']], axis=1)
    df['tp'] = np.nan
    return df


# #Load single file from specific path using cdflib
# def get_soloplas(fp):
#     """raw = rtn"""
#     try:
#         cdf = cdflib.CDF(fp)
#         t1 = cdflib.cdfepoch.to_datetime(cdf.varget('EPOCH'))
#         df = pd.DataFrame(t1, columns=['time'])
#         df['np'] = cdf['N']
#         df['tp'] = cdf['T']
#         vx, vy, vz = cdf['V_RTN'][:].T
#         df['vx'] = vx
#         df['vy'] = vy
#         df['vz'] = vz
#         df['vt'] = np.linalg.norm(df[['vx', 'vy', 'vz']], axis=1)
#         cols = ['np', 'tp', 'vx', 'vy', 'vz', 'vt']
#         for col in cols:
#             df[col].mask(df[col] < -9.999E29 , pd.NA, inplace=True)
#     except Exception as e:
#         print('ERROR:', e, fp)
#         df = None
#     return df


#Load range of files using specified start and end dates/ timestamps
def get_soloplas_range_l2(start_timestamp, end_timestamp, path=f'{solo_path}'+'swa/plas/l2'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        fn = f'{path}/solo_L2_swa-pas-grnd-mom_{date_str}.cdf'
        _df = get_soloplas_l2(fn)
        if _df is not None:
            if df is None:
                df = _df.copy(deep=True)
            else:
                df = pd.concat([df, _df])
        start += timedelta(days=1)
    return df


#Loads range between specified dates, returns combined df. 
#BE CAREFUL: path where data is saved must be carefully maintained: only keep one file per date (the longest time range)
def get_soloplas_range_ll(start_timestamp, end_timestamp, path=f'{solo_path}'+'swa/plas/ll'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given."""
    df = None
    start = start_timestamp.date()
    end = end_timestamp.date() + timedelta(days=1)
    while start < end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        fn = glob.glob(f'{path}/solo_LL02_swa-pas-mom_{date_str}*.cdf')
        try:
            _df = get_soloplas_ll(fn[0])
            if _df is not None:
                if df is None:
                    df = _df.copy(deep=True)
                else:
                    df = pd.concat([df, _df])
        except Exception as e:
            print('ERROR', e, "DataFrame is empty") 
        start += timedelta(days=1)
    return df


#combined soloplas range function to specify level and resolution of data 
def get_soloplas_range(start_timestamp, end_timestamp, level="l2", res="full"): #res col is filler for consistency with mag for now
    if level == "l2":
        df = get_soloplas_range_l2(start_timestamp, end_timestamp)
    elif level == "ll":
        df = get_soloplas_range_ll(start_timestamp, end_timestamp)
    return df 

"""
COMBINED SOLO MAG AND PLAS
"""

def get_solomagplas(start_timestamp, end_timestamp, level='l2', res='1min'):
    df_mag = get_solomag_range(start_timestamp, end_timestamp, level, res)
    if df_mag is None:
        print(f'SolO MAG data is empty for this timerange')
        df_mag = pd.DataFrame({'time':[], 'bt':[], 'bx':[], 'by':[], 'bz':[]})
        mag_rdf = df_mag.drop(columns=['time'])
    else:
        mag_rdf = df_mag.set_index('time').resample('1min').mean().reset_index(drop=False)
        mag_rdf.set_index(pd.to_datetime(mag_rdf['time']), inplace=True)
        
    df_plas = get_soloplas_range(start_timestamp, end_timestamp, level)
    if df_plas is None:
        print(f'SolO SWA data is empty for this timerange')
        df_plas = pd.DataFrame({'time':[], 'vt':[], 'vx':[], 'vy':[], 'vz':[], 'np':[], 'tp':[]})
        plas_rdf = df_plas
    else:
        plas_rdf = df_plas.set_index('time').resample('1min').mean().reset_index(drop=False)
        plas_rdf.set_index(pd.to_datetime(plas_rdf['time']), inplace=True)
        if mag_rdf.shape[0] != 0:
            plas_rdf = plas_rdf.drop(columns=['time'])

    magplas_rdf = pd.concat([mag_rdf, plas_rdf], axis=1)
    magplas_rdf = magplas_rdf.drop(columns=['time'])
    magplas_rdf['time'] = magplas_rdf.index
    magplas_rdf = magplas_rdf.reset_index(drop=True)

    return magplas_rdf


"""""
SOLO PAD FUNCTIONS:
PAD data from PAS EAS, L3
"""""


def get_solopad_range(start_timestamp, end_timestamp, path=f'{solo_path}'+'swa/eas/l3'):

    epoch_arr  = np.empty((0,), dtype='datetime64[ns]') 
    energy_arr = np.empty((64)) 
    pangle_arr = np.empty((36))
    psd_arr   = np.empty((0,36,64))

    start = start_timestamp.date()
    end = end_timestamp.date()

    while start <= end:
        fp = path+'/solo_L3_swa-eas-nmpad-psd_'+start.strftime("%Y")+start.strftime("%m")+start.strftime("%d")+'_V01.cdf'
        cdf = pycdf.CDF(fp)
        data = {df_col: cdf[cdf_col][:] for cdf_col, df_col in zip(['EPOCH', 'SWA_EAS_ENERGY', 'SWA_EAS_NMPAD_PSD_DATA', 'SWA_EAS_PA'], ['time', 'energy', 'psd_data', 'pangle'])}
        epoch  = pd.to_datetime(data['time'])  # array of dim X  (time entries)
        energy = data['energy']  # array of dim [X,Y] (time entries and 15 energy channels) 
        pangle = data['pangle']  # array of dim [X,Z] (time entries and 8 angles) 
        psd_data   = data['psd_data']    # array of dim [X,Z,Y] (time entries, 8 angles, 15 energy channels) 
        epoch_arr  = np.concatenate((epoch_arr, epoch))
        energy_arr = energy #np.concatenate((energy_arr, energy))
        pangle_arr = pangle #np.concatenate((pangle_arr, pangle))
        psd_arr   = np.concatenate((psd_arr, psd_data))
        start+=timedelta(days=1)
    epoch_arr = pd.to_datetime(epoch_arr)
    return epoch_arr, energy_arr, pangle_arr, psd_arr


def get_solopad_array(epoch_arr, pangle_arr, psd_arr, energy_arr, min_energy=70, max_energy=5300):
    x_arr = epoch_arr.to_numpy()
    y_arr = pangle_arr
    i = len(energy_arr)-np.sum(energy_arr<=max_energy)
    j = np.sum(energy_arr>min_energy)-1
    z = np.sum(psd_arr[:, :, i:j], axis=2)
    z_arr_unfiltered = z.T
    z_arr = z_arr_unfiltered.copy()
    mask = (np.abs(z_arr_unfiltered) >= 1e30) | (z_arr_unfiltered == 0) #bad value given as >1e30/<-1e30, and remove 0 for log scale plotting
    z_arr[mask] = np.nan
    return x_arr, y_arr, z_arr


def plot_solopad_array(x_arr, y_arr, z_arr):
    xsize = 12
    ysize = 1
    pad   = 0.04
    labelpad = 12
    fig, axes  =  plt.subplots(nrows=1, ncols=1, figsize=(xsize,ysize))

    im = axes.pcolormesh(x_arr,y_arr,z_arr, 
                    norm=colors.LogNorm(vmin=np.nanmin(z_arr), vmax=np.nanmax(z_arr)),
                    cmap=plt.get_cmap('jet'))
    axes.set_ylim(0, 180)
    axes.set_yticks([0, 90, 180])
    # colorbar
    divider = make_axes_locatable(axes)
    cax = divider.append_axes('right', size='5%', pad=pad)
    cbar = fig.colorbar(im, cax=cax, orientation='vertical')
    cbar.ax.set_ylabel('PSD', rotation=90, labelpad=labelpad)
    plt.show()


"""""
SOLO EPD FUNCTIONS:
EPT and HET functions
"""""


def get_soloept_l2(fp):
    try:
        cdf = pycdf.CDF(fp)
        e_data = cdf["Electron_Flux"][:]
        e_df = pd.DataFrame.from_dict(e_data)
        e_df.mask((e_df < -9.99e+30), inplace=True)
        e_flux = e_df.to_numpy() 
        e_bins = cdf["Electron_Bins_Text"][:]
        i_data = cdf["Ion_Flux"][:]
        i_df = pd.DataFrame.from_dict(i_data)
        i_df.mask((i_df < -9.99e+30), inplace=True)
        i_flux = i_df.to_numpy() 
        i_bins = cdf["Ion_Bins_Text"][:]
        time = cdf["EPOCH"][:]
    except Exception as e:
        print('ERROR:', e, fp)
        time = None 
        e_flux = None
        e_bins = None 
        i_flux = None 
        i_bins = None
    return time, e_flux, e_bins, i_flux, i_bins


def get_soloept_range_l2(start_timestamp, end_timestamp, direction:str, path=f'{solo_path}'+'epd/ept/l2'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given. Possible directions = asun, sun, north, south"""
    time_arr  = np.empty((0,), dtype='datetime64[ns]') 
    e_flux_arr = np.empty((0,34), dtype='float32')
    e_bins_arr = np.empty((0,), dtype='<U25')
    i_flux_arr = np.empty((0,64), dtype='float32')
    i_bins_arr = np.empty((0,), dtype='<U25')
    
    start = start_timestamp.date()
    end = end_timestamp.date() 
    while start <= end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        fn = glob.glob(f'{path}/solo_L2_epd-ept-{direction}-rates_{date_str}_*.cdf')
        time, e_flux, e_bins, i_flux, i_bins = get_soloept_l2(fn[0])
        time_arr  = np.concatenate((time_arr, time))
        e_flux_arr = np.concatenate((e_flux_arr, e_flux))
        e_bins_arr = e_bins
        i_flux_arr = np.concatenate((i_flux_arr, i_flux))
        i_bins_arr = i_bins
        start += timedelta(days=1)
    return time_arr, e_flux_arr, e_bins_arr, i_flux_arr, i_bins_arr


def make_soloept_dataframes(time_arr, e_flux_arr, e_bins_arr, i_flux_arr, i_bins_arr, cadence=60):
        e_df = pd.DataFrame(e_flux_arr, columns=e_bins_arr)
        e_df['time'] = time_arr
        e_rdf = e_df.set_index('time').resample(f'{cadence}min').mean().reset_index(drop=False)
        i_df = pd.DataFrame(i_flux_arr, columns=i_bins_arr)
        i_df['time'] = time_arr
        i_rdf = i_df.set_index('time').resample(f'{cadence}min').mean().reset_index(drop=False)
        return e_rdf, i_rdf


def get_solohet_l2(fp):
    try:
        cdf = pycdf.CDF(fp)
        e_data = cdf["Electron_Flux"][:]
        e_df = pd.DataFrame.from_dict(e_data)
        e_df.mask((e_df < -9.99e+30), inplace=True)
        e_flux = e_df.to_numpy() 
        e_bins = cdf["Electron_Bins_Text"][:]
        e_time = cdf["EPOCH_4"][:]
        h_data = cdf["H_Flux"][:]
        h_df = pd.DataFrame.from_dict(h_data)
        h_df.mask((h_df < -9.99e+30), inplace=True)
        h_flux = h_df.to_numpy() 
        h_bins = cdf["H_Bins_Text"][:]
        h_time = cdf["EPOCH"][:]
    except Exception as e:
        print('ERROR:', e, fp)
        e_time = None 
        e_flux = None
        e_bins = None 
        h_time = None
        h_flux = None 
        h_bins = None
    return e_time, e_flux, e_bins, h_time, h_flux, h_bins


def get_solohet_range_l2(start_timestamp, end_timestamp, direction:str, path=f'{solo_path}'+'epd/het/l2'):
    """Pass two datetime objects and grab .cdf files between dates, from
    directory given. Possible directions = asun, sun, north, south"""
    e_time_arr  = np.empty((0,), dtype='datetime64[ns]') 
    e_flux_arr = np.empty((0,4), dtype='float32')
    e_bins_arr = np.empty((0,), dtype='<U25')
    h_time_arr  = np.empty((0,), dtype='datetime64[ns]') 
    h_flux_arr = np.empty((0,36), dtype='float32')
    h_bins_arr = np.empty((0,), dtype='<U25')
    
    start = start_timestamp.date()
    end = end_timestamp.date() 
    while start <= end:
        date_str = f'{start.year}{start.month:02}{start.day:02}'
        fn = glob.glob(f'{path}/solo_L2_epd-het-{direction}-rates_{date_str}_*.cdf')
        e_time, e_flux, e_bins, h_time, h_flux, h_bins = get_solohet_l2(fn[0])
        e_time_arr  = np.concatenate((e_time_arr, e_time))
        e_flux_arr = np.concatenate((e_flux_arr, e_flux))
        e_bins_arr = e_bins
        h_time_arr  = np.concatenate((h_time_arr, h_time))
        h_flux_arr = np.concatenate((h_flux_arr, h_flux))
        h_bins_arr = h_bins
        start += timedelta(days=1)
    return e_time_arr, e_flux_arr, e_bins_arr, h_time_arr, h_flux_arr, h_bins_arr


def make_solohet_dataframes(e_time_arr, e_flux_arr, e_bins_arr, h_time_arr, h_flux_arr, h_bins_arr, cadence=60):
        e_df = pd.DataFrame(e_flux_arr, columns=e_bins_arr)
        e_df['time'] = e_time_arr
        e_rdf = e_df.set_index('time').resample(f'{cadence}min').mean().reset_index(drop=False)
        h_df = pd.DataFrame(h_flux_arr, columns=h_bins_arr)
        h_df['time'] = h_time_arr
        h_rdf = h_df.set_index('time').resample(f'{cadence}min').mean().reset_index(drop=False)
        return e_rdf, h_rdf


"""
SOLO POSITION FUNCTIONS: coord maths, furnish kernels, and call position for each timestamp
Currently set to HEEQ, but will implement options to change
"""


def cart2sphere(x,y,z):
    r = np.sqrt(x**2+ y**2 + z**2) /1.495978707E8         
    theta = np.arctan2(z,np.sqrt(x**2+ y**2)) * 360 / 2 / np.pi
    phi = np.arctan2(y,x) * 360 / 2 / np.pi                   
    return (r, theta, phi)


#http://spiftp.esac.esa.int/data/SPICE/SOLAR-ORBITER/kernels/fk/ for solo_ANC_soc-sci-fk_V08.tf
#http://spiftp.esac.esa.int/data/SPICE/SOLAR-ORBITER/kernels/spk/ for solo orbit .bsp


def solo_furnish():
    """Main"""
    solo_path = kernels_path+'solo/'
    generic_path = kernels_path+'generic/'
    solo_kernels = os.listdir(solo_path)
    generic_kernels = os.listdir(generic_path)
    for kernel in solo_kernels:
        spiceypy.furnsh(os.path.join(solo_path, kernel))
    for kernel in generic_kernels:
        spiceypy.furnsh(os.path.join(generic_path, kernel))


def get_solo_pos(t, coord_sys='ECLIPJ2000'): 
    if spiceypy.ktotal('ALL') < 1:
        solo_furnish()
    if coord_sys == 'GSE':
        try:
            pos = spiceypy.spkpos("SOLAR ORBITER", spiceypy.datetime2et(t), f"{coord_sys}", "NONE", "EARTH")[0] 
            r, lat, lon = cart2sphere(pos[0],pos[1],pos[2])
            position = t, pos[0], pos[1], pos[2], r, lat, lon
            return position
        except Exception as e:
            print(e)
            return [t, None, None, None, None, None, None]
    else:
        try:
            pos = spiceypy.spkpos("SOLAR ORBITER", spiceypy.datetime2et(t), f"{coord_sys}", "NONE", "SUN")[0] 
            r, lat, lon = cart2sphere(pos[0],pos[1],pos[2])
            position = t, pos[0], pos[1], pos[2], r, lat, lon
            return position
        except Exception as e:
            print(e)
            return [t, None, None, None, None, None, None]


def get_solo_positions(time_series, coord_sys):
    positions = []
    for t in time_series:
        position = get_solo_pos(t, coord_sys)
        positions.append(position)
    df_positions = pd.DataFrame(positions, columns=['time', 'x', 'y', 'z', 'r', 'lat', 'lon'])
    return df_positions



def get_solo_positions_daily(start, end, coord_sys, cadence, dist_unit='au', ang_unit='deg'):
    t = start
    positions = []
    while t < end:
        position = get_solo_pos(t, coord_sys)
        positions.append(position)
        t += timedelta(days=cadence)
    df_positions = pd.DataFrame(positions, columns=['time', 'x', 'y', 'z', 'r', 'lat', 'lon'])
    if dist_unit == 'au':
        df_positions.x = df_positions.x/1.495978707E8 
        df_positions.y = df_positions.y/1.495978707E8
        df_positions.z = df_positions.z/1.495978707E8
    if ang_unit == 'rad':
        df_positions.lat = df_positions.lat * np.pi / 180
        df_positions.lon = df_positions.lon * np.pi / 180
    spiceypy.kclear()
    return df_positions


def get_solo_positions_hourly(start, end, coord_sys, cadence, dist_unit='au', ang_unit='deg'):
    t = start
    positions = []
    while t < end:
        position = get_solo_pos(t, coord_sys)
        positions.append(position)
        t += timedelta(hours=cadence)
    df_positions = pd.DataFrame(positions, columns=['time', 'x', 'y', 'z', 'r', 'lat', 'lon'])
    if dist_unit == 'au':
        df_positions.x = df_positions.x/1.495978707E8 
        df_positions.y = df_positions.y/1.495978707E8
        df_positions.z = df_positions.z/1.495978707E8
    if ang_unit == 'rad':
        df_positions.lat = df_positions.lat * np.pi / 180
        df_positions.lon = df_positions.lon * np.pi / 180
    spiceypy.kclear()
    return df_positions


def get_solo_positions_minute(start, end, coord_sys, cadence, dist_unit='au', ang_unit='deg'):
    t = start
    positions = []
    while t < end:
        position = get_solo_pos(t, coord_sys)
        positions.append(position)
        t += timedelta(minutes=cadence)
    df_positions = pd.DataFrame(positions, columns=['time', 'x', 'y', 'z', 'r', 'lat', 'lon'])
    if dist_unit == 'au':
        df_positions.x = df_positions.x/1.495978707E8 
        df_positions.y = df_positions.y/1.495978707E8
        df_positions.z = df_positions.z/1.495978707E8
    if ang_unit == 'rad':
        df_positions.lat = df_positions.lat * np.pi / 180
        df_positions.lon = df_positions.lon * np.pi / 180
    spiceypy.kclear()
    return df_positions


"""
OUTPUT COMBINED PICKLE FILE
including MAG, PLAS, and POSITION data
"""


def create_solo_pkl(start_timestamp, end_timestamp, pos_coord_sys, level='l2', res='1min', output_path='/Users/emmadavies/Documents/Projects/SolO_Realtime_Preparation/March2024/'):
    
    # #download solo mag and plasma data up to now 
    # download_solomag_1min(start_timestamp)
    # download_soloplas(start_timestamp)

    #load in mag data to DataFrame and resample, create empty mag and resampled DataFrame if no data
    # if empty, drop time column ready for concat
    df_mag = get_solomag_range(start_timestamp, end_timestamp, level, res)
    if df_mag is None:
        print(f'SolO MAG data is empty for this timerange')
        df_mag = pd.DataFrame({'time':[], 'bt':[], 'bx':[], 'by':[], 'bz':[]})
        mag_rdf = df_mag.drop(columns=['time'])
    else:
        mag_rdf = df_mag.set_index('time').resample('1min').mean().reset_index(drop=False)
        mag_rdf.set_index(pd.to_datetime(mag_rdf['time']), inplace=True)
        
    #load in plasma data to DataFrame and resample, create empty plasma and resampled DataFrame if no data
    #only drop time column if MAG DataFrame is not empty
    df_plas = get_soloplas_range(start_timestamp, end_timestamp, level)
    if df_plas is None:
        print(f'SolO SWA data is empty for this timerange')
        df_plas = pd.DataFrame({'time':[], 'vt':[], 'vx':[], 'vy':[], 'vz':[], 'np':[], 'tp':[]})
        plas_rdf = df_plas
    else:
        plas_rdf = df_plas.set_index('time').resample('1min').mean().reset_index(drop=False)
        plas_rdf.set_index(pd.to_datetime(plas_rdf['time']), inplace=True)
        if mag_rdf.shape[0] != 0:
            plas_rdf = plas_rdf.drop(columns=['time'])

    #need to combine mag and plasma dfs to get complete set of timestamps for position calculation
    magplas_rdf = pd.concat([mag_rdf, plas_rdf], axis=1)
    #some timestamps may be NaT so after joining, drop time column and reinstate from combined index col
    magplas_rdf = magplas_rdf.drop(columns=['time'])
    magplas_rdf['time'] = magplas_rdf.index
     
    #get solo positions for corresponding timestamps
    solo_furnish()
    solo_pos = get_solo_positions(magplas_rdf['time'], pos_coord_sys)
    solo_pos.set_index(pd.to_datetime(solo_pos['time']), inplace=True)
    solo_pos = solo_pos.drop(columns=['time'])
    spiceypy.kclear()

    #produce final combined DataFrame with correct ordering of columns 
    comb_df = pd.concat([magplas_rdf, solo_pos], axis=1)

    #produce recarray with correct datatypes
    time_stamps = comb_df['time']
    dt_lst= [element.to_pydatetime() for element in list(time_stamps)] #extract timestamps in datetime.datetime format

    solo=np.zeros(len(dt_lst),dtype=[('time',object),('bx', float),('by', float),('bz', float),('bt', float),\
                ('vx', float),('vy', float),('vz', float),('vt', float),('np', float),('tp', float),\
                ('x', float),('y', float),('z', float), ('r', float),('lat', float),('lon', float)])
    solo = solo.view(np.recarray) 

    solo.time=dt_lst
    solo.bx=comb_df['bx']
    solo.by=comb_df['by']
    solo.bz=comb_df['bz']
    solo.bt=comb_df['bt']
    solo.vx=comb_df['vx']
    solo.vy=comb_df['vy']
    solo.vz=comb_df['vz']
    solo.vt=comb_df['vt']
    solo.np=comb_df['np']
    solo.tp=comb_df['tp']
    solo.x=comb_df['x']
    solo.y=comb_df['y']
    solo.z=comb_df['z']
    solo.r=comb_df['r']
    solo.lat=comb_df['lat']
    solo.lon=comb_df['lon']
    
    #dump to pickle file
    header='Internal low latency solar wind magnetic field (MAG) from Solar Orbiter, ' + \
    'provided directly by Imperial College London SolO MAG Team '+ \
    'Timerange: '+solo.time[0].strftime("%Y-%b-%d %H:%M")+' to '+solo.time[-1].strftime("%Y-%b-%d %H:%M")+\
    ', resampled to a time resolution of 1 min. '+\
    'The data are available in a numpy recarray, fields can be accessed by solo.time, solo.bx, solo.r etc. '+\
    'Total number of data points: '+str(solo.size)+'. '+\
    'Units are btxyz [nT, RTN], heliospheric position x/y/z/r/lon/lat [AU, degree, HEEQ]. '+\
    'Made with script by E.E. Davies (github @ee-davies, twitter @spacedavies). File creation date: '+\
    datetime.utcnow().strftime("%Y-%b-%d %H:%M")+' UTC'

    t_now_date = datetime.utcnow().strftime("%Y-%m-%d")
    pickle.dump([solo,header], open(output_path+f'solo_rtn_{t_now_date}.p', "wb"))