#
# (C) Copyright 2017- ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.
#

import copy
import datetime
import logging
import os
from pathlib import Path

from . import utils
import numpy as np
import pandas as pd
import yaml

# logging.basicConfig(level=logging.DEBUG)
# logging.basicConfig(level=logging.INFO)
LOG = logging.getLogger(__name__)

NEWER = True


class GribIndexer:
    VECTOR_PARAMS = {
        "wind10m": ["10u", "10v"],
        "wind100m": ["100u", "100v"],
        "wind200m": ["200u", "200v"],
        "wind": ["u", "v"],
        "wind3d": ["u", "v", "w"],
    }

    # tuple-> 0: ecCodes type, 1: pandas type, 2: Python type 3: use in duplicate check
    DEFAULT_KEYS = {
        "shortName": ("s", str, str, False),
        "paramId": ("l", "Int32", int, False),
        "date": ("l", "Int64", int, True),
        "time": ("l", "Int64", int, True),
        "step": ("l", "Int32", int, True),
        "level": ("l", "Int32", int, True),
        "typeOfLevel": ("s", str, str, False),
        "number": ("s", str, str, True),
        "experimentVersionNumber": ("s", str, str, False),
        "marsClass": ("s", str, str, False),
        "marsStream": ("s", str, str, False),
        "marsType": ("s", str, str, False),
    }

    DEFAULT_ECC_KEYS = [f"{k}:{v[0]}" for k, v in DEFAULT_KEYS.items()]
    BLOCK_KEYS = ["shortName", "typeOfLevel"]

    DEFAULT_SORT_KEYS = [
        "date",
        "time",
        "step",
        "number",
        "level",
        "paramId",
    ]
    DATE_KEYS = {
        k: ("l", "Int64", int)
        for k in ["date", "dataDate", "validityDate", "mars.date", "marsDate"]
    }
    TIME_KEYS = {
        k: ("l", "Int64", int)
        for k in ["time", "dataTime", "validityTime", "mars.time", "marsTime"]
    }

    DATETIME_KEYS = {
        "_dateTime": ("date", "time"),
        "_dataDateTime": ("dataDate", "dataTime"),
        "_validityDateTime": ("validityDate", "validityTime"),
    }

    KEYS_TO_REPLACE = {
        ("type", "mars.type"): "marsType",
        ("stream", "mars.stream"): "marsStream",
        ("class", "mars.class", "class_"): "marsClass",
        ("perturbationNumber"): "number",
        ("mars.date", "marsDate"): "date",
        ("mars.time", "marsTime"): "time",
    }

    PREDEF_KEYS = copy.deepcopy(DEFAULT_KEYS)
    PREDEF_KEYS.update(DATE_KEYS)
    PREDEF_KEYS.update(TIME_KEYS)

    PREDEF_PD_TYPES = {k: v[1] for k, v in PREDEF_KEYS.items()}
    PREDEF_PT_TYPES = {k: v[2] for k, v in PREDEF_KEYS.items()}

    def __init__(self, db):
        self.db = db
        assert self.db is not None
        self.ref_column_count = None
        self.keys = []
        self.keys_ecc = []
        for k, v in GribIndexer.DEFAULT_KEYS.items():
            name = k
            self.keys.append(name)
            if v[0]:
                name = f"{k}:{v[0]}"
            self.keys_ecc.append(name)

        self.keys_duplicate_check = [
            k for k, v in GribIndexer.DEFAULT_KEYS.items() if v[3] == True
        ]

        self.shortname_index = self.keys.index("shortName")
        self.levtype_index = self.keys.index("typeOfLevel")
        self.type_index = self.keys.index("marsType")
        self.number_index = self.keys.index("number")
        self.param_id_index = self.keys.index("paramId")

        self.wind_check_index = []
        for v in [
            "date",
            "time",
            "step",
            "level",
            "typeOfLevel",
            "level",
            "number",
            "experimentVersionNumber",
            "marsClass",
            "marsStream",
            "marsType",
        ]:
            self.wind_check_index.append(self.keys.index(v) + 1)

        # self.block_key_index = [self.keys.index(v) for v in GribIndexer.BLOCK_KEYS]

        self.pd_types = {k: v[1] for k, v in GribIndexer.DEFAULT_KEYS.items()}
        self.pt_types = {k: v[2] for k, v in GribIndexer.DEFAULT_KEYS.items()}

    def update_keys(self, keys):
        ret = False
        for k in keys:
            name = k
            # we do not add datetime keys (they are pseudo keys, and their value
            # is always generated on the fly)
            if name not in self.keys and name not in GribIndexer.DATETIME_KEYS:
                self.keys.append(name)
                p = GribIndexer.PREDEF_KEYS.get(name, ("", str, str))
                ecc_name = name if p[0] == "" else name + ":" + p[0]
                self.keys_ecc.append(ecc_name)
                self.pd_types[name] = p[1]
                self.pt_types[name] = p[2]
                ret = True
        return ret

    def _check_duplicates(self, name, df):
        dup = df.duplicated(subset=self.keys_duplicate_check)
        first_dup = True
        cnt = 0
        for i, v in dup.items():
            if v:
                if first_dup:
                    LOG.error(
                        f"{name}: has duplicates for key group: {self.keys_duplicate_check}!"
                    )
                    first_dup = False
                    LOG.error(f" first duplicate: {df.iloc[i]}")
                cnt += 1

        if cnt > 1:
            LOG.error(f"  + {cnt-1} more duplicate(s)!")

    def _build_vector_index(self, df, v_name, v_comp):
        # LOG.debug(f"v_name={v_name} v_comp={v_comp}")
        comp_num = len(v_comp)

        # filter components belonging together
        comp_df = []
        for i, comp_name in enumerate(v_comp):
            query = f"shortName == '{comp_name}'"
            r = df.query(query, engine="python")
            # if we do not use copy, the assignment below as:
            # comp_df[0].loc[...
            # generates the SettingWithCopyWarning warning!!!
            if i == 0:
                r = df.query(query, engine="python").copy()
            else:
                r = df.query(query, engine="python")
            if r.empty:
                return []
            else:
                comp_df.append(r)

        assert comp_num == len(comp_df)

        # pair up components within a 2D vector field. This
        # version proved to be the fastest!
        # LOG.debug(" pair up collected components:")

        # print(f"v_name={v_name} {len(comp_df[1].index)}")
        # print("view=", comp_df[0]._is_view)
        r = []
        used1 = np.full(len(comp_df[1].index), False, dtype="?")
        comp_df[0].loc[:, "shortName"] = v_name
        # 2D
        if comp_num == 2:
            for row0 in comp_df[0].itertuples(name=None):
                i = 0
                for row1 in comp_df[1].itertuples(name=None):
                    if not used1[i]:
                        b = True
                        for x in self.wind_check_index:
                            if row0[x] != row1[x]:
                                b = False
                                break
                        if b:
                            d = list(row0[1:])
                            d.extend(row1[-self.ref_column_count :])
                            r.append(d)
                            used1[i] = True
                            break
                    i += 1
        # 3D
        elif comp_num == 3:
            used2 = np.full(len(comp_df[2].index), False, dtype="?")
            for row0 in comp_df[0].itertuples(name=None):
                i = 0
                for row1 in comp_df[1].itertuples(name=None):
                    if not used1[i]:
                        b = True
                        for x in self.wind_check_index:
                            if row0[x] != row1[x]:
                                b = False
                                break
                        if b:
                            j = 0
                            for row2 in comp_df[2].itertuples(name=None):
                                if not used2[j]:
                                    b = True
                                    for x in self.wind_check_index:
                                        if row0[x] != row2[x]:
                                            b = False
                                            break
                                    if b:
                                        d = list(row0[1:])
                                        d.extend(row1[-self.ref_column_count :])
                                        d.extend(row2[-self.ref_column_count :])
                                        r.append(d)
                                        used1[i] = True
                                        used2[j] = True
                                        j = -1
                                        break
                                j += 1
                            if j == -1:
                                break
                    i += 1
        return r

    def _make_dataframe(self, data, sort=False, columns=None):
        if columns is not None:
            df = pd.DataFrame(data, columns=columns)
        else:
            df = pd.DataFrame(data)

        for c in df.columns:
            if self.pd_types.get(c, "") in ["Int32", "Int64"]:
                df.fillna(value={c: np.nan}, inplace=True)
            df = df.astype(self.pd_types)
        if sort:
            df = GribIndexer._sort_dataframe(df)

        return df

    @staticmethod
    def _sort_dataframe(df, columns=None, ascending=True):
        if columns is None:
            columns = list(df.columns)
        elif not isinstance(columns, list):
            columns = [columns]

        # mergesoft is a stable sorting algorithm
        df = df.sort_values(by=columns, ascending=ascending, kind="mergesort")
        df = df.reset_index(drop=True)
        return df

    def _write_dataframe(self, df, name, out_dir):
        f_name = os.path.join(out_dir, f"{name}.csv.gz")
        df.to_csv(path_or_buf=f_name, header=True, index=False, compression="gzip")

    @staticmethod
    def read_dataframe(key, dir_name):
        # assert len(key) == len(GribIndexer.BLOCK_KEYS)
        name = key
        f_name = os.path.join(dir_name, f"{name}.csv.gz")
        # LOG.debug("f_name={}".format(f_name))
        return pd.read_csv(f_name, index_col=None, dtype=GribIndexer.PREDEF_PD_TYPES)

    @staticmethod
    def get_storage_key_list(dir_name):
        r = []
        # LOG.debug(f"dir_name={dir_name}")
        suffix = ".csv.gz"
        for f in utils.get_file_list(os.path.join(dir_name, f"*{suffix}")):
            name = os.path.basename(f)
            # LOG.debug(f"name={name}")
            r.append(name[: -len(suffix)])
        return r

    @staticmethod
    def is_key_wind(key):
        return key in GribIndexer.VECTOR_PARAMS

    @staticmethod
    def _convert_query_value(v, col_type):
        # print(f"v={v} {type(v)} {col_type}")
        return v if col_type != "object" else str(v)

    @staticmethod
    def _check_datetime_in_filter_input(keys):
        for k, v in GribIndexer.DATETIME_KEYS.items():
            name = k[1:]
            name_date = v[0]
            name_time = v[1]
            if keys.get(name, []) and (
                keys.get(name_date, []) or keys.get(name_time, [])
            ):
                raise Exception(
                    f"Cannot specify {name} together with {name_date} and {name_time}!"
                )

    @staticmethod
    def _convert_filter_value(name, val):
        """
        Analyse the filter key-value pairs and perform the necessary conversions
        """
        valid_name = name.split(":")[0] if ":" in name else name

        # datetime keys are pseudo keys, they start with _. Their value is converted to
        # datetime. The key itself is not added to the scan!
        if ("_" + valid_name) in GribIndexer.DATETIME_KEYS:
            valid_name = "_" + valid_name
            name_date = GribIndexer.DATETIME_KEYS[valid_name][0]
            name_time = GribIndexer.DATETIME_KEYS[valid_name][1]
            for i, t in enumerate(val):
                val[i] = GribIndexer._to_datetime(name, t)
                # print(f"t={t} -> {val[i]}")
            # We add the date and time components with an empty value. So they will be
            # added to the scan, but they will be ignored by the query. Conversely,
            # the datetime key itself will be ignored in the scan, but will be used
            # in the query.
            return [("_" + name, val), (name_date, []), (name_time, [])]
        # we convert dates to int
        elif valid_name in GribIndexer.DATE_KEYS:
            for i, t in enumerate(val):
                d = GribIndexer._to_date(name, t)
                # for daily climatologies dates where the year is missing the
                # the a tuple is returned
                if not isinstance(d, tuple):
                    val[i] = int(d.strftime("%Y%m%d"))
                else:
                    val[i] = d[0] * 100 + d[1]
        # we convert times to int
        elif valid_name in GribIndexer.TIME_KEYS:
            for i, t in enumerate(val):
                val[i] = int(GribIndexer._to_time(name, t).strftime("%H%M"))
                # print(f"t={t} -> {val[i]}")
        else:
            pt_type = GribIndexer.PREDEF_PT_TYPES.get(name, None)
            # print(f"name={name} {pt_type}")
            if pt_type is not None:
                for i, t in enumerate(val):
                    val[i] = pt_type(t)
                    # print(f" t={t} -> {val[i]}")

        # remap some names to the ones already in the default set of indexer keys
        for k, v in GribIndexer.KEYS_TO_REPLACE.items():
            if name in k:
                name = v

        return [(name, val)]

    @staticmethod
    def _to_datetime(param, val):
        try:
            if isinstance(val, datetime.datetime):
                return val
            elif isinstance(val, str):
                return utils.date_from_str(val)
            elif isinstance(val, (int, float)):
                return utils.date_from_str(str(val))
            else:
                raise
        except:
            raise Exception(f"Invalid datetime value={val} specified for key={param}")

    @staticmethod
    def _to_date(param, val):
        try:
            if isinstance(val, datetime.datetime):
                return val.date()
            elif isinstance(val, datetime.date):
                return val
            elif isinstance(val, str):
                d = utils.date_from_str(val)
                return d.date() if not isinstance(d, tuple) else d
            elif isinstance(val, (int, float)):
                d = utils.date_from_str(str(val))
                return d.date() if not isinstance(d, tuple) else d
            else:
                raise
        except:
            raise Exception(f"Invalid date value={val} specified for key={param}")

    @staticmethod
    def _to_time(param, val):
        try:
            if isinstance(val, (datetime.datetime)):
                return val.time()
            elif isinstance(val, datetime.time):
                return val
            elif isinstance(val, str):
                return utils.time_from_str(val)
            elif isinstance(val, int):
                return utils.time_from_str(str(val))
            else:
                raise
        except:
            raise Exception(f"Invalid time value={val} specified for key={param}")


class FieldsetIndexer(GribIndexer):
    def __init__(self, *args):
        super().__init__(*args)
        self.ref_column_count = 1

    def scan(self, vector=False):
        data = self._scan(self.db.fs, mapped_params=self.db.mapped_params)
        if data:
            df = self._make_dataframe(data, sort=False)
            self.db.blocks["scalar"] = df
            if vector:
                self._scan_vector()

    def _scan(self, fs, mapped_params={}):
        LOG.info(f" scan fields ...")
        data = {}
        # print(f"fs_len={len(fs)}")
        # print(f"keys_ecc={self.keys_ecc}")
        if utils.is_fieldset_type(fs) and len(fs) > 0:
            md_vals = fs.grib_get(self.keys_ecc, "key")
            if mapped_params:
                for i in range(len(fs)):
                    v = md_vals[self.param_id_index][i]
                    if v in mapped_params:
                        short_name = mapped_params[v]
                        md_vals[self.shortname_index][i] = short_name

            assert len(self.keys) == len(self.keys_ecc)
            data = {k: md_vals[i] for i, k in enumerate(self.keys)}
            data["_msgIndex1"] = list(range(len(fs)))
            LOG.info(f" {len(fs)} GRIB messages processed")
        return data

    def _scan_vector(self):
        df = self.db.blocks["scalar"]
        if df is not None and not df.empty:
            for v_name, v_comp in GribIndexer.VECTOR_PARAMS.items():
                r = self._build_vector_index(df, v_name, v_comp)
                comp_num = len(v_comp)
                if r:
                    cols = [*self.keys]
                    for i in range(comp_num):
                        cols.extend([f"_msgIndex{i+1}"])
                    w_df = self._make_dataframe(r, sort=False, columns=cols)
                    self.db.blocks[v_name] = w_df
                    # self._write_dataframe(w_df, v_name, out_dir)
                else:
                    LOG.debug(" No paired fields found!")
                    continue


class ExperimentIndexer(GribIndexer):
    def __init__(self, *args):
        super().__init__(*args)
        self.ref_column_count = 2

    def scan(self):
        out_dir = self.db.db_dir
        Path(out_dir).mkdir(exist_ok=True, parents=True)
        LOG.info(f"scan {self.db} out_dir={out_dir} ...")

        data = {k: [] for k in [*self.keys, "_msgIndex1", "_fileIndex1"]}
        input_files = []

        # print(f"out_dir={out_dir}")
        # merge existing experiment objects
        if self.db.merge_conf:
            ds = []
            # simple merge
            if isinstance(self.db.merge_conf, list):
                for c_name in self.db.merge_conf:
                    ds.append(
                        {"data": db.dataset.find(c_name), "name": c_name, "ens": {}}
                    )
            # explicit ENS merge
            else:
                assert "pf" in self.db.merge_conf
                # control forecast
                c_name = self.db.merge_conf.get("cf", "")
                if c_name != "":
                    ds.append(
                        {
                            "data": self.db.dataset.find(c_name, comp="field"),
                            "name": c_name,
                            "ens": {"type": "cf", "number": 0},
                        }
                    )
                for i, c_name in enumerate(self.db.merge_conf.get("pf", [])):
                    ds.append(
                        {
                            "data": self.db.dataset.find(c_name, comp="field"),
                            "name": c_name,
                            "ens": {"type": "pf", "number": i + 1},
                        }
                    )

            for c in ds:
                if c["data"] is None:
                    c_name = d["name"]
                    raise Exception(
                        f"Cannot merge experiments as {self.db}! Experiment {c_name} is not found!"
                    )
                else:
                    input_files = self._scan_one(
                        input_dir=c["data"].path,
                        file_name_pattern=c["data"].file_name_pattern,
                        input_files=input_files,
                        mapped_params=self.db.mapped_params,
                        ens=c["ens"],
                        data=data,
                        rootdir_placeholder_value=c["data"].rootdir_placeholder_value,
                        rootdir_placeholder_token=self.db.ROOTDIR_PLACEHOLDER_TOKEN,
                    )
        # index a single experiment
        else:
            input_files = self._scan_one(
                input_dir=self.db.path,
                file_name_pattern=self.db.file_name_pattern,
                input_files=[],
                mapped_params=self.db.mapped_params,
                ens={},
                data=data,
                rootdir_placeholder_value=self.db.rootdir_placeholder_value,
                rootdir_placeholder_token=self.db.ROOTDIR_PLACEHOLDER_TOKEN,
            )

        # print(f"input_files={input_files}")
        if len(input_files) > 0 and len(data["shortName"]) > 0:
            # write config file for input file list
            LOG.info(f"generate datafiles.yaml ...")
            f_name = os.path.join(out_dir, "datafiles.yaml")
            r = yaml.dump(input_files, default_flow_style=False)
            with open(f_name, "w") as f:
                f.write(r)
            self.db.input_files = input_files

            # scalar
            LOG.info(f"generate scalar fields index ...")
            df = self._make_dataframe(data, sort=True)
            self.db.blocks["scalar"] = df
            self._write_dataframe(df, "scalar", out_dir)

            # vector (2D)
            LOG.info(f"generate vector fields index ...")
            for v_name, v_comp in GribIndexer.VECTOR_PARAMS.items():
                r = self._build_vector_index(df, v_name, v_comp)
                comp_num = len(v_comp)
                if r:
                    cols = [*self.keys]
                    for i in range(comp_num):
                        cols.extend([f"_msgIndex{i+1}", f"_fileIndex{i+1}"])
                    w_df = self._make_dataframe(r, sort=True, columns=cols)
                    # print(f"wind_len={len(w_df.index)}")
                    self.db.blocks[v_name] = w_df
                    self._write_dataframe(w_df, v_name, out_dir)
                else:
                    LOG.debug(" No paired fields found!")
                    continue

    def _scan_one(
        self,
        input_dir="",
        file_name_pattern="",
        input_files=[],
        mapped_params={},
        ens={},
        data={},
        rootdir_placeholder_value="",
        rootdir_placeholder_token=None,
    ):
        LOG.info("scan fields ...")
        LOG.info(f" input_dir={input_dir} file_name_pattern={file_name_pattern}")
        # print(f" input_dir={input_dir} file_name_pattern={file_name_pattern}")

        # for f_path in glob.glob(f_pattern):
        cnt = 0
        input_files_tmp = []
        for f_path in utils.get_file_list(
            input_dir, file_name_pattern=file_name_pattern
        ):
            # LOG.debug(f"  f_path={f_path}")
            fs = self.db.fieldset_class(path=f_path)
            if utils.is_fieldset_type(fs) and len(fs) > 0:
                cnt += 1
                input_files_tmp.append(f_path)
                file_index = len(input_files) + len(input_files_tmp) - 1
                md_vals = fs.grib_get(self.keys_ecc, "key")

                if mapped_params:
                    for i in range(len(fs)):
                        v = md_vals[self.param_id_index][i]
                        if v in mapped_params:
                            short_name = mapped_params[v]
                            md_vals[self.shortname_index][i] = short_name
                if ens:
                    for i in range(len(fs)):
                        md_vals[self.type_index][i] = ens["type"]
                        md_vals[self.number_index][i] = ens["number"]

                assert len(self.keys) == len(self.keys_ecc)
                for i, c in enumerate(self.keys):
                    data[c].extend(md_vals[i])
                data["_msgIndex1"].extend(list(range(len(fs))))
                data["_fileIndex1"].extend([file_index] * len(fs))

                # print({k: len(v) for k, v in data.items()})

        if rootdir_placeholder_value:
            input_files_tmp = [
                x.replace(rootdir_placeholder_value, rootdir_placeholder_token)
                for x in input_files_tmp
            ]

        input_files.extend(input_files_tmp)

        LOG.info(f" {cnt} GRIB files processed")
        return input_files

    def allowed_keys(self):
        r = list(self.keys)
        r.extend(GribIndexer.DATE_KEYS)
        r.extend(GribIndexer.TIME_KEYS)
        r.extend(list(GribIndexer.DATETIME_KEYS.keys()))
        return set(r)