import pandas as pd

file = 'example_txt/vortex.times.677619.6m 100m UTC+02.0 ERA5.txt'
df = pd.read_csv(file, sep=r"\s+", skiprows=3, header=0, names=None,
                parse_dates={'time': [0, 1]}, index_col='time')
df

import numpy as np
from typing import Dict, Tuple


def read_vortex_times_pandas(filename: str) -> (Tuple)[pd.DataFrame, Dict[str, float]]:
   """
   Read typical vortex time series from TIMES product


   We return a pandas dataframe at utc 0 (universal time zone)


   Parameters
   ----------
   filename: str


   Returns
   -------
   data: pd.DataFrame
       Time series with time index
   coords: Dict[float]
       Coordinates (keys: lat, lon and lev)
   """
   patterns = {'Lat=': 'lat', 'Lon=': 'lon', 'Hub-Height=': 'lev'}


   # Look at the first line of the text file (the header)
   # and manually look for the values of Lat= Lon= Timezone=
   # and Hub-Height to save the
   utc = np.nan
   coords = {}
   with open(filename, 'r') as f:
       header = f.readline()
       for info in header.split(' '):
           for pattern, keyword in patterns.items():
               if pattern in info:
                   coords[keyword] = float(info.replace(pattern, ''))
               elif 'Timezone=' in info:
                   utc = float(info.replace('Timezone=', ''))
   if np.isnan(utc):
       raise ValueError('Could not read UTC from the header: ' + header)


   data = pd.read_csv(filename, sep=r"\s+", skiprows=3, header=0, names=None,
                      parse_dates={'time': [0, 1]}, index_col='time')
   data.dropna(inplace=True)
   data.index = data.index - pd.Timedelta(utc, unit='h')


   vars_new_names = {
       'M(m/s)': 'M', 'D(deg)': 'Dir', 'SD(m/s)': 'SD', 'T(C)': 'T',
       'De(k/m3)': 'D', 'PRE(hPa)': 'P', 'RiNumber': 'RI',
       'RH(%)': 'RH', 'RMOL(1/m)': 'RMOL', 'VertM(m/s)': 'W'
   }
   data = data.rename(columns=vars_new_names)

   return data, coords

df, coords = read_vortex_times_pandas(file)
print(coords)
df

df['M'].describe()

import matplotlib.pyplot as plt

# Plot the 300 first timestamps of the wind speed
df['M'].iloc[:300].plot()
plt.show()

# Plot the wind speed histogram for the full period for 0.2m/s bins
df['M'].plot.hist(bins=np.arange(0, 30, 0.2), density=True)
plt.show()

import xarray as xr

def convert_to_xarray(df: pd.DataFrame,
                      coords: Dict[str, float] = None) -> xr.Dataset:
    """
    Convert a dataframe to a xarray object.

    Parameters
    ----------
    df: pd.DataFrame
    coords: Dict[str, float]
        Info about lat, lon, lev so that the new dimensions can be added

    Returns
    -------
    xr.Dataset
        With added coordinates
    """
    # Simply use the to_xarray function to convert from dataframe to dataset
    # It sets the dataframe's index as the dataset's coordinates
    ds: xr.Dataset = df.to_xarray()
    
    # Add the other coordinates (which were not the dataframe indices: lat, lon, lev)
    if coords is not None:
        coords_dict = {name: [float(val)] for name, val in coords.items()
                       if name not in ds.dims}
        ds = ds.assign_coords(coords_dict)
        
    return ds

ds = convert_to_xarray(df, coords=coords)
print(ds)

# I have downloaded files for the Sample TIMES run for heights 10, 50, 100, 150 and 200m
# All of them at local UTC, and stored in the folder `example_txt` with the default names
# Therefore, all 5 files that I want to merge have this name:
file_template = 'example_txt/vortex.times.677619.6m <LEV>m UTC+02.0 ERA5.txt'
# with <LEV> representing the height in meters

# We will do the dataset for each height, and store them in this list
all_heights = []
for lev in [10, 50, 100, 150, 200]:
    # the name of the file can be deduced by changing <LEV> to the height value
    file_lev = file_template.replace('<LEV>', str(lev))
    # I use the read_vortex_times_pandas pandas to read the dataframe and the coordinates
    df_lev, coords_lev = read_vortex_times_pandas(file_lev)
    # Then, I convert the dataframe to an xarray dataset (that includes info about the lev)
    ds_lev = convert_to_xarray(df_lev, coords=coords_lev)
    # Append this dataset to the list
    all_heights.append(ds_lev)
    
# Concatenate each dataset of the list over the 'lev' dimension
ds_full = xr.concat(all_heights, dim='lev')
print(ds_full)

# Plot the firsts timestamps for all heights
ds_full['M'].isel(time=range(200)).plot(hue='lev')
plt.show()

# Plot the mean wind speed for each month at all heights
ds_full['M'].groupby('time.month').mean(dim='time').plot.scatter(hue='lev')
plt.show()

# Interpolate at height 90m + do the monthly means + plot the time series
ds_full['M'].interp(lev=90).resample(time='M').mean().plot()
plt.show()

# Plot the wind standard deviation at all heights for one day
ds_full['SD'].sel(time='2012-08-01').plot()
plt.show()