analysis

#!/usr/bin/env python3
"""
Analyzes particle decay events by generating signal and background data, computing factors such as Q-factors and sPlot weights, and visualizing the data through plots and fits. The available options enable customization of the dataset size, plot types, and factor calculations.

Usage:
    analysis [options]

Options:
    -h --help               Show this screen.
    --num-sig=<nsig>        Number of signal events to generate. [default: 10000]
    --num-bkg=<nbkg>        Number of background events to generate. [default: 10000]
    --parallel              Use parallel processing for event generation.
    --knn=<knn>             Number of nearest neighbors for kNN calculations. [default: 100]
    --density-knn           Compute kNN calculations based off on local density for each event
    --radius-knn=<radius>   Use radius-based neighbors calculations with specified radius. [default: None]
    --t-dep                 Use t-dependence in mass variable
    --num-iter=<niter>      Specifies the number of iterations to run the analysis. [default: 1]
    --seed=<seed>           Starting seed for all iterations. [default: 0]
    --append                Don't overwrite fit_results.tsv, create it if it doesn't exist or append if it does.
    -o --output=<outdir>    Specify output directory. [default: studies]
"""

from __future__ import annotations

from pathlib import Path
import sys
import os

import numpy as np
from docopt import docopt
from sqfactors import console, plot, set_seed, r
from sqfactors.analysis import (
    bounds,
    calculate_inplot,
    calculate_q_factors,
    calculate_sideband_weights,
    calculate_splot_weights,
    calculate_theoretical_q_factors,
    get_results,
    truths,
)
from sqfactors.event import (
    gen_bkg,
    gen_bkg_event_partial,
    gen_event_partial,
    gen_sig,
    parallel_event_generation,
)
from sqfactors.utils import Results

if __name__ == '__main__':
    args = docopt(__doc__)

    num_sig = int(args['--num-sig'])
    num_bkg = int(args['--num-bkg'])
    num_knn = int(args['--knn'])
    use_density_knn = args['--density-knn']
    use_radius_knn = args['--radius-knn']
    t_dep = args['--t-dep']
    num_iterations = int(args['--num-iter'])

    if use_radius_knn != 'None':
        try:
            use_radius_knn = float(use_radius_knn)
        except ValueError:
            msg = f'Invalid value for --radius_knn: {use_radius_knn}'
            raise ValueError(msg) from None
    else:
        use_radius_knn = None

    outdir = Path(args['--output'])
    directory = str(outdir / 'study')
    if use_density_knn:
        directory += '_density'
    if use_radius_knn:
        directory += '_radius'
    if t_dep:
        directory += '_t_dep'
    Path(directory).resolve().mkdir(parents=True, exist_ok=True)
    tsv_filename = 'fit_results.tsv'
    tsv_filepath = Path(directory) / tsv_filename
    if not args['--append'] or not tsv_filepath.is_file():
        tsv_filepath.write_text(Results.tsv_header())

    parallel = args['--parallel']

    # Determine if SLURM is being used
    running_in_slurm = 'SLURM_JOB_ID' in os.environ

    # Specify the iteration to run in SLURM from bash
    iteration_to_run = int(args['--num-iter']) if running_in_slurm else None
   
    if running_in_slurm and iteration_to_run is not None:
        iteration_range = [iteration_to_run]
    else:
        iteration_range = range(1, num_iterations + 1)

    for iteration in iteration_range:

        set_seed(int(args['--seed']))
        seed = r().integers(100000)
        set_seed(int(seed + iteration))
        it_dir_name = args['--seed'] + '_' + str(iteration)
        it_directory = Path(directory) / it_dir_name
        it_directory.resolve().mkdir(parents=True, exist_ok=True)
        if parallel:
            # Generate events in parallel.
            console.print('Generating signal and background events in parallel ...')
            events_sig = parallel_event_generation(gen_event_partial, n=num_sig, num_workers=4)
            events_bkg = parallel_event_generation(gen_bkg_event_partial, n=num_bkg, num_workers=4)
        else:
            # Default to sequential generation.
            console.print('Generating signal and background events sequentially ...')
            events_sig = gen_sig(n=num_sig)
            events_bkg = gen_bkg(n=num_bkg)

        with console.status('Plotting events'):
            plot.plot_all_events(
                events_sig, events_bkg, filename='all_events.png', directory=it_directory
            )
        events_all = events_sig + events_bkg

        # Define weights functions and descriptions for various analyses
        analysis_config = {
            'No Weights': {
                'weight_func': lambda _: np.ones(len(events_all)),
                'description': 'No Weights Analysis',
                'compare_q_factors_required': False,
            },
            'Sideband': {
                'weight_func': calculate_sideband_weights,
                'description': 'Sideband Subtraction Analysis',
                'compare_q_factors_required': False,
            },
            'InPlot': {
                'weight_func': calculate_inplot,
                'description': 'InPlot Analysis',
                'compare_q_factors_required': False,
            },
            'sPlot': {
                'weight_func': lambda events: calculate_splot_weights(events)[:, 0],
                'description': 'sPlot Analysis',
                'compare_q_factors_required': False,
            },
            'Q-Factor': {
                'weight_funcs': lambda events: calculate_q_factors(
                    events,
                    phase_space=np.array(
                        [[e.costheta / (2 / 3), e.phi / (2 * np.pi**3 / 3)] for e in events_all]
                    ),
                    name='angles',
                    num_knn=num_knn,
                    use_density_knn=use_density_knn,
                    use_radius_knn=use_radius_knn,
                    directory=str(it_directory),
                    plot_indices=[0, 1, 2, num_sig, num_sig + 1, num_sig + 2],
                ),
                'descriptions': ['Q-Factor Analysis', 'sQ-Factor Analysis'],
                'compare_q_factors_required': True,
            },
            'Q-Factor_t': {
                'weight_funcs': lambda events: calculate_q_factors(
                    events,
                    phase_space=np.array(
                        [
                            [
                                e.costheta / (2 / 3),
                                e.phi / (2 * np.pi**3 / 3),
                                e.t / ((bounds['t_max'] ** 3 - bounds['t_min'] ** 3) / 3),
                            ]
                            for e in events
                        ]
                    ),
                    name='angles_t',
                    num_knn=num_knn,
                    use_density_knn=use_density_knn,
                    use_radius_knn=use_radius_knn,
                    directory=it_directory,
                    plot_indices=[0, 1, 2, num_sig, num_sig + 1, num_sig + 2],
                ),
                'descriptions': ['Q-Factor Analysis (with t)', 'sQ-Factor Analysis (with t)'],
                'compare_q_factors_required': True,
            },
            'Q-Factor_g': {
                'weight_funcs': lambda events: calculate_q_factors(
                    events,
                    phase_space=np.array(
                        [
                            [
                                e.costheta / (2 / 3),
                                e.phi / (2 * np.pi**3 / 3),
                                e.g / ((bounds['g_max'] ** 3 - bounds['g_min'] ** 3) / 3),
                            ]
                            for e in events
                        ]
                    ),
                    name='angles_g',
                    num_knn=num_knn,
                    use_density_knn=use_density_knn,
                    use_radius_knn=use_radius_knn,
                    directory=it_directory,
                    plot_indices=[0, 1, 2, num_sig, num_sig + 1, num_sig + 2],
                ),
                'descriptions': ['Q-Factor Analysis (with g)', 'sQ-Factor Analysis (with g)'],
                'compare_q_factors_required': True,
            },
            'Q-Factor_t_g': {
                'weight_funcs': lambda events: calculate_q_factors(
                    events,
                    phase_space=np.array(
                        [
                            [
                                e.costheta / (2 / 3),
                                e.phi / (2 * np.pi**3 / 3),
                                e.t / ((bounds['t_max'] ** 3 - bounds['t_min'] ** 3) / 3),
                                e.g / ((bounds['g_max'] ** 3 - bounds['g_min'] ** 3) / 3),
                            ]
                            for e in events_all
                        ]
                    ),
                    name='angles_t_g',
                    num_knn=num_knn,
                    use_density_knn=use_density_knn,
                    use_radius_knn=use_radius_knn,
                    directory=it_directory,
                    plot_indices=[0, 1, 2, num_sig, num_sig + 1, num_sig + 2],
                ),
                'descriptions': [
                    'Q-Factor Analysis (with t & g)',
                    'sQ-Factor Analysis (with t & g)',
                ],
                'compare_q_factors_required': True,
            },
        }

        # Directory for storing results
        results_dir = Path(directory) / 'results'
        results_dir.mkdir(exist_ok=True)

        # Handle each iteration
        console.print(f'Starting Iteration {iteration}', style='bold yellow')
        results = Results()

        # Inside each iteration, process each analysis type
        for analysis_name, config in analysis_config.items():
            console.log(f'Processing {analysis_name} for iteration {iteration}...')

            # Calculate weights for this analysis type
            if (w_func := config.get('weight_func')) is not None:
                weights = w_func(events_all)
                sq_weights = None
            elif (w_funcs := config.get('weight_funcs')) is not None:
                weights, sq_weights = w_funcs(events_all)
            else:
                msg = f'Weights not defined for {analysis_name}'
                raise Exception(msg)

            # Save row to outputs
            if sq_weights is not None:
                results.add_row(
                    get_results(config['descriptions'][0], it_dir_name, events_all, weights)
                )
                plot.plot_events(
                    events_all,
                    events_sig,
                    weights=weights,
                    filename=f'events_{analysis_name}_{iteration}.png',
                    directory=it_directory,
                )
                if config.get('compare_q_factors_required', True):
                    # Theoretical model remains constant across variants
                    q_factors_theoretical = calculate_theoretical_q_factors(events_all, truths['b'])
                    plot.compare_q_factors(
                        weights,
                        q_factors_theoretical,
                        title=f"{config['descriptions'][0]} Comparison",
                        q_factor_type=analysis_name.lower().replace(' ', '_'),
                        directory=it_directory,
                    )
                results.add_row(
                    get_results(config['descriptions'][1], it_dir_name, events_all, sq_weights),
                )
                plot.plot_events(
                    events_all,
                    events_sig,
                    weights=sq_weights,
                    filename=f'events_s{analysis_name}_{iteration}.png',
                    directory=it_directory,
                )
                if config.get('compare_q_factors_required', True):
                    # Theoretical model remains constant across variants
                    q_factors_theoretical = calculate_theoretical_q_factors(events_all, truths['b'])
                    plot.compare_q_factors(
                        sq_weights,
                        q_factors_theoretical,
                        title=f"{config['descriptions'][1]} Comparison",
                        q_factor_type='s' + analysis_name.lower().replace(' ', '_'),
                directory=it_directory,
                    )
            else:
                results.add_row(
                    get_results(config['description'], it_dir_name, events_all, weights)
                )
                plot.plot_events(
                    events_all,
                    events_sig,
                    weights=weights,
                    filename=f'events_{analysis_name}_{iteration}.png',
                    directory=it_directory,
                )
                if config.get('compare_q_factors_required', True):
                    # Theoretical model remains constant across variants
                    q_factors_theoretical = calculate_theoretical_q_factors(events_all, truths['b'])
                    plot.compare_q_factors(
                        weights,
                        q_factors_theoretical,
                        title=f"{config['description']} Comparison",
                        q_factor_type=analysis_name.lower().replace(' ', '_'),
                        directory=it_directory,
                    )

        console.print(results)
        with open(tsv_filepath, 'a') as tsv:
            tsv.write(results.as_tsv(header=False))

        if use_radius_knn:
            selected_event_index = 0  # Index of the event you want to inspect
            plot.plot_radius_knn_visualization(
                events_all, selected_event_index, use_radius_knn, directory=it_directory
            )