import os

import numpy as np
import pandas as pd
from geonum import atmosphere as atm
from tqdm import tqdm

from pyaerocom import const
from pyaerocom.aux_var_helpers import concx_to_vmrx
from pyaerocom.exceptions import DataRetrievalError
from pyaerocom.helpers import varlist_aerocom
from pyaerocom.io import ReadUngriddedBase
from pyaerocom.molmasses import get_molmass
from pyaerocom.ungriddeddata import UngriddedData

P_STD = atm.p0  # standard atmosphere pressure
T_STD = atm.T0_STD  # standard atmosphere temperature


def _conc_to_vmr_marcopolo_stats(
    data, to_var, from_var, p_pascal=None, T_kelvin=None, mmol_air=None
):
    if p_pascal is None:
        p_pascal = P_STD
    if T_kelvin is None:
        T_kelvin = T_STD
    if mmol_air is None:
        mmol_air = get_molmass("air_dry")

    for stat in data:
        if not from_var in stat:
            continue

        concdata = stat[from_var]

        mmol_var = get_molmass(to_var)
        from_unit = stat["var_info"][from_var]["units"]
        to_unit = const.VARS[to_var].units
        vmrvals = concx_to_vmrx(
            concdata,
            p_pascal=p_pascal,
            T_kelvin=T_kelvin,
            mmol_var=mmol_var,
            mmol_air=mmol_air,
            conc_unit=from_unit,
            to_unit=to_unit,
        )
        stat[to_var] = vmrvals
        vi = {}
        vi.update(stat["var_info"][from_var])
        vi["computed"] = True
        vi["units"] = to_unit
        vi["units_info"] = (
            f"The original data is provided as mass conc. in units of ug m-3 and "
            f"was converted to ppb assuming a standard atmosphere "
            f"(p={p_pascal/100}hPa, T={T_kelvin}K) assuming dry molar mass of air "
            f"M_Air_dry={mmol_air}g/mol and total molecular mass of "
            f"{from_var}={mmol_var}g/mol"
        )
        stat["var_info"][to_var] = vi

    return data


class ReadMarcoPolo(ReadUngriddedBase):
    """
    Reading routine for Chinese MarcoPolo observations
    """

    #: Data type of files
    _FILETYPE = ".csv"

    #: Filemask for glob data file search
    _FILEMASK = f"*{_FILETYPE}"

    #: Column delimiter
    FILE_COL_DELIM = ","

    #: Version log of this class (for caching)
    __version__ = "0.02"

    #: Default data ID
    DATA_ID = "MarcoPolo"

    #: List of all datasets supported by this interface
    SUPPORTED_DATASETS = [DATA_ID]

    #: Indices of columns in data files
    FILE_COLS = {"station_id": 0, "datetime": 1, "var_name": 2, "value": 3}
    #: Mapping of columns in station metadata file to pyaerocom standard
    STATION_META_MAP = {
        "stationId": "station_id",
        "stationName": "station_name",
        "cityId": "city_id",
        "latitude": "latitude",
        "longitude": "longitude",
        "cityName": "city",
    }

    #: conversion functions for metadata dtypes
    STATION_META_DTYPES = {
        "station_id": str,
        "station_name": str,
        "city_id": str,
        "latitude": float,
        "longitude": float,
        "city": str,
    }

    #: Variable names in data files
    VAR_MAP = {
        "concco": "co",
        "concso2": "so2",
        "concno2": "no2",
        "conco3": "o3",
        "concpm10": "pm10",
        "concpm25": "pm2_5",
    }

    #: Variable units (not provided in data files)
    VAR_UNITS = {
        "concco": "unknown",
        "concso2": "unknown",
        "concno2": "ug m-3",  # mail from Henk Eskes, 1 Feb. 2021
        "conco3": "ug m-3",  # mail from Henk Eskes, 1 Feb. 2021
        "concpm10": "ug m-3",  # mail from Henk Eskes, 1 Feb. 2021
        "concpm25": "ug m-3",  # mail from Henk Eskes, 1 Feb. 2021
    }

    #: Variables that are computed (cannot be read directly)
    AUX_REQUIRES = {"vmro3": ["conco3"], "vmrno2": ["concno2"]}

    #: functions used to convert variables that are computed
    AUX_FUNS = {"vmro3": _conc_to_vmr_marcopolo_stats, "vmrno2": _conc_to_vmr_marcopolo_stats}

    PROVIDES_VARIABLES = list(VAR_MAP)
    DEFAULT_VARS = ["concpm10", "concpm25", "concno2", "conco3"]

    TS_TYPE = "hourly"

    STAT_METADATA_FILENAME = "station_info.xlsx"

    def __init__(self, data_id=None, data_dir=None):
        super().__init__(data_id=data_id, data_dir=data_dir)

        try:
            import openpyxl
        except ImportError:
            raise ImportError(
                "ReadMarcoPolo class needs package openpyxl "
                "which is not part of the pyaerocom standard installation. "
                "Please install openpyxl via conda or pip."
            )

    def _read_metadata_file(self):
        fn = os.path.join(self.data_dir, self.STAT_METADATA_FILENAME)
        cfg = pd.read_excel(fn, engine="openpyxl")
        return cfg

    def _init_station_metadata(self):
        """
        Initiate metadata for all stations

        Returns
        -------
        dict
            dictionary with metadata dictionaries for all stations

        """

        cfg = self._read_metadata_file()

        meta_map = self.STATION_META_MAP
        dtypes = self.STATION_META_DTYPES

        cols = list(cfg.columns.values)
        col_idx = {}
        for from_meta, to_meta in meta_map.items():
            col_idx[to_meta] = cols.index(from_meta)

        arr = cfg.values

        stats = {}
        for row in arr:
            stat = {}
            for meta_key, col_num in col_idx.items():
                stat[meta_key] = dtypes[meta_key](row[col_num])
            stat["data_id"] = self.data_id
            stat["ts_type"] = self.TS_TYPE
            stats[stat["station_id"]] = stat

        return stats

    def _read_file(self, file):
        df = pd.read_csv(file, sep=self.FILE_COL_DELIM)
        return df

    def _make_station_data(
        self, var, var_stats_unique, stat_ids, var_stats, var_dtime, var_values, stat_meta
    ):

        units = self.VAR_UNITS
        stats = []
        for stat_id in var_stats_unique:
            if not stat_id in stat_ids:
                continue

            statmask = var_stats == stat_id

            if statmask.sum() == 0:
                continue

            timestamps = var_dtime[statmask]
            vals = var_values[statmask]

            stat = {}
            stat.update(stat_meta[stat_id])

            stat["dtime"] = timestamps
            stat["timezone"] = "unknown"
            stat[var] = vals
            unit = units[var]
            stat["var_info"] = {}
            stat["var_info"][var] = dict(units=unit)
            stats.append(stat)
        return stats

    def _read_files(self, files, vars_to_retrieve):

        filecols = self.FILE_COLS

        stat_meta = self._init_station_metadata()

        # get all station IDs found in the metadata file
        stat_ids = list(stat_meta)

        arrs = []

        varcol = filecols["var_name"]
        for i in tqdm(range(len(files))):
            filedata = self._read_file(files[i])
            arr = filedata.values

            for i, var in enumerate(vars_to_retrieve):
                filevar = self.VAR_MAP[var]
                if i == 0:
                    mask = arr[:, varcol] == filevar
                else:
                    mask = np.logical_or(mask, arr[:, varcol] == filevar)
            matches = mask.sum()
            if matches:
                vardata = arr[mask]
                arrs.append(vardata)
        if len(arrs) == 0:
            raise DataRetrievalError("None of the input variables could be found in input list")
        data = np.concatenate(arrs)

        dtime = data[:, filecols["datetime"]].astype("datetime64[s]")

        varis = data[:, varcol].astype(str)
        values = data[:, filecols["value"]].astype(np.float64)
        statids = data[:, filecols["station_id"]].astype(str)

        stats = []
        for var in vars_to_retrieve:
            var_in_file = self.VAR_MAP[var]

            varmask = varis == var_in_file

            var_values = values[varmask]
            var_stats = statids[varmask]
            var_dtime = dtime[varmask]

            var_stats_unique = np.unique(var_stats)

            var_stats = self._make_station_data(
                var, var_stats_unique, stat_ids, var_stats, var_dtime, var_values, stat_meta
            )
            stats.extend(var_stats)

        return stats

    def read_file(self):
        """
        This method is not implemented (but needs to be declared for template)

        Raises
        ------
        NotImplementedError
        """
        raise NotImplementedError("Not needed for these data since the format is unsuitable...")

    def compute_additional_vars(self, statlist_from_file, vars_to_compute):

        for var in vars_to_compute:
            fun = self.AUX_FUNS[var]
            requires = self.AUX_REQUIRES[var]
            # this will add the variable data to each station data in
            # statlist_from_file
            statlist_from_file = fun(statlist_from_file, var, *requires)

        return statlist_from_file

    def read(self, vars_to_retrieve=None, first_file=None, last_file=None):
        """
        Read variable data

        Parameters
        ----------
        vars_to_retrieve : str or list, optional
            List of variables to be retrieved. The default is None.
        first_file : int, optional
            Index of first file to be read. The default is None, in which case
            index 0 in file list is used.
        last_file : int, optional
            Index of last file to be read. The default is None, in which case
            last index in file list is used.

        Returns
        -------
        data : UngriddedData
            loaded data object.

        """
        if vars_to_retrieve is None:
            vars_to_retrieve = self.DEFAULT_VARS
        elif isinstance(vars_to_retrieve, str):
            vars_to_retrieve = [vars_to_retrieve]

        # make sure to use AeroCom variable names in output data
        vars_to_retrieve = varlist_aerocom(vars_to_retrieve)

        vars_to_read, vars_to_compute = self.check_vars_to_retrieve(vars_to_retrieve)

        files = self.get_file_list()
        if first_file is None:
            first_file = 0
        if last_file is None:
            last_file = len(files)
        files = files[first_file:last_file]

        stats = self._read_files(files, vars_to_read)

        stats = self.compute_additional_vars(stats, vars_to_compute)

        data = UngriddedData.from_station_data(stats)

        return data
