import random
import matplotlib.pyplot as plt

random.seed(0)

"""
Warm Up Questions
"""

class Animal:
  def __init__(self, name, age):
    self.name = name
    self.age = age
    
  def get_age(self):
    return self.age
  
  def eat(self):
    return "that was yummy"

class Dog(Animal):
  def __init__(self, name, age):
    super().__init__(name, age)
    self.num_legs = 4

  def get_age(self):
    return self.age * 7
  
  def drink(self):
    return "that was refreshing"
  
   
dog = Dog("Max", 3)
zebra = Animal("Daisy", 3)

# Q: What will dog.get_age() output?
# Q: What will zebra.get_age() output?
print(f"Max's age: {dog.get_age()}")
print(f"Daisy's age: {zebra.get_age()}")

"""
Q: Which of the following will throw an error (can be more than one)?
Answer choices:
A - dog.drink()
B - zebra.drink()
C - dog.eat()
D - zebra.eat()
E - dog.num_legs
F - zebra.num_legs
G - None of the above

Answer: C, eat() only exists in Dog class, F, num legs is only an atrtibute of Dog
"""







### Stochastic Thinking – Random Walk in Boston
### Exploring N Friends Visiting M Locations
"""
Problem Setup
You find yourself in Boston with N-1 other people. There are M landmarks scattered around
the city like Fenway Park, MIT Dome, etc. (modeled as a 2D rectangular grid).

- Each person starts at a random location
- Each step, a person moves up, down, left, or right (bounded by the grid), according to their move function
- Goal: Simulate how long it takes until every one of the M locations has been visited
        by at least one of the N people (collectively)
"""

class City:
  """
  City has a set of M landmarks and N people
  Defined on a 2D rectangular grid from (0, 0) to (max_x, max_y)
  """
  
  def __init__(self, max_x, max_y, landmarks=None, people=None):
    """
    max_x, max_y (int)
    landmarks (set)           
    people (set)
    """
    if not landmarks:
      self.landmarks = set()
    else:
      self.landmarks = landmarks
    if not people:
      self.people = set()
    else:
      self.people = people

    self.min_x = 0
    self.max_x = max_x

    self.min_y = 0
    self.max_y = max_y

  def add_landmark(self, landmark):
    if type(landmark) == Landmark:
      self.landmarks.add(landmark)

  def add_person(self, person):
    if isinstance(person, Person):
      self.people.add(person)

  def get_random_location(self):
    return (random.randint(0, self.max_x), random.randint(0, self.max_y))

  def get_landmarks(self):
    return self.landmarks

  def get_people(self):
    return self.people
  
  def all_landmarks_visited(self):
    """
    Return true if all landmarks are visited
    Otherwise, return false
    """
    for landmark in self.get_landmarks():
      if not landmark.has_visited():
          return False
    return True

  def get_distance(self, current_pos, new_pos):
    """
    Assume this is movable distance not shortest path distance
    They can move up, down, left, right, not diagonal
    """

    curr_x, curr_y = current_pos
    new_x, new_y = new_pos

    diff_x = abs(new_x - curr_x)
    diff_y = abs(new_y - curr_y)

    return diff_x + diff_y

  def reset(self):
    for landmark in self.landmarks:
      landmark.reset_visited()

    for person in self.people:
      person.set_random_location()

class Landmark:
  """
  Each landmark is at a stationary location in a city  
  """
  def __init__(self, name, loc):
    self.is_visited = False
    self.name = name
    self.loc = loc

  def mark_visited(self):
    self.is_visited = True

  def reset_visited(self):
    self.is_visited = False

  def has_visited(self):
    return self.is_visited

  def get_location(self):
    return self.loc

# Now we will implement classes for different kinds of people in a city

class Person:
    person_id = 0

    def __init__(self, env):
      # Task 1: initalize environment attribute and make loc attribute a random location
      self.env = env
      self.loc = self.env.get_random_location()

      # Assigns unique ID to each person
      # Q: Why can we not just do person_id = 1?
      # A: It is a class attribute, we want the updated value to be the same for all
      #    instances of Person class
      Person.person_id += 1
      self.id = Person.person_id

    def set_random_location(self):
      self.loc = self.env.get_random_location()

    def get_location(self):
      return self.loc
    
    def set_location(self,loc):
      self.loc = loc

    def get_x(self):
      return self.loc[0]

    def get_y(self):
      return self.loc[1]
     
    def get_closest_landmark(self):
      """
      Return the Landmark object corresponding to the closest unvisited landmark
      """
      closest_landmark = None
      min_distance = float("inf")
      my_loc = self.get_location()
      for landmark in self.env.get_landmarks():
          if not landmark.has_visited():
              landmark_loc = landmark.get_location()
              distance = self.env.get_distance(my_loc, landmark_loc)
              if distance < min_distance:
                  min_distance = distance
                  closest_landmark = landmark
      return closest_landmark



# Task 2: Implement Local Person
"""
Implement a local city dweller
They should inherit from the Person class
"""
class LocalPerson(Person):
    def __init__(self, env):
        super().__init__(env)

    def move(self):
        """
        Moves randomly at each time step
        """
        directions = [(0,1), (0,-1), (1,0), (-1,0)] # Right, Left, Down, Up

        dx, dy = random.choice(directions)
        new_x = max(0, min(self.env.max_x, self.get_x() + dx))
        new_y = max(0, min(self.env.max_y, self.get_y() + dy))

        closest_landmark = self.get_closest_landmark()
        if closest_landmark is not None:
          if (new_x, new_y) == closest_landmark.get_location():
              closest_landmark.mark_visited()
        return (new_x, new_y)


# Task 3: Implement Tourist Person
"""
Implement a tourist in the city
They should inherit from the Person class
"""
class TouristPerson(Person):
    def __init__(self, env):
        super().__init__(env)

    def move(self):
        """
        Moves one step (priortizing x) towards the nearest unvisited landmark
        """
        closest_landmark = self.get_closest_landmark()

        if closest_landmark is not None:
          dx = closest_landmark.get_location()[0] - self.get_x()
          dy = closest_landmark.get_location()[1] - self.get_y()

          new_x = self.get_x()
          new_y = self.get_y()
          if abs(dx) >= abs(dy):
              if closest_landmark.get_location()[0] > self.get_x():
                  new_x += 1
              else:
                  new_x -= 1
          elif abs(dx) < abs(dy):
              if closest_landmark.get_location()[1] > self.get_y():
                  new_y += 1
              else:
                  new_y -= 1

          if (new_x, new_y) == closest_landmark.get_location():
              closest_landmark.mark_visited()

          return (new_x, new_y)
        return self.get_location()

# Task 4: Implement Simulation Function	        
"""
Simulation Function
This function simulates the exploration process and returns the average number of steps
needed to visit all M locations by at least one person (essentially: collectively with N people).
"""
def simulate_exploration(N, p, trials=100, max_steps=1000):
    """
    N is number of total people in the city
    p is proportion of tourists
    """
    steps_list = []
    landmark_names = ["Lobby 7", "Killian Court", "Walker Memorial", "Next House", "Fenway Park", "Boston Commons"]
    for _ in range(trials):

        boston = City(150, 200)

        #Add landmark objects to city object at a random location
        for landmark_name in landmark_names:
          boston.add_landmark(Landmark(landmark_name, boston.get_random_location()))

        # Add person objects to city object at a random location
        num_tourists = int(N*p)
        num_locals = N - num_tourists
        for _ in range(num_locals):
          boston.add_person(LocalPerson(boston))
        for _ in range(num_tourists):
          boston.add_person(TouristPerson(boston))

        step = 0

        while not boston.all_landmarks_visited() and step < max_steps:
            for person in boston.get_people():
                new_pos = person.move()
                person.set_location(new_pos)
            step += 1
        steps_list.append(step)
        
    return sum(steps_list) / len(steps_list)

"""
Demo 1: Varying proportion of tourists (p)
"""

N = 20
p_vals = [.1*i for i in range(11)]
avg_steps = [simulate_exploration(N, p) for p in p_vals]

plt.plot(p_vals, avg_steps, marker='o')
plt.title(f"Expected Steps vs Proportion of Tourists (N={N})")
plt.xlabel("Proportion of Tourists (p)")
plt.ylabel("Expected Steps")
plt.grid(True)
plt.show()

"""
What happens to the average number of steps to visit all M locations when M increases? Is it a linear increase?
As the number of locations (M) increases, the average number of steps required increases as well.
- The relationship is not linear. As M increases, it becomes increasingly difficult to stumble upon all new locations.
- Some locations may be missed for a long time due to randomness.
"""


"""
Demo 2: Varying Number of Total People (N)
"""

p = 0.5 # Fix 10 locations of interest
N_vals = list(range(1, 21))
avg_steps_n = [simulate_exploration(N, p) for N in N_vals]

plt.plot(N_vals, avg_steps_n, marker='o', color='orange')
plt.title(f"Expected Steps vs Number of People (p={p})")
plt.xlabel("Number of People (N)")
plt.ylabel("Expected Steps")
plt.grid(True)
plt.show()

"""
Q: What happens to the average number of steps to visit all M locations when N increases?
A: As the number of friends increases, the expected number of steops to cover all M landmarks
   drops significantly, especially at first.
   - However, there are diminishing returns. Adding more friends beyond a certain point
    doesn't help as much.
"""