analysis_solution.py

import os
import pandas as pd
import time
import google.generativeai as genai

# ---------------------------
# Setup Gemini
# ---------------------------
try:
    with open("api_keys/gemini_key.txt", "r") as f:
        gemini_key = f.read().strip()
    genai.configure(api_key=gemini_key)
    print("Gemini API configured.")
except Exception:
    print("Error: Gemini API key not found.")
    exit()

model = genai.GenerativeModel(model_name="gemini-2.5-pro")

# ---------------------------
# Prompt Template
# ---------------------------
def build_prompt(urls):
    return f"""
You are classifying solution types based on URLs.

For each URL, classify into ONE of:
- "market_product": a commercially sold product
- "diy": something that requires building or assembling by the user
- "life_hack": a simple trick or informal workaround

Return ONLY JSON in this format:
[
  {{"url": "...", "category": "..."}},
  ...
]

URLs:
{urls}
"""

# ---------------------------
# Gemini Call
# ---------------------------
def classify_urls(url_batch):
    prompt = build_prompt(url_batch)

    try:
        response = model.generate_content(
            prompt,
            generation_config=genai.types.GenerationConfig(
                response_mime_type="application/json"
            ),
        )
        return response.text
    except Exception as e:
        print("Gemini error:", e)
        return None

# ---------------------------
# Process One File
# ---------------------------
def process_file(filepath, batch_size=10):
    df = pd.read_csv(filepath)

    if "page_url" not in df.columns:
        print(f"Skipping {filepath}: no page_url column")
        return

    urls = df["page_url"].dropna().unique().tolist()

    results = []

    # Batch processing (important for rate limits + cost)
    for i in range(0, len(urls), batch_size):
        batch = urls[i:i + batch_size]
        print(f"Processing batch {i} - {i + len(batch)}")

        response_text = classify_urls(batch)

        if response_text:
            try:
                parsed = pd.read_json(response_text)
                results.append(parsed)
            except Exception as e:
                print("JSON parse error:", e)
                print(response_text)

        time.sleep(2)  # avoid rate limits

    if results:
        result_df = pd.concat(results, ignore_index=True)

        # Merge back to original dataframe
        df = df.merge(result_df, left_on="page_url", right_on="url", how="left")
        df.drop(columns=["url"], inplace=True)

        output_path = filepath.replace("_rated.csv", "_classified.csv")
        df.to_csv(output_path, index=False)
        print(f"Saved: {output_path}")

# ---------------------------
# Process All Participants
# ---------------------------
def process_all(folder_path):
    for file in os.listdir(folder_path):
        if file.endswith("_rated.csv"):
            filepath = os.path.join(folder_path, file)
            print(f"\nProcessing {file}")
            process_file(filepath)
            

def plot_solution_type_distribution(folder_path, save_dir="figures"):
    import os
    import pandas as pd
    import seaborn as sns
    import matplotlib.pyplot as plt

    os.makedirs(save_dir, exist_ok=True)

    # Initialize counts
    category_counts = {
        "market_product": 0,
        "diy": 0,
        "life_hack": 0
    }

    # Aggregate across all classified CSVs
    for file in os.listdir(folder_path):
        if file.endswith("_classified.csv"):
            filepath = os.path.join(folder_path, file)
            df = pd.read_csv(filepath)

            if "category" not in df.columns:
                print(f"Skipping {file}: no 'category' column")
                continue

            counts = df["category"].value_counts()
            for category in category_counts:
                if category in counts:
                    category_counts[category] += counts[category]

    # Convert to DataFrame
    summary_df = pd.DataFrame(
        list(category_counts.items()),
        columns=["category", "count"]
    )

    # Optional: nicer labels for paper
    label_map = {
        "market_product": "Market Product",
        "diy": "DIY",
        "life_hack": "Life Hack"
    }
    summary_df["category"] = summary_df["category"].map(label_map)

    print("\nSolution type distribution:")
    print(summary_df.to_string(index=False))

    # 🔥 Match your seaborn style
    sns.set(style="whitegrid", context="talk")

    plt.figure(figsize=(6,5))
    ax = sns.barplot(
        data=summary_df,
        x="category",
        y="count",
        palette="Set2"
    )

    # ✅ Add numbers on bars (same style as your code)
    for p in ax.patches:
        height = p.get_height()
        ax.text(
            p.get_x() + p.get_width()/2,
            height + 0.1,
            int(height),
            ha="center",
            va="bottom",
            fontsize=12
        )

    plt.xlabel("Solution Type")
    plt.ylabel("Number of Solutions")
    plt.title("Distribution of Solution Types")
    plt.tight_layout()

    # Save
    save_path = os.path.join(save_dir, "solution_type_distribution.pdf")
    plt.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.show()
    plt.close()

    print(f"Plot saved to {save_path}")
    


def plot_solution_type_grouped(folder_path, save_dir="figures"):
    import os
    import pandas as pd
    import matplotlib.pyplot as plt
    import seaborn as sns
    import numpy as np

    os.makedirs(save_dir, exist_ok=True)

    categories = ["market_product", "diy", "life_hack"]
    counts = {
        cat: {"accepted": 0, "neutral": 0, "rejected": 0}
        for cat in categories
    }

    # Aggregate data
    for file in os.listdir(folder_path):
        if file.endswith("_classified.csv"):
            filepath = os.path.join(folder_path, file)
            df = pd.read_csv(filepath)

            # detect willingness column (robust)
            willingness_col = None
            for col in df.columns:
                if "wilingness" in col.lower() or "willingness" in col.lower():
                    willingness_col = col
                    break

            if "category" not in df.columns or willingness_col is None:
                print(f"Skipping {file}")
                continue

            for _, row in df.iterrows():
                cat = str(row["category"]).strip().lower()
                score = row[willingness_col]

                if cat not in counts or pd.isna(score):
                    continue

                try:
                    score = int(score)
                except:
                    continue

                if score in [4, 5]:
                    counts[cat]["accepted"] += 1
                elif score == 3:
                    counts[cat]["neutral"] += 1
                elif score in [1, 2]:
                    counts[cat]["rejected"] += 1

    df_plot = pd.DataFrame(counts).T.fillna(0)

    label_map = {
        "market_product": "Market Product",
        "diy": "DIY",
        "life_hack": "Life Hack"
    }
    df_plot.index = df_plot.index.map(label_map)

    # Compute totals
    totals = df_plot.sum(axis=1)

    print("\nGrouped counts:")
    print(df_plot)

    # Plot
    sns.set(style="whitegrid", context="talk")
    plt.figure(figsize=(9, 5))

    x = np.arange(len(df_plot.index))
    width = 0.25

    bars1 = plt.bar(x - width, df_plot["accepted"], width, label="Accepted", color="green")
    bars2 = plt.bar(x, df_plot["neutral"], width, label="Neutral", color="orange")
    bars3 = plt.bar(x + width, df_plot["rejected"], width, label="Rejected", color="red")

    # ✅ Add value labels on each bar
    def add_labels(bars):
        for bar in bars:
            height = bar.get_height()
            if height > 0:
                plt.text(
                    bar.get_x() + bar.get_width() / 2,
                    height + 0.5,
                    f"{int(height)}",
                    ha="center",
                    va="bottom",
                    fontsize=11
                )

    add_labels(bars1)
    add_labels(bars2)
    add_labels(bars3)

    # ✅ Add totals to x-axis labels
    x_labels = [
        f"{cat}\n(N={int(total)})"
        for cat, total in zip(df_plot.index, totals)
    ]

    plt.xticks(x, x_labels)

    plt.ylabel("Count")
    plt.xlabel("Solution Type")
    plt.title("User Willingness by Solution Type")
    plt.legend()

    plt.tight_layout()

    save_path = os.path.join(save_dir, "solution_type_grouped.pdf")
    plt.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.show()
    plt.close()

    print(f"Plot saved to {save_path}")
    


def plot_solution_type_ratio_grouped(folder_path, save_dir="figures"):
    import os
    import pandas as pd
    import matplotlib.pyplot as plt
    import seaborn as sns
    import numpy as np

    os.makedirs(save_dir, exist_ok=True)

    categories = ["market_product", "diy", "life_hack"]
    counts = {
        cat: {"accepted": 0, "neutral": 0, "rejected": 0}
        for cat in categories
    }

    # Aggregate data
    for file in os.listdir(folder_path):
        if file.endswith("_classified.csv"):
            filepath = os.path.join(folder_path, file)
            df = pd.read_csv(filepath)

            # detect willingness column (robust)
            willingness_col = None
            for col in df.columns:
                if "wilingness" in col.lower() or "willingness" in col.lower():
                    willingness_col = col
                    break

            if "category" not in df.columns or willingness_col is None:
                print(f"Skipping {file}")
                continue

            for _, row in df.iterrows():
                cat = str(row["category"]).strip().lower()
                score = row[willingness_col]

                if cat not in counts or pd.isna(score):
                    continue

                try:
                    score = int(score)
                except:
                    continue

                if score in [4, 5]:
                    counts[cat]["accepted"] += 1
                elif score == 3:
                    counts[cat]["neutral"] += 1
                elif score in [1, 2]:
                    counts[cat]["rejected"] += 1

    df_plot = pd.DataFrame(counts).T.fillna(0)

    # Compute ratios
    df_ratio = df_plot.div(df_plot.sum(axis=1), axis=0) * 100

    label_map = {
        "market_product": "Market Product",
        "diy": "DIY",
        "life_hack": "Life Hack"
    }
    df_ratio.index = df_ratio.index.map(label_map)

    print("\nGrouped ratios (%):")
    print(df_ratio.round(1))

    # Plot
    sns.set(style="whitegrid", context="talk")
    plt.figure(figsize=(9, 5))

    x = np.arange(len(df_ratio.index))
    width = 0.25

    bars1 = plt.bar(x - width, df_ratio["accepted"], width, label="Accepted", color="green")
    bars2 = plt.bar(x, df_ratio["neutral"], width, label="Neutral", color="orange")
    bars3 = plt.bar(x + width, df_ratio["rejected"], width, label="Rejected", color="red")

    # Add value labels (as percentages)
    def add_labels(bars):
        for bar in bars:
            height = bar.get_height()
            if height > 0:
                plt.text(
                    bar.get_x() + bar.get_width() / 2,
                    height + 1,
                    f"{height:.0f}%",
                    ha="center",
                    va="bottom",
                    fontsize=11
                )

    add_labels(bars1)
    add_labels(bars2)
    add_labels(bars3)

    # Add totals to x-axis labels (optional)
    totals = df_plot.sum(axis=1)
    x_labels = [
        f"{cat}\n(N={int(total)})"
        for cat, total in zip(df_ratio.index, totals)
    ]
    plt.xticks(x, x_labels)

    plt.ylabel("Percentage (%)")
    plt.xlabel("Solution Type")
    plt.title("User Willingness Ratios by Solution Type")
    plt.legend()

    plt.tight_layout()

    save_path = os.path.join(save_dir, "solution_type_ratio_grouped.pdf")
    plt.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.show()
    plt.close()

    print(f"Plot saved to {save_path}")
    
# ---------------------------
# Run
# ---------------------------

if __name__ == "__main__":
    # process_all("/Users/kwon/owa_dataset/result/rated")
    # plot_solution_type_distribution("/Users/kwon/owa_dataset/result/rated")
    # plot_solution_type_grouped("/Users/kwon/owa_dataset/result/classified")
    plot_solution_type_ratio_grouped("/Users/kwon/owa_dataset/result/classified")