#!/usr/bin/env python
# encoding: utf-8
# TODO : Remove this file, integrate filters directly to core api
"""
dataset.py : Interaction with DSS datasets
Copyright (c) 2013-2014 Dataiku SAS. All rights reserved.
"""

from typing import Optional
import numpy as np
from dataiku.base.utils import check_base_package_version
from dataiku.core.dataset import Schema, create_sampling_argument, unique
import warnings
import logging
import uuid

try:
    import pandas as pd

    check_base_package_version(
        pd, "pandas", "0.23.0", None, "DSS requires version 0.23 or above"
    )  # keep the version number in sync with install-python-packages.sh
except ImportError as e:
    logging.exception("Pandas import failure")
    warnings.warn(
        "Could not import pandas (%s). Pandas support will be disabled. To enable get_dataframe and other methods, please install the pandas package"
        % (e),
        Warning,
    )

from dataiku.core import default_project_key
import json, os

# Module code
from dataiku.core import flow, schema_handling
from dataiku.core import intercom

FULL_SAMPLING = {"samplingMethod": "FULL"}
CSV_SEP = "\t"
CSV_DOUBLE_QUOTE = True
CSV_QUOTE_CHAR = '"'


DEFAULT_TIMEOUT = 30
if flow.FLOW is not None:
    # Timeout has been introduced to cope with ipython leaks.
    # As a default value, we have an infinite timeout when in flow.
    DEFAULT_TIMEOUT = -1

# We want to stderr something on DeprecationWarning
# But don't reset everything because pandas has set up some filters
warnings.filterwarnings("default", category=DeprecationWarning)

(
    GENERATING,  # underlying  generator is currently working
    IDLE,  # waiting for the generator user to call .next()
    TIMEOUT_REACHED,  # timeout has been reached
    END_OF_ITERATOR,
    TERMINATED,
) = range(
    5
)  # we reached the generator last element.


class Dataset:
    """This is a handle to obtain readers and writers on a dataiku Dataset.
    From this Dataset class, you can:

    * Read a dataset as a Pandas dataframe
    * Read a dataset as a chunked Pandas dataframe
    * Read a dataset row-by-row
    * Write a pandas dataframe to a dataset
    * Write a series of chunked Pandas dataframes to a dataset
    * Write to a dataset row-by-row
    * Edit the schema of a dataset"""

    def __init__(self, name, project_key=None, ignore_flow=False):
        self.name = name
        self.cols = None
        self.partitions = None
        self.read_partitions = None
        self.writePartition = None
        self.writable = False
        self.readable = False
        self.preparation_steps = None
        self.preparation_requested_output_schema = None
        self.preparation_context_project_key = None
        self.ignore_flow = ignore_flow

        # Flow mode, initialize partitions to read and write and read/write flags
        if flow.FLOW is not None and ignore_flow == False:
            for input_dataset in flow.FLOW["in"]:
                if (
                    input_dataset["smartName"] == self.name
                    or input_dataset["fullName"] == self.name
                ):
                    self.readable = True
                    self.name = input_dataset["fullName"]
                    if "partitions" in input_dataset:
                        self.read_partitions = input_dataset["partitions"]
            for output_dataset in flow.FLOW["out"]:
                if (
                    output_dataset["smartName"] == self.name
                    or output_dataset["fullName"] == self.name
                ):
                    self.name = output_dataset["fullName"]
                    self.writable = True
                    self.spec_item = output_dataset
                    if "partition" in output_dataset:
                        self.writePartition = output_dataset["partition"]
            if not self.readable and not self.writable:
                raise Exception(
                    "Dataset %s cannot be used : declare it as input or output of your recipe"
                    % self.name
                )
            (self.project_key, self.short_name) = self.name.split(".", 1)

        else:
            if "." not in name:
                try:
                    self.project_key = project_key or default_project_key()
                    self.short_name = name
                    self.name = self.project_key + "." + name
                except:
                    logging.exception("Failure happened")
                    raise Exception(
                        "Dataset %s is specified with a relative name, "
                        "but no default project was found. Please use complete name"
                        % self.name
                    )
            else:
                # use gave a full name
                (self.project_key, self.short_name) = self.name.split(".", 1)
                if project_key is not None and self.project_key != project_key:
                    raise ValueError(
                        "Project key %s incompatible with fullname dataset %s."
                        % (project_key, name)
                    )
            self.readable = True
            self.writable = True
            self.spec_item = {"appendMode": False}  # notebook always overwrites

    @property
    def full_name(
        self,
    ):
        return self.project_key + "." + self.short_name

    def _repr_html_(
        self,
    ):
        s = "Dataset[   <b>%s</b>   ]</br>" % self.name
        s += self.read_schema()._repr_html_()
        return s

    def read_schema(self, raise_if_empty=True):
        """Gets the schema of this dataset, as an array of objects like this one:
        { 'type': 'string', 'name': 'foo', 'maxLength': 1000 }.
        There is more information for the map, array and object types.
        """
        if self.cols is None:
            flow_forced_schemas = os.getenv("FLOW_FORCED_SCHEMAS")
            if flow_forced_schemas is not None:
                ffs = json.loads(flow_forced_schemas)
                if self.full_name in ffs:
                    logging.info("Forcing schema: %s" % ffs[self.full_name])
                    return ffs[self.full_name]["columns"]

            self.cols = intercom.jek_or_backend_json_call(
                "datasets/get-schema/",
                data={"fullDatasetName": self.full_name},
                err_msg="Unable to fetch schema for %s" % (self.name),
            ).get("columns")

        if raise_if_empty and len(self.cols) == 0:
            raise Exception(
                "No column in schema of %s."
                " Have you set up the schema for this dataset?" % self.name
            )
        return Schema(
            self.cols,
        )

    def _stream(
        self,
        infer_with_pandas=True,
        sampling="head",
        sampling_column=None,
        limit=None,
        ratio=None,
        columns=None,
        read_session_id=None,
        filter=None,
    ):
        if not self.readable:
            raise Exception(
                "You cannot read dataset %s, "
                "it is not declared as an input" % self.name
            )
        if flow.FLOW is not None:
            add_env = {"DKU_FLOW": "1"}
        else:
            add_env = {}

        sampling_params = create_sampling_argument(
            sampling=sampling,
            sampling_column=sampling_column,
            limit=limit,
            ratio=ratio,
        )

        if self.preparation_steps is not None:
            data = {
                "fullDatasetName": self.full_name,
                "script": json.dumps({"steps": self.preparation_steps}),
                "requestedOutputSchema": json.dumps(
                    self.preparation_requested_output_schema
                ),
                "contextProjectKey": self.preparation_context_project_key,
                "sampling": json.dumps(sampling_params),
                "readSessionId": read_session_id,
            }
            if self.read_partitions is not None:
                data["partitions"] = json.dumps(self.read_partitions)

            return intercom.jek_or_backend_stream_call(
                "datasets/stream-prepared-dataset/",
                data=data,
                err_msg="Failed to read prepared data",
            )

        else:
            data = {
                "projectKey": self.project_key,
                "datasetName": self.short_name,
                "sampling": json.dumps(sampling_params)
                if sampling_params is not None
                else None,
                "columns": ",".join(columns) if columns is not None else None,
                "format": "tsv-excel-noheader",
                "partitions": ",".join(self.read_partitions)
                if self.read_partitions is not None
                else None,
                "readSessionId": read_session_id,
                "filter": filter,
            }
            return intercom.jek_or_backend_stream_call(
                "datasets/read-data/",
                data=data,
                err_msg="Failed to read dataset stream data",
            )

    def _verify_read(self, read_session_id):
        intercom.jek_or_backend_void_call(
            "datasets/verify-read/",
            data={"readSessionId": read_session_id},
            err_msg="Reading dataset failed",
        )

    @staticmethod
    def get_dataframe_schema_st(
        schema,
        columns=None,
        parse_dates=True,
        infer_with_pandas=False,
        bool_as_str=False,
        int_as_float=False,
    ):
        names = []
        dtypes = {}
        for col in schema:
            n = col["name"]
            t = col["type"]
            if bool_as_str and t == "boolean":
                dtypes[n] = "str"  # see df_from_split_desc
            elif int_as_float and t in ["tinyint", "smallint", "int", "bigint"]:
                dtypes[n] = "float64"
            elif t in schema_handling.DKU_PANDAS_TYPES_MAP:
                dtypes[n] = schema_handling.DKU_PANDAS_TYPES_MAP[t]
            else:
                dtypes[n] = np.object_
            names.append(n)
        if columns is not None:
            columns = list(unique(columns))
            names = columns
            dtypes = {
                column_name: dtypes[column_name]
                for column_name in dtypes
                if column_name in columns
            }

        # if parse_dates is set to True,
        # list up the index of the columns set up as dates by DSS
        # and forward them to pandas.
        if parse_dates is True:
            parse_dates = [
                col_id
                for (col_id, col_schema) in enumerate(schema)
                if col_schema["type"] == "date"
                and (columns is None or col_schema["name"] in columns)
            ]
            if len(parse_dates) == 0:
                parse_dates = False
        if infer_with_pandas:
            if bool_as_str:
                dtypes = dict(
                    (c["name"], "str") for c in schema if c["type"] == "boolean"
                )
            else:
                dtypes = None
        return (names, dtypes, parse_dates)

    def _get_dataframe_schema(
        self, columns=None, parse_dates=True, infer_with_pandas=False, bool_as_str=False
    ):
        if self.preparation_steps is not None:
            return Dataset.get_dataframe_schema_st(
                self.preparation_requested_output_schema["columns"],
                columns,
                parse_dates,
                infer_with_pandas,
                bool_as_str,
            )
        else:
            return Dataset.get_dataframe_schema_st(
                self.read_schema(), columns, parse_dates, infer_with_pandas, bool_as_str
            )

    def iter_dataframes(
        self,
        chunksize=10000,
        infer_with_pandas=True,
        sampling="head",
        sampling_column=None,
        parse_dates=True,
        limit=None,
        ratio=None,
        columns=None,
        bool_as_str=False,
        float_precision=None,
        na_values=None,
        keep_default_na=True,
        filter: Optional[str] = None,
    ):
        """Read the dataset to Pandas dataframes by chunks of fixed size.

        Returns a generator over pandas dataframes.

        Useful is the dataset doesn't fit in RAM."""
        if not self.readable:
            raise Exception(
                "You cannot read dataset %s, "
                "it is not declared as an input" % self.name
            )
        (names, dtypes, parse_date_columns) = self._get_dataframe_schema(
            columns=columns,
            parse_dates=parse_dates,
            infer_with_pandas=infer_with_pandas,
            bool_as_str=bool_as_str,
        )
        read_session_id = str(uuid.uuid4())
        with self._stream(
            infer_with_pandas=infer_with_pandas,
            sampling=sampling,
            sampling_column=sampling_column,
            limit=limit,
            ratio=ratio,
            columns=columns,
            read_session_id=read_session_id,
            filter=filter,
        ) as dku_output:
            df_it = pd.read_table(
                dku_output,
                dtype=dtypes,
                names=names,
                low_memory=True,
                header=None,
                sep=CSV_SEP,
                doublequote=CSV_DOUBLE_QUOTE,
                chunksize=chunksize,
                iterator=True,
                parse_dates=parse_date_columns,
                float_precision=float_precision,
                na_values=na_values,
                keep_default_na=keep_default_na,
            )
            logging.info("Starting dataframes iterator")
            for df in df_it:
                yield df

        # stream seems to have run fine. 'Seems'. Verify that.
        # note to self: this call has to be made after the dataframe creation, because it is streamed
        self._verify_read(read_session_id)
