# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import pandas as pd

from ... import opcodes as OperandDef
from ...core import OutputType
from ...core.operand.base import SchedulingHint
from ...serialization.serializables import StringField
from ...tensor.datastore.to_vineyard import resolve_vineyard_socket
from ..operands import DataFrameOperand, DataFrameOperandMixin
from ..utils import parse_index

try:
    import vineyard
    from vineyard.data.dataframe import make_global_dataframe
    from vineyard.data.utils import to_json
except ImportError:
    vineyard = None


class DataFrameToVineyardChunk(DataFrameOperand, DataFrameOperandMixin):
    _op_type_ = OperandDef.DATAFRAME_STORE_VINEYARD_CHUNK

    # vineyard ipc socket
    vineyard_socket = StringField('vineyard_socket')

    def __init__(self, vineyard_socket=None, dtypes=None, **kw):
        super().__init__(vineyard_socket=vineyard_socket, _dtypes=dtypes,
                         _output_types=[OutputType.dataframe], **kw)

    def __call__(self, df):
        return self.new_dataframe([df], shape=(0, 0), dtypes=df.dtypes,
                                  index_value=df.index_value,
                                  columns_value=df.columns_value)

    @classmethod
    def _process_out_chunks(cls, op, out_chunks):
        dtypes = pd.Series([np.dtype('O')], index=pd.Index([0]))
        merge_op = DataFrameToVinyardStoreMeta(
                vineyard_socket=op.vineyard_socket,
                chunk_shape=op.inputs[0].chunk_shape,
                shape=(1, 1), dtypes=dtypes)
        return merge_op.new_chunks(out_chunks, shape=(1, 1), dtypes=dtypes,
                                   index=(0, 0))

    @classmethod
    def tile(cls, op):
        out_chunks = []
        scheduling_hint = SchedulingHint(fuseable=False)
        dtypes = pd.Series([np.dtype('O')], index=pd.Index([0]))
        for idx, chunk in enumerate(op.inputs[0].chunks):
            chunk_op = op.copy().reset_key()
            chunk_op.scheduling_hint = scheduling_hint
            out_chunk = chunk_op.new_chunk([chunk], shape=(1, 1), dtypes=dtypes,
                                           index_value=chunk.index_value,
                                           columns_value=chunk.columns_value,
                                           index=(idx, 0))
            out_chunks.append(out_chunk)
        out_chunks = cls._process_out_chunks(op, out_chunks)

        in_df = op.inputs[0]
        new_op = op.copy().reset_key()
        return new_op.new_dataframes(op.inputs, shape=(len(out_chunks), 1),
                                     dtypes=dtypes,
                                     index_value=in_df.index_value,
                                     columns_value=in_df.columns_value,
                                     chunks=out_chunks,
                                     nsplits=((np.prod(op.inputs[0].chunk_shape),),))

    @classmethod
    def execute(cls, ctx, op):
        if vineyard is None:
            raise RuntimeError('vineyard is not available')

        socket, needs_put = resolve_vineyard_socket(ctx, op)
        client = vineyard.connect(socket)

        # some op might be fused and executed twice on different workers
        if not needs_put:
            # might be fused
            try:  # pragma: no cover
                meta = ctx.get_chunks_meta([op.inputs[0].key])[0]
                df_id = vineyard.ObjectID(meta['object_ref'])
                if not client.exists(df_id):
                    needs_put = True
            except KeyError:
                needs_put = True
        if needs_put:
            df_id = client.put(ctx[op.inputs[0].key],
                               partition_index=op.inputs[0].index)
        else:  # pragma: no cover
            meta = client.get_meta(df_id)
            new_meta = vineyard.ObjectMeta()
            for k, v in meta.items():
                if k not in ['id', 'signature', 'instance_id']:
                    if isinstance(v, vineyard.ObjectMeta):
                        new_meta.add_member(k, v)
                    else:
                        new_meta[k] = v
            new_meta['partition_index_'] = to_json(op.inputs[0].index)
            df_id = client.create_metadata(new_meta).id

        client.persist(df_id)
        ctx[op.outputs[0].key] = pd.DataFrame({0: [df_id]})


class DataFrameToVinyardStoreMeta(DataFrameOperand, DataFrameOperandMixin):
    _op_type_ = OperandDef.DATAFRAME_STORE_VINEYARD_META

    # vineyard ipc socket
    vineyard_socket = StringField('vineyard_socket')

    def __init__(self, vineyard_socket=None, dtypes=None, **kw):
        super().__init__(vineyard_socket=vineyard_socket, dtypes=dtypes,
                         _output_types=[OutputType.dataframe], **kw)

    @classmethod
    def tile(cls, op):
        dtypes = pd.Series([np.dtype('O')], index=pd.Index([0]))
        chunk_op = op.copy().reset_key()
        out_chunk = chunk_op.new_chunk(op.inputs[0].chunks, shape=(1, 1), dtypes=dtypes,
                                       index_value=parse_index(pd.Index([-1])),
                                       columns_value=parse_index(pd.Index([0])),
                                       index=(0, 0))
        new_op = op.copy().reset_key()
        return new_op.new_dataframes(op.inputs, shape=(1, 1),
                                     dtypes=dtypes,
                                     index_value=parse_index(pd.Index([0])),
                                     columns_value=parse_index(pd.Index([0])),
                                     chunks=[out_chunk],
                                     nsplits=((1,), (1,)))

    @classmethod
    def execute(cls, ctx, op):
        if vineyard is None:
            raise RuntimeError('vineyard is not available')

        socket, _ = resolve_vineyard_socket(ctx, op)
        client = vineyard.connect(socket)

        # # store the result object id to execution context
        chunks = [ctx[chunk.key][0][0] for chunk in op.inputs]
        ctx[op.outputs[0].key] = pd.DataFrame(
                {0: [make_global_dataframe(client, chunks).id]})


def to_vineyard(df, vineyard_socket=None):
    op = DataFrameToVineyardChunk(vineyard_socket=vineyard_socket)
    return op(df)
