import pandas as pd
import re
import sys

def extract_subreddit(url):
    match = re.search(r'reddit\.com/r/([^/]+)/', url)
    return match.group(1) if match else None

def main():
    if len(sys.argv) != 4:
        print("Usage:")
        print("  python jtbd_analysis.py <comments.csv> <jtbd_comments.csv> <relevance.csv>")
        sys.exit(1)

    comments_path = sys.argv[1]
    jtbd_path = sys.argv[2]
    relevance_path = sys.argv[3]

    # Load data
    df_comments = pd.read_csv(comments_path)
    df_jtbd = pd.read_csv(jtbd_path)
    df_relevance = pd.read_csv(relevance_path)

    # Extract subreddit from post_url
    df_comments['subreddit'] = df_comments['post_url'].apply(extract_subreddit)
    df_jtbd['subreddit'] = df_comments.loc[df_jtbd.index, 'subreddit']

    # Filter only jtbd == "YES"
    df_jtbd_yes = df_jtbd[df_jtbd['jtbd'] == 'YES']

    # -------------------------------
    # ✅ ANALIZA 1: JTBD count per subreddit (brut, fără relevanță)
    # -------------------------------
    total_jtbd_per_subreddit = df_jtbd_yes['subreddit'].value_counts()

    total_jtbd_count = total_jtbd_per_subreddit.sum()

    print("=== Total JTBDs per subreddit ===")
    for subreddit, count in total_jtbd_per_subreddit.items():
        percent = (count / total_jtbd_count) * 100
        print(f"- r/{subreddit}: {count} JTBDs ({percent:.2f}%)")
    
    # -------------------------------
    # ✅ ANALIZA 2: JTBD relevance percentage per subreddit
    # -------------------------------
    # Merge cu relevanța doar pe cele relevante
    df_merged = pd.merge(df_jtbd_yes, df_relevance[['situation', 'relevance']], on='situation', how='left')
    df_relevant = df_merged[df_merged['relevance'].isin(['✅ Directly relevant', '⚙️ Partially relevant'])]

    # Count relevant JTBDs per subreddit
    relevant_counts = df_relevant.groupby(['subreddit', 'relevance']).size().unstack(fill_value=0)

    # Ensure all relevant labels exist as columns
    for label in ['✅ Directly relevant', '⚙️ Partially relevant']:
        if label not in relevant_counts.columns:
            relevant_counts[label] = 0

    # Calculate percentages
    result = relevant_counts.div(total_jtbd_per_subreddit, axis=0) * 100
    result = result.fillna(0)

    # Print percentages
    if '✅ Directly relevant' in result:
        print("\n✅ Directly relevant JTBDs (percent per subreddit):")
        for subreddit, value in result['✅ Directly relevant'].items():
            print(f"- r/{subreddit}: {value:.2f}%")

    if '⚙️ Partially relevant' in result:
        print("\n⚙️ Partially relevant JTBDs (percent per subreddit):")
        for subreddit, value in result['⚙️ Partially relevant'].items():
            print(f"- r/{subreddit}: {value:.2f}%")

if __name__ == "__main__":
    main()