"""
This file holds the interaction sites class used in simulation.py.
"""
import warnings
from random import random
from copy import deepcopy
from itertools import combinations
from math import comb

import numpy as np

class InteractionSites:
    """A class designed to host interactions between persons within specific locations.

    There are currently 7 different locations that can host interactions between
    person objects.

    All attributes are passed through the sim_obj, which accesses the simulation
    configuration file. Outlined below are the main object attributes that provide the
    interaction functionality.

    Attributes
    ----------
    grade_A_sites : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to represent resturants, gas stations, retail stores, etc. Any location where you
        do not visit often, but attend a wide variety of them.
    grade_B_sites : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to represent a gym, grocery store, etc. Any location where
        you visit semi-often, and are likly to visit the same one, but this may varry.
    grade_C_sites : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to represent offices or schools. Any location where they are
        visited almost every workday, and you almost always visit the same one.
    house_sites : :obj:`np.array` of :obj:`list` of :obj:`int`
        Visited by every person each day, and hosts interactions between members
        of the same household. Infection spread at home is not defined by explicit contacts,
        but by a known spread factor.
    lect_sites : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to replicate university lecture hall interactions. They are only visited by students.
    study_sites : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to replicate study environments at university, on-campus (library, bookable rooms, ...).
        They are only visited by students.
    food_sites : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to replicate cafeteria and restaurant interactions on-campus. Only visited by students.
    res_sites : :obj:`list` of :obj:`np.array` of :obj:`int`
        Designed to replicate the student residences on campus. They are only visited by first year students.
    stud_house_sites : :obj:`np.array` of :obj:`list` of :obj:`int`
        Visited by every student each day, and hosts interactions between members
        of the same household. Infection spread at home is not defined by explicit contacts,
        but by a known spread factor.
    """

    def __init__(self, sim_obj):
        """ __init__ method docstring.

        Parameters
        ----------
        sim_obj : :obj:`cv19.simulation.simulation`
            The encompassing simulation obejct hosting the simulation.
        """

        # Set attributes from config file
        self.load_attributes_from_sim_obj(sim_obj)

        self.daily_interactions = {"HOUSE_GENERAL": np.zeros(self.nDays),
                                   "HOUSE_STUDENT": np.zeros(self.nDays)}

        # Generates a list of people that go to different grade X sites
        # len(grade_X_sites) is how many sites there are; len(grade_X_sites[i]) is how many people go to that site

        self.grade_A_sites = self.init_grade(grade_code="A")
        self.grade_B_sites = self.init_grade(grade_code="B")
        self.grade_C_sites = self.init_grade(grade_code="C")
        self.house_sites = deepcopy(self.pop.household)
        self.house_indices = deepcopy(self.pop.house_ppl_i)

        # Students Stuff #
        self.stud_house_sites = deepcopy(self.pop.stud_houses)
        self.stud_house_indices = deepcopy(self.pop.house_stud_i)
        self.lect_sites = self.init_uni(grade_code="LECT")
        self.study_sites = self.init_uni(grade_code="STUDY")
        self.food_sites = self.init_uni(grade_code="FOOD")
        self.res_sites = self.init_res(grade_code="RES")

        self.daily_new_infections = 0

    def load_attributes_from_sim_obj(self, sim_obj):
        """Method to load in attributes from the provided simulation class object.

        Sets all objects in the "interaction_sites_data" dictionary key as self
        attributes of the InteractionSites class.

        Parameters
        ----------
        sim_obj : :obj:`cv19.simulation.simulation`
            The encompassing simulation obejct hosting the simulation.
        """

        attributes = sim_obj.parameters["interaction_sites_data"].keys()
        for attr in attributes:
            setattr(self, attr, sim_obj.parameters["interaction_sites_data"][attr])

        # Get the disease parameters
        d_attributes = sim_obj.disease_parameters["spread_data"].keys()
        for attr in d_attributes:
            setattr(self, attr, sim_obj.disease_parameters["spread_data"][attr])

        # Get the virus type names
        self.variant_codes = sim_obj.variant_codes
        self.variant_code_map = {v_id: v_name for v_name, v_id in self.variant_codes.items()}  # virus ids

        # Set the actual objects now
        self.pop = sim_obj.pop
        self.policy = sim_obj.policy

        self.nDays = sim_obj.parameters["simulation_data"]["nDays"]

    def init_grade(self, grade_code):
        """Method designed to associate members of the population with interaction sites.

        This method initializes all non-student interaction sites by creating a list
        of person indices for each interaction site, for that type of interaction type.

        Parameters
        ----------
        grade_code : str
            Code used to index the values to create this type of site from the config file.

        Returns
        -------
        grade_sites : :obj:`np.array` of :obj:`np.array` of :obj:`int`
            An array holding one array for each interaction site of this type. Each nested
            array holds the index of people that are associated with that site (can visit it).
        """

        loyalty_mean = self.grade_loyalty_means[grade_code]
        loyalty_std = self.grade_loyalty_stds[grade_code]
        students_interact = self.students_participate[grade_code]

        # Calculate number of sites
        num_sites = self.calculate_num_sites(grade_code=grade_code)
        grade_sites = [[] for _ in range(num_sites)]

        for person in self.pop.get_population():
            if students_interact or not (self.students_on and person.job == 'Student'):
                # if students are meant to go to this site
                # Assign people to this specific site
                num_diff_sites = abs(round(np.random.normal(loyalty_mean, loyalty_std)))
                num_diff_sites = num_diff_sites if num_diff_sites <= num_sites else num_sites
                # Get a list of len num_diff_sites for this person to be associated with now
                person_sites = np.random.choice(num_sites, num_diff_sites, replace=False)
                for site in person_sites:
                    # Assign this person to that site
                    grade_sites[site].append(person.get_index())

        # Convert everything to numpy arrays
        grade_sites = [np.asarray(site) for site in grade_sites]

        # Initialize the number of interactions dictionary
        self.daily_interactions[grade_code] = np.zeros(self.nDays)

        return grade_sites

    def init_uni(self, grade_code):
        """Method designed to associate members of the student population with interaction sites.

        This method initializes all student interaction sites by creating a list
        of person indices for each interaction site, for that type of interaction type.

        Parameters
        ----------
        grade_code : str
            Code used to index the values to create this type of site from the config file.

        Returns
        -------
        grade_sites : :obj:`np.array` of :obj:`np.array` of :obj:`int`
            An array holding one array for each interaction site of this type. Each nested
            array holds the index of people that are associated with that site (can visit it).
        """

        loyalty_mean = self.grade_loyalty_means[grade_code]
        loyalty_std = self.grade_loyalty_stds[grade_code]

        # Calculate number of sites
        num_sites = self.calculate_num_sites(grade_code=grade_code)
        grade_sites = [[] for _ in range(num_sites)]

        for student in self.pop.get_population():
            if student.job == 'Student':
                # Assign people to this specific site
                num_diff_sites = abs(round(np.random.normal(loyalty_mean, loyalty_std)))
                num_diff_sites = num_diff_sites if num_diff_sites <= num_sites else num_sites
                # Get a list of len num_diff_sites for this person to be associated with now
                student_sites = np.random.choice(num_sites, num_diff_sites, replace=False)
                for site in student_sites:
                    # Assign this person to that site
                    grade_sites[site].append(student.get_index())

        # Convert everything to numpy arrays
        grade_sites = [np.asarray(site) for site in grade_sites]

        # Initialize the number of interactions dictionary
        self.daily_interactions[grade_code] = np.zeros(self.nDays)

        return grade_sites

    def init_res(self, grade_code):
        """Method designed to associate students with the residence interaction site.

        This method initializes the residence interaction sites by creating a list
        of person indices for each interaction site.

        Parameters
        ----------
        grade_code : str
            Code used to index the values to create this type of site from the config file.

        Returns
        -------
        grade_sites : :obj:`np.array` of :obj:`np.array` of :obj:`int`
            An array holding one array for each interaction site of this type. Each nested
            array holds the index of people that are associated with that site (can visit it)
        """

        loyalty_mean = self.grade_loyalty_means[grade_code]
        loyalty_std = self.grade_loyalty_stds[grade_code]

        # Calculate number of sites
        num_sites = self.calculate_num_sites(grade_code=grade_code)
        grade_sites = [[] for _ in range(num_sites)]

        for room in self.pop.get_residences():
            for student_i in self.stud_house_indices[room]:
                # Assign people to this specific site
                num_diff_sites = abs(round(np.random.normal(loyalty_mean, loyalty_std)))
                num_diff_sites = num_diff_sites if num_diff_sites <= num_sites else num_sites
                # Get a list of len num_diff_sites for this person to be associated with now
                student_sites = np.random.choice(num_sites, num_diff_sites, replace=False)
                for site in student_sites:
                    # Assign this person to that site
                    grade_sites[site].append(student_i)

        # Convert everything to numpy arrays
        grade_sites = [np.asarray(site) for site in grade_sites]

        # Initialize the number of interactions dictionary
        self.daily_interactions[grade_code] = np.zeros(self.nDays)

        return grade_sites

    def daily_reset(self):
        """Method used to reset the interaction sites at the end of each day.

        This function is currently used to clean up dead agents from interaction sites,
        and to reset daily counts (such as the daily infection count).

        Parameters
        ----------
        None

        Returns
        -------
        None
        """

        self.remove_dead()
        self.daily_new_infections = 0

    def calculate_num_sites(self, grade_code):
        """Method used to calculate the number of sites for an interaction site grade.

        Parameters
        ----------
        grade_code : str
            Code used to index the values to create this type of site from the config file.

        Returns
        -------
        num_sites : int
            The number of sites to be used for that interaction site grade.
        """

        if grade_code in self.site_num and self.site_num[grade_code] == 0:
            # Raise a warning
            warnings.warn(f"Site type '{grade_code}' size set to 0. No interaction sites of this type created.")
            return 0
        else:
            return self.site_num[grade_code] if grade_code in self.site_num else \
                max(round(self.pop.get_population_size() / self.site_size[grade_code]), 1)

    def remove_dead(self):
        """Method to remove dead agents from interaction site arrays.

        Iterates through each type of site array, and will remove all agents that are
        dead from each array.

        Parameters
        ----------
        None

        Returns
        -------
        None
        """

        # Create list of all dead agents
        dead_agents = self.pop.get_dead()

        # Site type A
        for i, site_array in enumerate(self.grade_A_sites):
            # Mask where True indicates alive and False indicates dead (note the invert argument)
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.grade_A_sites[i] = site_array[mask_alive]

        # Site type B
        for i, site_array in enumerate(self.grade_B_sites):
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.grade_B_sites[i] = site_array[mask_alive]

        # Site type C
        for i, site_array in enumerate(self.grade_C_sites):
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.grade_C_sites[i] = site_array[mask_alive]

        # Site type lecture
        for i, site_array in enumerate(self.lect_sites):
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.lect_sites[i] = site_array[mask_alive]

        # Site type study
        for i, site_array in enumerate(self.study_sites):
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.study_sites[i] = site_array[mask_alive]

        # Site type food
        for i, site_array in enumerate(self.food_sites):
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.food_sites[i] = site_array[mask_alive]

        # Site type res
        for i, site_array in enumerate(self.res_sites):
            mask_alive = np.isin(site_array, dead_agents, invert=True)
            self.res_sites[i] = site_array[mask_alive]

    def will_visit_site(self, site_array, will_go_prob):
        """Method to determine who will visit a site on a given day.

        Generates a boolean list for each individual interaction site in site_array,
        indicating what people from the list will visit that site on a given day. Accounts
        for quarantined people by setting their will_go_prob value to
        self.quarantine_isolation_factor.

        Parameters
        ----------
        site_array : :obj:`np.array` of :obj:`np.array` of :obj:`int`
            An array holding lists (one for each interaction site) of the index of each person
            associated with each of the individual sites.
        will_go_prob : float
            The probability that any given person in site_array will visit this type of site.

        Returns
        -------
        will_visit_grade : :obj:`np.array` of :obj:`np.array` of :obj:`int`
            An array holding an array for each site of this interaction site type.
            Each individual list holds the indexes of people that will visit that site for this day.
        """

        # Figure out who is going to go to this site type today.
        person_ids = np.unique(np.concatenate(site_array))

        # Create array of attendence probabilities.
        prob_attendence = [self.quarantine_isolation_factor
                           if self.pop.get_person(person).is_quarantined() else will_go_prob
                           for person in person_ids]

        # Select a subset of people who will actually choose to go to the site.
        person_will_go_mask = np.random.binomial(1, p=prob_attendence).astype(bool)
        person_ids = person_ids[person_will_go_mask]

        # Create a Boolean array of people (rows) and sites (columns).
        # Each entry corresponds to whether or not
        # a given person can go to the given site.
        person_site_array = np.zeros(shape=(person_ids.shape[0], len(site_array)), dtype=bool)
        for s, site in enumerate(site_array):
            mask = np.isin(person_ids, site)
            person_site_array[mask, s] = True

        # Choose a random number for each person,
        # with an upper bound as the number of available sites for that person.
        high = person_site_array.sum(axis=-1)
        random_site_index = np.random.randint(low=0, high=high)

        # argsort the array (descending) along sites and use random number above to select
        # one of the available sites (first sites up to high[i] are available for person i).
        site_indexes_argsorted = np.argsort(person_site_array, axis=-1)[..., ::-1]
        person_site_index = site_indexes_argsorted[np.arange(site_indexes_argsorted.shape[0]),
                                                   random_site_index]

        will_visit_grade = [person_ids[np.where(person_site_index == s)[0]]
                            for s in range(len(site_array))]

        return will_visit_grade

    def site_interaction(self, will_go_array, day, personal, grade_code):
        """Method that hosts interactions between people for an interaction site type.

        This method manages interactions between people going to the same interaction
        site this day. Currently, all people that visit the same site on a day have a
        chance to interact with each other. Does not provide a return value, all
        infections are managed within the function.

        Parameters
        ----------
        will_go_array : :obj:`np.array` of :obj:`np.array` of :obj:`int`
            An array holding an array for each site of this interaction site type. Each individual list
            holds the indexes of people that will visit that site for this day.
        day : int
            The day value that this function is being called on in the encompassing simulation class.
            Used as input to the infect function after infections have been determined.
        personal : bool
            Used to indicate if the type of interaction at this site is personal, which relates to
            contact tracing abilities.
        grade_code : str
            Code used to index the values to create this type of site from the config file.
        """

        new_infections = np.zeros(self.pop.get_population_size(), dtype=bool)
        new_infection_type = np.zeros(self.pop.get_population_size(), dtype=int)

        total_interactions_count = 0

        for ppl_going in will_go_array:

            infected_persons = [index for index in ppl_going if self.pop.get_person(index).is_infected()]
            recovered_persons = [index for index in ppl_going if self.pop.get_person(index).is_recovered()]

            # Generate a list of how many interactions ppl have at the site
            num_interactions = self.calc_interactions(site_day_pop=len(ppl_going))
            total_interactions_count += np.sum(num_interactions) // 2

            if len(infected_persons) == 0 or (len(infected_persons) + len(recovered_persons) == len(ppl_going)):
                continue  # No ppl to infect here or no one already infected

            while np.sum(num_interactions > 0) > 1:
                # grab the highest interactor
                person_1 = np.argmax(num_interactions)
                # find a random interactor for them to pair with (that is not them)
                person_2 = np.random.randint(num_interactions.shape[0])
                while person_2 == person_1 or num_interactions[person_2] <= 0:
                    person_2 = np.random.randint(num_interactions.shape[0])

                # Get the actual people at these indexes
                person_1_index = ppl_going[person_1]
                person_2_index = ppl_going[person_2]

                # Getting the Person objects and logging the contacts
                p1_obj = self.pop.get_person(person_1_index)
                p2_obj = self.pop.get_person(person_2_index)
                p1_obj.log_contact(p2_obj, day=day, personal=personal)
                p2_obj.log_contact(p1_obj, day=day, personal=personal)

                # Check to make sure one is infected
                person_1_infected = p1_obj.is_infected()
                person_2_infected = p2_obj.is_infected()

                if person_1_infected != person_2_infected:
                    # Have an interaction between those people
                    did_infect = self.interact(p1_obj, p2_obj)
                    if did_infect:
                        if person_1_infected:
                            new_infections[person_2_index] = True
                            new_infection_type[person_2_index] = self.pop.get_person(person_1_index).get_virus_type()
                        else:
                            new_infections[person_1_index] = True
                            new_infection_type[person_1_index] = self.pop.get_person(person_2_index).get_virus_type()

                # Lower the interaction count for those people
                num_interactions[person_1] -= 1
                num_interactions[person_2] -= 1

        #  Update people who get infected only at the end. Assuming if I get CV19 at work, I probably won't spread at the store that night.
        new_infection_indexes = np.where(new_infections)[0]
        self.daily_new_infections += len(new_infection_indexes)
        for new_infection in new_infection_indexes:
            self.pop.infect(index=new_infection, virus_type=new_infection_type[new_infection], day=day)

        # Update total daily interactions count
        self.daily_interactions[grade_code][day] = total_interactions_count

    def calc_interactions(self, site_day_pop):
        """Method to determine how many interactions a person will have.

        Note
        ----
        Currently the distribution for the number of interactions a given person will have is
        a "triangular" distribution with only one side (a linear distribution). The distribution
        output spans from 0 to site_day_pop/day_hours_scaler, where it is much more likely to have 0
        interactions than the max. day_hours_scaler takes into account that people will not all be
        at the interaction site at the same time, but will be dispersed throughout the 12 hour day.

        As it stands, day_hours_scaler is not a config file parameter, as the hours in the day should not be
        adjusted between simulations. If the need is felt for an adjustable scaling factor, a new (second)
        variable should be introduced.

        Parameters
        ----------
        site_day_pop : `int`
            The total number of people at that specific interaction site this day.

        Returns
        -------
        number_of_interactions : :obj:`np.array` of :obj:`int`
            The number of interactions all people will have within this interaction site.
        """

        day_hours_scaler = 12

        if site_day_pop == 0:
            return np.array([])
        else:
            # Generate a linaer distribution from
            number_of_interactions = np.round(np.random.triangular(left=0, mode=0, right=site_day_pop / day_hours_scaler,
                                                                   size=site_day_pop)).astype(int)

        return number_of_interactions

    def interact(self, person_1, person_2):
        """Method that models the interaction between two people.

        Parameters
        ----------
        person_1 : :obj:`cv19.person.Person`
            First person in the two-way interaction.
        person_2 : :obj:`cv19.person.Person`
            Second person in the two-way interaction.

        Returns
        -------
        : :obj:`bool`
            Whether or not the interaction caused the spread of the infection.
        """

        p1_infected = person_1.is_infected()
        p2_infected = person_2.is_infected()

        virus_type = person_1.get_virus_type() if p1_infected else person_2.get_virus_type()
        spread_prob = self.base_infection_spread_prob[self.variant_code_map[virus_type]]

        if self.policy.get_mask_mandate():
            p1_mask = person_1.wear_mask()
            p2_mask = person_2.wear_mask()

            P1_INWARD_EFF, P1_OUTWARD_EFF = person_1.mask_type_efficiency()
            P2_INWARD_EFF, P2_OUTWARD_EFF = person_2.mask_type_efficiency()

            if p1_infected:
                if p1_mask:
                    spread_prob *= (1 - P1_OUTWARD_EFF)
                if p2_mask:
                    spread_prob *= (1 - P2_INWARD_EFF)

            elif p2_infected:
                if p1_mask:
                    spread_prob *= (1 - P1_INWARD_EFF)
                if p2_mask:
                    spread_prob *= (1 - P2_OUTWARD_EFF)

        p1_vaccinated1 = person_1.is_vaccinated()
        p2_vaccinated1 = person_2.is_vaccinated()

        p1_vaccine_eff = person_1.vaccine_type_efficiency() if p1_vaccinated1 else 0
        p2_vaccine_eff = person_2.vaccine_type_efficiency() if p2_vaccinated1 else 0

        spread_prob *= ((1 - p1_vaccine_eff) * (1 - p2_vaccine_eff))

        return random() < spread_prob

    def house_interact(self, day):
        """Method to manage interactions between members of the same household.

        Determines if any infection will spread among members of the same household. Different
        from interaction sites in the fact that contacts are not calculated, but assumed to happen
        between all house members. Does not have a return value, infections are managed internally.

        Parameters
        ----------
        day : int
            The day value that this function is being called on in the encompassing simulation class.
            Used as input to the infect function after infections have been determined.
        """

        total_house_interactions = 0
        for house_indices in self.house_indices:
            # Get people in house
            house_size = len(house_indices)
            housemembers = [self.pop.get_population()[ind] for ind in house_indices]
            virus_types = [person.get_virus_type() for person in housemembers]

            total_house_interactions += comb(len(housemembers), 2)

            # Do interactions between the housemates
            for member1, member2 in combinations(housemembers, 2):
                member1.log_contact(member2, day=day, personal=True)
                member2.log_contact(member1, day=day, personal=True)

            # Check if anyone in the house is infected
            if any(housemembers[i].is_infected() for i in range(house_size)):
                infected_housemembers = [i for i in range(house_size) if housemembers[i].is_infected()]
                virus_types = [virus_types[i] for i in infected_housemembers]
                healthy_housemembers = [i for i in range(house_size) if not housemembers[i].is_infected()]

                for person in healthy_housemembers:
                    virus_id = np.random.choice(a=virus_types)
                    virus_name = self.variant_code_map[virus_id]

                    infection_chance = self.base_infection_spread_prob[virus_name] * self.house_infection_spread_factor
                    person_vaccinated = housemembers[person].is_vaccinated()
                    person_vaccine_eff = housemembers[person].vaccine_type_efficiency() if person_vaccinated else 0
                    infection_chance *= (1 - person_vaccine_eff)
                    caught_infection = random() < infection_chance

                    if caught_infection:
                        self.daily_new_infections += 1
                        if virus_id is None:
                            raise ValueError("House infection has incorrect virus type.")
                        self.pop.infect(index=housemembers[person].get_index(), day=day, virus_type=virus_id)

        self.daily_interactions["HOUSE_GENERAL"][day] = total_house_interactions

    def student_house_interact(self, day):
        """Method to manage interactions between members of the same student household.

        Determines if any infection will spread among members of the same household. Different
        from interaction sites in the fact that contacts are not calculated, but assumed to happen
        between all house members. Does not have a return value, infections are managed internally.

        Parameters
        ----------
        day : int
            The day value that this function is being called on in the encompassing simulation class.
            Used as input to the infect function after infections have been determined.
        """

        total_house_interactions = 0
        for house_indices in self.stud_house_indices:
            # Get people in house
            house_size = len(house_indices)
            housemembers = [self.pop.get_population()[ind] for ind in house_indices]
            virus_types = [person.get_virus_type() for person in housemembers]

            total_house_interactions += comb(len(housemembers), 2)

            # Do interactions between the housemates
            for member1, member2 in combinations(housemembers, 2):
                member1.log_contact(member2, day=day, personal=True)
                member2.log_contact(member1, day=day, personal=True)

            # Check if anyone in the house is infected
            if any(housemembers[i].is_infected() for i in range(house_size)):
                infected_housemembers = [i for i in range(house_size) if housemembers[i].is_infected()]
                virus_types = [virus_types[i] for i in infected_housemembers]
                healthy_housemembers = [i for i in range(house_size) if not housemembers[i].is_infected()]

                for person in healthy_housemembers:
                    virus_id = np.random.choice(a=virus_types)
                    virus_name = self.variant_code_map[virus_id]

                    infection_chance = self.base_infection_spread_prob[virus_name] * self.house_infection_spread_factor
                    person_vaccinated = housemembers[person].is_vaccinated()
                    person_vaccine_eff = housemembers[person].vaccine_type_efficiency() if person_vaccinated else 0
                    infection_chance *= (1 - person_vaccine_eff)
                    caught_infection = random() < infection_chance
                    if caught_infection:
                        self.daily_new_infections += 1
                        if virus_id is None:
                            raise ValueError("House infection has incorrect virus type.")
                        self.pop.infect(index=housemembers[person].get_index(), day=day, virus_type=virus_id)

        self.daily_interactions["HOUSE_STUDENT"][day] = total_house_interactions

    def testing_site(self, tests_per_day, day):
        """Method to update status of symptoms and run the testing sites code.

        Parameters
        ----------
        tests_per_day : int
            The max number of available tests for this given day.
        day : int
            The day value that this function is being called on in the encompassing simulation class.
        """

        self.pop.update_uninfected_symptomatics()
        self.pop.update_infected_symptomatics(day)
        self.pop.get_tested(tests_per_day, day)

    def get_grade_A_sites(self):
        """Method to return a copy of the grade_A_sites attribute.

        Returns
        -------
        self.grade_A_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.grade_A_sites)

    def get_grade_B_sites(self):
        """Method to return a copy of the grade_B_sites attribute.

        Returns
        -------
        self.grade_B_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.grade_B_sites)

    def get_grade_C_sites(self):
        """Method to return a copy of the grade_C_sites attribute.

        Returns
        -------
        self.grade_C_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.grade_C_sites)

    def get_lect_sites(self):
        """Method to return a copy of the lect_sites attribute.

        Returns
        -------
        self.lect_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.lect_sites)

    def get_study_sites(self):
        """Method to return a copy of the study_sites attribute.

        Returns
        -------
        self.study_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.study_sites)

    def get_food_sites(self):
        """Method to return a copy of the food_sites attribute.

        Returns
        -------
        self.food_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.food_sites)

    def get_res_sites(self):
        """Method to return a copy of the res_sites attribute.

        Returns
        -------
        self.res_sites.copy() : :obj:`np.array` of :obj:`list` of :obj:`np.array` of :obj:`int`
        """
        return deepcopy(self.res_sites)


import os
import re
import shutil
from glob import glob

from osrf_pycommon.process_utils import AsyncSubprocessProtocol

from catkin_tools.common import mkdir_p
from catkin_tools.terminal_color import fmt

from .events import ExecutionEvent

MAX_LOGFILE_HISTORY = 10

class IOBufferContainer(object):

    """A simple buffer container for use in logging.

    This class will open a logfile for a given job stage and write to it
    continuously while receiving stdout and stderr.
    """

    def __init__(self, label, job_id, stage_label, event_queue, log_path):
        self.label = label
        self.job_id = job_id
        self.stage_label = stage_label
        self.event_queue = event_queue
        self.log_path = log_path

        self.is_open = False
        self.stdout_buffer = b""
        self.stderr_buffer = b""
        self.interleaved_buffer = b""

        # Construct the logfile path for this job and stage
        logfile_dir_path = os.path.join(log_path, self.job_id)
        self.logfile_basename = os.path.join(logfile_dir_path, '.'.join([self.label, self.stage_label]))
        self.logfile_name = '{}.log'.format(self.logfile_basename)

        # Create the logfile dir if it doesn't exist
        if not os.path.exists(logfile_dir_path):
            mkdir_p(logfile_dir_path)

        # Get the existing number of logfiles
        # TODO: Make this number global across all build stages
        existing_logfile_indices = sorted([int(lf.split('.')[-2])
                                           for lf in glob('{}.*.log'.format(self.logfile_basename))])
        if len(existing_logfile_indices) == 0:
            self.logfile_index = 0
        else:
            self.logfile_index = 1 + existing_logfile_indices[-1]

        # Generate the logfile name
        self.unique_logfile_name = '{}.{:0>{}}.log'.format(self.logfile_basename, self.logfile_index, 3)

        # Remove colliding file if necessary
        if os.path.exists(self.logfile_name):
            os.unlink(self.logfile_name)

        # Open logfile
        self.log_file = open(self.logfile_name, 'wb')
        self.is_open = True

    def close(self):
        # Close logfile
        self.log_file.close()
        self.is_open = False

        # Copy logfile to unique name
        shutil.copy(self.logfile_name, self.unique_logfile_name)

        # Remove older logfiles
        for logfile_name in glob('{}.*.log'.format(self.logfile_basename)):
            if (self.logfile_index - int(logfile_name.split('.')[-2])) >= MAX_LOGFILE_HISTORY:
                os.unlink(logfile_name)

        # Save output from stderr (these don't get deleted until cleaning the logfile directory)
        if len(self.stderr_buffer) > 0:
            with open(self.unique_logfile_name + '.stderr', 'wb') as logfile:
                logfile.write(self.stderr_buffer)

    def get_interleaved_log(self):
        """get decoded interleaved log."""
        try:
            return self._decode(self.interleaved_buffer)
        except UnicodeDecodeError:
            return "interleaved_log: some output cannot be displayed.\n"

    def get_stdout_log(self):
        """get decoded stdout log."""
        try:
            return self._decode(self.stdout_buffer)
        except UnicodeDecodeError:
            return "stdout_log: some output cannot be displayed.\n"

    def get_stderr_log(self):
        """get decoded stderr log."""
        try:
            return self._decode(self.stderr_buffer)
        except UnicodeDecodeError:
            return "stderr_log: some output cannot be displayed.\n"

    @staticmethod
    def _encode(data):
        """Encode a Python str into bytes.
        :type data: str
        """
        return data.encode('utf8')

    @staticmethod
    def _decode(data):
        """Decode bytes into Python str.
        :type data: bytes
        """
        return data.decode('utf-8', 'replace')

    def __del__(self):
        if self.is_open:
            self.close()

    @classmethod
    def factory(cls, label, job_id, stage_label, event_queue, log_path):
        """Factory method for constructing with job metadata."""

        def init_proxy(*args, **kwargs):
            return cls(label, job_id, stage_label, event_queue, log_path, *args, **kwargs)

        return init_proxy

class IOBufferLogger(IOBufferContainer):

    """This is a logging class to be used instead of sys.stdout and sys.stderr
    in FunctionStage operations.

    This class also generates `stdout` and `stderr` events.
    """

    def __init__(self, label, job_id, stage_label, event_queue, log_path, *args, **kwargs):
        IOBufferContainer.__init__(self, label, job_id, stage_label, event_queue, log_path)

    def out(self, data, end='\n'):
        """
        :type data: str
        :type end: str
        """
        # Buffer the encoded data
        data += end
        encoded_data = self._encode(data)
        self.stdout_buffer += encoded_data
        self.interleaved_buffer += encoded_data

        # Save the encoded data
        self.log_file.write(encoded_data)

        # Emit event with decoded Python str
        self.event_queue.put(ExecutionEvent(
            'STDOUT',
            job_id=self.job_id,
            stage_label=self.stage_label,
            data=data))

    def err(self, data, end='\n'):
        """
        :type data: str
        :type end: str
        """
        # Buffer the encoded data
        data += end
        encoded_data = self._encode(data)
        self.stderr_buffer += encoded_data
        self.interleaved_buffer += encoded_data

        # Save the encoded data
        self.log_file.write(encoded_data)

        # Emit event with decoded Python str
        self.event_queue.put(ExecutionEvent(
            'STDERR',
            job_id=self.job_id,
            stage_label=self.stage_label,
            data=data))

class IOBufferProtocol(IOBufferContainer, AsyncSubprocessProtocol):

    """An asyncio protocol that collects stdout and stderr.

    This class also generates `stdout` and `stderr` events.

    Since the underlying asyncio API constructs the actual protocols, this
    class provides a factory method to inject the job and stage information
    into the created protocol.
    """

    def __init__(self, label, job_id, stage_label, event_queue, log_path, *args, **kwargs):
        IOBufferContainer.__init__(self, label, job_id, stage_label, event_queue, log_path)
        AsyncSubprocessProtocol.__init__(self, *args, **kwargs)

        self.intermediate_stdout_buffer = b''
        self.intermediate_stderr_buffer = b''

    @staticmethod
    def _split(data):
        try:
            last_break = data.rindex(b'\n') + 1
            return data[0:last_break], data[last_break:]
        except ValueError:
            return b'', data

    def on_stdout_received(self, data):
        """
        :type data: encoded bytes
        """

        data, self.intermediate_stdout_buffer = self._split(self.intermediate_stdout_buffer + data)

        self.stdout_buffer += data
        self.interleaved_buffer += data
        self.log_file.write(data)

        # Get the decoded Python str
        decoded_data = self._decode(data)

        # Emit event with decoded Python str
        self.event_queue.put(ExecutionEvent(
            'STDOUT',
            job_id=self.job_id,
            stage_label=self.stage_label,
            data=decoded_data))

    def on_stderr_received(self, data):
        """
        :type data: encoded bytes
        """

        data, self.intermediate_stderr_buffer = self._split(self.intermediate_stderr_buffer + data)

        self.stderr_buffer += data
        self.interleaved_buffer += data
        self.log_file.write(data)

        # Get the decoded Python str
        decoded_data = self._decode(data)

        # Emit event with decoded Python str
        self.event_queue.put(ExecutionEvent(
            'STDERR',
            job_id=self.job_id,
            stage_label=self.stage_label,
            data=decoded_data))

    def on_process_exited2(self, returncode):
        """
        Dump anything remaining in the intermediate buffers.
        """

        if len(self.intermediate_stdout_buffer) > 0:
            self.on_stdout_received(self.intermediate_stdout_buffer + b'\n')
        if len(self.intermediate_stderr_buffer) > 0:
            self.on_stderr_received(self.intermediate_stderr_buffer + b'\n')

class CatkinTestResultsIOBufferProtocol(IOBufferProtocol):
    """An IOBufferProtocol which parses the output of catkin_test_results"""
    def on_stdout_received(self, data):
        lines = data.decode().splitlines()
        clines = []
        for line in lines:
            match = re.match(r'(.*): (\d+) tests, (\d+) errors, (\d+) failures, (\d+) skipped', line)
            if match:
                line = fmt('@!{}@|: {} tests, @{rf}{} errors@|, @{rf}{} failures@|, @{kf}{} skipped@|')
                line = line.format(*match.groups())
            clines.append(line)

        cdata = '\n'.join(clines) + '\n'

        super(CatkinTestResultsIOBufferProtocol, self).on_stdout_received(cdata.encode())


from scripts.util.joystick import Joystick
from carla_env import CarlaEnv
from wrapped_carla_env import BiasedAction

import time
import carla
import pygame

import numpy as np

class ManualInterface:
    def __init__(self, env: CarlaEnv):
        # create env
        self.env = env

        self.obs = None
        self.rew = None
        self.done = False
        self.reset_env()

        self.total_reward = None

        # init pygame
        pygame.init()

        self.running = True
        self.surface = None

        # font
        self.font = pygame.font.Font(pygame.font.get_default_font(), 30)

        # control
        self.joysticks = None
        self.reset_flag = False

        # tps counter
        self.tps_total_time = 0
        self.tps_total_frame = 0

    def __del__(self):
        pygame.quit()

    def reset_env(self):
        self.obs = self.env.reset()
        self.rew = None
        self.done = False

        self.total_reward = None

    def get_action(self):
        # init joysticks
        if self.joysticks is None:
            self.joysticks = [Joystick(fn) for fn in Joystick.list_devices()]

        # pump events before get
        pygame.event.pump()

        # get joysticks
        act = []
        for js in self.joysticks:
            accel = -js.axes["ry"]  # Left lever, L <--> R
            steer = js.axes["x"]   # Right lever, U <--> D
            reverse = js.buttons["tl"]    # LB

            # act.append(carla.VehicleControl(
            #     throttle=max(0.0, accel),
            #     brake=-min(0.0, accel),
            #
            #     steer=steer,
            #     reverse=reverse
            # ))
            act.append(np.array([accel, steer], dtype=np.float32))

        # check if reset
        is_reset = sum([js.buttons["y"] for js in self.joysticks])
        is_reset |= self.done
        if is_reset:
            if not self.reset_flag:
                self.reset_env()
                self.reset_flag = True
        else:
            self.reset_flag = False

        return act

    def on_event(self, event):
        if event.type == pygame.QUIT:
            self.running = False
            return

    def on_update(self):
        # update env
        act = self.get_action()

        start_time = time.time()
        self.obs, self.rew, self.done, _, = self.env.step(act)
        elapsed = time.time() - start_time

        # update total reward
        if self.total_reward is None:
            self.total_reward = np.array(self.rew)
        else:
            self.total_reward += np.array(self.rew)

        # tps counter
        self.tps_total_time += elapsed
        self.tps_total_frame += 1
        if self.tps_total_frame >= 100:
            print("TPS: {}".format(self.tps_total_frame / self.tps_total_time))

            self.tps_total_frame = 0
            self.tps_total_time = 0

    def on_render(self):
        _, h, w = self.obs[0][0].shape

        # init surface
        if self.surface is None:
            n_cars = len(self.obs)
            n_cameras = len(self.obs[0])

            self.surface = pygame.display.set_mode((w * n_cameras, h * n_cars), pygame.HWSURFACE | pygame.DOUBLEBUF)

        # show images
        y = 0
        for cam, rew in zip(self.obs, self.total_reward):
            # draw car cam images
            x = 0
            for cam_img in cam:
                if cam_img.shape[0] < 3:
                    # pad channel
                    padded_cam_img = np.concatenate([
                        cam_img,
                        np.zeros((1, *cam_img.shape[1:]), dtype=cam_img.dtype)], axis=0)
                else:
                    padded_cam_img = cam_img

                cam_surf = pygame.surfarray.make_surface(padded_cam_img.transpose(2, 1, 0))
                self.surface.blit(cam_surf, (x, y))

                x += w

            # draw reward
            rew_surf = self.font.render("Reward: {:.2f}".format(rew), True, (0, 0, 255))
            self.surface.blit(rew_surf, (10, y + 10))

            y += h

        # update display
        pygame.display.update()

    def run(self, fps: int = None):
        clock = pygame.time.Clock()

        while self.running:
            for event in pygame.event.get():
                self.on_event(event)

            self.on_update()
            self.on_render()

            if fps:
                clock.tick(fps)

def main():
    dt = 0.1
    global_options = {
        "world": {
            "dt": dt
        }
    }
    env = CarlaEnv(global_options, 0)

    ui = ManualInterface(env)
    ui.run(int(1.0 / dt))

if __name__ == '__main__':
    main()


import os
import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt

model_type = "DPT_Large"

midas = torch.hub.load("intel-isl/MiDaS", model_type)

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()

midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")

if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
    transform = midas_transforms.dpt_transform
else:
    transform = midas_transforms.small_transform

input_dir = 'dataset/dop/speaker/output/'  # directory with your input images
output_dir = 'dataset/validation/depth/speaker/'  # directory where you want to save depth maps
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

for filename in os.listdir(input_dir):
    if filename.endswith(".jpg") or filename.endswith(".png"):
        img_path = os.path.join(input_dir, filename)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        input_batch = transform(img).to(device)

        with torch.no_grad():
            prediction = midas(input_batch)

            prediction = torch.nn.functional.interpolate(
                prediction.unsqueeze(1),
                size=img.shape[:2],
                mode="bicubic",
                align_corners=False,
            ).squeeze()

        output = prediction.cpu().numpy()
        output = np.sqrt((output - np.min(output)) / np.ptp(output))  # apply square root to decrease contrast

        # Apply color map
        output_colored = cv2.applyColorMap((output * 255).astype(np.uint8), cv2.COLORMAP_JET)

        # Save depth map
        output_colored_bgr = cv2.cvtColor(output_colored, cv2.COLOR_RGB2BGR)
        output_filename = os.path.join(output_dir, filename.split('.')[0] + '_depth.jpg')
        cv2.imwrite(output_filename, output_colored_bgr)

        # Show depth map
        # plt.imshow(output_colored)
        # plt.show()

print('Depth maps created and saved.')


import time, datetime
import numpy as np
import os
import os.path as osp
import torch
import torchvision
import matplotlib.pyplot as plt
import torchvision.utils as vutils
import torch.nn.functional as F
import cv2
import glob
import random

from lib.utils.eval_utils import (
    batch_compute_similarity_transform_torch,
)
from lib.utils.geometry import batch_rodrigues

def inverse_normalize(tensor):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    return tensor

def normalize(tensor):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    for t, m, s in zip(tensor, mean, std):
        t.sub_(m).div_(s)
    return tensor

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        if self.count != 0:
            self.avg = self.sum / self.count

def get_HHMMSS_from_second(seconds):
    HHMMSS = time.strftime("%H:%M:%S", time.gmtime(seconds))
    return HHMMSS

def save_checkpoint(state, output_dir, filename):
    """ save model parameter """

    save_path = "results/{}/save_pth".format(output_dir)
    if not osp.exists(save_path):
        os.makedirs(save_path)

    torch.save(state, osp.join(save_path, filename))

def save_img(img, output_dir, filename, epoch=None, img_path=None, test_dataset=None, test_idx=None, vflip=False, hflip=False):
    """ save image """
    if isinstance(img, torch.Tensor):
        img = img.cpu()

    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)

    elif epoch and img_path:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    elif test_dataset:
        save_path = "results/{}/eval_data/{}_{}/{}".format(output_dir, test_dataset, epoch, test_idx)

    if not osp.exists(save_path):
        os.makedirs(save_path)

    # For jigsaw image with tile
    if len(img.shape) == 4: # tile, channel, width, height
        for tile_idx, _img_tile in enumerate(img):
            _filename = "_".join([filename.replace(".jpg", ""), str(tile_idx)]) + ".jpg"

            # _img = np.transpose(vutils.make_grid(_img_tile, normalize=True), (1, 2, 0))
            # _img = _img.numpy()
            _img = inverse_normalize(_img_tile).detach().numpy()
            _img = np.transpose(_img, (1, 2, 0))
            if vflip:
                _img = _img[::-1, :, :]
            if hflip:
                _img = _img[:, ::-1, :]
            plt.imsave(osp.join(save_path, _filename), _img)

    # For a image
    elif len(img.shape) == 3: # channel, width, height
        if isinstance(img, torch.Tensor):
            # if filename == "reconst_img.jpg":
            #     img = np.transpose(vutils.make_grid(img, normalize=True), (1, 2, 0))
            #     img = img.numpy()
            #     plt.imsave(osp.join(save_path, filename), img)
            if filename in ["detach.jpg", "rendering.jpg", "rendering_ren.jpg", "rendering_bg.jpg", "rendering1.jpg", "rendering_ren1.jpg",
                            "rendering2.jpg", "rendering_ren2.jpg", "rendering3.jpg", "rendering_ren3.jpg",
                            "rendering4.jpg", "rendering_ren4.jpg"]:
                img = np.transpose(img, (1, 2, 0))
                img = img.numpy()
                cv2.imwrite(osp.join(save_path, filename), 255*img[:, :, ::-1])
            else:
                img = inverse_normalize(img).detach().numpy()
                img = np.transpose(img, (1, 2, 0))
                plt.imsave(osp.join(save_path, filename), img)
        # if vflip:
        #     img = img[::-1, :, :]
        # if hflip:
        #     img = img[:, ::-1, :]

    elif len(img.shape) == 2: # width, height
        if isinstance(img, torch.Tensor):
            img = img[None, :, :]
            img = np.transpose(vutils.make_grid(img, normalize=True), (1, 2, 0))
            img = img.numpy()
            # img = inverse_normalize(img).detach().numpy()[:, :, None]
            # img = np.transpose(img, (1, 2, 0))
        if vflip:
            img = img[::-1, :, :]
        if hflip:
            img = img[:, ::-1, :]

        plt.imsave(osp.join(save_path, filename), img)

def save_all_img(img_dict, output_dir, epoch=None, img_path=None, test_dataset=None, test_idx=None, vflip=False, hflip=False):
    """
        img_dict keys: filename, value: image
    """
    for filename, img in img_dict.items():
        save_img(img, output_dir, filename, epoch, img_path, test_dataset, test_idx, vflip, hflip)

def save_mesh(verts, faces, output_dir, epoch, img_path=None, test_dataset=None, test_idx=None):
    """ save verts """
    filename = "mesh.obj"
    img = verts.cpu().numpy()
    faces = faces.cpu().numpy()

    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)

    elif test_dataset is not None and test_idx is not None:
        save_path = "results/{}/eval_data/{}_{}/{}".format(output_dir, test_dataset, epoch, test_idx)

    else:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    if not osp.exists(save_path):
        os.makedirs(save_path)

    if len(verts.shape) == 2:
        with open(osp.join(save_path, filename), "w") as f:
            for verts_xyz in verts:
                f.write("v {} {} {}\n".format(verts_xyz[0], verts_xyz[1], verts_xyz[2]))

            for face in faces:
                f.write("f {} {} {}\n".format(face[0]+1, face[1]+1, face[2]+1))

def save_templates_info(test_templates, output_dir, filename):
    """ save templates
        save results information or train information """

    save_path = "results/{}/save_txt".format(output_dir)
    if not osp.exists(save_path):
        os.makedirs(save_path)

    _filename = osp.join(save_path, filename)

    with open(_filename, "w") as f:
        f.writelines("\n".join(test_templates))

def save_joints2d_img(gt_keypoints_2d, pred_keypoints_2d, output_dir, epoch, img_path):
    """
        save image of joints2d on 2d coordinate while training
    """
    line_list = [
        [ 0, 1 ],
        [ 1, 2 ],
        [ 3, 4 ],
        [ 4, 5 ],
        [ 6, 7 ],
        [ 7, 8 ],
        [ 8, 2 ],
        [ 8, 9 ],
        [ 9, 3 ],
        [ 2, 3 ],
        [ 8, 12],
        [ 9, 10],
        [12, 9 ],
        [10, 11],
        [12, 13],
    ]

    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)
    else:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    if not osp.exists(save_path):
        os.makedirs(save_path)

    plt.figure()
    ax = plt.subplot()
    ax.invert_yaxis()

    gt_2ds = gt_keypoints_2d.clone()
    gt_2ds = gt_2ds.cpu()

    pred_2ds = pred_keypoints_2d.clone()
    pred_2ds = pred_2ds.cpu().detach()

    for joint_idx, (gt_2d, pred_2d) in enumerate(zip(gt_2ds, pred_2ds)):
        ax.scatter(gt_2d[0], gt_2d[1], marker='o', s=2, c="r")
        ax.text(gt_2d[0], gt_2d[1], joint_idx+1)
        ax.scatter(pred_2d[0], pred_2d[1], marker='o', s=2, c="b")
        ax.text(pred_2d[0], pred_2d[1], joint_idx+1)

    for start_point, end_point in line_list:
        start_point = start_point
        end_point = end_point
        ax.plot([gt_2ds[start_point][0], gt_2ds[end_point][0]], [gt_2ds[start_point][1], gt_2ds[end_point][1]], "r", linewidth=2)
        ax.plot([pred_2ds[start_point][0], pred_2ds[end_point][0]], [pred_2ds[start_point][1], pred_2ds[end_point][1]], "b", linewidth=2)

    plt.savefig(osp.join(save_path, "joints2d.jpg"))
    plt.close()

def save_joints3d_img(gt_keypoints_3d, pred_keypoints_3d, output_dir, epoch=None, img_path=None, test_dataset=None, test_idx=None):
    """
        save image of joints3d on 3d coordinate while training
    """
    line_list = [
        [ 0, 1 ],
        [ 1, 2 ],
        [ 3, 4 ],
        [ 4, 5 ],
        [ 6, 7 ],
        [ 7, 8 ],
        [ 8, 2 ],
        [ 8, 9 ],
        [ 9, 3 ],
        [ 2, 3 ],
        [ 8, 12],
        [ 9, 10],
        [12, 9 ],
        [10, 11],
        [12, 13],
    ]
    if isinstance(img_path, int):
        img_path = str(img_path)
        root = "results/{}/save_output_train".format(output_dir)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, img_path)

    elif epoch and img_path:
        dataset = img_path.split("/")[2]
        root = "results/{}/save_output_train/{}".format(output_dir, dataset)
        epoch = "epoch"+str(epoch)
        save_path = osp.join(root, epoch, ("_".join(img_path.split("/")[-2:])).replace(".jpg", ""))

    elif test_dataset is not None and test_idx is not None:
        save_path = "results/{}/eval_data/{}_{}/{}".format(output_dir, test_dataset, epoch, test_idx)

    if not osp.exists(save_path):
        os.makedirs(save_path)

    plt.figure()
    ax = plt.subplot(1, 1, 1, projection='3d')
    gt_3ds = gt_keypoints_3d.clone()
    gt_3ds = gt_3ds.cpu()
    pred_3ds = pred_keypoints_3d.clone()
    pred_3ds = pred_3ds.cpu().detach()
    for joint_idx, (gt_3d, pred_3d) in enumerate(zip(gt_3ds, pred_3ds)):
        ax.scatter(gt_3d[0], gt_3d[1], gt_3d[2], marker='o', s=2, c="r")
        ax.text(gt_3d[0], gt_3d[1], gt_3d[2], joint_idx+1)
        ax.scatter(pred_3d[0], pred_3d[1], pred_3d[2], marker='o', s=2, c="b")
        ax.text(pred_3d[0], pred_3d[1], pred_3d[2], joint_idx+1)

    for start_point, end_point in line_list:
        start_point = start_point
        end_point = end_point
        ax.plot(
            [gt_3ds[start_point][0], gt_3ds[end_point][0]],
            [gt_3ds[start_point][1], gt_3ds[end_point][1]],
            [gt_3ds[start_point][2], gt_3ds[end_point][2]],
            color="r",
            linewidth=2)
        ax.plot(
            [pred_3ds[start_point][0], pred_3ds[end_point][0]],
            [pred_3ds[start_point][1], pred_3ds[end_point][1]],
            [pred_3ds[start_point][2], pred_3ds[end_point][2]],
            color="b",
            linewidth=2)

    plt.savefig(osp.join(save_path, "joints3d.jpg"))
    plt.close()

def get_acc(output, label):
    batch_size = output.shape[0]
    pred = torch.argmax(output, dim=1)
    correct = (pred==label).sum().item()
    acc = correct/batch_size * 100
    return acc

def spin2h36m_joint(spins, device):
    """
        Get h36m 14 joints from spin 49 joints
    """
    convert_matrix = torch.zeros((14, 49)).to(device)
    h36m_index_list = [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38]
    for idx, h36m_index in enumerate(h36m_index_list):
        convert_matrix[idx, h36m_index] = 1
    h36m_joints = torch.matmul(convert_matrix, spins)
    return h36m_joints

def get_spherical_coords(X):
    # X is N x 3
    rad = np.linalg.norm(X, axis=1)
    # Inclination
    theta = np.arccos(X[:, 2] / rad)
    # Azimuth
    phi = np.arctan2(X[:, 1], X[:, 0])

    # Normalize both to be between [-1, 1]
    vv = (theta / np.pi) * 2 - 1
    uu = ((phi + np.pi) / (2*np.pi)) * 2 - 1
    # Return N x 2
    return np.stack([uu, vv],1)

def compute_uvsampler(verts, faces, tex_size=2):
    """
    For this mesh, pre-computes the UV coordinates for
    F x T x T points.
    Returns F x T x T x 2
    """
    alpha = np.arange(tex_size, dtype=np.float) / (tex_size-1)
    beta = np.arange(tex_size, dtype=np.float) / (tex_size-1)
    import itertools
    # Barycentric coordinate values
    coords = np.stack([p for p in itertools.product(*[alpha, beta])])
    vs = verts[faces]
    # Compute alpha, beta (this is the same order as NMR)
    v2 = vs[:, 2]
    v0v2 = vs[:, 0] - vs[:, 2]
    v1v2 = vs[:, 1] - vs[:, 2]
    # F x 3 x T*2
    samples = np.dstack([v0v2, v1v2]).dot(coords.T) + v2.reshape(-1, 3, 1)
    # F x T*2 x 3 points on the sphere
    samples = np.transpose(samples, (0, 2, 1))

    # Now convert these to uv.
    uv = get_spherical_coords(samples.reshape(-1, 3))
    # uv = uv.reshape(-1, len(coords), 2)

    uv = uv.reshape(-1, tex_size, tex_size, 2)
    return uv

def train_only_3task_network(
        HMR,
        context_encoder_net,
        discriminator,
        jigsaw_puzzle_net,
        rotation_net,

        loss_fn_BCE,
        loss_fn_MSE,
        loss_fn_CE,

        losses_CE,
        losses_DC,
        acces_JP,
        losses_JP,
        acces_ROT,
        losses_ROT,

        discriminator_optimizer,
        context_encoder_optimizer,
        jigsaw_puzzle_optimizer,
        rotation_optimizer,

        img,
        context_encoder_input,
        center_crop_img,
        jigsaw_input,
        jigsaw_order,
        rotation_img,
        rotation_idx,

        num_patch,
        ones,
        zeros,
        batch_size,
    ):
    ### Context Encoder ###
    # Update Discriminator

    feature_ce = HMR(context_encoder_input, return_only_features=True)
    feature_ce = feature_ce.reshape(-1, 2048, 1, 1)
    output_ce = context_encoder_net(feature_ce)
    output_ce_224 = context_encoder_input.clone()
    output_ce_224[:, :, 80:144, 80:144] = output_ce.clone()
    output_fake = discriminator(output_ce_224)
    output_real = discriminator(img)
    loss_BCE_fake = loss_fn_BCE(output_fake, zeros)
    loss_BCE_real = loss_fn_BCE(output_real, ones)
    loss_BCE = loss_BCE_fake + loss_BCE_real
    losses_DC.update(loss_BCE.item(), batch_size)
    discriminator_optimizer.zero_grad()
    loss_BCE.backward()
    discriminator_optimizer.step()

    # Update Decoder
    feature_ce = HMR(context_encoder_input, return_only_features=True)
    feature_ce = feature_ce.reshape(-1, 2048, 1, 1)
    output_ce = context_encoder_net(feature_ce)
    output_ce_224 = context_encoder_input.clone()
    output_ce_224[:, :, 80:144, 80:144] = output_ce.clone()
    output_fake = discriminator(output_ce_224)
    loss_BCE = loss_fn_BCE(output_fake, ones)
    loss_MSE = loss_fn_MSE(output_ce, center_crop_img)
    loss_ce = 0.001 * loss_BCE + 0.999 * loss_MSE
    losses_CE.update(loss_ce.item(), batch_size)
    context_encoder_optimizer.zero_grad()
    loss_ce.backward()
    context_encoder_optimizer.step()

    ### Jigsaw Puzzle ###
    # Update classifier
    _jigsaw_input = jigsaw_input.permute(1, 0, 2, 3, 4) # tile, batch, c, w, h
    feature_jp = list()
    for i in range(num_patch):
        feature_jp.append(HMR(_jigsaw_input[i], return_only_features=True))

    feature_jp.append(HMR(img, return_only_features=True))
    feature_jp = torch.cat(feature_jp, 1)
    output_jp = jigsaw_puzzle_net(feature_jp)
    acc_jp = get_acc(output_jp, jigsaw_order)
    acces_JP.update(acc_jp, batch_size)
    loss_jp = loss_fn_CE(output_jp, jigsaw_order)
    losses_JP.update(loss_jp.item(), batch_size)
    jigsaw_puzzle_optimizer.zero_grad()
    loss_jp.backward()
    jigsaw_puzzle_optimizer.step()

    ### Rotation ###
    # Update rotation net
    feature_rot = HMR(rotation_img, return_only_features=True)
    output_rot = rotation_net(feature_rot)
    acc_rot = get_acc(output_rot, rotation_idx)
    acces_ROT.update(acc_rot, batch_size)
    loss_rot = loss_fn_CE(output_rot, rotation_idx)
    losses_ROT.update(loss_rot, batch_size)
    rotation_optimizer.zero_grad()
    loss_rot.backward()
    rotation_optimizer.step()

    return output_ce_224

def train_hmr_using_3task(
        HMR,
        context_encoder_net,
        discriminator,
        jigsaw_puzzle_net,
        rotation_net,

        loss_fn_BCE,
        loss_fn_MSE,
        loss_fn_CE,

        losses_CE,
        acces_JP,
        losses_JP,
        acces_ROT,
        losses_ROT,
        losses_HMR_3task,

        img,
        context_encoder_input,
        center_crop_img,
        jigsaw_input,
        jigsaw_order,
        rotation_img,
        rotation_idx,

        num_patch,
        ones,
        zeros,
        batch_size,
        args,
    ):
    # loss for HMR - ce
    feature_ce = HMR(context_encoder_input, return_only_features=True)
    feature_ce = feature_ce.reshape(-1, 2048, 1, 1)
    output_ce = context_encoder_net(feature_ce)
    output_ce_224 = context_encoder_input.clone()
    output_ce_224[:, :, 80:144, 80:144] = output_ce.clone()
    output_fake = discriminator(output_ce_224)
    loss_BCE = loss_fn_BCE(output_fake, ones)
    loss_MSE = loss_fn_MSE(output_ce, center_crop_img)
    loss_ce = 0.001 * loss_BCE + 0.999 * loss_MSE
    losses_CE.update(loss_ce.item(), batch_size)

    # loss for HMR - jp
    _jigsaw_input = jigsaw_input.permute(1, 0, 2, 3, 4) # tile, batch, c, w, h
    feature_jp = list()
    for i in range(num_patch):
        feature_jp.append(HMR(_jigsaw_input[i], return_only_features=True))
    feature_jp.append(HMR(img, return_only_features=True))
    feature_jp = torch.cat(feature_jp, 1)
    output_jp = jigsaw_puzzle_net(feature_jp)
    acc_jp = get_acc(output_jp, jigsaw_order)
    acces_JP.update(acc_jp, batch_size)
    loss_jp = loss_fn_CE(output_jp, jigsaw_order)
    losses_JP.update(loss_jp.item(), batch_size)

    # loss for HMR - rot
    feature_rot = HMR(rotation_img, return_only_features=True)
    output_rot = rotation_net(feature_rot)
    acc_rot = get_acc(output_rot, rotation_idx)
    acces_ROT.update(acc_rot, batch_size)
    loss_rot = loss_fn_CE(output_rot, rotation_idx)
    losses_ROT.update(loss_rot, batch_size)

    loss_HMR = args.ce_weight * loss_ce + args.jp_weight * loss_jp + args.rot_weight * loss_rot
    loss_HMR = args.total_weight * loss_HMR
    losses_HMR_3task.update(loss_HMR.item(), batch_size)
    return loss_HMR, output_ce_224

def train_hmr_using_joints(
        HMR,

        loss_fn_keypoints,

        losses_HMR_joints3d,

        img,
        gt_keypoints_2d,
        gt_keypoints_3d,
        has_joints3d,

        joint_mapper_gt,
        batch_size,
        device,
        args,
    ):

    ### training HMR resnet update using joints info
    output = HMR(img)
    output = output[-1]

    ### calcuate loss of 2d joints ###
    pred_keypoints_2d = output["kp_2d"]

    conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone()
    joints2d_loss = (conf * loss_fn_keypoints(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).mean()
    joints2d_loss = 5*joints2d_loss

    ### calcuate loss of 3d joints ###
    pred_keypoints_3d = output["kp_3d"][:, 25:, :]

    conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
    gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
    gt_keypoints_3d = gt_keypoints_3d[has_joints3d==1]
    conf = conf[has_joints3d==1]
    pred_keypoints_3d = pred_keypoints_3d[has_joints3d==1]

    if len(gt_keypoints_3d) > 0:
        gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
        gt_keypoints_3d = gt_keypoints_3d - gt_pelvis[:, None, :]

        pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
        pred_keypoints_3d = pred_keypoints_3d - pred_pelvis[:, None, :]
        joints3d_loss = (conf*loss_fn_keypoints(pred_keypoints_3d, gt_keypoints_3d)).mean()

        pred_j3ds = pred_keypoints_3d[:, joint_mapper_gt, :].clone().detach()
        target_j3ds = gt_keypoints_3d[:, joint_mapper_gt, :].clone().detach()

        errors = torch.sqrt(((pred_j3ds - target_j3ds) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
        S1_hat = batch_compute_similarity_transform_torch(pred_j3ds, target_j3ds)
        errors_pa = torch.sqrt(((S1_hat - target_j3ds) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
        m2mm = 1000

        mpjpe = np.mean(errors) * m2mm
        pa_mpjpe = np.mean(errors_pa) * m2mm
        num_data = len(gt_keypoints_3d)
    else:
        joints3d_loss = torch.FloatTensor(1).fill_(0.).to(device).mean()
        mpjpe = np.array(0)
        pa_mpjpe = np.array(0)
        num_data = len(gt_keypoints_3d)

    total_loss = joints2d_loss

    joints3d_loss = 5*joints3d_loss
    total_loss += joints3d_loss
    losses_HMR_joints3d.update(joints3d_loss.item(), batch_size)

    total_loss *= 60
    return total_loss, mpjpe, pa_mpjpe, num_data

def train_texture_net(
        HMR,
        texture_net,
        img_renderer,

        loss_fn_MSE,
        loss_fn_mask,

        losses_texture_ori_img,
        losses_seg,
        losses_texture_total,

        texture_net_optimizer,

        img,
        black_img,

        batch_size,
        args,
        gt_mask,
        has_mask,
        train_first_stage
    ):
    output = HMR(img)
    output = output[-1]
    vertices = output['verts']
    cam = output['theta'][:, :3]

    textures = texture_net(img)
    textures = textures.expand(-1, -1, 2, 2, 2, -1)

    mask = gt_mask.clone().detach()
    mask_est, rendering = img_renderer(vertices, cam, textures)

    valid_mask = mask > 0
    valid_mask = valid_mask[:, None, :, :].type(torch.int)
    detach_images = img * valid_mask + black_img * (1-valid_mask)

    for i in range(batch_size):
        detach_images[i] = inverse_normalize(detach_images[i])

    #======================================================================================#
    # loss_texture => MSE loss (texture images with bg, original img)
    # loss_seg => MSE loss (segmentation images, target seg)
    # loss_texture_BCE => BCE loss (texture images with bg, Real(1))
    # loss_texture_total => SUM(loss_texture, loss_seg, loss_texture_BCE)
    #======================================================================================#
    loss_texture_ori_img = loss_fn_MSE(detach_images[has_mask==1], rendering[has_mask==1])
    losses_texture_ori_img.update(loss_texture_ori_img.item(), batch_size)
    loss_all = args.rendering_weight*loss_texture_ori_img

    if train_first_stage:
        texture_net_optimizer.zero_grad()
        loss_all.backward()
        texture_net_optimizer.step()
        return mask_est, detach_images, rendering, vertices
    else:
        _mask = mask_est[has_mask == 1]
        _gt_mask = gt_mask[has_mask == 1]

        if len(_mask) != 0 and len(_gt_mask) != 0:
            loss_seg = loss_fn_mask(_mask, _gt_mask)
            losses_seg.update(loss_seg.item(), batch_size)
            loss_all += args.seg_weight * loss_seg

        loss_all = args.texture_total_weight * loss_all
        losses_texture_total.update(loss_all.item(), batch_size)

        return loss_all, mask_est, detach_images, rendering, vertices

def train_hmr_using_adv_loss(
        HMR,
        texture_discriminator,
        texture_net,
        img_renderer,

        losses_disc_e,
        losses_disc,
        losses_disc_real,
        losses_disc_fake,

        img,

        batch_size,
    ):

    output = HMR(img)[-1]
    vertices = output['verts']
    cam = output['theta'][:, :3]

    textures = texture_net(img)
    textures = textures.expand(-1, -1, 2, 2, 2, -1)

    mask, rendering = img_renderer(vertices, cam, textures)

    rendering_img = rendering.clone()

    bg_idx = random.randint(1, 18)
    bg_list = glob.glob("/data/indoor_bg/train/LR/{}/color/*.png".format(bg_idx))

    cropped_bg_list = torch.zeros(batch_size, 3, 224, 224)
    for i in range(batch_size):
        random_idx = random.randint(0, len(bg_list)-1)
        bg_path = bg_list[random_idx]
        bg = cv2.imread(bg_path)
        bg_w, bg_h, _ = bg.shape
        h = w = 224
        rand_idx_w = int(np.random.randint(bg_w-w))
        rand_idx_h = int(np.random.randint(bg_h-h))
        cropped_bg = bg[rand_idx_w:rand_idx_w+w, rand_idx_h:rand_idx_h+h, :]/255.0
        cropped_bg = torch.from_numpy(cropped_bg).permute(2, 0, 1)
        cropped_bg_list[i] = cropped_bg

    cropped_bg_list = cropped_bg_list.to(rendering_img.device)
    valid_mask = mask > 0
    valid_mask = valid_mask[:, None, :, :].type(torch.int)

    rendering_img_input = valid_mask * rendering_img + (1-valid_mask) * cropped_bg_list
    rendering_bg = rendering_img_input.clone().detach()

    for i in range(batch_size):
        rendering_img_input[i] = normalize(rendering_img_input[i])

    e_disc_loss = batch_encoder_disc_l2_loss(texture_discriminator(rendering_img_input))
    losses_disc_e.update(e_disc_loss.item(), batch_size)

    fake_rendering_img_input = rendering_img_input.clone().detach()

    real = texture_discriminator(img)
    fake = texture_discriminator(fake_rendering_img_input)

    d_disc_real, d_disc_fake, d_disc_loss = batch_adv_disc_l2_loss(real, fake)
    losses_disc.update(d_disc_loss.item(), batch_size)
    losses_disc_real.update(d_disc_real.item(), batch_size)
    losses_disc_fake.update(d_disc_fake.item(), batch_size)
    return e_disc_loss, d_disc_loss, rendering_bg

def batch_encoder_disc_l2_loss(disc_value):
    '''
        Inputs:
            disc_value: N x 25
    '''
    k = disc_value.shape[0]
    return torch.sum((disc_value - 1.0) ** 2) * 1.0 / k

def batch_adv_disc_l2_loss(real_disc_value, fake_disc_value):
    '''
        Inputs:
            disc_value: N x 25
    '''
    ka = real_disc_value.shape[0]
    kb = fake_disc_value.shape[0]
    lb, la = torch.sum(fake_disc_value ** 2) / kb, torch.sum((real_disc_value - 1) ** 2) / ka
    return la, lb, la + lb


"""
Test read submittal chain

NOTE: this just makes sure the chain executes properly but DOES NOT assess the quality of the agent's analysis. That is done in the ipython notebooks in the evals/ folder
"""
import pytest
from meche_copilot.schemas import Session
from meche_copilot.chains.read_submittal_chain import ReadSubmittalChain

# Used for verbose output of langchain prompts and responses
# from langchain.callbacks import StdOutCallbackHandler

@pytest.fixture(scope="session")
def read_submittal_chain():
    return ReadSubmittalChain()

@pytest.mark.skip(reason="TODO - write when done")
def test_read_submittal_data(session: Session, read_submittal_chain: ReadSubmittalChain, visualize: bool):
    chain = read_submittal_chain
    chain.read_submittal_data(scoped_eq=session.equipments, show_your_work=visualize)

@pytest.mark.skip(reason="TODO - write when done")
def test_read_submittal_data_with_scoped_eq(session: Session, read_submittal_chain: ReadSubmittalChain, visualize: bool):
    """
    Test can read submittal data for a specific equipment (e.g. pump)
    """
    chain = read_submittal_chain
    pump_eq = session.equipments[0]
    pump_eq.instances[0].design_uid = 'P-1A'
    chain.read_submittal_data(scoped_eq=[pump_eq], show_your_work=visualize)
    res = chain(scoped_eq=session.equipments)


#!/usr/bin/env python3

"""
Autonomy node for the TurtleBot3.

This script relies on a YAML file of potential navigation locations,
which is listed as a `location_file` ROS parameter.

Example usage:
  ros2 run tb3_autonomy autonomy_node.py
  ros2 run tb3_autonomy autonomy_node.py --ros-args -p location_file:=/path/to/my/file.yaml
  ros2 run tb3_autonomy autonomy_node.py --ros-args -p tree_type:=queue -p target_color:=green
"""

import os
import yaml
import random
import rclpy
from rclpy.node import Node
import time
import py_trees
import py_trees_ros
from py_trees.common import OneShotPolicy
from ament_index_python.packages import get_package_share_directory

from tb3_behaviors.navigation import GoToPose, GetLocationFromQueue
from tb3_behaviors.vision import LookForObject

default_location_file = os.path.join(
    get_package_share_directory("tb3_worlds"), "maps", "sim_house_locations.yaml"
)

class AutonomyBehavior(Node):
    def __init__(self):
        super().__init__("autonomy_node")
        self.declare_parameter("location_file", value=default_location_file)
        self.declare_parameter("tree_type", value="queue")
        self.declare_parameter("enable_vision", value=True)
        self.declare_parameter("target_color", value="blue")

        # Parse locations YAML file and shuffle the location list.
        location_file = self.get_parameter("location_file").value
        with open(location_file, "r") as f:
            self.locations = yaml.load(f, Loader=yaml.FullLoader)
        self.loc_list = list(self.locations.keys())
        random.shuffle(self.loc_list)

        # Create and setup the behavior tree
        self.tree_type = self.get_parameter("tree_type").value
        self.enable_vision = self.get_parameter("enable_vision").value
        self.target_color = self.get_parameter("target_color").value
        self.create_behavior_tree(self.tree_type)

        self.tree.node.get_logger().info(f"Using location file: {location_file}")
        self.tree.node.get_logger().info(f"Looking for color {self.target_color}...")

    def create_behavior_tree(self, tree_type):
        if tree_type == "naive":
            self.tree = self.create_naive_tree()
        elif tree_type == "queue":
            self.tree = self.create_queue_tree()
        else:
            self.get_logger().info(f"Invalid behavior tree type {tree_type}.")

    def create_naive_tree(self):
        """Create behavior tree with explicit nodes for each location."""
        if self.enable_vision:
            selector = py_trees.composites.Selector(name="navigation", memory=True)
            root = py_trees.decorators.OneShot(
                name="root",
                child=selector,
                policy=OneShotPolicy.ON_SUCCESSFUL_COMPLETION,
            )
            tree = py_trees_ros.trees.BehaviourTree(root, unicode_tree_debug=False)
            tree.setup(timeout=15.0, node=self)

            for loc in self.loc_list:
                pose = self.locations[loc]
                selector.add_child(
                    py_trees.decorators.OneShot(
                        name=f"try_{loc}",
                        child=py_trees.composites.Sequence(
                            name=f"search_{loc}",
                            children=[
                                GoToPose(f"go_to_{loc}", pose, tree.node),
                                LookForObject(
                                    f"find_{self.target_color}_{loc}",
                                    self.target_color,
                                    tree.node,
                                ),
                            ],
                            memory=True,
                        ),
                        policy=OneShotPolicy.ON_COMPLETION,
                    )
                )

        else:
            seq = py_trees.composites.Sequence(name="navigation", memory=True)
            root = py_trees.decorators.OneShot(
                name="root", child=seq, policy=OneShotPolicy.ON_SUCCESSFUL_COMPLETION
            )
            tree = py_trees_ros.trees.BehaviourTree(root, unicode_tree_debug=False)
            tree.setup(timeout=15.0, node=self)

            for loc in self.loc_list:
                pose = self.locations[loc]
                seq.add_child(GoToPose(f"go_to_{loc}", pose, self))

        return tree

    def create_queue_tree(self):
        """Create behavior tree by picking a next location from a queue"""
        bb = py_trees.blackboard.Blackboard()
        bb.set("loc_list", self.loc_list)

        seq = py_trees.composites.Sequence(name="search", memory=True)
        root = py_trees.decorators.OneShot(
            name="root", child=seq, policy=OneShotPolicy.ON_SUCCESSFUL_COMPLETION
        )
        tree = py_trees_ros.trees.BehaviourTree(root, unicode_tree_debug=False)
        tree.setup(timeout=15.0, node=self)

        seq.add_children(
            [
                GetLocationFromQueue("get_next_location", self.locations),
                GoToPose("go_to_location", None, tree.node),
            ]
        )
        if self.enable_vision:
            seq.add_child(
                LookForObject(f"find_{self.target_color}", self.target_color, tree.node)
            )
        return tree

    def execute(self, period=0.5):
        """Executes the behavior tree at the specified period."""
        self.tree.tick_tock(period_ms=period * 1000.0)
        rclpy.spin(self.tree.node)
        rclpy.shutdown()

if __name__ == "__main__":
    rclpy.init()
    behavior = AutonomyBehavior()
    behavior.execute()


import torch
import numpy as np
import torch.utils.data
import matplotlib.pyplot as plt

n_train = 1000
class DS(torch.utils.data.Dataset):
    def __init__(self, n):
        self.n = n
        self.y = torch.rand(n)*21-10.5
        self.x = torch.sin(0.75*self.y)*7.0+self.y*0.5+torch.randn(n)
    def __len__(self):
        return self.n
    def __getitem__(self,i):
        return (self.x[i],self.y[i])

train_ds = DS(n_train)
plt.scatter(train_ds.x.numpy(),train_ds.y.numpy(), s=2)
plt.show()
# train_dl = torch.utils.data.DataLoader(train_ds, batch_size=batch_size)

class HandwritingGenerator(Module):
    def __init__(self, alphabet_size, hidden_size, num_window_components, num_mixture_components):
        super(HandwritingGenerator, self).__init__()
        self.alphabet_size = alphabet_size + 1
        # First LSTM layer, takes as input a tuple (x, y, eol)
        self.lstm1_layer = LSTM(input_size=3,
                                hidden_size=hidden_size,
                                batch_first=True)

        self.lstm2_layer = LSTM(input_size=3 + hidden_size + alphabet_size + 1,
                                hidden_size=hidden_size,
                                batch_first=True)

        # Third LSTM layer, takes as input the concatenation of the output of the first LSTM layer,
        # the output of the second LSTM layer
        # and the output of the Window layer

        self.lstm3_layer = LSTM(input_size=hidden_size,
                                hidden_size=hidden_size,
                                batch_first=True)

        self.output_layer = mdn(input_size=hidden_size,
                                num_mixtures=num_mixture_components)

        self.hidden1 = None
        self.hidden2 = None
        self.hidden3 = None

        # Initiliaze parameters
        self.reset_parameters()

    def forward(self, strokes, onehot):
        # First LSTM Layer
        input_ = strokes
        output1, self.hidden1 = self.lstm1_layer(input_, self.hidden1)
        # Gaussian Window Layer

        eos, pi, mu1, mu2, sigma1, sigma2, rho = self.output_layer(output1, bias)
        return (eos, pi, mu1, mu2, sigma1, sigma2, rho)