# %%
import simpy
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

RANDOM_SEED = 42
random.seed(RANDOM_SEED)

WALK_RADIUS = 50
# TO-DO:
# parametrize speeds
# include gis maps

#####    INPUT OF DATA #####
data = {
    # Size of the grid
    "grid": {"xlim": [0, 100], "ylim": [0, 100]},
    # Docking stations with their location and capacity
    "stations": [{"location": [5, 5], "capacity": 3}, {"location": [10, 10], "capacity": 3}, {"location": [90, 90], "capacity": 2}, {"location": [95, 95], "capacity": 2},],
    # All the bikes (0=SB; 1=Dockless,2=Autonomous) and their initial locations
    "bikes": [
        {"mode": 0, "station": 0},
        {"mode": 0, "station": 0},
        {"mode": 0, "station": 1},
        {"mode": 0, "station": 2},
        {"mode": 0, "station": 3},
        {"mode": 1, "location": [10, 5]},
        {"mode": 1, "location": [10, 5]},
        {"mode": 1, "location": [10, 5]},
        {"mode": 2, "location": [5, 15]},
        {"mode": 2, "location": [5, 15]},
        {"mode": 2, "location": [5, 15]},
    ],
    # Number of agents
    "nagents": 5,
    # 'agents': [{'source','target','time_source','time_target','mode'}]
}

####### DEFINITION OF CLASSES (City, Staion, Bike, Agent) ######
class City:
    def __init__(self, env, data):  # Constructor is automatically called when you create a new instance of a class
        self.env = env  # Takes simulation environment
        self.data = data  # Its the city that takes the initial data and sets it

        self.stations = []  # defines arrays of objects of different classes within itself
        self.bikes = []
        self.agents = []

        self.start()

    def start(self):  # Start calls to the initialization of station, bikes, agents
        self.init_stations()
        self.init_bikes()
        self.init_agents()

    def init_stations(self):
        # Takes data of stations and goes one by one adjudicating the data to the station objects
        for station_id, station_data in enumerate(self.data["stations"]):
            station = Station(self.env, station_id)  # Takes data of stations and goes one by one adjudicating the data to the station objects (that now have an id)
            station.set_capacity(station_data["capacity"])  # asjudicating data
            station.set_location(np.array(station_data["location"]))  # " np.array() helps to structure data ???
            station.init_station()  # initialize that station (resource and container)
            self.stations.append(station)  # Adds station to the staions array in city

    def init_bikes(self):
        for bike_id, bike_data in enumerate(self.data["bikes"]):  # There is not Dockless 1 and Autonomous 1, they share ids !
            mode = bike_data["mode"]  # Takes the mode
            if mode == 0:
                bike = StationBike(self.env, bike_id)
                station_id = bike_data["station"]
                self.stations[station_id].push_bike(bike_id)  # Saves the bike in the station
                bike.set_station(station_id)  # saves the station in the bike
                bike.set_location(self.stations[station_id].location)  # Saves the location of the station as its location
            elif mode == 1:
                bike = DocklessBike(self.env, bike_id)  # creates object of class Dockless Bike with the env and the id as imput
                bike.set_location(np.array(bike_data["location"]))  # initial location is given, so it takes it and saves it
            elif mode == 2:
                bike = AutonomousBike(self.env, bike_id)  # creates object of class Autonomous Bike with the env and the id as imput
                bike.set_location(np.array(bike_data["location"]))  # initial location is given, so it takes it and saves it
            self.bikes.append(bike)  # Adds bike to the array of Bikes in city

    def init_agents(self):
        for agent_id in range(self.data["nagents"]):  # here the only input is the number of agents
            mode = np.random.randint(0, 2)  # assigns a random mode
            # mode = 2
            # Creates a different class of agent depending in the mode
            if mode == 0:
                agent = StationAgent(self.env, agent_id)
            elif mode == 1:
                agent = DocklessAgent(self.env, agent_id)
            elif mode == 2:
                agent = AutonomousAgent(self.env, agent_id)
            agent.set_data(self.data["grid"], self.stations, self.bikes)  # To 'know' the city / It needs this data to operate # TO-DO ???¿
            agent.start()  # Defines structure for location, bike id and event setup ?
            self.agents.append(agent)  # Adds the agent to the array of agents in city

class Station:
    def __init__(self, env, station_id):  # constructor
        self.env = env
        self.station_id = station_id

        self.location = None
        self.capacity = None
        self.bikes = []

        self.history = []
        self.init_state()

    def init_state(self):
        state_keys = ["time", "bikes", "nbikes"]
        self.state = dict.fromkeys(state_keys)

    def init_station(self):  # parameters that help to manage the station
        self.resource = simpy.Resource(self.env, capacity=1)  # Put some order in case 2 arrite at the same time Do we need this ? If all slots are taken, requests are enqueued. Once a usage request is released, a pending request will be triggered
        self.container = simpy.Container(self.env, self.capacity, init=0)  # It supports requests to put or get matter into/from the container

    def set_location(self, location):  # takes location defined in City class definition and saves it
        self.location = location

    def set_capacity(self, capacity):  # takes capacity defined in City class definition and saves it
        self.capacity = capacity

    def has_bikes(self):
        return self.container.level > 0

    def has_docks(self):
        return self.capacity - self.container.level > 0

    def empty(self):
        return self.container.level == 0

    def full(self):
        return self.container.level == self.capacity

    def choose_bike(self):  # Selects any bike (Random package)
        return random.choice(self.bikes)

    def pull_bike(self, bike_id):
        if self.has_bikes():  # Check that it has bikes
            self.container.get(1)  # Counts that there is one less
            self.bikes.remove(bike_id)  # removes that bike from the list
            self.save_state()  # saves time/bikes/nbikes in history -> FORMAT? and then init_state??
            self.print_info()  # prints time/station/nbikes/capacity
        else:
            print("[%.2f] Station %d has no bikes available" % (self.env.now, self.station_id))

    def push_bike(self, bike_id):
        if self.has_docks():  # check that there are docks available
            self.container.put(1)  # add one item
            self.bikes.append(bike_id)  # add the bike's id
            self.save_state()  # saves time/bikes/nbikes in history
            self.print_info()  # prints time/station/nbikes/capacity
        else:
            print("[%.2f] Station %d has no docks available" % (self.env.now, self.station_id))

    def save_state(self):  # FORMAT?
        # print(self.station_id, np.round(self.env.now), self.bikes, len(self.bikes))
        self.state["time"] = np.round(self.env.now)
        self.state["bikes"] = self.bikes.copy()
        self.state["nbikes"] = len(self.bikes)
        self.history.append(self.state)
        self.init_state()

    def print_info(self):  # Prints time, station, and nbikes/capacity
        print("[%.2f] Station %d has %d bikes out of %d" % (self.env.now, self.station_id, self.container.level, self.capacity))

class Bike:
    def __init__(self, env, bike_id):  # constructor
        self.env = env
        self.bike_id = bike_id

        self.agent_id = None
        self.location = None

        self.history = []
        self.print = False

    def init_state(self):
        state_keys = [
            "agent_id",
            "mode",
            "station_pull",
            "time_pull",
            "station_push",
            "time_push",
            "location_unlock",
            "time_unlock",
            "location_drop",
            "time_drop",
            "location_call",
            "time_call",
            "location_meet",
            "time_meet",
        ]
        self.state = dict.fromkeys(state_keys)

    def set_agent(self, agent_id):
        self.agent_id = agent_id

    def set_location(self, location):
        self.location = location

    def pop_agent(self):
        self.agent_id = None

    def rented(self):
        return self.agent_id is not None

    def dist(self, a, b):
        return np.linalg.norm(a - b)

    def ride(self, location):
        distance = self.dist(self.location, location)
        yield self.env.timeout(distance)
        self.location = location

class StationBike(Bike):
    def __init__(self, env, bike_id):
        super().__init__(env, bike_id)
        self.station_id = None
        self.init_state()

    def init_state(self):
        state_keys = [
            "agent_id",
            "station_pull",
            "time_pull",
            "station_push",
            "time_push",
        ]
        self.state = dict.fromkeys(state_keys)

    def save_state(self):
        self.history.append(self.state)
        self.init_state()

    def set_station(self, station_id):
        self.station_id = station_id

    def pop_station(self):
        self.station_id = None

    def register_pull(self, agent_id):
        if self.print:
            print("[%.2f] Bike %d pulled from station %d by agent %d" % (self.env.now, self.bike_id, self.station_id, agent_id))

        self.state["agent_id"] = agent_id
        self.state["station_pull"] = self.station_id
        self.state["time_pull"] = np.round(self.env.now, 2)
        self.set_agent(agent_id)
        self.pop_station()

    def register_push(self, station_id):
        if self.print:
            print("[%.2f] Bike %d pushed on station %d by agent %d" % (self.env.now, self.bike_id, station_id, self.agent_id))

        self.state["station_push"] = station_id
        self.state["time_push"] = np.round(self.env.now, 2)
        self.save_state()
        self.pop_agent()
        self.set_station(station_id)

    def docked(self):
        return self.station_id is not None

class DocklessBike(Bike):
    def __init__(self, env, bike_id):
        super().__init__(env, bike_id)
        self.init_state()

    def init_state(self):
        state_keys = [
            "agent_id",
            "location_unlock",
            "time_unlock",
            "location_lock",
            "time_lock",
        ]
        self.state = dict.fromkeys(state_keys)

    def save_state(self):
        self.history.append(self.state)
        self.init_state()

    def unlock(self, agent_id):
        self.set_agent(agent_id)
        self.state["agent_id"] = agent_id
        self.state["time_unlock"] = np.round(self.env.now, 2)
        self.state["location_unlock"] = self.location

        if self.print:
            print("[%.2f] Dockless bike %d unlocked by agent %d on location [%.2f, %.2f]" % (self.env.now, self.bike_id, agent_id, *self.location))

    def lock(self):
        if self.print:
            print("[%.2f] Dockless bike %d locked by agent %d on location [%.2f, %.2f]" % (self.env.now, self.bike_id, self.agent_id, *self.location))

        self.state["location_lock"] = self.location
        self.state["time_lock"] = np.round(self.env.now, 2)
        self.pop_agent()
        self.save_state()

class AutonomousBike(Bike):
    def __init__(self, env, bike_id):
        super().__init__(env, bike_id)
        self.init_state()

    def init_state(self):
        state_keys = [
            "agent_id",
            "location_call",
            "time_call",
            "location_meet",
            "time_meet",
            "location_drop",
            "time_drop",
        ]
        self.state = dict.fromkeys(state_keys)

    def save_state(self):
        self.history.append(self.state)
        self.init_state()

    def call(self, agent_id):
        self.set_agent(agent_id)
        self.state["agent_id"] = agent_id
        self.state["time_call"] = np.round(self.env.now, 2)
        self.state["location_call"] = self.location

        if self.print:
            print("[%.2f] Autonomous bike %d called by agent %d" % (self.env.now, self.bike_id, agent_id))

    def autonomous_move(self, location):
        if self.print:
            print("[%.2f] Autonomous bike %d moving from [%.2f, %.2f] to [%.2f, %.2f]" % (self.env.now, self.bike_id, *self.location, *location))

        distance = self.dist(self.location, location)
        yield self.env.timeout(distance)

        if self.print:
            print("[%.2f] Autonomous bike %d moved from [%.2f, %.2f] to [%.2f, %.2f]" % (self.env.now, self.bike_id, *self.location, *location))

        self.location = location
        self.state["location_meet"] = self.location
        self.state["time_meet"] = np.round(self.env.now, 2)

    def drop(self):
        if self.print:
            print("[%.2f] Autonomous bike %d dropped by agent %d on location [%.2f, %.2f]" % (self.env.now, self.bike_id, self.agent_id, *self.location))

        self.state["location_drop"] = self.location
        self.state["time_drop"] = np.round(self.env.now, 2)
        self.pop_agent()
        self.save_state()

class Agent:
    def __init__(self, env, agent_id):
        self.env = env
        self.agent_id = agent_id
        self.grid = None
        self.stations = None
        self.bikes = None

        self.print = True
        self.history = []
        self.init_state()

    def init_state(self):
        state_keys = [
            "source",
            "target",
            "time_source",
            "time_target",
            "bike_id",
            "mode",
            "source_history",
            "target_history",
            "dockless_history",
            "autonomous_history",
        ]
        self.state = dict.fromkeys(state_keys)

    def set_data(self, grid, stations, bikes):
        self.grid = grid
        self.stations = stations
        self.bikes = bikes

    def start(self):
        self.location = None
        self.bike_id = None
        self.event_setup_task = self.env.event()

    def process(self):
        # 0-Setup
        self.setup_task()
        yield self.event_setup_task

        # 1-Init on source
        yield self.env.process(self.init_agent())

    def setup_task(self):
        # self.source = self.random_location()
        # self.target = self.random_location()
        self.source = np.array([0, 0]) + np.random.normal(0, 1, size=2)
        self.target = np.array([100, 100]) + np.random.normal(0, 1, size=2)

        self.time_source = np.round(self.random_time(), 2)
        self.time_target_desired = np.round(self.time_source + self.random_time(), 2)

        self.event_setup_task.succeed()

        if self.print:
            print("[%.2f] Agent %d task created" % (self.env.now, self.agent_id))
            print("[%.2f] Agent %d source location: (%.2f, %.2f) at %.2f" % (self.env.now, self.agent_id, *self.source, self.time_source))
            print("[%.2f] Agent %d target location: (%.2f, %.2f) at %.2f" % (self.env.now, self.agent_id, *self.target, self.time_target_desired))

    def init_agent(self):
        yield self.env.timeout(self.time_source)
        self.location = self.source
        if self.print:
            print("[%.2f] Agent %d initialized at location [%.2f, %.2f]" % (self.env.now, self.agent_id, *self.location))

    def random_location(self):
        x = np.random.randint(*self.grid["xlim"])
        y = np.random.randint(*self.grid["ylim"])
        return np.array([x, y])

    def random_time(self):
        return np.random.randint(0, 10)  # 10

    def dist(self, a, b):
        return np.linalg.norm(a - b)

    def walk_to(self, location):
        distance = self.dist(self.location, location)
        yield self.env.timeout(distance)
        self.location = location

        if self.print:
            print("[%.2f] Agent %d walked to location [%.2f, %.2f]" % (self.env.now, self.agent_id, *location))

    def ride_bike_to(self, location):
        if self.print:
            print("[%.2f] Agent %d heads from location [%.2f, %.2f] to location [%.2f, %.2f] on bike %d" % (self.env.now, self.agent_id, *self.location, *location, self.bike_id))

        bike = self.bikes[self.bike_id]
        yield self.env.process(bike.ride(location))

        if self.print:
            print("[%.2f] Agent %d reached from location [%.2f, %.2f] to location [%.2f, %.2f] on bike %d" % (self.env.now, self.agent_id, *self.location, *location, self.bike_id))

        self.location = location

class StationAgent(Agent):
    def __init__(self, env, agent_id):
        super().__init__(env, agent_id)
        self.init_state()

    def init_state(self):
        state_keys = [
            "source",
            "target",
            "time_source",
            "time_target",
            "bike_id",
            "station_history",
        ]
        self.state = dict.fromkeys(state_keys)

    def save_state(self):
        self.state["source"] = self.source.tolist()
        self.state["target"] = self.target.tolist()
        self.state["time_source"] = self.time_source
        self.state["time_target"] = self.time_target
        self.state["bike_id"] = self.bike_id
        self.state["station_history"] = self.station_history
        self.history.append(self.state)
        self.init_state()

    def save_station_history(self, aim):
        history = {
            "aim": aim,
            "station_id": self.station_id,
            "time_select_station": self.time_select_station,
            "time_in_station": self.time_in_station,
            "time_interact_bike": self.time_interact_bike,
        }
        self.station_history.append(history)

    def start(self):
        super().start()

        # STATION-BASED
        self.station_id = None
        self.event_select_station = self.env.event()
        self.event_interact_bike = self.env.event()
        self.visited_stations = []
        self.station_history = []

        self.env.process(self.process())

    def process(self):
        # 0-Setup
        # 1-Init on source
        yield self.env.process(super().process())

        self.event_interact_bike = self.env.event()
        while not self.event_interact_bike.triggered:
            # 2-Select source station
            self.select_station(aim="source")

            # 3-Walk to source station
            yield self.event_select_station
            station = self.stations[self.station_id]
            yield self.env.process(self.walk_to(station.location))
            self.time_in_station = np.round(self.env.now, 2)

            # 4-Pull bike
            yield self.env.process(self.interact_bike(action="pull"))
            self.save_station_history(aim="source")

        self.event_interact_bike = self.env.event()
        while not self.event_interact_bike.triggered:
            # 5-Select target station
            self.select_station(aim="target")

            # 6-Ride bike
            yield self.event_select_station
            station = self.stations[self.station_id]
            yield self.env.process(self.ride_bike_to(station.location))
            self.time_in_station = np.round(self.env.now, 2)

            # 7-Push bike
            yield self.env.process(self.interact_bike(action="push"))
            self.save_station_history(aim="target")

        # 8-Walk to target
        yield self.env.process(self.walk_to(self.target))
        self.time_target = np.round(self.env.now, 2)

        # 9-Save state
        self.save_state()

        # 10-Finish
        yield self.env.timeout(10)
        if self.print:
            print("[%.2f] Agent %d working" % (self.env.now, self.agent_id))

    # STATION BASED

    def update_station_info(self, location):
        values = []
        for station in self.stations:
            station_id = station.station_id
            has_bikes = station.has_bikes()
            has_docks = station.has_docks()
            visited = station_id in self.visited_stations
            distance = self.dist(location, station.location)
            walkable = distance < WALK_RADIUS
            values.append((station_id, has_bikes, has_docks, visited, distance, walkable))
        labels = [
            "station_id",
            "has_bikes",
            "has_docks",
            "visited",
            "distance",
            "walkable",
        ]
        types = [int, int, int, int, float, int]
        dtype = list(zip(labels, types))
        self.station_info = np.array(values, dtype=dtype)

    def select_station(self, aim):
        self.event_select_station = self.env.event()
        location = self.location if aim == "source" else self.target  ###?????? aim== source

        self.update_station_info(location)
        for e in np.sort(self.station_info, order="distance"):
            if aim == "source":
                valid = e["has_bikes"] and not e["visited"] and e["walkable"]
            else:
                valid = e["has_docks"] and not e["visited"] and e["walkable"]
            if valid:
                self.station_id = e["station_id"]
                self.visited_stations.append(self.station_id)
                self.event_select_station.succeed()
                self.time_select_station = np.round(self.env.now, 2)
                if self.print:
                    print("[%.2f] Agent %d selected %s station %d" % (self.env.now, self.agent_id, aim, self.station_id))
                break
        if not self.event_select_station.triggered:
            print("estoy in da shit")
            # if we already pulled a bike, we are f***ed

            # call taxi / uber
            # walk to random station
            # walk towards target
            # wait and try again
            # change mode and start again
            # increase walkable distance
            # create new bike
            # stop

    def interact_bike(self, action):
        station = self.stations[self.station_id]
        valid = station.has_bikes() if action == "pull" else station.has_docks()
        if valid:
            if action == "pull":
                self.bike_id = station.choose_bike()
                self.bikes[self.bike_id].register_pull(self.agent_id)
                station.pull_bike(self.bike_id)  # HERE prints state of station
            else:
                self.bikes[self.bike_id].register_push(self.station_id)
                station.push_bike(self.bike_id)

            self.event_interact_bike.succeed()
            self.time_interact_bike = np.round(self.env.now, 2)
            if self.print:  # AND HERE the action, the order is counfusing
                print("[%.2f] Agent %d %sed bike %d on station %d" % (self.env.now, self.agent_id, action, self.bike_id, self.station_id,))
            yield self.env.timeout(1)
        else:
            self.time_interact_bike = None
            if self.print:
                print("[%.2f] Station %d has zero %s available" % (self.env.now, self.station_id, "bikes" if action == "pull" else "docks",))
            yield self.env.timeout(3)

class DocklessAgent(Agent):
    def __init__(self, env, agent_id):
        super().__init__(env, agent_id)
        self.init_state()

    def init_state(self):
        state_keys = [
            "source",
            "target",
            "time_source",
            "time_target",
            "bike_id",
            "dockless_history",
        ]
        self.state = dict.fromkeys(state_keys)

    def save_state(self):
        self.state["source"] = self.source.tolist()
        self.state["target"] = self.target.tolist()
        self.state["time_source"] = self.time_source
        self.state["time_target"] = self.time_target
        self.state["bike_id"] = self.bike_id
        self.state["dockless_history"] = self.dockless_history
        self.history.append(self.state)
        self.init_state()

    def save_dockless_history(self):
        history = {
            "dockless_bike_id": self.dockless_bike_id,
            "time_select_bike": self.time_select_bike,
            "time_in_dockless_bike": self.time_in_dockless_bike,
            "time_unlock_bike": self.time_unlock_bike,
        }
        self.dockless_history.append(history)

    def start(self):
        super().start()

        # DOCKLESS
        self.dockless_bike_id = None
        self.dockless_history = []

        self.event_select_dockless_bike = self.env.event()
        self.event_unlock_bike = self.env.event()

        self.env.process(self.process())

    def process(self):
        # 0-Setup
        # 1-Init on source
        yield self.env.process(super().process())

        while not self.event_unlock_bike.triggered:
            # 2-Select dockless bike
            self.select_dockless_bike()

            # 3-Walk to dockless bike
            yield self.event_select_dockless_bike
            dockless_bike = self.bikes[self.dockless_bike_id]
            yield self.env.process(self.walk_to(dockless_bike.location))
            self.time_in_dockless_bike = np.round(self.env.now, 2)

            # 4-Unlock bike
            yield self.env.process(self.unlock_bike())
            self.save_dockless_history()

        # 5-Ride bike
        yield self.env.process(self.ride_bike_to(self.target))

        # 6-Drop bike
        self.lock_bike()
        self.time_target = np.round(self.env.now, 2)

        # 7-Save state
        self.save_state()

        # 8-Finish
        yield self.env.timeout(10)
        if self.print:
            print("[%.2f] Agent %d working" % (self.env.now, self.agent_id))

    def update_bike_info(self):
        values = []
        for bike in self.bikes:
            if isinstance(bike, DocklessBike):
                bike_id = bike.bike_id
                rented = bike.rented()
                distance = self.dist(self.location, bike.location)
                walkable = distance < WALK_RADIUS
                values.append((bike_id, rented, distance, walkable))
        labels = ["bike_id", "rented", "distance", "walkable"]
        types = [int, int, float, int]
        dtype = list(zip(labels, types))
        self.bike_info = np.array(values, dtype=dtype)

    def select_dockless_bike(self):
        self.event_select_dockless_bike = self.env.event()
        self.update_bike_info()
        for e in np.sort(self.bike_info, order="distance"):
            # print(e)
            valid = not e["rented"] and e["walkable"]
            if valid:
                self.dockless_bike_id = e["bike_id"]
                self.event_select_dockless_bike.succeed()
                self.time_select_bike = np.round(self.env.now, 2)
                if self.print:
                    print("[%.2f] Agent %d selected dockless bike %d" % (self.env.now, self.agent_id, self.dockless_bike_id))
                break
        if not self.event_select_dockless_bike.triggered:
            print("estoy in da shit")
            # call taxi / uber
            # walk towards target
            # wait and try again
            # change mode and start again
            # increase walkable distance
            # create new bike
            # stop

    def unlock_bike(self):
        dockless_bike = self.bikes[self.dockless_bike_id]
        if not dockless_bike.rented():
            yield self.env.timeout(1)
            self.bike_id = dockless_bike.bike_id
            dockless_bike.unlock(self.agent_id)
            # HERE IT DOESNT SAVE THAT IT HAS BEEN RENTED
            self.event_unlock_bike.succeed()
            self.time_unlock_bike = np.round(self.env.now, 2)
            if self.print:
                print("[%.2f] Agent %d unlocked bike %d" % (self.env.now, self.agent_id, self.bike_id))
        else:
            yield self.env.timeout(3)
            self.time_unlock_bike = None
            if self.print:
                print("[%.2f] Bike %d has already been rented" % (self.env.now, self.dockless_bike_id))

    def lock_bike(self):
        bike = self.bikes[self.bike_id]
        bike.lock()

        if self.print:
            print("[%.2f] Agent %d locked bike %d" % (self.env.now, self.agent_id, self.bike_id))

class AutonomousAgent(Agent):
    def __init__(self, env, agent_id):
        super().__init__(env, agent_id)
        self.init_state()

    def init_state(self):
        state_keys = [
            "source",
            "target",
            "time_source",
            "time_target",
            "bike_id",
            "autonomous_history",
        ]
        self.state = dict.fromkeys(state_keys)

    def save_state(self):
        self.state["source"] = self.source.tolist()
        self.state["target"] = self.target.tolist()
        self.state["time_source"] = self.time_source
        self.state["time_target"] = self.time_target
        self.state["bike_id"] = self.bike_id
        self.state["autonomous_history"] = {
            "time_call_bike": self.time_call_bike,
            "time_meet_bike": self.time_meet_bike,
        }
        self.history.append(self.state)
        self.init_state()

    def start(self):
        super().start()

        # AUTONOMOUS
        self.event_call_autonomous_bike = self.env.event()
        self.autonomous_history = []

        self.env.process(self.process())

    def process(self):
        # 0-Setup
        # 1-Init on source
        yield self.env.process(super().process())

        # 2-Call autonomous bike
        self.call_autonomous_bike()
        yield self.event_call_autonomous_bike

        # 3-Wait for autonomous bike
        autonomous_bike = self.bikes[self.bike_id]
        yield self.env.process(autonomous_bike.autonomous_move(self.location))
        self.time_meet_bike = np.round(self.env.now, 2)

        # 4-Ride bike
        yield self.env.process(self.ride_bike_to(self.target))

        # 5-Drop bike
        self.drop_bike()
        self.time_target = np.round(self.env.now, 2)

        # 6-Save state
        self.save_state()

        # 7-Finish
        yield self.env.timeout(10)
        if self.print:
            print("[%.2f] Agent %d working" % (self.env.now, self.agent_id))

    def update_bike_info(self):
        values = []
        for bike in self.bikes:
            if isinstance(bike, AutonomousBike):
                bike_id = bike.bike_id
                rented = bike.rented()
                distance = self.dist(self.location, bike.location)
                walkable = distance < WALK_RADIUS
                values.append((bike_id, rented, distance, walkable))
        labels = ["bike_id", "rented", "distance", "walkable"]
        types = [int, int, float, int]
        dtype = list(zip(labels, types))
        self.bike_info = np.array(values, dtype=dtype)

    def call_autonomous_bike(self):
        self.event_call_autonomous_bike = self.env.event()
        self.update_bike_info()
        for e in np.sort(self.bike_info, order="distance"):
            # print(e)
            valid = not e["rented"] and e["walkable"]
            if valid:
                self.bike_id = e["bike_id"]
                self.bikes[self.bike_id].call(self.agent_id)
                self.event_call_autonomous_bike.succeed()
                self.time_call_bike = np.round(self.env.now, 2)
                if self.print:
                    print("[%.2f] Agent %d called autonomous bike %d" % (self.env.now, self.agent_id, self.bike_id))
                break
        if not self.event_call_autonomous_bike.triggered:
            print("estoy in da shit")
            # call taxi / uber
            # walk to random station
            # walk towards target
            # wait and try again
            # change mode and start again
            # increase walkable distance
            # create new bike
            # stop

    def drop_bike(self):
        bike = self.bikes[self.bike_id]
        bike.drop()

        if self.print:
            print("[%.2f] Agent %d dropped bike %d" % (self.env.now, self.agent_id, self.bike_id))

env = simpy.Environment()
city = City(env, data)
env.run(until=1000)

# %%

stations = []
for station_id, station in enumerate(city.stations):
    df = pd.DataFrame(station.history).assign(station_id=station_id)
    stations.append(df)
df_stations = pd.concat(stations).set_index("station_id").reset_index()

bikes = []
for bike_id, bike in enumerate(city.bikes):
    bike_type = type(bike).__name__
    df = pd.DataFrame(bike.history).assign(bike_id=bike_id, bike_type=bike_type)
    bikes.append(df)
df_bikes = pd.concat(bikes, sort=False).set_index("bike_id").reset_index()

agents = []
for agent_id, agent in enumerate(city.agents):
    agent_type = type(agent).__name__
    df = pd.DataFrame(agent.history).assign(agent_id=agent_id, agent_type=agent_type)
    agents.append(df)

df_agents = pd.concat(agents).set_index("agent_id").reset_index()


from setuptools import setup, find_packages
import sys, os

here = os.path.abspath(os.path.dirname(__file__))
try:
    README = open(os.path.join(here, 'README.rst')).read()
except IOError:
    README = ''

version = "0.3.0"

test_requirements = [
  'nose',
  'webtest<=2.0.24',
]

setup(name='tgext.mailer',
      version=version,
      description="TurboGears extension for sending emails with transaction manager integration",
      long_description=README,
      classifiers=[
        "Environment :: Web Environment",
        "Framework :: TurboGears"
      ],
      keywords='turbogears2.extension',
      author='Alessandro Molina',
      author_email='amol@turbogears.org',
      url='https://github.com/amol-/tgext.mailer',
      license='MIT',
      packages=find_packages(exclude=['ez_setup', 'examples', 'tgext.mailer.tests']),
      namespace_packages = ['tgext'],
      include_package_data=True,
      zip_safe=False,
      install_requires=[
          "TurboGears2 >= 2.3.2",
          "repoze.sendmail == 4.3",
      ],
      extras_require={
           # Used by Travis and Coverage due to setup.py nosetests
           # causing a coredump when used with coverage
           'testing': test_requirements,
      },
      test_suite='nose.collector',
      tests_require=test_requirements,
      entry_points="""
      # -*- Entry points: -*-
      """,
      )


import torch
from torch import nn
from torch.nn import functional as F
import lightning as L

from typing import Dict, Tuple
import warnings
from torch.distributions import Normal
from torch.distributions import kl_divergence as kl

def apply_scaled_init(model):
    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.Linear):
            nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

class Encoder(nn.Module):
    def __init__(self,
                 latent_dim: int,
                 num_input_channels: int,
                 base_channel_size: int,
                 variational: bool=False,
                 label_latent_dim: int=0,
                 BatchNorm = None,
                 act_fn: object = nn.GELU,
                 model=None,
                 width=64,
                 height=64,
                 scale_factor=0.1,
                 *args,
                 **kwargs):
        """
        Args:
           num_input_channels : Number of input channels of the image.
           base_channel_size : Number of channels we use in the first convolutional layers.
           latent_dim : Dimensionality of latent representation z
           act_fn : Activation function used throughout the encoder network
        """
        super().__init__()
        self.variational = variational
        c_hid = base_channel_size
        if model == 'uhler':
            self.net = nn.Sequential(
                nn.Conv2d(num_input_channels, c_hid, 4, 2, 1, bias=False),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Conv2d(c_hid, c_hid * 2, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 2),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Conv2d(c_hid * 2, c_hid * 4, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 4),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Conv2d(c_hid * 4, c_hid * 8, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 8),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Conv2d(c_hid * 8, c_hid * 8, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 8),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Flatten(),  # Image grid to single feature vector
            )
        elif model == 'test':
            self.net = nn.Sequential(
                nn.Conv2d(num_input_channels, c_hid, kernel_size=4, stride=2, padding=1, bias=False),  # 96x96 => 48x48
                nn.LayerNorm([c_hid, 48, 48]),
                nn.GELU(),
                nn.Conv2d(c_hid, c_hid * 2, kernel_size=4, stride=2, padding=1, bias=False),  # 48x48 => 24x24
                nn.LayerNorm([c_hid * 2, 24, 24]),
                nn.GELU(),
                nn.Conv2d(c_hid * 2, c_hid * 4, kernel_size=3, stride=2, padding=1, bias=False),  # 24x24 => 12x12
                nn.LayerNorm([c_hid * 4, 12, 12]),
                nn.GELU(),
                nn.Conv2d(c_hid * 4, c_hid * 8, kernel_size=3, stride=2, padding=1, bias=False),  # 12x12 => 6x6
                nn.LayerNorm([c_hid * 8, 6, 6]),
                nn.GELU(),
                nn.Conv2d(c_hid * 8, c_hid * 8, kernel_size=3, stride=2, padding=1, bias=False),  # 6x6 => 3x3
                nn.LayerNorm([c_hid * 8, 3, 3]),
                nn.GELU(),
                nn.Flatten()  # Image grid to single feature vector
            )
        else:
            if width == 96:
                self.net = nn.Sequential(
                    nn.Conv2d(num_input_channels, c_hid, kernel_size=3, padding=1, stride=2),  # 96x96 => 48x48
                    act_fn(),
                    nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.Conv2d(c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 48x48 => 24x24
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 24x24 => 12x12
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 12x12 => 6x6
                    act_fn(),
                    nn.Flatten(),  # Image grid to single feature vector
                )
            elif width == 64:
                self.net = nn.Sequential(
                    nn.Conv2d(num_input_channels, c_hid, kernel_size=3, padding=1, stride=2),  # 64x64 => 32x32
                    act_fn(),
                    nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.Conv2d(c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 32x32 => 16x16
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 16x16 => 8x8
                    act_fn(),
                    nn.Flatten(),  # Image grid to single feature vector
                )
        apply_scaled_init(self.net)

        if self.variational:
            if model is not None:
                input_size = c_hid * 8 * 3 * 3
                self.fc_mu = nn.Linear(input_size, latent_dim)
                self.fc_log_var = nn.Linear(input_size, latent_dim)
            else:
                if width == 96:
                    self.fc_mu = nn.Linear(2 * 6 * 6 * c_hid, latent_dim)
                    self.fc_log_var = nn.Linear(2 * 6 * 6 * c_hid, latent_dim)
                elif width == 64:
                    self.fc_mu = nn.Linear(2 * 8 * 8 * c_hid, latent_dim)
                    self.fc_log_var = nn.Linear(2 * 8 * 8 * c_hid, latent_dim)
        else:
            self.net = nn.Sequential(self.net, nn.Linear(input_size, latent_dim))
        apply_scaled_init(self.fc_mu)
        apply_scaled_init(self.fc_log_var)

    def forward(self, x, **kwargs):
        if self.variational:
            x = self.net(x)
            mu = self.fc_mu(x)
            log_var = self.fc_log_var(x)
            return mu, log_var
        else:
            x = self.net(x)
            return x

class Decoder(nn.Module):
    def __init__(self,
                 latent_dim: int,
                 num_input_channels: int,
                 base_channel_size: int,
                 batch_latent_dim: int=0,
                 BatchNorm = None,
                 act_fn: object = nn.GELU,
                 model=None,
                 width=64,
                 height=64,
                 *args,
                 **kwargs):
        """
        Args:
           num_input_channels : Number of channels of the image to reconstruct.
           base_channel_size : Number of channels we use in the last convolutional layers. Early layers might use a duplicate of it.
           latent_dim : Dimensionality of latent representation z
           act_fn : Activation function used throughout the decoder network
        """
        super().__init__()
        c_hid = base_channel_size

        if model == 'uhler':
            print('using uhler decoder')
            self.linear = nn.Sequential(
                nn.Linear(latent_dim + batch_latent_dim, 2 * 6 * 6 * c_hid),
                act_fn(),
                nn.Unflatten(1, (2 * c_hid, 6, 6)),
            )
            self.net = nn.Sequential(
                nn.ConvTranspose2d(c_hid * 8, c_hid * 8, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 8),
                nn.LeakyReLU(0.2, inplace=True),
                nn.ConvTranspose2d(c_hid * 8, c_hid * 4, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 4),
                nn.LeakyReLU(0.2, inplace=True),
                nn.ConvTranspose2d(c_hid * 4, c_hid * 2, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid * 2),
                nn.LeakyReLU(0.2, inplace=True),
                nn.ConvTranspose2d(c_hid * 2, c_hid, 4, 2, 1, bias=False),
                nn.BatchNorm2d(c_hid),
                nn.LeakyReLU(0.2, inplace=True),
                nn.ConvTranspose2d(c_hid, num_input_channels, 4, 2, 1, bias=False),
                nn.Tanh(),
        )
        elif model == 'test':
            print('using test decoder')
            self.linear = nn.Sequential(
                nn.Linear(latent_dim + batch_latent_dim, 8 * 3 * 3 * c_hid),
                nn.LayerNorm(8 * 3 * 3 * c_hid),
                act_fn(),
                nn.Unflatten(1, (8 * c_hid, 3, 3)),
            )

            self.net = nn.Sequential(
                nn.ConvTranspose2d(8 * c_hid, 4 * c_hid, kernel_size=4, padding=1, stride=2), # 3x3 => 6x6
                nn.LayerNorm([4 * c_hid, 6, 6]),
                act_fn(),
                nn.ConvTranspose2d(4 * c_hid, 2 * c_hid, kernel_size=4, padding=1, stride=2), # 6x6 => 12x12
                nn.LayerNorm([2 * c_hid, 12, 12]),
                act_fn(),
                nn.ConvTranspose2d(2 * c_hid, c_hid, kernel_size=3, output_padding=1, padding=1, stride=2), # 12x12 => 24x24
                nn.LayerNorm([c_hid, 24, 24]),
                act_fn(),
                nn.ConvTranspose2d(c_hid, c_hid // 2, kernel_size=5, output_padding=1, padding=2, stride=2), # 24x24 => 48x48, using a larger kernel
                nn.LayerNorm([c_hid // 2, 48, 48]),
                act_fn(),
                nn.ConvTranspose2d(c_hid // 2, num_input_channels, kernel_size=3, output_padding=1, padding=1, stride=2), # 48x48 => 96x96
                nn.Tanh(),
            )
        else:
            if width == 96:
                print('using width 96 decoder')
                self.linear = nn.Sequential(
                    nn.Linear(latent_dim + batch_latent_dim, 2 * 6 * 6 * c_hid),
                    act_fn(),
                    nn.Unflatten(1, (2 * c_hid, 6, 6)),
                )
                self.net = nn.Sequential(
                    nn.ConvTranspose2d(2 * c_hid, 2 * c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 8x8 => 16x16
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.ConvTranspose2d(2 * c_hid, c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 16x16 => 32x32
                    act_fn(),
                    nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.ConvTranspose2d(c_hid, c_hid // 2, kernel_size=3, output_padding=1, padding=1, stride=2),  # 32x32 => 64x64
                    act_fn(),
                    nn.ConvTranspose2d(c_hid // 2, num_input_channels, kernel_size=3, output_padding=1, padding=1, stride=2),  # 64x64 => 96x96
                    nn.Tanh(),
                )
            elif width == 64:
                print('using width 64 decoder')
                self.linear = nn.Sequential(
                    nn.Linear(latent_dim + batch_latent_dim, 2 * 8 * 8 * c_hid),
                    act_fn(),
                    nn.Unflatten(1, (-1, 8, 8)),
                    )
                self.net = nn.Sequential(
                    nn.ConvTranspose2d(2 * c_hid, 2 * c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 8x8 => 16x16
                    act_fn(),
                    nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.ConvTranspose2d(2 * c_hid, c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 16x16 => 32x32
                    act_fn(),
                    nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
                    act_fn(),
                    nn.ConvTranspose2d(c_hid, num_input_channels, kernel_size=3, output_padding=1, padding=1, stride=2),  # 32x32 => 64x64
                    nn.Tanh(),
                )
        apply_scaled_init(self.linear)
        apply_scaled_init(self.net)

    def forward(self, x, **kwargs):
        x = self.linear(x)
        x = self.net(x)
        return x

class BaseModel(L.LightningModule):
    def __init__(
        self,
        model_name: str,
        optimizer_param: dict,
        latent_dim: int=32,
        base_channel_size: int=32,
        num_input_channels: int = 4,
        image_size: int = 64,
        act_fn=nn.GELU,
        *args,
        **kwargs,
        ):
        super().__init__()

        self.model_name = model_name
        self.num_input_channels = num_input_channels
        self.width = image_size
        self.height = image_size
        self.base_channel_size = base_channel_size
        self.latent_dim = latent_dim
        self.network_param = {'latent_dim': latent_dim, 'num_input_channels':num_input_channels,
                              'base_channel_size':base_channel_size, 'act_fn':act_fn}

        # Example input array needed for visualizing the graph of the network
        self.optimizer_param = optimizer_param

    def forward(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat

    def _get_loss(self, x):
        x_hat = self.forward(x)
        loss = F.mse_loss(x, x_hat, reduction="none")
        loss = loss.sum(dim=[1, 2, 3]).mean(dim=[0])
        return loss

    def configure_optimizers(self):
        lr = self.optimizer_param['lr']
        if self.optimizer_param['optimizer'] == 'Adam':
            optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        elif self.optimizer_param['optimizer'] == 'SGD':
            momentum = self.optimizer_param['momentum']
            nesterov = self.optimizer_param['nesterov']
            optimizer = torch.optim.SGD(self.parameters(), lr=lr, momentum=momentum, nesterov=nesterov)
        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.2, patience=20, min_lr=5e-5)
        return {"optimizer": optimizer, "monitor": "train_loss"}

    def training_step(self, batch, batch_idx):
        loss = self._get_loss(batch)
        self.log("train_loss", loss.detach())
        while True:  # Inside your training loop
            current_memory_allocated = torch.cuda.memory_allocated() / (1024.0 ** 3)  # Convert bytes to GB
            max_memory_allocated = torch.cuda.max_memory_allocated() / (1024.0 ** 3)  # Convert bytes to GB
            self.log_dict({'Current GPU Memory (GB)': current_memory_allocated, 'Max GPU Memory (GB)': max_memory_allocated})
            return loss

    def validation_step(self, batch, batch_idx):
        loss = self._get_loss(batch)
        self.log("val_loss", loss.detach())

    def test_step(self, batch, batch_idx):
        loss = self._get_loss(batch)
        self.log("test_loss", loss.detach())

class AEmodel(BaseModel):
    def __init__(self,
                 encoder_class: object = Encoder,
                 decoder_class: object = Decoder,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.encoder = encoder_class(variational=False, **self.network_param)
        self.decoder = decoder_class(**self.network_param)
        self.example_input_array = torch.zeros(2, self.num_input_channels, self.width, self.height)

    def get_image_embedding(self, x):
        return self.encoder(x)

class VAEmodel(BaseModel):
    def __init__(self,
                 step_size: int,
                 latent_dim = 64,
                 encoder_class: object = Encoder,
                 decoder_class: object = Decoder,
                 *args, **kwargs):
        super().__init__(latent_dim=latent_dim, *args, **kwargs)

        # Initialize the cyclic weight scheduler
        self.kl_weight_scheduler = CyclicWeightScheduler(step_size=step_size)

        # for the gaussian likelihood
        self.log_scale = nn.Parameter(torch.Tensor([0.0]))
        self.encoder = encoder_class(variational=True, **self.network_param)
        self.decoder = decoder_class(**self.network_param)
        self.example_input_array = torch.zeros(2, self.num_input_channels, self.width, self.height)
        # Saving hyperparameters of autoencoder
        self.save_hyperparameters()

    def forward(self, x):
        # encode x to get the mu and variance parameters
        mu, log_var = self.encoder(x)
        # sample z
        log_var = torch.maximum(log_var, torch.tensor(-20)) #clipping to prevent going to -inf
        std = torch.exp(log_var / 2)
        z = self.sampling(mu, std)
        # decoded
        x_hat = self.decoder(z)
        return x_hat, mu, std

    def get_image_embedding(self, x):
        mu, _ = self.encoder(x)
        return mu

    @staticmethod
    def sampling(mu, std):
        q = Normal(mu, std)
        return q.rsample()

    @staticmethod
    def reconstruction_loss(sample: torch.Tensor,
                            mean: torch.Tensor,
                            logscale: torch.Tensor,
                            ):
        scale = torch.exp(logscale)
        dist = Normal(mean, scale)
        log_pxz = dist.log_prob(sample)
        return -log_pxz.sum(dim=(1, 2, 3))

    def _generic_loss(self, x, x_hat, mu, std):
        # reconstruction probability
        recon_loss = self.reconstruction_loss(x, x_hat, self.log_scale)

        # kl
        kl = self.latent_kl_divergence(mu, std)
        kl_term_weight = self.kl_weight_scheduler.step()

        # elbo
        elbo = (kl_term_weight*kl + recon_loss)
        elbo = elbo.mean()
        self.log_dict({
            'elbo': elbo.detach(),
            'kl': kl.mean().detach(),
            'recon_loss': recon_loss.mean().detach(),
            'kl_term_weight': kl_term_weight,
        })
        return elbo

    @staticmethod
    def latent_kl_divergence(variational_mean,
                             variational_std,
                             prior_mean=None,
                             prior_std=None) -> torch.Tensor:
        """
        Compute KL divergence between a variational posterior and standard Gaussian prior.
        Args:
        ----
            variational_mean: Mean of the variational posterior Gaussian.
            variational_var: Variance of the variational posterior Gaussian.
        Returns
        -------
            KL divergence for each data point. If number of latent samples == 1,
            the tensor has shape `(batch_size, )`. If number of latent
            samples > 1, the tensor has shape `(n_samples, batch_size)`.
        """
        if prior_mean is None:
            prior_mean = torch.zeros_like(variational_mean)
        prior_std = torch.ones_like(variational_std)
        return kl(
            Normal(variational_mean, variational_std),
            Normal(prior_mean, prior_std),
        ).sum(dim=-1)

    def _get_loss(self, x):
        # get reconstruction, mu and std
        x_hat, mu, std = self.forward(x)
        elbo = self._generic_loss(x, x_hat, mu, std)
        return elbo

class ContrastiveVAEmodel(BaseModel):
    """
    Args:
    ----
        n_z_latent: Dimensionality of the background latent space.
        n_s_latent: Dimensionality of the salient latent space.
        wasserstein_penalty: Weight of the Wasserstein distance loss that further
            discourages shared variations from leaking into the salient latent space.
    """

    def __init__(self,
                 n_z_latent: int = 32,
                 n_s_latent: int = 32,
                 encoder_class: object = Encoder,
                 decoder_class: object = Decoder,
                 step_size: float=2000,
                 ngene: int=None,
                 adjust_prior_s: bool=False,
                 adjust_prior_z: bool=False,
                 classify_s: bool=False,
                 classify_z: bool=False,
                 wasserstein_penalty: float = 0,
                 BatchNorm = None,
                 n_unique_batch: int = 34,
                 model = None,
                 batch_size: int=1024,
                 tc_penalty: float=1,
                 classification_weight: float=1,
                 scale_factor: float=0.1,
                 max_kl_weight: float=1,
                 batch_latent_dim: int=32,
                 reg_type: str=None,
                 total_steps: int=3000,
                 klscheduler: str='cyclic',
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.step_size = step_size
        self.ngene = ngene
        self.adjust_prior_s = adjust_prior_s
        self.adjust_prior_z = adjust_prior_z
        self.n_s_latent = n_s_latent
        self.n_z_latent = n_z_latent
        self.wasserstein_penalty = wasserstein_penalty
        self.BatchNorm = BatchNorm
        self.n_unique_batch = n_unique_batch
        self.model = model
        self.batch_size = batch_size
        self.tc_penalty = tc_penalty
        self.classify_s = classify_s
        self.classify_z = classify_z
        self.classification_weight = classification_weight
        self.scale_factor = scale_factor
        self.batch_latent_dim = batch_latent_dim
        self.reg_type = reg_type
        self.total_steps = total_steps
        self.klscheduler = klscheduler

        if n_s_latent != n_z_latent:
            warnings.warn('Target latent dim does not equal background latent dim')

        # Initialize the weight scheduler
        if self.klscheduler == 'cyclic':
            self.kl_weight_scheduler = CyclicWeightScheduler(step_size=self.step_size, max_weight=max_kl_weight)
        elif self.klscheduler == 'ramp':
            self.kl_weight_scheduler = KLRampScheduler(total_steps=self.total_steps, max_weight=max_kl_weight)

        # Background encoder
        self.coder_param = {'num_input_channels': self.num_input_channels, "scale_factor": self.scale_factor,
                         'base_channel_size': self.base_channel_size, 'variational': True, 'width': self.width, 'height':self.height,
                          'BatchNorm': self.BatchNorm, 'n_unique_batch': self.n_unique_batch, 'model': self.model,}
        self.z_encoder = encoder_class(
                                       latent_dim=self.n_z_latent,
                                       **self.coder_param,
                                       )
        # Salient encoder
        self.s_encoder = encoder_class(
                                       latent_dim=self.n_s_latent,
                                       **self.coder_param,
                                       )

        # Decoder from latent variable to distribution parameters in data space.
        self.n_total_latent = self.n_z_latent + self.n_s_latent
        self.decoder = decoder_class(
                                    latent_dim=self.n_total_latent,
                                    batch_latent_dim=self.batch_latent_dim,
                                    **self.coder_param,
                                    )

        if self.adjust_prior_z:
            self.zprior_embedding = nn.Embedding(self.ngene, self.n_z_latent)
        if self.adjust_prior_s:
            self.sprior_embedding = nn.Embedding(self.ngene, self.n_s_latent)

        # for the gaussian likelihood
        self.log_scale = nn.Parameter(torch.Tensor([0.0]))

        # Example input array needed for visualizing the graph of the network
        self.example_input_array = {'background': torch.zeros(2, self.num_input_channels, self.width, self.height),
                                        'target': torch.zeros(2, self.num_input_channels, self.width, self.height)}
        if self.adjust_prior_s or self.adjust_prior_z:
            self.example_input_array['background_label'] = torch.zeros(2, dtype=torch.int32)
            self.example_input_array['target_label'] = torch.zeros(2, dtype=torch.int32)
        if self.batch_latent_dim > 0:
            self.batch_embedding = nn.Embedding(self.n_unique_batch, self.batch_latent_dim)
            self.example_input_array.update({'background_batch': torch.zeros(2, dtype=torch.int32),
                                            'target_batch': torch.zeros(2, dtype=torch.int32)})
        # Saving hyperparameters of autoencoder
        self.save_hyperparameters()

    def forward(self, background, target, **kwargs):
        background_label = kwargs.get('background_label')
        target_label = kwargs.get('target_label')
        prior_mu_background = {'zprior_m': None,  'sprior_m': None}
        prior_mu_target = {'zprior_m': None, 'sprior_m': None}
        # zlabel_embedding = None
        # slabel_embedding = None
        if self.adjust_prior_s:
            prior_mu_background['sprior_m'] = self.sprior_embedding(background_label.int())
            prior_mu_target['sprior_m'] = self.sprior_embedding(target_label.int())
            # slabel_embedding = torch.cat([prior_mu_background['sprior_m'],
            #                               prior_mu_target['sprior_m']], dim=0)
        if self.adjust_prior_z:
            prior_mu_background['zprior_m'] = self.zprior_embedding(background_label.int())
            prior_mu_target['zprior_m'] = self.zprior_embedding(target_label.int())
            # zlabel_embedding = torch.cat([prior_mu_background['zprior_m'],
            #                               prior_mu_target['zprior_m']], dim=0)
        inference_outputs = self.inference(background=background,
                                           target=target)
        background_batch = kwargs.get('background_batch')
        target_batch = kwargs.get('target_batch')
        generative_outputs = self.generative(inference_outputs['background'],
                                             inference_outputs['target'],
                                             background_batch=background_batch,
                                             target_batch=target_batch)
        recon = {'bg':generative_outputs['background']["px_m"],
                 "tg":generative_outputs['target']["px_m"]}
        inference_outputs['background'].update(prior_mu_background)
        inference_outputs['target'].update(prior_mu_target)

        return recon, inference_outputs, generative_outputs

    def get_image_embedding(self, img, label=None):
        qz_m, _ = self.z_encoder(img)
        qs_m, _ = self.s_encoder(img)
        return torch.cat((qs_m, qz_m), dim=1)

    def _generic_inference(self,
                             x: torch.Tensor,
                             ):
        qz_m, qz_lv = self.z_encoder(x)
        qs_m, qs_lv = self.s_encoder(x)

        # sample from latent distribution
        qz_lv = torch.maximum(qz_lv, torch.tensor(-20)) #clipping to prevent going to -inf
        qs_lv = torch.maximum(qs_lv, torch.tensor(-20)) #clipping to prevent going to -inf
        qz_s = torch.exp(qz_lv / 2)
        qs_s = torch.exp(qs_lv / 2)
        qz = Normal(qz_m, qz_s)
        qs = Normal(qs_m, qs_s)
        z = qz.rsample()
        s = qs.rsample()

        outputs = dict(
            qz_m=qz_m,
            qz_s=qz_s,
            z=z,
            qs_m=qs_m,
            qs_s=qs_s,
            s=s,)
        return outputs

    def inference(
            self,
            background: torch.Tensor,
            target: torch.Tensor,
        ) -> Dict[str, Dict[str, torch.Tensor]]:
            background_batch_size = background.shape[0]
            target_batch_size = target.shape[0]
            inference_input = torch.cat([background, target], dim=0)
            outputs = self._generic_inference(x=inference_input)
            background_outputs, target_outputs = {}, {}
            for key in outputs.keys():
                if outputs[key] is not None:
                    background_tensor, target_tensor = torch.split(
                        outputs[key],
                        [background_batch_size, target_batch_size],
                        dim=0,
                    )
                else:
                    background_tensor, target_tensor = None, None
                background_outputs[key] = background_tensor
                target_outputs[key] = target_tensor
            background_outputs["s"] = torch.zeros_like(background_outputs["s"])
            return dict(background=background_outputs, target=target_outputs)

    def _generic_generative(self,
                              z: torch.Tensor,
                              s: torch.Tensor,
                              batch_embedding: torch.Tensor=None,):
        latent = torch.cat([z, s], dim=-1)
        if batch_embedding is not None:
            latent = torch.cat([latent, batch_embedding], dim=-1)
        px_m = self.decoder(latent)
        return dict(px_m=px_m, px_s=self.log_scale)

    def generative(
        self,
        background: Dict[str, torch.Tensor],
        target: Dict[str, torch.Tensor],
        **kwargs,
        ) -> Dict[str, Dict[str, torch.Tensor]]:
        latent_z_shape = background["z"].shape
        batch_size_dim = 0 if len(latent_z_shape) == 2 else 1
        background_batch_size = background["z"].shape[batch_size_dim]
        target_batch_size = target["z"].shape[batch_size_dim]
        generative_input = {}
        for key in ["z", "s"]:
            generative_input[key] = torch.cat(
                [background[key], target[key]], dim=batch_size_dim
            )
        background_batch = kwargs.get("background_batch")
        target_batch = kwargs.get("target_batch")
        if background_batch is not None and target_batch is not None:
            generative_input["batch_embedding"] = torch.cat(
                [self.batch_embedding(background_batch),
                 self.batch_embedding(target_batch)], dim=batch_size_dim
                )
        outputs = self._generic_generative(**generative_input)
        background_outputs, target_outputs = {}, {}
        if outputs["px_m"] is not None:
            background_tensor, target_tensor = torch.split(
                outputs["px_m"],
                [background_batch_size, target_batch_size],
                dim=batch_size_dim,
            )
        else:
            background_tensor, target_tensor = None, None
        background_outputs["px_m"] = background_tensor
        target_outputs["px_m"] = target_tensor
        background_outputs["px_s"] = outputs["px_s"]
        target_outputs["px_s"] = outputs["px_s"]
        return dict(background=background_outputs, target=target_outputs)

    def _generic_loss(self,
                      tensors: torch.Tensor,
                      inference_outputs: Dict[str, torch.Tensor],
                      generative_outputs: Dict[str, torch.Tensor],
                    )-> Dict[str, torch.Tensor]:

        qz_m = inference_outputs["qz_m"]
        qz_s = inference_outputs["qz_s"]
        qs_m = inference_outputs["qs_m"]
        qs_s = inference_outputs["qs_s"]
        zprior_m = inference_outputs["zprior_m"]
        sprior_m = inference_outputs["sprior_m"]
        px_m = generative_outputs["px_m"]
        px_s = generative_outputs["px_s"]

        recon_loss = VAEmodel.reconstruction_loss(tensors, px_m, px_s)
        kl_z = VAEmodel.latent_kl_divergence(qz_m, qz_s, prior_mean=zprior_m)
        kl_s = VAEmodel.latent_kl_divergence(qs_m, qs_s, prior_mean=sprior_m)
        return dict(recon_loss=recon_loss, kl_z=kl_z, kl_s=kl_s)

    def compute_independent_loss(self, zb, zc):
        reg_type = self.reg_type
        if reg_type == "TC":
            return self.compute_tc(zb, zc)
        elif reg_type == "HSIC":
            return self.compute_HSIC(zb, zc)
        else:
            raise ValueError("reg_type should be TC or HSIC")

    @staticmethod
    def rbf_kernel(X, sigma=1.0):
        # Compute the pairwise squared Euclidean distances
        pairwise_dists = torch.cdist(X, X, p=2) ** 2
        # Apply the RBF kernel function
        values = torch.div(-pairwise_dists, (2 * sigma**2))
        return values.exp()

    @staticmethod
    def compute_HSIC(Z_b, Z_c):
        n = Z_b.shape[0]
        # Compute kernel matrices
        K = ContrastiveVAEmodel.rbf_kernel(Z_b)
        L = ContrastiveVAEmodel.rbf_kernel(Z_c)
        # print(K.shape, L.shape)
        # Implement the HSIC formula
        term1 = (1 / (n**2)) * torch.sum(K * L)
        term2 = (1 / (n**4)) * torch.sum(K) * torch.sum(L)
        term3 = (2 / (n**3)) * torch.sum(K @ L)
        HSIC_n = term1 + term2 - term3
        return HSIC_n * n

    @staticmethod
    def compute_tc(zb, zc):
        # Calculate the empirical means
        mean_zb = torch.mean(zb, dim=0)
        mean_zc = torch.mean(zc, dim=0)
        # Calculate the centered variables
        centered_zb = zb - mean_zb
        centered_zc = zc - mean_zc
        # Calculate the covariance matrix of the concatenated latent variables
        z_concat = torch.cat([centered_zb, centered_zc], dim=1)
        cov_matrix = torch.matmul(z_concat.T, z_concat) / z_concat.shape[0]
        # Calculate the covariance matrices for zb and zc individually
        cov_zb = torch.matmul(centered_zb.T, centered_zb) / centered_zb.shape[0]
        cov_zc = torch.matmul(centered_zc.T, centered_zc) / centered_zc.shape[0]
        # Calculate total correlation loss
        tc_loss = torch.logdet(cov_matrix) - (torch.logdet(cov_zb) + torch.logdet(cov_zc))
        # Multiply by the weighting factor
        return -tc_loss

    def _get_loss(self,
             concat_tensors: Dict[str, Tuple[Dict[str, torch.Tensor], int]],
             ):
        _, inference_outputs, generative_outputs = self.forward(**concat_tensors)

        background_losses = self._generic_loss(
            concat_tensors["background"],
            inference_outputs["background"],
            generative_outputs["background"],
        )
        target_losses = self._generic_loss(
            concat_tensors["target"],
            inference_outputs["target"],
            generative_outputs["target"],
        )
        recon_loss = background_losses["recon_loss"] + target_losses["recon_loss"]
        kl_divergence_z = background_losses["kl_z"] + target_losses["kl_z"]
        kl_divergence_s = target_losses["kl_s"]

        wasserstein_loss = (
            torch.norm(inference_outputs["background"]["qs_m"], dim=-1)**2
            + torch.sum(inference_outputs["background"]["qs_s"]**2, dim=-1)
        )

        if self.reg_type is not None:
            zb = torch.concat([inference_outputs["target"]["qz_m"], inference_outputs["background"]["qz_m"]], axis=0)
            zs = torch.concat([inference_outputs["target"]["qs_m"], inference_outputs["background"]["qs_m"]], axis=0)
            tc_loss = self.compute_independent_loss(zb, zs)
        else:
            tc_loss = torch.zeros(1, device=self.device)

        kl_term_weight = self.kl_weight_scheduler.step()

        elbo = torch.mean(recon_loss +
                          kl_term_weight * (kl_divergence_s + kl_divergence_z +
                                            self.wasserstein_penalty * wasserstein_loss +
                                            self.tc_penalty * tc_loss))

        self.log_dict({
            'kl_divergence_z': kl_divergence_z.mean().detach(),
            'kl_divergence_s': kl_divergence_s.mean().detach(),
            'total_recon_loss': recon_loss.mean().detach(),
            'wasserstein_loss': wasserstein_loss.mean().detach(),
            'tc_loss': tc_loss.mean().detach(),
            # 'background_recon_loss': background_losses["recon_loss"].mean().detach(),
            # 'target_recon_loss': target_losses["recon_loss"].mean().detach(),
            'kl_term_weight': kl_term_weight,
        })
        return elbo

class CyclicWeightScheduler:
    def __init__(self, step_size, base_weight=0, max_weight=1):
        self.base_weight = base_weight
        self.max_weight = max_weight
        self.step_size = step_size
        self.cycle = 0
        self.step_count = 0

    def step(self):
        # Compute the current position in the cycle
        cycle_position = self.step_count / self.step_size

        if cycle_position <= 1:
            weight = self.base_weight + (self.max_weight - self.base_weight) * cycle_position
        else:
            weight = self.max_weight
            # weight = self.max_weight - (self.max_weight - self.base_weight) * (cycle_position - 1)

        self.step_count = (self.step_count + 1) % (self.step_size * 2)

        return weight

class KLRampScheduler:
    def __init__(self, start_weight=0, max_weight=1, total_steps=3000):
        self.start_weight = start_weight
        self.max_weight = max_weight
        self.total_steps = total_steps
        self.current_step = 0
        self.current_weight = start_weight

    def step(self):
        self.current_step += 1
        progress = self.current_step / self.total_steps
        self.current_weight = self.start_weight + (self.max_weight - self.start_weight) * progress
        # Clip the weight to be within the specified range
        self.current_weight = min(max(self.current_weight, self.start_weight), self.max_weight)
        return self.current_weight

    def get_weight(self):
        return self.current_weight

class LinearDiscriminator(nn.Module):
    def __init__(self, input_features):
        super(LinearDiscriminator, self).__init__()
        self.linear = nn.Linear(input_features, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.linear(x)
        x = self.sigmoid(x)
        return x

class ConditionalBatchNorm1d(nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.num_features = num_features
        self.bn = nn.BatchNorm1d(num_features, affine=False)
        self.embed = nn.Embedding(num_classes, num_features * 2)
        self.embed.weight.data[:, :num_features].normal_(1, 0.02)  # Initialise scale at N(1, 0.02)
        self.embed.weight.data[:, num_features:].zero_()  # Initialise bias at 0

    def forward(self, x, y):
        # y is the condition, i.e. batch number
        out = self.bn(x)
        gamma, beta = self.embed(y).chunk(2, 1)
        gamma = gamma.expand_as(out)
        beta = beta.expand_as(out)
        return gamma * out + beta

class ConditionalBatchNorm2d(nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.num_features = num_features
        self.bn = nn.BatchNorm2d(num_features, affine=False)
        self.embed = nn.Embedding(num_classes, num_features * 2)
        self.embed.weight.data[:, :num_features].normal_(1, 0.02)  # Initialise scale at N(1, 0.02)
        self.embed.weight.data[:, num_features:].zero_()  # Initialise bias at 0

    def forward(self, x, y):
        # y is the condition, i.e. batch number
        out = self.bn(x)
        gamma, beta = self.embed(y).chunk(2, 1)
        gamma = gamma.unsqueeze(2).unsqueeze(3).expand_as(out)
        beta = beta.unsqueeze(2).unsqueeze(3).expand_as(out)
        return gamma * out + beta

def add_encoder_batch_norm(model,
                           BatchNorm,
                           n_conditions=None
                           ):
    new_layers = []
    for layer in model:
        new_layers.append(layer)
        if isinstance(layer, nn.Conv2d):
            if BatchNorm == ConditionalBatchNorm2d:
                new_layers.append(BatchNorm(layer.out_channels, n_conditions))
            else:
                new_layers.append(BatchNorm(layer.out_channels))
    return nn.Sequential(*new_layers)

def add_decoder_batch_norm(linear,
                           net,
                           BatchNorm1d=nn.BatchNorm1d,
                           BatchNorm2d=nn.BatchNorm2d,
                           n_conditions=None):
    new_linear = []
    for layer in linear:
        new_linear.append(layer)
        if isinstance(layer, nn.Linear):
            if BatchNorm1d == ConditionalBatchNorm1d:
                new_linear.append(BatchNorm1d(layer.out_features, n_conditions))
            else:
                new_linear.append(BatchNorm1d(layer.out_features))

    new_net = []
    for i, layer in enumerate(net):
        new_net.append(layer)
        if isinstance(layer, (nn.Conv2d, nn.ConvTranspose2d)) and i != len(net):
            if BatchNorm2d == ConditionalBatchNorm2d:
                new_net.append(BatchNorm2d(layer.out_channels, n_conditions))
            else:
                new_net.append(BatchNorm2d(layer.out_channels))

    return nn.Sequential(*new_linear), nn.Sequential(*new_net)


# Simple Python Pong
import turtle
import time

# Screen creation and definition
wind = turtle.Screen()
wind.title("Pong Game by @Angelion879")
wind.bgcolor("black")
wind.setup(width=800, height=600)
wind.tracer(0)

# Objects in the screen:

class Paddle(turtle.Turtle):
    def __init__(self, position):
        super().__init__(shape='square')
        self.shapesize(stretch_wid=5, stretch_len=1)
        self.speed(0)
        self.color("white")
        self.penup()
        self.goto(position)

class Ball(turtle.Turtle):
    def __init__(self, position):
        super().__init__(shape='square')
        self.speed(0)
        self.color("white")
        self.penup()
        self.goto(position)
        self.dx = 0.3
        self.dy = 0.3

class Text(turtle.Turtle):
    def __init__(self, position, message, l_size):
        super().__init__()
        self.speed(0)
        self.color("white")
        self.penup()
        self.hideturtle()
        self.goto(position)
        self.write(message, align="center", font=("Courier", l_size, "normal"))

class SelectionArrow(turtle.Turtle):
    def __init__(self, position, direction):
        super().__init__(shape='arrow')
        self.speed(0)
        self.color('white')
        self.penup()
        self.goto(position)
        self.left(direction)

class Option(turtle.Turtle):
    def __init__(self, position, message):
        super().__init__()
        self.speed(0)
        self.color("white")
        self.penup()
        self.hideturtle()
        self.goto(position)
        self.write(message, align="center", font=("Courier", 20, "normal"))

menu = Text((0, 20), "PRESS SPACE TO START", 35)
exit_message = Text((0, -290), "PRESS ESC TO EXIT", 10)
left_paddle = Paddle((-360, 0))
right_paddle = Paddle((360, 0))
ball = Ball((0, 0))
start = False

# AI Starter
single = Option((-150, -50), "SINGLEPLAYER")
multi = Option((150, -50), "MULTIPLAYER")
right_arrow = SelectionArrow((-35, -35), 180)
left_arrow = SelectionArrow((-265, -35), 0)
pc_play = False

# Point Writing
score_left = 0
score_right = 0
points = Text(
    (0, 260), f"A: {score_left}\t\t\t\t\tB: {score_right}", 20)

# Functions

def start_game():
    global start
    start = True

def exit_game():
    turtle.bye()

def left_paddle_up():
    y = left_paddle.ycor()
    y += 20
    left_paddle.sety(y)

def right_paddle_up():
    y = right_paddle.ycor()
    y += 20
    right_paddle.sety(y)

def left_paddle_down():
    y = left_paddle.ycor()
    y -= 20
    left_paddle.sety(y)

def right_paddle_down():
    y = right_paddle.ycor()
    y -= 20
    right_paddle.sety(y)

def ai_paddle_down():
    y = left_paddle.ycor()
    y -= 10
    left_paddle.sety(y)

def ai_paddle_up():
    y = left_paddle.ycor()
    y += 10
    left_paddle.sety(y)

def selection_move_right():
    rx = 255
    lx = 40
    right_arrow.setx(rx)
    left_arrow.setx(lx)

def selection_move_left():
    rx = -35
    lx = -265
    right_arrow.setx(rx)
    left_arrow.setx(lx)

def game_mode():
    if right_arrow.xcor() < 0:
        global pc_play
        pc_play = True
    else:
        pc_play = False

# Key binding
wind.listen()
wind.onkeypress(left_paddle_up, 'w')
wind.onkeypress(right_paddle_up, 'Up')
wind.onkeypress(left_paddle_down, 's')
wind.onkeypress(right_paddle_down, 'Down')
wind.onkeypress(start_game, 'space')
wind.onkeypress(exit_game, 'Escape')
wind.onkeypress(selection_move_right, 'Right')
wind.onkeypress(selection_move_left, 'Left')

# Main game loop
while True:
    wind.update()
    game_mode()

    # Ball moving
    if start:
        menu.clear()
        single.clear()
        multi.clear()
        right_arrow.hideturtle()
        left_arrow.hideturtle()
        wind.onkeypress(None, 'Right')
        wind.onkeypress(None, 'Left')
        ball.setx(ball.xcor() + ball.dx)
        ball.sety(ball.ycor() + ball.dy)

    # Border colision
    if ball.ycor() > 290:
        ball.sety(290)
        ball.dy *= -1
    elif ball.ycor() < -285:
        ball.sety(-285)
        ball.dy *= -1
    elif ball.xcor() > 390:
        ball.goto(0, 0)
        ball.dx *= -1
        score_left += 1
        points.clear()
        points.write(f"A: {score_left}\t\t\t\t\tB: {score_right}",
                     align="center", font=("Courier", 20, "normal"))
    elif ball.xcor() < -390:
        ball.goto(0, 0)
        ball.dx *= -1
        score_right += 1
        points.clear()
        points.write(f"A: {score_left}\t\t\t\t\tB: {score_right}",
                     align="center", font=("Courier", 20, "normal"))

    # Paddle colision
    if (ball.xcor() > 340 and ball.xcor() < 350) and (ball.ycor() < (right_paddle.ycor() + 50) and ball.ycor() > (right_paddle.ycor() - 50)):
        ball.setx(340)
        ball.dx *= -1
    elif (ball.xcor() < -340 and ball.xcor() > -350) and (ball.ycor() < (left_paddle.ycor() + 50) and ball.ycor() > (left_paddle.ycor() - 50)):
        ball.setx(-340)
        ball.dx *= -1

    # AI Player
    if pc_play == True:
        if ball.xcor() < 0 and left_paddle.ycor() < ball.ycor() and abs(left_paddle.ycor() - ball.ycor()) > 20:
            ai_paddle_up()
        elif ball.xcor() < 0 and left_paddle.ycor() > ball.ycor() and abs(left_paddle.ycor() - ball.ycor()) > 20:
            ai_paddle_down()

    # Winning condition
    if (score_left >= 5) or (score_right >= 5):
        start = False
        if score_left > score_right:
            menu.write(f"PLAYER A WINS", align="center",
                       font=("Courier", 35, "normal"))
        else:
            menu.write(f"PLAYER B WINS", align="center",
                       font=("Courier", 35, "normal"))