from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
import sys


def savepdf_tex(fig, name, **kwargs):
    import subprocess, os

    fig.savefig("temp.pdf", format="pdf", **kwargs)
    incmd = [
        "inkscape",
        "temp.pdf",
        "--export-pdf={}.pdf".format(name),
        "--export-latex",
    ]  # "--export-ignore-filters",
    subprocess.check_output(incmd)
    os.remove("temp.pdf")


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
    """
    f = open(filename, "r")
    data = np.loadtxt(f, dtype=float)
    f.close()
    return data


def plot_heatmap(t_values, title, vmax_multiplier, shading):
    # Create a meshgrid of tr and rr values
    TR, RR = np.meshgrid(tr_values, rr_values)

    # Plot the heatmap using pcolormesh
    if shading == True:
        plt.pcolormesh(
            TR,
            RR,
            t_values,
            cmap="seismic",
            shading="gouraud"
            # shading='flat'
        )
    else:
        plt.pcolormesh(
            TR,
            RR,
            t_values,
            cmap="seismic",
        )
    plt.colorbar(label="Mean Path Duration $<t_{TPT}>$")
    plt.ylabel("tumble rate $\lambda_{1,0}$")
    plt.xlabel("run rate $\lambda_{0,1}$")
    plt.title("Heatmap")

    vmax_value = np.nanmax(t_values)
    min_value = np.nanmin(t_values)
    plt.clim(min_value, vmax_value * vmax_multiplier)
    if shading == True:
        title = title + "_shaded"
    plt.savefig(f"{title}.svg", format="svg", bbox_inches="tight", dpi=600)
    plt.savefig(f"{title}.png", format="png", bbox_inches="tight", dpi=600)
    plt.savefig(f"{title}.pdf", format="pdf", bbox_inches="tight", dpi=600)
    plt.show()


def plot_heatmaps(t_values, title, vmax_multiplier):
    tr_values = np.arange(0, 42, 2)
    rr_values = np.arange(0, 42, 2)
    TR, RR = np.meshgrid(tr_values * tau, rr_values * tau)

    Z1 = t_values

    vmax_value = np.nanmax(t_values)
    min_value = np.nanmin(t_values)

    # Create subplots
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))

    # Plot the first heatmap
    im1 = axes[0].pcolormesh(
        TR, RR, Z1, cmap="seismic", vmin=min_value, vmax=vmax_value * vmax_multiplier
    )
    # axes[0].set_title('Heatmap 1')
    axes[0].set_xlabel("run rate $\lambda_{0,1}$")
    axes[0].set_ylabel("tumble rate $\lambda_{1,0}$")
    cbar1 = fig.colorbar(im1, ax=axes[0])

    # Plot the second heatmap
    im2 = axes[1].pcolormesh(
        TR,
        RR,
        Z1,
        cmap="seismic",
        shading="gouraud",
        vmin=min_value,
        vmax=vmax_value * vmax_multiplier,
    )
    # axes[1].set_title('Heatmap 2')
    axes[1].set_xlabel("run rate $\lambda_{0,1}$")
    axes[1].set_ylabel("tumble rate $\lambda_{1,0}$")
    cbar2 = fig.colorbar(im2, ax=axes[1])

    plt.tight_layout()  # Adjust layout to prevent overlap
    plt.show()


dt = 0.005
tr_values = np.arange(0, 42, 2)
rr_values = np.arange(0, 42, 2)
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

# Create an empty array to store t_values
t_values = np.zeros((len(tr_values), len(rr_values)))
t_values2 = np.zeros((len(tr_values), len(rr_values)))
t_rel_err = np.zeros((len(tr_values), len(rr_values)))
for i, tr in enumerate(tr_values):
    for j, rr in enumerate(rr_values):
        if tr == 0 and rr == 0:
            t_values[i, j] = np.nan
            t_rel_err[i, j] = np.nan
        else:
            # Construct the filename dynamically
            filename = f"data/length_distributions/length_distribution_histo_TPS_rtp_step0.005pe10tumblerate{tr}runrate{rr}samplesize100000000.dat"
            # Read data from file
            data = readfile(filename)
            print(np.sum(data), "tr", tr, "rr", rr)
            # print(np.arange(0, len(data)))
            t_values[i, j] = np.sum(
                0.005 / tau * data * np.arange(0, len(data))
            ) / np.sum(data)
            t_values2[i, j] = np.sum(
                (0.005 / tau * data * np.arange(0, len(data))) ** 2
            ) / np.sum(data)

            t_rel_err[i, j] = (
                np.sqrt(np.sum((t_values2[i][j] - t_values[i][j] ** 2)) / np.sum(data))
                / t_values[i, j]
            )
            # print(t_rel_err[i, j])
# Write to the output file
print("min value of mean times")
print(np.nanmin(t_values))
print("max value of mean times")
print(np.nanmax(t_values))
print("value of passive mean times")
for i, tr in enumerate(tr_values):
    for j, rr in enumerate(rr_values):
        if tr != 0 and rr == 0:
            print(t_values[i, j])
output_filename = f"data/meanvalues_tau{tau}.dat"
with open(output_filename, "w") as f:
    f.write(f"1->0, 0->1, <t>, rel_err\n")
    for i, tr in enumerate(tr_values):
        for j, rr in enumerate(rr_values):
            f.write(f"{tr*tau}, {rr*tau}, {t_values[i, j]}, {t_rel_err[i, j]}\n")
# plot_heatmap(t_values, 'data/imgs/meanTPT', 0.5, True)
# plot_heatmap(t_values, 'data/imgs/meanTPT', 0.5, False)

# plot_heatmap(t_rel_err, 'data/imgs/rel_error_TPT' ,1, True)
# plot_heatmap(t_rel_err, 'data/imgs/rel_error_TPT' ,1, False)


plot_heatmaps(t_values, "data/imgs/mean", 0.5)
plot_heatmaps(t_rel_err, "data/imgs/rel_err", 1)