from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scipy.special import factorial
from scipy.stats import poisson
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.ticker import MultipleLocator
import os
from matplotlib.lines import Line2D
from matplotlib.ticker import FormatStrFormatter

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

axislabelfontsize = 2 * 10
axisticslabelfontsize = 2 * 8
colorbarfontsize = 2 * 8

# Set LaTeX font and styles for the entire figure
plt.rc("text", usetex=True)
plt.rc("font", family="serif", weight="normal")


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


import os
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages


def adjust_and_save_figure(fig, columns, filename):
    # Adjust the figure size
    axislabeltickfontsize = 8
    axislabelfontsize = 10
    fig_width = columns * 3.5  # inches
    fig_height = fig.get_size_inches()[1] * (fig_width / fig.get_size_inches()[0])
    fig.set_size_inches(fig_width, fig_height)

    # Enable LaTeX font
    plt.rc("text", usetex=True)
    plt.rc("font", family="serif", weight="normal", size=10)

    # Adjust tick and label sizes
    for ax in fig.axes:
        ax.tick_params(axis="both", labelsize=axislabeltickfontsize)
        ax.xaxis.label.set_size(axislabelfontsize)
        ax.yaxis.label.set_size(axislabelfontsize)
        ax.title.set_size(axislabelfontsize)

        # If there's a colorbar, adjust its ticks and label
        if hasattr(ax, "collections") and ax.collections:
            for coll in ax.collections:
                if hasattr(coll, "colorbar") and coll.colorbar:
                    coll.colorbar.ax.tick_params(labelsize=axislabeltickfontsize)
                    coll.colorbar.ax.yaxis.label.set_size(axislabelfontsize)

    # Ensure target directory exists
    target_dir = ""
    os.makedirs(target_dir, exist_ok=True)

    # Full path to the PDF file
    filename
    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}")
    # Show the plot
    plt.show()


# Example usage (assuming `fig` is your existing figure):
# fig = plt.figure()  # Your existing figure here
# adjust_and_save_figure(fig)

# Define the base directory
base_directory = "data/length_distributions/"

# Define the common parts of the file names
common_file_name_RTP = "length_distribution_histo_md_rtp_step{}pe10tumblerate1runrate999samplesize100000.dat"
common_file_name_RTPE = (
    "length_distribution_histo_md_rtpE_step{}pe10tumblerate1samplesize100000.dat"
)


# Define the list of dt values
dts = [0.1, 0.001, 0.00001]

# Function to read data from file
def readfile(filename):
    return np.loadtxt(filename)


# Create a figure with 6 subplots
fig, axes = plt.subplots(1, 3, figsize=(14, 7), constrained_layout=True)
axes = axes.flatten()
plt.rc("text", usetex=True)
plt.rc("font", family="serif", weight="bold", size=2 * 8)

# Iterate over dt values and plot data
for i, dt in enumerate(dts[:3]):
    histosize = int(10 / dt)
    file_to_open_RTP = base_directory + common_file_name_RTP.format(dt)
    file_to_open_RTPE = base_directory + common_file_name_RTPE.format(dt)

    # Read data from files
    data_RTP_raw = readfile(file_to_open_RTP)[:histosize]
    data_RTPE_raw = readfile(file_to_open_RTPE)[:histosize]
    data_RTP_sum = sum(data_RTP_raw)
    data_RTPE_sum = sum(data_RTPE_raw)

    # Normalize the data by dividing each element by the sum
    normalized_data_RTP = [value / data_RTP_sum for value in data_RTP_raw]
    normalized_data_RTPE = [value / data_RTPE_sum for value in data_RTPE_raw]

    # Chunk the data and calculate average frequency
    num_chunks = 50
    chunk_size = len(normalized_data_RTP) // num_chunks

    chunked_data_RTP = []
    chunked_data_RTPE = []
    x_values = []

    for j in range(num_chunks):
        start_index = int(j * chunk_size)
        end_index = int((j + 1) * chunk_size)
        chunk_RTP = normalized_data_RTP[start_index:end_index]
        chunk_RTPE = normalized_data_RTPE[start_index:end_index]

        chunk_avg_RTP = np.mean(chunk_RTP)
        chunk_avg_RTPE = np.mean(chunk_RTPE)

        chunked_data_RTP.append(chunk_avg_RTP * chunk_size)
        chunked_data_RTPE.append(chunk_avg_RTPE * chunk_size)
        x_values.append((start_index + end_index) / 2 * dt / 0.12)  # Adjust x-values

    # Normalize chunked data by the interval size
    interval_size = x_values[2] - x_values[1]
    for j in range(num_chunks):
        chunked_data_RTP[j] /= interval_size
        chunked_data_RTPE[j] /= interval_size

    # Automatically compute the bar width based on the spacing between x-values
    bar_width = np.min(np.diff(x_values))

    print("normalized?")
    print(np.sum(chunked_data_RTP) * interval_size)

    # Plot data in the subplot
    ax = axes[i]
    power = int(np.log10(dt))
    formatted_dt = f"10^{{{power}}}"
    t = ax.text(
        0.3,
        1.03,
        rf"$\Delta t = {8.5} \cdot {formatted_dt} / \tau$",
        fontsize=axislabelfontsize,
        transform=ax.transAxes,
    )
    t.set_bbox(dict(facecolor="white", alpha=0.9, edgecolor="black", pad=1.5))
    print("anzahl der x values")
    print(len(x_values))

    ax.bar(
        x_values,
        chunked_data_RTP,
        width=bar_width,
        label=r"RTP $\lambda_{0\rightarrow 1}\rightarrow \infty$",
        edgecolor="red",
        fill=False,
        linewidth=1,
    )
    ax.bar(
        x_values,
        chunked_data_RTPE,
        width=bar_width,
        label=r"RTP instant tumbling",
        color="blue",
        alpha=0.6,
    )
    if i == 0:
        ax.legend(
            loc="upper right", fontsize=axislabelfontsize  # , bbox_to_anchor=(0.8, 0.9)
        )
    ax.set_xlim([0, 50])
    ax.set_xlabel(rf"$t/\tau$", fontsize=axislabelfontsize)
    ax.set_yticks([])
    if i == 0:
        ax.set_ylabel("P(t_\text{{tpt}})", fontsize=axislabelfontsize)
        ax.set_yticks([0, 0.02, 0.04, 0.06, 0.08, 0.1])

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

    plt.rc("text", usetex=True)
    plt.rc("font", family="serif", weight="normal", size=10)  # Adjust size as needed

    ax.set_xlabel(r"$t/\tau$", fontsize=axislabelfontsize)
    ax.set_ylabel(r"$P(t_\mathrm{{TPT}})$", fontsize=axislabelfontsize)

    ax.tick_params(axis="both", which="major", labelsize=8)
    # ax.set_yticks()
    ax.set_xticks([0, 25, 50])
    if i != 0:
        ax.set_ylabel(rf"", fontsize=2 * 3)
    # ax.legend()

    for ax in fig.axes:
        ax.tick_params(axis="both", labelsize=axisticslabelfontsize)
        ax.xaxis.label.set_size(axislabelfontsize)
        ax.yaxis.label.set_size(axislabelfontsize)
        ax.title.set_size(axislabelfontsize)


target_dir = "newfolder/"
# Adjust subplot layout
# fig.subplots_adjust(wspace=0.1, hspace=0.0)  # Fine-tune as needed
plt.tight_layout(pad=1.0)  # Adjust pad as needed
# Ensure target directory exists
os.makedirs(target_dir, exist_ok=True)
# Full path to the PDF file
filename = f"P(tTPT)_RTP-RTPE.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}")
# Show the plot
plt.show()