"""
PSF parameter optimisation and fitting routines for mirror alignment and reflection parameters.

This module provides functions for loading PSF data, generating random parameter sets,
running PSF simulations, calculating RMSD, and finding the best-fit parameters for a given
telescope model.

PSF (Point Spread Function) describes how a point source of light is spread out by the
optical system, and RMSD (Root Mean Squared Deviation) is used as the optimization metric
to quantify the difference between measured and simulated PSF curves.
"""

import logging
from collections import OrderedDict

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_pdf import PdfPages

from simtools.data_model import model_data_writer as writer
from simtools.model import model_utils
from simtools.ray_tracing.ray_tracing import RayTracing
from simtools.utils import general as gen
from simtools.visualization import visualize

logger = logging.getLogger(__name__)

# Constants
RADIUS_CM = "Radius [cm]"
CUMULATIVE_PSF = "Cumulative PSF"

MRRA_RANGE_DEFAULT = 0.004  # Mirror reflection random angle range
MRF_RANGE_DEFAULT = 0.1  # Mirror reflection fraction range
MRRA2_RANGE_DEFAULT = 0.03  # Second mirror reflection random angle range
MAR_RANGE_DEFAULT = 0.005  # Mirror alignment random range
MAX_OFFSET_DEFAULT = 4.5  # Maximum off-axis angle in degrees
OFFSET_STEPS_DEFAULT = 0.1  # Step size for off-axis angle sampling


def load_psf_data(data_file):
    """
    Load data from a text file containing cumulative PSF measurements.

    Parameters
    ----------
    data_file : str
        Name of the data file with the measured cumulative PSF.
        Expected format:
        Column 0: radial distance in mm
        Column 2: cumulative PSF values

    Returns
    -------
    numpy.ndarray
        Loaded and processed data with radius in cm and normalized cumulative PSF.
    """
    d_type = {"names": (RADIUS_CM, CUMULATIVE_PSF), "formats": ("f8", "f8")}
    data = np.loadtxt(data_file, dtype=d_type, usecols=(0, 2))
    data[RADIUS_CM] *= 0.1  # Convert from mm to cm
    data[CUMULATIVE_PSF] /= np.max(np.abs(data[CUMULATIVE_PSF]))  # Normalize to max = 1.0
    return data


def calculate_rmsd(data, sim):
    """Calculate Root Mean Squared Deviation to be used as metric to find the best parameters."""
    return np.sqrt(np.mean((data - sim) ** 2))


def add_parameters(
    all_parameters,
    mirror_reflection,
    mirror_align,
    mirror_reflection_fraction=0.15,
    mirror_reflection_2=0.035,
):
    """
    Transform and add parameters to the all_parameters list.

    Parameters
    ----------
    mirror_reflection : float
        The random angle of mirror reflection.
    mirror_align : float
        The random angle for mirror alignment (both horizontal and vertical).
    mirror_reflection_fraction : float, optional
        The fraction of the mirror reflection. Default is 0.15.
    mirror_reflection_2 : float, optional
        A secondary random angle for mirror reflection. Default is 0.035.

    Returns
    -------
    None
        Updates the all_parameters list in place.
    """
    pars = {
        "mirror_reflection_random_angle": [
            mirror_reflection,
            mirror_reflection_fraction,
            mirror_reflection_2,
        ],
        "mirror_align_random_horizontal": [mirror_align, 28.0, 0.0, 0.0],
        "mirror_align_random_vertical": [mirror_align, 28.0, 0.0, 0.0],
    }
    all_parameters.append(pars)


def get_previous_values(tel_model):
    """
    Retrieve previous parameter values from the telescope model.

    Parameters
    ----------
    tel_model : TelescopeModel
        Telescope model object.

    Returns
    -------
    tuple
        Tuple containing the previous values of mirror_reflection_random_angle (first entry),
        mirror_reflection_fraction, second entry), mirror_reflection_random_angle (third entry),
        and mirror_align_random_horizontal/vertical.
    """
    split_par = tel_model.get_parameter_value("mirror_reflection_random_angle")
    mrra_0, mfr_0, mrra2_0 = split_par[0], split_par[1], split_par[2]
    mar_0 = tel_model.get_parameter_value("mirror_align_random_horizontal")[0]
    logger.debug(
        "Previous parameter values:\n"
        f"MRRA = {mrra_0!s}\n"
        f"MRF = {mfr_0!s}\n"
        f"MRRA2 = {mrra2_0!s}\n"
        f"MAR = {mar_0!s}\n"
    )
    return mrra_0, mfr_0, mrra2_0, mar_0


def generate_random_parameters(
    all_parameters, n_runs, args_dict, mrra_0, mfr_0, mrra2_0, mar_0, tel_model
):
    """
    Generate random parameters for tuning.

    The parameter ranges around the previous values are configurable via module constants.

    Parameters
    ----------
    all_parameters : list
        List to store all parameter sets.
    n_runs : int
        Number of random parameter combinations to test.
    args_dict : dict
        Dictionary containing parsed command-line arguments.
    mrra_0 : float
        Initial value of mirror_reflection_random_angle.
    mfr_0 : float
        Initial value of mirror_reflection_fraction.
    mrra2_0 : float
        Initial value of the second mirror_reflection_random_angle.
    mar_0 : float
        Initial value of mirror_align_random_horizontal/vertical.
    tel_model : TelescopeModel
        Telescope model object to check if it's a dual mirror telescope.
    """
    if args_dict["fixed"]:
        logger.debug("fixed=True - First entry of mirror_reflection_random_angle is kept fixed.")

    is_dual_mirror = model_utils.is_two_mirror_telescope(tel_model.name)
    if is_dual_mirror:
        mar_fixed_value = 0.0
    else:
        mar_fixed_value = None

    for _ in range(n_runs):
        mrra_range = MRRA_RANGE_DEFAULT if not args_dict["fixed"] else 0
        mrf_range = MRF_RANGE_DEFAULT
        mrra2_range = MRRA2_RANGE_DEFAULT
        mar_range = MAR_RANGE_DEFAULT
        rng = np.random.default_rng(seed=args_dict.get("random_seed"))
        mrra = rng.uniform(max(mrra_0 - mrra_range, 0), mrra_0 + mrra_range)
        mrf = rng.uniform(max(mfr_0 - mrf_range, 0), mfr_0 + mrf_range)
        mrra2 = rng.uniform(max(mrra2_0 - mrra2_range, 0), mrra2_0 + mrra2_range)

        # Set mar to 0 for dual mirror telescopes, otherwise use random value
        if mar_fixed_value is not None:
            mar = mar_fixed_value
        else:
            mar = rng.uniform(max(mar_0 - mar_range, 0), mar_0 + mar_range)

        add_parameters(all_parameters, mrra, mar, mrf, mrra2)


def _run_ray_tracing_simulation(tel_model, site_model, args_dict, pars):
    """
    Run a ray tracing simulation with the given telescope parameters.

    Parameters
    ----------
    tel_model : TelescopeModel
        Telescope model object.
    site_model : SiteModel
        Site model object.
    args_dict : dict
        Dictionary containing parsed command-line arguments.
    pars : dict
        Parameter set dictionary.

    Returns
    -------
    tuple
        (d80, simulated_data) - D80 value and simulated data from ray tracing.
    """
    if pars is not None:
        tel_model.change_multiple_parameters(**pars)
    else:
        raise ValueError("No best parameters found")

    ray = RayTracing(
        telescope_model=tel_model,
        site_model=site_model,
        simtel_path=args_dict["simtel_path"],
        zenith_angle=args_dict["zenith"] * u.deg,
        source_distance=args_dict["src_distance"] * u.km,
        off_axis_angle=[0.0] * u.deg,
    )
    ray.simulate(test=args_dict.get("test", False), force=True)
    ray.analyze(force=True, use_rx=False)
    im = ray.images()[0]
    d80 = im.get_psf()

    return d80, im


def _create_psf_simulation_plot(data_to_plot, pars, d80, rmsd, is_best, pdf_pages):
    """
    Create a plot for PSF simulation results.

    Parameters
    ----------
    data_to_plot : dict
        Data dictionary for plotting.
    pars : dict
        Parameter set dictionary.
    d80 : float
        D80 value.
    rmsd : float
        RMSD value.
    is_best : bool
        Whether this is the best parameter set.
    pdf_pages : PdfPages
        PDF pages object for saving plots.
    """
    fig = visualize.plot_1d(
        data_to_plot,
        plot_difference=True,
        no_markers=True,
    )
    ax = fig.get_axes()[0]
    ax.set_ylim(0, 1.05)
    ax.set_ylabel(CUMULATIVE_PSF)

    title_prefix = "* " if is_best else ""
    ax.set_title(
        f"{title_prefix}refl_rnd = "
        f"{pars['mirror_reflection_random_angle'][0]:.5f}, "
        f"{pars['mirror_reflection_random_angle'][1]:.5f}, "
        f"{pars['mirror_reflection_random_angle'][2]:.5f}\n"
        f"align_rnd = {pars['mirror_align_random_vertical'][0]:.5f}, "
        f"{pars['mirror_align_random_vertical'][1]:.5f}, "
        f"{pars['mirror_align_random_vertical'][2]:.5f}, "
        f"{pars['mirror_align_random_vertical'][3]:.5f}"
    )

    d80_color = "red" if is_best else "black"
    d80_weight = "bold" if is_best else "normal"
    d80_text = f"D80 = {d80:.5f} cm"

    ax.text(
        0.5,
        0.3,
        f"{d80_text}\nRMSD = {rmsd:.4f}",
        verticalalignment="center",
        horizontalalignment="left",
        transform=ax.transAxes,
        color=d80_color,
        weight=d80_weight,
        bbox={"boxstyle": "round,pad=0.3", "facecolor": "yellow", "alpha": 0.7}
        if is_best
        else None,
    )

    if is_best:
        fig.text(
            0.02,
            0.02,
            "* Best parameter set (lowest RMSD)",
            fontsize=8,
            style="italic",
            color="red",
        )

    pdf_pages.savefig(fig, bbox_inches="tight")
    plt.clf()


def run_psf_simulation(
    tel_model,
    site_model,
    args_dict,
    pars,
    data_to_plot,
    radius,
    pdf_pages=None,
    is_best=False,
    return_simulated_data=False,
):
    """
    Run the simulation for one set of parameters and return D80, RMSD.

    Parameters
    ----------
    tel_model : TelescopeModel
        Telescope model object.
    site_model : SiteModel
        Site model object.
    args_dict : dict
        Dictionary containing parsed command-line arguments.
    pars : dict
        Parameter set dictionary.
    data_to_plot : dict
        Data dictionary for plotting.
    radius : array-like
        Radius data.
    pdf_pages : PdfPages, optional
        PDF pages object for plotting. If None, no plotting is done.
    is_best : bool, optional
        Whether this is the best parameter set for highlighting in plots.
    return_simulated_data : bool, optional
        If True, returns simulated data as third element in return tuple.

    Returns
    -------
    tuple
        (d80, rmsd) if return_simulated_data=False
        (d80, rmsd, simulated_data) if return_simulated_data=True
    """
    d80, im = _run_ray_tracing_simulation(tel_model, site_model, args_dict, pars)

    if radius is None:
        raise ValueError("Radius data is not available.")

    simulated_data = im.get_cumulative_data(radius * u.cm)
    rmsd = calculate_rmsd(data_to_plot["measured"][CUMULATIVE_PSF], simulated_data[CUMULATIVE_PSF])

    # Handle plotting if requested
    if pdf_pages is not None and args_dict.get("plot_all", False):
        data_to_plot["simulated"] = simulated_data
        _create_psf_simulation_plot(data_to_plot, pars, d80, rmsd, is_best, pdf_pages)
        del data_to_plot["simulated"]

    return (d80, rmsd, simulated_data) if return_simulated_data else (d80, rmsd)


def load_and_process_data(args_dict):
    """
    Load and process data if specified in the command-line arguments.

    Returns
    -------
    - data_to_plot: OrderedDict containing loaded and processed data.
    - radius: Radius data from loaded data (if available).
    """
    data_to_plot = OrderedDict()
    radius = None
    if args_dict["data"] is not None:
        data_file = gen.find_file(args_dict["data"], args_dict["model_path"])
        data_to_plot["measured"] = load_psf_data(data_file)
        radius = data_to_plot["measured"][RADIUS_CM]
    return data_to_plot, radius


def _create_plot_for_parameters(pars, rmsd, d80, simulated_data, data_to_plot, is_best, pdf_pages):
    """
    Create a single plot for a parameter set.

    Parameters
    ----------
    pars : dict
        Parameter set dictionary
    rmsd : float
        RMSD value for this parameter set
    d80 : float
        D80 value for this parameter set
    simulated_data : array
        Simulated data for plotting
    data_to_plot : dict
        Data dictionary for plotting
    is_best : bool
        Whether this is the best parameter set
    pdf_pages : PdfPages
        PDF pages object to save the plot
    """
    original_simulated = data_to_plot.get("simulated")
    data_to_plot["simulated"] = simulated_data

    fig = visualize.plot_1d(
        data_to_plot,
        plot_difference=True,
        no_markers=True,
    )
    ax = fig.get_axes()[0]
    ax.set_ylim(0, 1.05)
    ax.set_ylabel(CUMULATIVE_PSF)

    title_prefix = "* " if is_best else ""

    ax.set_title(
        f"{title_prefix}reflection = "
        f"{pars['mirror_reflection_random_angle'][0]:.5f}, "
        f"{pars['mirror_reflection_random_angle'][1]:.5f}, "
        f"{pars['mirror_reflection_random_angle'][2]:.5f}\n"
        f"align_vertical = {pars['mirror_align_random_vertical'][0]:.5f}, "
        f"{pars['mirror_align_random_vertical'][1]:.5f}, "
        f"{pars['mirror_align_random_vertical'][2]:.5f}, "
        f"{pars['mirror_align_random_vertical'][3]:.5f}\n"
        f"align_horizontal = {pars['mirror_align_random_horizontal'][0]:.5f}, "
        f"{pars['mirror_align_random_horizontal'][1]:.5f}, "
        f"{pars['mirror_align_random_horizontal'][2]:.5f}, "
        f"{pars['mirror_align_random_horizontal'][3]:.5f}"
    )

    d80_color = "red" if is_best else "black"
    d80_weight = "bold" if is_best else "normal"

    ax.text(
        0.5,
        0.3,
        f"D80 = {d80:.5f} cm\nRMSD = {rmsd:.4f}",
        verticalalignment="center",
        horizontalalignment="left",
        transform=ax.transAxes,
        color=d80_color,
        weight=d80_weight,
        bbox={"boxstyle": "round,pad=0.3", "facecolor": "yellow", "alpha": 0.7}
        if is_best
        else None,
    )

    if is_best:
        fig.text(
            0.02,
            0.02,
            "* Best parameter set (lowest RMSD)",
            fontsize=8,
            style="italic",
            color="red",
        )

    pdf_pages.savefig(fig, bbox_inches="tight")
    plt.clf()

    if original_simulated is not None:
        data_to_plot["simulated"] = original_simulated


def _create_all_plots(results, best_pars, data_to_plot, pdf_pages):
    """
    Create plots for all parameter sets if requested.

    Parameters
    ----------
    results : list
        List of (pars, rmsd, d80, simulated_data) tuples
    best_pars : dict
        Best parameter set for highlighting
    data_to_plot : dict
        Data dictionary for plotting
    pdf_pages : PdfPages
        PDF pages object to save plots
    """
    logger.info("Creating plots for all parameter sets...")

    for i, (pars, rmsd, d80, simulated_data) in enumerate(results):
        is_best = pars is best_pars
        logger.info(f"Creating plot {i + 1}/{len(results)}{' (BEST)' if is_best else ''}")

        _create_plot_for_parameters(
            pars, rmsd, d80, simulated_data, data_to_plot, is_best, pdf_pages
        )


def find_best_parameters(
    all_parameters, tel_model, site_model, args_dict, data_to_plot, radius, pdf_pages=None
):
    """
    Find the best parameters by running simulations for all parameter sets.

    Loop over all parameter sets, run the simulation, compute RMSD,
    and return the best parameters and their RMSD.
    """
    best_rmsd = float("inf")
    best_pars = None
    best_d80 = None
    results = []  # Store (pars, rmsd, d80, simulated_data)

    logger.info(f"Running {len(all_parameters)} simulations...")

    for i, pars in enumerate(all_parameters):
        try:
            logger.info(f"Running simulation {i + 1}/{len(all_parameters)}")
            d80, rmsd, simulated_data = run_psf_simulation(
                tel_model,
                site_model,
                args_dict,
                pars,
                data_to_plot,
                radius,
                return_simulated_data=True,
                pdf_pages=None,
            )
        except (ValueError, RuntimeError) as e:
            logger.warning(f"Simulation failed for parameters {pars}: {e}")
            continue

        results.append((pars, rmsd, d80, simulated_data))
        if rmsd < best_rmsd:
            best_rmsd = rmsd
            best_pars = pars
            best_d80 = d80

    logger.info(f"Best RMSD found: {best_rmsd:.5f}")

    # Create all plots if requested
    if pdf_pages is not None and args_dict.get("plot_all", False) and results:
        _create_all_plots(results, best_pars, data_to_plot, pdf_pages)

    return best_pars, best_d80, results


def create_d80_vs_offaxis_plot(tel_model, site_model, args_dict, best_pars, output_dir):
    """
    Create D80 vs off-axis angle plot using the best parameters.

    Parameters
    ----------
    tel_model : TelescopeModel
        Telescope model object.
    site_model : SiteModel
        Site model object.
    args_dict : dict
        Dictionary containing parsed command-line arguments.
    best_pars : dict
        Best parameter set.
    output_dir : Path
        Output directory for saving plots.
    """
    logger.info("Creating D80 vs off-axis angle plot with best parameters...")

    # Apply best parameters to telescope model
    tel_model.change_multiple_parameters(**best_pars)

    # Create off-axis angle array
    max_offset = args_dict.get("max_offset", MAX_OFFSET_DEFAULT)
    offset_steps = args_dict.get("offset_steps", OFFSET_STEPS_DEFAULT)
    off_axis_angles = np.linspace(
        0,
        max_offset,
        int(max_offset / offset_steps) + 1,
    )

    ray = RayTracing(
        telescope_model=tel_model,
        site_model=site_model,
        simtel_path=args_dict["simtel_path"],
        zenith_angle=args_dict["zenith"] * u.deg,
        source_distance=args_dict["src_distance"] * u.km,
        off_axis_angle=off_axis_angles * u.deg,
    )

    logger.info(f"Running ray tracing for {len(off_axis_angles)} off-axis angles...")
    ray.simulate(test=args_dict.get("test", False), force=True)
    ray.analyze(force=True)

    for key in ["d80_cm", "d80_deg"]:
        plt.figure(figsize=(10, 6), tight_layout=True)

        ray.plot(key, marker="o", linestyle="-", color="blue", linewidth=2, markersize=6)

        plt.title(
            f"PSF D80 vs Off-axis Angle - {tel_model.name}\n"
            f"Best Parameters: \n"
            f"reflection=[{best_pars['mirror_reflection_random_angle'][0]:.4f},"
            f"{best_pars['mirror_reflection_random_angle'][1]:.4f},"
            f"{best_pars['mirror_reflection_random_angle'][2]:.4f}],\n"
            f"align_horizontal={best_pars['mirror_align_random_horizontal'][0]:.4f}\n"
            f"align_vertical={best_pars['mirror_align_random_vertical'][0]:.4f}\n"
        )
        plt.xlabel("Off-axis Angle (degrees)")
        plt.ylabel("D80 (cm)" if key == "d80_cm" else "D80 (degrees)")
        plt.ylim(bottom=0)
        plt.xticks(rotation=45)
        plt.xlim(0, max_offset)
        plt.grid(True, alpha=0.3)

        plot_file_name = f"tune_psf_{tel_model.name}_best_params_{key}.pdf"
        plot_file = output_dir.joinpath(plot_file_name)
        visualize.save_figure(plt, plot_file, log_title=f"D80 vs off-axis ({key})")

    plt.close("all")


def write_tested_parameters_to_file(results, best_pars, best_d80, output_dir, tel_model):
    """
    Write all tested parameters and their metrics to a text file.

    Parameters
    ----------
    results : list
        List of (pars, rmsd, d80, simulated_data) tuples
    best_pars : dict
        Best parameter set
    best_d80 : float
        Best D80 value
    output_dir : Path
        Output directory path
    tel_model : TelescopeModel
        Telescope model object for filename generation
    """
    param_file = output_dir.joinpath(f"psf_optimization_{tel_model.name}.log")
    with open(param_file, "w", encoding="utf-8") as f:
        f.write("# PSF Parameter Optimization Log\n")
        f.write(f"# Telescope: {tel_model.name}\n")
        f.write(f"# Total parameter sets tested: {len(results)}\n")
        f.write("#" + "=" * 60 + "\n\n")

        f.write("PARAMETER TESTING RESULTS:\n")
        for i, (pars, rmsd, d80, _) in enumerate(results):
            is_best = pars is best_pars
            status = "BEST" if is_best else "TESTED"
            f.write(f"[{status}] Set {i + 1:03d}: RMSD={rmsd:.5f}, D80={d80:.5f} cm\n")
            for par, value in pars.items():
                f.write(f"    {par}: {value}\n")
            f.write("\n")

        f.write("OPTIMIZATION SUMMARY:\n")
        f.write(f"Best RMSD: {min(result[1] for result in results):.5f}\n")
        f.write(f"Best D80: {best_d80:.5f} cm\n")
        f.write("\nOPTIMIZED PARAMETERS:\n")
        for par, value in best_pars.items():
            f.write(f"{par}: {value}\n")
    return param_file


def _add_units_to_psf_parameters(best_pars):
    """
    Add proper astropy units to PSF parameters based on their schemas.

    Parameters
    ----------
    best_pars : dict
        Dictionary with PSF parameter names as keys and values as lists

    Returns
    -------
    dict
        Dictionary with same keys but values converted to astropy quantities with units
    """
    psf_pars_with_units = {}

    for param_name, param_values in best_pars.items():
        if param_name == "mirror_reflection_random_angle":
            psf_pars_with_units[param_name] = [
                param_values[0] * u.deg,
                param_values[1] * u.dimensionless_unscaled,
                param_values[2] * u.deg,
            ]
        elif param_name in ["mirror_align_random_horizontal", "mirror_align_random_vertical"]:
            psf_pars_with_units[param_name] = [
                param_values[0] * u.deg,
                param_values[1] * u.deg,
                param_values[2] * u.dimensionless_unscaled,
                param_values[3] * u.dimensionless_unscaled,
            ]
        else:
            psf_pars_with_units[param_name] = param_values

    return psf_pars_with_units


def export_psf_parameters(best_pars, tel_model, parameter_version, output_dir):
    """
    Export PSF parameters as simulation model parameter files.

    Parameters
    ----------
    best_pars : dict
        Best parameter set
    tel_model : TelescopeModel
        Telescope model object
    parameter_version : str
        Parameter version string
    output_dir : Path
        Output directory path
    """
    try:
        logger.info("Exporting best PSF parameters as simulation model parameter files")
        psf_pars_with_units = _add_units_to_psf_parameters(best_pars)
        parameter_output_path = output_dir / tel_model.name
        for parameter_name, parameter_value in psf_pars_with_units.items():
            writer.ModelDataWriter.dump_model_parameter(
                parameter_name=parameter_name,
                value=parameter_value,
                instrument=tel_model.name,
                parameter_version=parameter_version,
                output_file=f"{parameter_name}-{parameter_version}.json",
                output_path=parameter_output_path,
                use_plain_output_path=True,
            )
        logger.info(f"simulation model parameter files exported to {output_dir}")
    except ImportError as e:
        logger.warning(f"Could not export simulation parameters: {e}")
    except (ValueError, KeyError, OSError) as e:
        logger.error(f"Error exporting simulation parameters: {e}")


def run_psf_optimization_workflow(tel_model, site_model, args_dict, output_dir):
    """
    Run the complete PSF parameter optimization workflow.

    This function consolidates the main optimization logic to make the application lighter.

    Parameters
    ----------
    tel_model : TelescopeModel
        Telescope model object
    site_model : SiteModel
        Site model object
    args_dict : dict
        Dictionary containing parsed command-line arguments
    output_dir : Path
        Output directory path

    Returns
    -------
    None
        All results are saved to files and printed to console
    """
    # Generate parameter sets
    all_parameters = []
    mrra_0, mfr_0, mrra2_0, mar_0 = get_previous_values(tel_model)

    n_runs = args_dict.get("n_runs")
    generate_random_parameters(
        all_parameters, n_runs, args_dict, mrra_0, mfr_0, mrra2_0, mar_0, tel_model
    )

    data_to_plot, radius = load_and_process_data(args_dict)

    # Preparing figure name and PDF pages for plotting
    plot_file_name = "_".join(("tune_psf", tel_model.name + ".pdf"))
    plot_file = output_dir.joinpath(plot_file_name)
    pdf_pages = PdfPages(plot_file)

    # Find best parameters
    best_pars, best_d80, results = find_best_parameters(
        all_parameters, tel_model, site_model, args_dict, data_to_plot, radius, pdf_pages
    )

    plt.close()
    pdf_pages.close()

    # Write all tested parameters and their metrics to a file
    param_file = write_tested_parameters_to_file(
        results, best_pars, best_d80, output_dir, tel_model
    )
    print(f"\nParameter results written to {param_file}")

    # Automatically create D80 vs off-axis angle plot for best parameters
    create_d80_vs_offaxis_plot(tel_model, site_model, args_dict, best_pars, output_dir)
    print("D80 vs off-axis angle plots created successfully")

    print("\nBest parameters:")
    for par, value in best_pars.items():
        print(f"{par} = {value}")

    # Export best parameters as simulation model parameter files (if flag is provided)
    if args_dict.get("write_psf_parameters", False):
        export_psf_parameters(
            best_pars,
            tel_model,
            args_dict.get("parameter_version", "0.0.0"),
            output_dir.parent,
        )
