import asyncio
import datetime
import json
from math import isnan
from pathlib import Path

import numpy as np
import pandas as pd
from context import ClinicContext
from mathbox.app.signal.filter import moving_median
from mathbox.calculus.differential import diff
from mathbox.statistics.estimator import mean, pearson_correlation_coefficient
from utils import cal_imbalance, get_all_metrics
from yaml import Loader, load


def cal_mean(all_metrics):
    for metrics in all_metrics:
        mean = {}
        _max = {}
        _min = {}
        try:
            metrics_des = metrics.signals.describe()
        except:
            raise Exception("No data in metrics")
        if metrics.signals.shape[1] == 1:
            if "by_instance" in metrics.name or "by_type" in metrics.name:
                mean[metrics.name] = {metrics.signals.columns[0]: metrics_des.loc["mean"].values[0]}
                _max[metrics.name] = {metrics.signals.columns[0]: metrics_des.loc["max"].values[0]}
                _min[metrics.name] = {metrics.signals.columns[0]: metrics_des.loc["min"].values[0]}
            else:
                mean[metrics.name] = metrics_des.loc["mean"].values[0]
                _max[metrics.name] = metrics_des.loc["max"].values[0]
                _min[metrics.name] = metrics_des.loc["min"].values[0]
        elif metrics.signals.shape[1] > 1:
            mean[metrics.name] = {}
            _max[metrics.name] = {}
            _min[metrics.name] = {}
            for col in metrics.signals.columns:
                mean[metrics.name][col] = metrics_des.loc["mean"][col]
                _max[metrics.name][col] = metrics_des.loc["max"][col]
                _min[metrics.name][col] = metrics_des.loc["min"][col]
        metrics.stats = {"mean": mean, "max": _max, "min": _min}


def reshape(metric_obj, stats="mean"):
    metric = metric_obj.name.split(":")[0]
    label = metric_obj.name.split(":")[1]
    for k, v in metric_obj.stats[stats].items():
        if isinstance(v, dict):
            for kk, vv in v.items():
                if isnan(vv):
                    v[kk] = None
        elif isnan(v):
            metric_obj.stats[stats][k] = None
    return {
        "id": metric_obj.id,
        "metrics": metric,
        "label": label,
        "unit": metric_obj.unit,
        "component": metric_obj.component,
        "type": metric_obj.type,
        "sub_metrics": metric_obj.sub_metrics,
        "only_sub": metric_obj.only_sub,
        "description": {"key": metric_obj.description, "variables": None},
        "grafana": metric_obj.panel,
        "value": metric_obj.stats[stats],
    }


async def handler(context: ClinicContext):
    start: int = int(context.start)
    end: int = int(context.stop)

    print("job start, perpare_data ...")

    template = "promql_op.yaml"
    if context.org_type == "cloud":
        template = "promql_cloud.yaml"
    all_metrics = await get_all_metrics(context, template, start, end)  # a list of panel dicts

    print("perpare_data done, start processing ...")
    cal_mean(all_metrics)

    stats = []
    for metric in all_metrics:
        stats.append(reshape(metric))
        # add additional non-mean stats
        if metric.name == "TPS:total":
            tps_min = reshape(metric, stats="min")
            tps_min["metrics"] = "TPS_min"
            tps_min["value"]["TPS_min:total"] = tps_min["value"].pop("TPS:total")
            stats.append(tps_min)
        if metric.name == "QPS:total":
            qps_min = reshape(metric, stats="min")
            qps_min["metrics"] = "QPS_min"
            qps_min["value"]["QPS_min:total"] = qps_min["value"].pop("QPS:total")
            qps_min["type"] = ["Workload"]
            qps_min["sub_metrics"] = []
            stats.append(qps_min)
        if metric.name == "Avg_cluster_duration:total":
            latency_max = reshape(metric, stats="max")
            latency_max["metrics"] = "Avg_cluster_duration_max"
            latency_max["value"]["Avg_cluster_duration_max:total"] = latency_max["value"].pop(
                "Avg_cluster_duration:total"
            )
            latency_max["type"] = ["Perf_basic"]
            stats.append(latency_max)

    return {"start_ts": end, "end_ts": start, "stats": stats}