#!/usr/bin/env python3
"""
Benchmark for comparing different trajectory formats:
- TACO format
- ASE trajectory
- XYZ format
- GROMACS XTC (using MDAnalysis)

This script measures compression ratios, read speeds, and write speeds.
"""

import os
import time
import argparse
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import tempfile
import shutil

# Import ASE
import ase
from ase.io import read, write
from ase import Atoms

# Import taco_format
try:
    import taco_format
except ImportError:
    print("Warning: taco_format not found. Make sure it's installed.")

# Try to import MDAnalysis for XTC support
try:
    import MDAnalysis as mda

    has_mdanalysis = True
except ImportError:
    print("Warning: MDAnalysis not installed. XTC comparison will be skipped.")
    has_mdanalysis = False


def generate_trajectory(
    num_atoms, num_frames, frame_interval=0.001, with_velocities=True
):
    """Generate a synthetic trajectory for testing"""
    print(f"Generating trajectory with {num_atoms} atoms and {num_frames} frames...")

    # Create atoms list with random elements
    elements = ["C", "H", "O", "N", "S"]
    symbols = np.random.choice(elements, num_atoms)

    # Create trajectory
    traj = []
    positions = np.random.rand(num_atoms, 3) * 10.0
    for i in range(num_frames):
        # Base positions with small random offset for each frame
        positions += np.random.normal(0, 0.01, (num_atoms, 3)) * i

        # Create periodic box
        cell = [10.0, 10.0, 10.0]

        # Create velocities (optional)
        velocities = (np.random.rand(num_atoms, 3) - 0.5) * 2.0

        # Create Atoms object
        atoms = Atoms(symbols=symbols, positions=positions, cell=cell, pbc=True)
        if with_velocities:
            atoms.set_velocities(velocities)

        # Store
        traj.append(atoms)

    return traj


def write_ase_traj(traj, filepath):
    """Write an ASE trajectory file"""
    start_time = time.time()
    write(filepath, traj)
    end_time = time.time()
    return end_time - start_time


def write_xyz(traj, filepath):
    """Write an XYZ trajectory file"""
    start_time = time.time()
    write(filepath, traj)
    end_time = time.time()
    return end_time - start_time


def write_taco(
    traj, filepath, compression_level=3, lossless=False, full_frame_interval=None
):
    """Write a TACO trajectory file"""
    start_time = time.time()
    kwargs = {"compression_level": compression_level, "lossless": lossless}
    if full_frame_interval is not None:
        kwargs["full_frame_interval"] = full_frame_interval
    taco_format.write(filepath, traj, **kwargs)
    end_time = time.time()
    return end_time - start_time


def write_xtc(traj, filepath):
    """Write a GROMACS XTC file using MDAnalysis"""
    if not has_mdanalysis:
        return None

    # First write to PDB to convert
    temp_dir = tempfile.mkdtemp()
    temp_pdb = os.path.join(temp_dir, "temp.pdb")
    write(temp_pdb, traj[0])

    # Load with MDAnalysis
    u = mda.Universe(temp_pdb)

    # Set up XTC writer
    start_time = time.time()
    with mda.Writer(filepath, n_atoms=u.atoms.n_atoms) as W:
        for i, atoms in enumerate(traj):
            # Update positions
            u.atoms.positions = atoms.get_positions()

            # Update cell dimensions
            dimensions = list(atoms.get_cell().lengths()) + [90.0, 90.0, 90.0]
            u.dimensions = dimensions

            # Write frame
            W.write(u.atoms)
    end_time = time.time()

    # Clean up
    shutil.rmtree(temp_dir)

    return end_time - start_time


def read_ase_traj(filepath):
    """Read an ASE trajectory file"""
    start_time = time.time()
    traj = read(filepath, index=":")
    end_time = time.time()
    return end_time - start_time


def read_xyz(filepath):
    """Read an XYZ trajectory file"""
    start_time = time.time()
    traj = read(filepath, index=":")
    end_time = time.time()
    return end_time - start_time


def read_taco(filepath):
    """Read a TACO trajectory file"""
    start_time = time.time()
    traj = taco_format.read(filepath)
    end_time = time.time()
    return end_time - start_time


def read_xtc(filepath, reference_pdb):
    """Read a GROMACS XTC file using MDAnalysis"""
    if not has_mdanalysis:
        return None

    start_time = time.time()
    universe = mda.Universe(reference_pdb, filepath)

    # Force reading all frames
    n_frames = len(universe.trajectory)
    for ts in universe.trajectory:
        positions = universe.atoms.positions

    end_time = time.time()
    return end_time - start_time


def get_file_size(filepath):
    """Get file size in MB"""
    return os.path.getsize(filepath) / (1024 * 1024)


def calculate_errors(original_traj, loaded_traj):
    """Calculate RMSE, MAE and max error between original and loaded trajectories"""
    if len(original_traj) != len(loaded_traj):
        print(
            f"Warning: Different number of frames: {len(original_traj)} vs {len(loaded_traj)}"
        )
        min_frames = min(len(original_traj), len(loaded_traj))
        original_traj = original_traj[:min_frames]
        loaded_traj = loaded_traj[:min_frames]

    # Calculate position errors
    rmse_list = []
    mae_list = []
    max_error_list = []

    for i in range(len(original_traj)):
        orig_pos = original_traj[i].get_positions()
        loaded_pos = loaded_traj[i].get_positions()

        # Calculate errors
        diff = orig_pos - loaded_pos
        squared_diff = np.square(diff)
        abs_diff = np.abs(diff)

        rmse = np.sqrt(np.mean(squared_diff))
        mae = np.mean(abs_diff)
        max_error = np.max(abs_diff)

        rmse_list.append(rmse)
        mae_list.append(mae)
        max_error_list.append(max_error)

    # Average over all frames
    avg_rmse = np.mean(rmse_list)
    avg_mae = np.mean(mae_list)
    avg_max_error = np.mean(max_error_list)

    return {"rmse": avg_rmse, "mae": avg_mae, "max_error": avg_max_error}


def run_benchmark(num_atoms=1000, num_frames=1000, output_dir="benchmark_results"):
    """Run the benchmark with specified parameters"""
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)

    # Generate trajectories (with and without velocities)
    traj_with_vel = generate_trajectory(num_atoms, num_frames, with_velocities=True)
    traj_without_vel = generate_trajectory(num_atoms, num_frames, with_velocities=False)

    # File paths
    ase_traj_file = os.path.join(output_dir, "trajectory.traj")
    xyz_file = os.path.join(output_dir, "trajectory.xyz")
    taco_file = os.path.join(output_dir, "trajectory.taco")
    taco_file_10 = os.path.join(output_dir, "trajectory_ffi10.taco")
    taco_file_50 = os.path.join(output_dir, "trajectory_ffi50.taco")
    taco_file_100 = os.path.join(output_dir, "trajectory_ffi100.taco")
    xtc_file = os.path.join(output_dir, "trajectory.xtc")
    temp_pdb = os.path.join(output_dir, "reference.pdb")

    # File paths for trajectory without velocities
    ase_traj_file_nv = os.path.join(output_dir, "trajectory_novel.traj")
    xyz_file_nv = os.path.join(output_dir, "trajectory_novel.xyz")
    taco_file_nv_10 = os.path.join(output_dir, "trajectory_novel_ffi10.taco")
    taco_file_nv_50 = os.path.join(output_dir, "trajectory_novel_ffi50.taco")
    taco_file_nv_100 = os.path.join(output_dir, "trajectory_novel_ffi100.taco")

    # Write reference PDB for MDAnalysis
    write(temp_pdb, traj_with_vel[0])

    # Results dictionary
    results = {}

    # ASE Trajectory format (with velocities)
    print("Testing ASE trajectory format (with velocities)...")
    write_time = write_ase_traj(traj_with_vel, ase_traj_file)
    read_time = read_ase_traj(ase_traj_file)
    file_size = get_file_size(ase_traj_file)
    results["ASE Traj (with vel)"] = {
        "write_time": write_time,
        "read_time": read_time,
        "file_size": file_size,
    }
    print(
        f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
    )

    # ASE Trajectory format (without velocities)
    print("Testing ASE trajectory format (without velocities)...")
    write_time = write_ase_traj(traj_without_vel, ase_traj_file_nv)
    read_time = read_ase_traj(ase_traj_file_nv)
    file_size = get_file_size(ase_traj_file_nv)
    loaded_traj = read(ase_traj_file_nv, index=":")
    errors = calculate_errors(traj_without_vel, loaded_traj)
    results["ASE Traj"] = {
        "write_time": write_time,
        "read_time": read_time,
        "file_size": file_size,
        "rmse": errors["rmse"],
        "mae": errors["mae"],
        "max_error": errors["max_error"],
    }
    print(
        f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
    )
    print(
        f"  RMSE: {errors['rmse']:.6f}, MAE: {errors['mae']:.6f}, Max Error: {errors['max_error']:.6f}"
    )

    # XYZ format (with velocities)
    print("Testing XYZ format (with velocities)...")
    write_time = write_xyz(traj_with_vel, xyz_file)
    read_time = read_xyz(xyz_file)
    file_size = get_file_size(xyz_file)
    results["XYZ (with vel)"] = {
        "write_time": write_time,
        "read_time": read_time,
        "file_size": file_size,
    }
    print(
        f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
    )

    # XYZ format
    print("Testing XYZ format (without velocities)...")
    write_time = write_xyz(traj_without_vel, xyz_file_nv)
    read_time = read_xyz(xyz_file_nv)
    file_size = get_file_size(xyz_file_nv)
    loaded_traj = read(xyz_file_nv, index=":")
    errors = calculate_errors(traj_without_vel, loaded_traj)
    results["XYZ"] = {
        "write_time": write_time,
        "read_time": read_time,
        "file_size": file_size,
        "rmse": errors["rmse"],
        "mae": errors["mae"],
        "max_error": errors["max_error"],
    }
    print(
        f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
    )
    print(
        f"  RMSE: {errors['rmse']:.6f}, MAE: {errors['mae']:.6f}, Max Error: {errors['max_error']:.6f}"
    )

    try:
        # TACO format (with velocities, default settings)
        print("Testing TACO format (with velocities)...")
        write_time = write_taco(traj_with_vel, taco_file)
        read_time = read_taco(taco_file)
        file_size = get_file_size(taco_file)
        results["TACO (with vel)"] = {
            "write_time": write_time,
            "read_time": read_time,
            "file_size": file_size,
        }
        print(
            f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
        )

        # TACO with full_frame_interval=10
        print("Testing TACO format with full_frame_interval=10...")
        write_time = write_taco(
            traj_without_vel, taco_file_nv_10, full_frame_interval=10
        )
        read_time = read_taco(taco_file_nv_10)
        file_size = get_file_size(taco_file_nv_10)
        loaded_traj = taco_format.read(taco_file_nv_10)
        errors = calculate_errors(traj_without_vel, loaded_traj)
        results["TACO (FFI=10)"] = {
            "write_time": write_time,
            "read_time": read_time,
            "file_size": file_size,
            "rmse": errors["rmse"],
            "mae": errors["mae"],
            "max_error": errors["max_error"],
        }
        print(
            f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
        )
        print(
            f"  RMSE: {errors['rmse']:.6f}, MAE: {errors['mae']:.6f}, Max Error: {errors['max_error']:.6f}"
        )

        # TACO with full_frame_interval=50
        print("Testing TACO format with full_frame_interval=50...")
        write_time = write_taco(
            traj_without_vel, taco_file_nv_50, full_frame_interval=50
        )
        read_time = read_taco(taco_file_nv_50)
        file_size = get_file_size(taco_file_nv_50)
        loaded_traj = taco_format.read(taco_file_nv_50)
        errors = calculate_errors(traj_without_vel, loaded_traj)
        results["TACO (FFI=50)"] = {
            "write_time": write_time,
            "read_time": read_time,
            "file_size": file_size,
            "rmse": errors["rmse"],
            "mae": errors["mae"],
            "max_error": errors["max_error"],
        }
        print(
            f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
        )
        print(
            f"  RMSE: {errors['rmse']:.6f}, MAE: {errors['mae']:.6f}, Max Error: {errors['max_error']:.6f}"
        )

        # TACO with full_frame_interval=100
        print("Testing TACO format with full_frame_interval=100...")
        write_time = write_taco(
            traj_without_vel, taco_file_nv_100, full_frame_interval=100
        )
        read_time = read_taco(taco_file_nv_100)
        file_size = get_file_size(taco_file_nv_100)
        loaded_traj = taco_format.read(taco_file_nv_100)
        errors = calculate_errors(traj_without_vel, loaded_traj)
        results["TACO (FFI=100)"] = {
            "write_time": write_time,
            "read_time": read_time,
            "file_size": file_size,
            "rmse": errors["rmse"],
            "mae": errors["mae"],
            "max_error": errors["max_error"],
        }
        print(
            f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
        )
        print(
            f"  RMSE: {errors['rmse']:.6f}, MAE: {errors['mae']:.6f}, Max Error: {errors['max_error']:.6f}"
        )

    except Exception as e:
        print(f"Error with TACO format: {e}")

    # XTC format (GROMACS via MDAnalysis)
    if has_mdanalysis:
        print("Testing XTC format...")
        write_time = write_xtc(
            traj_without_vel, xtc_file
        )  # XTC doesn't store velocities
        read_time = read_xtc(xtc_file, temp_pdb)
        file_size = get_file_size(xtc_file)

        # For XTC we need to read with MDAnalysis then convert to ASE
        temp_dir = tempfile.mkdtemp()
        try:
            u = mda.Universe(temp_pdb, xtc_file)
            loaded_traj = []

            for ts in u.trajectory:
                atoms = traj_without_vel[0].copy()
                atoms.set_positions(u.atoms.positions)
                loaded_traj.append(atoms)

            errors = calculate_errors(traj_without_vel, loaded_traj)
            results["XTC"] = {
                "write_time": write_time,
                "read_time": read_time,
                "file_size": file_size,
                "rmse": errors["rmse"],
                "mae": errors["mae"],
                "max_error": errors["max_error"],
            }
        finally:
            shutil.rmtree(temp_dir)

        print(
            f"  Write: {write_time:.3f}s, Read: {read_time:.3f}s, Size: {file_size:.2f} MB"
        )
        print(
            f"  RMSE: {errors['rmse']:.6f}, MAE: {errors['mae']:.6f}, Max Error: {errors['max_error']:.6f}"
        )

    # Generate plots
    plot_results(
        results,
        os.path.join(output_dir, f"benchmark_{num_atoms}atoms_{num_frames}frames"),
    )

    return results


def plot_results(results, output_prefix):
    """Generate comparison plots"""
    # Reorganize formats to show TACO first
    all_formats = list(results.keys())
    formats = []

    # Add TACO formats first
    for fmt in all_formats:
        if "TACO" in fmt:
            formats.append(fmt)

    # Add other formats
    for fmt in all_formats:
        if "TACO" not in fmt:
            formats.append(fmt)

    plt.rcParams.update({"font.size": 8})

    # Create fixed color mapping as specified
    color_mapping = {}
    for fmt in formats:
        if "TACO" in fmt:
            color_mapping[fmt] = "C3"  # TACO always C3
        elif "ASE" in fmt or "Traj" in fmt:
            color_mapping[fmt] = "C0"  # ASE always C0
        elif "XYZ" in fmt:
            color_mapping[fmt] = "C1"  # XYZ always C1
        elif "XTC" in fmt:
            color_mapping[fmt] = "C2"  # XTC always C2
        else:
            # Fallback
            color_mapping[fmt] = "C0"

    # File size comparison
    sizes = [results[fmt]["file_size"] for fmt in formats]
    colors = [color_mapping[fmt] for fmt in formats]
    plt.figure(figsize=(3.5, 3))
    plt.bar(formats, sizes, color=colors)
    plt.ylabel("File Size (MB)")
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()
    plt.savefig(f"{output_prefix}_filesize.svg", bbox_inches="tight", dpi=300)
    plt.savefig(f"{output_prefix}_filesize.pdf", bbox_inches="tight", dpi=300)

    # Write speed comparison (excluding XYZ)
    formats_no_xyz = [fmt for fmt in formats if "XYZ" not in fmt]
    write_times = [results[fmt]["write_time"] for fmt in formats_no_xyz]
    colors_no_xyz = [color_mapping[fmt] for fmt in formats_no_xyz]
    plt.figure(figsize=(3.5, 3))
    plt.bar(formats_no_xyz, write_times, color=colors_no_xyz)
    plt.ylabel("Write Time (s)")
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()
    plt.savefig(f"{output_prefix}_writespeed.svg", bbox_inches="tight", dpi=300)
    plt.savefig(f"{output_prefix}_writespeed.pdf", bbox_inches="tight", dpi=300)

    # Read speed comparison (excluding XYZ)
    read_times = [results[fmt]["read_time"] for fmt in formats_no_xyz]
    plt.figure(figsize=(3.5, 3))
    plt.bar(formats_no_xyz, read_times, color=colors_no_xyz)
    plt.ylabel("Read Time (s)")
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()
    plt.savefig(f"{output_prefix}_readspeed.svg", bbox_inches="tight", dpi=300)
    plt.savefig(f"{output_prefix}_readspeed.pdf", bbox_inches="tight", dpi=300)

    # Extract formats with error metrics
    formats_with_errors = [fmt for fmt in formats if "rmse" in results[fmt]]

    if formats_with_errors:
        # RMSE comparison
        rmse_values = [results[fmt]["rmse"] for fmt in formats_with_errors]
        colors_with_errors = [color_mapping[fmt] for fmt in formats_with_errors]
        plt.figure(figsize=(3.5, 3))
        plt.bar(formats_with_errors, rmse_values, color=colors_with_errors)
        plt.ylabel("RMSE (Å)")
        plt.grid(axis="y", linestyle="--", alpha=0.7)
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        plt.savefig(f"{output_prefix}_rmse.svg", bbox_inches="tight", dpi=300)
        plt.savefig(f"{output_prefix}_rmse.pdf", bbox_inches="tight", dpi=300)

        # MAE comparison
        mae_values = [results[fmt]["mae"] for fmt in formats_with_errors]
        plt.figure(figsize=(3.5, 3))
        plt.bar(formats_with_errors, mae_values, color=colors_with_errors)
        plt.ylabel("MAE (Å)")
        plt.grid(axis="y", linestyle="--", alpha=0.7)
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        plt.savefig(f"{output_prefix}_mae.svg", bbox_inches="tight", dpi=300)
        plt.savefig(f"{output_prefix}_mae.pdf", bbox_inches="tight", dpi=300)

        # Max Error comparison
        max_error_values = [results[fmt]["max_error"] for fmt in formats_with_errors]
        plt.figure(figsize=(3.5, 3))
        plt.bar(formats_with_errors, max_error_values, color=colors_with_errors)
        plt.ylabel("Max Error (Å)")
        plt.grid(axis="y", linestyle="--", alpha=0.7)
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        plt.savefig(f"{output_prefix}_maxerror.svg", bbox_inches="tight", dpi=300)
        plt.savefig(f"{output_prefix}_maxerror.pdf", bbox_inches="tight", dpi=300)

    # Write results to CSV
    csv_path = f"{output_prefix}_results.csv"
    with open(csv_path, "w") as f:
        f.write(
            "Format,File Size (MB),Write Time (s),Read Time (s),RMSE (Å),MAE (Å),Max Error (Å)\n"
        )
        for fmt in formats:
            rmse = results[fmt].get("rmse", "N/A")
            mae = results[fmt].get("mae", "N/A")
            max_error = results[fmt].get("max_error", "N/A")

            # Format appropriately based on whether it's a number or N/A
            rmse_str = f"{rmse:.6f}" if isinstance(rmse, (int, float)) else rmse
            mae_str = f"{mae:.6f}" if isinstance(mae, (int, float)) else mae
            max_error_str = (
                f"{max_error:.6f}" if isinstance(max_error, (int, float)) else max_error
            )

            f.write(
                f"{fmt},{results[fmt]['file_size']:.4f},{results[fmt]['write_time']:.4f},"
                f"{results[fmt]['read_time']:.4f},{rmse_str},{mae_str},{max_error_str}\n"
            )

    # Generate HTML table for documentation
    try:
        import sys
        import os

        sys.path.append(os.path.dirname(os.path.abspath(__file__)))
        from csv_to_html_table import convert_csv_to_html_table

        html_path = f"{output_prefix}_table.html"
        convert_csv_to_html_table(csv_path, html_path)
        print(f"Generated HTML table: {html_path}")
    except Exception as e:
        print(f"Error generating HTML table: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Benchmark different trajectory file formats"
    )
    parser.add_argument(
        "--atoms", type=int, default=1000, help="Number of atoms in test trajectory"
    )
    parser.add_argument(
        "--frames", type=int, default=1000, help="Number of frames in test trajectory"
    )
    parser.add_argument(
        "--output", type=str, default="benchmark_results", help="Output directory"
    )
    args = parser.parse_args()

    print(f"Starting benchmark with {args.atoms} atoms and {args.frames} frames")
    print(f"Results will be saved to: {args.output}")

    # Run benchmark
    run_benchmark(args.atoms, args.frames, args.output)

    print("Benchmark completed. Results are saved to the output directory.")
