#!/usr/bin/env python3

import glob
import os
import time
import argparse
from pathlib import Path

from pyctbn.legacy import JsonImporter
from pyctbn.legacy import SamplePath
from pyctbn.legacy import StructureConstraintBasedEstimator

def structure_constraint_based_estimation_example(network_file_path, jobs):
    print("Importing %s..." % (network_file_path))
    # <initialize a JsonImporter object for the first file>
    importer = JsonImporter(
        file_path=network_file_path,
        samples_label='samples',
        structure_label='dyn.str',
        variables_label='variables',
        time_key='Time',
        variables_key='Name'
    )
    start_time = time.time()
    # <import the data at index 0 of the outer json array>
    importer.import_data(0)
    print("Data imported in %d seconds." % (time.time() - start_time))
    # construct a SamplePath Object passing a filled AbstractImporter object
    s1 = SamplePath(importer=importer)
    # build the trajectories
    print("Building trajectories...")
    s1.build_trajectories()
    # build the information about the net
    print("Building structure...")
    s1.build_structure()
    # construct a StructureEstimator object passing a correctly build SamplePath object
    # and the independence tests significance, if you have prior knowledge about
    # the net structure create a list of tuples
    # that contains them and pass it as known_edges parameter
    se1 = StructureConstraintBasedEstimator(
        sample_path=s1,
        exp_test_alfa=1e-6,
        chi_test_alfa=1e-4,
        known_edges=[],
        thumb_threshold=25
    )
    # call the algorithm to estimate the structure
    print("Estimating structure...")
    start_estimating_time = time.time()
    if jobs == 0:
        disable_multiprocessing=False
        processes_number = None
    elif jobs == 1:
        disable_multiprocessing = True
        processes_number = None
    else:
        disable_multiprocessing = False
        processes_number = jobs
    se1.estimate_structure(
        disable_multiprocessing=disable_multiprocessing,
        processes_number=processes_number
    )
    print("Structure estimated in %d seconds." % (time.time() - start_estimating_time))
    end_time = time.time()
    print("Total elaspsed time for %s: %d seconds" % (network_file_path, end_time - start_time))
    # obtain the adjacency matrix of the estimated structure
    #print(se1.adjacency_matrix())
    Path("./res").mkdir(parents=True, exist_ok=True)
    # save the estimated structure  to a json file
    # (remember to specify the path AND the .json extension)....
    se1.save_results("./res/results_%s" % (os.path.basename(network_file_path)))
    # ...or save it also in a graphical model fashion
    # (remember to specify the path AND the .png extension)
    #se1.save_plot_estimated_structure_graph("./res/result0.png")


class JobsAction(argparse.Action):

    def __call__(self, parser, namespace, values, option_string=None):
        if values < 0:
            parser.error(
                "Minimum jobs for {0} is 0 (default), which means to "
                "use all cores available.".format(option_string)
            )
        setattr(namespace, self.dest, values)

parser = argparse.ArgumentParser(
    prog='pyCTBN - Benchmark',
    description="This benchmark program is tailored to work with the "
    "\"modernized\" and polished version of pyCTBN."
)
parser.add_argument(
    'network_path',
    help="path of the network file in json format"
)
parser.add_argument(
    '-j',
    '--jobs',
    action=JobsAction,
    type=int,
    default=0,
    required=False,
    help="number of jobs (processes) to use (0 by default, it uses all available cores)"
)
args = parser.parse_args()

structure_constraint_based_estimation_example(args.network_path, args.jobs)