main.py

import os
import pandas as pd
import json
from tqdm import tqdm
import helper_functions
import cruise_creation
import persona_creation
import persona_system_interaction
from FeatureSet import FeatureSet
from PersonaSet import PersonaSet
import csv


demographic_features = {
    "age": {
        "Young Adult (18-24)": 0.2,
        "Early Career (25-34)": 0.2,
        "Midlife Adult (35-44)": 0.2,
        "Established Adult (45-54)": 0.2,
        "Senior Adult (55+)": 0.2
    },
    "race": {
        "White": 0.2,
        "Black or African American": 0.2,
        "American Indian or Alaska Native": 0.2,
        "Asian": 0.2,
        "Hispanic or Latinx": 0.2,
    },
    "sex": {
        "female": 0.5,
        "male": 0.5
    }
}


cruise_features = {
    'mobility': ['none', 'cane', 'wheel-chair', 'limited stairs'],
    'budget_level': ['budget', 'mid-range', 'luxury', 'ultra-luxury'],
    'fitness_interest': ['none', 'stretching', 'regular at gym', 'strength training'],
    'destination_preference': ['caribbean', 'mediterranean', 'europe', 'asia', 'alaska'],
    'travel_group': ['solo', 'couple', 'family', 'friends', 'multi-generation family'],
    'length_in_days': ['2-3', '4-6', '7-10', '10+'],
    'type_of_alcohol_drinker': ['daily', 'often', 'social', 'rarely', 'none'],
    'excursion': ['snorkeling', 'hiking', 'historical tour', 'shopping', 'none'],
    'chronotype': ['night', 'morning']
}


statement_templates = {
    "age": "You belong to the {age} age group.",
    "race": "You identify as {race}.",
    "sex": "You are {sex}.",
    "mobility": "Your mobility level is best described as {mobility}.",
    "budget_level": "You prefer a {budget_level} travel experience.",
    "fitness_interest": "Your fitness interest includes {fitness_interest}.",
    "destination_preference": "You are drawn to destinations like the {destination_preference}.",
    "travel_group": "You usually travel with a {travel_group} group.",
    "length_in_days": "Your ideal trip length is {length_in_days} days.",
    "type_of_alcohol_drinker": "You are a {type_of_alcohol_drinker} alcohol drinker.",
    "excursion": "Your preferred excursion activity is {excursion}.",
    "chronotype": "You are a natural {chronotype} person."
}


def main():
   
    ## GENERATE NARRATIVES TO COMPUTE FREQUENCIES ###
    # Generates all possible pairs of demographic features and queries the LLM to infer assumed preferences for each cruise feature.

    print("Generating narratives to compute frequencies ...\n\n")

    user_feature_set = cruise_creation.generate_personas(demographic_features)
    
    folder_name = "mean_fixed_narratives"
    #our results in included folder: "mean fixed narratives github"
    try:
        os.mkdir(folder_name)
        print(f"Folder '{folder_name}' created successfully.")
    except FileExistsError:
        print(f"Folder '{folder_name}' already exists.")

    directory = "./mean_fixed_narratives"
    filename = '{age}_{race}_{sex}_{i}.txt'

    for user_features in tqdm(user_feature_set):
        for i in range(15):
            file_path = os.path.join(directory, filename.format(i=i, age=user_features['age'][:2], sex=user_features['sex'][:2], race=user_features['race'][:2]))
            if not os.path.isfile(file_path):
                helper_functions.log_to_file('Persona Set', str(user_features), file_path)
                cruise_creation.narrative(
                    user_features,
                    file_path,
                    cruise_creation.prompt,
                    cruise_creation.question)
            else:
                print(f"File {file_path} already exists, skipping...\n")
    



    ## COMPUTE FREQUENCIES ###
    # Calculates how often each cruise option appears within a demographic group, normalizing so that the frequencies across all options in a category sum to one.

    print("Computing demographic conditioned frequencies ...\n\n")

    directory = './mean_fixed_narratives/'
    output_filename = './cond_narrative_fixed_features.csv'
    cruise_creation.process_cond_narratives(directory, output_filename)


    df = pd.read_csv(output_filename)
    out = cruise_creation.get_features_frequency(df, cruise_features, demographic_features)

    with open('./cond_narrative_fixed_frequencies.json', 'w') as file:
        json.dump(out, file, indent=4)



    ### GENERATE PERSONAS ###
    # Generates four persona types based on the computed frequency distributions.


    print("Generating surprising, flat, mean, and random personas ...\n\n")

    with open('./cond_narrative_fixed_frequencies.json', 'r') as file:
        features_cond = json.load(file)

    with open('./possible_cruises.txt', "r", encoding="utf-8") as f:
        cruise_list = [line.strip() for line in f if line.strip()]

    feature_set = FeatureSet(feature_set=features_cond,
                             demographics=demographic_features,
                             statement_templates=statement_templates)
    
    
    personas = persona_creation.optimal_persona_set(feature_set, [], max_set=100, tau_=0.05, lambda_=0.5, iter=4)

    mean_personas = persona_creation.get_mean_personas(personas, feature_set)

    flat_personas = [PersonaSet(p.get_demographics(), {}, feature_set) for p in personas]

    random_personas = persona_creation.get_random_personas(feature_set, 200, seed=42)

    results = []


    ### RUN MULTI-TURN CONVERSATION ###
    # Simulates a multi-turn dialogue between the cruise recommender agent and the generated persona.

    print("Running multi-turn conversation...\n\n")


    folder_name = "recommendation_diversity"
    try:
        os.mkdir(folder_name)
        print(f"Folder '{folder_name}' created successfully.")
    except FileExistsError:
        print(f"Folder '{folder_name}' already exists.")


    for set_name, set in {"surprising":personas, "mean":mean_personas, "flat":flat_personas, "random":random_personas}.items():
        filename = './recommendation_diversity/{type}_{i}.txt'
        for i, persona in tqdm(enumerate(set)):
            persona_system_interaction.match_cruise(persona, cruise_list)
            persona_system_interaction.simulate_self_conversation(
                persona=persona,
                persona_prompt=persona_system_interaction.persona_prompt,
                task_prompt=persona_system_interaction.task_prompt.format(cruise_options=", ".join(cruise_list)),
                initial_message=persona_system_interaction.user_initial_message,
                log_file=filename.format(i=i, type=set_name),
                temp=0.8,
                start_with_persona=True,
                num_turns=10)
            predicted = persona_system_interaction.extract_final_cruise_recommendation(filename.format(i=i, type=set_name))
            results.append({
                "demographics": persona.get_demographics(),
                "features": persona.get_features(),
                "assigned_cruise": persona.get_recommendation(),
                "predicted_cruise": predicted
            })
        with open(f'./recommendation_diversity_{set_name}.csv', 'w', newline='') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=["demographics", "features", "assigned_cruise", "predicted_cruise"])
            writer.writeheader()
            writer.writerows(results)






if __name__ == "__main__":
   main()