import numpy as np
import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.ticker import MultipleLocator
import os
from matplotlib.lines import Line2D
from matplotlib.ticker import FormatStrFormatter


axislabelfontsize = 10
axisticslabelfontsize = 8
colorbarfontsize = 8


def readfile(filename):
    """
    opens the tsv files with the data in it
    :param filename : string, name of the file
    :return data    : array, data of the file
    """
    with open(filename, "r") as f:
        data = np.loadtxt(f, dtype=float)
    return data


D = 0.1
mu = 0.1
kx = 6
x0 = 1
Fmax = 8 * kx * x0**3 / np.sqrt(27)
L = D / (mu * Fmax)
tau = L * L / D


# TODO
# Define the base directory
base_directory = "./data/length_distributions/"
# Define the common parts of the file names
common_file_name_md = "length_distribution_histo_md_rtp_step0.005pe10tumblerate{}runrate{}samplesize1000000.dat"
common_file_name_TPS = "length_distribution_histo_TPS_rtp_step0.005pe10tumblerate{}runrate{}samplesize10000000.dat"
# common_file_name_TPS = "length_distribution_histo_TPS_rtp_step0.005pe10tumblerate{}runrate{}samplesize100000000.dat"
# Initialize empty lists to store data
data_md = []
data_TPS = []
tr_array = [0.1, 1, 10]
rr_array = [0.1, 1, 10]
## Iterate over tumble rate
for tr in tr_array:
    # Iterate over run rate
    for rr in rr_array:
        # Skip the combination where both tr and rr are 0
        if tr == 0 and rr == 0:
            continue

        file_to_open_md = base_directory + common_file_name_md.format(tr, rr)
        file_to_open_TPS = base_directory + common_file_name_TPS.format(tr, rr)

        # Read data from files using readfile function and assign to variables
        data_md_raw = readfile(file_to_open_md)
        data_TPS_raw = readfile(file_to_open_TPS)
        # print("sum von data =", np.sum(data_TPS_raw))
        # Calculate the sum of all elements in the data_md and data_TPS arrays
        data_md_sum = sum(data_md_raw)
        data_TPS_sum = sum(data_TPS_raw)

        # Normalize the data by dividing each element by the sum
        normalized_data_md = [100 * value / data_md_sum for value in data_md_raw]
        normalized_data_TPS = [100 * value / data_TPS_sum for value in data_TPS_raw]

        # Append the normalized data to the respective lists
        data_md.append(normalized_data_md)
        data_TPS.append(normalized_data_TPS)

# Now, for each tumble rate and run rate combination, we'll plot the histograms
num_plots = len(data_md)  # Add 1 for the 0, 0 combination
num_rows = int(np.ceil(np.sqrt(num_plots)))
num_cols = int(np.ceil(num_plots / num_rows))

plt.rc("text", usetex=True)
plt.rc("font", family="serif", weight="bold", size=axislabelfontsize)

fig, axs = plt.subplots(figsize=(3.5, 2.79))

# TODO
file_to_open = (
    "data/trajectories/tps_trajectories_rtp_step0.005pe10tumblerate10runrate2.dat"
)

positions_at_time = readfile(file_to_open)
x = positions_at_time[:, 1]
y = positions_at_time[:, 2]
theta = positions_at_time[:, 3]
nu = positions_at_time[:, 4]
xnew = positions_at_time[:, 6]
ynew = positions_at_time[:, 7]
thetanew = positions_at_time[:, 8]
nunew = positions_at_time[:, 9]
nu = nu[x != 0]
theta = theta[x != 0]
x = x[x != 0]
y = y[y != 0]
nunew = nunew[xnew != 0]
thetanew = thetanew[xnew != 0]
xnew = xnew[xnew != 0]
ynew = ynew[ynew != 0]
xtumble = x[nu != 1]
ytumble = y[nu != 1]
xrun = x[nu == 1]
yrun = y[nu == 1]
xrunnew = xnew[nunew == 1]
yrunnew = ynew[nunew == 1]
xtumblenew = xnew[nunew == 0]
ytumblenew = ynew[nunew == 0]

# print potential landscape
a = np.linspace(-3.5, 3.5, 1000)
b = np.linspace(-3.5, 3.5, 1000)
X, Y = np.meshgrid(a, b)
Z = (6 * ((X**2 - 1) ** 2)) + (10 * Y**2)
# Z[Z>20.1] = np.NaN
CS = axs.contourf(X, Y, Z, levels=np.linspace(0, 38, 55), cmap="viridis")
CSS = axs.contour(X, Y, Z, [2, 4, 6, 8, 14], colors="white", linewidths=0.8)
CSSS = axs.contour(X, Y, Z, [2], colors="lime", linewidths=0.8)
axs.plot(
    x[x != 0],
    y[y != 0],
    linestyle="-",
    linewidth=1.5,
    marker=" ",
    label="passive phase",
    color="cyan",
)
axs.plot(
    xrun,
    yrun,
    linestyle="-",
    marker=" ",
    label="active phase",
    markersize=6,
    color="red",
)
r = axs.text(0.15, 0.46, "R", fontsize=15, transform=axs.transAxes, color="white")
t = axs.text(0.77, 0.46, "T", fontsize=15, transform=axs.transAxes, color="white")
# t.set_bbox(dict(facecolor='white', alpha=0.9, edgecolor='none',pad=1.5))

axs.set_xlabel(r"${x}/{x_0}$")
axs.set_ylabel(r"${y}/{x_0}$", labelpad=10, rotation=0)
axs.set_xlim(-1.6, 1.6)
axs.set_ylim(-1.6, 1.6)
ticks = [-1, 0, 1]
axs.set_xticks(ticks)
axs.set_yticks(ticks)
# Example if you have a single axis
# axs.set_aspect('equal', adjustable='box')  # Use 'box' instead of 'datalim'

axs.tick_params(top=True, right=True, direction="in")
clb = plt.colorbar(
    CS, ax=axs, shrink=0.8, ticks=[0, 4, 8, 12, 16, 20, 24, 28, 32, 36], pad=0.2  # 0.18
)
clb.ax.tick_params(labelsize=colorbarfontsize)
clb.ax.set_title("$U/k_{\mathrm{B}}T_{\mathrm{eff}}$", fontsize=axislabelfontsize)
axs.legend()
axs.tick_params(axis="both", which="major", labelsize=axisticslabelfontsize)

plt.tight_layout()  # Ensures tight layout of the plot


top_axis = axs.twiny()


top_axis.set_xlabel(r"$x / L$", fontsize=axislabelfontsize, labelpad=5)
top_axis.set_xlim(-2, 2)
top_axis.xaxis.set_ticks([1 * (-1.082532), 0, 1 * 1.082532])
top_axis.xaxis.set_ticklabels(["-10", "0", "10"], fontsize=axisticslabelfontsize)
for tick in top_axis.xaxis.get_major_ticks():
    tick.label.set_fontsize(axisticslabelfontsize)

right_axis = axs.twinx()
right_axis.set_ylabel(r"$y / L$", fontsize=axislabelfontsize, labelpad=5, rotation=0)
right_axis.set_ylim(-2, 2)
right_axis.yaxis.set_ticks([1 * (-1.082532), 0, 1 * 1.082532])
right_axis.yaxis.set_ticklabels(["-10", "0", "10"], fontsize=axisticslabelfontsize)
for tick in right_axis.yaxis.get_major_ticks():
    tick.label.set_fontsize(axisticslabelfontsize)

# axs.yaxis.set_label_coords(-0.05, 0.5)  # Adjust the coordinates

right_axis.yaxis.set_label_coords(1.21, 0.566)  # Same coordinates to align
# Adjust the layout
plt.subplots_adjust(
    left=0.1, right=0.9, bottom=0.1, top=0.9
)  # Ensure margins are balanced

plt.tight_layout()

plt.savefig(
    "histo_and_trajectory_in_landscape_redgreen_onlytumble.pdf",
    format="pdf",
    compress="false",
    bbox_inches="tight",
    dpi=600,
)
# plt.show()

target_dir = "data/imgs"

# Ensure target directory exists
os.makedirs(target_dir, exist_ok=True)
# Full path to the PDF file
filename = f"U(r).pdf"
pdf_path = os.path.join(target_dir, filename)

# Save the figure to a PDF
with PdfPages(pdf_path) as pdf:
    pdf.savefig(fig)
    plt.close(fig)  # Close the figure after saving

print(f"Figure saved to {pdf_path}")

# todo todo todo


######################################################################################################################################################################################


fig, ax = plt.subplots(figsize=(3.5, 3.5))

# TODO
# Define the file paths
file_to_open_md = "data/length_distributions/length_distribution_histo_md_rtp_step0.005pe10tumblerate2runrate2samplesize1000000.dat"
file_to_open_TPS = "data/length_distributions/length_distribution_histo_TPS_rtp_step0.005pe10tumblerate2runrate2samplesize100000000.dat"

# Read data from files
data_md_raw = readfile(file_to_open_md)
data_TPS_raw = readfile(file_to_open_TPS)

# Calculate the sum of all elements in the data arrays
data_md_sum = sum(data_md_raw[:-1])
data_TPS_sum = sum(data_TPS_raw[:-1])

# Normalize the data by dividing each element by the sum
normalized_data_md = [value / data_md_sum for value in data_md_raw[:-1]]
normalized_data_TPS = [value / data_TPS_sum for value in data_TPS_raw[:-1]]

# Chunk the data and calculate average frequency
num_chunks = 300

chunk_size = len(normalized_data_md) / num_chunks
chunked_data_md = []
chunked_data_TPS = []

x_values = []
for j in range(num_chunks):
    start_index = j * chunk_size
    end_index = (j + 1) * chunk_size
    start_index = int(start_index)
    end_index = int(end_index)
    chunk_md = normalized_data_md[start_index:end_index]
    chunk_TPS = normalized_data_TPS[start_index:end_index]

    chunk_avg_md = np.mean(chunk_md)
    chunk_avg_TPS = np.mean(chunk_TPS)

    chunked_data_md.append(chunk_avg_md * chunk_size)
    chunked_data_TPS.append(chunk_avg_TPS * chunk_size)
    x_values.append((start_index + end_index) / 2 * 0.005 / tau)  # Adjust x-values

chunked_data_md = chunked_data_md / (x_values[3] - x_values[2])
chunked_data_TPS = chunked_data_TPS / (x_values[3] - x_values[2])
# Automatically compute the bar width based on the spacing between x-values
bar_width = np.min(np.diff(x_values))
print(np.sum(chunked_data_md) * (x_values[3] - x_values[2]))
# Plotting the bars
ax.bar(
    x_values,
    chunked_data_md,
    bar_width,
    label="MD",
    edgecolor="black",
    fill=False,
    linewidth=2,
)
ax.bar(x_values, chunked_data_TPS, bar_width, label="TPS", color="black", alpha=0.5)
ax.set_xlim([0, 50])
ax.set_ylim([0, 0.14])

# Set the axis labels
ax.set_xlabel(r"$t/\tau$", fontsize=10)
ax.set_ylabel(r"P(t_\text{tpt})", fontsize=10)

# Set the legend labels
ax.legend()

# TODO
# File paths
file_to_open_md_2 = "data/length_distributions/length_distribution_histo_md_rtp_step0.005pe10tumblerate0runrate2samplesize1000000.dat"
file_to_open_md_3 = "data/length_distributions/length_distribution_histo_md_rtp_step0.005pe10tumblerate2runrate0samplesize1000000.dat"

# Read data from files
data_md_2_raw = readfile(file_to_open_md_2)[:-1]  # Exclude the last line
data_md_3_raw = readfile(file_to_open_md_3)[:-1]  # Exclude the last line

# Calculate the sum of all elements in the data arrays
data_md_2_sum = sum(data_md_2_raw)
data_md_3_sum = sum(data_md_3_raw)

# Normalize the data by dividing each element by the sum
normalized_data_md_2 = [value / data_md_2_sum for value in data_md_2_raw]
normalized_data_md_3 = [value / data_md_3_sum for value in data_md_3_raw]

# Chunk the data and calculate average frequency
num_chunks = 300

chunk_size = len(normalized_data_md) / num_chunks  # Integer division to get chunk size
chunked_data_md_2 = []
chunked_data_md_3 = []
x_values = []

for j in range(num_chunks):
    start_index = j * chunk_size
    end_index = (j + 1) * chunk_size
    start_index = int(start_index)
    end_index = int(end_index)

    chunk_md_2 = normalized_data_md_2[start_index:end_index]
    chunk_md_3 = normalized_data_md_3[start_index:end_index]

    chunk_avg_md_2 = np.mean(chunk_md_2)
    chunk_avg_md_3 = np.mean(chunk_md_3)

    chunked_data_md_2.append(chunk_avg_md_2 * chunk_size)
    chunked_data_md_3.append(chunk_avg_md_3 * chunk_size)

    x_values.append((start_index + end_index) / 2 * 0.005 / tau)  # Adjust x-values
chunked_data_md_2 /= x_values[3] - x_values[2]
chunked_data_md_3 /= x_values[3] - x_values[2]
print(np.sum(chunked_data_md_2) * (x_values[3] - x_values[2]))
print(np.sum(chunked_data_md_3) * (x_values[3] - x_values[2]))
# Automatically compute the bar width based on the spacing between x-values
bar_width = np.min(np.diff(x_values))

# Plotting the bars
fig, ax = plt.subplots(figsize=(3.5, 3.8))
fig.subplots_adjust(left=0.15)  # Adjust as needed

plt.rc("text", usetex=True)
plt.rc("font", family="serif", weight="bold", size=axislabelfontsize)
# Plot MD_2 as a dashed green line

ax.bar(
    x_values,
    chunked_data_md,
    bar_width,
    label="RTP (MD)",
    edgecolor="black",
    fill=False,
    linewidth=2,
)
ax.bar(
    x_values, chunked_data_TPS, bar_width, label="RTP (TPS)", color="black", alpha=0.5
)

ax.plot(x_values, chunked_data_md_2, linestyle="--", color="red", label="ABP (MD)")

# Plot MD_3 as a dashed purple line
ax.plot(x_values, chunked_data_md_3, linestyle=":", color="blue", label="BP (MD)")

ax.set_xlim([0, 50])
ax.set_ylim([0, 0.14])

# Set the tick labels font size
# ax.tick_params(axis='both', which='major', labelsize=10)  # Adjust label size as needed

# Set the axis labels
# Set global settings for LaTeX usage and serif font family
plt.rc("text", usetex=True)
plt.rc("font", family="serif", weight="normal", size=10)  # Adjust size as needed

ax.set_xlabel(rf"$t/\tau$", fontsize=10, fontweight="normal")
ax.set_ylabel(rf"$P(t_{{TPT}})$", fontsize=10, fontweight="normal")

ax.tick_params(axis="both", which="major", labelsize=8)
# ax.set_yticks()
ax.set_xticks([0, 25, 50])


# Set the legend with explicit order
handles, labels = ax.get_legend_handles_labels()
order = [2, 3, 1, 0]  # Adjust this order based on your preference
ax.legend([handles[idx] for idx in order], [labels[idx] for idx in order])
# Show the plot
# plt.show()

target_dir = "data/imgs"

# Ensure target directory exists
os.makedirs(target_dir, exist_ok=True)
# Full path to the PDF file
filename = f"P(tTPT).pdf"
pdf_path = os.path.join(target_dir, filename)

# Save the figure to a PDF
with PdfPages(pdf_path) as pdf:
    pdf.savefig(fig)
    plt.close(fig)  # Close the figure after saving

print(f"Figure saved to {pdf_path}")