text: from planner import *
  from logic.pref_scltl import *
  from collections import defaultdict
  import pickle
  import time
  
  
  class Gridworld(MdpGame):
      def __init__(self, name):
          super(Gridworld, self).__init__(name)
          self.max_x = None
          self.max_y = None
          self.slip_cells = dict()
          self.obs = set()
          self.pred_map = dict()
  
      def delta(self, u, a):
          successors = set()
          # Decouple state
          u_x, u_y, u_b = u
  
          # If battery is 0, don't move.
          if u_b == 0:
              successors.add((u_x, u_y, 0))
              return successors
  
          # Apply action
          if a == "N":
              v = (u_x, u_y + 1, u_b - 1)
          elif a == "E":
              v = (u_x + 1, u_y, u_b - 1)
          elif a == "S":
              v = (u_x, u_y - 1, u_b - 1)
          elif a == "W":
              v = (u_x - 1, u_y, u_b - 1)
          else:
              v = (u_x, u_y, u_b - 1)
  
          # If cell is obstacle or out of bounds, bounce back.
          v_x, v_y, v_b = v
          if (v_x, v_y) in self.obs or not (0 <= v_x < self.max_x) or not (0 <= v_y < self.max_y):
              v = (u_x, u_y, u_b - 1)
  
          # Add cell to successors
          successors.add(v)
  
          # If cell is slippery, add neighboring cells.
          v_b = u_b
          if (v_x, v_y) in self.slip_cells.keys():
              for d in self.slip_cells[(v_x, v_y)]:
                  successors = set.union(successors, self.delta((v_x, v_y, v_b), d))
  
          return successors
  
      def pred(self, v):
          return self.pred_map[v]
  
      def succ(self, u):
          successors = set()
          for a in actions:
              successors = set.union(successors, {(v, a) for v in self.delta(u, a)})
          return successors
  
      def construct(self, max_x, max_y, max_batt, actions, obs, slip_cells):
          self.max_x = max_x
          self.max_y = max_y
          self.obs = obs
          self.slip_cells = slip_cells
          states = set(itertools.product(range(max_x), range(max_y), range(max_batt + 1)))
          self.construct_symbolic(states=states, actions=actions)
  
          # Cache predecessor states
          # FIXME: same set() is shared to all keys. mutation problem!
          # self.pred_map = dict(zip(states, [set()] * len(states)))
  
          self.pred_map = dict()
          # Initialize dictionary
          for v in states:
              self.pred_map[v] = set()
  
          for u in states:
              for a in actions:
                  succ_u_a = self.delta(u, a)
                  for v in succ_u_a:
                      self.pred_map[v].add((u, a))
  
  
  def mdp_construction_test():
      # Manual testing [DONE OK.]
      # Move into slip cell
  
      print(gw.delta((1, 0, 1), "N"))
      print(gw.delta((2, 1, 4), "N"))
      print(gw.delta((2, 1, 4), "W"))
      print(gw.delta((2, 3, 4), "E"))
  
      # Move to non-slip cell
      print(gw.delta((3, 3, 4), "E"))
      print(gw.delta((2, 2, 4), "N"))
      print(gw.delta((2, 4, 3), "N"))
      print(gw.delta((2, 4, 3), "W"))
  
      # Check predecessor function
      print("predecessor test -- ")
      print(gw.pred((1, 1, 0)))
      print(gw.pred((4, 3, 3)))
      print(gw.pred((3, 4, 2)))
  
  
  def solve1(prod_mdp):
      strategy = synthesize_strategy(prod_mdp)
      with open(os.path.join("output", "comp_strategy1.pkl"), "wb") as pkl_file:
          pickle.dump(strategy, pkl_file)
  
  
  def solve2(prod_mdp):
      strategy = synthesize_strategy(prod_mdp)
      with open(os.path.join("output", "comp_strategy2.pkl"), "wb") as pkl_file:
          pickle.dump(strategy, pkl_file)
  
  
  def run_example(mdp, formula, fname):
      # Get preference automaton
      pref_dfa = formula.translate()
      print(pref_dfa)
  
      # Construct product MDP
      prod_mdp = ProductMdp(f"Product({mdp}, {pref_dfa})")
      prod_mdp.construct(mdp, pref_dfa)
  
      # Solve for competitive strategy
      strategy, win, rank_map = synthesize_strategy(prod_mdp)
  
      with open(os.path.join("output", fname), "wb") as pkl_file:
          save_dict = {
              "strategy": strategy,
              "win": win,
              "rank_map": rank_map
          }
          pickle.dump(save_dict, pkl_file)
  
  
  def analyze1(prod_mdp):
      with open(os.path.join("output", "comp_strategy1.pkl"), "rb") as pkl_file:
          strategy = pickle.load(pkl_file)
      # with open("win_0.pkl", "rb") as pkl_file:
      #     win0 = pickle.load(pkl_file)
      # with open("win_1.pkl", "rb") as pkl_file:
      #     win1 = pickle.load(pkl_file)
  
      print(f"pi[((2, 1, 4), (1, 1))]: {strategy[((2, 1, 4), (1, 1))]}")
      print(f"delta(((2, 1, 4), (1, 1))): {prod_mdp.delta(((2, 1, 4), (1, 1)), 'W')}")
      # print(f"pi[((3, 1, 3), (1, 1))]: {strategy[((3, 1, 3), (1, 1))]}")
      # print(f"delta(((3, 1, 3), (1, 1))): {prod_mdp.delta(((3, 1, 3), (1, 1)), 'E')}")
      # print(f"pi[((4, 1, 2), (1, 1))]: {strategy[((4, 1, 2), (1, 1))]}")
      # print(f"delta(((4, 1, 2), (1, 1))): {prod_mdp.delta(((4, 1, 2), (1, 1)), 'S')}")
  
      # print()
      # print(f"((2, 1, 4), (1, 1)) in win0: {((2, 1, 4), (1, 1)) in win0}. Expected: True.")
      # print(f"((1, 1, 3), (1, 1)) in win0: {((1, 1, 3), (1, 1)) in win0}. Expected: True")
      # print(f"((1, 2, 3), (1, 1)) in win0: {((1, 2, 3), (1, 1)) in win0}. Expected: True")
      # print(f"((1, 0, 3), (1, 1)) in win0: {((1, 0, 3), (1, 1)) in win0}. Expected: True")
      # print(f"((2, 1, 4), (1, 1)) in win0: {((1, 1, 3), (1, 1)) in win0}. If False -- this is unexpected.")
      # print(f"delta(((1, 2, 3), (1, 1)), 'W'): {prod_mdp.delta(((1, 2, 3), (1, 1)), 'W')}")
  
  
  def analyze2(prod_mdp):
      with open(os.path.join("output", "comp_strategy2.pkl"), "rb") as pkl_file:
          strategy = pickle.load(pkl_file)
      # with open("win_0.pkl", "rb") as pkl_file:
      #     win0 = pickle.load(pkl_file)
      # with open("win_1.pkl", "rb") as pkl_file:
      #     win1 = pickle.load(pkl_file)
  
      print(f"pi[((2, 1, 4), ((1, 1), (1, 1)))]: {strategy[((2, 1, 4), ((1, 1), (1, 1)))]}")
      print(f"delta(((2, 1, 4), ((1, 1), (1, 1))): {prod_mdp.delta(((2, 1, 4), ((1, 1), (1, 1))), 'W')}")
      for v in strategy:
          if len(strategy[v]) > 0:
              print(v, strategy[v])
  
  
  if __name__ == '__main__':
      # Define parameters.
      actions = {"N", "E", "S", "W"}
      obstacles = {(2, 0), (0, 1), (3, 2), (4, 2), (0, 3), (1, 3), (0, 4), (1, 4)}
      slippery_cells = {(1, 1): {"N", "S"}, (2, 2): {"N", "W"}, (3, 3): {"N", "W"}}
      label_map = {(0, 0): "e", (4, 0): "a", (0, 2): "b", (4, 3): "d", (3, 4): "c"}
  
      # Define gridworld.
      gw = Gridworld("GW-MDP")
      gw.construct(5, 5, 4, actions, obstacles, slippery_cells)
      gw.make_labeled({"a", "b", "c", "d", "e"}, lambda v: label_map[v[:2]] if v[:2] in label_map else set())
  
      # Solve
      # f = PrefScLTL("(Fb > Fa) & (Fe > Fa)")
      # run_example(mdp=gw, formula=f, fname="pi-simple-and")
  
      f = PrefScLTL("(Fa & (F(c & Fd))) > (Fa & F(b | c))")
      run_example(mdp=gw, formula=f, fname="pi-complex-temporal")
  
      # # Preference specification (part I)
      # f1 = PrefScLTL("F(b | c | e) > Fa")
      # pref_dfa1 = f1.translate()
      #
      # # Construct product MDP
      # prod_mdp1 = ProductMdp("product_mdp1")
      # prod_mdp1.construct(gw, pref_dfa1)
      #
      # # # Solve or analyze.
      # # solve1(prod_mdp1)
      # # analyze1(prod_mdp1)
  
      # Preference specification (part II)
      # f2 = PrefScLTL("(F(b | c | e) > Fa) & (F(d & F(c)) >= Fc)")
      # f2 = PrefScLTL("(F(b | c | e) > Fa) & (F(d & F(c)) >= Fc)")
      # pref_dfa2 = f2.translate()
      #
      # # Construct product MDP
      # prod_mdp2 = ProductMdp("product_mdp2")
      # prod_mdp2.construct(gw, pref_dfa2)
      #
      # # Solve or analyze.
      # # solve2(prod_mdp2)
      # analyze2(prod_mdp2)
  