import csv
import re

# Function to read the transformed CSV and return a dictionary of email to topics
def read_transformed_csv(file_path):
    data = {}
    with open(file_path, mode="r", encoding="utf-8") as csvfile:
        reader = csv.reader(csvfile)
        header = next(reader)  # Skip the header row
        emails = header[1:]  # Extract email addresses from the header
        for row in reader:
            topic = row[0]
            for i, cell in enumerate(row[1:]):
                if cell == "X":
                    if emails[i] not in data:
                        data[emails[i]] = set()
                    data[emails[i]].add(topic)
    return data

# Function to clean topics (remove text in parentheses)
def clean_topic(topic):
    return re.sub(r"\s*\(.*?\)", "", topic).strip()


def find_best_pairings(mentors, mentees):
    pairings = []
    used_mentors = {}
    used_mentees = set()

    # Mentors allowed to have up to 2 matches
    multi_match_mentors = set()

    # Prioritize specific topics
    priority_topics = {"Generative AI (for Software Engineering)", "Generative AI (for Data Science)"}
    priority_topics_cleaned = {clean_topic(topic) for topic in priority_topics}

    # Create a list of all possible mentor-mentee pairs with their overlap count
    all_pairs = []
    for mentor_email, mentor_topics in mentors.items():
        for mentee_email, mentee_topics in mentees.items():
            overlap_topics = mentor_topics & mentee_topics
            overlap_count = len(overlap_topics)
            if overlap_count > 0:
                all_pairs.append((mentor_email, mentee_email, overlap_count, overlap_topics))

    # Sort all pairs by overlap count in descending order
    all_pairs.sort(key=lambda x: x[2], reverse=True)

    # Greedily select pairs for priority topics first
    for mentor_email, mentee_email, overlap_count, overlap_topics in all_pairs:
        if mentee_email not in used_mentees:
            if mentor_email not in used_mentors:
                used_mentors[mentor_email] = 0

            # Allow multi-match mentors to have up to 2 matches
            if mentor_email in multi_match_mentors and used_mentors[mentor_email] < 2:
                if overlap_topics & priority_topics_cleaned:
                    pairings.append((mentor_email, mentee_email, overlap_topics))
                    used_mentors[mentor_email] += 1
                    used_mentees.add(mentee_email)
            elif mentor_email not in multi_match_mentors and used_mentors[mentor_email] == 0:
                if overlap_topics & priority_topics_cleaned:
                    pairings.append((mentor_email, mentee_email, overlap_topics))
                    used_mentors[mentor_email] += 1
                    used_mentees.add(mentee_email)

    # Greedily select the remaining best pairs
    for mentor_email, mentee_email, overlap_count, overlap_topics in all_pairs:
        if mentee_email not in used_mentees:
            if mentor_email not in used_mentors:
                used_mentors[mentor_email] = 0

            # Allow multi-match mentors to have up to 2 matches
            if mentor_email in multi_match_mentors and used_mentors[mentor_email] < 2:
                pairings.append((mentor_email, mentee_email, overlap_topics))
                used_mentors[mentor_email] += 1
                used_mentees.add(mentee_email)
            elif mentor_email not in multi_match_mentors and used_mentors[mentor_email] == 0:
                pairings.append((mentor_email, mentee_email, overlap_topics))
                used_mentors[mentor_email] += 1
                used_mentees.add(mentee_email)

    return pairings

def write_pairings_to_csv(pairings, output_file):
    with open(output_file, mode="w", encoding="utf-8", newline="") as csvfile:
        writer = csv.writer(csvfile)
        # Write the header row
        writer.writerow(["Mentor Email", "Mentee Email", "Overlapped Topics"])
        # Write the pairings
        for mentor, mentee, overlapped_topics in pairings:
            writer.writerow([mentor, mentee, ", ".join(overlapped_topics)])

# Main function
def main():
    mentor_file = "mentor_result_transformed.csv"
    mentee_file = "mentee_result_transformed.csv"
    output_file = "mentor_mentee_pairings.csv"

    # Read mentor and mentee data
    mentors = read_transformed_csv(mentor_file)
    mentees = read_transformed_csv(mentee_file)

    # Find the best pairings
    pairings = find_best_pairings(mentors, mentees)

    # Write the pairings to a CSV file
    write_pairings_to_csv(pairings, output_file)
    print(f"Mentor-Mentee Pairings have been saved to {output_file}")

if __name__ == "__main__":
    main()