import matplotlib.pyplot as plt
import numpy as np
import os
import typer

app = typer.Typer()

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(energy_files: list[str]=[], labels: list[str]=[], mass: float=1.0):
    """Plot kinetic energy from velocity 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}"
        plt.plot(time, total_energy, label=labels[i], color=color, linestyle="-")
        plt.plot(time, potential_energy, color=color, linestyle="--")
        plt.plot(time, kinetic_energy, color=color, linestyle=":")
    plt.xlabel('Time (s)')
    plt.ylabel('Energy (J)')
    plt.title('Energy Evolution')
    plt.legend()
    plt.show()


if __name__ == "__main__":
    app()