"""Workflow for optimization and analysis of PEtab problems"""

from snakemake.utils import min_version, validate
min_version("5.0.0")

import os

import matplotlib
# Allow running without tkinter
matplotlib.use('agg')
import matplotlib.pyplot as plt

configfile: "parpe_optimize_petab.yaml"
validate(config, "config.schema.yaml")

def process_config(config):
    """Process configuration entries"""

    if 'root' in config['petab']:
        # Expand environment variables
        config['petab']['root'] = os.path.expandvars(config['petab']['root'])

        # Make prepend root path if not absolute path
        for x in ['yaml_file']:
            if not config['petab'][x].startswith('/'):
                config['petab'][x] = os.path.join(config['petab']['root'],
                                                  config['petab'][x])

    # Expand environment variables
    config['amici_build_dir'] = os.path.expandvars(config['amici_build_dir'])
    config['amici_src_dir'] = os.path.expandvars(config['amici_src_dir'])
    config['parpe_build_dir'] = os.path.expandvars(config['parpe_build_dir'])
    config['parpe_src_dir'] = os.path.expandvars(config['parpe_src_dir'])

    # Set additional variables
    # TODO: allow setting in config file
    config['amici_model_dir'] = os.path.join(
        'results', f'amici_{config["model_name"]}')
    config['parpe_model_dir'] = os.path.join(
        'results', f'parpe_{config["model_name"]}')
    config['hdf5_training_file'] = os.path.join(
        'results', config['model_name'] + '.h5')

process_config(config)


rule show_config:
    """Print configuration"""
    run:
        from pprint import pprint
        pprint(config)

rule generate_amici_model:
    """Generate AMICI model from petab files"""
    input:
        yaml_file=config['petab']['yaml_file'],
    output:
        directory(config['amici_model_dir'])
    shell:
        "amici_import_petab --verbose "
        "-n {config[model_name]} "
        "-y {input.yaml_file} "
        "-o {output}"

rule create_hdf5_file:
    """Create parPE hdf5 file"""
    # TODO: apply num_starts, etc.
    input:
        yaml_file=config['petab']['yaml_file'],
        amici_model_dir=config['amici_model_dir']
    output:
        h5=config['hdf5_training_file']
    shell:
        "parpe_petab_to_hdf5 "
        "-o {output.h5} "
        "-y {input.yaml_file} "
        "-d {input.amici_model_dir} "
        "-n {config[model_name]} "
        "&& {config[parpe_src_dir]}/misc/optimizationOptions.py {output.h5} -s numStarts {config[optimization][num_starts]} "
        "&& {config[parpe_src_dir]}/misc/optimizationOptions.py {output.h5} -s maxIter {config[optimization][max_iter]} "
        "&& {config[parpe_src_dir]}/misc/optimizationOptions.py {output.h5} -s ipopt/max_iter "

rule setup_parpe_model:
    """Setup AMICI model for use with parPE and compile"""
    input:
        config['amici_model_dir']
    output:
        #directory(config['parpe_model_dir']),
        estimate_exe=os.path.join(
            config['parpe_model_dir'], 'build',
            f'estimate_amici_{config["model_name"]}')
    shell:
        # delete since this will be created by snakemake and will
        # result in failure of following script
        "rm -r {config[parpe_model_dir]};"
        "{config[parpe_src_dir]}/misc/setup_amici_model.sh "
        "\"{input}\" \"{config[parpe_model_dir]}\""

rule optimize:
    """Run optimization"""
    input:
        optim_exe=rules.setup_parpe_model.output.estimate_exe,
        hdf5_training=rules.create_hdf5_file.output
    output:
        first_file="results/optimization/_rank00000.h5"
    params:
        out_dir=lambda wildcards, output:
        f"{os.path.dirname(output.first_file)}{os.path.sep}"
    shell:
        "PARPE_NO_DEBUG=1 {input.optim_exe} -o {params.out_dir} {input.hdf5_training}"

rule preprocess:
    """Prepare for optimization"""
    input:
        rules.optimize.input

rule gradient_check:
    """Run finite difference gradient check"""
    input:
        optim_exe=rules.setup_parpe_model.output.estimate_exe,
        #parpe_model_dir=config['parpe_model_dir'],
        hdf5_training=config['hdf5_training_file']
    output:
        first_file="results/gradient_check/_rank00000.h5"
    params:
        out_dir=lambda wildcards, output:
        f"{os.path.dirname(output.first_file)}{os.path.sep}"
    shell:
        "{input.optim_exe} -t gradient_check -o {params.out_dir} {input.hdf5_training}"


# TODO: separate by per-optimization and combined
rule postprocess:
    """Postprocessing of results"""
    input:
        first_file=rules.optimize.output.first_file
    # TODO dynamic input
    # TODO: data analysis script
    output:
        cost_trajectory_file="results/figures/cost_trajectory.png"
    run:
        filename = input.first_file
        import parpe

        trajectories = parpe.getCostTrajectories(filename)
        #print(repr(trajectories))
        parpe.plotting.plotCostTrajectory(trajectories, log=False)
        plt.savefig(output.cost_trajectory_file)

"""
# rule merge_results

rule analyze_results:
    input:
    output:
"""
"""
rule import_and_run:
    output:
        directory('{model_name}')
    input:
        model_name=lambda wildcards: os.path.join(benchmark_model_dir, wildcards.model_name)
    shell:
        './import_and_run.sh {benchmark_model_dir}/{wildcards.model_name}'

rule clean:
    shell:
        "ls -d */ | grep -P '^(parpe_)?[A-Z]\w+_\w+.*\d{{4}}/$' | xargs -d\"\n\" rm -r"
"""