#!/usr/bin/env python3
from collections import OrderedDict
from functools import partial
from joblib import delayed, parallel_backend, Parallel
import logging
import os
from pathlib import Path
import psutil
from .load_data import DataNotFoundError, read_netatmo_data_for_dtg
from .logs import get_logger, logcolor
from .mpi import mpi_parallel

logger = logging.getLogger(__name__)


def save_df_as_netatmo_csv(df, path, overwrite=False):
    """Saves a dataframe in the same CSV format netatmoqc reads"""

    path = Path(path).resolve()
    if path.exists() and not overwrite:
        raise FileExistsError("File {} exists".format(path))
    path.parent.mkdir(parents=True, exist_ok=True)

    if pd.api.types.is_datetime64_any_dtype(df["time_utc"]):
        df["time_utc"] = datetime2epoch(df["time_utc"])
    elif not pd.api.types.is_integer_dtype(df["time_utc"]):
        raise ValueError("time_utc is nether an integer nor a datetime")

    # It hurts a bit to label pressures as mslp again, but this is needed
    # so that the csv files generated by this routine are consistent
    # in data type with the original ones provided by netatmo
    if "pressure" in df.columns and "mslp" in df.columns:
        # In this case, we have recovered pressure from netatmo original file
        # when reading the df
        df = df.drop(["pressure"], axis=1)
    df = df.rename(columns=dict(mslp="pressure"))

    # Saving the adjusted df
    logger.info("Saving file %s", path)
    df.to_csv(path, index=None, mode="w")


def save_df_as_obsoul(
    df,
    fpath=None,
    export_params=["pressure", "temperature", "humidity", "sum_rain_1"],
):
    """Saves a dataframe returned by read_netatmo_data_for_dtg in obsoul format

    For further reference on the obsoul file format, please take a look at the
    documentation file available at
    <http://www.rclace.eu/File/Data_Assimilation/2007/lace_obspp.pdf>
    """

    dtg = df._parent_dtg

    # Define n_params, the "number of bodies" in obsoul lingo
    df_data_cols = [_ for _ in export_params if _ in df.columns]
    n_params = len(df_data_cols)
    if n_params == 0:
        logger.warning("Cannot produce file '%s': No valid export_params.")
        return
    # Remove duplicates from df_data_cols. Keep order so output files
    # can be compared using "diff" (we'll iterate over df_data_cols)
    df_data_cols = list(OrderedDict.fromkeys(df_data_cols))

    # More on comparability via diff: Sort df rows as we'll iterate over them
    key_cols = ["time_utc", "lat", "lon", "alt", "id"]
    sort_by = key_cols + [_ for _ in df_data_cols if _ not in key_cols]
    df = df.sort_values(by=sort_by, ignore_index=True)

    # Validate fpath and create parent dirs tree if not in place
    if fpath is None:
        fpath = Path() / "OBSOUL{}".format(dtg.strftime("%Y%m%d%H"))
    fpath = Path(fpath)
    fpath.parent.mkdir(parents=True, exist_ok=True)

    # Convert pressure from hPa to Pa
    if "pressure" in df_data_cols:
        df["pressure"] *= 100.0
    if "mslp" in df_data_cols:
        df["mslp"] *= 100.0

    # Define some constants
    obs_type = 1  # 1: SYNOP
    # Which obs_code to use?
    #  * obs_code=14: Automated Land Synop (seems reasonable)
    #  * obs_code=199: New codetype (used as codetype in Roel's csv2obsoul.jl)
    obs_code = 199
    obs_quality_flag = "1111"  # Got from Jelena
    site_dependent_flag = 100000  # Got from Jelena
    var2varcode = dict(
        # I got these from https://apps.ecmwf.int/odbgov/varno/
        mslp=108,  # varname="pmsl"
        pressure=110,  # varname="ps"
        temperature=39,  # varname="t2m"
        humidity=58,  # varname="rh2m". Or should I use 7 (varname=q)?
        # sum_rain_1 is prob not OK. I got it from
        # https://apps.ecmwf.int/odbgov/accumulationkind/
        sum_rain_1=39001,
    )

    # Start file format conversion
    analysis_date = dtg.strftime("%Y%m%d")
    analysis_time = dtg.strftime("%H")
    with open(fpath, "w") as f:
        # The "date time" first line in the file
        f.write("{} {}\n".format(analysis_date, analysis_time))

        for row in df.itertuples(index=False):
            # Construct the records. Each record has a header and multiple
            # "bodies", all in the same line.
            #
            # 1) Construct header.
            # We should have one header for each (id, lat, lon, alt). Since we
            # forbid, in this package, that stations move (lat, lon, alt) in
            # any given DTG, this implies that we should have one -- and only
            # one - header for each id. Moreover, as we are also removing
            # duplicate station IDs within any given DTG, this therefore
            # implies that there should be one and only one header for each
            # row in our input dataframe.
            obs_date = row.time_utc.strftime("%Y%m%d")
            # obs_hour must not have leading zeros for the hour
            obs_hour = row.time_utc.strftime("%k%M%S").strip()
            header = (
                17,  # Got this from Jelena
                obs_type,
                obs_code,
                row.lat,
                row.lon,
                "'{}'".format(row.id[:8]),
                obs_date,
                obs_hour,
                row.alt,
                n_params,
                obs_quality_flag,
                site_dependent_flag,
            )

            # Write the header. According to the obsoul file format docs, each
            # observation record (header + bodies) should occupy only one line
            # in the file, so let's skip putting a newline between header and
            # bodies and just put a space at the beginning of each body
            f.write(" ".join(map(str, header)))

            # Construct bodies
            for param_name in df_data_cols:
                # param_type is the same as "varid@body" in odb lingo
                param_type = var2varcode[param_name]

                # param_quality_flag:
                #   * Jelena said it could maybe be "ne"
                #   * Roel uses paramqcflag=2048 in his csv2obsoul.jl
                # I'm using the same as Roel for consistency
                param_quality_flag = 2048

                # Set first_vert_coord similarly as in  Roel's csv2obsoul.jl
                # g0: Globally-averaged gravity acceleration in m*s^-2
                g0 = 9.81
                first_vert_coord = g0 * row.alt

                # vertco_reference_2: 2nd vertical coord reference. We don't
                # have a value for that. Signal missing value with a sentinel
                # value of 0.1699999976E+39 (according to the obsoul doc file)
                vertco_reference_2 = 0.1699999976e39

                body = (
                    param_type,
                    first_vert_coord,
                    vertco_reference_2,
                    getattr(row, param_name),
                    param_quality_flag,
                )
                # Again, no space between the bodies of the same record,
                # but just a space at the beginning of each body
                f.write(" " + " ".join(map(str, body)))

            # All bodies for this record are written. Newline is finally needed
            f.write("\n")


def netatmo_csv2obsoul(
    dtg,
    netatmo_data_rootdir,
    selected_stations=None,
    dropna=True,
    fillna=dict(sum_rain_1=0.0),
    rm_duplicate_stations=True,
    rm_moving_stations=True,
    outdir=Path(),
    loglevel="INFO",
):
    logger = get_logger(__name__, loglevel)

    logger.info(
        "Reading data for %sDTG=%s%s...", logcolor.cyan, dtg, logcolor.reset,
    )
    try:
        # read_netatmo_data_for_dtg will raise DataNotFoundError if
        # there's no data for the selected DTG
        df = read_netatmo_data_for_dtg(
            dtg,
            rootdir=netatmo_data_rootdir,
            dropna=dropna,
            fillna=fillna,
            recover_pressure_from_mslp=True,
            drop_mslp=True,
            remove_duplicate_stations=rm_duplicate_stations,
            remove_moving_stations=rm_moving_stations,
        )
        if selected_stations is not None:
            # Leave only selected stations
            df = df[df["id"].isin(selected_stations)]
            if len(df.index) == 0:
                raise DataNotFoundError
    except (DataNotFoundError):
        logger.warning("No stations for DTG=%s", dtg)
        return

    fpath = Path(outdir) / "OBSOUL{}".format(dtg.strftime("%Y%m%d%H"))
    logger.info(
        "Saving OBSOUL: %sDTG=%s, %d obs%s, file %s",
        logcolor.cyan,
        dtg,
        len(df.index),
        logcolor.reset,
        fpath,
    )
    save_df_as_obsoul(df, fpath)


def netatmo_csvs2obsouls(
    dtgs,
    netatmo_data_rootdir,
    selected_stations=None,
    dropna=True,
    fillna=dict(sum_rain_1=0.0),
    rm_duplicate_stations=True,
    rm_moving_stations=True,
    outdir=Path(),
    loglevel="INFO",
    use_mpi=False,
):
    """Save in OBSOUL files the data from the CSV files that netatmoqc reads

    If a list of selected stations is passed in the "selected_stations", then
    only these are kept.

    Used in the "select" and "csv2obsoul" commands.
    """

    logger.info(
        "%sSaving selected data in obsoul format...%s",
        logcolor.cyan,
        logcolor.reset,
    )
    if use_mpi:
        logger.info("Parallelising tasks over DTGs using MPI")
        func = partial(
            netatmo_csv2obsoul,
            netatmo_data_rootdir=netatmo_data_rootdir,
            selected_stations=selected_stations,
            dropna=dropna,
            fillna=fillna,
            rm_duplicate_stations=rm_duplicate_stations,
            rm_moving_stations=rm_moving_stations,
            outdir=outdir,
            loglevel=loglevel,
        )
        mpi_parallel(func, dtgs)
    else:
        n_procs = int(
            os.getenv("NETATMOQC_MAX_PYTHON_PROCS", psutil.cpu_count())
        )
        if n_procs > 1:
            logger.info("Parallelising tasks over DTGs within a single host")
        # Using the "loky" backend helps avoiding oversubscription
        # of cpus in child processes
        with parallel_backend("loky"):
            Parallel(n_jobs=n_procs)(
                delayed(netatmo_csv2obsoul)(
                    dtg,
                    netatmo_data_rootdir=netatmo_data_rootdir,
                    selected_stations=selected_stations,
                    dropna=dropna,
                    fillna=fillna,
                    rm_duplicate_stations=rm_duplicate_stations,
                    rm_moving_stations=rm_moving_stations,
                    outdir=outdir,
                    loglevel=loglevel,
                )
                for dtg in dtgs
            )
