"""Compare the efficiency of different solver for different domain sizes."""

import argparse
import platform
import subprocess
from itertools import cycle
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from tabulate import tabulate

from .utils import ChdirTemporaryDirectory
from ..post import extract

JyJz = 24
DOMAIN_SIZES: tuple[int, ...] = (2**7, 2**8, 2**9, 2**10, 2**11, 2**12)
CSV_FILENAME: str = "bench_efficiency.csv"


def get_num_iterations(Jx: int) -> int:
    """Get the number of iterations for the benchmark based on domain size."""
    return max(2**17 // Jx, 1)


def get_gpu_name() -> str:
    """Get the name of the OpenCL device to use for the benchmark."""
    from llg3d.solvers.opencl import get_context_and_device

    _, device = get_context_and_device("gpu")
    try:
        # AMD GPUs have a more user-friendly board name
        return str(device.board_name_amd)
    except AttributeError:
        # Fallback to the device name
        return str(device.name)


def get_cpu_name() -> str | None:
    """Get the CPU name of the current machine."""
    try:
        os = platform.system()
        if os == "Linux":
            with open("/proc/cpuinfo") as f:
                for line in f:
                    if "model name" in line:
                        return line.split(":", 1)[1].strip()
        if os == "Darwin":
            return subprocess.check_output(
                ["sysctl", "-n", "machdep.cpu.brand_string"], text=True
            ).strip()
        if os == "Windows":
            return (
                subprocess.check_output(["wmic", "cpu", "get", "Name"], text=True)
                .splitlines()[1]
                .strip()
            )
        else:
            return platform.processor()
    except Exception:
        return None


def get_cpu_gpu_legend() -> str:
    """Get a legend string with CPU and GPU names."""
    cpu_name = get_cpu_name()
    gpu_name = get_gpu_name()
    return f"CPU: {cpu_name} | GPU: {gpu_name}"


Results = dict[str, list[float]]


def run_benchmark(mpi_nprocs: int, repeats: int = 1) -> tuple[Results, Results | None]:
    """
    Run the benchmark for different solvers and domain sizes.

    Args:
        mpi_nprocs: Number of MPI processes to use for the MPI solver.
        repeats: Number of times to repeat each measurement. If >1, function
            returns both means and std-devs per solver/domain.

    Returns:
        (means, stds): two mappings from display-name (e.g. 'NumPy (1 CPU core)')
            to lists of mean and std values per domain size. If `repeats == 1`,
            `stds` is None.
    """
    solvers = {
        "numpy": "NumPy (1 CPU core)",
        "mpi": f"MPI ({mpi_nprocs} CPU cores)",
        "opencl": "OpenCL (1 GPU)",
    }
    # Use display names as keys in results for simpler downstream handling
    means: Results = {display: [] for display in solvers.values()}
    stds: Results | None = (
        {display: [] for display in solvers.values()} if repeats > 1 else None
    )

    for Jx in DOMAIN_SIZES:
        N = get_num_iterations(Jx)
        for solver in solvers:
            # collect `repeats` measurements
            values: list[float] = []
            for r in range(repeats):
                # Move to a temporary directory to avoid overwriting files
                with ChdirTemporaryDirectory():
                    cmd = [
                        "llg3d",
                        "--N",
                        str(N),
                        "--Jx",
                        str(Jx),
                        "--Jy",
                        str(JyJz),
                        "--Jz",
                        str(JyJz),
                        "--n_mean",
                        "0",
                        "--precision",
                        "single",
                        "--result_file",
                        "run.npz",
                        "--solver",
                        solver,
                    ]
                    if solver == "mpi":
                        cmd = ["mpirun", "-np", str(mpi_nprocs)] + cmd
                    print(f"Running: {' '.join(cmd)} (run {r + 1}/{repeats})")
                    result = subprocess.run(cmd, capture_output=True, text=True)
                    if result.stderr:
                        # show output for debugging if there's an error
                        print("STDOUT:", result.stdout)
                        print("STDERR:", result.stderr)
                    (eff,) = extract.extract_values(
                        Path("run.npz"), "results/metrics/efficiency"
                    )
                    values.append(float(eff))

            display_key = solvers[solver]
            mean_val = float(np.mean(values))
            means[display_key].append(mean_val)
            if stds is not None:
                stds[display_key].append(float(np.std(values, ddof=0)))

    return means, stds


def save_as_csv(
    results: Results,
    filename: str,
    stds: Results | None = None,
    legend: str | None = None,
) -> None:
    """
    Save the benchmark results as a CSV file.

    Args:
        results: The benchmark results.
        filename: The output CSV filename.
        stds: Optional standard deviations for the results.
        legend: An optional legend string to include as metadata.
    """
    import json

    solver_keys = list(results.keys())
    with open(filename, "w") as f:
        # write metadata as JSON comment on the first line if provided
        if legend is not None:
            metadata = {"legend": legend}
            f.write("#META " + json.dumps(metadata) + "\n")
        # If stds provided, write two columns per solver: <solver>, <solver> std
        header = ["Domain size"]
        for solver in solver_keys:
            header.append(solver)
            if stds is not None:
                header.append(f"{solver} std")
        f.write(",".join(header) + "\n")
        for i, J in enumerate(DOMAIN_SIZES):
            row = [str(J)]
            for solver in solver_keys:
                val = results[solver][i]
                row.append(f"{val:.6e}")
                if stds is not None:
                    std_val = stds[solver][i]
                    row.append(f"{std_val:.6e}")
            f.write(",".join(row) + "\n")


def results_as_table(
    results: Results, stds: Results | None = None, legend: str | None = ""
) -> str:
    """
    Format the benchmark results as a table.

    Add an acceleration column showing speedup relative to the NumPy solver.

    Args:
        results: The benchmark results.
        stds: Optional standard deviations for the results.
        legend: An optional legend string to include above the table.

    Returns:
        A string representing the results in table format.
    """
    solver_keys = list(results.keys())
    # legend should be provided by the caller (run or metadata); do not auto-fetch
    if legend is None:
        legend = ""
    # Ensure NumPy is the reference for acceleration
    numpy_key = next((k for k in solver_keys if k.startswith("NumPy")), solver_keys[0])
    # Reorder keys so numpy_key is first
    ordered_keys = [numpy_key] + [k for k in solver_keys if k != numpy_key]
    headers = ["Domain size", numpy_key]
    for solver in ordered_keys[1:]:
        headers.append(solver + " (Accel)")

    table = []
    numpy_times = results[numpy_key]
    for i, J in enumerate(DOMAIN_SIZES):
        row = [str(J)]
        val_numpy = results[numpy_key][i]
        if stds is not None:
            std_numpy = stds[numpy_key][i]
            row.append(f"{val_numpy:.1e} ± {std_numpy:.1e}")
        else:
            row.append(f"{val_numpy:.1e}")
        for solver in ordered_keys[1:]:
            val = results[solver][i]
            accel = numpy_times[i] / val if val != 0 else float("inf")
            if stds is not None:
                std_val = stds[solver][i]
                row.append(f"{val:.1e} ± {std_val:.1e} ({accel:5.1f}x)")
            else:
                row.append(f"{val:.1e} ({accel:5.1f}x)")
        table.append(row)
    tab = tabulate(
        table, headers=headers, tablefmt="simple", numalign="right", stralign="right"
    )
    return f"{legend}\n{tab}"


def plot(
    results: Results,
    show: bool = False,
    legend: str | None = None,
    errorbars: bool = False,
    stds: Results | None = None,
) -> None:
    """
    Plot the benchmark results.

    Args:
        results: The benchmark results.
        show: Whether to display the plot interactively .
        legend: An optional legend string to include in the plot title.
        errorbars: Whether to plot error bars using stds if provided.
        stds: Optional standard deviations for the results.
    """
    fig, ax = plt.subplots(figsize=(7, 5))
    fig.suptitle("LLG3D efficiency vs domain size")
    markers = cycle(["o", "s", "^", "D", "v", "x", "*", "+"])
    for solver in results.keys():
        if errorbars and stds is not None and solver in stds:
            ax.errorbar(
                DOMAIN_SIZES,
                results[solver],
                yerr=stds[solver],
                marker=next(markers),
                label=solver,
                capsize=3,
                linestyle="-",
            )
        else:
            ax.plot(DOMAIN_SIZES, results[solver], marker=next(markers), label=solver)
    ax.set_xlabel(f"Domain size $(J_x \\times {JyJz} \\times {JyJz})$")
    ax.set_ylabel("Efficiency [s/iteration/point]")
    ax.set_title(
        legend if legend is not None else get_cpu_gpu_legend(),
        fontsize=8,
    )
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.legend()
    ax.grid(True, which="both", ls=":")
    fig.tight_layout()
    fig.savefig("bench_efficiency.png", dpi=300)
    if show:
        plt.show()


def load_csv(filename: str) -> tuple[Results, Results | None, dict]:
    """
    Load benchmark results from a CSV file.

    Returns results keyed by display-name (same format as run_benchmark output).

    Args:
        filename: The CSV filename to load.

    Returns:
        A tuple (results, stds, metadata) where `results` is the benchmark results,
        `stds` are the standard deviations if present (else None), and `metadata` is
        any metadata extracted from the CSV file.
    """
    import csv
    import json

    results: Results = {}
    metadata: dict = {}
    stds: Results | None = None
    with open(filename, "r") as f:
        # read possible metadata comment lines
        pos = f.tell()
        first = f.readline()
        if first.startswith("#META "):
            try:
                metadata = json.loads(first[len("#META ") :])
            except Exception:
                metadata = {"raw": first.strip()}
            header_line = f.readline()
        else:
            # no metadata, rewind
            f.seek(pos)
            header_line = f.readline()
        reader = csv.reader([header_line] + f.readlines())
        header = next(reader)
        domain_sizes = []
        cols = header[1:]
        # detect pairs: solver, solver std
        solver_names: list[str] = []
        solver_has_std: dict[str, bool] = {}
        i = 0
        while i < len(cols):
            name = cols[i]
            if i + 1 < len(cols) and cols[i + 1].strip().endswith(" std"):
                solver_names.append(name)
                solver_has_std[name] = True
                if stds is None:
                    stds = {}
                results[name] = []
                if name not in stds:
                    stds[name] = []
                i += 2
            else:
                solver_names.append(name)
                solver_has_std[name] = False
                results[name] = []
                i += 1

        for row in reader:
            domain_sizes.append(int(row[0]))
            j = 0
            for name in solver_names:
                val = row[1 + j]
                val = val.split(" ")[0]
                results[name].append(float(val))
                j += 1
                if solver_has_std.get(name, False):
                    if 1 + j < len(row):
                        std_val = row[1 + j]
                        std_val = std_val.split(" ")[0]
                        if stds is not None:
                            stds[name].append(float(std_val))
                    elif stds is not None:
                        # missing std value, append 0
                        stds[name].append(0.0)
                    j += 1
    # Remap keys if needed to match original order
    # Here, we assume the order is the same
    return results, stds, metadata


def main():
    """Main function to run the benchmark, plot, or report results."""
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="mode", required=True)

    parser_run = subparsers.add_parser(
        "run",
        help="Run the benchmark and save the CSV file",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser_run.add_argument(
        "--mpi-nprocs",
        type=int,
        default=32,
        help="Number of MPI processes",
    )
    parser_run.add_argument(
        "--repeats",
        type=int,
        default=1,
        help="Number of repetitions per measurement",
    )
    parser_run.add_argument(
        "--csv", type=str, default=CSV_FILENAME, help="CSV file to save"
    )

    parser_plot = subparsers.add_parser(
        "plot",
        help="Plot the graph from the CSV file",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser_plot.add_argument(
        "--csv", type=str, default=CSV_FILENAME, help="CSV file to load"
    )
    parser_plot.add_argument(
        "--show", action="store_true", help="Show the plot interactively", default=False
    )
    parser_plot.add_argument(
        "--errorbars",
        action="store_true",
        default=False,
        help="Plot error bars using std columns if present in the CSV",
    )

    parser_report = subparsers.add_parser(
        "report",
        help="Print a results table from a CSV file",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser_report.add_argument(
        "--csv", type=str, default=CSV_FILENAME, help="CSV file to load"
    )
    parser_report.add_argument(
        "--show-std",
        action="store_true",
        default=False,
        help="Include std deviations in the report if present in the CSV",
    )

    args = parser.parse_args()

    if args.mode == "run":
        results, stds = run_benchmark(args.mpi_nprocs, repeats=args.repeats)
        legend = get_cpu_gpu_legend()
        save_as_csv(results, filename=args.csv, stds=stds, legend=legend)
        print(results_as_table(results, stds=stds, legend=legend))
    elif args.mode == "plot":
        results, stds, metadata = load_csv(args.csv)
        legend = metadata.get("legend") if metadata else None
        plot(
            results, show=args.show, legend=legend, errorbars=args.errorbars, stds=stds
        )
    elif args.mode == "report":
        results, stds, metadata = load_csv(args.csv)
        legend = metadata.get("legend") if metadata else None
        if args.show_std:
            print(results_as_table(results, stds=stds, legend=legend))
        else:
            print(results_as_table(results, stds=None, legend=legend))
