from lazylinop.quantization import qbutterfly, qmonarch
from lazylinop.butterfly import Chain, ksm
import json
import time
import torch
import matplotlib
import matplotlib.pyplot as plt
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)


def benchmark(L, x, n1, n2):
    """
    To benchmark L @ x.
    """
    # Warmup
    for _ in range(n2):
        y = L @ x
    # Benchmark
    _type = x.device.type
    start = torch.Event(enable_timing=True)
    end = torch.Event(enable_timing=True)
    duration = [None] * n1
    for i in range(n1):
        if _type == 'cuda':
            start.record()
        else:
            start = time.time()
        for _ in range(n2):
            L @ x
        if _type == 'cuda':
            end.record()
            end.synchronize()
            duration[i] = 1e-3 * start.elapsed_time(end) / n2
            # print(L.ksm_data.duration["all"], duration[i])
        else:
            duration[i] = (time.time() - start) / n2
    return duration, y


duration = {"N": [], "batch": [], "base": {}, "target": {}}
n1, n2 = 100, 10
for n in [10, 11, 12]:
    # for b in [0, 4, 7, 9, 10, 11, 12, 13, 14, 15, 16]:
    for b in [1, 10, 100, 200, 500, 1000, 2000, 5000, 10000, 20000, 50000]:
        # All the (N, batch) tuples use the same seed.
        torch.manual_seed(1)
        backend = "xp"
        # N, batch = 2 ** n, 2 ** b
        N, batch = 2 ** n, b
        # Base and target dtypes.
        base = torch.float64
        target = torch.bfloat16
        # device = 'cpu'
        device = 'cuda:0'
        # Quantization of random Butterfly.
        chain = Chain.square_dyadic((N, N))
        # chain = Chain.monarch((N, N))
        if batch == 1:
            ksv = [torch.randn(*p).to(dtype=base, device=device) for p in chain.ks_patterns]
            ksvq, _ = qbutterfly(ksv, target, "l2r")
            # ksvq, _ = qmonarch(ksv, target)
            L = ksm(ksv, backend=backend)
            Q = ksm(ksvq, backend=backend)
        # xb with base dtype.
        xb = torch.randn(N, batch).to(dtype=base, device=device)
        # xt with target dtype.
        xt = xb.to(dtype=target)
        # Benchmark L @ xb versus Lq @ xt.
        durationb, yb = benchmark(L, xb, n1, n2)
        durationt, yt = benchmark(Q, xt, n1, n2)
        duration["N"].append(N)
        duration["batch"].append(batch)
        duration["base"][f"N={N}-batch={batch}"] = durationb
        duration["target"][f"N={N}-batch={batch}"] = durationt
        print(f"N={N} batch={batch}\n" +
              f"time(no quantization)={sum(durationb) / n1}\n" +
              f"time(quantization)={sum(durationt) / n1}\n" +
              f"relative error={torch.linalg.norm(yb - yt) / torch.linalg.norm(yb)}")
        if backend != "xp":
            Lq = ksm(ksvq, backend="xp")
            tmp, y = benchmark(Lq, xt, n1, n2)
            print(f"time(xp)={sum(tmp) / n1}")

    xx, yy, zz = [], [], []
    for k in duration["base"].keys():
        if f"N={N}" in k:
            xx.append(int(k.replace(f"N={N}-batch=", "")))
            yy.append(sum(duration["base"][k]) / n1)
            zz.append(sum(duration["target"][k]) / n1)
    fig = plt.figure("benchmark")
    abox = [0.15, 0.15, 0.7, 0.7]
    ax = fig.add_axes(
        abox,
        xlabel="batch",
        ylabel=r"$\Longleftarrow~\text{duration in second}$",
        xlim=(9e-1, 1e5),
        ylim=(1e-3, 1e-1),
        xscale="log",
        yscale="log",
    )
    ax.plot(xx, yy, color="black", linestyle="", marker="s", label="no quantization")
    ax.plot(xx, zz, color="violet", linestyle="", marker="^", label="quantization")
    ax.legend(ncol=1, loc="best", columnspacing=0.01,
              handletextpad=0.04, title_fontsize=10.0,
              fontsize=10.0, title=f"N={N}")
    for e in ["png", "svg"]:
        plt.savefig(f"benchmark_quantization_N{N}.{e}",
                    dpi=600, transparent=False, bbox_inches='tight')
    fig.clf()
    plt.close("benchmark")