# jugs.py
# Andrew Davison, Dec 2023, ad@coe.psu.ac.th
# https://en.wikipedia.org/wiki/Water_pouring_puzzle
'''
Start with two jugs of size MAX_JUG1 and MAX_JUG2, 
and use the operations to get a jug containing GOAL.

In this case, MAX_JUG1 = 5
MAX_JUG2 = 3
GOAL = 4

The operations: 
  * Empty a jug (X, 0)->(0, 0) 
  * Fill a Jug, (0, 0)->(X, 0)
  * Pour water from one jug to the other until one of the jugs is either empty or full, (X, Y) -> (X-d, Y+d)

  state = (jug1, jug2)
  start state: (0,0)
  final state: (0, GOAL) or (GOAL, 0)
'''
import sys
import searchUtils
import matplotlib.pyplot as plt


# problem-specific ---------------

def isGoal(state):
  jug1, jug2 = state
  return ((jug1 == GOAL) and (jug2 == 0)) or \
         ((jug1 == 0) and (jug2 == GOAL))


def nextStates(state):
  jug1, jug2 = state
  sts = []

  # fill a jug
  sts.append((MAX_JUG1, jug2)) # Fill jug1
  sts.append((jug1, MAX_JUG2)) # Fill jug2

  # empty a jug
  sts.append((0, jug2)) # empty jug1
  sts.append((jug1, 0)) # empty jug2

  # pour from jug1 to jug2
  filljug2 = MAX_JUG2 - jug2
  if filljug2 > jug1:
    sts.append((0, jug2+jug1)) # Fill jug2 until jug1 empty
  else:
    sts.append((jug1-filljug2, MAX_JUG2)) # Fill jug2

  # pour from jug2 to jug1
  filljug1 = MAX_JUG1 - jug1
  if filljug1 > jug2:
    sts.append((jug1+jug2, 0)) # Fill jug1 until jug2 empty
  else:
    sts.append((MAX_JUG1, jug2-filljug1)) # Fill jug1

  return sts


def printPath(path):
  prevState = path[0]
  print("Path:")
  print(prevState)

  for state in path[1:]:
    pjug1, pjug2 = prevState
    jug1, jug2 = state

    # report decreases before increases
    if jug1 < pjug1:
      print("Decrease jug1 by", pjug1-jug1)
    if jug2 < pjug2:
      print("Decrease jug2 by", pjug2-jug2)

    if jug1 > pjug1:
      print("Increase jug1 by", jug1-pjug1)
    if jug2 > pjug2:
      print("Increase jug2 by", jug2-pjug2)

    print(state)
    prevState = state


# override the imported dummy functions
searchUtils.isGoal = isGoal
searchUtils.nextStates = nextStates


# ------------- main ---------------

MAX_JUG1 = 5 # 4
MAX_JUG2 = 3 # 3
GOAL = 4 # 2

print("Max of jug1:", MAX_JUG1)
print("Max of jug2:", MAX_JUG2)
print("Goal:", GOAL)


path, numVisited = searchUtils.bfs((0,0))

print("No. of states visited:", numVisited)
if path == None:
  print("No path found");
  sys.exit()

printPath(path)


# ---------- plot state space  ---------------

xs, ys = zip(*path)

plt.plot(xs, ys)
plt.title("Jugs Puzzle")
plt.xticks(range(6))
plt.yticks(range(4))
plt.xlabel("5L Jug")
plt.ylabel("3L Jug")
plt.show()