from experiments.synthetic_data.analysis.processing_time_analysis import mean_l
from experiments.synthetic_data.analysis.processing_time_analysis import regression_k3
from experiments.synthetic_data.analysis.processing_time_analysis import regression_m
from experiments.synthetic_data.analysis.call_number_analysis import (
    nonlinear_regression,
)
import numpy as np
import pandas as pd
import os


def k3(V):
    return k3_a * np.exp(-k3_b * V) + k3_c


def m(V):
    return m_a * np.log2(V) + m_b


def relax_success(E, S, AVG_DEG):
    relax_attempts = E
    relax_success_ratio = (rs_a * np.log(S) + rs_b) * AVG_DEG**rs_c
    return relax_attempts * relax_success_ratio


def runtime_predict(V, D, S):
    E = V * (V - 1) * D
    AVG_DEG = (V - 1) * D
    if AVG_DEG < 1:
        return False, False

    runtime = V * l + E * k3(V) + relax_success(E, S, AVG_DEG) * m(V)
    print(relax_success(E, S, AVG_DEG))
    # runtime_eq = f"[V * {l}] + [E * {k3_a} * exp(-{k3_b} * V) + {k3_c}] + [E * ({rs_a} * log(S) + {rs_b}) * AVG_DEG ** {rs_c} * {m_a} * log_2(V) + {m_b}]"

    return runtime, relax_success(E, S, AVG_DEG)


l = mean_l()
k3_a, k3_b, k3_c = regression_k3()
m_a, m_b = regression_m()
rs_a, rs_b, rs_c = nonlinear_regression()

csv_file = "results/real_data/derived/dimacs_graph_distribution/distribution.csv"
save_folder = f"results/real_data/derived/{os.path.splitext(os.path.basename(csv_file))[0]}/dimacs_prediction_result"
os.makedirs(save_folder, exist_ok=True)

df = pd.read_csv(csv_file)

records = []

for row in df.itertuples():
    nodes = row.nodes
    density = row.density
    sigma = np.sqrt(np.log((row.std / row.mean) ** 2 + 1))

    pred_time, pred_relax_success = runtime_predict(nodes, density, sigma)

    if pred_time is False:
        continue
    records.append(
        {
            "file": row.file,
            "nodes": nodes,
            "density": density,
            "sigma": sigma,
            "predict_time": pred_time,
            "relax_success": pred_relax_success,
        }
    )
    print(f"finished {row.file}")

result = pd.DataFrame(records)
out_path = os.path.join(save_folder, "prediction.csv")
result.to_csv(out_path, index=False)
print(f"Saved {len(result)} rows → {out_path}")