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

@author: ormondt
"""

import os

import numpy as np
import pandas as pd
import xarray as xr
from pyproj import CRS
from siphon.catalog import TDSCatalog
from xarray.backends import NetCDF4DataStore


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/catalog/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
    # 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")
        url = base_url + month_string + "/" + date_string + "/catalog.xml"
        try:
            gfs = TDSCatalog(url)
        except Exception:
            gfs = []
        cstr = "0000_000"
        name = "gfsanl_4_" + date_string + "_" + cstr + ".grb2"
        okay = False
        if gfs:
            for j, ds in enumerate(gfs.datasets):
                if ds == name:
                    okay = True
                    break
        if not okay:
            # Try the next time
            continue
        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.load_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

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

    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 Exception:
            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)
                        with xr.open_dataset(NetCDF4DataStore(data)) as 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)
                        with xr.open_dataset(NetCDF4DataStore(data)) as 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 Exception:
                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)
