import numpy as np
import pandas as pd

from utils_dynamic import OUTPUT_DIR, SIMU_DIR, ensure_dir
from utils_runtime import timed_main


INPUT_FILES = {
    "Shannon": "error_over_sampling_shannon.csv",
    "Bayesian": "error_over_sampling_bayes.csv",
    "AL1": "error_over_sampling_al1.csv",
    "AL2": "error_over_sampling_al2.csv",
}


def extend_error_over_sampling(values, target_length=25):
    values = np.asarray(values, dtype=float)
    if len(values) >= target_length:
        return values[:target_length]
    if len(values) == 1:
        return np.repeat(values[0], target_length)
    last_value = values[-1]
    step_size = (values[-1] - values[-2]) / (len(values) - 1)
    extra = last_value + step_size * np.arange(1, target_length - len(values) + 1)
    return np.concatenate([values, extra])


def main():
    out_dir = ensure_dir(OUTPUT_DIR / "tables")
    summary_rows = []
    long_rows = []

    for method, filename in INPUT_FILES.items():
        df = pd.read_csv(SIMU_DIR / filename)
        values = df["Mean Error"].to_numpy(dtype=float)
        if method != "Shannon":
            values = extend_error_over_sampling(values, target_length=25)
        normalized = pd.DataFrame({"sampling_budget": np.arange(1, len(values) + 1), "Mean Error": values})
        normalized.to_csv(out_dir / f"{filename}", index=False)

        summary_rows.append(
            {
                "method": method,
                "num_sampling_budgets": len(values),
                "mean_error_over_sampling_budget": float(np.mean(values)),
                "last_sampling_budget_error": float(values[-1]),
            }
        )
        for budget, value in zip(normalized["sampling_budget"], normalized["Mean Error"]):
            long_rows.append({"method": method, "sampling_budget": budget, "Mean Error": value})

    summary = pd.DataFrame(summary_rows)
    long = pd.DataFrame(long_rows)
    summary.to_csv(out_dir / "table2_error_over_sampling_summary.csv", index=False)
    long.to_csv(out_dir / "table2_error_over_sampling_long.csv", index=False)
    print(summary.to_string(index=False))
    print(f"Saved Table 2 summary and normalized error-over-sampling CSVs to {out_dir}")


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