#!/usr/bin/python
#
# Copyright 2018-2020 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 pandas as pd

from typing import Dict, List, Set, Union

import plotly.express as px

from polyaxon.client import RunClient
from polyaxon.client.decorators import check_no_op
from polyaxon.polyboard.artifacts import V1ArtifactKind
from polyaxon.polyboard.events import V1Events


class RunPlot(RunClient):
    @check_no_op
    def __init__(
        self, owner=None, project=None, run_uuid=None, client=None,
    ):
        super().__init__(owner=owner, project=project, run_uuid=run_uuid, client=client)
        self.metrics = {}
        self.metric_names = set([])

    @check_no_op
    def refresh_data(self):
        super().refresh_data()
        if self.metrics:
            self.get_metrics(self.metric_names)

    @check_no_op
    def get_metrics(self, names: Union[Set[str], List[str]]) -> Dict:
        events = self.get_events(
            kind=V1ArtifactKind.METRIC, names=names, orient=V1Events.ORIENT_DICT,
        ).data
        for e in events:
            self.metrics[e["name"]] = e
            self.metric_names.add(e["name"])
        return self.metrics

    @check_no_op
    def get_tidy_df(self):
        dfs = []
        for m in self.metric_names:
            data = self.metrics[m]
            df = V1Events.read(**data).df
            df["name"] = m
            dfs.append(df)
        return pd.concat(dfs)

    @check_no_op
    def get_wide_df(self):
        dfs = []
        for m in self.metric_names:
            data = self.metrics[m]
            df = V1Events.read(**data).df
            df[m] = df.metric
            df = df[["step", "timestamp", m]]
            dfs.append(df)

        return pd.concat(dfs, axis=1)

    @check_no_op
    def bar(
        self,
        x: str = "timestamp",
        y: str = "metric",
        color: str = "name",
        barmode: str = "group",
    ):
        df = self.get_tidy_df()
        return px.bar(df, x=x, y=y, color=color, barmode=barmode)

    @check_no_op
    def line(self, x: str = "timestamp", y: str = "metric", color: str = "name"):
        df = self.get_tidy_df()
        return px.line(df, x=x, y=y, color=color)

    @check_no_op
    def scatter(self, x: str = None, y: str = None, color: str = None, **kwargs):
        if len(self.metric_names) < 2:
            raise ValueError("You need at least 2 metrics to use this plot.")

        df = self.get_wide_df()
        return px.scatter(df, x=x, y=y, color=color, **kwargs)


class MultiRunPlot(RunClient):
    @check_no_op
    def __init__(
        self, owner=None, project=None, client=None,
    ):
        super().__init__(owner=owner, project=project, client=client)
        self.runs = {}
        self.run_uuids = set([])
        self.metric_names = set([])

    @check_no_op
    def refresh_data(self):
        super().refresh_data()
        if self.runs:
            self.get_runs(query="uuid:".format("|".join(self.run_uuids)))
            if self.metric_names:
                self.get_metrics(self.metric_names)

    @check_no_op
    def get_metrics(self, names: Union[Set[str], List[str]]) -> Dict:
        events = self.get_events(
            kind=V1ArtifactKind.METRIC, names=names, orient=V1Events.ORIENT_DICT,
        ).data
        for e in events:
            self.metrics[e["name"]] = e
            self.metric_names.add(e["name"])

        return self.metrics

    @check_no_op
    def get_runs(
        self, query: str = None, sort: str = None, limit: int = None, offset: int = None
    ) -> Dict:
        runs = self.list(query=query, sort=sort, limit=limit, offset=offset,).data
        for r in runs:
            self.run_uuids.add(r.uuid)
            self.runs[r.uuid] = r
        return self.runs

@@ -4268,8 +4268,858 @@
 lf.runs%0A
+%0A    def get_run_io(%0A        self, query: str = None, sort: str = None, limit: int = None, offset: int = None%0A    ):%0A        runs = self.get_runs(query=query, sort=sort, limit=limit, offset=offset)%0A        data = %5B%5D%0A        for r in runs:%0A            values = r.inputs or %7B%7D%0A            values.update(r.outputs or %7B%7D)%0A            data.append(%7B'uid': r.uuid, 'values': values%7D)%0A        return data%0A%0A    @check_no_op%0A    def get_hiplot(%0A        self, query: str = None, sort: str = None, limit: int = None, offset: int = None%0A    ):%0A        import hiplot%0A%0A        data = self.get_run_io(query=query, sort=sort, limit=limit, offset=offset)%0A        exp = hiplot.Experiment()%0A        for d in data:%0A            dp = hiplot.Datapoint(%0A                uid=d%5B%22uid%22%5D, values=d%5B%22values%22%5D,%0A            )%0A            exp.datapoints.append(dp)%0A        return exp%0A

Update runplot client with hiplot