#!/usr/bin/python3

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

    This includes parameters mirror_reflection_random_angle, \
    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.
    fixed (activation mode, optional)
        Keep the first entry of mirror_reflection_random_angle fixed.
    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 6.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 6.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_reflection_random_angle, 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(
        "--fixed",
        help=("Keep the first entry of mirror_reflection_random_angle fixed."),
        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.01,
    )
    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"),
        db_config=app_context.db_config,
        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()
