import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

def fit_op_cost_functional(df, algo):
    sub = df[(df["algorithm"] == algo) & (df["reached"] == True)].copy()

    sub["logN"] = np.log(sub["nodes"])
    sub = sub.rename(columns={"relax_success":"decrease_key_calls",
                              "nodes":"add_calls"})

    if algo == "binary":
        sub["f_extract"] = sub["extract_min_calls"] * sub["logN"]
        sub["f_decrease"] = sub["decrease_key_calls"] * sub["logN"]
    else:  # fibonacci
        sub["f_extract"] = sub["extract_min_calls"] * sub["logN"]
        sub["f_decrease"] = sub["decrease_key_calls"]

    X = sub[["add_calls", "f_extract", "relax_attempts", "f_decrease"]].to_numpy()
    y = sub["time"].to_numpy()

    model = LinearRegression()
    model.fit(X, y)

    y_pred = model.predict(X)
    r2 = r2_score(y, y_pred)

    return {
        "algo": algo,
        "intercept": model.intercept_,
        "coef_add": model.coef_[0],
        "coef_extract": model.coef_[1],
        "coef_relax": model.coef_[2],
        "coef_decrease": model.coef_[3],
        "r2": r2,
        "n_samples": len(sub)
    }

def op_cost_functional_analysis(df):
    return pd.DataFrame([
        fit_op_cost_functional(df, "binary"),
        fit_op_cost_functional(df, "fibonacci")
    ])

df = pd.read_csv("results/synthetic_data/raw/20260303_083239.csv")
summary = op_cost_functional_analysis(df)
print(summary)