import pickle
import os
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.ticker
from scipy import stats
import numpy as np

import plot_utils
import run_expr_6

config_order = ["counter", "unroll", "adaptive"]  
bench_order = ["fft", "basicmath", "crc", "sha", "stringsearch", "aes", "conv2d", "matmul"]

def validate_and_cleanup_df(df: pd.DataFrame):

    # mask = df["is_correct"] == False
    # if mask.any():
    #     print(f"drop {mask.sum()} rows with incorrect output")
    # df = df[~mask].copy()
    # assert(df["is_correct"].all())
    
    # bench_name = df.iloc[0,0]
    # z_score = stats.zscore(df["time_taken"])
    # mask2 = np.abs(z_score) > 4.35
    # if mask2.any():
    #     print(f"({bench_name}) remove {mask2.sum()} outliers")
    #     time = float(df[mask2]["time_taken"].iloc[0])
    #     average = df["time_taken"].mean()
    #     print(f"time: {time:.2f}, average: {average:.2f}")

    # df = df[~mask2].copy()

    if df.empty:
        return df

    def rename_bench_name(row):
        return row["bench_name"].lower()[1:]

    df["bench_name"] = df.apply(rename_bench_name, axis=1)

    def process_stats(row):
        lines = row["stats"]
        result = {}
        try:
            for line in lines:
                key, val = line.split(":")
                val = int(val.strip())
                if key.startswith("checkpoint triggered"):
                    result["ckpt_trig"] = val
                if key.startswith("checkpoint executed"):
                    result["ckpt_exec"] = val
                if key.startswith("no forward progress"):
                    result["no_forward_progress"] = int(val)
        except:
            result["ckpt_trig"] = -2
            result["ckpt_exec"] = -2
            result["no_forward_progress"] = -2
        return result

    df_stats = df.apply(process_stats, axis=1, result_type="expand")
    df = pd.concat([df, df_stats], axis=1)

    drop_columns = [
        # "start",
        # "end",
        "recovery",
        "outputs",
        "stats",
        # "ckpt",
        # "VC",
        # "PC",
        "is_correct"
    ]
    df = df.drop(columns = drop_columns) 
    return df


def get_base_df(benchmark, configs):
    output_dir = "/home/ybkim/workspace/imc/imc_freertos_app_m33/imc/exprs/date2025/6_varying_power/output"

    all_dfs = []
    for config in configs:
        pickle_filename = f"{output_dir}/{config}/{benchmark}.pickle"

        with open(pickle_filename, "rb") as f:
            orig_df = pickle.load(f)
            df = validate_and_cleanup_df(orig_df)
        df["config"] = config
        all_dfs.append(df)

    orig_df = pd.concat(all_dfs)

    def rename_config(row):
        name = row["config"]
        d = {
            "adaptive": "adaptive",
            "pass_count": "counter",
            "unroll": "unroll"
        }
        return d[name]
    orig_df["config"] = orig_df.apply(rename_config, axis=1)

    return orig_df


def remove_error_stats(orig_df):
    mask = orig_df["ckpt_trig"].abs() > 1e5 
    df = orig_df[~mask]
    mask = df["ckpt_exec"].abs() > 1e6
    df = df[~mask]
    mask = df["ckpt_trig"] < 0
    df = df[~mask]
    mask = df["ckpt_exec"] < 0
    df = df[~mask]
    return df


def add_secondary_power_graph(ax, wait_time):
    # ax2 = ax.twinx()
    power_data = "data/2015-04-23-Sacramento.csv"
    power_data = "data/2015-04-22-Sacramento.csv"
    power_df = run_expr_6.parse_energy_trace_nrel(power_data)
    df = power_df.reset_index()
    df["index"] = df["index"] * wait_time
    df = df.rename(columns={
        "Global Horizontal [W/m^2]": "current (mA)"
    })

    g = sns.lineplot(
        df,
        x="index",
        y="current (mA)",
        color="dimgray",
        linewidth = 0,
        alpha=0,
        ax=ax,
    )
    ax.fill_between(df["index"], df["current (mA)"], alpha=0.1)
    ax.grid(False)
    ax.set_ylabel("Current (mA)")
    return ax


def draw_graph(benchmark="vSha"):
    configs = ["adaptive"]
    orig_df = get_base_df(benchmark, configs)
    # orig_df = remove_error_stats(orig_df)

    wait_time = orig_df["wait_time"].unique()[0]
    # wait_time = 1

    start_time = orig_df.iloc[0]["start"]
    start_time = orig_df.attrs["start_time"]
    df = orig_df.copy()
    df["time_passed"] = orig_df["end"] - start_time
    drop_columns = ["start", "end", "time_taken", "debug_outputs", "wait_time", "config"]
    df = df.drop(columns=drop_columns)

    fig_size = (12, 3)
    n_rows = 1
    n_cols = 1
    hspace = 0.05
    fig = plt.figure(figsize=fig_size)
    ax = fig.subplots(n_rows, n_cols, squeeze=True, sharex=True, gridspec_kw={"hspace": hspace})

    linewidth = 1.8

    df2 = df.melt(id_vars=["bench_name", "time_passed"], value_vars=["ckpt_trig", "ckpt_exec"])
    mask = df["no_forward_progress"] > 0
    # display(df)
    print(df[mask])

    plot_utils.set_theme_seaborn(kind="line")
    sns.lineplot(
        df,
        x="time_passed",
        y="ckpt_trig",
        marker="x",
        markeredgecolor="tomato",
        color="tomato",
        markersize=0,
        markeredgewidth=1.5,
        linewidth = linewidth,
        # hue="config",
        ax=ax
    )
    ax.set_ylim([-10, 1000])
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Ckpt tirggered")
    ax.set_title("")

    add_secondary_power_graph(ax.twinx(), wait_time)
    return fig

def draw_completeions(benchmark="vSha"):
    configs = [
        "adaptive",
        "pass_count",
        "unroll"
    ]
    orig_df = get_base_df(benchmark, configs)
    # orig_df = remove_error_stats(orig_df)
    wait_time = orig_df["wait_time"].unique()[0]

    start_time = orig_df.iloc[0]["start"]
    df = orig_df.copy()
    df["time_passed"] = orig_df["end"] - start_time
    # drop_columns = ["start", "end", "time_taken", "wait_time", "ckpt_trig", "ckpt_exec", "bench_name"]
    # df = df.drop(columns=drop_columns)
    start_times = {}
    for config in configs:
        start_time = get_base_df(benchmark, [config]).attrs["start_time"]
        key = config if config != "pass_count" else "counter"
        start_times[key] = start_time

    def count_mapper(d):
        start_time = d.iloc[0]["start"]
        config = d["config"].unique()[0]
        start_time = start_times[config]
        d["time_passed"] = d["end"] - start_time
        total_finished = len(d)
        d["finished"] = list(range(1, total_finished+1))
        return d
    df = df.groupby(["config"], as_index=False).apply(count_mapper)

    fig_size = (8, 4)
    n_rows = 1
    n_cols = 1
    hspace = 0.05
    fig = plt.figure(figsize=fig_size)
    ax = fig.subplots(n_rows, n_cols, squeeze=True, sharex=True, gridspec_kw={"hspace": hspace})

    linewidth = 1.8
    plot_utils.set_theme_seaborn(kind="line")
    g = sns.lineplot(
        df,
        x="time_passed",
        y="finished",
        # marker="x",
        # markeredgecolor="none",
        # markers=["x", "x", "x"],
        # markersize=6,
        # markeredgewidth=1.5,
        linewidth = linewidth,
        hue="config",
        # style="config",
        # dashes=False,
        hue_order=["counter", "unroll", "adaptive"],
        # legend_labels=["counter", "unroll", "FastTrack"],
        ax=ax
    )

    colors = {
        "counter": "dimgray",
        "unroll": "cornflowerblue",
        "adaptive": "tomato"
    }


    num_finished = {}

    for key, d in df.groupby("config"):
        key_color = colors[key]
        sns.scatterplot(
            d,
            x="time_passed",
            y="finished",
            marker="x",
            color=key_color,
            s=8,
            legend=False,
            linewidth=0.5,
            ax=ax
        )
        num_finished[key] = len(d)
    
    print("number of finished tasks:")
    print(num_finished)
    print("ratio to adaptive:")
    print("counter:", num_finished["adaptive"]/num_finished["counter"])
    print("unroll:", num_finished["adaptive"]/num_finished["unroll"])

    for key, d in df.groupby("config"):
        mask = d["time_passed"] < 200
        columns = ["time_passed", "finished"]
        d2 = d[mask]
        print(key)
        # print(d2[columns].tail(1))
        finished = len(d2)
        rate = finished / 200
        print("finished during 0~200 secs:", finished, rate)

        mask = (d["time_passed"] >= 300) & (d["time_passed"] < 500)
        d2 = d[mask]
        finished = len(d2)
        rate = finished / 200
        print("finished during 300~500 secs:", finished, rate)
    

    # ax.set_ylim([-50, 1000])
    ax.legend(ncol=3, loc="upper center", bbox_to_anchor=(0.5, 1.4), labelspacing=0.2)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Finished Tasks")
    add_secondary_power_graph(ax.twinx(), wait_time)
    
    legend_labels = ["counter", "unroll", "FastTrack"]
    for t in ax.get_legend().get_texts():
        if t.get_text() == "adaptive":
            t.set_text("FastTrack")

    return fig


def draw_power_trace():
    ratio = 0.9
    fig_width = 4.5
    fig_size = (fig_width, fig_width * ratio)
    n_rows = 1
    n_cols = 1
    hspace = 0.05
    fig = plt.figure(figsize=fig_size)
    ax = fig.subplots(n_rows, n_cols, squeeze=True, sharex=True, gridspec_kw={"hspace": hspace})

    power_data = "data/2015-04-23-Sacramento.csv"
    power_data = "data/2015-04-22-Sacramento.csv"
    power_df = run_expr_6.parse_energy_trace_nrel(power_data, include_datetime=True)
    df = power_df.reset_index()
    df = df.rename(columns={
        "Global Horizontal [W/m^2]": "current (mA)"
    })

    print(df)

    sns.lineplot(
        df,
        x="PST",
        y="current (mA)",
        color="dimgray",
        linewidth = 1,
        ax=ax,
    )
    ax.grid(True)
    ax.set_ylabel("Current (mA)")
    ax.set_xlabel("Time")
    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator(3))
    ax.yaxis.grid(visible=True, which="both")
    ax.xaxis.set_major_locator(matplotlib.ticker.MaxNLocator(4))
    return fig
