import matplotlib.pyplot as plt
import numpy as np
import os
import typer
from typing import Literal, List
from pathlib import Path

app = typer.Typer()

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

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

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).name.replace(".dat", f"_{plottername}.pdf")
    return str(filename)

def load_trajectory_data(filename: Path):
    data = np.loadtxt(filename)
    step_count = int(float(filename.name.split("_")[5]))
    particle_count = int(float(filename.name.split("_")[2]))
    single_particle = particle_count == 1
    if single_particle:
        data = data.reshape((data.shape[0], 1))
    data = data.reshape(step_count, 3, particle_count).transpose(1, 2, 0)
    return data, step_count, particle_count

@app.command()
def plot_z_evolution(trajectory_file: Path):
    """Plot the z-coordinate evolution from a trajectory file."""
    plt.rcParams.update(get_rc_params())
    fig, ax = plt.subplots()
    data, step_count, particle_count = load_trajectory_data(trajectory_file)
    time = np.linspace(0, 50e-6, step_count)
    for i in range(particle_count):
        ax.plot(time, data[2, i, :], label=f'Particle {i+1}')
    ax.set_xlabel('Time (s)')
    ax.set_ylabel(r'Z Position ($\mathrm{\mu m}$)')
    ax.legend()
    fig.savefig(get_filename(trajectory_file, "z_evolution"))
    if SHOW_PLOTS:
        fig.show()

@app.command()
def plot_xy_trajectory(trajectory_file: Path, plot_start_end: bool = True, plot_time: bool = False):
    """Plot the XY trajectory from a trajectory file."""
    plt.rcParams.update(get_rc_params())
    fig, ax = plt.subplots()
    data, step_count, particle_count = load_trajectory_data(trajectory_file)
    for i in range(particle_count):
        if plot_start_end:
            ax.scatter(data[0, i, 0], data[1, i, 0], marker="o", color="black", label="Start" if i == 0 else "")
            ax.scatter(data[0, i, -1], data[1, i, -1], marker="x", color="red", label="End" if i == 0 else "")
        if plot_time:
            marker = [".", ",", "o", "v", "^", "<", ">", "1", "2", "3", "4", "8"][i % 12]
            cm = plt.get_cmap('viridis')
            c = cm(np.linspace(0, 1, step_count))
            ax.scatter(data[0, i, :], data[1, i, :], label=f'Particle {i+1}', marker=marker, c=c)
        else:
            ls = "-"
            c = f"C{i % 10}"
            ax.plot(data[0, i, :], data[1, i, :], label=f'Particle {i+1}', linestyle=ls, color=c)
    ax.set_xlabel(r'X Position ($\mathrm{\mu m}$)')
    ax.set_ylabel(r'Y Position ($\mathrm{\mu m}$)')
    ax.axis('equal')
    if plot_time:
        #sm = plt.cm.ScalarMappable(cmap=cm, norm=plt.Normalize(vmin=0, vmax=50e-6))
        #sm.set_array([])
        #cbar = plt.colorbar(sm)
        #cbar.set_label('Time')
        pass
    ax.legend()
    fig.savefig(get_filename(trajectory_file, "xy_trajectory"))
    if SHOW_PLOTS:
        fig.show()

@app.command()
def phase_space_plot(trajectory_file: Path, velocity_file: Path, axis: Literal['x', 'y', 'z'] = 'z'):
    """Plot phase space (Z vs Vz) from trajectory and velocity files."""
    plt.rcParams.update(get_rc_params())
    fig, ax = plt.subplots()
    axis_index = {'x': 0, 'y': 1, 'z': 2}[axis]
    traj_data, step_count, particle_count = load_trajectory_data(trajectory_file)
    vel_data, _, _ = load_trajectory_data(velocity_file)
    for i in range(particle_count):
        ax.plot(traj_data[axis_index, i, :], vel_data[axis_index, i, :], label=f'Particle {i+1}')
    ax.set_xlabel(rf'{axis.upper()} Position ($\mathrm{{\mu m}}$)')
    ax.set_ylabel(f'{axis.upper()} Velocity (m/s)')
    ax.set_title(f'Phase Space Plot ({axis.upper()} vs V{axis.upper()})')
    ax.legend()
    fig.savefig(get_filename(trajectory_file, f"phase_space_{axis}"))
    if SHOW_PLOTS:
        fig.show()

def get_solver_name(path: Path) -> str:
    entire_path = str(path).lower()
    if "verlet" in entire_path:
        return "Velocity-Verlet"
    elif "rk4" in entire_path:
        return "Runge-Kutta 4"
    elif "boris" in entire_path:
        return "Boris"
    elif "euler" in entire_path:
        return "Euler"
    return "Unknown"

def get_step_count(path: Path) -> int:
    try:
        return int(float(path.name.split("_")[5]))
    except (IndexError, ValueError):
        return -1


@app.command()
def plot_error(trajectory_files: list[Path], analytical_file: Path):
    """Plot error between numerical and analytical trajectories."""
    plt.rcParams.update(get_rc_params())
    fig, ax = plt.subplots()
    analytical_data, step_count, particle_count = load_trajectory_data(analytical_file)
    time = np.linspace(0, 50e-6, step_count)
    for traj_file in trajectory_files:
        numerical_data, _, _ = load_trajectory_data(traj_file)
        error = numerical_data - analytical_data
        relative_error = np.mean(np.linalg.norm(error, axis=0) / np.linalg.norm(analytical_data, axis=0), axis=0)
        ax.plot(time, relative_error, label=f'{get_solver_name(traj_file)} ({get_step_count(traj_file)} steps)')
    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Relative Error')
    ax.set_yscale('log')
    ax.legend()
    fig.savefig(get_filename(analytical_file, "error_plot"))
    if SHOW_PLOTS:
        fig.show()

if __name__ == "__main__":
    app()