import pandas as pd
import numpy as np
import os
import pathlib
import re
import typer
from uncertainties import ufloat

app = typer.Typer()

hr_solver_labels = {
    'euler': 'Euler',
    'velocityverlet': 'Velocity Verlet',
    'rk4': 'Runge-Kutta 4',
    'boris': 'Boris',
}

def load_performance_data(file_path: pathlib.Path) -> dict:
    results = {}
    results['file_path'] = file_path
    results['file_name'] = file_path.name

    entire_path = str(file_path.resolve())

    particles_match = re.search(r'(\d+)_particles', entire_path)
    if particles_match:
        results['num_particles'] = int(particles_match.group(1))
    else:
        results['num_particles'] = np.nan

    steps_match = re.search(r'(\d+)_steps', entire_path)
    if steps_match:
        results['num_steps'] = int(steps_match.group(1))
    else:
        results['num_steps'] = np.nan

    interactions_match = re.search(r'_(enabled|disabled)_interactions', entire_path)
    if interactions_match:
        results['interactions'] = interactions_match.group(1) == 'enabled'
    else:
        results['interactions'] = None

    results['oscillating'] = 'oscillating_potential' in entire_path

    solver = "unknown"
    if 'euler' in entire_path.lower():
        solver = 'euler'
    elif 'velocityverlet' in entire_path.lower():
        solver = 'velocityverlet'
    elif 'rk4' in entire_path.lower():
        solver = 'rk4'
    elif 'boris' in entire_path.lower():
        solver = 'boris'
    results['solver'] = solver

    with open(file_path, 'r') as file:
        # Example: CPU time used: 0.02746 seconds
        cpu_time_match = re.search(r'CPU time used:\s+([\d.]+)\s+seconds', file.read())
        if cpu_time_match:
            results['cpu_time'] = float(cpu_time_match.group(1))
        else:
            results['cpu_time'] = np.nan
    return results

def aggregate_performance_data(directory: pathlib.Path) -> pd.DataFrame:
    all_results = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith('_timing.txt'):
                file_path = pathlib.Path(root) / file
                result = load_performance_data(file_path)
                all_results.append(result)
    return pd.DataFrame(all_results)

def conv2siunitx(val, err, err_points=2):
    val = f'{val:.20e}'.split('e')
    err = f'{err:.20e}'.split('e') 
    first_uncertain = int(val[1]) - int(err[1]) + err_points

    my_val = f'{np.round(float(val[0]), first_uncertain-1):.10f}'
    my_err = f'{np.round(float(err[0]), err_points-1):.10f}'.replace('.','')
    # Avoid 1. and write 1 instead
    if first_uncertain > 1:
        first_uncertain = first_uncertain + 1

    content_string = f'{my_val[:first_uncertain]}({my_err[:err_points]})e{val[1]}'
    return r"\num{" + content_string + "}"

@app.command()
def print_performance_summary(directory: pathlib.Path):
    df = aggregate_performance_data(directory)  
    summary = df.loc[:, ['solver', 'num_particles', 'num_steps', 'interactions', 'oscillating', 'cpu_time']]
    print(summary)

@app.command()
def statistical_solver_comparison(directory: pathlib.Path, save_path: pathlib.Path = pathlib.Path("tables/solver_comparison.tex")):
    df = aggregate_performance_data(directory)
    df = df.loc[df.oscillating == True]
    print(f"Statistical Summary of CPU Time by Solver and Interaction Setting for {directory}: {len(df)} records")
    summary = df.groupby(['solver', 'interactions']).agg(
        mean_cpu_time=('cpu_time', 'mean'),
        std_cpu_time=('cpu_time', 'std'),
        count=('cpu_time', 'count'),
        min_cpu_time=('cpu_time', 'min'),
        max_cpu_time=('cpu_time', 'max')
    )
    print(summary)
    for solver in ["velocityverlet", "rk4", "boris"]:
        avg_speedup = ufloat(0, 0)
        for interactions in [True, False]:
            boris_row = summary.loc[('boris', interactions)]
            solver_row = summary.loc[(solver, interactions)]
            boris_ufloat = ufloat(boris_row['mean_cpu_time'], boris_row['std_cpu_time']) # type: ignore
            solver_ufloat = ufloat(solver_row['mean_cpu_time'], solver_row['std_cpu_time']) # type: ignore
            speedup = solver_ufloat / boris_ufloat - 1 # type: ignore
            avg_speedup += speedup / 2
            print(f"{hr_solver_labels[solver]} (interactions={'enabled' if interactions else 'disabled'}): {speedup:.2%} slower than Boris")
        print(f"Average: {hr_solver_labels[solver]} is {avg_speedup:.2%} slower than Boris")
    summary['siunitx'] = summary.apply(lambda row: conv2siunitx(row['mean_cpu_time'], row['std_cpu_time']), axis=1)
    summary = summary.reset_index()
    summary = summary.pivot(index='solver', columns='interactions', values='siunitx')
    summary.columns = ['Interactions Disabled', 'Interactions Enabled']
    summary = summary.reindex(index=['velocityverlet', 'rk4', 'boris'])
    summary.index = summary.index.map(hr_solver_labels)
    summary.index.name = 'Solver'
    os.makedirs(save_path.parent, exist_ok=True)
    summary.to_latex(save_path, escape=False, column_format='lcc')
    print(f"Saved LaTeX table to {save_path}")



@app.command()
def solver_scaling_analysis(directory: pathlib.Path, save_path: pathlib.Path = pathlib.Path("tables/solver_scaling_analysis.tex")):
    df = aggregate_performance_data(directory)
    df = df.loc[(df.oscillating == False) & (df.interactions == True)]
    df.loc[:, "us_per_particle_step"] = (df.cpu_time / (df.num_particles * df.num_steps)) * 1e6
    print(f"Solver Scaling Analysis for {directory}: {len(df)} records")
    summary = df.groupby(['solver', 'num_particles']).agg(
        mean_us_per_particle_step=('us_per_particle_step', 'mean'),
        std_us_per_particle_step=('us_per_particle_step', 'std'),
        count=('us_per_particle_step', 'count'),
        min_us_per_particle_step=('us_per_particle_step', 'min'),
        max_us_per_particle_step=('us_per_particle_step', 'max')
    ).reset_index()
    print(summary)
    summary['siunitx'] = summary.apply(lambda row: conv2siunitx(row['mean_us_per_particle_step'], row['std_us_per_particle_step']), axis=1)
    summary = summary.pivot(index='solver', columns='num_particles', values='siunitx')
    summary.columns.name = 'Num. Particles'
    summary.index.name = 'Solver'
    summary = summary.reindex(index=['euler', 'velocityverlet', 'rk4', 'boris'])
    summary.index = summary.index.map(hr_solver_labels)
    os.makedirs(save_path.parent, exist_ok=True)
    summary.to_latex(save_path, escape=False, column_format='l' + 'c' * len(summary.columns))
    print(f"Saved LaTeX table to {save_path}")


if __name__ == "__main__":
    app()