import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.ticker import MultipleLocator
import numpy as np
import os
from matplotlib.lines import Line2D
from matplotlib.ticker import FormatStrFormatter


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


def readfile_(file_path):
    # Placeholder for reading file function, replace with actual implementation
    # Assumes the file format is compatible with numpy's genfromtxt or loadtxt.
    try:
        return np.genfromtxt(file_path)
    except Exception as e:
        print(f"Error reading file {file_path}: {e}")
        return None


def plotfig5(Pe, rr1, tr1, rr2, tr2, rr3, tr3, dt):
    plt.rc("text", usetex=True)
    plt.rc("font", family="serif", weight="bold", size=axislabelfontsize)

    def prepare_data(rr, tr):
        source_dir = "./data"
        numberoffiles = 1900
        counter = 0
        data_reactive_current = None
        if rr == tr == 1:
            for i in range(numberoffiles):
                readfilename = f"ReactiveCurrent_RTP_Pe{Pe}_RR{rr}_TR{tr}_timestep{dt:.3f}_seed{i}_newversion.dat"
                file_to_open = os.path.join(source_dir, readfilename)
                if os.path.exists(file_to_open):
                    data = readfile(file_to_open)
                    if data is not None:
                        if data_reactive_current is None:
                            data_reactive_current = data
                        else:
                            data_reactive_current += data
                        counter += 1

            if data_reactive_current is not None:
                data_reactive_current /= counter
                print(f"Number of files for RR={rr}, TR={tr} = {counter}")
            else:
                print(f"No data found for RR={rr}, TR={tr}.")
        else:
            readfilename = f"ReactiveCurrent_RTP_Pe{Pe}_RR{rr}_TR{tr}_timestep{dt:.3f}_newversion.dat"
            file_to_open = os.path.join(source_dir, readfilename)
            if os.path.exists(file_to_open):
                data_reactive_current = readfile(file_to_open)
        return data_reactive_current

    # Prepare data for each combination
    data1 = prepare_data(rr1, tr1)
    data2 = prepare_data(rr2, tr2)
    data3 = prepare_data(rr3, tr3)

    if data1 is None or data2 is None or data3 is None:
        print("One or more datasets are missing. Exiting the plotting function.")
        return

    def plot_panel(panel, data, label):
        x = data[:, 0]
        y = data[:, 1]
        z = data[:, 5]
        if np.sum(z) < 0:
            z = -z

        x_direct = data[:, 2]
        y_direct = data[:, 3]

        x_unique = np.unique(x)
        y_unique = np.unique(y)
        N1 = len(x_unique)
        N2 = len(y_unique)

        z_grid = np.zeros((N1, N2))
        x_direct_grid = np.zeros((N1, N2))
        y_direct_grid = np.zeros((N1, N2))

        for i in range(len(z)):
            x_idx = np.where(x_unique == x[i])[0][0]
            y_idx = np.where(y_unique == y[i])[0][0]
            z_grid[x_idx, y_idx] = z[i]
            x_direct_grid[x_idx, y_idx] = x_direct[i]
            y_direct_grid[x_idx, y_idx] = y_direct[i]

        pcm = panel.pcolormesh(
            x_unique, y_unique, z_grid.T, shading="gouraud", cmap="jet"
        )
        if panel == panels[2]:
            cbar = plt.colorbar(pcm, ax=panel, location="bottom", pad=0.2, shrink=0.7)
        else:
            cbar = plt.colorbar(pcm, ax=panel, location="bottom", pad=0.2, shrink=0.7)
        cbar.ax.tick_params(labelsize=colorbarfontsize)
        cbar.set_label(
            r"$m(\mathbf{r})$",
            fontsize=axislabelfontsize,
            labelpad=-0,
            rotation=0,
            verticalalignment="bottom",
        )
        cbar.ax.yaxis.set_label_coords(0.5, 1.02)

        # cbar.set_ticks([pcm.get_array().min(), pcm.get_array().max()])
        ticks = [pcm.get_array().min(), pcm.get_array().max()]
        cbar.set_ticks(ticks)
        cbar.set_ticklabels([f"{tick:.4f}" for tick in ticks])

        panel.streamplot(
            x_unique,
            y_unique,
            x_direct_grid,
            y_direct_grid,
            density=1,
            linewidth=z_grid.T * 100,
            color="cyan",
            arrowsize=1,
            arrowstyle="-|>",
            minlength=0.1,
            maxlength=4.0,
        )

        panel.set_xlabel(r"$x/x_0$", fontsize=axislabelfontsize)

        if panel == panels[0]:
            panel.set_ylabel(r"$y/x_0$", fontsize=axislabelfontsize, y=0.5)
        panel.set_xlim(-2, 2)
        panel.xaxis.set_major_locator(MultipleLocator(1))
        panel.set_ylim(-2, 2)
        panel.yaxis.set_major_locator(MultipleLocator(1))

        custom_handle = Line2D([0], [0], linestyle="None", label=label)

        # legend = panel.legend(handles=[custom_handle], loc='upper left', fontsize=axisticslabelfontsize, frameon=True)
        # legend.get_frame().set_facecolor('white')
        # legend.get_frame().set_edgecolor('none')

        t = panel.text(0.05, 0.89, label, fontsize=5, transform=panel.transAxes)
        t.set_bbox(dict(facecolor="white", alpha=0.9, edgecolor="none", pad=1.5))

        top_axis = panel.twiny()

        top_axis.set_xlabel(r"$x / L$", fontsize=axislabelfontsize)
        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)

        if panel == panels[2]:
            right_axis = panel.twinx()
            right_axis.set_ylabel(r"$y / L$", fontsize=axislabelfontsize)
            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)
        if panel != panels[0]:
            panel.set_yticks([])
        else:
            panel.set_yticks([-2, 0, 2], fontsize=axisticslabelfontsize)
        panel.set_xticks([-2, 0, 2], fontsize=axisticslabelfontsize)
        for tick in panel.xaxis.get_major_ticks():
            tick.label.set_fontsize(axisticslabelfontsize)
        for tick in panel.yaxis.get_major_ticks():
            tick.label.set_fontsize(axisticslabelfontsize)

    # Set up the figure and subplots
    target_dir = "data/imgs"
    fig, panels = plt.subplots(1, 3, figsize=(7, 3))  # Create a 1x3 grid of subplots
    # Anpassen der Abstände zwischen den Subplots
    plt.subplots_adjust(wspace=0.08, hspace=0.3)
    """fig, panels = plt.subplot_mosaic([[0,1], # alternativ für poster
                                  [2, 2]], figsize=(5, 6))"""

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

    # Plot each panel
    plot_panel(
        panels[0],
        data1,
        rf"A) $\lambda_{{0\rightarrow 1}}$=$\lambda_{{1\rightarrow 0}}$={tr1*0.12}$/ \tau$",
    )

    # panels[0].set_title(f'a) $\lambda_{0\rightarrow 1}$={rr1}, TR={tr1}', fontsize=axisticslabelfontsize)

    plot_panel(
        panels[1],
        data2,
        rf"B) $\lambda_{{0\rightarrow 1}}= 1.2/ \tau $;\quad   $\lambda_{{1\rightarrow 0}}$ = {tr2*0.12}$/ \tau$",
    )
    # panels[1].set_title(f'b) RR={rr2}, TR={tr2}', fontsize=axisticslabelfontsize)

    plot_panel(
        panels[2],
        data3,
        rf"C) $\lambda_{{0\rightarrow 1}}$=$\lambda_{{1\rightarrow 0}}$={tr3*0.12}$/ \tau$",
    )
    # panels[2].set_title(f'c) RR={rr3}, TR={tr3}', fontsize=axisticslabelfontsize)

    # Ensure target directory exists
    #
    filename = rf"ReactiveJ_Pe{Pe}RR{rr2}_TR{tr2}_RR{rr3}_TR{tr3}_RR{rr2}_TR{tr2}_RR{rr3}_TR{tr3}_timestep{dt}.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}")


def make_figure4(x, y, z, x_pos, y_pos, x_direct, y_direct):
    plt.rc("text", usetex=True)
    plt.rc("font", family="serif", weight="bold", size=axislabelfontsize)

    target_dir = "data/imgs"
    fig, panel = plt.subplots(
        figsize=(5, 4)
    )  # Increase the figure size for better visibility

    # Set up the main panel
    plt.rc("text", usetex=True)
    N1 = len(x)
    N2 = len(y)
    cmap = plt.get_cmap("jet")

    pcm = panel.pcolormesh(x, y, z.T, shading="gouraud", cmap=cmap)
    if panel == panel[2]:
        cbar = plt.colorbar(pcm, ax=panel, location="bottom", pad=0.15, shrink=20)
    else:
        cbar = plt.colorbar(pcm, ax=panel, location="bottom", pad=0.15, shrink=20)
    cbar.ax.tick_params(labelsize=colorbarfontsize)

    cbar.set_label(
        r"$m(\mathbf{r})$",
        fontsize=axislabelfontsize,
        labelpad=-30,
        rotation=0,
        verticalalignment="bottom",
    )
    cbar.ax.yaxis.set_label_coords(0.5, 1.02)
    cbar.set_ticks([pcm.get_array().min(), pcm.get_array().max()])

    x_direct = x_direct.reshape(N1, N2)
    y_direct = y_direct.reshape(N1, N2)
    panel.streamplot(
        x,
        y,
        x_direct,
        y_direct,
        density=1,
        linewidth=z.T * 100,
        color="cyan",
        arrowsize=1,
        arrowstyle="-|>",
        minlength=0.1,
        maxlength=4.0,
    )

    # Primary X and Y axis
    panel.set_xlabel(r"$x/x_0$", fontsize=axislabelfontsize)
    panel.set_ylabel(r"$y/x_0$", fontsize=axislabelfontsize, y=0.46)

    # Tick label font size
    for tick in panel.xaxis.get_major_ticks():
        tick.label.set_fontsize(axisticslabelfontsize)
    for tick in panel.yaxis.get_major_ticks():
        tick.label.set_fontsize(axisticslabelfontsize)

    panel.set_xlim(-2, 2)
    panel.xaxis.set_major_locator(MultipleLocator(1))
    panel.set_ylim(-2, 2)
    panel.yaxis.set_major_locator(MultipleLocator(1))

    # Secondary X and Y axis
    top_axis = panel.twiny()
    right_axis = panel.twinx()

    # Setting up the top X axis
    top_axis.set_xlabel(r"$x / L$", fontsize=axislabelfontsize)
    top_axis.set_xlim(-3, 3)
    top_axis.xaxis.set_ticks([2 * (-1.082532), -1.082532, 0, 1.082532, 2 * 1.082532])
    top_axis.xaxis.set_ticklabels(
        ["-20", "-10", "0", "10", "20"], fontsize=axisticslabelfontsize
    )
    for tick in top_axis.xaxis.get_major_ticks():
        tick.label.set_fontsize(axisticslabelfontsize)

    # Setting up the right Y axis
    right_axis.set_ylabel(r"$y / L$", fontsize=axislabelfontsize)
    right_axis.set_ylim(-3, 3)
    right_axis.yaxis.set_ticks([2 * (-1.082532), -1.082532, 0, 1.082532, 2 * 1.082532])
    right_axis.yaxis.set_ticklabels(
        ["-20", "-10", "0", "10", "20"], fontsize=axisticslabelfontsize
    )
    for tick in right_axis.yaxis.get_major_ticks():
        tick.label.set_fontsize(axisticslabelfontsize)
        # Ensuring target directory exists
        os.makedirs(target_dir, exist_ok=True)

    # Full path to the PDF file
    filename = f"ReactiveJ_Pe{Pe}_RR{rr}_TR{tr}_timestep{dt}.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


# Schriftgrößen definieren
axislabelfontsize = 10
axisticslabelfontsize = 8
colorbarfontsize = 8


# Definieren der Funktion


def plotfig4(Pe, rr, tr, dt):

    source_dir = "data/"
    readfilename = (
        f"data/ReactiveCurrent_RTP_Pe{Pe}_RR{rr}_TR{tr}_timestep{dt:.3f}_newversion.dat"
    )
    file_to_open = os.path.join(source_dir, readfilename)
    data_reactive_current = readfile(file_to_open)

    numberoffiles = 1900
    counter = 0
    if rr == 1 and tr == 1:

        for i in range(numberoffiles):
            readfilename = f"data/ReactiveCurrent_RTP_Pe{Pe}_RR{rr}_TR{tr}_timestep{dt:.3f}_seed{i}_newversion.dat"
            file_to_open = os.path.join(source_dir, readfilename)
            if os.path.exists(file_to_open):
                data_reactive_current += readfile(file_to_open)
                counter += 1
            # else: print(file_to_open)
    if rr == 1 and tr == 1:
        data_reactive_current /= counter
        print(f"number of files = {counter}")

    x = data_reactive_current[:, 0]
    y = data_reactive_current[:, 1]
    z = data_reactive_current[:, 5]
    if np.sum(z) < 0:
        z = -z
    n = data_reactive_current[:, 4]

    # z = n / np.sum(n)

    x_direct = data_reactive_current[:, 2]
    y_direct = data_reactive_current[:, 3]
    # Gitter erstellen
    x_unique = np.unique(x)
    y_unique = np.unique(y)

    # Dimensionen des Gitters
    N1 = len(x_unique)
    N2 = len(y_unique)

    # z in 2D-Array umformen
    z_grid = np.zeros((N1, N2))
    for i in range(len(z)):
        x_idx = np.where(x_unique == x[i])[0][0]
        y_idx = np.where(y_unique == y[i])[0][0]
        z_grid[x_idx, y_idx] = z[i]

        # Richtungsvektoren umformen
    x_direct_grid = np.zeros((N1, N2))
    y_direct_grid = np.zeros((N1, N2))
    for i in range(len(z)):
        x_idx = np.where(x_unique == x[i])[0][0]
        y_idx = np.where(y_unique == y[i])[0][0]
        x_direct_grid[x_idx, y_idx] = x_direct[i]
        y_direct_grid[x_idx, y_idx] = y_direct[i]

    # `x_pos` und `y_pos` können hier z.B. mit Nullen belegt werden
    # Sie können bei Bedarf auch echte Positionen zuweisen
    x_pos = np.zeros(len(x))
    y_pos = np.zeros(len(y))

    # Aufruf der Visualisierungsfunktion
    make_figure4(x_unique, y_unique, z_grid, x_pos, y_pos, x_direct_grid, y_direct_grid)


# definieren der parameter
Pe = 10
dt = 0.005

values = [0.1, 1, 10]
rr = 1
tr = 1
# plotfig4(Pe, tr, rr, dt)

rr3 = 1
tr3 = 1

rr2 = 10
tr2 = 1

rr1 = 10
tr1 = 10
plotfig5(Pe, rr1, tr1, rr2, tr2, rr3, tr3, dt)