import asyncio
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import List

import aiohttp
import numpy as np
import pandas as pd
from aiohttp_retry import ExponentialRetry, RetryClient
from context import ClinicContext
from mathbox.app.signal.filter import moving_median
from mathbox.statistics.estimator import pearson_correlation_coefficient
from yaml import Loader, load


def timer(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        res = func(*args, **kwargs)
        end = time.perf_counter()
        print(f"Running time：{(end-start):.3f} s")
        return res

    return wrapper


def async_timer(func):
    async def process(func, *args, **params):
        if asyncio.iscoroutinefunction(func):
            return await func(*args, **params)
        else:
            print("this is not a coroutine")
            return func(*args, **params)

    async def helper(*args, **params):
        start = time.perf_counter()
        result = await process(func, *args, **params)
        print("Running time：", time.perf_counter() - start, "s")
        return result

    return helper


@dataclass
class Metric:
    # One PromQL query with name, timestamp, query response and signals
    name: str
    unit: str
    query: str
    component: str = ""
    rawdata: bool = False
    operators: list = field(default_factory=list, init=True)
    status: bool = field(
        default=False, init=True
    )  # lazy, if not init=True, __repr__ will not show this field
    err_msg: str = field(default="", init=True)
    stats: pd.DataFrame = field(default=pd.DataFrame(), init=True)
    insight: dict = field(default_factory=dict, init=True)
    signals: pd.DataFrame = field(default=pd.DataFrame(), init=True)
    imputed_signals: pd.DataFrame = field(default=pd.DataFrame(), init=True)

    def init_stats(self):
        self.stats = pd.DataFrame(index=self.signals.columns)

    def cal_sum(self):
        col = "sum"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(np.nansum(item))
        self.stats.insert(0, col, tmp)

    def cal_mean(self):
        col = "mean"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(np.nanmean(item))
        self.stats.insert(0, col, tmp)

    def cal_var(self):
        col = "var"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(np.nanvar(item))
        self.stats.insert(0, col, tmp)

    def cal_std(self):
        col = "std"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(np.nanstd(item))
        self.stats.insert(0, col, tmp)

    def cal_max(self):
        col = "max"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(np.nanmax(item))
        self.stats.insert(0, col, tmp)

    def cal_min(self):
        col = "min"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(np.nanmin(item))
        self.stats.insert(0, col, tmp)

    def cal_quantile(self, q):
        col = f"P{int(q * 100)}"
        tmp = []
        for _col_name, item in self.imputed_signals.iteritems():
            data = item.tolist()
            data.sort()
            tmp.append(data[int(len(data) * q)])
        self.stats.insert(0, col, tmp)

    def cal_overall_trend(self, step=15):
        col = "over_alltrend"
        tmp = []
        for _col_name, item in self.imputed_signals.iteritems():
            tmp.append(cal_trend(item, step))
        self.stats.insert(0, col, tmp)

    def cal_latest_trend(self, time_range=600, step=15):
        col = "latest_trend"
        tmp = []
        points = int(time_range / step) + 1
        points = min(points, self.imputed_signals.shape[0])
        for _col_name, item in self.imputed_signals.iteritems():
            tmp.append(cal_trend(item[-points:], step))
        self.stats.insert(0, col, tmp)

    def cal_pcc(self, other):
        index = f'pcc_{other.name.split(":")[0]}'
        if len(self.imputed_signals.columns) == 1 and len(other.imputed_signals.columns) == 1:
            self.insight[index] = pearson_correlation_coefficient(
                self.imputed_signals.iloc[:, 0].tolist(), other.imputed_signals.iloc[:, 0].tolist()
            )

    def cal_imbalance(self):
        if "by_instance" in self.name:
            index = "imbalance"
            tmp = []
            for i, row in self.imputed_signals.iterrows():
                # ratio = 1 - np.nanmin(row) / (np.nanmax(row) + np.nanmin(row)/1000000) # cannot handle min=0 and max~0
                ratio = np.nanstd(row) / (np.nansum(row) + 1e-10)
                if np.isnan(ratio):
                    ratio = 0
                tmp.append(ratio)
            # self.insight[index] = np.mean(tmp)
            self.insight[index] = np.quantile(tmp, 0.5)
            if len(self.signals.columns) == 1:
                self.insight[index] = 0

    def cal_latest(self):
        col = "latest"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(item.tolist()[-2])
        self.stats.insert(0, col, tmp)

    def cal_first(self):
        col = "first"
        tmp = []
        for _col_name, item in self.signals.iteritems():
            tmp.append(item.tolist()[0])
        self.stats.insert(0, col, tmp)

    def cal_anomaly(self):
        pass

    def cal_changepoint(self):
        pass

    def cal_falldown(self):
        start = next(idx for idx, row in self.imputed_signals.iterrows() if (row != 0).all())
        end = next(
            idx for idx, row in self.imputed_signals.iloc[::-1].iterrows() if (row != 0).all()
        )
        temp = self.imputed_signals.loc[start - 15 : end + 15]
        col = "falldown"
        tmp = []
        for col_name, item in temp.iteritems():
            col_fall = []
            for i in range(len(item.tolist()) - 2):
                if item.tolist()[i + 1] - item.tolist()[i] < -10:
                    col_fall.append(f"{temp.index[i]} UTC")
            tmp.append(col_fall)
        self.stats.insert(0, col, tmp)

    def check_stationary(self):
        pass

    def check_continuous_high(self, threshold, timerange):
        col = f"range_gt_{int(threshold*100)}%"
        tmp = []
        for col_name, item in self.imputed_signals.iteritems():
            time = []
            item_sign = [1 if v > threshold * 100 else 0 for v in item]
            res = cal_continuous_flag_slice(item_sign, timerange * 60 / 15 + 1)
            time_range = [
                (
                    self.imputed_signals.index.values.tolist()[item[0]],
                    self.imputed_signals.index.values.tolist()[item[1] - 1],
                )
                for item in res
            ]
            tmp.append(time_range)
        self.stats.insert(0, col, tmp)

    def aggr_insight(self):
        self.stats.replace(np.nan, None, inplace=True)
        self.stats.replace(np.inf, None, inplace=True)
        if "total" not in self.name:
            self.insight = self.stats.to_dict()
        else:
            for i, col in self.stats.iteritems():
                self.insight[i] = np.mean(col)


# @async_timer
async def preprocess(session, metric: Metric, context, desired_timestamp):
    """
    Query PromQL and write the result,which is not imputed as pandas.DataFrame into metric.signals.
    """
    try:
        resp = await context.get_metrics(
            session, metric.query, context.start, context.end
        )  # list of dicts
    except Exception as e:
        resp = {"ok": False, "reason": str(e), "metrics": None}

    metric.status = resp["ok"]
    metric.err_msg = resp["reason"]

    if metric.status:
        res = resp["metrics"]
    else:
        res = []

    step = 15
    timestamp = np.arange(context.start, context.end + step, step)
    cols = len(res) if len(res) > 0 else 1

    data = np.zeros((cols, len(timestamp)), dtype=np.float32)
    # data[:] = np.nan # prometheus is rediculous, request data is NaN, but treats as zero.
    col_names = []

    if len(res) > 0:
        for i, sig in enumerate(res):
            # sig: {"metric":{"name": "tidb-01"}, "values":[[t,v],[t,v],...,[t,v]]}
            if sig["metric"].get("name") is not None:
                sig_name = sig["metric"]["name"]
            elif sig["metric"].get("name3") is not None:
                sig_name = f'{sig["metric"].get("name1")}-{sig["metric"].get("name2")}-{sig["metric"].get("name3")}'
            else:
                sig_name = f'{sig["metric"].get("name1")}-{sig["metric"].get("name2")}'
            t, v = zip(*sig["values"])
            index = np.rint((np.array(t) - context.start) / step).astype(np.int)
            data[i, index] = [float(val) for val in v]
            col_names.append(sig_name)
        # metric.signals = pd.DataFrame(data.T, columns=col_names, index=pd.to_datetime(timestamp, unit="s"))
        metric.signals = pd.DataFrame(data.T, columns=col_names, index=timestamp)
        metric.imputed_signals = metric.signals.copy()
        metric.imputed_signals.interpolate(inplace=True)

        # step 1: check if too many NaNs, too many raise an errror
        # metric.signals.interpolate(method='index', inplace=True, axis=0, limit_direction='both') # step 2: imputation
    else:
        # metric.signals = pd.DataFrame(data.T, columns=[metric.name], index=pd.to_datetime(timestamp, unit="s"))
        metric.signals = pd.DataFrame(data.T, columns=[metric.name], index=timestamp)
        metric.signals.replace(0, np.nan, inplace=True)

    return metric


async def get_all_metrics(
    context: ClinicContext, fname: str, start: int, end: int, extra_promqls=None
):
    p = Path(__file__).with_name(fname)
    # promql 去空格去重和 validate 之类的事情
    with p.open("r") as f:
        metrics_req = load(f, Loader=Loader)

        desired_timestamp = list(range(start, end + 15, 15))

        promqls = []
        for item in metrics_req:
            promqls.append(Metric(item["name"], item["unit"], item["expr"], item.get("component")))
        for item in extra_promqls or []:
            promqls.append(
                Metric(
                    item["name"],
                    item["unit"],
                    item["expr"],
                    item.get("component"),
                    item.get("rawdata"),
                    item.get("operators"),
                )
            )

        conn = aiohttp.TCPConnector(limit=2)
        async with RetryClient(raise_for_status=False, connector=conn) as session:
            metrics = await asyncio.gather(
                *[preprocess(session, item, context, desired_timestamp) for item in promqls],
                return_exceptions=True,
            )
        return metrics


def search_metric(metrics, mname):
    results = [metric for metric in metrics if metric.name == mname]
    if len(results) == 0:
        return None
    else:
        return results[0]


def cal_trend(alist, step, window=9):
    if np.nan in alist:
        return np.nan
    alist_filtered = moving_median(alist, window)
    return (
        np.polynomial.polynomial.polyfit(list(range(0, len(alist_filtered))), alist_filtered, 1)[1]
        / step
    )


def cal_continuous_flag_slice(flag_slice, continuous_num=40):
    rising_edge = []
    falling_edge = []
    res = []
    for i in range(len(flag_slice) - 1):
        if flag_slice[i] == 0 and flag_slice[i + 1] == 1:
            rising_edge.append(i)
        if flag_slice[i] == 1 and flag_slice[i + 1] == 0:
            falling_edge.append(i)
    i = len(rising_edge)
    j = len(falling_edge)
    if i == 0 and j == 0:
        if flag_slice[0] == 1:
            res.append((0, len(flag_slice)))
    elif i == 0 and j == 1:
        res.append((0, falling_edge[0] + 1))
    elif i == 1 and j == 0:
        res.append((rising_edge[0] + 1, len(flag_slice)))
    elif rising_edge[0] < falling_edge[0]:  # rising fisrt
        for s in range(j):
            if falling_edge[s] - rising_edge[s] >= continuous_num:
                res.append((rising_edge[s] + 1, falling_edge[s] + 1))
        if j < i:
            if len(flag_slice) - rising_edge[-1] >= continuous_num:
                res.append((rising_edge[-1] + 1, len(flag_slice)))
    elif rising_edge[0] > falling_edge[0]:  # falling first
        if falling_edge[0] + 1 >= continuous_num:
            res.append((0, falling_edge[0] + 1))
        for s in range(j - 1):
            if falling_edge[s + 1] - rising_edge[s] >= continuous_num:
                res.append((rising_edge[s] + 1, falling_edge[s + 1] + 1))
        if i == j:
            if len(flag_slice) - rising_edge[-1] >= continuous_num:
                res.append((rising_edge[-1] + 1, len(flag_slice)))
    return res