#!/usr/bin/python
#
# Copyright 2018-2021 Polyaxon, Inc.
#
# 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 polyaxon_sdk

from marshmallow import fields

from polyaxon.polyflow.params.params import ParamSchema, ParamValueMixin
from polyaxon.schemas.base import BaseCamelSchema, BaseConfig
from polyaxon.schemas.fields.ref_or_obj import RefOrObject


class JoinParamSchema(ParamSchema):
    ref = None

    @staticmethod
    def schema_config():
        return V1JoinParam


class V1JoinParam(BaseConfig, ParamValueMixin, polyaxon_sdk.V1JoinParam):
    SCHEMA = JoinParamSchema
    IDENTIFIER = "join_param"
    REDUCED_ATTRIBUTES = [
        "contextOnly",
        "connection",
        "toInit",
    ]

    @property
    def is_literal(self):
        return False

    @property
    def is_ref(self):
        return True

    @property
    def is_join_ref(self):
        return True

    @property
    def is_runs_ref(self):
        return False

    @property
    def is_ops_ref(self):
        return False


class JoinSchema(BaseCamelSchema):
    query = fields.Str(required=True)
    sort = fields.Str(allow_none=True)
    limit = RefOrObject(fields.Int(allow_none=True))
    offset = RefOrObject(fields.Int(allow_none=True))
    params = fields.Dict(
        keys=fields.Str(), values=fields.Nested(JoinParamSchema), allow_none=True
    )

    @staticmethod
    def schema_config():
        return V1Join


class V1Join(BaseConfig, polyaxon_sdk.V1Join):
    """Joins allow to query several runs based on a search specification.

    The result of the join will be a list of values based on the results from executing the search.

    A Join corresponds to a valid search that can be resolved by Polyaxon,
    the result will be injected to resolve the params referencing the join.

    ```yaml
    >>> joins:
    >>>   join1:
    >>>     query: "metrics.loss: <0.01"
    >>>     sort: "metrics.loss"
    >>>     limit: 5
    >>>     params:
    >>>       param1:
    >>>         value: inputs.param1
    >>>       result1:
    >>>         value: outputs.result1
    >>>       new_name:
    >>>         value: outputs.result2
    >>>       tensorboad_paths:
    >>>         value: artifacts.tensorboard
    >>>   join2:
    >>>     query: "metrics.accuracy: >.9"
    >>>     sort: "-metrics.accuracy"
    >>>     params:
    >>>       all_inputs:
    >>>         value: inputs
    >>>       all_outputs:
    >>>         value: outputs
    >>>       run_artifact_paths:
    >>>         value: artifacts
    >>>       uuids:
    >>>         value: globals.uuid
    >>>         contextOnly: true
    >>>       files:
    >>>         value: {files: ["subath/files", "subpath2/file2"]}
    >>>         toInit: true
    ```

    ```python
    >>> joins = {
    >>>     "join1": V1Join(
    >>>         query="metrics.loss: <0.01",
    >>>         sort="metrics.loss",
    >>>         limit=5,
    >>>         params={
    >>>             "param1": "inputs.param1",
    >>>             "result1": "outputs.result1",
    >>>             "new_name": "outputs.result2",
    >>>             "tensorboards": "artifacts.tensorboard.assets",
    >>>         },
    >>>     ),
    >>>     "join2": V1Join(
    >>>         query="metrics.accuracy: >.9",
    >>>         sort="-metrics.accuracy",
    >>>         params={
    >>>             "inputs": "inputs",
    >>>             "outputs": "outputs",
    >>>             "artifacts": "artifacts",
    >>>             "uuids": "globals.uuid",
    >>>         }
    >>>     )
    >>> }
    ```

    This will expose the artifacts generated by
    the top 5 runs where the loss is less than 0.01 ascending.
    """

    SCHEMA = JoinSchema
    IDENTIFIER = "join"
    REDUCED_ATTRIBUTES = [
        "sort",
        "limit",
        "offset",
    ]
