# -*- coding: utf-8 -*-
import os
import json
import time
import logging
from func_set.trans_sys import compute_trans_graph
from func_set.trans_sys2 import compute_trans_graph2
from func_set.gltl2ba import ltl_formula_to_ba
from func_set.sampling_func import sampling_ltl_planning, samp_get_robot_path
from func_set.show_graph import nx_to_graphviz_tree, nx_to_graphviz_trans
from func_set.general_func import pts2AP, AP2pts
import pickle  


if __name__ == '__main__':
    LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
    DATE_FORMAT = "%m/%d/%Y %H:%M:%S %p"
    logging.basicConfig(filename="samp_planning.log",
                        filemode="w",
                        format=LOG_FORMAT,
                        datefmt=DATE_FORMAT,
                        level=logging.DEBUG)
    logging.info("   ")

    for txt_name in ['/home/luojiongts/sampling_ltl_planning/robot_ts/war_task.txt']:
        logging.info('-------------------------------------------------')

        with open(txt_name, 'r') as f:
            config = json.loads(f.read())

        itera_pre_num = config['itera_pre_num']
        itera_suf_num = config['itera_suf_num']
        logging.info(f'iteration number: {itera_pre_num,itera_suf_num}')
        
        # path weight
        path_weight = config['path_weight']
        logging.info(f'path weight: {path_weight}')

        # Transition system
        trans_graph_list = []
        robot_id =  1
        for robot_config_file in config['robots_config_file']:
            if robot_id <= 2:
                trans_graph_list.append(compute_trans_graph2(robot_config_file))
            elif robot_id > 2 and robot_id <=6:
                trans_graph_list.append(compute_trans_graph(robot_config_file))
            print(robot_id)
            robot_id = robot_id + 1
        pts_state_num = 1
        for trans_graph in trans_graph_list:
            pts_state_num *= len(trans_graph)
        print(f'pts state number: {pts_state_num}')
        logging.info(f'pts state number: {pts_state_num}')
       
        # Init position 
        robots_init_pos = config['robots_init_pos']
        robot_number = len(robots_init_pos)
        print(f'robot number: {robot_number}')
        logging.info(f'robot number: {robot_number}')

        # LTL task formula
        # if config['is_task_surveillance'] == 'False':
        #     is_surveillance = False
        # else:
        #     is_surveillance = True
        task = config['ltl_task']
        logging.info(f'ltl task: {task}')
        if "[]<>" in task or "[] <>" in task:
            is_surveillance = True
        else:
            is_surveillance = False

        # os.getcwd get current work directory
        LTL_FILE_POS = os.getcwd()+'/clasc_ltlFile.txt'

        # convert ltl to buchi automaton
        print("start converting buchi automaton:")
        [buchi_graph, buchi_init_states, buchi_accept_states,
            buchi_dot_graph] = ltl_formula_to_ba(task, LTL_FILE_POS)
        # buchi_init_states = ['T0_S2']
        print(f'buchi graph size: {len(buchi_graph)}')
        logging.info(f'buchi graph size: {len(buchi_graph)}')
        print(f'product automaton size: {len(buchi_graph)*pts_state_num}')
        logging.info(
            f'product automaton size: {len(buchi_graph)*pts_state_num}')

        print("start sampling planning:")
        samp_start_time = time.time()
        # a = AP2pts(['s1','g1c','g2c'], trans_graph_list, 2)
        # print(a)
        ok_num = 0
        itera_test = [20,50,80,100,150,200,300,400,500]
        time_cost = [] 
        # with open('data/cost_lists2.pkl', 'wb') as file:
        with open('cost_lists.pkl', 'wb') as file:
            # for w_cost data
            for i in range(1):
                
                [optimal_path, optimal_path_cost,pre_search_tree,history_min_cost] = sampling_ltl_planning(robots_init_pos, trans_graph_list,
                                                                        buchi_graph, buchi_init_states,
                                                                        is_surveillance,
                                                                        itera_pre_num, itera_suf_num,
                                                                        path_weight, samp_start_time)
                if len(optimal_path[0]) != 0:  
                    pickle.dump(history_min_cost, file) 
                    ok_num = ok_num+1 
                
            # for time_cost data    
            # for i in range(len(itera_test)):
            #     ok_flag = 0
            #     while ok_flag == 0:
            #         itera_pre_num = itera_test[i]
            #         samp_start_time = time.time()
            #         [optimal_path, optimal_path_cost,pre_search_tree,history_min_cost] = sampling_ltl_planning(robots_init_pos, trans_graph_list,
            #                                                                 buchi_graph, buchi_init_states,
            #                                                                 is_surveillance,
            #                                                                 itera_pre_num, itera_suf_num,
            #                                                                 path_weight, samp_start_time)
            #         samp_end_time = time.time()
            #         if len(optimal_path[0]) != 0 or i<2:
            #             ok_flag = 1
            #     time_cost.append(samp_end_time-samp_start_time)
            # time_cost=[0.010216951370239258, 0.10536646842956543, 0.7453296184539795, 1.0438737869262695, 0.6006045341491699, 2.837425708770752, 4.431213855743408, 6.666805744171143, 8.35905933380127]    
            # pickle.dump(time_cost, file)
        print(time_cost)        
        # record path for every robot
        a = 1
        robot_path = samp_get_robot_path(optimal_path, robot_number)

        samp_end_time = time.time()
        print(f'computation time: {samp_end_time-samp_start_time} seconds.')
        logging.info(
            f'computation time: {samp_end_time-samp_start_time} seconds.')

        for i, path in enumerate(robot_path):
            print(f'Robot{i}: {path[0]} + {path[1]}')
            logging.info(f'Robot{i}: {path[0]} + {path[1]}')
        print(f'best path weight: {sum(optimal_path_cost)}')
        print(f'pre weight: {optimal_path_cost[0]}')
        print(f'suf weight: {optimal_path_cost[1]}')
        logging.info(f'best path weight: {sum(optimal_path_cost)}')
        logging.info(f'pre weight: {optimal_path_cost[0]}')
        logging.info(f'suf weight: {optimal_path_cost[1]}')
        print ('ok_num:',ok_num )

        # trans_dot_graph1 = nx_to_graphviz_trans(trans_graph_list[0])
        # trans_dot_graph1.show('samp_trans_graph1')
        # trans_dot_graph2 = nx_to_graphviz_trans(trans_graph_list[1])
        # trans_dot_graph2.show('samp_trans_graph2')
        buchi_dot_graph.show('clasc_buchi_graph')
        search_dot_pre_tree = nx_to_graphviz_tree(pre_search_tree)
        search_dot_pre_tree.show('samp_search_tree')
        
        # 保存到pickle文件  
        # with open('data/cost_lists.pkl', 'wb') as file:  
        #     pickle.dump(history_min_cost, file)  
            # pickle.dump(list2, file)  
            # pickle.dump(list3, file)