import json
import numpy as np
import matplotlib.pyplot as plt
import csv
import re
import readBC
import glob


file_path = "TZ_data.json"

with open(file_path, "r") as f:
    data = json.load(f)

T_values = {}

for func in data["funcs"]:
    for point in func["pts"]:
        T = point["x"]
        Z = point["y"]
        if T not in T_values:
            T_values[T] = []
        T_values[T].append(Z)


T_sorted = sorted(T_values.keys())
Z_avg = [np.mean(T_values[T]) for T in T_sorted]


# Save to CSV
csv_file_path = "../average_TtoZ_curve.csv"
with open(csv_file_path, "w", newline="") as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["T", "Z_avg"])
    writer.writerows(zip(T_sorted, Z_avg))

# Plotting feature
# plt.figure(figsize=(11, 8))

# for func in data["funcs"]:
#     fName = func["fName"]
#     print(fName)
#     match = re.match(r"Ne\s*=\s*([\d\.]+)E([+-]?\d+)", fName)
#     if match:
#         Ne, X = match.groups()
#         formatted_label = rf"$n_e = {float(Ne):.1f} \times 10^{{{int(X)}}}$"
#     else:
#         formatted_label = fName  # Fallback in case of no match
#     x = np.array([point["x"] for point in func["pts"]])
#     y = np.array([point["y"] for point in func["pts"]])
    # plt.plot(x, y, linestyle="--", alpha=0.6, label=formatted_label)

# Plot the average curve
# plt.errorbar(x_sorted, y_avg, y_std, capsize=3, color="black", linewidth=2, label="Average")
# Basko = 22.5 * np.power(np.divide(x_sorted, 100.0), 0.6)
# plt.plot(x_sorted, Basko, color="blue", linewidth=2, label="Basko's power law")

# plt.xscale("log")  
# plt.xlim([0.5, 100])
# plt.ylim([0, 30])
# plt.xlabel(f"Temperature ($\mathrm{{eV}}$)")
# plt.ylabel(f"$\overline{{Z}}$")
# plt.title(f"{data['title']}")
# plt.legend()
# plt.grid(True)
# plt.savefig("TtoZ.pdf")
# plt.show()

# Collect all unique x values
x_values = sorted(set(point["x"] for func in data["funcs"] for point in func["pts"]))

# Initialize a dictionary to store y values for each Ne
y_values_dict = {}

# Process each function's data
ne_list = []
for func in data["funcs"]:
    # Extract only the numerical part of Ne
    ne_value = re.search(r"[\d\.E\+\-]+", func["fName"]).group(0)
    ne_value = float(ne_value) * 1e6  # using conversion to m^-3
    ne_list.append(float(ne_value))
    y_values_dict[ne_value] = {x: None for x in x_values}
    
    for point in func["pts"]:
        y_values_dict[ne_value][point["x"]] = point["y"]


header = ["x"] + list(y_values_dict.keys())
csv_data = []

for x in x_values:
    row = [x] + [y_values_dict[ne].get(x, "") for ne in y_values_dict]
    csv_data.append(row)

# Save to CSV
csv_file_path = "../Zave_values_per_Ne.csv"
with open(csv_file_path, "w", newline="") as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(header)  # Write heade
    writer.writerows(csv_data)  # Write data rows





# Plotting contour of T to Z per ne
# import pandas as pd

# # Define file path
# file_path = "Zave_values_per_Ne.csv"

# # Read the file and inspect its contents
# df = pd.read_csv(file_path)

# # Display the first few rows
# df.head()
# # Extract x values and Ne values (column headers)
# x_values = df.iloc[:10, 0].values  # First column
# Ne_values = np.array([float(col) for col in df.columns[1:]])  # Convert column names to float
# # Ne_values = np.array([col for col in Ne_values])

# # Extract y values as a 2D array
# y_values = df.iloc[:10, 1:].values

# # Create meshgrid for contour plot
# X, Y = np.meshgrid(Ne_values, x_values)

# # Plot contour
# plt.figure(figsize=(8, 6))
# cp = plt.contourf(X, Y, y_values, levels=20, cmap='viridis')
# plt.colorbar(cp, label="Zave values")

# # fileBC = glob.glob('../2025-03-10_11.23.25/bc.csv')
# # time, n, u, T, TtoZ, Zinj = readBC.read(fileBC[0])
# # plt.plot(n, T, color="black", linestyle="dashed", marker="o", markersize=4, label="Boundary Condition")
# # plt.scatter(n[1], T[1], color="white", edgecolors="black", s=80, label="t = 0")  # White dot with black edge
# # plt.scatter(n[-1], T[-1], color="cyan", edgecolors="black", s=80, label="t = end") 

# # Labels and title
# # plt.yscale("log")
# plt.legend()
# plt.xscale("log")  # Since Ne spans many orders of magnitude
# plt.xlabel("Electron Density (ne)")
# plt.ylabel("Temperature (eV)")
# plt.title("Zave")
# plt.savefig("Contour_Tin_Zave_BC.pdf")
# plt.show()