"""
Set of utilities for decoding and analysing grib files.
"""

import glob
import os
import logging
import logging.config
import numpy as np
import pygrib as pg
import astropy.units as u
from astropy.coordinates import Angle, Latitude, Longitude
from astropy.time import Time
from astropy.table import Table, join, vstack
from molecularprofiles.utils.constants import (
    STD_NUMBER_DENSITY,
    STD_AIR_PRESSURE,
    STD_AIR_TEMPERATURE,
    DENSITY_SCALE_HEIGHT,
    STD_GRAVITATIONAL_ACCELERATION,
    STD_EARTH_RADIUS,
)

ROOTDIR = os.path.dirname(os.path.abspath(__file__))
log_config_file = f"{ROOTDIR}/mdps_log.conf"
logging.config.fileConfig(fname=log_config_file, disable_existing_loggers=False)
logger = logging.getLogger(__name__)


def get_altitude_from_geopotential_height(geopotential_height, latitude):
    """
    Function to compute the real altitude from the geopotential value at
    a certain coordinate on Earth.
    Uses expression 20 from http://dx.doi.org/10.1029/2002JB002333

    Parameters
    ----------
    geopotential : astropy.units.Quantity [u.m]
        Geopotential height
    latitude : astropy.coordinates.Latitude
        Geographical latitude of interest.

    Returns
    -------
    astropy.units.Quantity [u.m]
        Real altitude as fGeoidOffset.
    """

    lat = (np.asarray(latitude) * latitude.unit).to_value(u.rad)
    return (
        (1.0 + 0.002644 * np.cos(2 * lat)) * (geopotential_height)
        + (1 + 0.0089 * np.cos(2 * lat)) * ((geopotential_height) ** 2 / STD_EARTH_RADIUS.to(u.m))
    ).to(u.m)


def get_gribfile_variables(filename):
    """
    Returns all the different variable names in a grib file.

    Parameters
    ----------
    filename : str
       Path to the grib file.

    Returns
    -------
    list
        varname (str): variable names.
    list
        varshortname (str): variable short names.
    """
    logger.debug("Opening %s grib file", filename)
    with pg.open(filename) as grib_file:
        varshortname = []
        varname = []
        grib_file.read(10)[0]
        while True:
            v = grib_file.read(1)[0]
            variable_name = v.name.replace(" ", "")
            variable_short_name = v.shortName
            if variable_name not in varname:
                varname.append(variable_name)
                varshortname.append(variable_short_name)
            else:
                break
    return varname, varshortname


def get_plevels(atmo_parameter):
    """
    The measurements of atmospheres parameters are taking place in various pressure levels.
    This functions returns the indices of these levels where we have measurements
    for the given atmospheric parameter.
    E.g. for GDAS data that should be a range of values [1, 23].

    Parameters
    ----------
    atmo_parameter : str FIXME: this description doesn't
        Atmospheric parameter from grib file.

    Returns
    -------
    list
        List (int) with the indices of the pressure levels
    """

    plevels = []
    index = []
    for i in range(len(atmo_parameter)):
        pressure_level = atmo_parameter[i].level
        if pressure_level not in plevels:
            plevels.append(pressure_level)
            index.append(i + 1)
        else:
            break
    return index


def create_table(grib_var):
    """
    Create astropy.table.Table from grib record

    Parameters
    ----------
    grib_var :
        Grib data record

     Returns
     -------
     astropy.table.Table
         Table with the different measurables together with their dimensions
    """
    res = Table()
    for v in grib_var:
        unit = u.Unit(v.units)
        timestamp = Time(
            {"year": v.year, "month": v.month, "day": v.day, "hour": v.hour}, scale="utc"
        )
        pressure_level = v.level * u.hPa
        latitudes = Latitude(v.latlons()[0].ravel() * u.deg)
        longitudes = Longitude(v.latlons()[1].ravel() * u.deg, wrap_angle=180 * u.deg)
        if isinstance(v.values, float):
            vals = np.array([v.values]) * unit
        else:
            vals = v.values.ravel() * unit
        t = Table([latitudes, longitudes], names=["Latitude", "Longitude"])
        t["Timestamp"] = timestamp
        t["Pressure"] = pressure_level
        t[v.name] = vals
        res = vstack([res, t])
    return res


def get_grib_file_data(filename):
    """
    This function opens a grib file, selects the parameters
    (all available: Temperature, Geopotential, RH, ...),
    and creates an astropy.table.Table with them.
    Parameters
    ----------
    filename : str
        Path to the grib file.

    Returns
    -------
    astropy.table.Table
        Table with grib file data gridded in isobaric levels
    """

    logger.debug("Working with %s", filename)
    logger.debug("Getting all variable names...")
    variable_names, variable_short_names = get_gribfile_variables(filename)
    logger.info("Indexing the file %s (this might take a while...)", filename)
    grib_file = pg.index(filename, "shortName", "typeOfLevel")
    logger.debug(
        "Selecting the parameters information for %s (this might take a while...)", filename
    )
    gpm = u.def_unit("gpm", u.m)
    u.add_enabled_units([gpm])
    data = Table()
    for short_name in variable_short_names:
        if short_name == "unknown":
            continue
        var = grib_file.select(shortName=short_name, typeOfLevel="isobaricInhPa")
        try:
            t = create_table(var)
        except ValueError:
            logger.warning("Grib message can't be parsed, skipping")
            continue
        if len(data) == 0:
            data = t
        else:
            data = join(data, t, join_type="outer")
    return data


def extend_grib_data(data):
    """
    Extends grib data table by filling the gaps in data and calculating additional quantities:
        - altitude
        - density
        - exponential density
        - wind direction

    Parameters
    ----------
    astropy.table.Table
        Table with grib data
    Returns
    -------
    astropy.table.Table
        Extended table with grib data and additional quantities
    """
    logger.debug("Check for gaps in relative humidity and fill them if neccessary")
    data["Relative humidity"] = data["Relative humidity"].filled(0)
    logger.debug("Compute altitude from geopotential")
    if "Geopotential Height" in data.keys():
        data["Altitude"] = get_altitude_from_geopotential_height(
            data["Geopotential Height"], data["Latitude"]
        )
    else:
        data["Altitude"] = get_altitude_from_geopotential_height(
            data["Geopotential"].quantity / STD_GRAVITATIONAL_ACCELERATION, data["Latitude"]
        )
    logger.debug("Compute density")
    data["Density"] = (
        STD_NUMBER_DENSITY
        * data["Pressure"]
        / STD_AIR_PRESSURE
        * STD_AIR_TEMPERATURE
        / data["Temperature"]
    )
    logger.debug("Compute exponential density")
    data["Exponential Density"] = (
        data["Density"] / STD_NUMBER_DENSITY * np.exp(data["Altitude"] / DENSITY_SCALE_HEIGHT)
    )
    logger.debug("Compute wind speed")
    data["Wind Speed"] = np.sqrt(
        data["U component of wind"] ** 2 + data["V component of wind"] ** 2
    )
    logger.debug("Compute wind direction")
    data["Wind Direction"] = Angle(
        np.array(
            np.arctan2(-data["V component of wind"].value, -data["U component of wind"].value)
        ),
        unit=u.rad,
    ).wrap_at(360 * u.deg)
    return data


def save_grib_table(data, filename, fmt="ecsv"):
    """
    Save grib data in a file according to provided format.

    Parameters
    ----------
    data: astropy.table.Table
        Grib data in astropy table
    filename: str
        Path to the file
    fmt: str
        Desired format
    """
    if fmt == "ecsv":
        data.write(filename, format="ascii.ecsv", overwrite=True)
    elif fmt == "magic":
        raise NotImplementedError("Magic txt format writing is not implemented yet")
    else:
        raise ValueError("Not recognized format")


def convert_to_text(data_path):
    """
    Converts GRIB/GRIB2 files to ecsv

    Parameters
    ----------
    data_path: path-like
        Path to the folder with GRIB/GRIB2 data files
    """

    grib_file_list = list(glob.glob(os.path.join(data_path, "*.grib2"))) + list(
        glob.glob(os.path.join(data_path, "*.grib"))
    )
    for grib_file in grib_file_list:
        data = get_grib_file_data(grib_file)
        data = extend_grib_data(data)
        save_grib_table(data, os.path.splitext(grib_file)[0] + ".ecsv")


def merge_ecsv_files(data_path):
    """
    Merge ecsv files

    Parameters
    ----------
    data_path: path-like
        Path to the folder with GRIB/GRIB2 data files
    """
    merged_data = Table()
    for tfile in os.listdir(data_path):
        if tfile.endswith(".ecsv"):
            data = Table.read(f"{data_path}/{tfile}")
            merged_data = vstack([merged_data, data])
    save_grib_table(merged_data, f"{data_path}/merged_file.ecsv")
