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

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])
    # 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 load_data(total_energy_file):
    total_energy = np.loadtxt(total_energy_file)
    kin_file = total_energy_file.replace("total_energies", "kinetic_energies")
    pot_file = total_energy_file.replace("total_energies", "potential_energies")
    kinetic_energy = np.loadtxt(kin_file)
    potential_energy = np.loadtxt(pot_file)
    step_count = len(total_energy)
    return total_energy, kinetic_energy, potential_energy, step_count

@app.command()
def plot_kinetic_energy(total_energy_files: List[Path], output_file: Path):
    """Plot kinetic energy from velocity files."""
    plt.rcParams.update(get_rc_params())
    fig, ax = plt.subplots()
    energy_files = [str(fn) for fn in total_energy_files if fn.is_file()]
    labels = [f"{fn.split('/')[-2]} ({fn.split('/')[-1].split('_')[4]})" for fn in energy_files]
    for i, energy_file in enumerate(energy_files):
        total_energy, kinetic_energy, potential_energy, step_count = load_data(energy_file)
        time = np.linspace(0, 50e-6, step_count)
        color = f"C{i % 10}"
        ax.plot(time, total_energy, label=labels[i], color=color, linestyle="-")
        ax.plot(time, potential_energy, color=color, linestyle="--")
        ax.plot(time, kinetic_energy, color=color, linestyle=":")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel(r"Energy (Da $\cdot$ m$^2$/s$^2$)")
    ax.legend()
    fig.savefig(output_file)
    if SHOW_PLOTS:
        fig.show()


if __name__ == "__main__":
    app()