# This script plot the result generated by profile_factorization.rs
from re import A
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt

def plot_n_stats():
    table = pd.read_csv("profile_stats.csv")
    table.drop(columns=["n"], inplace=True)
    pollard_cols = list(k for k in table.columns if k.startswith("pollard"))
    squfof_cols = list(k for k in table.columns if k.startswith("squfof"))
    oneline_cols = list(k for k in table.columns if k.startswith("one_line"))

    # MAXITER = 1 << 20
    # table[table >= MAXITER] = np.nan

    mean_table = table.groupby(table['n_bits'] // 4).agg(np.nanmean)
    fig, ax = plt.subplots()
    mean_table.plot("n_bits", pollard_cols, ax=ax)
    mean_table.plot("n_bits", squfof_cols, ax=ax)
    mean_table.plot("n_bits", oneline_cols, ax=ax)
    ax.set_yscale("log")

    min_table = table.groupby(table['n_bits'] // 4).agg(np.nanmin)
    fig, ax = plt.subplots()
    ax.plot(min_table["n_bits"], np.mean(min_table[pollard_cols], axis=1), label="pollard")
    ax.plot(min_table["n_bits"], np.mean(min_table[squfof_cols], axis=1), label="squfof")
    ax.plot(min_table["n_bits"], np.mean(min_table[oneline_cols], axis=1), label="one_line")
    ax.legend()
    ax.set_yscale("log")

def plot_n_min_stats():
    table = pd.read_csv("profile_stats.csv")
    table.drop(columns=["n"], inplace=True)
    for k in table.columns: # caculate average time
        if k.startswith("time_"):
            table[k] = table[k] / table[k[5:]]
            print(table[k])
    min_table = table.groupby(table['n_bits'] // 4).agg(np.nanmean)

    # MAXITER = 1 << 24
    # table[table >= MAXITER] = np.nan

    ax = min_table.plot("n_bits", ["pollard_rho", "squfof", "one_line"])
    ax.set_yscale("log")
    ax.set_ylabel("min iters")
    
    ax = min_table.plot("n_bits", ["time_pollard_rho", "time_squfof", "time_one_line"])
    ax.set_yscale("log")
    ax.set_ylabel("avg time per iter")

if __name__ == "__main__":
    # plot_n_stats()
    plot_n_min_stats()
    plt.show()