#!/usr/bin/env python3

import os
import xml.etree.ElementTree as etree
import subprocess
import re
import argparse


def create_folder(path):
    """Creates a folder if it does not exist

    Parameters
    ----------
    path : str
        Path of the new folder

    Examples
    --------

    >>> create_folder('./results')
    """
    if not os.path.exists(path):
        os.makedirs(path)


def execute_command(command):
    """Run a command in the shell"""
    output = None
    try:
        output = subprocess.check_output(
            command, shell=True, stderr=subprocess.STDOUT).decode("utf-8")
    except Exception as e:
        output = e.output.decode("utf-8")
    except KeyboardInterrupt:
        pass

    return output


def get_atomic_propositions(pnml_file):
    """Returns the atomic propositions from a PNML file"""
    xml_tree = etree.parse(pnml_file)
    root = xml_tree.getroot()
    ns = {'pnml': root.tag[1:].split("}")[0]}

    places = [p.get('id') for p in root.findall('.//pnml:place', ns)]
    transitions = [t.get('id') for t in root.findall('.//pnml:transition', ns)]

    return {'places': places, 'transitions': transitions}


def generate_spot_formulas(model_name, model_instance, n, paths):
    """Generates a number of formulas for a specific PNML model

    Parameters
    ----------
    model_name : str
        Name of model
    model_instance : str
        Instance of the model
    n : int
        Number of formulas to be generated
    paths : dict of str
        Paths of the project

    Returns
    -------
    list of str
        list with formulas
    """
    ignored_operators = "X=0,xor=0,W=0,M=0,R=0,U=0"

    pnml_file = os.path.join(paths['models'], model_name,
                             f'{model_instance}.pnml')

    atomic_propositions = get_atomic_propositions(pnml_file)
    atomic_propositions = " ".join(
        [f"{ap}==1" for ap in atomic_propositions['places']])

    command = f'{paths["randltl"]} -L -s -p -n{n} --ltl-priorities="{ignored_operators}" {atomic_propositions}'

    output = execute_command(command)
    formulas = output.split('\n')[:-1]

    return formulas


def check_formula(model_name, model_instance, formula, timeout, min_states,
                  paths):
    """Check whether a property is satisfied or not"""
    tool = paths['pnml2lts-mc']

    model_file = os.path.join(paths['models'], model_name,
                              f"{model_instance}.pnml")

    command = f'{tool} --strategy=ndfs --timeout={timeout} --ltl="{formula}" {model_file}'
    ltsmin_out = execute_command(command)

    # timeout or error found
    if (ltsmin_out is None):
        return None

    match_false = re.search(r'Accepting cycle FOUND', ltsmin_out)
    match_time = re.search(r'Total exploration time (\d+(\.\d+)?)', ltsmin_out)
    match_states = re.search(r'State space has (\d+) states', ltsmin_out)

    # ignore formulas with time zero
    if match_time:
        nb_states = int(match_states.group(1))

        if (nb_states >= min_states):
            print("# states (state space):", nb_states)
            return not match_false

    return None


def check_explored_states(model_name, model_instance, formula, time_out,
                          min_explored_states, paths):
    """Checks the number of states visited by the model checker"""
    tool = paths['pnml2lts-mc']
    model_file = os.path.join(paths['models'], model_name,
                              f"{model_instance}.pnml")

    command = f'{tool} --strategy=dfs --timeout={timeout*60} --ltl="{formula}" {model_file}'
    ltsmin_out = execute_command(command)

    # timeout or error found
    if (ltsmin_out is None):
        return False

    match_regex_explored = re.search(r'Explored (\d+) states', ltsmin_out)

    if match_regex_explored:
        nb_explored_states = int(match_regex_explored.group(1))
        print("# explored states: ", nb_explored_states)
        return nb_explored_states >= min_explored_states

    return False


def get_default_paths():
    """Returns the default path of the project"""

    # Absolute path where are stored the formulas, models
    base_folder = os.path.abspath(
        os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))

    paths = {
        # randltl path: this tool generates random formulas
        "randltl": "randltl",
        # tool to verify the property
        'pnml2lts-mc': os.path.join(base_folder, "tools", "pnml2lts-mc"),
        # Folder where the models are saved
        'models': os.path.join(base_folder, "models"),
        # Folder where the formulas are saved
        'formulas': os.path.join(base_folder, "formulas"),
    }

    return paths


def write_formula(result, model_name, model_instance, formula, identifier, paths):
    """write a formula into a file"""
    base = os.path.join(paths['formulas'], model_name, model_instance)
    create_folder(base)
    filename_ltsmin = f"{model_instance}-{identifier}.ltl"
    filename_pmcsog = f"{filename_ltsmin}.reduced"
    filename_results = "formula_results"

    with open(os.path.join(base, filename_ltsmin), 'w') as ltsmin, \
        open(os.path.join(base, filename_pmcsog), 'w') as pmcsog, \
        open(os.path.join(base, filename_results), 'a') as results:

        # remove ==1 and add " to atomic propositions
        pmcsog_formula = re.sub(pattern=r'\((\w*\d*)==1\)',
                                repl='("\\1")',
                                string=formula)

        ltsmin.write(formula)
        pmcsog.write(pmcsog_formula)
        results.write(f"{identifier};{result}\n")

        
def get_instance_stats(model_name, model_instance):
    """get the next formula id for a model instance and the number of true
    and false formulas already computed."""
    base = os.path.join(paths['formulas'], model_name, model_instance)
    filename_results = os.path.join(base, "formula_results")
    if not os.path.exists(filename_results):
        return 1, 0, 0
    true_formulas = 0
    false_formulas = 0
    with open(filename_results, 'r') as fd:
        formulas = fd.read().split('\n')
        for line in formulas:
            id_result = line.split(';')
            if len(id_result) != 2:
                break
            if id_result[1].lower() == 'true':
                true_formulas += 1
            else:
                false_formulas += 1
        return len(formulas), true_formulas, false_formulas


def generate_formulas(model_name,
                      model_instance,
                      n,
                      timeout,
                      min_states,
                      min_explored_states,
                      paths,
                      save=True,
                      max_retries=5):
    """Generates N satisfied and N violated formulas"""
    true_formulas = []
    false_formulas = []
    unknown_formulas = []  #  subset of true_formulas for which ltsmin
                           #  timed out max_retries times

    random_formulas = generate_spot_formulas(model_name, model_instance,
                                             100000, paths)

    formula_id, nb_true, nb_false = get_instance_stats(model_name, model_instance)
    print(f"I still have {max(0, n - nb_true)} true formulas, " +
          f"and {max(0, n - nb_false)} false formulas to generate.\n" +
          f"I will start from formula id {formula_id}.")
    for f in random_formulas:
        if len(true_formulas) + nb_true >= n and len(false_formulas) + nb_false >= n:
            break

        #  launch ltsmin at most max_retries times until a run terminates
        for _ in range(max_retries):
            is_satisfied = check_formula(model_name, model_instance, f, timeout,
                                         min_states, paths)

            #  run terminated successfuly
            if is_satisfied is not None:
                break

        #  no run terminated after max_retries attempts => claim the
        #  property is verified (even if we're not sure)
        is_unknown = is_satisfied is None
        if is_satisfied is None:
            is_satisfied = True

        # add to formulas satisfied
        if is_satisfied and len(true_formulas) + nb_true < n:
            # if number of explored states is not satisfied
            has_explored_states = check_explored_states(
                model_name, model_instance, f, timeout, min_explored_states,
                paths)
            if not has_explored_states:
                continue

            print('### true formula found: ', f)
            true_formulas.append(f)
            if is_unknown:
                unknown_formulas.append(f)
            write_formula(True, model_name, model_instance, f, formula_id, paths)
            formula_id = formula_id + 1

        # add to formulas violated
        if not is_satisfied and len(false_formulas) + nb_false < n:
            print('### false formula found: ', f)
            false_formulas.append(f)
            write_formula(False, model_name, model_instance, f, formula_id, paths)
            formula_id = formula_id + 1

    return {'satisfied': true_formulas,
            'violated': false_formulas,
            'unknown': unknown_formulas}


def generate_all_formulas(model_name,
                          model_instances,
                          n,
                          timeout,
                          min_states,
                          min_explored_states,
                          paths,
                          save=True):
    for model_instance in model_instances:
        formulas = generate_formulas(model_name, model_instance, n, timeout,
                                     min_states, min_explored_states, paths,
                                     save)
        print(model_instance, formulas)


def parse_arguments():
    """Creates the command line interface"""
    parser = argparse.ArgumentParser(
        description='Generates formulas for a model')

    parser.add_argument('--n',
                        type=int,
                        nargs=1,
                        required=True,
                        help='number of formulas')

    parser.add_argument('--timeout',
                        type=int,
                        nargs=1,
                        required=True,
                        help='timeout of verification (s)')

    parser.add_argument('--min_states',
                        type=int,
                        nargs=1,
                        required=True,
                        help='minimum number of states')

    parser.add_argument('--min_explored_states',
                        type=int,
                        nargs=1,
                        required=True,
                        help='minimum number of explored states')

    parser.add_argument('--model-name',
                        type=str,
                        nargs=1,
                        required=True,
                        help='Name of the model (e.g. philo)')

    parser.add_argument('--model-instance',
                        type=str,
                        nargs='+',
                        required=True,
                        help='Instances of the model (e.g. philo10)')

    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()

    # default paths of the project
    paths = get_default_paths()

    n_formulas = args.n[0]
    timeout = args.timeout[0]
    min_states = args.min_states[0]
    min_explored_states = args.min_explored_states[0]
    model_name = args.model_name[0]
    model_instances = args.model_instance

    generate_all_formulas(model_name, model_instances, n_formulas, timeout,
                          min_states, min_explored_states, paths)