#! /usr/bin/env python

# SPDX-FileCopyrightText: 2021 Helmholtz-Zentrum für Umweltforschung GmbH - UFZ
#
# SPDX-License-Identifier: GPL-3.0-or-later

# -*- coding: utf-8 -*-

from __future__ import annotations

import json
from functools import reduce

import numpy as np
import pandas as pd

from saqc import BAD, DOUBTFUL, GOOD, UNFLAGGED
from saqc.core import Flags, History
from saqc.core.translation.basescheme import BackwardMap, ForwardMap, MappingScheme
from saqc.lib.tools import getUnionIndex

_QUALITY_CAUSES = [
    "",
    "BATTERY_LOW",
    "BELOW_MINIMUM",
    "ABOVE_MAXIMUM",
    "BELOW_OR_ABOVE_MIN_MAX",
    "ISOLATED_SPIKE",
    "DEFECTIVE_SENSOR",
    "LEFT_CENSORED_DATA",
    "RIGHT_CENSORED_DATA",
    "OTHER",
    "AUTOFLAGGED",
]

_QUALITY_LABELS = [
    "quality_flag",
    "quality_cause",
    "quality_comment",
]


class DmpScheme(MappingScheme):

    """
    Implements the translation from and to the flagging scheme implemented in
    the UFZ - Datamanagementportal
    """

    ARGUMENTS = {"comment": "", "cause": "OTHER"}

    DFILTER_DEFAULT = GOOD + 1

    _FORWARD: ForwardMap = {
        "NIL": UNFLAGGED,
        "OK": GOOD,
        "DOUBTFUL": DOUBTFUL,
        "BAD": BAD,
    }

    _BACKWARD: BackwardMap = {
        UNFLAGGED: "NIL",
        np.nan: "NIL",
        GOOD: "OK",
        DOUBTFUL: "DOUBTFUL",
        BAD: "BAD",
    }

    def __init__(self):
        super().__init__(forward=self._FORWARD, backward=self._BACKWARD)

    def toHistory(self, field_flags: pd.DataFrame):
        """
        Translate a single field of external ``Flags`` to a ``History``
        """
        field_history = History(field_flags.index)

        for (flag, cause, comment), values in field_flags.groupby(_QUALITY_LABELS):
            try:
                comment = json.loads(comment)
            except json.decoder.JSONDecodeError:
                comment = {"test": "unknown", "comment": ""}

            histcol = pd.Series(np.nan, index=field_flags.index)
            histcol.loc[values.index] = self(flag)

            meta = {
                "func": comment["test"],
                "kwargs": {"comment": comment["comment"], "cause": cause},
            }
            field_history.append(histcol, meta=meta)
        return field_history

    def toInternal(self, df: pd.DataFrame) -> Flags:
        """
        Translate from 'external flags' to 'internal flags'

        Parameters
        ----------
        df : pd.DataFrame
            The external flags to translate

        Returns
        -------
        Flags object
        """

        self.validityCheck(df)

        data = {}

        for field in df.columns.get_level_values(0).drop_duplicates():
            data[str(field)] = self.toHistory(df[field])

        return Flags(data)

    def toExternal(
        self, flags: Flags, attrs: dict | None = None, **kwargs
    ) -> pd.DataFrame:
        """
        Translate from 'internal flags' to 'external flags'

        Parameters
        ----------
        flags : The external flags to translate

        attrs : dict or None, default None
            global meta information of saqc-object

        Returns
        -------
        translated flags
        """
        tflags = super().toExternal(flags, attrs=attrs)

        out = pd.DataFrame(
            index=getUnionIndex(tflags),
            columns=pd.MultiIndex.from_product([flags.columns, _QUALITY_LABELS]),
        )

        for field in tflags.columns:
            df = pd.DataFrame(
                {
                    "quality_flag": tflags[field],
                    "quality_cause": "",
                    "quality_comment": "",
                }
            )

            history = flags.history[field]
            for col in history.columns:
                valid = (history.hist[col] != UNFLAGGED) & history.hist[col].notna()

                # extract from meta
                meta = history.meta[col]
                keywords = meta.get("kwargs", {})
                comment = json.dumps(
                    {
                        "test": meta.get("func", "unknown"),
                        "comment": keywords.get("comment", self.ARGUMENTS["comment"]),
                    }
                )
                cause = keywords.get("cause", self.ARGUMENTS["cause"])
                df.loc[valid, "quality_comment"] = comment
                df.loc[valid, "quality_cause"] = cause

            out[field] = df.reindex(out.index)

        self.validityCheck(out)
        return out

    @classmethod
    def validityCheck(cls, df: pd.DataFrame) -> None:
        """
        Check wether the given causes and comments are valid.

        Parameters
        ----------
        df : external flags
        """

        cols = df.columns
        if not isinstance(cols, pd.MultiIndex):
            raise TypeError("DMP-Flags need multi-index columns")

        if not cols.get_level_values(1).isin(_QUALITY_LABELS).all(axis=None):
            raise TypeError(
                f"DMP-Flags expect the labels {list(_QUALITY_LABELS)} in the secondary level"
            )

        for field in df.columns.get_level_values(0):
            # we might have NaN injected by DictOfSeries -> DataFrame conversions
            field_df = df[field].dropna(how="all", axis="index")
            flags = field_df["quality_flag"]
            causes = field_df["quality_cause"]
            comments = field_df["quality_comment"]

            if not flags.isin(cls._FORWARD.keys()).all(axis=None):
                raise ValueError(
                    f"invalid quality flag(s) found, only the following values are supported: {set(cls._FORWARD.keys())}"
                )

            if not causes.isin(_QUALITY_CAUSES).all(axis=None):
                raise ValueError(
                    f"invalid quality cause(s) found, only the following values are supported: {_QUALITY_CAUSES}"
                )

            if (~flags.isin(("OK", "NIL")) & (causes == "")).any(axis=None):
                raise ValueError(
                    "quality flags other than 'OK and 'NIL' need a non-empty quality cause"
                )

            if ((causes == "OTHER") & (comments == "")).any(axis=None):
                raise ValueError(
                    "quality cause 'OTHER' needs a non-empty quality comment"
                )
