import numpy as np
import pandas as pd
from datetime import datetime
import xarray as xr
import urllib.request
import concurrent.futures
import re
from tqdm import tqdm
import os
from os.path import isfile, join, exists
from dataclasses import dataclass
import warnings

folder = "../data/raw/gdp-v2.00/"
aoml_https_url = "https://www.aoml.noaa.gov/ftp/pub/phod/lumpkin/hourly/v2.00/netcdf/"
file_pattern = "drifter_{id}.nc"
os.makedirs(folder, exist_ok=exists(folder))  # create raw data folder


def fetch_netcdf(url, file):
    """
    Download and save file from the given url (if not present)
    """
    if not isfile(file):
        req = urllib.request.urlretrieve(url, file)


def download(drifter_ids: list = None, n_random_id: int = None):
    """
    Download individual NetCDF files from the AOML server

    :param drifter_ids [list]: list of drifter to retrieve (Default: all)
    :param n_random_id [int]: randomly select n drifter NetCDF files
    :return drifters_ids [list]: list of retrived drifter
    """
    # retrieve all drifter ID numbers
    if drifter_ids == None:
        urlpath = urllib.request.urlopen(aoml_https_url)
        string = urlpath.read().decode("utf-8")
        pattern = re.compile("drifter_[0-9]*.nc")
        filelist = pattern.findall(string)
        drifter_ids = np.unique([int(f.split("_")[-1][:-3]) for f in filelist])

    # retrieve only a subset of n_random_id trajectories
    if n_random_id:
        if n_random_id > len(drifter_ids):
            warnings.warn(
                f"Retrieving all listed trajectories because {n_random_id} is larger than the {len(drifter_ids)} listed trajectories."
            )
        else:
            rng = np.random.RandomState(42)
            drifter_ids = sorted(rng.choice(drifter_ids, n_random_id, replace=False))

    with concurrent.futures.ThreadPoolExecutor() as exector:
        # create list of urls and paths
        urls = []
        files = []
        for i in drifter_ids:
            file = file_pattern.format(id=i)
            urls.append(join(aoml_https_url, file))
            files.append(join(folder, file))

        # parallel retrieving of individual netCDF files
        list(
            tqdm(
                exector.map(fetch_netcdf, urls, files),
                total=len(files),
                desc="Downloading files",
                ncols=80,
            )
        )

    return drifter_ids


def decode_date(t):
    """
    The date format is specified in 'seconds since 1970-01-01 00:00:00' but the missing values
    are stored as -1e+34 which is not supported by the default parsing mechanism in xarray

    This function returns replaced the missing valye by NaT and return a datetime object.
    :param t: date
    :return: datetime object
    """
    nat_index = np.logical_or(np.isclose(t, -1e34), np.isnan(t))
    t[nat_index] = np.nan
    return t


def fill_values(var, default=np.nan):
    """
    Change fill values (-1e+34, inf, -inf) in var array to value specifed by default
    """
    missing_value = np.logical_or(np.isclose(var, -1e34), ~np.isfinite(var))
    if np.any(missing_value):
        var[missing_value] = default
    return var


def str_to_float(value, default=np.nan):
    """
    :param value: string
    :return: bool
    """
    try:
        fvalue = float(value)
        if np.isnan(fvalue):
            return default
        else:
            return fvalue
    except ValueError:
        return default


def cut_str(value, max_length):
    """
    Cut a string to a specific length.
    :param value: string
           max_length: length of the output
    :return: string with max_length chars
    """
    charar = np.chararray(1, max_length)
    charar[:max_length] = value
    return charar


def drogue_presence(lost_time, time):
    """
    Create drogue status from the drogue lost time and the trajectory time
    :params lost_time: timestamp of the drogue loss (or NaT)
            time[obs]: observation time
    :return: bool[obs]: 1 drogued, 0 undrogued
    """
    if pd.isnull(lost_time) or lost_time >= time[-1]:
        return np.ones_like(time, dtype="bool")
    else:
        return time < lost_time


def rowsize(index: int) -> int:
    return xr.open_dataset(
        join(folder, file_pattern.format(id=index)),
        decode_cf=False,
        decode_times=False,
        concat_characters=False,
        decode_coords=False,
    ).dims["obs"]


def preprocess(index: int) -> xr.Dataset:
    """
    Mandatory function that extract and preprocess the Lagragangian data and attributes. The function takes and
    identification number that can be used to: create a file or url pattern or select data from a Dataframe. It
    then preprocess the data and return a clean xarray Dataset.

    :param index: drifter's identification number
    :return: xr.Dataset containing the data and attributes
    """
    ds = xr.load_dataset(
        join(folder, file_pattern.format(id=index)),
        decode_times=False,
        decode_coords=False,
    )

    # parse the date with custom function
    ds["deploy_date"].data = decode_date(np.array([ds.deploy_date.data[0]]))
    ds["end_date"].data = decode_date(np.array([ds.end_date.data[0]]))
    ds["drogue_lost_date"].data = decode_date(np.array([ds.drogue_lost_date.data[0]]))
    ds["time"].data = decode_date(np.array([ds.time.data[0]]))

    # convert fill values to nan
    ds["err_lon"].data = fill_values(ds["err_lon"].data)
    ds["err_lat"].data = fill_values(ds["err_lat"].data)
    ds["err_ve"].data = fill_values(ds["err_ve"].data)
    ds["err_vn"].data = fill_values(ds["err_vn"].data)
    ds["sst"].data = fill_values(ds["sst"].data)
    ds["sst1"].data = fill_values(ds["sst1"].data)
    ds["sst2"].data = fill_values(ds["sst2"].data)
    ds["err_sst"].data = fill_values(ds["err_sst"].data)
    ds["err_sst1"].data = fill_values(ds["err_sst1"].data)
    ds["err_sst2"].data = fill_values(ds["err_sst2"].data)

    # fix missing values stored as str
    for var in [
        "longitude",
        "latitude",
        "err_lat",
        "err_lon",
        "ve",
        "vn",
        "err_ve",
        "err_vn",
        "sst",
        "sst1",
        "sst2",
    ]:
        ds[var].encoding["missing value"] = -1e-34

    # convert type of some variable
    ds["ID"].data = ds["ID"].data.astype("int64")
    ds["WMO"].data = ds["WMO"].data.astype("int32")
    ds["expno"].data = ds["expno"].data.astype("int32")
    ds["typedeath"].data = ds["typedeath"].data.astype("int8")
    ds["flg_sst"].data = ds["flg_sst"].data.astype("int8")
    ds["flg_sst1"].data = ds["flg_sst1"].data.astype("int8")
    ds["flg_sst2"].data = ds["flg_sst2"].data.astype("int8")

    # new variables
    ds["ids"] = (["traj", "obs"], [np.repeat(ds.ID.values, ds.dims["obs"])])
    ds["drogue_status"] = (
        ["traj", "obs"],
        [drogue_presence(ds.drogue_lost_date.data, ds.time.data[0])],
    )

    # convert attributes to variable
    ds["location_type"] = (
        ("traj"),
        [False if ds.location_type == "Argos" else True],
    )  # 0 for Argos, 1 for GPS
    ds["DeployingShip"] = (("traj"), cut_str(ds.DeployingShip, 20))
    ds["DeploymentStatus"] = (("traj"), cut_str(ds.DeploymentStatus, 20))
    ds["BuoyTypeManufacturer"] = (("traj"), cut_str(ds.BuoyTypeManufacturer, 20))
    ds["BuoyTypeSensorArray"] = (("traj"), cut_str(ds.BuoyTypeSensorArray, 20))
    ds["CurrentProgram"] = (("traj"), np.int32([str_to_float(ds.CurrentProgram, -1)]))
    ds["PurchaserFunding"] = (("traj"), cut_str(ds.PurchaserFunding, 20))
    ds["SensorUpgrade"] = (("traj"), cut_str(ds.SensorUpgrade, 20))
    ds["Transmissions"] = (("traj"), cut_str(ds.Transmissions, 20))
    ds["DeployingCountry"] = (("traj"), cut_str(ds.DeployingCountry, 20))
    ds["DeploymentComments"] = (
        ("traj"),
        cut_str(ds.DeploymentComments.encode("ascii", "ignore").decode("ascii"), 20),
    )  # remove non ascii char
    ds["ManufactureYear"] = (("traj"), np.int16([str_to_float(ds.ManufactureYear, -1)]))
    ds["ManufactureMonth"] = (
        ("traj"),
        np.int16([str_to_float(ds.ManufactureMonth, -1)]),
    )
    ds["ManufactureSensorType"] = (("traj"), cut_str(ds.ManufactureSensorType, 20))
    ds["ManufactureVoltage"] = (
        ("traj"),
        np.int16([str_to_float(ds.ManufactureVoltage[:-6], -1)]),
    )  # e.g. 56 V
    ds["FloatDiameter"] = (
        ("traj"),
        [str_to_float(ds.FloatDiameter[:-3])],
    )  # e.g. 35.5 cm
    ds["SubsfcFloatPresence"] = (
        ("traj"),
        np.array([str_to_float(ds.SubsfcFloatPresence)], dtype="bool"),
    )
    ds["DrogueType"] = (("traj"), cut_str(ds.DrogueType, 7))
    ds["DrogueLength"] = (("traj"), [str_to_float(ds.DrogueLength[:-2])])  # e.g. 4.8 m
    ds["DrogueBallast"] = (
        ("traj"),
        [str_to_float(ds.DrogueBallast[:-3])],
    )  # e.g. 1.4 kg
    ds["DragAreaAboveDrogue"] = (
        ("traj"),
        [str_to_float(ds.DragAreaAboveDrogue[:-4])],
    )  # 10.66 m^2
    ds["DragAreaOfDrogue"] = (
        ("traj"),
        [str_to_float(ds.DragAreaOfDrogue[:-4])],
    )  # e.g. 416.6 m^2
    ds["DragAreaRatio"] = (("traj"), [str_to_float(ds.DragAreaRatio)])  # e.g. 39.08
    ds["DrogueCenterDepth"] = (
        ("traj"),
        [str_to_float(ds.DrogueCenterDepth[:-2])],
    )  # e.g. 20.0 m
    ds["DrogueDetectSensor"] = (("traj"), cut_str(ds.DrogueDetectSensor, 20))

    # vars attributes
    vars_attrs = {
        "ID": {"long_name": "Global Drifter Program Buoy ID", "units": "-"},
        "longitude": {"long_name": "Longitude", "units": "degrees_east"},
        "latitude": {"long_name": "Latitude", "units": "degrees_north"},
        "time": {"long_name": "Time", "units": "seconds since 1970-01-01 00:00:00"},
        "ids": {
            "long_name": "Global Drifter Program Buoy identification number repeated along observations",
            "units": "-",
        },
        "rowsize": {
            "long_name": "Number of observations per trajectory",
            "sample_dimension": "obs",
            "units": "-",
        },
        "location_type": {
            "long_name": "Satellite-based location system",
            "units": "-",
            "comments": "0 (Argos), 1 (GPS)",
        },
        "WMO": {
            "long_name": "World Meteorological Organization buoy identification number",
            "units": "-",
        },
        "expno": {"long_name": "Experiment number", "units": "-"},
        "deploy_date": {
            "long_name": "Deployment date and time",
            "units": "seconds since 1970-01-01 00:00:00",
        },
        "deploy_lon": {"long_name": "Deployment longitude", "units": "degrees_east"},
        "deploy_lat": {"long_name": "Deployment latitude", "units": "degrees_north"},
        "end_date": {
            "long_name": "End date and time",
            "units": "seconds since 1970-01-01 00:00:00",
        },
        "end_lon": {"long_name": "End latitude", "units": "degrees_north"},
        "end_lat": {"long_name": "End longitude", "units": "degrees_east"},
        "drogue_lost_date": {
            "long_name": "Date and time of drogue loss",
            "units": "seconds since 1970-01-01 00:00:00",
        },
        "typedeath": {
            "long_name": "Type of death",
            "units": "-",
            "comments": "0 (buoy still alive), 1 (buoy ran aground), 2 (picked up by vessel), 3 (stop transmitting), 4 (sporadic transmissions), 5 (bad batteries), 6 (inactive status)",
        },
        "typebuoy": {
            "long_name": "Buoy type (see https://www.aoml.noaa.gov/phod/dac/dirall.html)",
            "units": "-",
        },
        "DeployingShip": {"long_name": "Name of deployment ship", "units": "-"},
        "DeploymentStatus": {"long_name": "Deployment status", "units": "-"},
        "BuoyTypeManufacturer": {"long_name": "Buoy type manufacturer", "units": "-"},
        "BuoyTypeSensorArray": {"long_name": "Buoy type sensor array", "units": "-"},
        "CurrentProgram": {
            "long_name": "Current Program",
            "units": "-",
            "_FillValue": "-1",
        },
        "PurchaserFunding": {"long_name": "Purchaser funding", "units": "-"},
        "SensorUpgrade": {"long_name": "Sensor upgrade", "units": "-"},
        "Transmissions": {"long_name": "Transmissions", "units": "-"},
        "DeployingCountry": {"long_name": "Deploying country", "units": "-"},
        "DeploymentComments": {"long_name": "Deployment comments", "units": "-"},
        "ManufactureYear": {
            "long_name": "Manufacture year",
            "units": "-",
            "_FillValue": "-1",
        },
        "ManufactureMonth": {
            "long_name": "Manufacture month",
            "units": "-",
            "_FillValue": "-1",
        },
        "ManufactureSensorType": {"long_name": "Manufacture Sensor Type", "units": "-"},
        "ManufactureVoltage": {
            "long_name": "Manufacture voltage",
            "units": "V",
            "_FillValue": "-1",
        },
        "FloatDiameter": {"long_name": "Diameter of surface floater", "units": "cm"},
        "SubsfcFloatPresence": {"long_name": "Subsurface Float Presence", "units": "-"},
        "DrogueType": {"drogue_type": "Drogue Type", "units": "-"},
        "DrogueLength": {"long_name": "Length of drogue.", "units": "m"},
        "DrogueBallast": {
            "long_name": "Weight of the drogue's ballast.",
            "units": "kg",
        },
        "DragAreaAboveDrogue": {"long_name": "Drag area above drogue.", "units": "m^2"},
        "DragAreaOfDrogue": {"long_name": "Drag area drogue.", "units": "m^2"},
        "DragAreaRatio": {"long_name": "Drag area ratio", "units": "m"},
        "DrogueCenterDepth": {
            "long_name": "Average depth of the drogue.",
            "units": "m",
        },
        "DrogueDetectSensor": {"long_name": "Drogue detection sensor", "units": "-"},
        "ve": {"long_name": "Eastward velocity", "units": "m/s"},
        "vn": {"long_name": "Northward velocity", "units": "m/s"},
        "gap": {
            "long_name": "Time interval between previous and next location",
            "units": "s",
        },
        "err_lat": {
            "long_name": "95% confidence interval in latitude",
            "units": "degrees_north",
        },
        "err_lon": {
            "long_name": "95% confidence interval in longitude",
            "units": "degrees_east",
        },
        "err_ve": {
            "long_name": "95% confidence interval in eastward velocity",
            "units": "m/s",
        },
        "err_vn": {
            "long_name": "95% confidence interval in northward velocity",
            "units": "m/s",
        },
        "drogue_status": {
            "long_name": "Status indicating the presence of the drogue",
            "units": "-",
            "flag_values": "1,0",
            "flag_meanings": "drogued, undrogued",
        },
        "sst": {
            "long_name": "Fitted sea water temperature",
            "units": "Kelvin",
            "comments": "Estimated near-surface sea water temperature from drifting buoy measurements. It is the sum of the fitted near-surface non-diurnal sea water temperature and fitted diurnal sea water temperature anomaly. Discrepancies may occur because of rounding.",
        },
        "sst1": {
            "long_name": "Fitted non-diurnal sea water temperature",
            "units": "Kelvin",
            "comments": "Estimated near-surface non-diurnal sea water temperature from drifting buoy measurements",
        },
        "sst2": {
            "long_name": "Fitted diurnal sea water temperature anomaly",
            "units": "Kelvin",
            "comments": "Estimated near-surface diurnal sea water temperature anomaly from drifting buoy measurements",
        },
        "err_sst": {
            "long_name": "Standard uncertainty of fitted sea water temperature",
            "units": "Kelvin",
            "comments": "Estimated one standard error of near-surface sea water temperature estimate from drifting buoy measurements",
        },
        "err_sst1": {
            "long_name": "Standard uncertainty of fitted non-diurnal sea water temperature",
            "units": "Kelvin",
            "comments": "Estimated one standard error of near-surface non-diurnal sea water temperature estimate from drifting buoy measurements",
        },
        "err_sst2": {
            "long_name": "Standard uncertainty of fitted diurnal sea water temperature anomaly",
            "units": "Kelvin",
            "comments": "Estimated one standard error of near-surface diurnal sea water temperature anomaly estimate from drifting buoy measurements",
        },
        "flg_sst": {
            "long_name": "Fitted sea water temperature quality flag",
            "units": "-",
            "flag_values": "0, 1, 2, 3, 4, 5",
            "flag_meanings": "no-estimate, no-uncertainty-estimate, estimate-not-in-range-uncertainty-not-in-range, estimate-not-in-range-uncertainty-in-range estimate-in-range-uncertainty-not-in-range, estimate-in-range-uncertainty-in-range",
        },
        "flg_sst1": {
            "long_name": "Fitted non-diurnal sea water temperature quality flag",
            "units": "-",
            "flag_values": "0, 1, 2, 3, 4, 5",
            "flag_meanings": "no-estimate, no-uncertainty-estimate, estimate-not-in-range-uncertainty-not-in-range, estimate-not-in-range-uncertainty-in-range estimate-in-range-uncertainty-not-in-range, estimate-in-range-uncertainty-in-range",
        },
        "flg_sst2": {
            "long_name": "Fitted diurnal sea water temperature anomaly quality flag",
            "units": "-",
            "flag_values": "0, 1, 2, 3, 4, 5",
            "flag_meanings": "no-estimate, no-uncertainty-estimate, estimate-not-in-range-uncertainty-not-in-range, estimate-not-in-range-uncertainty-in-range estimate-in-range-uncertainty-not-in-range, estimate-in-range-uncertainty-in-range",
        },
    }

    # global attributes
    attrs = {
        "title": "Global Drifter Program hourly drifting buoy collection",
        "history": "Version 2.00.  Metadata from dirall.dat and deplog.dat",
        "Conventions": "CF-1.6",
        "date_created": datetime.now().isoformat(),
        "publisher_name": "GDP Drifter DAC",
        "publisher_email": "aoml.dftr@noaa.gov",
        "publisher_url": "https://www.aoml.noaa.gov/phod/gdp",
        "licence": "MIT License",
        "processing_level": "Level 2 QC by GDP drifter DAC",
        "metadata_link": "https://www.aoml.noaa.gov/phod/dac/dirall.html",
        "contributor_name": "NOAA Global Drifter Program",
        "contributor_role": "Data Acquisition Center",
        "institution": "NOAA Atlantic Oceanographic and Meteorological Laboratory",
        "acknowledgement": "Elipot et al. (2022) to be submitted. Elipot et al. (2016). Global Drifter Program quality-controlled hourly interpolated data from ocean surface drifting buoys, version 2.00. NOAA National Centers for Environmental Information. https://agupubs.onlinelibrary.wiley.com/doi/full/10.1002/2016JC011716TBA. Accessed [date].",
        "summary": "Global Drifter Program hourly data",
    }

    # set attributes
    for var in vars_attrs.keys():
        ds[var].attrs = vars_attrs[var]
    ds.attrs = attrs

    # remove vars
    ds = ds.drop_vars("lon360")

    # set coordinates
    ds = ds.set_coords(["ids", "longitude", "latitude", "time"])

    return ds
