import matplotlib.pyplot as plt
import numpy as np
import os
import argparse
from pathlib import Path
import typer

app = typer.Typer()

def get_rc_params():
    colors = ["FF220C", "70D6FF", "8AAA79", "666370", "1C1F33"]
    rcParams = plt.rcParams
    # Use LaTeX for rendering
    # Setup fonts
    rcParams["text.usetex"] = True
    rcParams["font.family"] = "serif"
    rcParams["font.size"] = 12
    rcParams["axes.labelsize"] = 12
    rcParams["axes.titlesize"] = 12
    rcParams["legend.fontsize"] = 10
    rcParams["xtick.labelsize"] = 10
    rcParams["ytick.labelsize"] = 10
    # Figure size and resolution
    rcParams["figure.figsize"] = (6, 4)
    rcParams["figure.dpi"] = 300
    # Use colors from the palette
    rcParams["axes.prop_cycle"] = plt.cycler(color=[f"#{color}" for color in colors]) # type: ignore
    # Grid
    rcParams["axes.grid"] = True
    rcParams["grid.alpha"] = 0.5
    rcParams["grid.linestyle"] = "--"
    # Point ticks to the inside of the axes
    rcParams["xtick.direction"] = "in"
    rcParams["ytick.direction"] = "in"
    rcParams["xtick.top"] = True
    rcParams["ytick.right"] = True
    return rcParams

SHOW_PLOTS = os.getenv("SHOW_PLOTS", "1") == "1"

def get_filename(path: Path, plottername: str) -> str:
    general_path = Path(__file__).parent / "plots"
    plotter_dir = general_path / plottername
    plotter_dir.mkdir(parents=True, exist_ok=True)
    filename = plotter_dir / Path(Path(path).parent.name + "_ " + Path(path).name.replace(".dat", f"_{plottername}.pdf"))
    return str(filename)


def get_filenames(directory):
    filenames = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith(".red")]
    return filenames

def get_pars(filename):
    amplitude = float(filename.split("_")[0][1:])
    omega = float(filename.split("_")[1][1:])
    return amplitude, omega

def load_data(filename):
    with open(filename, "r") as file:
        x = int(file.readline().split()[1])
    return x

def _plot_resonances(amplitudes, omegas, counts, fig, ax, linestyle='-'):
    unique_amps = sorted(set(amplitudes))
    for i, amp in enumerate(unique_amps):
        mask = np.array(amplitudes) == amp
        omegas_subset = np.array(omegas)[mask]
        counts_subset = np.array(counts)[mask]
        omegas_sorted, counts_sorted = zip(*sorted(zip(omegas_subset, counts_subset)))
        ax.plot(omegas_sorted, counts_sorted, label=f'Amplitude {amp}' if linestyle == '-' else "", linestyle=linestyle, color=f"C{i % 10}")
    ax.set_xlabel(r"$\omega_V$ (MHz)")
    ax.set_ylabel("Remaining particles")
    ax.legend()
    return fig, ax


@app.command()
def main(directories: list[Path]):
    """Plot resonance data from .red files in the specified directory."""
    plt.rcParams.update(get_rc_params())
    fig, ax = plt.subplots()

    lss = ['-', '--', ':', '-.']
    for ls, directory in zip(lss, directories):
        directory = Path(directory)
        if not directory.is_dir():
            print(f"Error: {directory} is not a valid directory.")
            continue

        filenames = get_filenames(directory)
        amplitudes, omegas, counts = [], [], []
        for filename in filenames:
            amplitude, omega = get_pars(os.path.basename(filename))
            count = load_data(filename)
            amplitudes.append(amplitude)
            omegas.append(omega)
            counts.append(count)
        _plot_resonances(amplitudes, omegas, counts, fig, ax, linestyle=ls)

    fig.savefig(get_filename(directories[0], "resonances"))
    if SHOW_PLOTS:
        fig.show()

if __name__ == "__main__":
    app()
