try:
    import grid2op
    import matplotlib
    import numpy as np
    import pandas as pd
    import matplotlib.pyplot as plt
    import warnings
    import os
    from grid2op.MakeEnv import make
    import itertools
    import os.path
except ImportError as exc_:
    raise ImportError("AsynchronousActorCritic baseline impossible to load the required dependencies. The error was: \n {}".format(exc_))


def prune_action_space(bus_no):
    # Initialize the env."case5_example",chronics_path=os.path.join("public_data", "chronics_5bus_example")
    if bus_no == "14":
        environment = make()
        num_substations = 14 # 14 for 14 bus
    elif bus_no == "5":
        environment = make("case5_example",chronics_path=os.path.join("public_data", "chronics_5bus_example"))
        num_substations = 5
    else:
        print("Unexpected bus size of the power system. Exiting code in Action_reduced_list.py")
        exit()

    from pandapower.plotting.plotly import simple_plotly
    # simple_plotly(environment.backend._grid)
    # Load ids
    print("\nInjection information:")
    load_to_subid = environment.action_space.load_to_subid
    print ('There are {} loads connected to substations with id: {}'.format(len(load_to_subid), load_to_subid))

    # Generators irds
    gen_to_subid = environment.action_space.gen_to_subid
    print ('There are {} generators, connected to substations with id: {}'.format(len(gen_to_subid), gen_to_subid))

    # Line id sender
    print("\nPowerline information:")
    lines_or_to_subid = environment.action_space.lines_or_to_subid
    lines_ex_to_subid = environment.action_space.lines_ex_to_subid

    print ('There are {} transmissions lines on this grid.'.format(len(lines_or_to_subid)))


    gen_at_sub = []
    load_at_sub = []
    lines_or_at_sub = []
    lines_ex_at_sub = []
    for k in range(num_substations):
        gen_at_sub.append(np.arange(len(gen_to_subid))[gen_to_subid==k])
        load_at_sub.append(np.arange(len(load_to_subid))[load_to_subid==k])
        lines_or_at_sub.append(np.arange(len(lines_or_to_subid))[lines_or_to_subid==k])
        lines_ex_at_sub.append(np.arange(len(lines_ex_to_subid))[lines_ex_to_subid==k])
    1;
    # Num of elements per SE
    print("\nSubstations information:")
    for i, nb_el in enumerate(environment.action_space.subs_info):
        print("On susbtation {} there are {} elements.".format(i, nb_el))

    # adding the change_actions at for 1 substation at a time . Initializing with no action
    gen_action_list =[[]] # list of gens acted
    load_action_list = [[]]
    line_or_action_list = [[]]
    line_ex_action_list = [[]]
    substation_acted = [[]] # substation acted
    gen_action_list_bus_2 = [[]]
    load_action_list_bus_2 = [[]]
    line_or_action_list_bus_2 = [[]]
    line_ex_action_list_bus_2 = [[]]
    do_nothing_act = environment.helper_action_player({})
    obs, reward, done, info = environment.step(do_nothing_act)
    for sub_id in range(num_substations):
        num_gen_at_sub = len(gen_at_sub[sub_id] )
        num_load_at_sub = len(load_at_sub[sub_id] )
        num_line_or_at_sub = len(lines_or_at_sub[sub_id] )
        num_line_ex_at_sub = len(lines_ex_at_sub[sub_id] )
        switching_patterns = ["".join(seq) for seq in itertools.product("01",repeat=environment.action_space.subs_info[sub_id] - 1)]  # reduce by 1 bit due to compliment being the same
        switching_patterns = [[int(sw_i_k) for sw_i_k in '0' + sw_i] for sw_i in switching_patterns]  # adding back the '0' at the beginning as we are fizing this bit
        switching_patterns.pop(0) # deleting first action as it is a no-action
        for sw_action in switching_patterns:
            switching_patterns_split = np.split(np.array(sw_action), np.cumsum([num_gen_at_sub,num_load_at_sub,num_line_or_at_sub,num_line_ex_at_sub]) )
            gen_action_list.append(gen_at_sub[sub_id][switching_patterns_split[0] == 1])
            load_action_list.append(load_at_sub[sub_id][switching_patterns_split[1] == 1])
            line_or_action_list.append(lines_or_at_sub[sub_id][switching_patterns_split[2] == 1])
            line_ex_action_list.append(lines_ex_at_sub[sub_id][switching_patterns_split[3] == 1])
            substation_acted.append(sub_id)
            # gen_action_list_bus_2.append(gen_at_sub[sub_id][switching_patterns_split[0] == 0])
            # load_action_list_bus_2.append(load_at_sub[sub_id][switching_patterns_split[1] == 0])
            # line_or_action_list_bus_2.append(lines_or_at_sub[sub_id][switching_patterns_split[2] == 0])
            # line_ex_action_list_bus_2.append(lines_ex_at_sub[sub_id][switching_patterns_split[3] == 0])
            # do_act = environment.action_space(
            #     {"set_bus": {"generators_id": [(g,1) for g in gen_action_list[-1]]+[(g,2) for g in gen_action_list_bus_2[-1]],
            #                  "loads_id": [(lo,1) for lo in load_action_list[-1]]+[(lo,2) for lo in load_action_list_bus_2[-1]],
            #                  "lines_or_id": [(li,1) for li in line_or_action_list[-1]]+[(li,2) for li in line_or_action_list_bus_2[-1]],
            #                  "lines_ex_id": [(li,1) for li in line_ex_action_list[-1]]+[(li,2) for li in line_ex_action_list_bus_2[-1]]}})
            # print(do_act)
            # obs_sim, reward_sim, is_done_sim, info_sim = obs.simulate(do_act)
            # print(obs_sim)
            # print(is_done_sim)
            # if is_done_sim:
            #     print("----------------------")
            # print(info_sim)
            # h_letters = [(letter, 2) for letter in 'human']
            # print(h_letters)

            # do_nothing_act = environment.action_space({"change_bus":{"generators_id": [0],"loads_id": [1],"lines_or_id":[3],"lines_ex_id":[7]}})
            # obs_sim, reward_sim, is_done_sim, info_sim = obs.simulate(do_nothing_act)
    #
    action_index = 1 # max is 60 for 5 bus
    do_act = environment.action_space(
        {"change_bus": {"generators_id": gen_action_list[action_index], "loads_id": load_action_list[action_index], "lines_or_id": line_or_action_list[action_index], "lines_ex_id": line_ex_action_list[action_index]}})
    print(do_act)
    obs_sim, reward_sim, is_done_sim, info_sim = obs.simulate(do_act)
    k = 1
    # environment.action_space({"change_bus":{"generators_id": [0],"loads_id": [1],"lines_or_id":[3],"lines_ex_id":[7]}})
    np.save("gen_action_list.npy", gen_action_list)
    np.save("load_action_list.npy", load_action_list)
    np.save("line_or_action_list.npy", line_or_action_list)
    np.save("line_ex_action_list.npy", line_ex_action_list)
    return gen_action_list, load_action_list, line_or_action_list, line_ex_action_list

def main_function(bus_no):
    # There are a total of 4 files.
    file_names = ["gen_action_list","load_action_list","line_or_action_list","line_ex_action_list"]
    flag_check = None
    for index_item, item in enumerate(file_names):
        flag_check = os.path.isfile(item+".npy")
        if flag_check == False:
            # run the action pruning code!
            gen_action_list, load_action_list, line_or_action_list, line_ex_action_list = prune_action_space(bus_no)
            return gen_action_list, load_action_list, line_or_action_list, line_ex_action_list
    if flag_check == True:
        # load the data from the saved files.
        gen_action_list = np.load("gen_action_list.npy",allow_pickle=True)
        load_action_list = np.load("load_action_list.npy",allow_pickle=True)
        line_or_action_list = np.load("line_or_action_list.npy",allow_pickle=True)
        line_ex_action_list = np.load("line_ex_action_list.npy",allow_pickle=True)
    return gen_action_list, load_action_list, line_or_action_list, line_ex_action_list
