from utils import *
from visualization import *
from parameters import *
import time
import ast
import tqdm

# Direct output the full node list

graph = 'sex'
g = graph_read(graph)
task = "IM_S"
if graph in ['facebook', 'sex', 'router']:
    seed_sizes = [10, 12, 15, 18, 20]
else:
    seed_sizes = [5, 6, 7, 8, 9, 10]
population_size = 10

paint_paras = painting_custimization(graph, task)
communities = g.community_fastgreedy().as_clustering()
# clear_folder(f"Figure/IM/{graph}")
# communities = load_community(graph)
# Membership = communities.membership

paint4IM_S = PaintClass(g, graph, communities, paint_paras)
image_file = f"Figure/IM/{graph}_plot.png"
clear_folder(f"result/population/{graph}")

for seed_size in seed_sizes:
    print('The seed size is:', seed_size)
    if graph in ['facebook', 'sex', 'router']:
        prompt_IM_agent1 = "You are an expert in network science and will be provided with one network in the form of an image. The network is divided into different communities and the nodes in the same community are of the same color. \n\nPlease help me intelligently select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. \n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"

        prompt_IM_agent2 = "You are an expert in network science and will be provided with one network in the form of an image. The network is divided into different communities and the nodes in the same community are of the same color. \n\nPlease help me intelligently select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. There are some tips for you: (1) Pick nodes spread across different parts of the network to cover various communities. (2) The number of nodes selected in a community should be proportional to the size of the community. \n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"

        prompt_IM_agent3 = "You are an expert in network science and will be provided with one network in the form of an image. The network is divided into different communities and the nodes in the same community are of the same color. \n\nPlease help me intelligently select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. Here are some tips for you: (1) Pay more attention to the nodes at the central place of the image. (2) Do not select nodes in the community near the edge of the image. \n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"

        prompt_IM_agent4 = "You are an expert in network science and will be provided with one network in the form of an image. The network is divided into different communities and the nodes in the same community are of the same color. \n\nPlease help me intelligently select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. Here are some tips for you: (1) Pay more attention to the nodes in the large communities. (2) Do not select nodes in the small communities. \n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"
    else:

        prompt_IM_agent1 = "You are an expert in network science and will be provided with one network in the form of an image.  \n\nPlease help me intelligently select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. \n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"

        prompt_IM_agent2 = "You are an expert in network science and will be provided with one network in the form of an image.  \n\nPlease help me wisely select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. Here are some tips for you: (1) Choose large-degree nodes. (2) Pick nodes spread across different parts of the network.\n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"

        prompt_IM_agent3 = "You are an expert in network science and will be provided with one network in the form of an image.  \n\nPlease help me wisely select " + str(
            seed_size) + " nodes as seeds in this network to achieve influence maximization. Here are some tips for you: (1) Choose large-degree nodes. (2) Pick nodes at the center place of the network.\n\nDo NOT output any other text or explanation. Just tell me the node IDs only. Your answer should be only a list as\n\n[node_id, ..., node_id]"

        #
    # Define the prompts for each agent
    if graph in ['facebook', 'sex', 'router']:
        prompts = [prompt_IM_agent1, prompt_IM_agent2, prompt_IM_agent3, prompt_IM_agent4]
    else:
        prompts = [prompt_IM_agent1, prompt_IM_agent2, prompt_IM_agent3]
    num_agents = len(prompts)

    populations = [[] for _ in range(num_agents)]


    # Create population paths dynamically
    population_paths = [
        f"result/population/{graph}/{graph}_agent{i}_{seed_size}" for i in range(1, num_agents + 1)
    ]

    # Process each population size
    for i in range(population_size):
        for j in range(num_agents):  # Iterate over each agent
            result = invoke_with_image(query=prompts[j], image_file=image_file)
            suggested_node = ast.literal_eval(result)
            populations[j].append(suggested_node)

    # Save each population to a file
    for idx, population in enumerate(populations):
        with open(population_paths[idx], 'w') as file:
            json.dump(population, file)
