import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np

from utils_discrete import mle_entropy
from utils_dynamic import OUTPUT_DIR, ensure_dir
from utils_runtime import timed_main


def mod_5(x):
    return np.asarray(x) % 5


def generate_data(n, m, func=None, density=None):
    if density is None:
        density = np.abs(np.random.rand(m))
        density = density / np.sum(density)
    x = np.random.choice(m, n, p=density)
    y = np.random.randint(0, m, n) if func is None else func(x)
    return x, y, density


def true_density(density, func, digit=10):
    out = np.zeros(digit)
    for i, p in enumerate(density):
        out[func(i)] += p
    return out


def main():
    np.random.seed(1010)
    out_dir = ensure_dir(OUTPUT_DIR / "figure10")
    tot_n = 3000
    n = np.arange(10, tot_n, 30)
    mc = 100
    h_y = np.zeros((len(n), mc))

    for j in range(mc):
        data_x, data_y, density_x = generate_data(tot_n, 100, mod_5)
        for i, n_i in enumerate(n):
            h_y[i, j] = mle_entropy(data_y[:n_i])

    h_y_std = np.std(h_y, axis=1)
    true_h_y_std = np.log(5) / np.sqrt(n)
    plt.figure(figsize=(8, 6))
    plt.plot(n, h_y_std, label="MI Estimation Std")
    plt.plot(n, true_h_y_std, label="Upper Bound for Estimation Std")
    plt.title("MI Estimation Std", fontsize=16)
    plt.legend(fontsize=16)
    plt.xlabel("n", fontsize=16)
    plt.ylabel("Std", fontsize=16)
    plt.savefig(out_dir / "MI_Estimation_Std.pdf", bbox_inches="tight")
    plt.close()
    print(f"Saved Figure 10 files to {out_dir}")


if __name__ == "__main__":
    timed_main("figure10.py", main)
