#!/usr/bin/env python3
"""Code related to saving output files."""
import logging
import os
from collections import OrderedDict
from functools import partial
from pathlib import Path

import numpy as np
import pandas as pd
import psutil
from joblib import Parallel, delayed, parallel_backend

from .dtgs import datetime2epoch
from .load_data import (
    DataNotFoundError,
    read_netatmo_data_for_dtg,
    remove_irregular_stations,
)
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):
    """Save dataframe in the same CSV format netatmoqc reads.

    Args:
        df (pandas.Dataframe): The dataframe to be saved.
        path (pathlib.Path): Path for the output file.
        overwrite (bool): Overwite or not the out file. Default value = False.

    Raises:
        FileExistsError: If path doesn't exist and overwrite==False.
        ValueError: If df["time_utc"] is nether an integer (assumed to be an
            epoch) nor a datetime object.

    """
    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.debug(
        "%sSaving CSV, %d obs%s: %s",
        logcolor.cyan,
        len(df.index),
        logcolor.reset,
        path,
    )
    df.to_csv(path, index=None, mode="w")


def save_df_as_obsoul(df, fpath=None, export_params=None):
    """Save 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>

    Args:
        df (netatmoqc.load_data.DataFrameWithDtgMetadata): Input dataframe.
            This type is inherited from pandas.DataFrame, but it has added
            DTG metadata info (df.metadata_dict["dtg"]).
        fpath (pathlib.Path):  Output file path (Default value = OBSOUL${DTG}).
        export_params (list): Parameters to save in the OBSOULD file.
            (Default value=["pressure","temperature","humidity","sum_rain_1"])

    """
    if export_params is None:
        export_params = ["pressure", "temperature", "humidity", "sum_rain_1"]

    dtg = df.metadata_dict["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 obsoul_file:
        # The "date time" first line in the file
        obsoul_file.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),
                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
            obsoul_file.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
                obsoul_file.write(" " + " ".join(map(str, body)))

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


@np.vectorize
def obs_timestamp2csv_rpath(timestamp):
    """Get appropriate relative CSV file path from observation timestamp.

    Args:
        timestamp (datetime.datetime): The timestamp of the observation.

    Returns:
        pathlib.Path: CSV file relative path to be used.

    """
    year = timestamp.strftime("%Y")
    month = timestamp.strftime("%m")
    day = timestamp.strftime("%d")
    minute = int(timestamp.strftime("%M"))
    if minute < 10:
        f_minute = "05"
    elif minute < 20:
        f_minute = "15"
    elif minute < 30:
        f_minute = "25"
    elif minute < 40:
        f_minute = "35"
    elif minute < 50:
        f_minute = "45"
    elif minute < 60:
        f_minute = "55"

    fname = "%s%s00Z.csv" % (timestamp.strftime("%Y%m%dT%H"), f_minute)
    return Path("%s/%s/%s/%s" % (year, month, day, fname))


def _input2output_single_dtg(
    dtg,
    netatmo_data_rootdir,
    selected_stations=None,
    dropna=True,
    fillna=None,
    rm_duplicate_stations=True,
    rm_moving_stations=True,
    outdir_csv=None,
    outdir_obsoul=None,
    obsoul_export_params=None,
    loglevel="INFO",
):
    """Read data from one DTG, filter selected obs and save results.

    Args:
        dtg (netatmoqc.dtgs.Dtg): Assimilation DTG for the data.
        netatmo_data_rootdir (pathlib.Path): Path to the top-level directory
            where the netatmo input csv files are stored.
        selected_stations (list): Selected station IDs (Default value = None).
        dropna (bool): Drop NaN values before saving (Default value = True).
        fillna (dict): Dictionary having data parameters as keys and the fill
            values as dict values. Passed to read_netatmo_csv
            (Default value = None).
        rm_duplicate_stations (bool): (Default value = True)
        rm_moving_stations (bool): (Default value = True)
        outdir_csv (pathlib.Path): Dir where out csv files will be put. If None
            is passed, then no csv files will be saved. (Default value = None)
        outdir_obsoul (pathlib.Path): Analogous to outdir_csv, but this time
            for obsoul files.
        obsoul_export_params (list): List of data params to be exported in the
            final obsoul file. Passed to save_df_as_obsoul.
            (Default value = None).
        loglevel (str): Level for logging (Default value = "INFO").

    """
    logger = get_logger(__name__, loglevel)

    logger.debug(
        "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,
        )
        df, _ = remove_irregular_stations(
            df, moving=rm_moving_stations, duplicates=rm_duplicate_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

    if outdir_csv is not None:
        df["_f_rpath"] = obs_timestamp2csv_rpath(df["time_utc"])
        for f_rpath, file_df in df.groupby("_f_rpath", sort=False):
            fpath = Path(outdir_csv) / f_rpath
            file_df = file_df.drop("_f_rpath", axis=1).sort_values(
                by=["time_utc"]
            )
            save_df_as_netatmo_csv(file_df, fpath, overwrite=True)
        df = df.drop("_f_rpath", axis=1)

    if outdir_obsoul is not None:
        fpath = Path(outdir_obsoul) / "OBSOUL{}".format(
            dtg.strftime("%Y%m%d%H")
        )
        logger.debug(
            "%sSaving OBSOUL: DTG=%s, %d obs%s, file %s",
            logcolor.cyan,
            dtg,
            len(df.index),
            logcolor.reset,
            fpath,
        )
        save_df_as_obsoul(df, fpath, export_params=obsoul_export_params)


def netatmoqc_input2output(
    dtgs,
    netatmo_data_rootdir,
    selected_stations=None,
    dropna=True,
    fillna=None,
    rm_duplicate_stations=True,
    rm_moving_stations=True,
    outdir=None,
    loglevel="INFO",
    use_mpi=False,
    save_csv=True,
    save_obsoul=False,
    obsoul_export_params=None,
):
    """Save NetAtmoQC output data in CSV or OBSOUL formats.

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

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

    Args:
        dtgs (netatmoqc.dtgs.DtgContainer): Assimilation DTGs for the data.
        netatmo_data_rootdir (pathlib.Path): Path to the top-level directory
            where the netatmo input csv files are stored.
        selected_stations (list): Selected station IDs (Default value = None).
        dropna (bool): Drop NaN values before saving (Default value = True).
        fillna (dict): Dictionary having data parameters as keys and the fill
            values as dict values. Passed to read_netatmo_csv
            (Default value = None).
        rm_duplicate_stations (bool): (Default value = True)
        rm_moving_stations (bool): (Default value = True)
        outdir (pathlib.Path): Dir where out files will be put.
            (Default value = ".")
        loglevel (str): Level for logging (Default value = "INFO").
        use_mpi (bool): Whether to use MPI or not (Default value = False).
        save_csv (bool): Whether to save csv files or not.
            (Default value = True).
        save_obsoul (bool): Whether to save obsoul files or not.
            (Default value = False)
        obsoul_export_params (list): List of data params to be exported in the
            final obsoul file. Passed to save_df_as_obsoul.
            (Default value = None).

    """
    if outdir is None:
        # Don't perform function calls in argument defaults
        outdir = Path()

    logger.info(
        "%sSaving selected observations...%s", logcolor.cyan, logcolor.reset,
    )

    outdir_csv = None
    outdir_obsoul = None
    if save_csv:
        outdir_csv = Path(outdir) / "csv_files"
        logger.info(
            "%s> CSV outdir:%s %s", logcolor.cyan, logcolor.reset, outdir_csv,
        )
    if save_obsoul:
        outdir_obsoul = Path(outdir) / "obsoul_files"
        logger.info(
            "%s> OBSOUL outdir:%s %s",
            logcolor.cyan,
            logcolor.reset,
            outdir_obsoul,
        )
        if obsoul_export_params is not None:
            logger.info(
                "    * OBSOUL files will contain: %s",
                ", ".join(obsoul_export_params),
            )

    in2out_fdtg = partial(
        _input2output_single_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,
        loglevel=loglevel,
        outdir_csv=outdir_csv,
        outdir_obsoul=outdir_obsoul,
        obsoul_export_params=obsoul_export_params,
    )

    if use_mpi:
        logger.info("Parallelising tasks over DTGs using MPI")
        mpi_parallel(in2out_fdtg, 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(in2out_fdtg)(dtg) for dtg in dtgs)
