import argparse
import json
import os

from abcsmc.loader import ABCLoader, SQLDataStore
from connectome.abctask import ABCConnectomeTask, execute_task_list, GroundTruthData
from connectome.builder import NoiseSubsamplingModifier, SubsamplingModifier
from connectome.builder import PipelineBuilder, ModelModifier
from connectome.model import default_definitions
from parallel import Slurm, slurm_available
from typing import Union


def sample(model_name: str, model_definition_modifier: ModelModifier, model_definitions) -> GroundTruthData:
    """

    Parameters
    ----------
    model_name: str
        Model name

    model_definition_modifier: ModelModifier
        Modify the model parameters from the file

    Returns
    -------

    gt_data: GroundTruthData
        Ground truth data
    """
    definitions = model_definitions()
    model_definition_modifier.modify(definitions)
    builder = PipelineBuilder(definitions)

    parameter = builder.build_prior(model_name).rvs()
    model = builder.build_model_with_noise(model_name)

    network = model(parameter)
    print("GT parameter", parameter, "Model", model)
    parameter.update(model_definition_modifier.parameters())  # this is a hack for recording purposes. not very nice.
    print("GT parameter to store", parameter, "Model", model)
    gt_data = GroundTruthData(network, model_name, parameter)
    return gt_data


def sample_ground_truth_starred(args, model_definitions):
    model_name, modifier = args
    print(f"Sample ground truth for model {model_name}")
    return sample(model_name, modifier, model_definitions)


def parse_noise(noise_str):
    try:
        return float(noise_str)
    except ValueError:
        dist_type = noise_str.split("(")[0]
        pars = list(map(float, noise_str.split("(")[1].split(")")[0].split(",")))
        return {"type": dist_type, "args": pars}


def main():
    argument_parser = argparse.ArgumentParser(description="Network ABCSMC. Samples the ground truth for each of the 7 models anda pplies noise and subsampling.")
    argument_parser.add_argument("db", type=str)
    argument_parser.add_argument("--no-parallel-jobs", type=int, default=30, help="Default: 30")
    argument_parser.add_argument("--noise", type=str, default="[0.0]", help="Default: [0.0]. A comma separated list of values: e.g. [0.0, 0.1]")
    argument_parser.add_argument("--noise-prior", type=str, default="", help="Default: Prior of noise model (see `--noise-model`). A scipy.stats distribution or a numberic value: e.g. beta(2,10) or just 0 or 0.5")
    argument_parser.add_argument("--subsampling", type=str, default="[1.0]", help="Default: [1.0]. A comma separated list of values: e.g. [1.0, 0.9]")
    argument_parser.add_argument("--memory", type=str, default="2G", help="Default: 2G. Can also be less")
    argument_parser.add_argument("--noise-model", type=str, default='remove_add_noise_and_subsample', help="Noise override")
    argument_parser.add_argument("--skip-true-model", action='store_true', help="Default: False. Remove true model from the domain of the posterior distribution.")
    args = argument_parser.parse_args()

    # NOTE(amotta): Load noise model (and priors) from JSON file
    noise_model_dir = os.path.dirname(__file__)
    noise_model_file = os.path.join(noise_model_dir, args.noise_model + '.json')

    with open(noise_model_file) as f:
        noise_model = json.load(f)

    assert 'noise' in noise_model
    assert 'noise_key' in noise_model
    if 'subsampling_key' not in noise_model:
        noise_model['subsampling_key'] = None

    def model_definitions():
        default = default_definitions()
        default['noise'] = noise_model['noise']

        if args.noise_prior != "":
            noise_prior = parse_noise(args.noise_prior)
            # NOTE(amotta): Override the noise parameters with the specified values
            default['noise']["parameters"][noise_model['noise_key']] = noise_prior

        return default

    mapper = Slurm(memory=args.memory, name="abc").map if slurm_available() else map

    model_names = model_definitions()['models']['model_specific_parameters'].keys()
    model_names = list(sorted(list(model_names)))

    print("Start generate ground truth samples", flush=True)
    # NOTE(amotta): Set noise to specified value and add subsampling
    gt_parameters = [(model_name, NoiseSubsamplingModifier(noise, subsampling,
                                                           noise_key=noise_model['noise_key'],
                                                           subsampling_key=noise_model['subsampling_key']))
                     for noise in json.loads(args.noise)
                     for subsampling in json.loads(args.subsampling)
                     for model_name in model_names]

    gt_samples = list(mapper(lambda gt_par: sample_ground_truth_starred(gt_par, model_definitions), gt_parameters))
    print("Finish generate ground truth samples", flush=True)

    abc_tasks = []
    for gt_sample, (model_name, modifier) in zip(gt_samples, gt_parameters):
        definitions = model_definitions()

        if args.skip_true_model:
            # NOTE(amotta): Remove true model from the domain of the posterior distribution
            del definitions['models']['model_specific_parameters'][model_name]

        # NOTE(amotta): Add subsampling, but leave the noise unchanged. This
        # way, the noise values will be drawn from the prior distribution
        # specified by the user.
        SubsamplingModifier(subsampling=modifier.subsampling,
                            subsampling_key=noise_model['subsampling_key']).modify(definitions)

        # NOTE(amotta): Allow arbitrarily small values for epsilon
        task = ABCConnectomeTask(definitions, gt_sample, args.db, mapper, minimum_epsilon=None)
        abc_tasks.append(task)

    print("Start ABC Tasks in parallel", flush=True)
    execute_task_list(abc_tasks, args.no_parallel_jobs)


if __name__ == "__main__":
    main()
