#!/usr/bin/env python3
# PYTHON_ARGCOMPLETE_OK

# Copyright (c) 2020-2025 Danny van Dyk
#
# This file is part of the EOS project. EOS is free software;
# you can redistribute it and/or modify it under the terms of the GNU General
# Public License version 2, as published by the Free Software Foundation.
#
# EOS is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 59 Temple
# Place, Suite 330, Boston, MA  02111-1307  USA

import argparse, argcomplete
import eos
from eos import debug, info, warn, error
import logging
import numpy as _np
import os
import pypmc
import scipy
import sys
import traceback
import yaml

try:
    from termcolor import colored
except ImportError:
    colored = lambda s, *args, **kwargs: s

# return the value of the environment variable, or a default value if the variable is unset.
def get_from_env(envvar, default):
    if not envvar in os.environ:
        return default
    if envvar == "EOS_VERBOSITY":
        return int(os.environ[envvar])

    return os.environ[envvar]


def _parser():
    parser = argparse.ArgumentParser(description='Carry out a Bayesian analysis using EOS')
    # 'parent' parser for common arguments
    common_subparser = argparse.ArgumentParser(add_help=False)
    # add verbosity arg and analysis-file arg to all commands
    common_subparser.add_argument('-v', '--verbose',
        help = 'Increases the verbosity of the script. Can also be set via the EOS_VERBOSITY environment variable.',
        dest = 'verbose', action = 'count', default = None
    )
    common_subparser.add_argument('-f', '--analysis-file',
        help = 'The analysis file. Defaults to \'analysis.yaml\'.',
        dest = 'analysis_file', action = 'store', default = 'analysis.yaml'
    )
    subparsers = parser.add_subparsers(title = 'commands')

    ## begin of commands

    # list-priors
    parser_list_priors = subparsers.add_parser('list-priors',
        parents = [common_subparser],
        description =
'''
Lists the named prior PDFs defined within the scope of this analysis file.
''',
        help = 'Lists the named priors defined in the analysis file.'
    )
    parser_list_priors.set_defaults(cmd = cmd_list_priors)


    # list-likelihoods
    parser_list_likelihoods = subparsers.add_parser('list-likelihoods',
        parents = [common_subparser],
        description =
'''
Lists the named likelihoods defined within the scope of this analysis file.
''',
        help = 'Lists the named likelihoods defined in the analysis file.'
    )
    parser_list_likelihoods.add_argument('-d', '--display-details',
        help = 'Whether to display further details for each likelihood.',
        dest = 'display', action = 'store_true', default = False
    )
    parser_list_likelihoods.set_defaults(cmd = cmd_list_likelihoods)


    # list-posteriors
    parser_list_posteriors = subparsers.add_parser('list-posteriors',
        parents = [common_subparser],
        description =
'''
Lists the named posterior PDFs defined within the scope of this analysis file.
''',
        help = 'Lists the named posteriors defined in the analysis file.'
    )
    parser_list_posteriors.set_defaults(cmd = cmd_list_posteriors)


    # list-predictions
    parser_list_predictions = subparsers.add_parser('list-predictions',
        parents = [common_subparser],
        description =
'''
Lists the named prediction sets defined within the scope of this analysis file.
''',
        help = 'Lists the named predictions define in the analysis file.'
    )
    parser_list_predictions.set_defaults(cmd = cmd_list_predictions)


    # sample-mcmc
    parser_sample_mcmc = subparsers.add_parser('sample-mcmc',
        parents = [common_subparser],
        description =
'''
Samples from a named posterior PDF using Markov Chain Monte Carlo (MCMC) methods.

The output file will be stored in EOS_BASE_DIRECTORY/POSTERIOR/mcmc-IDX.
''',
        help = 'Samples from a posterior using Marko Chain Monte Carlo methods.'
    )
    parser_sample_mcmc.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which to draw the samples.'
    )
    parser_sample_mcmc.add_argument('chain', metavar = 'CHAIN-IDX',
        help = 'The index assigned to the Markov chain. This value is used to seed the RNG for a reproducible analysis.',
        type = int
    )
    parser_sample_mcmc.add_argument('-N', '--number-of-samples',
        help = 'The number of samples to be stored in the output file.',
        dest = 'N', action = 'store', type = int, default = 1000
    )
    parser_sample_mcmc.add_argument('-S', '--stride',
        help = 'The ratio of samples drawn over samples stored. For every S samples, S - 1 will be discarded.',
        dest = 'stride', action = 'store', type = int, default = 5
    )
    parser_sample_mcmc.add_argument('-p', '--number-of-preruns',
        help = 'The number of prerun steps, which are used to adapt the MCMC proposal to the posterior.',
        dest = 'preruns', action = 'store', type = int, default = 3
    )
    parser_sample_mcmc.add_argument('-n', '--number-of-prerun-samples',
        help = 'The number of samples to be used for an adaptation in each prerun steps. These samples will be discarded.',
        dest = 'pre_N', action = 'store', type = int, default = 150
    )
    parser_sample_mcmc.add_argument('-s', '--start-point',
        help = 'Optional starting point for the chain; a comma separated list of floats',
        dest = 'start_point', action = 'store', type = lambda s: [float(item) for item in s.split(',')]
    )
    parser_sample_mcmc.add_argument('-c', '--cov-scale',
        help = 'Scale factor for the initial guess of the covariance matrix.',
        dest = 'cov_scale', action = 'store', type = float, default = 0.1
    )
    parser_sample_mcmc.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_sample_mcmc.set_defaults(cmd = cmd_sample_mcmc)


    # sample-pmc
    parser_sample_pmc = subparsers.add_parser('sample-pmc',
        parents = [common_subparser],
        description =
'''
Samples from a named posterior using the Population Monte Carlo (PMC) methods.

The results of the find-cluster command are expected in EOS_BASE_DIRECTORY/POSTERIOR/clusters.
The output will be stored in EOS_BASE_DIRECTORY/POSTERIOR/pmc.
In addition, an ImportanceSamples object is exported to EOS_BASE_DIRECTORY/POSTERIOR/samples.
''',
        help = 'Samples from a posterior using the Population Monte Carlo method.'
    )
    parser_sample_pmc.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which to draw the samples.'
    )
    parser_sample_pmc.add_argument('-n', '--number-of-adaptation-samples',
        help = 'The number of samples to be used in each adaptation step. These samples will be discarded.',
        dest = 'step_N', action = 'store', type = int, default = 500
    )
    parser_sample_pmc.add_argument('-s', '--number-of-adaptation-steps',
        help = 'The number of adaptation steps, which are used to adapt the PMC proposal to the posterior.',
        dest = 'steps', action = 'store', type = int, default = 10
    )
    parser_sample_pmc.add_argument('-t', '--perplexity-threshold',
        help = 'The threshold for the perplexity in the last step after which further adaptation steps are to be skipped.',
        dest = 'perplexity_threshold', action = 'store', type = float, default = 1.0
    )
    parser_sample_pmc.add_argument('-w', '--weight-threshold',
        help = 'Mixture components with a weight smaller than this threshold are pruned.',
        dest = 'weight_threshold', action = 'store', type = float, default = 1e-10
    )
    parser_sample_pmc.add_argument('-N', '--number-of-final-samples',
        help = 'The number of samples to be stored in the output file.',
        dest = 'final_N', action = 'store', type = int, default = 5000
    )
    parser_sample_pmc.add_argument('--pmc_iterations',
        help = 'The maximum number of update of the PMC, changing this value may make the update unstable',
        dest = 'pmc_iterations', action = 'store', type = int, default = 1
    )
    parser_sample_pmc.add_argument('--pmc-rel-tol',
        help = 'Relative tolerance of the PMC update.',
        dest = 'pmc_rel_tol', action = 'store', type = float, default = 1e-10
    )
    parser_sample_pmc.add_argument('--pmc-abs-tol',
        help = 'Absolute tolerance of the PMC update.',
        dest = 'pmc_abs_tol', action = 'store', type = float, default = 1e-5
    )
    parser_sample_pmc.add_argument('-l', '--pmc-lookback',
        help = 'The number of previous steps whose samples are used to update the PMC. 1 (default) means that only last step is used ; 0 means that all available steps are used.',
        dest = 'pmc_lookback', action = 'store', type = int, default = 1
    )
    parser_sample_pmc.add_argument('-p', '--initial-proposal',
        help = """Specify where the initial proposal should be taken from; 'clusters' (default): use the proposal obtained using `find-clusters`;
            'product': use the proposal obtained from `mixture_product`; 'pmc': continue sampling from the previous `sample-pmc` results.""",
        dest = 'initial_proposal', action = 'store', type = str, default = 'clusters'
    )
    parser_sample_pmc.add_argument('-S', '--sigma-test-stat',
        help = 'If provided, the inverse CDF of -2*log(PDF) will be evaluated, using the provided values as the respective significance.',
        dest = 'sigma_test_stat', action = 'store', type = lambda s: [float(item) for item in s.split(',')]
    )
    parser_sample_pmc.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_sample_pmc.set_defaults(cmd = cmd_sample_pmc)


    # sample-nested
    parser_sample_nested = subparsers.add_parser('sample-nested',
        parents = [common_subparser],
        description =
'''
Samples from a likelihood associated with a named posterior using dynamic nested sampling.

The output will be stored in EOS_BASE_DIRECTORY/POSTERIOR/nested.
In addition, an ImportanceSamples object is exported to EOS_BASE_DIRECTORY/POSTERIOR/samples.
''',
        help = 'Samples from a likelihood (associated with a posterior) using dynamic nested sampling.'
    )
    parser_sample_nested.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which to draw the samples.'
    )
    parser_sample_nested.add_argument('-B', '--target-bound',
        help = 'The method used to approximately bound the prior using the current set of live points. Conditions the sampling methods used to propose new live points. For valid values, see the dynesty documentation.',
        dest = 'bound', action = 'store', type = str, default = 'multi'
    )
    parser_sample_nested.add_argument('-n', '--number-of-live-points',
        help = 'The number of live points.',
        dest = 'nlive', action = 'store', type = int, default = 250
    )
    parser_sample_nested.add_argument('-d', '--evidence-tolerance',
        help = 'The relative tolerance for the remaining evidence.',
        dest = 'dlogz', action = 'store', type = float, default = 1.0
    )
    parser_sample_nested.add_argument('-m', '--max-number-iterations',
        help = 'The maximum number of iterations. Iterations may stop earlier if the termination condition is reached.',
        dest = 'maxiter', action = 'store', type = int, default = None
    )
    parser_sample_nested.add_argument('-l', '--min-number-iterations',
        help = 'The minimum number of iterations. If not provided, the sampler will run until the termination condition is reached.',
        dest = 'miniter', action = 'store', type = int, default = 0
    )
    parser_sample_nested.add_argument('-s', '--use-random-seed',
        help = 'The seed used to initialize the Mersenne Twister pseudo-random number generator.',
        dest = 'seed', action = 'store', type = int, default = 10
    )
    parser_sample_nested.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_sample_nested.add_argument('-M', '--sampling-method',
        help = 'The method used for sampling within the likelihood constraints. For valid values, see dynesty documentation.',
        dest = 'sample', action = 'store', type = str, default = 'auto'
    )
    parser_sample_nested.set_defaults(cmd = cmd_sample_nested)


    # plot-samples
    parser_plot_samples = subparsers.add_parser('plot-samples',
        parents = [common_subparser],
        description =
'''
Plots all samples obtained for a named posterior.

The results of either the sample-mcmc or the sample-pmc command are expected in
EOS_BASE_DIRECTORY/POSTERIOR/mcmc-* or EOS_BASE_DIRECTORY/POSTERIOR/pmc, respectively.
The plots will be stored as PDF files within the respective sample inputs.
''',
        help = 'Plots samples for a named posterior.'
    )
    parser_plot_samples.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which to draw the samples.'
    )
    parser_plot_samples.add_argument('-B', '--bins',
        help = 'The number of bins per histogram.',
        dest = 'bins', action = 'store', type = int, default = 50
    )
    parser_plot_samples.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_plot_samples.set_defaults(cmd = cmd_plot_samples)


    # find-mode
    parser_find_mode = subparsers.add_parser('find-mode',
        parents = [common_subparser],
        description =
'''
Finds the mode of the named posterior.

The optimization process can be initialised either by a provided parameter point or previously obtained
importance samples (recommended), a random point (default), or from a pre-existing Markov chain with provided index (deprecated). The optimization
can be iterated to increase the accuracy of the result. The output will be stored in
EOS_BASE_DIRECTORY/posterior/mode-LABEL.
''',
        help = 'Finds the mode of a named posterior.'
    )
    parser_find_mode.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF that will be maximized.',
    )
    parser_find_mode.add_argument('-o', '--optimizations',
        help = 'The number of calls to the optimization algorithm.',
        dest = 'optimizations', action = 'store', type = int, default = 3
    )
    parser_find_mode.add_argument('-c', '--from-mcmc',
        help = 'The chain index of an MCMC data file from which the maximization is started, this is now deprecated.',
        dest = 'chain', action = 'store', type = int, default = None
    )
    parser_find_mode.add_argument('-S', '--from-samples',
        help = 'If specified, the optimization will start from the importance samples stored in posterior/samples.',
        dest = 'importance_samples', action = 'store_true', default = None
    )
    parser_find_mode.add_argument('-p', '--from-point',
        help = 'The point from which the minimization is started; a comma separated list of floats.',
        dest = 'start_point', action = 'store', type = lambda s: [float(item) for item in s.split(',')]
    )
    parser_find_mode.add_argument('-s', '--use-random-seed',
        help = 'The seed used to generate the random starting point of the minimization.',
        dest = 'seed', action = 'store', type = int, default = None
    )
    parser_find_mode.add_argument('-L', '--label',
        help = 'The label used for the generated data file.',
        dest = 'label', action = 'store', type = str, default = 'default'
    )
    parser_find_mode.add_argument('-M', '--mask-name', metavar = 'MASK-NAME',
        help = 'The name of the mask to apply to the importance samples.',
        dest = 'mask_name', action = 'store', type = str
    )
    parser_find_mode.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_find_mode.set_defaults(cmd = cmd_find_mode)


    # mixture_product
    parser_mixture_product = subparsers.add_parser('mixture-product',
        parents = [common_subparser],
        description =
'''
    Compute the cartesian product of the densities listed in posteriors. Note that this product is not commutative.

    The input densities are read from EOS_BASE_DIRECTORY/POSTERIOR_i/pmc, where POSTERIOR_i is listed in posteriors.
    The output density will be stored in EOS_BASE_DIRECTORY/POSTERIOR/product.
''',
        help = 'Compute the cartesian product of named posteriors.'
    )
    parser_mixture_product.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF for which the product is computed',
        action = 'store', type = str
    )
    parser_mixture_product.add_argument('posteriors', metavar = 'POSTERIORS',
        help = 'The list of name of the posteriors PDF that will be concatenated.',
        action = 'store', type = lambda s: s.split(',')
    )
    parser_mixture_product.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_mixture_product.set_defaults(cmd = cmd_mixture_product)


    # find-clusters
    parser_find_clusters = subparsers.add_parser('find-clusters',
        parents = [common_subparser],
        description =
'''
Finds clusters among posterior MCMC samples, grouped by Gelman-Rubin R value, and creates a Gaussian mixture density.

Finding clusters and creating a Gaussian mixture density is a necessary intermediate step before using the sample-pmc subcommand.
The input files are expected in EOS_BASE_DIRECTORY/POSTERIOR/mcmc-*. All MCMC input files present will be used in the clustering.
The output files will be stored in EOS_BASE_DIRECTORY/POSTERIOR/clusters.
''',
        help = 'Finds clusters within MCMC samples of a named posterior.'
    )
    parser_find_clusters.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which MCMC samples have previously been drawn.',
        action = 'store', type = str
    )
    parser_find_clusters.add_argument('-t', '--threshold',
        help = 'The R value threshold. If two sample subsets have an R value larger than this threshold, they will be treated as two distinct clusters. (default: 2.0)',
        dest = 'threshold', action = 'store', type = float, default = 2.0
    )
    parser_find_clusters.add_argument('-c', '--clusters-per-group',
        help = 'The number of mixture components per cluster. (default: 1)',
        dest = 'K_g', action = 'store', type = int, default = 1
    )
    parser_find_clusters.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_find_clusters.set_defaults(cmd = cmd_find_clusters)


    # predict-observables
    parser_predict_observables = subparsers.add_parser('predict-observables',
        parents = [common_subparser],
        description =
'''
Predicts a set of observables based on previously obtained importance samples.

The input files are expected in EOS_BASE_DIRECTORY/POSTERIOR/samples.
The output files will be stored in EOS_BASE_DIRECTORY/POSTERIOR/pred-PREDICTION.
''',
        help = 'Predicts observables based on importance samples.'
    )
    parser_predict_observables.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which to draw the samples.'
    )
    parser_predict_observables.add_argument('prediction', metavar = 'PREDICTION',
        help = 'The name of the set of observables to predict.',
        action = 'store', type = str
    )
    parser_predict_observables.add_argument('-B', '--begin-index',
        help = 'The index of the first sample to use for the predictions.',
        dest = 'begin', action = 'store', type = int, default = 0
    )
    parser_predict_observables.add_argument('-E', '--end-index',
        help = 'The index beyond the last sample to use for the predictions.',
        dest = 'end', action = 'store', type = int, default = None
    )
    parser_predict_observables.add_argument('-M', '--mask-name', metavar = 'MASK-NAME',
        help = 'The name of the mask to apply to the samples.',
        dest = 'mask_name', action = 'store', type = str
    )
    parser_predict_observables.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_predict_observables.set_defaults(cmd = cmd_predict_observables)


    # corner-plot
    parser_corner_plot = subparsers.add_parser('corner-plot',
        parents = [common_subparser],
        description =
'''
Generate a corner plot of the 1-D and 2-D marginalized posteriors.

The input files are expected in EOS_BASE_DIRECTORY/POSTERIOR/samples.
The output files will be stored in EOS_BASE_DIRECTORY/POSTERIOR/plots.
''',
        help = 'Generate a corner plot of the 1-D and 2-D marginalized posteriors.'
    )
    parser_corner_plot.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which the samples were drawn.'
    )
    parser_corner_plot.add_argument('-B', '--begin-parameter',
        help = 'The index of the first parameter to plot.',
        dest = 'begin', action = 'store', type = int, default = 0
    )
    parser_corner_plot.add_argument('-E', '--end-parameter',
        help = 'The index beyond the last parameter to plot.',
        dest = 'end', action = 'store', type = int, default = None
    )
    parser_corner_plot.add_argument('-F', '--format',
        help = 'The plot output format. Can be a comma separated list of formats.',
        dest = 'format', action = 'store', type = lambda s: s.split(','), default = 'pdf'
    )
    parser_corner_plot.add_argument('-M', '--mask-name', metavar = 'MASK-NAME',
        help = 'The name of the mask to apply to the samples.',
        dest = 'mask_name', action = 'store', type = str
    )
    parser_corner_plot.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_corner_plot.set_defaults(cmd = cmd_corner_plot)


    # validate
    parser_validate = subparsers.add_parser('validate',
        parents = [common_subparser],
        description =
'''
Validates the analysis file by checking that all posteriors and all prediction sets can be created.
''',
        help = 'Validate the analysis file.'
    )
    parser_validate.set_defaults(cmd = cmd_validate)


    # list-figures
    parser_list_figures = subparsers.add_parser('list-figures',
        parents = [common_subparser],
        description =
'''
Lists the figures defined within the analysis file.
''',
        help = 'List the figures defined in the analysis file.'
    )
    parser_list_figures.set_defaults(cmd = cmd_list_figures)


    # list-steps
    parser_list_steps = subparsers.add_parser('list-steps',
        parents = [common_subparser],
        description =
'''
Lists the steps defined within the analysis file.
''',
        help = 'List the steps defined in the analysis file.'
    )
    parser_list_steps.set_defaults(cmd = cmd_list_steps)


    # list-step-dependencies
    parser_list_step_dependencies = subparsers.add_parser('list-step-dependencies',
        parents = [common_subparser],
        description =
'''
Lists the dependencies of a step defined within the analysis file.
''',
        help = 'List the dependencies of a step defined in the analysis file.'
    )
    parser_list_step_dependencies.add_argument('id', metavar = 'ID',
        help = 'The id of the step for which to list the dependencies.'
    )
    parser_list_step_dependencies.set_defaults(cmd = cmd_list_step_dependencies)


    # create-mask
    parser_create_mask = subparsers.add_parser('create-mask',
        parents = [common_subparser],
        description =
'''
Create a mask that can be applied to samples from a named posterior based on a set of observables.
''',
        help = 'Filters samples from a posterior.'
    )
    parser_create_mask.add_argument('posterior', metavar = 'POSTERIOR',
        help = 'The name of the posterior PDF from which to draw the samples.'
    )
    parser_create_mask.add_argument('mask_name', metavar = 'MASK-NAME',
        help = 'The name of the mask to create.',
        action = 'store', type = str
    )
    parser_create_mask.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_create_mask.set_defaults(cmd = cmd_create_mask)


    # report
    parser_report = subparsers.add_parser('report',
        parents = [common_subparser],
        description =
'''
Generates a report of the analysis based on the analysis file, the recorded results, and a report template file in Jinja2 format.
''',
        help = 'Generates a report of the analysis.'
    )
    parser_report.add_argument('template_file', metavar='TEMPLATE',
        help = 'The file name of the template that is used to generate the report.'
    )
    parser_report.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_report.add_argument('--no-pdf',
        help = 'Do not generate a PDF report.',
        dest = 'generate_pdf', action = 'store_false', default = True
    )
    parser_report.set_defaults(cmd = cmd_report)


    # run
    parser_run = subparsers.add_parser('run',
        parents = [common_subparser],
        description =
'''
Runs one out of a list of steps recorded within an analysis file.
''',
        help = 'Run a step by id.'
    )
    parser_run.add_argument('id', metavar = 'ID',
        help = 'The id of the step to run.'
    )
    parser_run.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_run.add_argument('-d', '--dry-run',
        help = 'Perform a dry run only. Outputs the list of subcommands that would be run instead of running them.',
        dest = 'dry_run', action = 'store_true', default = False
    )
    parser_run.set_defaults(cmd = cmd_run)


    # draw-figure
    parser_draw_figure = subparsers.add_parser('draw-figure',
        parents = [common_subparser],
        description =
'''
Draws a figure based on the specification in the analysis file and the recorded results.
''',
        help = 'Draws a figure of the analysis.'
    )
    parser_draw_figure.add_argument('figure_name', metavar='FIGURE',
        help = 'The name of the figure that shall be drawn.'
    )
    parser_draw_figure.add_argument('-b', '--base-directory',
        help = 'The base directory for the storage of data files. Can also be set via the EOS_BASE_DIRECTORY environment variable.',
        dest = 'base_directory', action = 'store', default = get_from_env('EOS_BASE_DIRECTORY', './')
    )
    parser_draw_figure.add_argument('-F', '--format',
        help = 'The additional formats in which the figure should be saved. Defaults to PDF only.',
        dest = 'format', action = 'store', default = ['pdf'], type = lambda f: [str(item) for item in f.split(',')]
    )
    parser_draw_figure.set_defaults(cmd = cmd_draw_figure)


    # init
    parser_init = subparsers.add_parser('init',
        parents = [common_subparser],
        description =
'''
Initializes a new analysis directory based on the EOS analysis template.
''',
        help = 'Initializes a new analysis directory.'
    )
    parser_init.set_defaults(cmd = cmd_init)

    ## end of commands

    return parser

class CustomLogFormatter(logging.Formatter):
    _MAP_LEVEL_TO_COLOR = {
        logging.ERROR:      ('✖', 'red'),
        logging.WARNING:    ('⚠', 'yellow'),
        logging.SUCCESS:    ('🗸', 'green'),
        logging.COMPLETED:  ('🗸', 'green'),
        logging.INPROGRESS: ('…', 'green'),
        logging.INFO:       ('ℹ', 'blue'),
        logging.DEBUG:      ('🤖', None),
    }
    def __init__(self):
        super().__init__(fmt='%(levelname)s %(message)s', datefmt=None, style='%')

    def format(self, record):
        levelno = record.levelno
        if record.levelno not in self._MAP_LEVEL_TO_COLOR:
            levelno = logging.ERROR

        symbol, color = self._MAP_LEVEL_TO_COLOR[record.levelno]
        record.levelname = colored(symbol, color, attrs=['bold'])
        record.msg = colored(record.msg, color)

        return super().format(record)

def main():
    parser = _parser()
    argcomplete.autocomplete(parser)
    args = parser.parse_args()

    try:
        if not 'cmd' in args:
            parser.print_help()
        elif not callable(args.cmd):
            parser.print_help()
        else:
            if not args.verbose:
                args.verbose = get_from_env('EOS_VERBOSITY', 0)
            if args.verbose > 5:
                args.verbose = 5

            levels = {
                0: logging.ERROR,
                1: logging.WARNING,
                2: logging.SUCCESS,
                3: logging.INPROGRESS,
                4: logging.INFO,
                5: logging.DEBUG
            }

            eos.stderr_handler.setLevel(levels[args.verbose])
            eos.stderr_handler.setFormatter(CustomLogFormatter())

            args.cmd(args)
    except Exception as e:
        print(colored('✖ Encountered an unrecoverable error:\n', 'red', attrs=['bold']), f'{e}')
        if not isinstance(e, ValueError):
            traceback.print_exception(e, e, e.__traceback__)
        sys.exit(1)


def make_analysis_file(args):
    analysis_file = eos.AnalysisFile(args.analysis_file)
    return analysis_file

def args_to_dict(args):
    result = vars(args)
    result.pop('verbose', None)
    result.pop('cmd', None)
    return result

class RNG:
    def __init__(self, seed):
        self._rng = _np.random.mtrand.RandomState(seed)

    def rand(self, *args):
        result = self._rng.rand(*args)
        print('rand: {}'.format(result))
        return result

    def normal(self, loc, scale=1, size=None):
        result = self._rng.normal(loc, scale, size)
        print('normal: {}'.format(result))
        return result

    def uniform(self, low, high, size=None):
        result = self._rng.uniform(low, high, size)
        print('uniform: {}'.format(result))
        return result


# Find mode
def cmd_find_mode(args):
    return eos.tasks.find_mode(**args_to_dict(args))


# Plot samples
def cmd_plot_samples(args):
    import pathlib
    input_path = pathlib.Path(os.path.join(args.base_directory, args.posterior))
    inputs  = [str(d) for d in input_path.glob('mcmc-*')]
    inputs += [str(d) for d in input_path.glob('samples')]
    for input in inputs:
        info('plotting samples in \'{}\''.format(input))
        basename = os.path.basename(os.path.normpath(input))
        if basename.startswith('mcmc-'):
            data = eos.data.MarkovChain(input)
        elif basename.startswith('samples'):
            data = eos.data.ImportanceSamples(input)
        else:
            raise RuntimeError('unsupported data set: {}'.format(input))

        parameters = eos.Parameters()
        for idx, p in enumerate(data.varied_parameters):
            info('plotting histogram for {}'.format(p['name']))
            if data.type in ['Prediction']:
                label = eos.Observables()[p['name']]
            elif data.type in ['MarkovChain', 'ImportanceSamples']:
                pp = parameters[p['name']]
                label = pp.latex()
            else:
                label = r'\verb+{}+'.format(p['name'])

            description = {
                'plot': {
                    'x': { 'label': label, 'range': [p['min'], p['max']] },
                    'y': { 'label': 'prob. density' }
                },
                'contents': [
                    {
                        'type': 'histogram', 'bins': args.bins,
                        'data': {
                            'samples': data.samples[:, idx],
                        }
                    }
                ]
            }
            plotter = eos.plot.Plotter(description, os.path.join(input, '{}.pdf'.format(idx)))
            plotter.plot()


# Sample MCMC
def cmd_sample_mcmc(args):
    return eos.sample_mcmc(**args_to_dict(args))


# Find clusters
def cmd_find_clusters(args):
    return eos.find_clusters(**args_to_dict(args))


# Sample PMC
def cmd_sample_pmc(args):
    return eos.sample_pmc(**args_to_dict(args))


# Nested sampling
def cmd_sample_nested(args):
    return eos.sample_nested(**args_to_dict(args))


# Corner plot
def cmd_corner_plot(args):
    return eos.corner_plot(**args_to_dict(args))


# Cartesian product
def cmd_mixture_product(args):
    return eos.mixture_product(**args_to_dict(args))


# Predict observables
def cmd_predict_observables(args):
    return eos.predict_observables(**args_to_dict(args))


# Validate analysis file
def cmd_validate(args):
    return eos.validate(**args_to_dict(args))


# List figures in analysis
def cmd_list_figures(args):
    result = eos.list_figures(**args_to_dict(args))
    if len(result) > 0:
        print('\n'.join(result))


# List steps in analysis file
def cmd_list_steps(args):
    result = eos.list_steps(**args_to_dict(args))
    print('\n'.join(result))


# List steps in analysis file
def cmd_list_step_dependencies(args):
    result = eos.list_step_dependencies(**args_to_dict(args))
    if isinstance(result, list) and len(result) > 0:
        print('\n'.join(result))
    elif isinstance(result, str):
        print(result)


# Filter samples
def cmd_create_mask(args):
    return eos.create_mask(**args_to_dict(args))


# Report
def cmd_report(args):
    return eos.report(**args_to_dict(args))


# Run steps
def cmd_run(args):
    return eos.run(**args_to_dict(args))


# Draw figure
def cmd_draw_figure(args):
    return eos.draw_figure(**args_to_dict(args))


# List priors
def cmd_list_priors(args):
    analysis_file = make_analysis_file(args)
    for name, prior in analysis_file.priors.items():
        print(name)


# List likelihoods
def cmd_list_likelihoods(args):
    analysis_file = make_analysis_file(args)
    for name, lh in analysis_file.likelihoods.items():
        print(name)
        if not args.display:
            continue

        if not 'constraints' in lh and not 'manual_constraints' in lh:
            error('Likelihoods {name} specifies neither the \'constraints\' nor the \'manual_constraints\' key'.format(name=name))
            continue

        for c in lh['constraints'] if 'constraints' in lh else []:
            print(' - {}'.format(c))

        for mc in lh['manual_constraints'] if 'manual_constraints' in lh else {}:
            print(' - {} [manual]'.format(mc))


# List predictions
def cmd_list_predictions(args):
    analysis_file = make_analysis_file(args)
    for name, pred in analysis_file.predictions.items():
        print(name)


# List posteriors
def cmd_list_posteriors(args):
    analysis_file = make_analysis_file(args)
    for name, lh in analysis_file.posteriors.items():
        print(name)


# Initialize analysis directory
def cmd_init(args):
    import io, requests, zipfile

    def download_template_zip():
        try:
            response = requests.get('https://github.com/eos/analysis-template/archive/main.zip')
        except OSError as e:
            eos.error(f'Failed to download analysis template: error = {e}')
            sys.exit(1)

        # validate response
        if not response.status_code == 200:
            eos.error(f'Failed to download analysis template: status code = {response.status_code}')
            sys.exit(2)

        return response.content

    def strip_leading_directory(zip):
        for name in zip.namelist():
            path = '/'.join(name.split('/')[1:])
            yield path

    with zipfile.ZipFile(io.BytesIO(download_template_zip())) as zip:
        for member in zip.namelist():
            path = '/'.join(member.split('/')[1:])
            zip.extract(member, path)


if __name__ == '__main__':
    main()
