#!/usr/bin/python3

r"""
    Derives the mirror alignment parameters using cumulative PSF measurement.

    This includes parameters mirror_align_random_horizontal and mirror_align_random_vertical.

    The measured cumulative PSF should be provided by using the command line argument data. \
    A file name is expected, in which the file should contain 3 columns: radial distance in mm, \
    differential value of photon intensity and its integral value.

    The derivation is performed through gradient descent optimization that minimizes either the \
    Root Mean Squared Deviation (RMSD) between measured and simulated PSF curves (default) or the \
    Kolmogorov-Smirnov (KS) statistic when the --ks_statistic flag is used.

    The optimization workflow includes:

    * Loading and preprocessing PSF data from measurement files
    * Running gradient descent optimization to minimize RMSD
    * Generating cumulative PSF plots for each iteration showing optimization progression
    * Logging parameter evolution through gradient descent steps
    * Creating convergence plots showing RMSD and D80 evolution
    * Automatically generating D80 vs off-axis angle analysis for best parameters
    * Optionally exporting optimized parameters as simulation model files

    The assumption are:

    a) mirror_align_random_horizontal and mirror_align_random_vertical are the same.

    b) mirror_align_random_horizontal/vertical have no dependence on the zenith angle.

    One example of the plot generated by this applications are shown below.

    .. _derive_psf_parameters_plot:
    .. image::  images/gradient_descent.png
      :width: 49 %

    Command line arguments
    ----------------------
    site (str, required)
        North or South.
    telescope (str, required)
        Telescope model name (e.g. LST-1, SST-D, ...).
    model_version (str, optional)
        Model version.
    parameter_version (str, optional)
        Parameter version for model parameter file export.
    src_distance (float, optional)
        Source distance in km.
    zenith (float, optional)
        Zenith angle in deg.
    data (str, optional)
        Name of the data file with the measured cumulative PSF.
    plot_all (activation mode, optional)
        If activated, plots will be generated for all values tested during tuning.
    test (activation mode, optional)
        If activated, application will be faster by simulating fewer photons.
    write_psf_parameters (activation mode, optional)
        Write the optimized PSF parameters as simulation model parameter files.
    rmsd_threshold (float, optional)
        RMSD threshold for gradient descent convergence (default: 0.007).
    learning_rate (float, optional)
        Learning rate for gradient descent optimization (default: 0.01).
    monte_carlo_analysis (activation mode, optional)
        Run Monte Carlo analysis to find statistical uncertainties.

    Example
    -------
    --telescope LSTN-01 --model_version 6.0.0



    Run the application:

    .. code-block:: console

        simtools-derive-psf-parameters --site North --telescope LSTN-01 \\
            --model_version 6.0.0 --data tests/resources/PSFcurve_data_v2.ecsv --plot_all --test

    Run with parameter export:

    .. code-block:: console

        simtools-derive-psf-parameters --site North --telescope LSTN-01 --model_version 7.0.0 \\
            --plot_all --test --rmsd_threshold 0.01 --learning_rate 0.001 \\
            --data tests/resources/PSFcurve_data_v2.ecsv \\
            --write_psf_parameters

    Run monte carlo analysis:

    .. code-block:: console

        simtools-derive-psf-parameters --site North --telescope LSTN-01 --model_version 7.0.0 \\
            --plot_all --test --monte_carlo_analysis \\
            --data tests/resources/PSFcurve_data_v2.ecsv \\
            --write_psf_parameters

    The output is saved in simtools-output/derive_psf_parameters.

    Output files include:

    * Gradient descent progression log in psf_gradient_descent_[telescope].log
    * Gradient descent convergence plots in gradient_descent_convergence_[telescope].png
    * PSF progression plots showing evolution through iterations (if --plot_all is specified)
    * D80 vs off-axis angle plots (d80_vs_offaxis_cm.png, d80_vs_offaxis_deg.png)
    * Optimized simulation model parameter files (if --write_psf_parameters is specified)

"""

from simtools.application_control import get_application_label, startup_application
from simtools.configuration import configurator
from simtools.model.model_utils import initialize_simulation_models
from simtools.ray_tracing import psf_parameter_optimisation as psf_opt


def _parse():
    config = configurator.Configurator(
        label=get_application_label(__file__),
        description=(
            "Derive mirror_align_random_horizontal and mirror_align_random_vertical "
            "using cumulative PSF measurement."
        ),
    )
    config.parser.add_argument(
        "--src_distance",
        help="Source distance in km",
        type=float,
        default=10,
    )
    config.parser.add_argument("--zenith", help="Zenith angle in deg", type=float, default=20)
    config.parser.add_argument(
        "--data", help="Data file name with the measured PSF vs radius [cm]", type=str
    )
    config.parser.add_argument(
        "--plot_all",
        help=(
            "On: plot cumulative PSF for all tested combinations, "
            "Off: plot it only for the best set of values"
        ),
        action="store_true",
    )
    config.parser.add_argument(
        "--write_psf_parameters",
        help=("Write the optimized PSF parameters as simulation model parameter files"),
        action="store_true",
        required=False,
    )
    config.parser.add_argument(
        "--rmsd_threshold",
        help=(
            "RMSD threshold for gradient descent convergence "
            "(not used with --monte_carlo_analysis)."
        ),
        type=float,
        default=0.01,
    )
    config.parser.add_argument(
        "--learning_rate",
        help=(
            "Learning rate for gradient descent optimization "
            "(not used with --monte_carlo_analysis)."
        ),
        type=float,
        default=0.0001,
    )
    config.parser.add_argument(
        "--monte_carlo_analysis",
        help="Run analysis to find monte carlo uncertainties.",
        action="store_true",
    )
    config.parser.add_argument(
        "--ks_statistic",
        help="Use KS statistic for monte carlo uncertainty analysis.",
        action="store_true",
    )
    config.parser.add_argument(
        "--fraction",
        help="PSF containment fraction for diameter calculation (e.g., 0.8 for D80, 0.95 for D95).",
        type=float,
        default=0.8,
    )
    config.parser.add_argument(
        "--cleanup",
        help="Remove intermediate *.log and *.lis* files after optimization.",
        action="store_true",
    )
    return config.initialize(
        db_config=True,
        simulation_model=["telescope", "model_version", "parameter_version"],
    )


def main():
    """Derive PSF parameters."""
    app_context = startup_application(_parse)

    tel_model, site_model, _ = initialize_simulation_models(
        label=app_context.args.get("label"),
        site=app_context.args["site"],
        telescope_name=app_context.args["telescope"],
        model_version=app_context.args["model_version"],
    )

    psf_opt.run_psf_optimization_workflow(
        tel_model,
        site_model,
        app_context.args,
        app_context.io_handler.get_output_directory(),
    )


if __name__ == "__main__":
    main()
