import numpy as np
from lazylinop.quantization.qbutterfly import _qbutterfly
from lazylinop.quantization.qmonarch import _qmonarch
from lazylinop.quantization.qrank_one import _qrank_one
from lazylinop.quantization.utils import finfo
from lazylinop.butterfly import Chain, ksm
import json
import time
import torch
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
import seaborn
from os import environ
plt.rcParams.update({"font.size": 12})
plt.rcParams.update({"lines.linewidth": 3})
plt.rcParams.update({"lines.markersize": 4})
if "OMP_NUM_THREADS" in environ.keys():
    torch.set_num_threads(int(environ["OMP_NUM_THREADS"]))
else:
    torch.set_num_threads(2)

base = torch.float64
targets = [torch.float8_e4m3fn, torch.bfloat16]
dev = ["cuda:0", "cpu"][0]
if dev == "cpu":
    start, end = 0, 0
else:
    start = torch.Event(enable_timing=True)
    end = torch.Event(enable_timing=True)
M = [128, 256, 512, 1024]
R, T = 100, len(targets)
xs, ys = torch.empty((3, len(M), R, T)), torch.empty((3, len(M), R, T))
ts = torch.empty((3, len(M), R, T))
xx = 1
colors = {
    str(targets[i]).replace("torch.", ""): [
        "orange", "violet", "darkblue", "darkgreen"][i] for i in range(T)}
markers = {
    str(targets[i]).replace("torch.", ""): [
        "o", "s", "*", "^"][i] for i in range(T)}

for w in range(3):
    # All the quantization functions use the same seed.
    np.random.seed(1)
    torch.manual_seed(1)
    what = ["qrank_one", "qbutterfly", "qmonarch"][w]
    for rnd in ["chop", "upcast_downcast"]:
        if rnd == "chop":
            continue
        for i, N in enumerate(M):
            for j in range(R):
                print(f"\n{what} M=N={N} repeat={j}")
                if what == "qbutterfly":
                    chain = Chain.square_dyadic((N, N))
                elif what == "qmonarch":
                    chain = Chain.monarch((N, N))
                elif what == "qrank_one":
                    x = (xx * torch.randn(N, 1)).to(dtype=base, device=dev)
                    y = (xx * torch.randn(N, 1)).to(dtype=base, device=dev)
                    chain = None
                if chain is not None:
                    ksv = [
                        (xx * torch.randn(*p)).to(
                            dtype=base, device=dev) for p in chain.ks_patterns]
                for k in range(T):
                    if dev == "cpu":
                        start = time.time()
                    else:
                        start.record()
                    if what == "qrank_one":
                        xh, yh, opt_rerr, rtn_rerr = _qrank_one(x, y, targets[k], rnd=rnd)
                    elif what == "qbutterfly":
                        ksvq, opt_rerr, rtn_rerr = _qbutterfly(ksv, targets[k], 'l2r', rnd=rnd)
                    elif what == "qmonarch":
                        ksvq, opt_rerr, rtn_rerr = _qmonarch(ksv, targets[k], rnd=rnd)
                    xs[w, i, j, k] = rtn_rerr
                    ys[w, i, j, k] = opt_rerr
                    if dev == "cpu":
                        ts[w, i, j, k] = time.time() - start
                    else:
                        end.record()
                        end.synchronize()
                        ts[w, i, j, k] = 1e-3 * start.elapsed_time(end)
                    print(f"{targets[k]} mantissa={finfo(targets[k]).nmant} bits duration={ts[w, i, j, k]}")

        try:
            with open(f"{what}_{rnd}_df.json", 'r') as in_file:
                df = json.load(in_file)
        except IOError:
            df = {"N": [], "rerr": [], "ratio": [], "duration": [], "algorithm": [], "target": [], "time": {}}
        for k in range(T):
            target = str(targets[k]).replace("torch.", "")
            df["time"][target] = {"N": [], "duration": []}
            for i, N in enumerate(M):
                df["N"] += [N] * (2 * R)
                df["rerr"] += xs[w, i, :, k].tolist()
                df["ratio"] += [None] * R
                df["duration"] += [None] * R
                df["algorithm"] += ["naive rounding"] * R
                df["rerr"] += ys[w, i, :, k].tolist()
                df["ratio"] += (100.0 * (1.0 - ys[w, i, :, k] / xs[w, i, :, k])).tolist()
                df["duration"] += ts[w, i, :, k].tolist()
                df["algorithm"] += ["optimal quantization"] * R
                df["target"] += [target] * (2 * R)
                df["time"][target]["N"].append(N)
                df["time"][target]["duration"].append(sum(ts[w, i, :, k].tolist()) / R)
        with open(f"{what}_{rnd}_df.json", "w") as fout:
            json.dump(df, fout, indent=1)
        # Draw duration.
        fig = plt.figure("duration")
        abox = [0.15, 0.15, 0.7, 0.7]
        ax = fig.add_axes(
            abox,
            xlabel="N",
            ylabel="duration",
            xlim=(10, 3e3),
            ylim=(1e-2, 1e2),
            xscale="log",
            yscale="log",
        )
        for k in df["time"].keys():
            x = np.asarray(df["time"][k]["N"])
            y = np.asarray(df["time"][k]["duration"])
            ax.plot(x, y, linestyle="",
                    marker=markers[k], color=colors[k], label=k)
            # y = a * x ^ b
            # log(y) = log(a) + b * log(x)
            z = np.polyfit(np.log(x), np.log(y), deg=1)
            b = z[0]
            a = np.exp(z[1])
            ax.plot(x, a * (x ** b), color=colors[k],
                    linestyle='-', linewidth=1.0, marker='',
                    label=r"${0:.1e}N^{{{1:.2f}}}$".format(a, b))
        ax.legend(ncol=2, loc="best", columnspacing=0.08,
                  handletextpad=0.08, title_fontsize=10.0, fontsize=10.0)
        for e in ["png", "svg"]:
            plt.savefig(f"{what}_{rnd}_duration.{e}",
                        dpi=600, transparent=False, bbox_inches='tight')
        fig.clf()
        plt.close("duration")
        # Draw accuracy gain.
        # del df["time"]
        df.pop("time", None)
        df = pd.DataFrame(df)
        fig = plt.figure("accuracy")
        abox = [0.15, 0.15, 0.7, 0.7]
        ax = fig.add_axes(
            abox,
            xlabel="N",
            ylabel=r"$\text{Accuracy gain (%)}~\Longrightarrow$",
            xlim=(0, 5),
            ylim=(40, 100) if what == "qbutterfly" else (0, 100),
            xscale="linear",
            yscale="linear",
        )
        ax = seaborn.boxplot(data=df, x="N", y="ratio", hue="target",
                             hue_order=[str(t).replace("torch.", "") for t in targets],
                             palette=colors)
        for e in ["png", "svg"]:
            plt.savefig(f"{what}_{rnd}_accuracy_gain.{e}",
                        dpi=600, transparent=False, bbox_inches='tight')
        fig.clf()
        plt.close("accuracy")