# -*- coding: utf-8 -*-
"""
Created on Thu May 20 10:32:33 2021

@author: ormondt
"""

import os
import xarray as xr
import pandas as pd
import numpy as np
from pyproj import CRS
from metpy.units import units

class Dataset():

    def __init__(self):

        self.quantity = None
        self.unit     = None
        self.time     = []
        self.x        = None
        self.y        = None
        self.crs      = None
        self.val      = None
        self.u        = None
        self.v        = None

def download(param_list, lon_range, lat_range, path, prefix,
             time_range=None, times=None):

    base_url = "https://www.ncei.noaa.gov/thredds/dodsC/model-gfs-g4-anl-files-old/"

    if times is not None:
        requested_times = times
        time_range = [times[0], times[-1]]
    else:    
        requested_times = pd.date_range(start=time_range[0],
                          end=time_range[1],
                          freq='3H').to_pydatetime().tolist()

    ntime = len(requested_times)

    datasets = []
    for param in param_list:        
        dataset = Dataset()
        dataset.crs = CRS.from_epsg(4326)
        dataset.quantity = param
        datasets.append(dataset)

    icont = False
    
    var_wind_u   = "u-component_of_wind_height_above_ground"
    var_wind_v   = "v-component_of_wind_height_above_ground"
    var_pressure = "Pressure_reduced_to_MSL_msl"
        
    # Get lat,lon
    for it, time in enumerate(requested_times):
        
        h            = requested_times[it].hour
        month_string = requested_times[it].strftime("%Y%m")
        date_string  = requested_times[it].strftime("%Y%m%d")
        time_string  = requested_times[it].strftime("%Y%m%d")
        url = base_url + month_string + "/" + date_string + "/"

        if h==0:
            cstr  = "0000_000"
            crstr = "0000_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==3:    
            cstr  = "0000_003"
            crstr = "0000_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        elif h==6:    
            cstr  = "0600_000"
            crstr = "0600_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==9:    
            cstr  = "0600_003"
            crstr = "0600_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        elif h==12:    
            cstr  = "1200_000"
            crstr = "1200_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==15:    
            cstr  = "1200_003"
            crstr = "1200_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        elif h==18:    
            cstr  = "1800_000"
            crstr = "1800_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==21:    
            cstr  = "1800_003"
            crstr = "1800_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        else:
            print('ERROR: cycle and cycle_last in xml-file can only be specified rounded off on 3 hour values!')
                
        url  = base_url + "gfsanl_4_" + date_string + "_" + cstr + ".grb2"

        ds = xr.open_dataset(url)  # or engine='pydap'

        i1 = np.where(ds.lat.values>lat_range[1])[0][-1]
        i2 = np.where(ds.lat.values<lat_range[0])[0][0]
        j1 = np.where(ds.lon.values<lon_range[0])[0][-1]
        j2 = np.where(ds.lon.values>lon_range[1])[0][0]

        okay = False        

        # Latitude and longitude
        lat = ds.lat.values[i1:i2]
        lon = ds.lon.values[j1:j2]

        nrows = len(lat)
        ncols = len(lon)
        
        # Now loop through parameters

        # initialize matrices
        for dataset in datasets:        
            dataset.x = lon
            dataset.y = lat
            if dataset.quantity == "wind":                    
                dataset.u    = np.empty((nrows, ncols))
                dataset.u[:] = np.NaN
                dataset.v    = np.empty((nrows, ncols))                
                dataset.v[:] = np.NaN
            else:
                dataset.val    = np.empty((nrows, ncols))
                dataset.val[:] = np.NaN

        # ncss  = gfs.datasets[j].subset()                
        # query = ncss.query()            
        # query.lonlat_box(north=lat_range[1], \
        #                  south=lat_range[0], \
        #                   east=lon_range[1],  \
        #                   west=lon_range[0]).vertical_level(10.0)
        # query.variables("u-component_of_wind_height_above_ground",
        #                 "v-component_of_wind_height_above_ground")
        # data = ncss.get_data(query)
        # data = xr.open_dataset(NetCDF4DataStore(data))        
        # lon = np.array(data['lon'])
        # lat = np.array(data['lat'])  
        # nrows = len(lat)
        # ncols = len(lon)
        # # Latitude and longitude found, so we can stop now
        # icont = True
        # break
    
    if not icont:
        # Could not find any data
        print("Could not find any data in requested range !")
        datasets = []
        return datasets
    
    
    for it, time in enumerate(requested_times):
        
        h            = time.hour
        month_string = time.strftime("%Y%m")
        date_string  = time.strftime("%Y%m%d")
        url = base_url + month_string + "/" + date_string + "/catalog.xml"
        
        # try:
        #     gfs   = TDSCatalog(url)
        # except:
        #     print("Could not fetch catalogue")
        #     continue
        
        if h==0:
            cstr  = "0000_000"
            crstr = "0000_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==3:    
            cstr  = "0000_003"
            crstr = "0000_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        elif h==6:    
            cstr  = "0600_000"
            crstr = "0600_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==9:    
            cstr  = "0600_003"
            crstr = "0600_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        elif h==12:    
            cstr  = "1200_000"
            crstr = "1200_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==15:    
            cstr  = "1200_003"
            crstr = "1200_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        elif h==18:    
            cstr  = "1800_000"
            crstr = "1800_003"
            var_prcp = "Total_precipitation_surface_3_Hour_Accumulation"
        elif h==21:    
            cstr  = "1800_003"
            crstr = "1800_006"
            var_prcp = "Total_precipitation_surface_6_Hour_Accumulation"
        else:
            print('ERROR: cycle and cycle_last in xml-file can only be specified rounded off on 3 hour values!')

        # Loop through requested parameters
        for ind, param in enumerate(param_list):
            
            dataset = datasets[ind]
            dataset.time.append(time)

            if param == "precipitation":
                name = "gfsanl_4_" + date_string + "_" + crstr + ".grb2" 
            else:    
                name = "gfsanl_4_" + date_string + "_" + cstr + ".grb2" 

            try:
                            
                okay = False
                makezeros = False
                for j, ds in enumerate(gfs.datasets):
                    if ds==name:
                        okay = True
                        break
                
                if not okay:
                    # File not found, on to the next parameter
                    print("Warning! " + name + " was not found on server")
                    makezeros = True
                    #continue
    
                print(name + " : " + param)    
    
                ncss  = gfs.datasets[j].subset()

                if okay:
                    if param == "wind":
                        query = ncss.query()
                        query.lonlat_box(north=lat_range[1], \
                                         south=lat_range[0], \
                                          east=lon_range[1],  \
                                          west=lon_range[0]).vertical_level(10.0)
                        query.variables("u-component_of_wind_height_above_ground",
                                        "v-component_of_wind_height_above_ground")
                        data = ncss.get_data(query)
                        data = xr.open_dataset(NetCDF4DataStore(data))
                        u = data['u-component_of_wind_height_above_ground']
                        v = data['v-component_of_wind_height_above_ground']
                        dataset.unit = u.units
                        u = u.metpy.unit_array.squeeze()
                        dataset.u[it,:,:] = np.array(u)
                        v = v.metpy.unit_array.squeeze()
                        dataset.v[it,:,:] = np.array(v)

                    else:

                        # Other scalar variables
        #                fac = 1.0
                        if param == "barometric_pressure":
                            var_name = "Pressure_reduced_to_MSL_msl"
                        elif param == "precipitation":
                            var_name = var_prcp
        #                    fac = 1.0
                        query = ncss.query()
                        query.lonlat_box(north=lat_range[1], \
                                         south=lat_range[0], \
                                         east=lon_range[1],  \
                                         west=lon_range[0])
                        query.variables(var_name)
                        data = ncss.get_data(query)
                        data = xr.open_dataset(NetCDF4DataStore(data))
                        val          = data[var_name]
                        dataset.unit = val.units
                        val          = np.array(val.metpy.unit_array.squeeze())
                        if param == "precipitation":
                            # Data is stored either as 3-hourly (at 03h) or 6-hourly (at 06h) accumulated rainfall
                            # For the first, just divide by 3 to get hourly precip
                            # For the second, first subtract volume that fell in the first 3 hours
                            if h==0 or h==6 or h==12 or h==18:
                                val = val/3 # Convert to mm/h
                            else:
                                val = (val - 3*np.squeeze(dataset.val[it-1,:,:]))/3
                        dataset.val[it,:,:]  = val

                elif makezeros: # add zeros
                    if param == "wind":
                        dataset.u[:] = 0
                        dataset.v[:] = 0
                        print(param + " was not found on server ... --> using 0.0 m/s instead !!!")

                    if param == "precipitation":
                        dataset.val[:] = 0
                        print(param + " was not found on server ... --> using 0.0 m/s instead !!!")

                    if param == "barometric_pressure":
                        dataset.val[:] = 102000.0
                        print(param + " was not found on server ... --> using 102000.0 Pa instead !!!")


            except:
                
                print("Could not download data")

        # Write data to netcdf
        time_string = time.strftime("%Y%m%d_%H%M")
        file_name = prefix + "." + time_string + ".nc"
        full_file_name = os.path.join(path, file_name)

        ds = xr.Dataset()
             
        okay = False
        for ind, dataset in enumerate(datasets):

            if dataset.quantity == "wind":

                uu = dataset.u[it,:,:]
                vv = dataset.v[it,:,:]
                
                if not np.any(np.isnan(uu)) and not np.any(np.isnan(vv)):
                    
                    okay = True

                    da = xr.DataArray(uu,
                                          coords=[("lat", dataset.y),
                                                  ("lon", dataset.x)])
                    ds["wind_u"] = da
    
                    da = xr.DataArray(vv,
                                          coords=[("lat", dataset.y),
                                                  ("lon", dataset.x)])
                    ds["wind_v"] = da

            else:

                val = dataset.val[it,:,:]

                if not np.any(np.isnan(val)):
                    
                    okay = True
                    
                    da = xr.DataArray(val,
                                          coords=[("lat", dataset.y),
                                                  ("lon", dataset.x)])
                    ds[dataset.quantity] = da
        if okay:
            # Only write to file if there is any data
            ds.to_netcdf(path=full_file_name)
    
    return datasets

# Helper function for finding proper time variable
def find_time_var(var, time_basename='time'):
    for coord_name in var.coords:
        if coord_name.startswith(time_basename):
            return var.coords[coord_name]
    raise ValueError('No time variable found for ' + var.name)

# Helper function for finding proper time variable
def find_height_var(var, time_basename='height'):
    for coord_name in var.coords:
        if coord_name.startswith(time_basename):
            return var.coords[coord_name]
    raise ValueError('No height variable found for ' + var.name)
