import argparse
import sys
import time
from functools import partial

from pypdevs.DEVS import CoupledDEVS
from pypdevs.simulator import Simulator

from devstone import LI, HI, HO, HOmod, LI2HI, HI2LI, dLI, dHI
from generator import Generator

sys.setrecursionlimit(10000)

def terminate(clock, model, timeLimit):
    # print(clock[0])
    if clock[0] > timeLimit:
        return True
    return False


class DEVStoneEnvironment(CoupledDEVS):

    def __init__(self, name, devstone_model, num_gen_outputs=1, gen_period=None):
        super(DEVStoneEnvironment, self).__init__(name=name)
        generator = Generator("generator", num_gen_outputs, gen_period)
        self.addSubModel(generator)
        self.addSubModel(devstone_model)

        self.connectPorts(generator.o_out, devstone_model.i_in)

        if isinstance(devstone_model, HO) or isinstance(devstone_model, HOmod):
            self.connectPorts(generator.o_out, devstone_model.i_in2)


MODEL_TYPES = ("LI", "HI", "HO", "HOmod", "LI2HI", "HI2LI", "dLI", "dHI")
GEN_METHODS = ("start", "uniform", "end")


def parse_args():
    parser = argparse.ArgumentParser(description='Script to compare DEVStone implementations with different engines')

    parser.add_argument('-m', '--model-type', required=True, help='DEVStone model type (LI, HI, HO, HOmod, LI2HI, HI2LI)')
    parser.add_argument('-d', '--depth', type=int, required=True, help='Number of recursive levels of the model.')
    parser.add_argument('-w', '--width', type=int, required=True, help='Width of each coupled model.')
    parser.add_argument('-i', '--int-cycles', type=int, default=0, help='Dhrystone cycles executed in internal transitions')
    parser.add_argument('-e', '--ext-cycles', type=int, default=0, help='Dhrystone cycles executed in external transitions')
    parser.add_argument('-f', '--flatten', action="store_true", help='Activate flattening on model')
    parser.add_argument('-c', '--chained', action="store_true", help='Use chained coordinator')

    parser.add_argument('-C', '--classic', action="store_true", default=False, help='Use classic DEVS simulator')
    parser.add_argument('-D', '--dynamic', action="store_true", default=False, help='Enables Dynamic Structure DEVS')
    parser.add_argument("-g", "--gen-mode", choices=GEN_METHODS, help="Determines the spread of model creation.")
    args = parser.parse_args()

    if args.model_type not in MODEL_TYPES:
        raise RuntimeError("Unrecognized model type.")

    return args


if __name__ == '__main__':

    args = parse_args()

    devstone_model = None

    start_time = time.time()
    if args.model_type == "LI":
        devstone_model = LI("LI_root", args.depth, args.width, args.int_cycles, args.ext_cycles, prep_time=1)
    elif args.model_type == "HI":
        devstone_model = HI("HI_root", args.depth, args.width, args.int_cycles, args.ext_cycles, prep_time=1)
    elif args.model_type == "HO":
        devstone_model = HO("HO_root", args.depth, args.width, args.int_cycles, args.ext_cycles, prep_time=1)
    elif args.model_type == "HOmod":
        devstone_model = HOmod("HOmod_root", args.depth, args.width, args.int_cycles, args.ext_cycles, prep_time=1)
    elif args.model_type == "LI2HI":
        devstone_model = LI2HI("LI2HI_root", args.depth, args.width, args.int_cycles, args.ext_cycles, prep_time=1)
    elif args.model_type == "HI2LI":
        devstone_model = HI2LI("HI2LI_root", args.depth, args.width, args.int_cycles, args.ext_cycles, prep_time=1)
    elif args.model_type == "dLI":
        devstone_model = dLI("dLI_root", args.depth, args.width, args.int_cycles, args.ext_cycles, args.gen_mode, prep_time=1)
    elif args.model_type == "dHI":
        devstone_model = dHI("dHI_root", args.depth, args.width, args.int_cycles, args.ext_cycles, args.gen_mode, prep_time=1)

    env = DEVStoneEnvironment("DEVStoneEnvironment", devstone_model)
    model_created_time = time.time()

    # devstone_model.dot(f"{args.model_type}-before")

    import cProfile
    import pstats
    import io

    sim = Simulator(env)
    if args.classic:
        sim.setClassicDEVS(True)
    term = partial(terminate, timeLimit=1.5)
    if args.dynamic:
        sim.setDSDEVS(True)
        term = partial(terminate, timeLimit=args.width + 0.5)
    sim.setTerminationCondition(term)
    # sim.setVerbose(None)
    # sim.setTerminationTime(10.0)
    # sim.setStateSaving("custom")
    engine_setup_time = time.time()
    # pr = cProfile.Profile()
    # pr.enable()
    sim.simulate()
    # pr.disable()
    sim_time = time.time()

    # s = io.StringIO()
    # sortby = pstats.SortKey.CUMULATIVE
    # ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
    # ps.print_stats()
    # print(s.getvalue())

    # devstone_model.dot(f"{args.model_type}-after")

    print("Model creation time: {}".format(model_created_time - start_time))
    print("Engine setup time: {}".format(engine_setup_time - model_created_time))
    print("Simulation time: {}".format(sim_time - engine_setup_time))