data.py

from pathlib import Path
from json import load as load_json
from numpy.random import permutation as perm

data_dir = Path(__file__).parent / "data" / "cmu" / "processed"


def get_cities():
    with open(data_dir / "cities.json", "r") as f:
        return load_json(f)


def get_score_threshold_per_city():
    with open(data_dir / "score_threshold_per_city.json", "r") as f:
        return load_json(f)


def get_city_to_hotel_id_map():
    with open(data_dir / "city_to_hotel_id_map.json", "r") as f:
        return load_json(f)


def get_hotel_id_to_name_map():
    with open(data_dir / "hotel_id_to_name_map.json", "r") as f:
        return load_json(f)


def get_hotel_id_to_review_map():
    with open(data_dir / "hotel_id_to_review_map.json", "r") as f:
        return load_json(f)


score_threshold_per_city = get_score_threshold_per_city()
city_to_hotel_id_map = get_city_to_hotel_id_map()
hotel_id_to_name_map = get_hotel_id_to_name_map()
hotel_id_to_review_map = get_hotel_id_to_review_map()


def get_reviews_for_prompt(city, preferences) -> dict:

    for hotel_id in perm(city_to_hotel_id_map[city]):
        hotel_id = str(hotel_id)
        res = {"hotel_name": hotel_id_to_name_map[hotel_id], 'positive': [], 'negative': []}
        try:
            hotel_reviews = hotel_id_to_review_map[hotel_id]['reviews']
        except KeyError:
            continue
        for review in perm(hotel_reviews):
            if (review['score'] == 5) & (len(res['positive']) < 3):
                res['positive'].append(review)
            if (review['score'] <= 2) & (len(res['negative']) < 1):
                res['negative'].append(review)
            if (len(res['positive']) >= 3) & (len(res['negative']) >= 1):
                return res

    return None