import asyncio
import time
from dataclasses import dataclass
from math import isnan
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 f_lowpass_filter, moving_median
from mathbox.app.signal.outlier import noise_outlier, simple_outlier
from mathbox.na.interpolation import series_linear_interp
from mathbox.statistics.estimator import mean, median, std
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"运行时间：{(end-start):.3f} s")
        return res

    return wrapper


def async_timer(func):
    async def process(func, *args, **params):
        if asyncio.iscoroutinefunction(func):
            # print('this function is a coroutine: {}'.format(func.__name__))
            return await func(*args, **params)
        else:
            print("this is not a coroutine")
            return func(*args, **params)

    async def helper(*args, **params):
        # print('{}.time'.format(func.__name__))
        start = time.perf_counter()
        result = await process(func, *args, **params)

        # Test normal function route...
        # result = await process(lambda *a, **p: print(*a, **p), *args, **params)

        print("运行时间：", time.perf_counter() - start, "s")
        return result

    return helper


@dataclass
class Signal:
    # One Signal with name, timestamp and value
    timestamp: List[int]
    value: List[float]
    change_points = None
    outliers = None
    details = {}

    def __truediv__(self, other):
        if self.timestamp != other.timestamp:
            raise ValueError("Signal timestamp not equal")
        value = []
        for i in range(len(self.timestamp)):
            value.append(self.value[i] / (other.value[i] + 1e-10))
        return Signal(self.timestamp, value)

    def set_name(self, name):
        self.name = name

    def interp(self, desired_timestamp):
        self.value = series_linear_interp(self.timestamp, self.value, desired_timestamp)
        self.timestamp = desired_timestamp

    # outlier point dectection, return pos index of outlier points
    def outlier_detect(self, window=5, n_top_feq=3, f_low_rate=0.01, level=3):
        # must be interpolated first!!!
        if self.outliers is not None:
            return
        sig_filtered = moving_median(self.value, window)
        moved_trend = [x - y for x, y in zip(self.value, sig_filtered)]
        sample_period = self.timestamp[1] - self.timestamp[0]
        _, seasonality = f_lowpass_filter(
            moved_trend, sample_period, n_top_feq, f_low_rate / sample_period
        )
        moved_seasonality = [x - y for x, y in zip(moved_trend, seasonality)]
        _, outlier_hi = noise_outlier(moved_seasonality, level)
        # self.outliers = [x[0] for x in outlier_hi]
        filtered_outlier = []  # modify change points: remove small changes
        for item in outlier_hi:  # modify change points: remove small changes
            if item[0] < 20:  # modify change points: remove small changes
                segment = moved_seasonality[
                    item[0] : item[0] + 40
                ]  # modify change points: remove small changes
            else:  # modify change points: remove small changes
                segment = moved_seasonality[
                    item[0] - 20 : item[0] + 20
                ]  # modify change points: remove small changes
            if max(segment) > 0.5:  # 0.1 modify change points: remove small changes
                filtered_outlier.append(item[0])  # modify change points: remove small changes
        self.outliers = filtered_outlier  # modify change points: remove small changes

    # for stationary data, return [[pos_included, pos_excluded), [pos_included, pos_excluded), ...]]
    def outlier_detect_stationary(self, window=20, density=0.9):
        seg_pos = anomaly_slice(self.value, window, density)
        self.stationary_anomaly_seg = seg_pos

    def get_slice(self, start, stop):
        slice = Signal(self.timestamp[start:stop], self.value[start:stop])
        slice.set_name(self.name)
        return slice

    def add_baseline(self, baseline):
        # baseline is a Signal(timestamp, value) and set_name(name)
        self.details["baseline"] = baseline


@dataclass
class Metric:
    # One PromQL query with name, timestamp, query response and signals
    name: str
    query: str
    signals = pd.DataFrame()  # list of Signal

    def trend(self):
        res = {}
        for col in self.signals.columns:
            res[col] = cal_trend(self.signals[col])
        return pd.DataFrame(res, index=["trend"])

    def to_signals(self, resp):
        # response is {"status": "xxx", "data": {"resultType": "matrix", "result": [{"metrics":{<label_key>: <label_value>, ...}, "value":[[<unix>, "<value>"],...]},...] }}
        resp = resp["data"]["result"]
        for item in resp:
            node = item["metric"]["name"]  # finally get how to write label_replace()
            timestamp = []
            value = []
            for pair in item["values"]:
                timestamp.append(pair[0])
                value.append(0.0 if isnan(float(pair[1])) else float(pair[1]))
            signal = Signal(timestamp, value)
            signal.set_name(node)
            self.signals.append(signal)


# @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.
    """
    res = await context.get_metrics(
        session, metric.query, context.start, context.stop
    )  # list of dicts

    step = 15
    timestamp = np.arange(context.start, context.stop + 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]]}
            sig_name = sig["metric"]["name"]
            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.fillna(0, 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.replace(0, np.nan, inplace=True)

    return metric


async def get_all_metrics(context: ClinicContext, fname: str, start: int, end: int):
    tidb_cluster_id = str(context.cluster_id)
    p = Path(__file__).with_name(fname)
    f = p.open("r")
    metrics_req = load(f, Loader=Loader)

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

    promqls = []
    for item in metrics_req:
        query = item["expr"].replace("{tidb_cluster_id}", context.cluster_id)
        promqls.append(Metric(item["name"], query))

    conn = aiohttp.TCPConnector(limit=2)
    # retry_options = ExponentialRetry(attempts=5)
    # async with RetryClient(raise_for_status=False, retry_options=retry_options, connector=conn) as session:
    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


# offset unit is hour
# not adjust to new metric class with dataframe
def set_historical_baseline(context: ClinicContext, metrics: Metric, offset: int, end: int):
    start = end - offset * 3600
    desired_timestamp = list(range(start, end + 15, 15))
    metrics_baseline = Metric(metrics.name, metrics.query)
    metrics_resp = context.get_metrics(metrics.query, start, end)
    # if query res error, return None
    metrics_baseline.to_signals(metrics_resp)
    for signal in metrics_baseline.signals:
        signal.interp(desired_timestamp)
    for signal in metrics.signals:
        for sig in metrics_baseline.signals:
            if sig.name == signal.name:
                signal.add_baseline(sig)
                break


# not adjust to new metric class with dataframe
def cal_metrics_anomaly(metrics, window=20, density=0.8):
    # return anomaly_slice
    for sig in metrics.signals:
        if "baseline" in sig.details.keys():  # sig.details['baseline'] is not None:
            avg = mean(sig.details["baseline"].value)
        else:
            avg = avg_threshold[metrics.name]
        # if metrics.name == "latch_wait_duration":
        #        print("avg: ", avg, " max: ", max(sig.value))
        if max(sig.value) <= avg:
            sig.stationary_anomaly_seg = []
            continue
        if min(sig.value) >= avg * 2:
            sig.stationary_anomaly_seg = [(0, len(sig.timestamp) - 1)]
            continue
        # to be continued: there is still some case that need to be handeled (e.g. Q1 low Q3 high and MAX approx to Q3)
        sig.outlier_detect_stationary(window, density)
        for seg in sig.stationary_anomaly_seg[:]:
            if max(sig.value[seg[0] : seg[1]]) <= avg:
                sig.stationary_anomaly_seg.remove(seg)


# not adjust to new metric class with dataframe
def anomaly_slice(signal, window=40, density=0.8):
    continuous_anomaly_flag = [0 for _ in range(len(signal))]
    # assume signal is stable.
    _, outlier_hi = simple_outlier(signal, bias=1.5)
    anomaly_flag = [1 if i in ([i for (i, _) in outlier_hi]) else 0 for i in range(len(signal))]
    for i in range(len(signal) - window):
        if sum(anomaly_flag[i : i + window]) / window <= density:
            continue
        else:
            continuous_anomaly_flag[i : i + window] = [1 for _ in range(window)]
    pos_slice = cal_continuous_flag_slice(continuous_anomaly_flag, 0)
    return pos_slice


# not adjust to new metric class with dataframe
def cal_continuous_flag_slice(flag_slice, continuous_num=9):
    # [[include, exclude),[include, exclude),[include, exclude),[include, exclude),[include, exclude),...]
    # flag_slice = [1,1,1,1,1,1,1] # mock
    # flag_slice = [0,0,0,0,0,0,0]
    # flag_slice = [0,0,0,0,0,1,1,1,1,1,1]
    # flag_slice = [1,1,1,1,1,0,0,0,0,0]
    # flag_slice = [0,0,1,1,0,1,1,0,0,1,1,1,1,0,0,0,0,0,1,1,1,1,1,1,1,1,1,0,0,0] # mock
    # flag_slice = [0,0,1,1,1,1,1,0,0,1,1,1,1,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1,1,1] # mock
    # flag_slice = [1,1,0,0,0,0,0,1,1,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,1,1,1] # mock
    # flag_slice = [1,1,0,0,0,0,0,1,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0] # mock
    # continuous_num = 3 # mock

    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


# not adjust to new metric class with dataframe
def cal_imbalance(pieces, ratio_threshold=0.5):
    ratio = 1 - min(pieces) / (max(pieces) + 1e-10)  # not break change but important
    if ratio >= ratio_threshold:
        return True
    else:
        return False


# not adjust to new metric class with dataframe
def cal_trend(alist, 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
    ]


# not adjust to new metric class with dataframe
def search_panel(name, dicts):
    return next((item for item in dicts if item["panel_id"] == name), False)


# not adjust to new metric class with dataframe
def quantile(data, q):
    data.sort()
    return data[int(len(data) * q)]


def max_key_value(dict):
    key = max(dict, key=dict.get)
    return key, dict[key]