import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit
import pandas as pd
import itertools
# Legend for colors
from matplotlib.patches import Patch

# ============================================================
# PART 1: Power Law Fit
# ============================================================
x = np.array([28, 112, 248, 432, 880, 1568, 1840, 740, 336, 1056, 2108, 4660, 2976, 3512])
y = np.array([71.074, 104.651, 116.899, 124.491, 134.543, 149.413, 147.272, 131.517, 111.469, 154.861, 171.535, 227.923, 175.423, 198.650])

if np.all(x == 0):
    raise ValueError("Cannot fit when all x values are zero.")


def power_law(x, a, b):
    return a * np.power(x, b)


p0 = [5.0, 0.4]
popt, pcov = curve_fit(power_law, x, y, p0=p0, maxfev=10000)
a_fit, b_fit = popt
perr = np.sqrt(np.diag(pcov))

print("Fitted parameters:")
print(f"  a = {a_fit:.6f} ± {perr[0]:.6f}")
print(f"  b = {b_fit:.6f} ± {perr[1]:.6f}")

y_fit = power_law(x, a_fit, b_fit)
rmse = np.sqrt(np.mean((y - y_fit) ** 2))
print(f"RMSE = {rmse:.4f} µs")

# Save power law plot
x_smooth = np.linspace(1, x.max(), 500)
y_smooth = power_law(x_smooth, a_fit, b_fit)
fig, ax = plt.subplots(figsize=(8, 4))
plt.scatter(x, y, color="blue", s=50, label="Data")
plt.plot(x_smooth, y_smooth, color="red", linewidth=2,
         label=f"Best fit: y = {a_fit:.4f} · x^{b_fit:.4f}\nRMSE = {rmse:.4f} µs")
plt.xlabel("Computational Gates")
plt.ylabel("Time (µs)")
ax.set_xticks(np.linspace(0, 4500, 10))
plt.grid(True)
plt.legend()
plt.savefig("power_law_fit.png", dpi=800, bbox_inches="tight")
plt.close()
print("\nPlot saved to ./power_law_fit.png")

alpha = a_fit

# ============================================================
# PART 2: Configuration
# ============================================================
MAPE_CSV = "mape_errors.csv"
GATES_CSV = "gates.csv"

GATES_MODEL_COL = "Model"
GATES_VALUE_COL = "gates"

T_max = 4500
V = 4 * 4 * (2 + 1)  # 48

GROUP_A = ["model1", "model2", "model3"]
GROUP_B = ["model4", "model5", "model6", "model7"]

SHORT = {
    "model1": r"$M_{1,1}$", "model2": r"$M_{1,2}$", "model3": r"$M_{1,3}$",
    "model4": r"$M_{2,1}$", "model5": r"$M_{2,2}$", "model6": r"$M_{2,3}$", "model7": r"$M_{2,4}$",
}

df_mape = pd.read_csv(MAPE_CSV)
df_gates = pd.read_csv(GATES_CSV)

mape_lookup = dict(zip(df_mape["Model"], df_mape["MAPE_%"]))
gates_lookup = dict(zip(df_gates[GATES_MODEL_COL], df_gates[GATES_VALUE_COL]))

# ============================================================
# PART 3: Precompute objectives for all combinations
# ============================================================
combinations = []
for a_model, b_model in itertools.product(GROUP_A, GROUP_B):
    mape_a = mape_lookup[a_model]
    mape_b = mape_lookup[b_model]
    gates_a = gates_lookup[a_model]
    gates_b = gates_lookup[b_model]

    exec_obj = ((1 / T_max) * V * alpha * gates_a ** b_fit) + \
               ((1 / T_max) * 1 * alpha * gates_b ** b_fit)
    pred_obj = (0.01 * mape_a + 0.01 * mape_b)

    label = SHORT[a_model] + " + " + SHORT[b_model]

    combinations.append({
        "Plant_Model": a_model,
        "Network_Model": b_model,
        "Label": label,
        "MAPE_A_%": mape_a,
        "MAPE_B_%": mape_b,
        "Gates_A": gates_a,
        "Gates_B": gates_b,
        "exec_obj": exec_obj,
        "pred_obj": pred_obj,
    })

combo_df = pd.DataFrame(combinations)

# ============================================================
# PART 4: Sensitivity Analysis — Sweep P from 0 to 1
# ============================================================
P_values = np.linspace(0, 1, 101)  # 0.00, 0.01, ..., 1.00

# Store scores for each P
sensitivity_results = []
best_per_P = []

for P in P_values:
    Q = 1 - P
    scores = P * combo_df["exec_obj"] + Q * combo_df["pred_obj"]
    combo_df_temp = combo_df.copy()
    combo_df_temp["Score"] = scores

    best_idx = scores.idxmin()
    best_per_P.append({
        "P": round(P, 4),
        "Q": round(Q, 4),
        "Best_Combination": combo_df.loc[best_idx, "Label"],
        "Best_Score": round(scores[best_idx], 6),
        "exec_obj": combo_df.loc[best_idx, "exec_obj"],
        "pred_obj": combo_df.loc[best_idx, "pred_obj"],
    })

    for idx, row in combo_df.iterrows():
        sensitivity_results.append({
            "P": round(P, 4),
            "Q": round(Q, 4),
            "Label": row["Label"],
            "Score": round(P * row["exec_obj"] + Q * row["pred_obj"], 6),
        })

sensitivity_df = pd.DataFrame(sensitivity_results)
best_df = pd.DataFrame(best_per_P)

# Save sensitivity results
best_df.to_csv("sensitivity_best_per_P.csv", index=False)
print("\n✅ Sensitivity results saved to sensitivity_best_per_P.csv")

# --- Sensitivity Plot 1: All combinations score vs P ---
fig, ax = plt.subplots(figsize=(12, 6))

for label in combo_df["Label"].unique():
    subset = sensitivity_df[sensitivity_df["Label"] == label]
    ax.plot(subset["P"], subset["Score"], linewidth=1.5, label=label)

ax.set_xlabel("P Value (Q = 1 − P)", fontsize=12)
ax.set_ylabel("Objective Function Score", fontsize=12)
ax.set_title("Sensitivity Analysis: Score vs P Value", fontsize=12)
ax.grid(True, alpha=0.3)
ax.legend(fontsize=8, loc="upper left", bbox_to_anchor=(1.02, 1), borderaxespad=0)
plt.savefig("sensitivity_all_combinations.png", dpi=600, bbox_inches="tight")
plt.close()
print("✅ Plot saved to sensitivity_all_combinations.png")

# --- Sensitivity Plot 2: Best combination vs P ---
fig, ax1 = plt.subplots(figsize=(8, 4))

# Color-code by best combination
unique_best = best_df["Best_Combination"].unique()
color_map = {name: plt.cm.tab10(i) for i, name in enumerate(unique_best)}
colors = [color_map[name] for name in best_df["Best_Combination"]]

ax1.scatter(best_df["P"], best_df["Best_Score"], c=colors, s=30, zorder=5)
ax1.set_xlabel("P Value (Q = 1 − P)", fontsize=12)
ax1.set_ylabel("Best Objective Function Score", fontsize=12)
ax1.set_title("Sensitivity Analysis: Best Solution vs P Value (colored by winner)", fontsize=12)
ax1.set_xticks(P_values[::10])  # Show every 10th tick for clarity
ax1.grid(True, alpha=0.3)

legend_elements = [Patch(facecolor=color_map[name], label=name) for name in unique_best]
# ax1.legend(handles=legend_elements, fontsize=13, loc="upper left", bbox_to_anchor=(0.8, 0.5), borderaxespad=0)
fig.legend(handles=legend_elements, fontsize=13, loc="outside right center", bbox_to_anchor=(1.12, 0.5), borderaxespad=0)

plt.savefig("sensitivity_best_vs_P.png", dpi=600, bbox_inches="tight")
plt.close()
print("✅ Plot saved to sensitivity_best_vs_P.png")

# --- Sensitivity Plot 3: Transition table ---
transitions = best_df[best_df["Best_Combination"] != best_df["Best_Combination"].shift(1)]
print("\n📊 Sensitivity Transitions (where the best combination changes):")
print(transitions[["P", "Q", "Best_Combination", "Best_Score"]].to_string(index=False))

# ============================================================
# PART 5: Pareto Front Analysis
# ============================================================
exec_objs = combo_df["exec_obj"].values
pred_objs = combo_df["pred_obj"].values
labels = combo_df["Label"].values


# Find Pareto-optimal solutions (non-dominated)
def find_pareto_front(exec_vals, pred_vals):
    """Find indices of Pareto-optimal points (minimizing both objectives)."""
    n = len(exec_vals)
    is_pareto = np.ones(n, dtype=bool)
    for i in range(n):
        for j in range(n):
            if i == j:
                continue
            # j dominates i if j is <= in both objectives and < in at least one
            if (exec_vals[j] <= exec_vals[i] and pred_vals[j] <= pred_vals[i]) and \
               (exec_vals[j] < exec_vals[i] or pred_vals[j] < pred_vals[i]):
                is_pareto[i] = False
                break
    return is_pareto


pareto_mask = find_pareto_front(exec_objs, pred_objs)

pareto_df = combo_df[pareto_mask].copy()
pareto_df = pareto_df.sort_values("exec_obj").reset_index(drop=True)

print("\n📊 Pareto-Optimal Solutions:")
print(pareto_df[["Label", "exec_obj", "pred_obj"]].to_string(index=False))
pareto_df.to_csv("pareto_front.csv", index=False)
print("✅ Pareto front saved to pareto_front.csv")

# --- Pareto Front Plot ---
fig, ax = plt.subplots(figsize=(8, 4), layout='constrained')

# Plot all combinations
ax.scatter(exec_objs[~pareto_mask], pred_objs[~pareto_mask],
           color="lightgray", s=80, edgecolors="gray", zorder=3, label="Dominated")

# Plot Pareto front points
ax.scatter(exec_objs[pareto_mask], pred_objs[pareto_mask],
           color="red", s=120, edgecolors="black", zorder=5, label="Pareto Front")

# Connect Pareto front with a line (sorted by exec_obj)
pareto_exec_sorted = pareto_df["exec_obj"].values
pareto_pred_sorted = pareto_df["pred_obj"].values
ax.plot(pareto_exec_sorted, pareto_pred_sorted, 'r--', linewidth=1.5, zorder=4)

# Annotate all points
for i in range(len(combo_df)):
    offset = (5, 5) if pareto_mask[i] else (5, -10)
    fontweight = "bold" if pareto_mask[i] else "normal"
    ax.annotate(labels[i], (exec_objs[i], pred_objs[i]),
                textcoords="offset points", xytext=offset,
                fontsize=13, fontweight=fontweight)

ax.set_xlabel("Cost of Computation", fontsize=12)
ax.set_ylabel("Cost of Prediction Error ", fontsize=12)
ax.set_title("Pareto Front: Cost of Computation vs Cost of Prediction Error", fontsize=14)
ax.set_ylim(-0.05, 1.8)
ax.grid(True, alpha=0.3)
# ax.legend(fontsize=11, bbox_to_anchor=(0.25, 0.65))
# ax.legend(fontsize=11, loc="outside right center")
fig.legend(fontsize=12, loc='outside right center')

plt.savefig("pareto_front.png", dpi=600, bbox_inches="tight")
plt.close()
print("✅ Plot saved to pareto_front.png")

# ============================================================
# PART 6: Combined Scores Plot (for a specific P, Q)
# ============================================================
P_final = 0.5
Q_final = 1 - P_final

combo_df["Score"] = P_final * combo_df["exec_obj"] + Q_final * combo_df["pred_obj"]
results_df = combo_df.sort_values("Score").reset_index(drop=True)
results_df = results_df.sample(frac=1, random_state=48).reset_index(drop=True)
results_df["ID"] = range(1, len(results_df) + 1)

results_df.to_csv("combined_scores.csv", index=False)

fig, ax = plt.subplots(figsize=(8, 2.5))
ax.plot(results_df["ID"], results_df["Score"], 'o', markersize=8,
        color='steelblue', markeredgecolor='black', zorder=5)

best = results_df.loc[results_df["Score"].idxmin()]
ax.plot(best["ID"], best["Score"], 'o', markersize=8,
        color='red', markeredgecolor='black', zorder=6)

ax.set_xlabel("Combination #", fontsize=12)
ax.set_ylabel("Objective Function Score", fontsize=12)
ax.set_title(f"Objective Function Scores (P={P_final}, Q={Q_final})", fontsize=12)
ax.set_xticks(results_df["ID"])
ax.grid(True, alpha=0.3)

legend_text = "\n".join(
    [f"{row['ID']:>2}. {row['Label']}" for _, row in results_df.iterrows()]
)
fig.subplots_adjust(right=0.70)
fig.text(0.72, 0.5, legend_text, fontsize=13, fontfamily='monospace',
         verticalalignment='center',
         bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.savefig("combined_scores_plot.png", dpi=600, bbox_inches="tight")
plt.close()
print(f"✅ Plot saved to combined_scores_plot.png (P={P_final}, Q={Q_final})")

print("\n🎯 Done! All outputs saved.")
